302 lines
12 KiB
Python
302 lines
12 KiB
Python
"""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"])
|