diff --git a/README.md b/README.md index 5f78a11..c84363f 100644 --- a/README.md +++ b/README.md @@ -87,6 +87,32 @@ with agentcompile.conversation(ticket.id, customer=ticket.customer_id): Streamed answers are captured too: your agent reads the stream as always, and the whole answer is captured once it ends. +## The Responses API and any framework + +A wrapped OpenAI client captures `client.responses.create(...)` as well (streamed or not). Those +calls always go to your model: compiled answers come in Chat Completions' and Anthropic's shapes. + +When a framework builds its own client (LangChain, LiteLLM, CrewAI, Pydantic AI, the OpenAI +Agents SDK), give it a capturing HTTP client instead of wrapping anything: + +```python +http = agentcompile.http_client(key="ack_...") # async: agentcompile.async_http_client() + +ChatOpenAI(model="gpt-5", http_client=http) # LangChain +litellm.client_session = http # LiteLLM, CrewAI +OpenAIProvider(http_client=agentcompile.async_http_client()) # Pydantic AI +set_default_openai_client(AsyncOpenAI(http_client=agentcompile.async_http_client())) # Agents SDK + +with agentcompile.conversation(ticket.id): + run_agent() +``` + +It captures the POSTs to `/chat/completions`, `/responses` and `/messages` that pass through +it, scrubbed the same way, and leaves everything else alone. It only captures: to have known +jobs answered compiled, use `wrap`. An SDK that takes only `httpx2` clients (newer `anthropic` +releases) takes `agentcompile.http_client(lib="httpx2")`. +`agentcompile.transport()` gives the bare transport, to wrap one of your own. + ## Privacy With `capture=True`, personal data is scrubbed on your machine before anything is sent: diff --git a/src/agentcompile/__init__.py b/src/agentcompile/__init__.py index 8218cae..5c47fe4 100644 --- a/src/agentcompile/__init__.py +++ b/src/agentcompile/__init__.py @@ -17,15 +17,40 @@ from ._capture import capturer_for, flush_all from ._conversation import conversation -from ._core import LIVE, SHADOW, Proxy, Router, async_create, sync_create +from ._core import ( + LIVE, + SHADOW, + Proxy, + Router, + async_create, + async_forward, + sync_create, + sync_forward, +) from ._decide import AsyncDecider, Decider, Settings from ._outcome import OUTCOMES, outcome from ._outcome import flush as _flush_outcomes from ._outcome import remember as _remember from ._scrub import load_key from ._trail import OnEvent, Trail - -__all__ = ["LIVE", "OUTCOMES", "SHADOW", "Settings", "conversation", "flush", "outcome", "wrap"] +from ._transport import AsyncCaptureTransport, CaptureTransport, classes, lib_of + +__all__ = [ + "LIVE", + "OUTCOMES", + "SHADOW", + "AsyncCaptureTransport", + "CaptureTransport", + "Settings", + "async_http_client", + "async_transport", + "conversation", + "flush", + "http_client", + "outcome", + "transport", + "wrap", +] try: from ._version import __version__ except ImportError: # running from a source checkout that was never built @@ -60,13 +85,7 @@ def wrap( this machine before anything is sent (default True; the key is AGENTCOMPILE_SCRUB_KEY or one created once in ~/.agentcompile/). """ - settings = Settings( - base_url=base_url or os.environ.get("AGENTCOMPILE_URL", DEFAULT_URL), - key=key or os.environ.get("AGENTCOMPILE_KEY"), - company=company or os.environ.get("AGENTCOMPILE_COMPANY"), - timeout=timeout, - mode=mode, - ) + settings = _settings(key, base_url, company, timeout, mode) the_trail = Trail(trail, on_event) _remember(settings) # outcome() sends with the client wrapped last if capture is None: @@ -90,10 +109,21 @@ def wrap( router = _router("openai", completions.create, settings, the_trail, mode, http_client) router.capturer = capturer create = _create(router, completions.create, _openai.build) - return Proxy( - client, - {"chat": Proxy(chat, {"completions": Proxy(completions, {"create": create})})}, - ) + overrides: dict[str, Any] = { + "chat": Proxy(chat, {"completions": Proxy(completions, {"create": create})}) + } + responses = getattr(client, "responses", None) + if responses is not None and hasattr(responses, "create"): + # The Responses API (the OpenAI Agents SDK's default): captured, always the model's. + forward = _router( + "openai", responses.create, settings, the_trail, mode, http_client + ) + forward.capturer = capturer + make = async_forward if _is_async(responses.create) else sync_forward + overrides["responses"] = Proxy( + responses, {"create": make(forward, responses.create)} + ) + return Proxy(client, overrides) messages = getattr(client, "messages", None) if messages is not None and hasattr(messages, "create"): @@ -107,6 +137,84 @@ def wrap( raise TypeError("agentcompile.wrap expects an OpenAI- or Anthropic-style client") +def transport( + wrapped: Any = None, + *, + lib: str | None = None, + key: str | None = None, + base_url: str | None = None, + company: str | None = None, + scrub: bool = True, + timeout: float = 2.0, +) -> Any: + """A transport that captures every model call through it, for frameworks whose client you + don't hold (LangChain, LiteLLM, CrewAI, Pydantic AI, ...). Capture only: calls are never + answered compiled. `wrapped`: the transport to send through (default the library's own). + `lib`: "httpx" (default) or "httpx2", the fork newer provider SDKs require; taken from + `wrapped` when it is given. The other arguments are wrap()'s; `timeout` bounds each + capture send, not your calls.""" + sync, _ = classes(lib_of(wrapped, lib)) + return sync(_capturer(key, base_url, company, scrub, timeout), wrapped) + + +def async_transport( + wrapped: Any = None, + *, + lib: str | None = None, + key: str | None = None, + base_url: str | None = None, + company: str | None = None, + scrub: bool = True, + timeout: float = 2.0, +) -> Any: + """The async version of transport().""" + _, async_ = classes(lib_of(wrapped, lib)) + return async_(_capturer(key, base_url, company, scrub, timeout), wrapped) + + +def http_client(*, lib: str = "httpx", **options: Any) -> Any: + """A client whose model calls are captured: pass it as the framework's `http_client`. + `lib="httpx2"` for SDKs built on httpx2. Takes transport()'s keyword arguments.""" + import importlib + + module = importlib.import_module(lib) + return module.Client(transport=transport(lib=lib, **options), timeout=_CALL_TIMEOUT) + + +def async_http_client(*, lib: str = "httpx", **options: Any) -> Any: + """An async client whose model calls are captured.""" + import importlib + + module = importlib.import_module(lib) + return module.AsyncClient( + transport=async_transport(lib=lib, **options), timeout=_CALL_TIMEOUT + ) + + +# The provider SDKs' own default: a model call can take minutes. +_CALL_TIMEOUT = 600.0 + + +def _settings( + key: str | None, base_url: str | None, company: str | None, timeout: float, mode: str +) -> Settings: + return Settings( + base_url=base_url or os.environ.get("AGENTCOMPILE_URL", DEFAULT_URL), + key=key or os.environ.get("AGENTCOMPILE_KEY"), + company=company or os.environ.get("AGENTCOMPILE_COMPANY"), + timeout=timeout, + mode=mode, + ) + + +def _capturer( + key: str | None, base_url: str | None, company: str | None, scrub: bool, timeout: float +) -> Any: + settings = _settings(key, base_url, company, timeout, LIVE) + _remember(settings) + return capturer_for(settings, None, load_key() if scrub else None) + + def flush(timeout: float = 5.0) -> None: """Send every captured call and outcome still queued (short scripts and tests; also runs at exit).""" diff --git a/src/agentcompile/_assemble.py b/src/agentcompile/_assemble.py index 5224441..8398a18 100644 --- a/src/agentcompile/_assemble.py +++ b/src/agentcompile/_assemble.py @@ -1,7 +1,7 @@ """Streamed answers, captured whole: each chunk passes to your agent untouched as it arrives, a copy is kept, and when the stream ends the pieces are assembled into the same shape a -non-streamed answer has (a chat completion, or an Anthropic message), so AgentCompile reads -streamed and non-streamed calls the same way.""" +non-streamed answer has (a chat completion, a Responses API response, or an Anthropic message), +so AgentCompile reads streamed and non-streamed calls the same way.""" from __future__ import annotations @@ -16,7 +16,11 @@ def assemble(provider: str, chunks: list[Any]) -> dict[str, Any] | None: try: - return _openai(chunks) if provider == "openai" else _anthropic(chunks) + if provider != "openai": + return _anthropic(chunks) + if any(str(c.get("type", "")).startswith("response.") for c in chunks): + return _responses(chunks) + return _openai(chunks) except Exception: # never break the agent over a shape we didn't expect return None @@ -60,6 +64,27 @@ def _openai(chunks: list[dict[str, Any]]) -> dict[str, Any]: return out +#: The events that end a Responses API stream, each carrying the whole response. +_RESPONSE_DONE = ("response.completed", "response.incomplete", "response.failed") + + +def _responses(events: list[dict[str, Any]]) -> dict[str, Any]: + """A Responses API stream: its last event carries the whole response. A stream stopped + early has none, so the response is rebuilt from the items it finished.""" + for event in reversed(events): + if event.get("type") in _RESPONSE_DONE and isinstance(event.get("response"), dict): + return dict(event["response"]) + response: dict[str, Any] = {} + items: dict[int, dict[str, Any]] = {} + for event in events: + if event.get("type") == "response.created" and isinstance(event.get("response"), dict): + response = dict(event["response"]) + elif event.get("type") == "response.output_item.done": + items[int(event.get("output_index", len(items)))] = dict(event.get("item") or {}) + response["output"] = [items[i] for i in sorted(items)] + return response + + def _anthropic(events: list[dict[str, Any]]) -> dict[str, Any]: message: dict[str, Any] = {} blocks: dict[int, dict[str, Any]] = {} diff --git a/src/agentcompile/_capture.py b/src/agentcompile/_capture.py index ddc6c6c..b930991 100644 --- a/src/agentcompile/_capture.py +++ b/src/agentcompile/_capture.py @@ -21,7 +21,7 @@ from ._assemble import assemble from ._decide import CONVERSATION_HEADER, Settings, _request -from ._payload import jsonable, payload +from ._payload import jsonable, request_of from ._scrub import scrub_call, scrub_value MAX_QUEUE = 2000 @@ -88,7 +88,7 @@ def add( """Queue one call. Never raises. `stream` without `streamed`: the answer wasn't kept (the record says so); `streamed`: the answer was assembled from its stream.""" try: - request, answer = payload(kwargs), None if stream else jsonable(response) + request, answer = request_of(kwargs), None if stream else jsonable(response) if self.scrub_key is not None: request = scrub_call(request, self.scrub_key) answer = scrub_call(answer, self.scrub_key) diff --git a/src/agentcompile/_core.py b/src/agentcompile/_core.py index 75a1118..7dc8659 100644 --- a/src/agentcompile/_core.py +++ b/src/agentcompile/_core.py @@ -250,6 +250,65 @@ async def create(*args: Any, **kwargs: Any) -> Any: return create +#: Why a Responses API call is never asked about: compiled answers come in Chat Completions' and +#: Anthropic's shapes, so these calls go to the model, captured like any other. +RESPONSES_API = "responses api" + + +def _forward_only(router: Router, conversation_id: str | None) -> tuple[str, str | None]: + return ( + ("no-conversation", None) if conversation_id is None else ("unsupported", RESPONSES_API) + ) + + +def sync_forward(router: Router, original: Callable[..., Any]) -> Callable[..., Any]: + """create() for an API we capture but don't answer: straight to the model.""" + + def create(*args: Any, **kwargs: Any) -> Any: + started = time.perf_counter() + conversation_id, real_kwargs = router.prepare(kwargs) + stream = bool(real_kwargs.get("stream")) + result = original(*args, **real_kwargs) + result = router.capture(conversation_id, real_kwargs, result, stream) + route, error = _forward_only(router, conversation_id) + router.record( + route=route, + conversation_id=conversation_id, + kwargs=real_kwargs, + decision=None, + decide_ms=None, + error=error, + started=started, + stream=stream, + ) + return result + + return create + + +def async_forward(router: Router, original: Callable[..., Any]) -> Callable[..., Any]: + async def create(*args: Any, **kwargs: Any) -> Any: + started = time.perf_counter() + conversation_id, real_kwargs = router.prepare(kwargs) + stream = bool(real_kwargs.get("stream")) + result = await original(*args, **real_kwargs) + result = router.capture(conversation_id, real_kwargs, result, stream, True) + route, error = _forward_only(router, conversation_id) + router.record( + route=route, + conversation_id=conversation_id, + kwargs=real_kwargs, + decision=None, + decide_ms=None, + error=error, + started=started, + stream=stream, + ) + return result + + return create + + class Proxy: """Everything delegates to the real object except the attributes we override.""" diff --git a/src/agentcompile/_payload.py b/src/agentcompile/_payload.py index 700ca5e..c30864d 100644 --- a/src/agentcompile/_payload.py +++ b/src/agentcompile/_payload.py @@ -6,6 +6,15 @@ # What the engine reads; everything else stays between the customer and their provider. _FIELDS = ("model", "messages", "system", "tools") +# A Responses API call (OpenAI's `responses.create`): the same, in that API's own words. +_RESPONSES_FIELDS = ( + "model", + "input", + "instructions", + "tools", + "previous_response_id", + "conversation", +) def jsonable(value: Any) -> Any: @@ -24,3 +33,14 @@ def jsonable(value: Any) -> Any: def payload(kwargs: dict[str, Any]) -> dict[str, Any]: return {k: jsonable(kwargs[k]) for k in _FIELDS if k in kwargs and kwargs[k] is not None} + + +def request_of(kwargs: dict[str, Any]) -> dict[str, Any]: + """What capture keeps of a call: `payload`, or a Responses API call's own fields.""" + if "input" in kwargs and "messages" not in kwargs: + return { + k: jsonable(kwargs[k]) + for k in _RESPONSES_FIELDS + if k in kwargs and kwargs[k] is not None + } + return payload(kwargs) diff --git a/src/agentcompile/_transport.py b/src/agentcompile/_transport.py new file mode 100644 index 0000000..5f36a19 --- /dev/null +++ b/src/agentcompile/_transport.py @@ -0,0 +1,294 @@ +"""Capture from any framework: an HTTP transport that sees the model calls a framework makes +(LangChain, LiteLLM, CrewAI, Pydantic AI, the OpenAI Agents SDK, your own loop) and captures +them, with no client to wrap. Capture only: every call goes to the model unchanged and is +never answered compiled (that takes `wrap`). + + http = agentcompile.http_client(key="ack_...") # async: agentcompile.async_http_client() + llm = ChatOpenAI(model="gpt-5", http_client=http) + with agentcompile.conversation(ticket.id): + agent.invoke(...) + +Seen: POSTs to `/chat/completions`, `/responses` (OpenAI and anything OpenAI-compatible) and +`/messages` (Anthropic). Everything else passes through untouched. Failed calls (not 2xx) are +not captured. Nothing here ever raises into the framework. + +Built for httpx and for httpx2 (the newer provider SDKs' fork, which refuses httpx objects): +the same classes, made once per library. +""" + +from __future__ import annotations + +import importlib +import json +import zlib +from collections.abc import AsyncIterator, Callable, Iterator +from functools import cache +from typing import Any + +import httpx + +from ._conversation import current, current_customer + + +def provider_of(request: Any) -> str | None: + """Which provider's shape a request is, or None for a request we don't capture.""" + if request.method != "POST": + return None + path = request.url.path.rstrip("/") + if path.endswith(("/chat/completions", "/responses")): + return "openai" + # Anthropic's create; not OpenAI's Assistants threads, which also end in /messages. + if path.endswith("/messages") and "/threads/" not in path: + return "anthropic" + return None + + +def body_of(request: Any) -> dict[str, Any] | None: + try: + body = json.loads(request.content) + except Exception: # a body not read yet, or not JSON: not a model call we read + return None + if isinstance(body, dict) and ("messages" in body or "input" in body): + return body + return None + + +def decode(raw: bytes, encoding: str) -> bytes | None: + """The body as sent, before any content encoding; None for one we can't undo. A body cut + short (a stream closed early) decodes as far as it goes.""" + for coding in reversed([c.strip().lower() for c in encoding.split(",") if c.strip()]): + if coding == "identity": + continue + if coding in ("gzip", "x-gzip", "deflate"): + try: + raw = zlib.decompressobj(47).decompress(raw) # gzip or zlib, by header + except zlib.error: + raw = zlib.decompressobj(-15).decompress(raw) # raw deflate + elif coding == "br": + try: + import brotli # type: ignore[import-not-found] + except ImportError: + return None + raw = brotli.decompress(raw) + else: + return None + return raw + + +def events(text: str) -> list[dict[str, Any]]: + """The JSON objects in a server-sent event stream, in order ([DONE] and comments + skipped).""" + found: list[dict[str, Any]] = [] + for block in text.replace("\r\n", "\n").split("\n\n"): + data = "\n".join( + line[5:].lstrip() for line in block.split("\n") if line.startswith("data:") + ) + if not data or data == "[DONE]": + continue + try: + value = json.loads(data) + except ValueError: + continue + if isinstance(value, dict): + found.append(value) + return found + + +_ENDS = ("response.completed", "response.incomplete", "response.failed", "message_stop") + + +def ended(text: str, chunks: list[dict[str, Any]]) -> bool: + """Whether a stream reached its own end marker. SDKs stop reading at the marker, before + the last bytes, so a stream can be whole without its transport stream running out.""" + if "data: [DONE]" in text or any(c.get("type") in _ENDS for c in chunks): + return True + last = chunks[-1] if chunks else {} + return any(choice.get("finish_reason") for choice in last.get("choices") or []) + + +class _Seen: + """One call in flight: what to capture once its answer is whole.""" + + def __init__(self, provider: str, body: dict[str, Any], response: Any) -> None: + self.provider = provider + self.body = body + self.encoding = response.headers.get("content-encoding", "") + self.raw: list[bytes] = [] + self.finished = False + + +class _Capture: + """What every transport shares: which calls to see, and handing them to the capturer.""" + + def __init__(self, capturer: Any) -> None: + self.capturer = capturer + + def seen(self, request: Any, response: Any) -> _Seen | None: + try: + provider = provider_of(request) + if provider is None or not 200 <= response.status_code < 300: + return None + body = body_of(request) + return None if body is None else _Seen(provider, body, response) + except Exception: + return None + + def whole(self, seen: _Seen, response: Any) -> None: + """A call answered in one body (already read).""" + try: + self.capturer.add( + seen.provider, + current(), + seen.body, + response.json(), + customer=current_customer(), + ) + except Exception: + self.capturer.dropped += 1 + + def streamed(self) -> Callable[[_Seen, bool], None]: + """The callback for a stream: its conversation is the one the call was made in, not + whatever is current when the framework finishes reading.""" + conversation, customer = current(), current_customer() + + def done(seen: _Seen, ran_out: bool) -> None: + if seen.finished: + return + seen.finished = True + try: + raw = decode(b"".join(seen.raw), seen.encoding) + if raw is None: + self.capturer.dropped += 1 + return + text = raw.decode("utf-8", "replace") + chunks = events(text) + self.capturer.add_stream( + seen.provider, + conversation, + seen.body, + chunks, + ran_out or ended(text, chunks), + customer, + ) + except Exception: + self.capturer.dropped += 1 + + return done + + +@cache +def classes(lib_name: str) -> tuple[type, type]: + """(CaptureTransport, AsyncCaptureTransport) for one HTTP library: "httpx" or "httpx2".""" + lib = httpx if lib_name == "httpx" else importlib.import_module(lib_name) + + def already_read(seen: _Seen, response: Any, done: Callable[[_Seen, bool], None]) -> bool: + """A stream the transport below already read whole (a mock, a cache, a replay): + captured from its body now, since nothing will flow through a tee.""" + try: + content = response.content + except lib.ResponseNotRead: + return False + seen.raw, seen.encoding = [content], "" # the library has already decoded it + done(seen, True) + return True + + class Tee(lib.SyncByteStream): # type: ignore[misc, name-defined] + def __init__(self, stream: Any, seen: _Seen, done: Callable[[_Seen, bool], None]): + self._stream, self._seen, self._done = stream, seen, done + + def __iter__(self) -> Iterator[bytes]: + ran_out = False + try: + for chunk in self._stream: + self._seen.raw.append(chunk) + yield chunk + ran_out = True + finally: + self._done(self._seen, ran_out) + + def close(self) -> None: + self._done(self._seen, False) + close = getattr(self._stream, "close", None) + if close is not None: + close() + + class AsyncTee(lib.AsyncByteStream): # type: ignore[misc, name-defined] + def __init__(self, stream: Any, seen: _Seen, done: Callable[[_Seen, bool], None]): + self._stream, self._seen, self._done = stream, seen, done + + async def __aiter__(self) -> AsyncIterator[bytes]: + ran_out = False + try: + async for chunk in self._stream: + self._seen.raw.append(chunk) + yield chunk + ran_out = True + finally: + self._done(self._seen, ran_out) + + async def aclose(self) -> None: + self._done(self._seen, False) + aclose = getattr(self._stream, "aclose", None) + if aclose is not None: + await aclose() + + class CaptureTransport(lib.BaseTransport): # type: ignore[misc, name-defined] + """A transport that captures the model calls passing through it.""" + + def __init__(self, capturer: Any, wrapped: Any = None) -> None: + self._capture = _Capture(capturer) + self._wrapped = wrapped or lib.HTTPTransport() + + def handle_request(self, request: Any) -> Any: + response = self._wrapped.handle_request(request) + seen = self._capture.seen(request, response) + if seen is None: + return response + if seen.body.get("stream"): + done = self._capture.streamed() + if not already_read(seen, response, done): + response.stream = Tee(response.stream, seen, done) + return response + response.read() + self._capture.whole(seen, response) + return response + + def close(self) -> None: + self._wrapped.close() + + class AsyncCaptureTransport(lib.AsyncBaseTransport): # type: ignore[misc, name-defined] + """The async version.""" + + def __init__(self, capturer: Any, wrapped: Any = None) -> None: + self._capture = _Capture(capturer) + self._wrapped = wrapped or lib.AsyncHTTPTransport() + + async def handle_async_request(self, request: Any) -> Any: + response = await self._wrapped.handle_async_request(request) + seen = self._capture.seen(request, response) + if seen is None: + return response + if seen.body.get("stream"): + done = self._capture.streamed() + if not already_read(seen, response, done): + response.stream = AsyncTee(response.stream, seen, done) + return response + await response.aread() + self._capture.whole(seen, response) + return response + + async def aclose(self) -> None: + await self._wrapped.aclose() + + return CaptureTransport, AsyncCaptureTransport + + +def lib_of(wrapped: Any, lib: str | None) -> str: + """The library a transport is for: as named, else the wrapped transport's own.""" + if lib: + return lib + module = type(wrapped).__module__.partition(".")[0] if wrapped is not None else "" + return module if module in ("httpx", "httpx2") else "httpx" + + +CaptureTransport, AsyncCaptureTransport = classes("httpx") diff --git a/tests/test_capture_everywhere.py b/tests/test_capture_everywhere.py new file mode 100644 index 0000000..8c9a9f1 --- /dev/null +++ b/tests/test_capture_everywhere.py @@ -0,0 +1,578 @@ +"""Capture beyond chat.completions and messages: the Responses API through wrap(), and any +framework through the capturing httpx transport (no client to wrap).""" + +from __future__ import annotations + +import gzip +import json +from pathlib import Path +from typing import Any + +import anthropic +import httpx +import openai +import pytest + +import agentcompile +from agentcompile import _capture +from agentcompile._decide import Settings +from agentcompile._transport import decode, events, provider_of + +RESPONSE = { + "id": "resp_1", + "object": "response", + "created_at": 0, + "model": "gpt-x", + "status": "completed", + "output": [ + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "get_order", + "arguments": '{"order_id": "#W1"}', + "status": "completed", + } + ], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], +} +COMPLETION = { + "id": "c1", + "object": "chat.completion", + "created": 0, + "model": "gpt-x", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "from the model"}, + "finish_reason": "stop", + } + ], +} +MESSAGE = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-x", + "content": [{"type": "text", "text": "from the model"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, +} +TOOLS = [{"type": "function", "name": "get_order", "parameters": {"type": "object"}}] + + +def sse(events_: list[dict[str, Any]], done: bool = False) -> bytes: + lines = [f"event: {e.get('type', 'message')}\ndata: {json.dumps(e)}\n\n" for e in events_] + if done: + lines.append("data: [DONE]\n\n") + return "".join(lines).encode() + + +RESPONSE_EVENTS = [ + { + "type": "response.created", + "response": {**RESPONSE, "status": "in_progress", "output": []}, + }, + {"type": "response.output_item.added", "output_index": 0, "item": RESPONSE["output"][0]}, + {"type": "response.output_item.done", "output_index": 0, "item": RESPONSE["output"][0]}, + {"type": "response.completed", "response": RESPONSE}, +] + + +class Service: + """AgentCompile: records capture batches; any decide call is a test failure.""" + + def __init__(self) -> None: + self.batches: list[dict[str, Any]] = [] + self.decides = 0 + + def __call__(self, request: httpx.Request) -> httpx.Response: + if request.url.path == "/v1/capture": + self.batches.append(json.loads(request.content)) + return httpx.Response(200, json={"kept": 1}) + self.decides += 1 + return httpx.Response(200, json={"action": "forward", "reason": "x"}) + + @property + def records(self) -> list[dict[str, Any]]: + return [r for b in self.batches for r in b["exchanges"]] + + +class Provider: + """The model: answers each path with its own shape, streamed when asked.""" + + def __init__(self, status: int = 200, encoding: str = "", preread: bool = False) -> None: + self.calls: list[httpx.Request] = [] + self.status, self.encoding, self.preread = status, encoding, preread + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.calls.append(request) + body = json.loads(request.content) if request.content else {} + path = request.url.path + if self.status != 200: + return httpx.Response(self.status, json={"error": {"message": "nope"}}) + if path.endswith("/responses"): + if body.get("stream"): + return self._sse(sse(RESPONSE_EVENTS)) + return self._json(RESPONSE) + if path.endswith("/messages"): + if body.get("stream"): + return self._sse(sse(_anthropic_events())) + return self._json(MESSAGE) + if path.endswith("/chat/completions"): + if body.get("stream"): + return self._sse(sse(_chat_chunks(), done=True)) + return self._json(COMPLETION) + return httpx.Response(200, json={"data": []}) + + def _json(self, data: dict[str, Any]) -> httpx.Response: + raw = json.dumps(data).encode() + headers = {"content-type": "application/json"} + if self.encoding == "gzip": + raw, headers["content-encoding"] = gzip.compress(raw), "gzip" + return httpx.Response(200, content=raw, headers=headers) + + def _sse(self, raw: bytes) -> httpx.Response: + headers = {"content-type": "text/event-stream"} + if self.encoding == "gzip": + raw, headers["content-encoding"] = gzip.compress(raw), "gzip" + if self.preread: + return httpx.Response(200, content=raw, headers=headers) + return httpx.Response(200, stream=Wire(raw), headers=headers) + + +class Wire(httpx.SyncByteStream, httpx.AsyncByteStream): + """A body still on the wire, in small pieces (bytes given to httpx.Response arrive read).""" + + def __init__(self, raw: bytes) -> None: + self.pieces = [raw[i : i + 7] for i in range(0, len(raw), 7)] + + def __iter__(self) -> Any: + yield from self.pieces + + async def __aiter__(self) -> Any: + for piece in self.pieces: + yield piece + + +def _chat_chunks() -> list[dict[str, Any]]: + base = {"id": "c1", "object": "chat.completion.chunk", "created": 0, "model": "gpt-x"} + return [ + {**base, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "from "}}]}, + {**base, "choices": [{"index": 0, "delta": {"content": "the model"}}]}, + {**base, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + ] + + +def _anthropic_events() -> list[dict[str, Any]]: + return [ + {"type": "message_start", "message": {**MESSAGE, "content": [], "stop_reason": None}}, + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "from the model"}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 3}, + }, + {"type": "message_stop"}, + ] + + +@pytest.fixture(autouse=True) +def _isolated(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Any: + monkeypatch.setenv("AGENTCOMPILE_SCRUB_KEY", "test-scrub-key") + yield + _capture._ALL.clear() + + +def _wrapped(service: Service, provider: Provider, tmp_path: Path, **kw: Any) -> Any: + real = openai.OpenAI( + api_key="sk-test", + base_url="http://provider.test/v1", + http_client=httpx.Client(transport=httpx.MockTransport(provider)), + ) + return agentcompile.wrap( + real, + key="ack_acme.k", + base_url="http://ac.test", + trail=tmp_path / "trail.jsonl", + http_client=httpx.Client(transport=httpx.MockTransport(service)), + capture=True, + **kw, + ) + + +def _trail(tmp_path: Path) -> list[dict[str, Any]]: + return [json.loads(x) for x in (tmp_path / "trail.jsonl").read_text().splitlines()] + + +def _capturing(service: Service) -> _capture.Capturer: + """A sender flush_all() reaches, unscrubbed so the tests can read what was sent.""" + return _capture.capturer_for( + Settings("http://ac.test", "ack_acme.k", None, 1.0), + httpx.Client(transport=httpx.MockTransport(service)), + None, + ) + + +# -- the Responses API through wrap() ------------------------------------------------------- + + +def test_a_responses_call_is_captured_in_its_own_shape_and_never_asked_about( + tmp_path: Path, +) -> None: + service, provider = Service(), Provider() + client = _wrapped(service, provider, tmp_path) + resp = client.responses.create( + model="gpt-x", + input=[{"role": "user", "content": "where is order #W1? I'm ann@example.com"}], + instructions="You are a support agent.", + tools=TOOLS, + previous_response_id="resp_0", + conversation_id="c1", + customer_id="ann@example.com", + ) + assert resp.output[0].name == "get_order" + sent = json.loads(provider.calls[0].content) + assert "conversation_id" not in sent and "customer_id" not in sent + agentcompile.flush(2) + (record,) = service.records + assert service.decides == 0 + assert record["provider"] == "openai" and record["conversation_id"] == "c1" + request = record["request"] + assert set(request) == {"model", "input", "instructions", "tools", "previous_response_id"} + assert "ann@example.com" not in json.dumps(record) # scrubbed here, before sending + assert record["end_user"].startswith(" None: + service, provider = Service(), Provider() + client = _wrapped(service, provider, tmp_path) + stream = client.responses.create( + model="gpt-x", input="where is my order?", stream=True, conversation_id="c1" + ) + kinds = [event.type for event in stream] + assert kinds[-1] == "response.completed" + agentcompile.flush(2) + (record,) = service.records + assert record["stream"] is True and record["stream_complete"] is True + assert record["response"]["id"] == "resp_1" + assert record["response"]["output"][0]["name"] == "get_order" + assert record["request"]["input"] == "where is my order?" + + +def test_a_responses_stream_stopped_early_keeps_the_items_it_finished() -> None: + from agentcompile._assemble import assemble + + built = assemble("openai", RESPONSE_EVENTS[:3]) + assert built is not None + assert built["id"] == "resp_1" and built["output"][0]["call_id"] == "call_1" + + +def test_without_a_conversation_a_responses_call_is_still_captured(tmp_path: Path) -> None: + service, provider = Service(), Provider() + _wrapped(service, provider, tmp_path).responses.create(model="gpt-x", input="hi") + agentcompile.flush(2) + assert service.records[0]["conversation_id"] is None + assert _trail(tmp_path)[0]["route"] == "no-conversation" + + +async def test_the_async_client_captures_responses_too(tmp_path: Path) -> None: + service, provider = Service(), Provider() + + async def answer(request: httpx.Request) -> httpx.Response: + return provider(request) + + real = openai.AsyncOpenAI( + api_key="sk-test", + base_url="http://provider.test/v1", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(answer)), + ) + client = agentcompile.wrap( + real, + key="ack_acme.k", + base_url="http://ac.test", + trail=False, + http_client=httpx.Client(transport=httpx.MockTransport(service)), + capture=True, + ) + await client.responses.create(model="gpt-x", input="hi", conversation_id="c9") + stream = await client.responses.create( + model="gpt-x", input="hi", stream=True, conversation_id="c9" + ) + async for _ in stream: + pass + agentcompile.flush(2) + assert [r["response"]["id"] for r in service.records] == ["resp_1", "resp_1"] + assert service.records[1]["stream_complete"] is True + + +def test_chat_completions_still_decide_on_a_client_that_has_responses(tmp_path: Path) -> None: + service, provider = Service(), Provider() + client = _wrapped(service, provider, tmp_path) + client.chat.completions.create( + model="gpt-x", messages=[{"role": "user", "content": "hi"}], conversation_id="c1" + ) + assert service.decides == 1 + + +# -- any framework, through the transport ----------------------------------------------------- + + +def _framework_client(service: Service, provider: Provider) -> openai.OpenAI: + """A framework's own OpenAI client, built with the capturing transport: no wrap().""" + transport = agentcompile.CaptureTransport( + _capturing(service), httpx.MockTransport(provider) + ) + return openai.OpenAI( + api_key="sk-test", + base_url="http://provider.test/v1", + http_client=httpx.Client(transport=transport), + ) + + +def test_the_transport_captures_a_framework_call_in_its_conversation() -> None: + service, provider = Service(), Provider() + client = _framework_client(service, provider) + with agentcompile.conversation("t1", customer="cust-7"): + resp = client.chat.completions.create( + model="gpt-x", messages=[{"role": "user", "content": "hi"}] + ) + assert resp.choices[0].message.content == "from the model" + _capture.flush_all(2) + (record,) = service.records + assert record["conversation_id"] == "t1" and record["end_user"] == "cust-7" + assert record["request"]["messages"][0]["content"] == "hi" + assert record["response"]["choices"][0]["message"]["content"] == "from the model" + assert service.decides == 0 # capture only + + +def test_the_transport_assembles_streams_in_every_shape() -> None: + service, provider = Service(), Provider() + client = _framework_client(service, provider) + with agentcompile.conversation("t2"): + chunks = list( + client.chat.completions.create( + model="gpt-x", messages=[{"role": "user", "content": "hi"}], stream=True + ) + ) + list(client.responses.create(model="gpt-x", input="hi", stream=True)) + assert len(chunks) == 3 + _capture.flush_all(2) + chat, responses = service.records + assert chat["response"]["choices"][0]["message"]["content"] == "from the model" + assert chat["stream_complete"] is True + assert responses["response"]["id"] == "resp_1" + + +def test_the_transport_reads_compressed_answers() -> None: + service, provider = Service(), Provider(encoding="gzip") + client = _framework_client(service, provider) + with agentcompile.conversation("t3"): + resp = client.chat.completions.create( + model="gpt-x", messages=[{"role": "user", "content": "hi"}] + ) + streamed = list( + client.chat.completions.create( + model="gpt-x", messages=[{"role": "user", "content": "hi"}], stream=True + ) + ) + assert resp.choices[0].message.content == "from the model" and len(streamed) == 3 + _capture.flush_all(2) + assert [r["response"]["choices"][0]["message"]["content"] for r in service.records] == [ + "from the model", + "from the model", + ] + + +def test_the_transport_captures_anthropic_calls_over_httpx2() -> None: + httpx2 = pytest.importorskip("httpx2") # only the newer provider SDKs bring it + + service, provider = Service(), Provider() + + def answer(request: Any) -> Any: + mine = provider( + httpx.Request(request.method, str(request.url), content=request.content) + ) + raw = mine.read() + on_the_wire = iter([raw[i : i + 7] for i in range(0, len(raw), 7)]) + return httpx2.Response( + mine.status_code, headers=dict(mine.headers), content=on_the_wire + ) + + sync, _ = agentcompile._transport.classes("httpx2") + transport = sync(_capturing(service), httpx2.MockTransport(answer)) + client = anthropic.Anthropic( + api_key="sk-test", + base_url="http://provider.test", + http_client=httpx2.Client(transport=transport), + ) + with agentcompile.conversation("t4"): + client.messages.create( + model="claude-x", max_tokens=10, messages=[{"role": "user", "content": "hi"}] + ) + with client.messages.stream( + model="claude-x", max_tokens=10, messages=[{"role": "user", "content": "hi"}] + ) as stream: + assert stream.get_final_text() == "from the model" + _capture.flush_all(2) + whole, streamed = service.records + assert whole["provider"] == streamed["provider"] == "anthropic" + assert streamed["response"]["content"][0]["text"] == "from the model" + assert streamed["stream_complete"] is True + + +def test_the_library_is_taken_from_the_wrapped_transport() -> None: + httpx2 = pytest.importorskip("httpx2") # only the newer provider SDKs bring it + + from agentcompile._transport import lib_of + + assert lib_of(httpx2.HTTPTransport(), None) == "httpx2" + assert lib_of(httpx.HTTPTransport(), None) == "httpx" + assert lib_of(None, None) == "httpx" and lib_of(None, "httpx2") == "httpx2" + made = agentcompile.transport(httpx2.HTTPTransport(), key="ack_acme.k") + assert isinstance(made, httpx2.BaseTransport) + + +def test_the_transport_leaves_other_calls_and_failures_alone() -> None: + service = Service() + client = _framework_client(service, Provider()) + client.models.list() # not a model call + failing = _framework_client(service, Provider(status=400)) + with pytest.raises(openai.BadRequestError): + failing.chat.completions.create( + model="gpt-x", messages=[{"role": "user", "content": "hi"}] + ) + _capture.flush_all(2) + assert service.records == [] + + +async def test_the_async_transport_captures_whole_and_streamed() -> None: + service, provider = Service(), Provider() + + async def answer(request: httpx.Request) -> httpx.Response: + return provider(request) + + transport = agentcompile.AsyncCaptureTransport( + _capturing(service), httpx.MockTransport(answer) + ) + client = openai.AsyncOpenAI( + api_key="sk-test", + base_url="http://provider.test/v1", + http_client=httpx.AsyncClient(transport=transport), + ) + with agentcompile.conversation("t5"): + await client.responses.create(model="gpt-x", input="hi") + stream = await client.chat.completions.create( + model="gpt-x", messages=[{"role": "user", "content": "hi"}], stream=True + ) + async for _ in stream: + pass + _capture.flush_all(2) + whole, streamed = service.records + assert whole["conversation_id"] == streamed["conversation_id"] == "t5" + assert streamed["response"]["choices"][0]["message"]["content"] == "from the model" + + +def test_a_stream_a_transport_already_read_is_captured_from_its_body() -> None: + service = Service() + client = _framework_client(service, Provider(preread=True)) + with agentcompile.conversation("t7"): + list( + client.chat.completions.create( + model="gpt-x", messages=[{"role": "user", "content": "hi"}], stream=True + ) + ) + _capture.flush_all(2) + (record,) = service.records + assert record["response"]["choices"][0]["message"]["content"] == "from the model" + + +def test_a_stream_belongs_to_the_conversation_it_was_started_in() -> None: + service, provider = Service(), Provider() + client = _framework_client(service, provider) + with agentcompile.conversation("started-here"): + stream = client.chat.completions.create( + model="gpt-x", messages=[{"role": "user", "content": "hi"}], stream=True + ) + with agentcompile.conversation("read-here"): + list(stream) + _capture.flush_all(2) + assert service.records[0]["conversation_id"] == "started-here" + + +def test_a_stream_closed_early_is_captured_once_and_marked() -> None: + service, provider = Service(), Provider() + client = _framework_client(service, provider) + with agentcompile.conversation("t6"): + stream = client.chat.completions.create( + model="gpt-x", messages=[{"role": "user", "content": "hi"}], stream=True + ) + next(iter(stream)) + stream.close() + _capture.flush_all(2) + (record,) = service.records + assert record["stream_complete"] is False + + +def test_http_client_helpers_build_capturing_clients(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("AGENTCOMPILE_KEY", "ack_acme.k") + sync, async_ = agentcompile.http_client(), agentcompile.async_http_client() + assert isinstance(sync._transport, agentcompile.CaptureTransport) + assert isinstance(async_._transport, agentcompile.AsyncCaptureTransport) + assert sync.timeout.read == 600.0 + httpx2 = pytest.importorskip("httpx2") # only the newer provider SDKs bring it + assert isinstance(agentcompile.http_client(lib="httpx2"), httpx2.Client) + assert isinstance(agentcompile.async_http_client(lib="httpx2"), httpx2.AsyncClient) + + +# -- the pieces ---------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("method", "path", "expected"), + [ + ("POST", "/v1/chat/completions", "openai"), + ("POST", "/v1beta/openai/chat/completions", "openai"), + ("POST", "/openai/deployments/d/chat/completions", "openai"), + ("POST", "/v1/responses", "openai"), + ("POST", "/v1/messages", "anthropic"), + ("POST", "/v1/messages/count_tokens", None), + ("POST", "/v1/threads/t1/messages", None), + ("GET", "/v1/responses", None), + ("POST", "/v1/embeddings", None), + ], +) +def test_which_requests_are_model_calls(method: str, path: str, expected: str | None) -> None: + assert provider_of(httpx.Request(method, f"http://x{path}")) == expected + + +def test_decode_undoes_gzip_deflate_and_refuses_the_unknown() -> None: + import zlib + + assert decode(gzip.compress(b"hi"), "gzip") == b"hi" + assert decode(zlib.compress(b"hi"), "deflate") == b"hi" + raw = zlib.compressobj(wbits=-15) + assert decode(raw.compress(b"hi") + raw.flush(), "deflate") == b"hi" + assert decode(b"hi", "") == b"hi" and decode(b"hi", "identity") == b"hi" + assert decode(b"hi", "zstd") is None + + +def test_events_reads_server_sent_events() -> None: + text = 'event: a\ndata: {"x": 1}\n\n: comment\n\ndata: [DONE]\n\ndata: not json\n\n' + assert events(text + 'data: {"y":\ndata: 2}\n\n') == [{"x": 1}, {"y": 2}]