fix: bound model stream aggregation and close failed responses
This commit is contained in:
1 parent
4879adc380
commit
7c0c2b59ca
13 files changed
+1315
-12
No files matched your search
@@ -33,7 +33,7 @@ def test_invalid_utf8_is_marked_without_failing_or_forging_a_digest():
|
||||
fingerprint = ArgumentFingerprint()
|
||||
for value in ["safe", "\ud800", "later"]:
|
||||
fingerprint.observe(value)
|
||||
assert fingerprint.snapshot() == {
|
||||
expected = {
|
||||
"argument_chars": 10,
|
||||
"argument_bytes": None,
|
||||
"argument_sha256": None,
|
||||
@@ -41,6 +41,24 @@ def test_invalid_utf8_is_marked_without_failing_or_forging_a_digest():
|
||||
"argument_prefixes": [],
|
||||
"argument_encoding_valid": False,
|
||||
}
|
||||
assert {key: fingerprint.snapshot()[key] for key in expected} == expected
|
||||
|
||||
|
||||
def test_shape_counts_cross_delta_whitespace_and_repetition_without_retaining_text():
|
||||
fingerprint = ArgumentFingerprint()
|
||||
for fragment in ['{"text":"private', " ", " \t", "", " \t", " \t", '\u00a0"}']:
|
||||
fingerprint.observe(fragment)
|
||||
stats = fingerprint.snapshot()
|
||||
assert stats["argument_json_whitespace_chars"] == 7
|
||||
assert stats["argument_json_whitespace_run_max"] == 7
|
||||
assert stats["argument_json_whitespace_run"] == 0
|
||||
assert stats["argument_delta_chars_min"] == 1
|
||||
assert stats["argument_delta_chars_max"] == len('{"text":"private')
|
||||
assert stats["argument_same_delta_run_max"] == 3
|
||||
assert stats["argument_same_delta_run"] == 1
|
||||
assert stats["argument_deltas"] == 6
|
||||
assert "private" not in repr(vars(fingerprint))
|
||||
assert "private" not in json.dumps(stats)
|
||||
|
||||
|
||||
def test_tool_tracking_is_bounded_but_total_count_keeps_all_received_arguments():
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
import itertools
|
||||
import json
|
||||
|
||||
from langchain_core.messages import AIMessageChunk
|
||||
import langchain_core.messages.ai as ai
|
||||
import pytest
|
||||
|
||||
from artemis.llm.chunk_accumulator import ChunkAccumulator
|
||||
from artemis.llm.response_guard import ResponseIntegrityError, ensure_response_complete
|
||||
|
||||
|
||||
def test_batching_preserves_interleaved_tools_unicode_metadata_and_usage():
|
||||
arguments = [json.dumps({"target": "中文🙂" * 80}), json.dumps({"target": 3})]
|
||||
chunks = [AIMessageChunk(content="Live text", id="resp_one")]
|
||||
for index in range(2):
|
||||
chunks.append(
|
||||
AIMessageChunk(
|
||||
content=[],
|
||||
tool_call_chunks=[
|
||||
{"name": "click", "args": "", "id": f"call_{index}", "index": index}
|
||||
],
|
||||
)
|
||||
)
|
||||
for pair in itertools.zip_longest(*arguments):
|
||||
for index, value in enumerate(pair):
|
||||
if value is not None:
|
||||
chunks.append(
|
||||
AIMessageChunk(content=[], tool_call_chunks=[{"args": value, "index": index}])
|
||||
)
|
||||
chunks.append(
|
||||
AIMessageChunk(
|
||||
content=[],
|
||||
response_metadata={"id": "resp_one", "status": "completed"},
|
||||
usage_metadata={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
|
||||
chunk_position="last",
|
||||
)
|
||||
)
|
||||
expected = chunks[0]
|
||||
accumulator = ChunkAccumulator()
|
||||
max_pending = 0
|
||||
for index, chunk in enumerate(chunks):
|
||||
if index:
|
||||
expected = expected + chunk
|
||||
accumulator.append(chunk)
|
||||
if accumulator.ready:
|
||||
accumulator.flush()
|
||||
max_pending = max(max_pending, len(accumulator.pending))
|
||||
result = accumulator.flush()
|
||||
assert result.model_dump() == expected.model_dump()
|
||||
assert [t["args"] for t in result.tool_calls] == [json.loads(a) for a in arguments]
|
||||
assert result.usage_metadata["total_tokens"] == 15
|
||||
assert max_pending < 256
|
||||
assert not accumulator.pending
|
||||
|
||||
|
||||
def test_partial_json_is_not_reparsed_for_every_single_character(monkeypatch):
|
||||
chunks = [
|
||||
AIMessageChunk(
|
||||
content=[], tool_call_chunks=[{"name": "click", "args": '{"target":x', "index": 0}]
|
||||
)
|
||||
]
|
||||
chunks.extend(
|
||||
AIMessageChunk(content=[], tool_call_chunks=[{"args": "x", "index": 0}]) for _ in range(800)
|
||||
)
|
||||
original = ai.parse_partial_json
|
||||
parsed_lengths = []
|
||||
|
||||
def counted(arguments, *args, **kwargs):
|
||||
parsed_lengths.append(len(arguments))
|
||||
return original(arguments, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(ai, "parse_partial_json", counted)
|
||||
accumulator = ChunkAccumulator()
|
||||
for chunk in chunks:
|
||||
accumulator.append(chunk)
|
||||
if accumulator.ready:
|
||||
accumulator.flush()
|
||||
result = accumulator.flush()
|
||||
assert result.tool_call_chunks[0]["args"] == '{"target":x' + "x" * 800
|
||||
assert len(parsed_lengths) < 20
|
||||
result.response_metadata["status"] = "completed"
|
||||
with pytest.raises(ResponseIntegrityError):
|
||||
ensure_response_complete(result)
|
||||
@@ -226,7 +226,7 @@ def test_tool_deltas_are_counted_once_without_retaining_arguments():
|
||||
[2, 1],
|
||||
strict=True,
|
||||
):
|
||||
assert tool == {
|
||||
expected = {
|
||||
"name": name,
|
||||
"call_id": call_id,
|
||||
"chunk_index": index,
|
||||
@@ -237,6 +237,7 @@ def test_tool_deltas_are_counted_once_without_retaining_arguments():
|
||||
"argument_prefixes": [],
|
||||
"argument_encoding_valid": True,
|
||||
}
|
||||
assert {key: tool[key] for key in expected} == expected
|
||||
assert attempt.data["response_status"] == "completed"
|
||||
assert "private" not in json.dumps(engine.records, default=str)
|
||||
assert "secret" not in json.dumps(engine.records, default=str)
|
||||
|
||||
@@ -0,0 +1,263 @@
|
||||
import json
|
||||
import gzip
|
||||
from types import SimpleNamespace
|
||||
|
||||
import httpx2
|
||||
import pytest
|
||||
|
||||
from artemis.llm.responses_diagnostics import MAX_FRAME_BYTES, ResponsesObservation
|
||||
from artemis.llm.diagnostics import ACTIVE_ATTEMPT, observe_response_headers
|
||||
from artemis.llm.tool_constraints import MAX_REQUEST_BYTES, request_constraints, tool_constraints
|
||||
|
||||
|
||||
def frame(kind, sequence, **fields):
|
||||
return (
|
||||
"data: "
|
||||
+ json.dumps({"type": kind, "sequence_number": sequence, **fields}, ensure_ascii=False)
|
||||
+ "\r\n\r\n"
|
||||
).encode()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("size", [1, 37, 100000])
|
||||
def test_raw_frames_split_unicode_and_done_snapshots_are_not_double_counted(size):
|
||||
arguments = '{"text":"私人🙂"}'
|
||||
data = b"".join(
|
||||
[
|
||||
frame("response.created", 0, response={"id": "resp_one"}),
|
||||
frame(
|
||||
"response.output_item.added",
|
||||
1,
|
||||
output_index=2,
|
||||
item={
|
||||
"type": "function_call",
|
||||
"id": "fc_one",
|
||||
"call_id": "call_one",
|
||||
"name": "click",
|
||||
"arguments": "",
|
||||
},
|
||||
),
|
||||
frame(
|
||||
"response.function_call_arguments.delta",
|
||||
2,
|
||||
output_index=2,
|
||||
item_id="fc_one",
|
||||
delta=arguments,
|
||||
),
|
||||
frame(
|
||||
"response.function_call_arguments.done",
|
||||
3,
|
||||
output_index=2,
|
||||
item_id="fc_one",
|
||||
arguments=arguments,
|
||||
),
|
||||
frame(
|
||||
"response.output_item.done",
|
||||
4,
|
||||
output_index=2,
|
||||
item={"type": "function_call", "id": "fc_one", "arguments": arguments},
|
||||
),
|
||||
frame(
|
||||
"response.completed",
|
||||
5,
|
||||
response={
|
||||
"id": "resp_one",
|
||||
"tools": [{"type": "function", "name": "click", "strict": True}],
|
||||
},
|
||||
),
|
||||
]
|
||||
)
|
||||
observer = ResponsesObservation()
|
||||
for offset in range(0, len(data), size):
|
||||
observer.feed(data[offset : offset + size])
|
||||
observer.end(eof=True)
|
||||
stats = observer.snapshot()
|
||||
assert stats["body_bytes"] == len(data)
|
||||
assert stats["parse_errors"] == 0
|
||||
assert stats["incomplete_sse_frame"] is False
|
||||
(tool,) = stats["tools"]
|
||||
assert tool["argument_chars"] == len(arguments)
|
||||
assert (
|
||||
tool["delta_events"]
|
||||
== tool["arguments_done_events"]
|
||||
== tool["output_item_done_events"]
|
||||
== 1
|
||||
)
|
||||
assert tool["call_id"] == "call_one"
|
||||
(response,) = stats["responses"]
|
||||
assert response["sequence_first"] == 0 and response["sequence_last"] == 5
|
||||
assert response["sequence_forward_gaps"] == response["sequence_non_increasing"] == 0
|
||||
assert response["returned_tool_constraints"]["tools"][0]["strict"] == "true"
|
||||
assert "私人" not in json.dumps(stats, ensure_ascii=False)
|
||||
assert not observer.line and not observer.lines
|
||||
|
||||
|
||||
def test_sequence_statistics_are_per_response_and_describe_observed_jumps():
|
||||
observer = ResponsesObservation()
|
||||
observer.event({"type": "response.in_progress", "sequence_number": 1})
|
||||
observer.event({"type": "response.created", "sequence_number": 0, "response": {"id": "resp_a"}})
|
||||
for seq in [2, 2, 1, 3, None]:
|
||||
observer.event({"type": "response.in_progress", "sequence_number": seq})
|
||||
observer.event({"type": "response.created", "sequence_number": 0, "response": {"id": "resp_b"}})
|
||||
a, b = observer.snapshot()["responses"]
|
||||
assert a["sequence_forward_gaps"] == a["sequence_forward_missing_total"] == 2
|
||||
assert a["sequence_non_increasing"] == 2
|
||||
assert a["sequence_equal_previous"] == 1
|
||||
assert a["sequence_missing_events"] == 1
|
||||
assert b["sequence_non_increasing"] == b["sequence_forward_gaps"] == 0
|
||||
assert observer.data["unknown_response_events"] == 1
|
||||
|
||||
|
||||
def test_oversized_frame_recovery_and_partial_eof_keep_memory_bounded():
|
||||
observer = ResponsesObservation()
|
||||
observer.feed(b"data: " + b"x" * MAX_FRAME_BYTES)
|
||||
observer.feed(b"\r\n\r\n")
|
||||
observer.feed(b"data: invalid\n\n")
|
||||
observer.feed(frame("response.created", 0, response={"id": "resp_valid"}))
|
||||
observer.feed(b'data: {"private":')
|
||||
observer.end(eof=True)
|
||||
stats = observer.snapshot()
|
||||
assert stats["oversized_frames"] == 1
|
||||
assert stats["parse_errors"] == 1
|
||||
assert len(stats["responses"]) == 1
|
||||
assert stats["incomplete_sse_frame"] is True
|
||||
assert "private" not in repr(vars(observer))
|
||||
assert not observer.line and not observer.lines
|
||||
observer.end() # HTTP close can follow EOF; retain the original coverage gap.
|
||||
assert observer.data["incomplete_sse_frame"] is True
|
||||
|
||||
|
||||
def test_raw_tool_and_response_statistics_have_fixed_caps():
|
||||
observer = ResponsesObservation()
|
||||
observer.event({"type": "response.created", "response": {"id": "resp_one"}})
|
||||
for index in range(100):
|
||||
observer.event(
|
||||
{
|
||||
"type": "response.function_call_arguments.delta",
|
||||
"item_id": f"fc_{index}",
|
||||
"output_index": index,
|
||||
"delta": " ",
|
||||
}
|
||||
)
|
||||
for index in range(20):
|
||||
observer.event({"type": "response.created", "response": {"id": f"resp_{index}"}})
|
||||
assert len(observer.tools) == len(observer.fingerprints) == len(observer.aliases) == 64
|
||||
assert len(observer.responses) == 16
|
||||
assert observer.data["statistics_truncated"] is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("size", [1, 37, MAX_FRAME_BYTES * 2])
|
||||
def test_oversized_crlf_frame_recovers_across_arbitrary_read_boundaries(size):
|
||||
observer = ResponsesObservation()
|
||||
data = b"data: " + b"x" * MAX_FRAME_BYTES + b"\r\n\r\n"
|
||||
data += frame("response.created", 0, response={"id": "resp_after_large"})
|
||||
for offset in range(0, len(data), size):
|
||||
observer.feed(data[offset : offset + size])
|
||||
assert observer.data["oversized_frames"] == 1
|
||||
assert observer.data["parse_errors"] == 0
|
||||
assert list(observer.responses) == ["resp_after_large"]
|
||||
assert not observer.line and not observer.lines and not observer.dropping
|
||||
|
||||
|
||||
def test_alias_statistics_are_bounded_even_when_item_keeps_changing_index():
|
||||
observer = ResponsesObservation()
|
||||
observer.event({"response": {"id": "resp_one"}})
|
||||
for index in range(1000):
|
||||
observer.event(
|
||||
{
|
||||
"type": "response.function_call_arguments.delta",
|
||||
"item_id": "fc_same",
|
||||
"output_index": index,
|
||||
"delta": " ",
|
||||
}
|
||||
)
|
||||
assert len(observer.tools) == 1
|
||||
assert len(observer.aliases) == 128
|
||||
assert observer.data["statistics_truncated"] is True
|
||||
|
||||
|
||||
def test_malformed_event_type_cannot_break_passive_observation():
|
||||
observer = ResponsesObservation()
|
||||
observer.feed(frame([], 0, response={"id": "resp_one"}))
|
||||
assert observer.snapshot()["responses"][0]["events"] == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_compressed_body_is_explicitly_skipped_and_sdk_decoding_is_unchanged():
|
||||
body = frame("response.created", 0, response={"id": "resp_one"})
|
||||
compressed = gzip.compress(body)
|
||||
|
||||
class Stream(httpx2.AsyncByteStream):
|
||||
async def __aiter__(self):
|
||||
yield compressed
|
||||
|
||||
attempt = SimpleNamespace(
|
||||
wire_streams=[],
|
||||
data={},
|
||||
headers=lambda _: None,
|
||||
stream_work=lambda *_: None,
|
||||
publish_if_due=lambda: None,
|
||||
)
|
||||
response = httpx2.Response(
|
||||
200,
|
||||
headers={"content-type": "text/event-stream", "content-encoding": "gzip"},
|
||||
stream=Stream(),
|
||||
)
|
||||
token = ACTIVE_ATTEMPT.set(attempt)
|
||||
try:
|
||||
await observe_response_headers(response)
|
||||
assert b"".join([chunk async for chunk in response.aiter_bytes()]) == body
|
||||
stats = attempt.wire_streams[0].observation.snapshot()
|
||||
assert stats["sse_observation"] == "skipped_content_encoding"
|
||||
assert stats["body_bytes"] == len(compressed)
|
||||
assert not stats["responses"] and not stats["parse_errors"]
|
||||
finally:
|
||||
await response.aclose()
|
||||
ACTIVE_ATTEMPT.reset(token)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("strict,expected", [(None, "omitted"), (False, "false"), (True, "true")])
|
||||
def test_actual_strict_state_and_optional_schema_are_observed_without_mutation(strict, expected):
|
||||
tool = {
|
||||
"type": "function",
|
||||
"name": "click",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"target": {"anyOf": [{"type": "string"}, {"type": "null"}]}},
|
||||
"required": [],
|
||||
},
|
||||
}
|
||||
if strict is not None:
|
||||
tool["strict"] = strict
|
||||
body = json.dumps(
|
||||
{"tools": [tool], "input": "private prompt", "max_output_tokens": 2048}
|
||||
).encode()
|
||||
result = request_constraints(body)
|
||||
assert result["max_output_tokens"] == 2048
|
||||
(observed,) = result["tools"]
|
||||
assert observed["strict"] == expected
|
||||
assert observed["schema"]["optional_properties"] == observed["schema"]["open_objects"] == 1
|
||||
assert "strict" not in tool if strict is None else tool["strict"] is strict
|
||||
assert "private" not in json.dumps(result)
|
||||
assert json.loads(body)["tools"][0] == tool
|
||||
|
||||
|
||||
def test_constraints_bound_large_requests_namespaces_and_deep_schemas():
|
||||
assert request_constraints(b"x" * (MAX_REQUEST_BYTES + 1)) == {"observation": "body_too_large"}
|
||||
assert request_constraints(b"no JSON") == {"observation": "invalid_json"}
|
||||
tools = [
|
||||
{
|
||||
"type": "namespace",
|
||||
"name": "functions",
|
||||
"tools": [{"type": "function", "name": "click"}],
|
||||
},
|
||||
{"type": "function", "name": "drag", "strict": False},
|
||||
]
|
||||
result = tool_constraints(tools)
|
||||
assert {t["name"]: t["strict"] for t in result["tools"]} == {
|
||||
"click": "omitted",
|
||||
"drag": "false",
|
||||
}
|
||||
schema = {"type": "object", "properties": {str(i): {"type": "string"} for i in range(2000)}}
|
||||
result = tool_constraints([{"type": "function", "name": "wide", "parameters": schema}])
|
||||
assert result["tools"][0]["schema"]["truncated"] is True
|
||||
assert result["tools"][0]["schema"]["nodes"] <= 1024
|
||||
@@ -0,0 +1,323 @@
|
||||
import asyncio
|
||||
from copy import deepcopy
|
||||
import json
|
||||
|
||||
import httpx2 as httpx
|
||||
from langchain_core.messages import AIMessageChunk
|
||||
from langchain_core.callbacks import AsyncCallbackHandler
|
||||
from pydantic import SecretStr
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from artemis.llm.diagnostics import ACTIVE_ATTEMPT
|
||||
from artemis.llm.reliability import LLMExhaustedError
|
||||
from artemis.llm.router import ModelEndpoint, ModelFactory
|
||||
from artemis.mcp.action_specs import tool_declaration
|
||||
from artemis.runtime.cancellation import CancellationReason
|
||||
from artemis.services import llm
|
||||
|
||||
|
||||
def event(kind, sequence, **fields):
|
||||
return (
|
||||
"data: " + json.dumps({"type": kind, "sequence_number": sequence, **fields}) + "\n\n"
|
||||
).encode()
|
||||
|
||||
|
||||
def initial_events():
|
||||
return event(
|
||||
"response.created",
|
||||
0,
|
||||
response={
|
||||
"id": "resp_test",
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "in_progress",
|
||||
"model": "test-model",
|
||||
"output": [],
|
||||
},
|
||||
) + event(
|
||||
"response.output_item.added",
|
||||
1,
|
||||
output_index=0,
|
||||
item={
|
||||
"type": "function_call",
|
||||
"id": "fc_test",
|
||||
"call_id": "call_test",
|
||||
"name": "action_sequence",
|
||||
"arguments": "",
|
||||
"status": "in_progress",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
class Stream(httpx.AsyncByteStream):
|
||||
def __init__(self):
|
||||
self.fragments = [
|
||||
event(
|
||||
"response.function_call_arguments.delta",
|
||||
2,
|
||||
output_index=0,
|
||||
item_id="fc_test",
|
||||
delta='{"steps":[',
|
||||
)
|
||||
]
|
||||
self.stalled = asyncio.Event()
|
||||
self.close_calls = 0
|
||||
|
||||
async def __aiter__(self):
|
||||
yield initial_events()
|
||||
for fragment in self.fragments:
|
||||
yield fragment
|
||||
self.stalled.set()
|
||||
await asyncio.Event().wait()
|
||||
|
||||
async def aclose(self):
|
||||
self.close_calls += 1
|
||||
|
||||
|
||||
class Engine:
|
||||
current_session_id = "wire-lifecycle"
|
||||
current_step_id = None
|
||||
|
||||
def __init__(self):
|
||||
self.records = []
|
||||
self.retrying = asyncio.Event()
|
||||
|
||||
def record_trace(self, **record):
|
||||
self.records.append(deepcopy(record))
|
||||
if record.get("payload", {}).get("phase") == "retry_wait":
|
||||
self.retrying.set()
|
||||
|
||||
def _publish(self, *_):
|
||||
pass
|
||||
|
||||
@property
|
||||
def data(self):
|
||||
return next(r["payload"] for r in reversed(self.records) if r["type"] == "llm_diagnostic")
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def case(monkeypatch, tmp_path):
|
||||
import openai
|
||||
from langchain_openai import ChatOpenAI
|
||||
|
||||
stream, engine, clients, requests = Stream(), Engine(), [], []
|
||||
|
||||
def handle(request):
|
||||
requests.append(json.loads(request.content))
|
||||
return httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=stream)
|
||||
|
||||
def client(**kwargs):
|
||||
value = httpx.AsyncClient(transport=httpx.MockTransport(handle), **kwargs)
|
||||
clients.append(value)
|
||||
return value
|
||||
|
||||
monkeypatch.setattr(openai, "DefaultAsyncHttpxClient", client)
|
||||
monkeypatch.setattr(llm.settings, "OPENAI_API_KEY", SecretStr("test-secret"))
|
||||
monkeypatch.setattr(llm.settings, "OPENAI_BASE_URL", "https://responses.example/v1")
|
||||
monkeypatch.delenv("ARTEMIS_FAKE_LLM", raising=False)
|
||||
monkeypatch.setattr(llm, "_get_current_data_engine", lambda: engine)
|
||||
monkeypatch.setattr(llm, "PAUSE_FILE", tmp_path / "pause")
|
||||
llm._ENDPOINT_BREAKER._states.clear()
|
||||
llm._NON_STREAMING_ENDPOINTS.clear()
|
||||
fallback = llm._FALLBACK_AVAILABLE.set(True)
|
||||
|
||||
def wrapper(timeout=30, max_retries=0):
|
||||
endpoint = ModelEndpoint(
|
||||
model_name="test-model", timeout_seconds=timeout, max_retries=max_retries
|
||||
)
|
||||
model = ModelFactory.create_model(endpoint)
|
||||
assert isinstance(model, ChatOpenAI)
|
||||
return llm.RobustChatModelWrapper(model, endpoint=endpoint).bind_tools(
|
||||
[tool_declaration(n) for n in ("click", "action_sequence", "drag")]
|
||||
)
|
||||
|
||||
yield stream, engine, requests, wrapper
|
||||
llm._FALLBACK_AVAILABLE.reset(fallback)
|
||||
for value in clients:
|
||||
await value.aclose()
|
||||
llm._ENDPOINT_BREAKER._states.clear()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("reason", ["outer_timeout", "request_timeout", "user_stop"])
|
||||
async def test_unfinished_arguments_below_length_limit_close_on_all_deadlines(case, reason):
|
||||
stream, engine, requests, wrapper = case
|
||||
if reason == "outer_timeout":
|
||||
running = asyncio.create_task(
|
||||
llm.invoke_llm_with_timeout_message(
|
||||
wrapper().complete([]), timeout_seconds=0.1, hard_timeout=0.1
|
||||
)
|
||||
)
|
||||
elif reason == "request_timeout":
|
||||
running = asyncio.create_task(wrapper(0.1).complete([]))
|
||||
else:
|
||||
running = asyncio.create_task(wrapper().complete([]))
|
||||
await asyncio.wait_for(stream.stalled.wait(), 2)
|
||||
running.cancel(CancellationReason.USER_STOP)
|
||||
with pytest.raises((TimeoutError, LLMExhaustedError, asyncio.CancelledError)):
|
||||
await asyncio.wait_for(running, 3)
|
||||
data = engine.data
|
||||
assert data["tool_argument_chars"] == len('{"steps":[') < 16384
|
||||
assert (
|
||||
data["failure_phase"] == "generating_tool_arguments"
|
||||
if reason == "request_timeout"
|
||||
else data["cancel_phase"] == "generating_tool_arguments"
|
||||
)
|
||||
assert (
|
||||
data["timeout_scope"] == ("outer" if reason == "outer_timeout" else "request")
|
||||
if reason != "user_stop"
|
||||
else data["cancel_reason"] == "user_stop"
|
||||
)
|
||||
assert data["cancel_observed_at"] <= data["ended_at"]
|
||||
(wire,) = data["wire_responses"]
|
||||
assert wire["closed_at"] >= wire["close_started_at"]
|
||||
assert wire["body_eof_observed"] is False
|
||||
assert wire["tools"][0].get("arguments_done_events", 0) == 0
|
||||
assert stream.close_calls == len(requests) == 1
|
||||
assert ACTIVE_ATTEMPT.get() is None
|
||||
assert "test-secret" not in repr(data)
|
||||
assert '{"steps":[' not in repr(data)
|
||||
assert not llm.PAUSE_FILE.exists()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_overflow_closes_transport_even_while_consumer_is_between_sdk_yields(case):
|
||||
stream, engine, requests, wrapper = case
|
||||
stream.fragments = [
|
||||
event(
|
||||
"response.function_call_arguments.delta",
|
||||
2,
|
||||
output_index=0,
|
||||
item_id="fc_test",
|
||||
delta=" " * 17000,
|
||||
)
|
||||
]
|
||||
with pytest.raises(LLMExhaustedError):
|
||||
await wrapper().complete([])
|
||||
assert engine.data["stream_guard"]["tool"] == "action_sequence"
|
||||
assert engine.data["wire_responses"][0]["closed_at"]
|
||||
assert stream.close_calls == len(requests) == 1
|
||||
assert not stream.stalled.is_set()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_overflow_closes_transport_before_waiting_for_retry(case):
|
||||
stream, engine, requests, wrapper = case
|
||||
stream.fragments = [
|
||||
event(
|
||||
"response.function_call_arguments.delta",
|
||||
2,
|
||||
output_index=0,
|
||||
item_id="fc_test",
|
||||
delta=" " * 17000,
|
||||
)
|
||||
]
|
||||
running = asyncio.create_task(wrapper(max_retries=1).complete([]))
|
||||
await asyncio.wait_for(engine.retrying.wait(), 2)
|
||||
assert stream.close_calls == len(requests) == 1
|
||||
assert engine.data["wire_responses"][0]["closed_at"]
|
||||
running.cancel(CancellationReason.USER_STOP)
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await running
|
||||
assert engine.data["cancel_phase"] == "retry_wait"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_actual_native_schemas_do_not_silently_force_upstream_strict(case):
|
||||
stream, engine, requests, wrapper = case
|
||||
running = asyncio.create_task(wrapper().complete([]))
|
||||
await asyncio.wait_for(stream.stalled.wait(), 2)
|
||||
running.cancel(CancellationReason.USER_STOP)
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await running
|
||||
constraints = engine.data["wire_request_constraints"]
|
||||
assert constraints["max_output_tokens"] is None
|
||||
assert {t["name"]: t["strict"] for t in constraints["tools"]} == {
|
||||
"click": "omitted",
|
||||
"action_sequence": "omitted",
|
||||
"drag": "omitted",
|
||||
}
|
||||
assert any(t["schema"]["optional_properties"] for t in constraints["tools"])
|
||||
assert all("strict" not in tool for tool in requests[0]["tools"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_live_deltas_are_emitted_before_a_batch_is_flushed(monkeypatch):
|
||||
emitted, resume = asyncio.Event(), asyncio.Event()
|
||||
events = []
|
||||
|
||||
class Model:
|
||||
async def astream(self, *_args, **_kwargs):
|
||||
yield AIMessageChunk(content="first live text")
|
||||
await resume.wait()
|
||||
|
||||
wrapper = llm.RobustChatModelWrapper(Model(), endpoint=ModelEndpoint(model_name="test"))
|
||||
monkeypatch.setattr(wrapper, "_emit_stream_delta", lambda *_: emitted.set())
|
||||
monkeypatch.setattr(llm, "_record_llm_event", lambda *args: events.append(args))
|
||||
running = asyncio.create_task(wrapper._complete_attempt(True, []))
|
||||
await asyncio.wait_for(emitted.wait(), 2)
|
||||
assert not running.done()
|
||||
running.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await running
|
||||
assert any(name == "llm_stream_reset" for name, *_ in events)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_callback_backpressure_can_leave_wire_ahead_of_sdk_without_prefetch(case):
|
||||
stream, engine, requests, wrapper = case
|
||||
blocked = asyncio.Event()
|
||||
stream.fragments = [
|
||||
b"".join(
|
||||
event(
|
||||
"response.function_call_arguments.delta",
|
||||
n + 2,
|
||||
output_index=0,
|
||||
item_id="fc_test",
|
||||
delta=" ",
|
||||
)
|
||||
for n in range(200)
|
||||
)
|
||||
]
|
||||
|
||||
class Callback(AsyncCallbackHandler):
|
||||
async def on_llm_new_token(self, token, *, chunk=None, **kwargs):
|
||||
if chunk is not None and any(t.get("args") for t in chunk.message.tool_call_chunks):
|
||||
attempt = ACTIVE_ATTEMPT.get()
|
||||
attempt.publish()
|
||||
blocked.set()
|
||||
await asyncio.Event().wait()
|
||||
|
||||
running = asyncio.create_task(wrapper().complete([], config={"callbacks": [Callback()]}))
|
||||
await asyncio.wait_for(blocked.wait(), 2)
|
||||
wire = engine.data["wire_responses"][0]
|
||||
assert wire["tools"][0]["argument_chars"] == 200
|
||||
assert engine.data["tool_argument_chars"] == 0
|
||||
assert not stream.stalled.is_set() # no second transport read or background reader
|
||||
running.cancel(CancellationReason.USER_STOP)
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(running, 2)
|
||||
assert stream.close_calls == len(requests) == 1
|
||||
assert engine.data["wire_responses"][0]["closed_at"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unresponsive_close_keeps_user_stop_and_records_cleanup_failure(case, monkeypatch):
|
||||
stream, engine, requests, wrapper = case
|
||||
|
||||
async def stuck_close():
|
||||
stream.close_calls += 1
|
||||
await asyncio.Event().wait()
|
||||
|
||||
monkeypatch.setattr(stream, "aclose", stuck_close)
|
||||
running = asyncio.create_task(wrapper().complete([]))
|
||||
await asyncio.wait_for(stream.stalled.wait(), 2)
|
||||
running.cancel(CancellationReason.USER_STOP)
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(running, 3)
|
||||
assert engine.data["cancel_reason"] == "user_stop"
|
||||
wire = engine.data["wire_responses"][0]
|
||||
assert wire["close_error_type"] == "TimeoutError"
|
||||
assert "closed_at" not in wire
|
||||
assert stream.close_calls == len(requests) == 1
|
||||
Reference in new issue
Block a user