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
@@ -66,6 +66,9 @@ def request_diagnostics(
|
||||
"headers_at",
|
||||
"first_output_at",
|
||||
"last_output_at",
|
||||
"first_tool_arguments_at",
|
||||
"last_tool_arguments_at",
|
||||
"cancel_observed_at",
|
||||
"ended_at",
|
||||
"retry_at",
|
||||
)
|
||||
|
||||
@@ -16,12 +16,38 @@ class ArgumentFingerprint:
|
||||
self.encoding_valid = True
|
||||
self.digest = hashlib.sha256()
|
||||
self.prefixes: list[dict] = []
|
||||
self.whitespace_chars = 0
|
||||
self.whitespace_run = 0
|
||||
self.max_whitespace_run = 0
|
||||
self.min_delta_chars: int | None = None
|
||||
self.max_delta_chars = 0
|
||||
self.same_delta_run = 0
|
||||
self.max_same_delta_run = 0
|
||||
self.previous_delta_digest: bytes | None = None
|
||||
|
||||
def observe(self, arguments: str) -> None:
|
||||
self.characters += len(arguments)
|
||||
if not arguments:
|
||||
return
|
||||
self.deltas += 1
|
||||
length = len(arguments)
|
||||
self.min_delta_chars = min(self.min_delta_chars or length, length)
|
||||
self.max_delta_chars = max(self.max_delta_chars, length)
|
||||
for character in arguments:
|
||||
if character in " \t\r\n":
|
||||
self.whitespace_chars += 1
|
||||
self.whitespace_run += 1
|
||||
self.max_whitespace_run = max(self.max_whitespace_run, self.whitespace_run)
|
||||
else:
|
||||
self.whitespace_run = 0
|
||||
# One digest, never the previous argument body. Surrogate handling here
|
||||
# compares SDK strings only; wire UTF-8 fingerprints remain invalid below.
|
||||
delta_digest = hashlib.sha256(arguments.encode("utf-8", "surrogatepass")).digest()
|
||||
self.same_delta_run = (
|
||||
self.same_delta_run + 1 if delta_digest == self.previous_delta_digest else 1
|
||||
)
|
||||
self.max_same_delta_run = max(self.max_same_delta_run, self.same_delta_run)
|
||||
self.previous_delta_digest = delta_digest
|
||||
if not self.encoding_valid:
|
||||
return
|
||||
try:
|
||||
@@ -50,4 +76,11 @@ class ArgumentFingerprint:
|
||||
"argument_deltas": self.deltas,
|
||||
"argument_prefixes": [dict(prefix) for prefix in self.prefixes],
|
||||
"argument_encoding_valid": self.encoding_valid,
|
||||
"argument_json_whitespace_chars": self.whitespace_chars,
|
||||
"argument_json_whitespace_run": self.whitespace_run,
|
||||
"argument_json_whitespace_run_max": self.max_whitespace_run,
|
||||
"argument_delta_chars_min": self.min_delta_chars,
|
||||
"argument_delta_chars_max": self.max_delta_chars,
|
||||
"argument_same_delta_run": self.same_delta_run,
|
||||
"argument_same_delta_run_max": self.max_same_delta_run,
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
"""Amortize LangChain's partial-JSON parsing without buffering live UI output."""
|
||||
|
||||
from langchain_core.messages import AIMessageChunk
|
||||
from langchain_core.messages.ai import add_ai_message_chunks
|
||||
|
||||
CHUNK_BATCH_SIZE = 128
|
||||
|
||||
|
||||
class ChunkAccumulator:
|
||||
def __init__(self):
|
||||
self.response = None
|
||||
self.pending = []
|
||||
|
||||
@property
|
||||
def has_output(self):
|
||||
return self.response is not None or bool(self.pending)
|
||||
|
||||
@property
|
||||
def ready(self):
|
||||
return len(self.pending) >= CHUNK_BATCH_SIZE
|
||||
|
||||
def append(self, chunk):
|
||||
self.pending.append(chunk)
|
||||
|
||||
def flush(self):
|
||||
chunks = ([self.response] if self.response is not None else []) + self.pending
|
||||
if not chunks:
|
||||
return None
|
||||
if all(isinstance(chunk, AIMessageChunk) for chunk in chunks):
|
||||
self.response = add_ai_message_chunks(chunks[0], *chunks[1:])
|
||||
else:
|
||||
result = chunks[0]
|
||||
for chunk in chunks[1:]:
|
||||
result = result + chunk
|
||||
self.response = result
|
||||
self.pending.clear()
|
||||
return self.response
|
||||
@@ -1,6 +1,7 @@
|
||||
"""Small, credential-free request snapshots in the existing task trace store."""
|
||||
|
||||
import base64
|
||||
import asyncio
|
||||
from contextvars import ContextVar
|
||||
from copy import deepcopy
|
||||
import io
|
||||
@@ -11,6 +12,8 @@ import time
|
||||
from uuid import uuid4
|
||||
|
||||
from PIL import Image
|
||||
import httpx
|
||||
import httpx2
|
||||
|
||||
from artemis.data_engine.context_vars import CURRENT_NODE_NAME, CURRENT_TRACE_ID
|
||||
from artemis.llm.argument_fingerprint import ArgumentFingerprint, MAX_TRACKED_TOOLS
|
||||
@@ -20,6 +23,8 @@ from artemis.llm.usage import usage_from_message
|
||||
from artemis.llm.stream_guard import ToolArgumentOverflow
|
||||
from artemis.llm.response_guard import ResponseIntegrityError
|
||||
from artemis.runtime.cancellation import CancellationReason, cancellation_reason
|
||||
from artemis.llm.responses_diagnostics import ObservedResponseStream
|
||||
from artemis.llm.tool_constraints import request_constraints
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
ACTIVE_ATTEMPT: ContextVar["RequestAttempt | None"] = ContextVar("llm_attempt", default=None)
|
||||
@@ -109,6 +114,7 @@ class RequestAttempt:
|
||||
self.last_publish = self.clock
|
||||
self.tools: dict[int | str, dict] = {}
|
||||
self.tool_fingerprints: dict[int | str, ArgumentFingerprint] = {}
|
||||
self.wire_streams: list[ObservedResponseStream] = []
|
||||
self.review = ACTIVE_FINAL_REVIEW.get()
|
||||
if self.review:
|
||||
self.review.request_started(request_id, attempt)
|
||||
@@ -138,6 +144,8 @@ class RequestAttempt:
|
||||
self.publish()
|
||||
|
||||
def publish(self):
|
||||
if self.wire_streams:
|
||||
self.data["wire_responses"] = [s.observation.snapshot() for s in self.wire_streams]
|
||||
self.data["sequence"] += 1
|
||||
self.data["updated_at"] = time.time()
|
||||
self.data["elapsed_seconds"] = round(time.monotonic() - self.clock, 3)
|
||||
@@ -168,6 +176,35 @@ class RequestAttempt:
|
||||
except (OSError, RuntimeError, ValueError, TypeError, AttributeError):
|
||||
logger.warning("Could not persist model request diagnostics", exc_info=True)
|
||||
|
||||
def publish_if_due(self):
|
||||
if time.monotonic() - self.last_publish >= 5:
|
||||
self.publish()
|
||||
|
||||
def stream_cancelled(self):
|
||||
self.data.setdefault("cancel_observed_at", time.time())
|
||||
self.data.setdefault("cancel_observed_phase", self.data["phase"])
|
||||
self.data.setdefault("cancel_phase", self.data["phase"])
|
||||
|
||||
def stream_work(self, stage: str, elapsed: float):
|
||||
stats = self.data.setdefault("stream_work", {}).setdefault(
|
||||
stage, {"calls": 0, "total_seconds": 0.0, "max_seconds": 0.0}
|
||||
)
|
||||
stats["calls"] += 1
|
||||
stats["total_seconds"] += elapsed
|
||||
stats["max_seconds"] = max(stats["max_seconds"], elapsed)
|
||||
|
||||
async def close_streams(self):
|
||||
try:
|
||||
async with asyncio.timeout(1):
|
||||
for stream in self.wire_streams:
|
||||
await stream.aclose()
|
||||
except (TimeoutError, OSError, RuntimeError, httpx.HTTPError, httpx2.HTTPError) as error:
|
||||
self.data["stream_close_error_type"] = type(error).__name__
|
||||
logger.warning("Model stream cleanup failed: %s", type(error).__name__)
|
||||
finally:
|
||||
if self.wire_streams:
|
||||
self.publish()
|
||||
|
||||
def headers(self, response):
|
||||
self.data.update(
|
||||
headers_at=time.time(),
|
||||
@@ -207,6 +244,9 @@ class RequestAttempt:
|
||||
arguments = tool.get("args")
|
||||
if isinstance(arguments, str):
|
||||
self.data["tool_argument_chars"] += len(arguments)
|
||||
if arguments:
|
||||
self.data.setdefault("first_tool_arguments_at", time.time())
|
||||
self.data["last_tool_arguments_at"] = time.time()
|
||||
index = tool.get("index")
|
||||
call_id = _identifier(tool.get("id"))
|
||||
key = index if index is not None else call_id or "unknown"
|
||||
@@ -297,6 +337,7 @@ class RequestAttempt:
|
||||
self.publish()
|
||||
|
||||
def cancel(self, error: BaseException | None = None):
|
||||
self.stream_cancelled()
|
||||
reason = cancellation_reason(error)
|
||||
if self.review and self.review.expired and reason is CancellationReason.UNKNOWN:
|
||||
self.data["timeout_scope"] = "final_review"
|
||||
@@ -321,6 +362,10 @@ async def observe_request_headers(request):
|
||||
attempt = ACTIVE_ATTEMPT.get()
|
||||
if attempt is None:
|
||||
return
|
||||
try:
|
||||
attempt.data["wire_request_constraints"] = request_constraints(request.content)
|
||||
except (httpx.RequestNotRead, httpx2.RequestNotRead):
|
||||
attempt.data["wire_request_constraints"] = {"observation": "streaming_body_unavailable"}
|
||||
request.headers.setdefault("X-Client-Request-Id", attempt.data["request_id"])
|
||||
request.headers.setdefault("X-Client-Attempt-Id", attempt.data["attempt_id"])
|
||||
routing = {}
|
||||
@@ -342,7 +387,25 @@ async def observe_request_headers(request):
|
||||
|
||||
|
||||
async def observe_response_headers(response):
|
||||
"""httpx hook: only allowlisted headers, never body/auth/request headers."""
|
||||
"""Correlate allowlisted headers and passively observe bounded SSE frames."""
|
||||
attempt = ACTIVE_ATTEMPT.get()
|
||||
if attempt is not None:
|
||||
attempt.headers(response)
|
||||
content_type = response.headers.get("content-type", "").split(";", 1)[0].lower()
|
||||
if (
|
||||
response.status_code == 200
|
||||
and content_type == "text/event-stream"
|
||||
and isinstance(response.stream, (httpx.AsyncByteStream, httpx2.AsyncByteStream))
|
||||
and not isinstance(response.stream, ObservedResponseStream)
|
||||
):
|
||||
if len(attempt.wire_streams) < 4:
|
||||
stream = ObservedResponseStream(
|
||||
response.stream,
|
||||
attempt,
|
||||
encoded=response.headers.get("content-encoding", "identity").lower().strip()
|
||||
not in {"", "identity"},
|
||||
)
|
||||
response.stream = stream
|
||||
attempt.wire_streams.append(stream)
|
||||
else:
|
||||
attempt.data["wire_streams_truncated"] = True
|
||||
@@ -0,0 +1,301 @@
|
||||
"""Bounded observations at the client's raw SSE read boundary; never prefetch."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import time
|
||||
|
||||
import httpx
|
||||
import httpx2
|
||||
|
||||
from artemis.llm.argument_fingerprint import ArgumentFingerprint, MAX_TRACKED_TOOLS
|
||||
from artemis.llm.tool_constraints import safe_label, tool_constraints
|
||||
|
||||
MAX_FRAME_BYTES = 256 * 1024
|
||||
MAX_RESPONSES = 16
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def identifier(value):
|
||||
return (
|
||||
value
|
||||
if isinstance(value, str) and re.fullmatch(r"[A-Za-z0-9_.:/-]{1,256}", value)
|
||||
else None
|
||||
)
|
||||
|
||||
|
||||
class ResponsesObservation:
|
||||
def __init__(self, *, observe_sse=True):
|
||||
self.observe_sse = observe_sse
|
||||
self.data = {
|
||||
"observation_boundary": "httpx_response_body_read",
|
||||
"body_bytes": 0,
|
||||
"sse_data_events": 0,
|
||||
"parse_errors": 0,
|
||||
"oversized_frames": 0,
|
||||
"unknown_response_events": 0,
|
||||
"statistics_truncated": False,
|
||||
"body_eof_observed": False,
|
||||
}
|
||||
if not observe_sse:
|
||||
self.data["sse_observation"] = "skipped_content_encoding"
|
||||
self.responses: dict[str, dict] = {}
|
||||
self.tools: dict[tuple, dict] = {}
|
||||
self.fingerprints: dict[tuple, ArgumentFingerprint] = {}
|
||||
self.aliases: dict[tuple, tuple] = {}
|
||||
self.current_response: str | None = None
|
||||
self.line = bytearray()
|
||||
self.lines: list[bytes] = []
|
||||
self.frame_bytes = 0
|
||||
self.dropping = False
|
||||
self.discarded_line = False
|
||||
|
||||
def feed(self, data: bytes):
|
||||
self.data["body_bytes"] += len(data)
|
||||
self.data["last_read_at"] = time.time()
|
||||
self.data.setdefault("first_read_at", self.data["last_read_at"])
|
||||
if not self.observe_sse:
|
||||
return
|
||||
offset = 0
|
||||
while offset < len(data):
|
||||
newline = data.find(b"\n", offset)
|
||||
end = len(data) if newline < 0 else newline
|
||||
size = end - offset
|
||||
if self.dropping:
|
||||
if size:
|
||||
if not self.discarded_line and not self.line and data[offset:end] == b"\r":
|
||||
self.line.extend(b"\r")
|
||||
else:
|
||||
self.discarded_line = True
|
||||
self.line.clear()
|
||||
elif self.frame_bytes + len(self.line) + size + (newline >= 0) > MAX_FRAME_BYTES:
|
||||
self.data["oversized_frames"] += 1
|
||||
self.dropping = True
|
||||
self.line.clear()
|
||||
self.lines.clear()
|
||||
self.discarded_line = True
|
||||
elif not self.dropping:
|
||||
self.line.extend(data[offset:end])
|
||||
if newline < 0:
|
||||
break
|
||||
# When dropping, a blank physical line still ends the oversized frame.
|
||||
if self.dropping:
|
||||
if not self.discarded_line:
|
||||
self.dropping = False
|
||||
self.frame_bytes = 0
|
||||
self.discarded_line = False
|
||||
self.line.clear()
|
||||
else:
|
||||
line = bytes(self.line).removesuffix(b"\r")
|
||||
self.frame_bytes += len(self.line) + 1
|
||||
self.line.clear()
|
||||
if not line:
|
||||
self._frame()
|
||||
self.lines.clear()
|
||||
self.frame_bytes = 0
|
||||
elif line.startswith(b"data:"):
|
||||
self.lines.append(line[5:].removeprefix(b" "))
|
||||
offset = newline + 1
|
||||
|
||||
def _frame(self):
|
||||
if not self.lines:
|
||||
return
|
||||
self.data["sse_data_events"] += 1
|
||||
body = b"\n".join(self.lines)
|
||||
if body == b"[DONE]":
|
||||
self.data["sse_done_events"] = self.data.get("sse_done_events", 0) + 1
|
||||
return
|
||||
try:
|
||||
event = json.loads(body)
|
||||
except (ValueError, UnicodeError, RecursionError):
|
||||
self.data["parse_errors"] += 1
|
||||
return
|
||||
if not isinstance(event, dict):
|
||||
self.data["parse_errors"] += 1
|
||||
return
|
||||
self.event(event)
|
||||
|
||||
def event(self, event: dict):
|
||||
kind = event.get("type")
|
||||
kind = kind if isinstance(kind, str) else None
|
||||
response = event.get("response")
|
||||
response = response if isinstance(response, dict) else {}
|
||||
response_id = identifier(event.get("response_id")) or identifier(response.get("id"))
|
||||
if response_id:
|
||||
self.current_response = response_id
|
||||
response_id = self.current_response
|
||||
if response_id is None:
|
||||
self.data["unknown_response_events"] += 1
|
||||
return
|
||||
if response_id not in self.responses and len(self.responses) >= MAX_RESPONSES:
|
||||
self.data["statistics_truncated"] = True
|
||||
return
|
||||
stats = self.responses.setdefault(
|
||||
response_id,
|
||||
{
|
||||
"response_id": response_id,
|
||||
"events": 0,
|
||||
"sequence_missing_events": 0,
|
||||
"sequence_forward_gaps": 0,
|
||||
"sequence_forward_missing_total": 0,
|
||||
"sequence_non_increasing": 0,
|
||||
"sequence_equal_previous": 0,
|
||||
},
|
||||
)
|
||||
stats["events"] += 1
|
||||
sequence = event.get("sequence_number")
|
||||
if type(sequence) is int and sequence >= 0:
|
||||
previous = stats.get("sequence_last")
|
||||
if previous is not None:
|
||||
if sequence <= previous:
|
||||
stats["sequence_non_increasing"] += 1
|
||||
stats["sequence_equal_previous"] += int(sequence == previous)
|
||||
elif sequence > previous + 1:
|
||||
stats["sequence_forward_gaps"] += 1
|
||||
stats["sequence_forward_missing_total"] += sequence - previous - 1
|
||||
stats.setdefault("sequence_first", sequence)
|
||||
stats["sequence_last"] = sequence
|
||||
else:
|
||||
stats["sequence_missing_events"] += 1
|
||||
if kind in {
|
||||
"response.completed",
|
||||
"response.failed",
|
||||
"response.incomplete",
|
||||
"response.cancelled",
|
||||
}:
|
||||
stats["terminal_event"] = kind
|
||||
if "tools" in response:
|
||||
stats["returned_tool_constraints"] = tool_constraints(response["tools"])
|
||||
if "max_output_tokens" in response:
|
||||
limit = response["max_output_tokens"]
|
||||
stats["returned_max_output_tokens"] = limit if type(limit) is int else None
|
||||
self._tool(event, response_id)
|
||||
|
||||
def _tool(self, event, response_id):
|
||||
kind = event.get("type")
|
||||
if not isinstance(kind, str):
|
||||
return
|
||||
item = event.get("item")
|
||||
item = item if isinstance(item, dict) else {}
|
||||
if kind in {"response.output_item.added", "response.output_item.done"}:
|
||||
if item.get("type") != "function_call":
|
||||
return
|
||||
elif kind not in {
|
||||
"response.function_call_arguments.delta",
|
||||
"response.function_call_arguments.done",
|
||||
}:
|
||||
return
|
||||
item_id = identifier(event.get("item_id")) or identifier(item.get("id"))
|
||||
index = event.get("output_index")
|
||||
index = index if type(index) is int and index >= 0 else None
|
||||
alias = (response_id, index)
|
||||
key = self.aliases.get(alias) if index is not None else None
|
||||
key = key or (response_id, item_id if item_id else index)
|
||||
if key not in self.tools and len(self.tools) >= MAX_TRACKED_TOOLS:
|
||||
self.data["statistics_truncated"] = True
|
||||
return
|
||||
if index is not None:
|
||||
if alias in self.aliases or len(self.aliases) < 2 * MAX_TRACKED_TOOLS:
|
||||
self.aliases[alias] = key
|
||||
else:
|
||||
self.data["statistics_truncated"] = True
|
||||
stats = self.tools.setdefault(key, {"response_id": response_id, "delta_events": 0})
|
||||
if not item_id and index is None:
|
||||
stats["identity_unknown"] = True
|
||||
if item_id and stats.get("item_id") not in (None, item_id):
|
||||
stats["identity_conflicts"] = stats.get("identity_conflicts", 0) + 1
|
||||
if item_id:
|
||||
stats["item_id"] = item_id
|
||||
if (call_id := identifier(item.get("call_id"))) is not None:
|
||||
stats["call_id"] = call_id
|
||||
if (name := safe_label(item.get("name"))) is not None:
|
||||
stats["name"] = name
|
||||
if index is not None:
|
||||
stats["output_index"] = index
|
||||
fingerprint = self.fingerprints.get(key)
|
||||
if fingerprint is None:
|
||||
fingerprint = self.fingerprints[key] = ArgumentFingerprint()
|
||||
arguments = None
|
||||
if kind == "response.function_call_arguments.delta":
|
||||
stats["delta_events"] += 1
|
||||
arguments = event.get("delta")
|
||||
elif kind == "response.output_item.added":
|
||||
arguments = item.get("arguments")
|
||||
elif kind == "response.function_call_arguments.done":
|
||||
stats["arguments_done_events"] = stats.get("arguments_done_events", 0) + 1
|
||||
elif kind == "response.output_item.done":
|
||||
stats["output_item_done_events"] = stats.get("output_item_done_events", 0) + 1
|
||||
if isinstance(arguments, str):
|
||||
fingerprint.observe(arguments)
|
||||
if arguments:
|
||||
stats.setdefault("first_argument_at", time.time())
|
||||
stats["last_argument_at"] = time.time()
|
||||
|
||||
def end(self, *, eof=False):
|
||||
self.data["body_eof_observed"] = self.data["body_eof_observed"] or eof
|
||||
self.data["incomplete_sse_frame"] = self.data.get("incomplete_sse_frame", False) or bool(
|
||||
self.line or self.lines or self.dropping
|
||||
)
|
||||
self.line.clear()
|
||||
self.lines.clear()
|
||||
|
||||
def snapshot(self):
|
||||
return {
|
||||
**self.data,
|
||||
"responses": [dict(r) for r in self.responses.values()],
|
||||
"tools": [{**r, **self.fingerprints[k].snapshot()} for k, r in self.tools.items()],
|
||||
}
|
||||
|
||||
|
||||
class ObservedResponseStream(httpx.AsyncByteStream, httpx2.AsyncByteStream):
|
||||
def __init__(
|
||||
self, stream: httpx.AsyncByteStream | httpx2.AsyncByteStream, attempt, *, encoded=False
|
||||
):
|
||||
self.stream = stream
|
||||
self.attempt = attempt
|
||||
self.observation = ResponsesObservation(observe_sse=not encoded)
|
||||
self.closed = False
|
||||
self.close_attempted = False
|
||||
|
||||
async def __aiter__(self):
|
||||
try:
|
||||
async for data in self.stream:
|
||||
started = time.monotonic()
|
||||
self.observation.feed(data)
|
||||
self.attempt.stream_work("wire_observe", time.monotonic() - started)
|
||||
self.attempt.publish_if_due()
|
||||
yield data
|
||||
self.observation.end(eof=True)
|
||||
except asyncio.CancelledError:
|
||||
self.attempt.stream_cancelled()
|
||||
raise
|
||||
|
||||
async def aclose(self):
|
||||
if self.close_attempted:
|
||||
return
|
||||
self.close_attempted = True
|
||||
stats = self.observation.data
|
||||
stats.setdefault("close_started_at", time.time())
|
||||
try:
|
||||
async with asyncio.timeout(1):
|
||||
await self.stream.aclose()
|
||||
except TimeoutError:
|
||||
# A cleanup deadline must not replace an in-flight user stop with
|
||||
# a retryable model timeout. Failure remains explicit in diagnostics.
|
||||
stats["close_error_type"] = "TimeoutError"
|
||||
logger.warning("Model HTTP stream close exceeded its cleanup deadline")
|
||||
except (
|
||||
asyncio.CancelledError,
|
||||
OSError,
|
||||
RuntimeError,
|
||||
httpx.HTTPError,
|
||||
httpx2.HTTPError,
|
||||
) as error:
|
||||
stats["close_error_type"] = type(error).__name__
|
||||
raise
|
||||
else:
|
||||
self.closed = True
|
||||
stats["closed_at"] = time.time()
|
||||
finally:
|
||||
self.observation.end(eof=stats["body_eof_observed"])
|
||||
@@ -0,0 +1,109 @@
|
||||
"""Summarize actual wire tool constraints without persisting schemas or prompts."""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import re
|
||||
|
||||
from artemis.llm.argument_fingerprint import MAX_TRACKED_TOOLS
|
||||
|
||||
MAX_REQUEST_BYTES = 4 * 1024 * 1024
|
||||
|
||||
|
||||
def safe_label(value):
|
||||
return value if isinstance(value, str) and re.fullmatch(r"[\w.:-]{1,128}", value) else None
|
||||
|
||||
|
||||
def strict_state(tool: dict) -> str:
|
||||
if "strict" not in tool:
|
||||
return "omitted"
|
||||
if tool["strict"] is True:
|
||||
return "true"
|
||||
if tool["strict"] is False:
|
||||
return "false"
|
||||
return "invalid"
|
||||
|
||||
|
||||
def schema_summary(schema: dict) -> dict:
|
||||
"""Count syntactic schema nodes, not expanded $refs or arbitrary examples."""
|
||||
result = {"objects": 0, "optional_properties": 0, "open_objects": 0, "nodes": 0}
|
||||
pending = [schema]
|
||||
while pending and result["nodes"] < 1024:
|
||||
node = pending.pop()
|
||||
result["nodes"] += 1
|
||||
properties = node.get("properties")
|
||||
if node.get("type") == "object" or isinstance(properties, dict):
|
||||
result["objects"] += 1
|
||||
required = node.get("required")
|
||||
required = required if isinstance(required, list) else []
|
||||
if isinstance(properties, dict):
|
||||
result["optional_properties"] += sum(k not in required for k in properties)
|
||||
if node.get("additionalProperties") is not False:
|
||||
result["open_objects"] += 1
|
||||
children = []
|
||||
for key in ("properties", "$defs", "definitions"):
|
||||
if isinstance(node.get(key), dict):
|
||||
children.extend(node[key].values())
|
||||
for key in ("anyOf", "oneOf", "allOf", "prefixItems"):
|
||||
if isinstance(node.get(key), list):
|
||||
children.extend(node[key])
|
||||
children.extend(node.get(k) for k in ("items", "additionalProperties", "not"))
|
||||
remaining = max(0, 1024 - result["nodes"] - len(pending))
|
||||
eligible = [v for v in children if isinstance(v, dict)]
|
||||
if len(eligible) > remaining:
|
||||
result["truncated"] = True
|
||||
pending.extend(eligible[:remaining])
|
||||
result["truncated"] = bool(pending) or result.get("truncated", False)
|
||||
try:
|
||||
result["sha256"] = hashlib.sha256(
|
||||
json.dumps(schema, sort_keys=True, separators=(",", ":"), ensure_ascii=True).encode()
|
||||
).hexdigest()
|
||||
except (RecursionError, TypeError, ValueError):
|
||||
result["sha256"] = None
|
||||
result["truncated"] = True
|
||||
return result
|
||||
|
||||
|
||||
def tool_constraints(tools) -> dict:
|
||||
records = []
|
||||
pending = tools[:MAX_TRACKED_TOOLS] if isinstance(tools, list) else []
|
||||
truncated = isinstance(tools, list) and len(tools) > MAX_TRACKED_TOOLS
|
||||
visited = 0
|
||||
while pending and visited < MAX_TRACKED_TOOLS:
|
||||
tool = pending.pop(0)
|
||||
visited += 1
|
||||
if not isinstance(tool, dict):
|
||||
continue
|
||||
if tool.get("type") == "namespace":
|
||||
nested = tool.get("tools")
|
||||
if isinstance(nested, list):
|
||||
remaining = MAX_TRACKED_TOOLS - len(pending) - visited
|
||||
truncated = truncated or len(nested) > remaining
|
||||
pending.extend(nested[:remaining])
|
||||
continue
|
||||
function = tool.get("function", tool)
|
||||
if not isinstance(function, dict):
|
||||
continue
|
||||
record = {
|
||||
"type": safe_label(tool.get("type")),
|
||||
"name": safe_label(function.get("name")),
|
||||
"strict": strict_state(function),
|
||||
}
|
||||
if isinstance(function.get("parameters"), dict):
|
||||
record["schema"] = schema_summary(function["parameters"])
|
||||
records.append(record)
|
||||
return {"tools": records, "truncated": truncated or bool(pending), "observation": "available"}
|
||||
|
||||
|
||||
def request_constraints(body: bytes) -> dict:
|
||||
if len(body) > MAX_REQUEST_BYTES:
|
||||
return {"observation": "body_too_large"}
|
||||
try:
|
||||
payload = json.loads(body)
|
||||
except (ValueError, UnicodeError, RecursionError):
|
||||
return {"observation": "invalid_json"}
|
||||
if not isinstance(payload, dict):
|
||||
return {"observation": "not_an_object"}
|
||||
result = tool_constraints(payload.get("tools"))
|
||||
limit = payload.get("max_output_tokens")
|
||||
result["max_output_tokens"] = limit if type(limit) is int and limit > 0 else None
|
||||
return result
|
||||
+36
-8
@@ -47,6 +47,7 @@ from artemis.config import (
|
||||
from artemis.context import ArtemisContext
|
||||
from artemis.data_engine.trace import CURRENT_TRACE_ID, DataEngineCallbackHandler
|
||||
from artemis.llm.diagnostics import ACTIVE_ATTEMPT, RequestAttempt, error_text, input_images
|
||||
from artemis.llm.chunk_accumulator import ChunkAccumulator
|
||||
from artemis.llm.final_review import (
|
||||
ACTIVE_FINAL_REVIEW,
|
||||
FinalReviewBudget,
|
||||
@@ -425,8 +426,13 @@ async def _run_with_recovery[T](
|
||||
observation_token = ACTIVE_ATTEMPT.set(observation)
|
||||
request_timer = asyncio.timeout(timeout)
|
||||
try:
|
||||
async with request_timer:
|
||||
result = await call_once()
|
||||
try:
|
||||
async with request_timer:
|
||||
result = await call_once()
|
||||
finally:
|
||||
# Release a failed stream before any retry backoff, even
|
||||
# when the SDK consumer stopped between yielded chunks.
|
||||
await observation.close_streams()
|
||||
except (KeyboardInterrupt, asyncio.CancelledError) as e:
|
||||
observation.cancel(e)
|
||||
raise
|
||||
@@ -818,24 +824,46 @@ class RobustChatModelWrapper:
|
||||
|
||||
stream_exec_id = uuid4()
|
||||
full_response = None
|
||||
accumulator = ChunkAccumulator()
|
||||
argument_guard = ToolArgumentGuard()
|
||||
try:
|
||||
async with aclosing(self.base_model.astream(*args, **kwargs)) as stream:
|
||||
wait_started = time.monotonic()
|
||||
async for chunk in stream:
|
||||
observation = ACTIVE_ATTEMPT.get()
|
||||
if observation is not None:
|
||||
observation.stream_work("await_sdk_chunk", time.monotonic() - wait_started)
|
||||
observe_started = time.monotonic()
|
||||
observation.chunk(chunk)
|
||||
observation.stream_work("sdk_observe", time.monotonic() - observe_started)
|
||||
argument_guard.observe(chunk)
|
||||
if full_response is None:
|
||||
full_response = chunk
|
||||
else:
|
||||
full_response = full_response + chunk
|
||||
accumulator.append(chunk)
|
||||
if accumulator.ready:
|
||||
# Buffered transport input may yield without suspending.
|
||||
# Give stop/deadline signals a turn before a batch merge.
|
||||
await asyncio.sleep(0)
|
||||
merge_started = time.monotonic()
|
||||
full_response = accumulator.flush()
|
||||
if observation is not None:
|
||||
observation.stream_work("merge", time.monotonic() - merge_started)
|
||||
if emit_deltas:
|
||||
emit_started = time.monotonic()
|
||||
self._emit_stream_delta(stream_exec_id, chunk)
|
||||
if observation is not None:
|
||||
observation.stream_work("emit", time.monotonic() - emit_started)
|
||||
wait_started = time.monotonic()
|
||||
merge_started = time.monotonic()
|
||||
full_response = accumulator.flush()
|
||||
observation = ACTIVE_ATTEMPT.get()
|
||||
if observation is not None:
|
||||
observation.stream_work("merge", time.monotonic() - merge_started)
|
||||
if full_response is not None:
|
||||
ensure_response_complete(full_response)
|
||||
except (KeyboardInterrupt, asyncio.CancelledError):
|
||||
if full_response is not None:
|
||||
observation = ACTIVE_ATTEMPT.get()
|
||||
if observation is not None:
|
||||
observation.stream_cancelled()
|
||||
if accumulator.has_output:
|
||||
_record_llm_event(
|
||||
"llm_stream_reset",
|
||||
{
|
||||
@@ -847,7 +875,7 @@ class RobustChatModelWrapper:
|
||||
)
|
||||
raise
|
||||
except Exception as stream_error:
|
||||
if full_response is not None or isinstance(
|
||||
if accumulator.has_output or isinstance(
|
||||
stream_error, (ToolArgumentOverflow, ResponseIntegrityError)
|
||||
):
|
||||
if isinstance(stream_error, ResponseIntegrityError):
|
||||
|
||||
@@ -2,6 +2,8 @@
|
||||
|
||||
小型内置动作的参数流仍受 16384 字符保护。超限时关闭输出流、丢弃不完整参数,不执行该次动作;一次新响应重试和可用备用端点耗尽后,只结束当前任务,不进入共享的模型暂停。批次仍遵循“失败即停止”、人工暂停和取消设置。网络故障、认证错误和最终验收保留各自的恢复规则。
|
||||
|
||||
长度保护不是生成进度判断:小于上限但一直没有闭合或终止的参数,仍可能持续到单次请求或外层任务超时。超时、取消和完整性校验共同约束这条路径,未完成的参数不能作为动作执行。
|
||||
|
||||
## 请求关联
|
||||
|
||||
每个模型请求尝试沿用现有 `llm_diagnostic` 记录,不另建日志数据库。
|
||||
@@ -27,6 +29,9 @@
|
||||
- `argument_deltas`:非空参数片段数,不能直接与网关原始事件数比较。
|
||||
- `argument_prefixes`:每 4096 字节的累计 SHA256,最多保留前 16 个。可以在分片大小不同、上游预读或中途取消时比较相同字节位置。
|
||||
- `argument_encoding_valid`:遇到不能编码为 UTF-8 的片段时为 false,当前总字节数和总哈希留空,避免伪造可比结果;诊断本身不改变执行流程。
|
||||
- `argument_json_whitespace_chars` / `argument_json_whitespace_run` / `argument_json_whitespace_run_max`:实际空格、制表、回车和换行字符的总数、当前尾部连续数和最长连续数。统计的是这四种 JSON 空白码点,也包含字符串内部的实际字符;转义文本 `\\n`、不换行空格不计入。不据此判断 JSON 是否有效。
|
||||
- `argument_delta_chars_min` / `argument_delta_chars_max`:非空参数增量的最小、最大字符数;尚无非空增量时最小值为空。
|
||||
- `argument_same_delta_run` / `argument_same_delta_run_max`:连续相同非空增量的当前、最长出现次数,首个增量算一次,空增量不打断。只保留上一增量的 SHA256,不保留正文。相同增量内容不等于重复的协议事件,例如连续空格也可能来自不同的合法事件。
|
||||
|
||||
最多记录 64 个工具;超过时设置 `tool_stats_truncated=true`,总参数字符计数仍继续。诊断不保留参数正文。哈希仅证明已收到的字节前缀,不能证明输出完整或动作成功;完整性仍由响应终止状态和解析校验决定。
|
||||
|
||||
@@ -34,6 +39,38 @@
|
||||
|
||||
历史日志没有保存的原始帧无法由新增诊断恢复。HTTP 200 只说明响应头已到达;客户端保护后出现的 499 或 Broken pipe 也不能证明最初原因。网关写入本机下游 writer 成功不等于客户端已消费或执行。
|
||||
|
||||
## 客户端原始 Responses 观测
|
||||
|
||||
`wire_responses` 被动观察 HTTP 客户端实际读取的 SSE 字节,在 SDK 转换之前统计。同一网络读取中可能包含多个事件,因此原始统计可以领先 SDK;SDK 回调阻塞也会扩大两者差距。这一边界不是服务器写出时刻,也不是工具执行时刻。诊断不增加预读任务、队列或重传,不过滤、改写交给 SDK 的字节。
|
||||
|
||||
每个响应记录 `response_id`、事件数及 `sequence_number` 的首末值。`sequence_forward_gaps` / `sequence_forward_missing_total` 统计相邻观测间的向前跳变次数及跳过的序号数;乱序事件可能随后补上,不能直接称为永久丢包。`sequence_non_increasing` 统计非递增,`sequence_equal_previous` 统计与前一个序号相等。`sequence_missing_events` 表示事件本身没有有效序号,不表示传输丢失。没有响应编号时不猜测,记录 `unknown_response_events`;标准流中没有编号的后续事件沿用当前响应编号。
|
||||
|
||||
每个函数调用记录与 SDK 相同的参数指纹和形态统计,另记原始 `delta_events`、首末参数到达时间,以及两类 done 事件的次数。只累计非空初始 arguments 和 arguments.delta,不重复累计 done 的完整快照。关联优先使用响应编号与实际 `call_id`,`item_id` / `output_index` 用于同一原始流内关联;`identity_unknown`、`identity_conflicts` 表明身份不足或冲突,此时不能据此归责。
|
||||
|
||||
每次尝试最多观察 4 条 HTTP 流;每条流最多 16 个响应、64 个工具、128 个索引别名,每工具最多 16 个固定前缀。单个待解析 SSE 帧最多 256KiB,超出后仅跳过该帧的诊断,继续原样交给 SDK。`oversized_frames`、`parse_errors`、`statistics_truncated`、`wire_streams_truncated` 说明覆盖缺口;`incomplete_sse_frame` 与 `body_eof_observed` 区分未闭合帧和正常字节流结束。待解析字节仅临时保留,关闭时释放,日志不保存帧或参数正文。
|
||||
|
||||
原始分帧支持 LF / CRLF。HTTP 体声明 gzip 等压缩编码时,标记 `sse_observation=skipped_content_encoding`,只统计读取字节与时间,不把压缩字节误当成 SSE;SDK 仍按原路径解压。这种记录缺少原始事件统计,不能把零计数当成没有收到参数。
|
||||
|
||||
## 实际工具约束与耗时
|
||||
|
||||
`wire_request_constraints` 检查 HTTP 请求已经序列化的正文,不记录提示、图片或 Schema 正文。逐工具记录 `strict` 为 `omitted`、`false`、`true`,异常值为 `invalid`;不能把省略当成 false,更不能把 Pydantic 本地严格校验当成上游约束。Schema 只保存 SHA256、对象数、非 required 的属性数及未关闭对象数。这是语法统计,不展开引用,也不是上游 Schema 兼容性验证。
|
||||
|
||||
当前 `click`、`action_sequence`、`drag` 的真实 SDK 离线请求均省略 strict,保留可选字段。不能只把 strict 开关改为 true:应先按目标接口支持的 Schema 子集处理对象和 required,保留可选值的 nullable 语义,再单独验证兼容性。
|
||||
|
||||
请求正文超过 4MiB 时跳过检查并记录 `body_too_large`;流式上传正文不可直接读取时记录 `streaming_body_unavailable`,不会为了诊断消费上传流。最多检查 64 个工具或命名空间节点,每个 Schema 最多统计 1024 个节点,达到上限标记截断。响应事件如果回传 tools,则保存同样的 `returned_tool_constraints`;回显不证明服务端实际执行了约束。
|
||||
|
||||
`max_output_tokens` 与 `returned_max_output_tokens` 只表示本层实际观察到的字段。本次核对的 CodexManager 通道会过滤输出 token 上限;短合成探针成功也不能证明上限被执行。不能把设置 max_tokens / max_output_tokens 当成这条链路可靠的止损办法,客户端长度、超时和取消保护仍然必要。
|
||||
|
||||
`stream_work` 记录每阶段的调用次数、总耗时和最长耗时:`wire_observe` 为原始帧诊断,`sdk_observe` 为 SDK 片段诊断,`merge` 为聚合,`emit` 为界面回调,`await_sdk_chunk` 为等待 SDK 交出下一个片段。最后一项包含网络读取、原始观测、SDK 转换和 SDK 回调等工作,与其他阶段可能重叠,不能相加或直接当成网络延迟。
|
||||
|
||||
聚合每 128 个片段合并一次,降低 SDK 反复解析半截 JSON 的开销;长度保护、诊断和界面输出仍逐片段执行,异步压缩继续使用原异步调用路径。分批合并不修复参数 JSON,最终仍须经过完整性和严格解析校验。
|
||||
|
||||
## 取消与关闭
|
||||
|
||||
`cancel_observed_at` / `cancel_observed_phase` 记录本次尝试首次在客户端观测到取消的时间与阶段,不等于用户点击停止的时间。`cancel_phase` / `cancel_reason` 沿用恢复层最终记录的取消阶段和来源,可能反映后续重试等待中的停止。首末 SDK 参数时间为 `first_tool_arguments_at` / `last_tool_arguments_at`。原始流内的读取、参数和关闭时间采用 Unix 秒。
|
||||
|
||||
SDK 读取中、回调中或两次 yield 之间触发长度保护后,尝试结束路径都会主动关闭已观察的 HTTP 流,不等待垃圾回收。`close_started_at` 记录关闭开始,只有实际成功才写 `closed_at`;失败只记录异常类型,关闭等待最多 1 秒。关闭超时不能覆盖用户停止而变成可重试的请求超时。未能关闭时保留错误证据,不声称连接已成功释放。
|
||||
|
||||
## 已核对的上游异常样本(2026-11-02)
|
||||
|
||||
一次真实 `action_sequence` 调用在任务超时前持续接收参数,没有完成该工具调用。该样本发生在 `action_sequence` 纳入参数保护之前。按同一 `response_id + call_id` 核对生产诊断,得到以下字符数;请求编号和设备信息不在文档中保留:
|
||||
@@ -48,10 +85,14 @@
|
||||
|
||||
五个边界在 4096 到 36864 字节的九个累计 SHA256 上一致。NewAPI 使用透传模式,收到和写出的完整参数哈希也一致。上游没有终止事件,也没有完成参数快照。该样本证明参数在官方上游返回时就已持续增长,不能归因于 NewAPI 转换、下游重复累计完成快照或 ARTEMIS 最终消息拼接。取消、预读和缓冲会使各层停止计数的位置不同;尾部总长度不同本身不能证明重复发送或字符丢失。
|
||||
|
||||
诊断没有参数正文,不能进一步断言是空白循环、重复描述或官方服务内部的重复事件。这个定位也不能追溯证明没有原始诊断的历史事故。另一次 `click` 异常同样在 NewAPI 接收时已超过四万字符,说明这种现象并非新动作序列独有。
|
||||
诊断没有参数正文,不能进一步断言是模型生成循环、空白循环、重复描述或官方服务内部的重复事件,也不能排除请求结构触发异常。这个定位只排除已核对样本的下游放大作为最初来源,不能追溯证明没有原始诊断的历史事故。另一次 `click` 异常同样在 NewAPI 接收时已超过四万字符,说明这种现象并非新动作序列独有。
|
||||
|
||||
CodexManager 另行修正了普通函数放入 functions 命名空间的兼容性问题。其短上下文、低思考强度合成测试通过,仅能作为兼容性对照,尚不能代替原事故的长上下文、多图、高思考强度复现,也不能称为参数膨胀的根因修复。
|
||||
|
||||
本地保护覆盖 `action_sequence`、`drag` 及其他小型内置动作:超限后关闭流、丢弃参数,并按既有策略有限重试;它限制异常对当前任务的影响,不修复上游生成行为,也不保证业务任务一定成功。缩放面板自动隐藏导致的定位失败由动作序列的 `ensure_visible` 准备步骤处理,与参数流异常分别诊断。
|
||||
|
||||
## 离线验证
|
||||
|
||||
`tests/unit/test_llm_stream_guard.py` 验证保护、有限重试、相同/不同备用端点,以及不会创建或清除其他请求的全局暂停。`tests/unit/test_request_diagnostics.py` 验证请求编号、超时重试、响应头及单次参数计数。`tests/unit/test_argument_fingerprint.py` 验证不同分片、中文和 emoji、固定字节前缀、有界统计与异常编码。这些测试不调用真实模型或设备。
|
||||
|
||||
`tests/unit/test_responses_diagnostics.py` 验证原始分帧、序号、done 去重、有界统计和 strict 三态。`tests/unit/test_responses_wire_lifecycle.py` 使用已安装的真实 OpenAI/LangChain SDK 与内存 HTTP 传输,验证小于长度阈值的未闭合流、SDK 回调背压、请求/任务超时、用户停止、超限后主动关闭和关闭失败,以及实际工具 Schema;不发网络请求。`tests/unit/test_chunk_accumulator.py` 对比逐片段聚合结果,验证交错工具、元数据和使用量保持一致,减少半截 JSON 重复解析,异常参数仍由最终校验拒绝。离线通过不代表原业务任务的稳定性已验收。
|
||||
@@ -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