diff --git a/README.md b/README.md index 0782103..5f78a11 100644 --- a/README.md +++ b/README.md @@ -29,11 +29,12 @@ with agentcompile.conversation(ticket.id): ## What happens on each call -1. With a conversation id, the SDK asks AgentCompile for a decision (`/v1/decide`; it waits at most `timeout`, default 2 s). -2. **compiled**: the answer (a tool call or a reply) comes back in exactly the provider's response shape; your loop runs the tool as usual. +1. With a conversation id, the SDK asks AgentCompile for a decision (`/v1/decide`). The whole call takes at most `timeout` (default 2 s), and AgentCompile is told that, so it answers in time. +2. **compiled**: the answer (a tool call or a reply) comes back in exactly the provider's response shape; your loop runs the tool as usual. A tool call is only ever one of the `tools` you passed. 3. **forwarded**: your model is called as normal. -4. **fail-open**: AgentCompile errored, was slow or sent something unusable, so your model is called as normal. -5. **no conversation id**: your model is called and no decision is made. +4. **fail-open**: AgentCompile errored, was slow or sent something unusable, so your model is called as normal. After 5 failures in a row the SDK stops asking for 30 s, then tries one call. +5. **unsupported**: the request asks for something a compiled answer couldn't honour (`n` > 1, a forced `tool_choice`, a `response_format`), so your model is called without asking. +6. **no conversation id**: your model is called and no decision is made. Your provider key never leaves your process. `mode="shadow"` decides but always calls your model, so you can see what it would have done before you switch it on. diff --git a/src/agentcompile/__init__.py b/src/agentcompile/__init__.py index c989444..8218cae 100644 --- a/src/agentcompile/__init__.py +++ b/src/agentcompile/__init__.py @@ -15,7 +15,7 @@ import os from typing import Any -from ._capture import Capturer, flush_all +from ._capture import capturer_for, flush_all from ._conversation import conversation from ._core import LIVE, SHADOW, Proxy, Router, async_create, sync_create from ._decide import AsyncDecider, Decider, Settings @@ -65,6 +65,7 @@ def wrap( key=key or os.environ.get("AGENTCOMPILE_KEY"), company=company or os.environ.get("AGENTCOMPILE_COMPANY"), timeout=timeout, + mode=mode, ) the_trail = Trail(trail, on_event) _remember(settings) # outcome() sends with the client wrapped last @@ -72,7 +73,7 @@ def wrap( capture = os.environ.get("AGENTCOMPILE_CAPTURE", "") in ("1", "true", "yes") # One sender per wrapped client; it always sends from its own thread (sync client). capturer = ( - Capturer( + capturer_for( settings, http_client if _is_sync_http(http_client) else None, load_key() if scrub else None, diff --git a/src/agentcompile/_assemble.py b/src/agentcompile/_assemble.py index bd59d1e..5224441 100644 --- a/src/agentcompile/_assemble.py +++ b/src/agentcompile/_assemble.py @@ -101,8 +101,9 @@ def __init__(self, stream: Any, on_done: OnDone) -> None: self._on_done = on_done self._chunks: list[Any] = [] self._finished = False + self._it = self._iterate() # one pass, shared by `for` and next() - def __iter__(self) -> Iterator[Any]: + def _iterate(self) -> Iterator[Any]: complete = False try: for item in self._stream: @@ -112,6 +113,12 @@ def __iter__(self) -> Iterator[Any]: finally: self._finish(complete=complete) + def __iter__(self) -> Iterator[Any]: + return self._it + + def __next__(self) -> Any: + return next(self._it) + def __enter__(self) -> CapturingStream: enter = getattr(self._stream, "__enter__", None) if enter is not None: @@ -146,6 +153,7 @@ def __init__(self, stream: Any, on_done: OnDone) -> None: self._on_done = on_done self._chunks: list[Any] = [] self._finished = False + self._it = self._iterate() async def _iterate(self) -> AsyncIterator[Any]: complete = False @@ -158,7 +166,10 @@ async def _iterate(self) -> AsyncIterator[Any]: self._finish(complete=complete) def __aiter__(self) -> AsyncIterator[Any]: - return self._iterate() + return self._it + + async def __anext__(self) -> Any: + return await self._it.__anext__() async def __aenter__(self) -> AsyncCapturingStream: enter = getattr(self._stream, "__aenter__", None) diff --git a/src/agentcompile/_capture.py b/src/agentcompile/_capture.py index 4e6cde5..ae03b79 100644 --- a/src/agentcompile/_capture.py +++ b/src/agentcompile/_capture.py @@ -11,6 +11,7 @@ import atexit import collections +import json import threading import time from datetime import datetime, timezone @@ -25,17 +26,35 @@ MAX_QUEUE = 2000 BATCH = 50 +# The server takes up to 8 MiB a request: batches stay under 4 MiB, and a call bigger than +# MAX_RECORD_BYTES on its own (a huge history) is dropped and counted. +MAX_BATCH_BYTES = 4 * 1024 * 1024 +MAX_RECORD_BYTES = 7 * 1024 * 1024 INTERVAL = 1.0 -_ALL: list[Capturer] = [] +_ALL: dict[tuple[Any, ...], Capturer] = {} +_ALL_LOCK = threading.Lock() def flush_all(timeout: float = 5.0) -> None: - for capturer in list(_ALL): + for capturer in list(_ALL.values()): capturer.flush(timeout) +def capturer_for( + settings: Settings, http: httpx.Client | None, scrub_key: bytes | None +) -> Capturer: + """One sender (one queue, one thread) per destination: wrapping a client per request, as + some apps do, reuses it instead of starting another each time.""" + key = (settings.base_url, settings.key, settings.company, scrub_key, id(http)) + with _ALL_LOCK: + found = _ALL.get(key) + if found is None: + found = _ALL[key] = Capturer(settings, http, scrub_key) + return found + + class Capturer: def __init__( self, @@ -55,7 +74,6 @@ def __init__( self.dropped = 0 self.failed = 0 atexit.register(self.flush, 2.0) - _ALL.append(self) def add( self, @@ -144,18 +162,43 @@ def _run(self) -> None: while self._queue: self._send_batch() - def _send_batch(self) -> None: + def _take(self) -> list[str]: + """Up to BATCH queued calls, encoded, under MAX_BATCH_BYTES together (the first one + alone may be up to MAX_RECORD_BYTES); what doesn't fit goes back to the front.""" with self._lock: - batch = [self._queue.popleft() for _ in range(min(BATCH, len(self._queue)))] + records = [self._queue.popleft() for _ in range(min(BATCH, len(self._queue)))] + batch: list[str] = [] + size = 0 + for i, record in enumerate(records): + try: + line = json.dumps(record, default=str) + except Exception: + self.dropped += 1 + continue + if len(line) > MAX_RECORD_BYTES: + self.dropped += 1 + continue + if batch and size + len(line) > MAX_BATCH_BYTES: + with self._lock: # the rest waits for the next batch, in order + self._queue.extendleft(reversed(records[i:])) + break + batch.append(line) + size += len(line) + return batch + + def _send_batch(self) -> None: + batch = self._take() if not batch: return - url, headers, _ = _request(self.settings, batch[0]["provider"], "capture", {}) + url, headers, _ = _request(self.settings, "openai", "capture", {}) headers.pop(CONVERSATION_HEADER, None) + headers["content-type"] = "application/json" url = url.rsplit("/v1/decide", 1)[0] + "/v1/capture" try: if self._http is None: self._http = httpx.Client(timeout=self.settings.timeout) - response = self._http.post(url, headers=headers, json={"exchanges": batch}) + body = '{"exchanges": [' + ", ".join(batch) + "]}" + response = self._http.post(url, headers=headers, content=body.encode("utf-8")) if response.status_code == 200: self.sent += len(batch) else: diff --git a/src/agentcompile/_core.py b/src/agentcompile/_core.py index f1bef0c..75a1118 100644 --- a/src/agentcompile/_core.py +++ b/src/agentcompile/_core.py @@ -15,6 +15,46 @@ LIVE, SHADOW = "live", "shadow" +def _given(value: Any) -> Any: + """None for an argument left out, however the SDK spells that (NOT_GIVEN, omit).""" + return None if type(value).__name__ in ("NotGiven", "Omit") else value + + +def unsupported(kwargs: dict[str, Any]) -> str | None: + """Why a compiled answer couldn't honour this request (then it goes to the model, and no + decision is asked for): several answers, a forced or forbidden tool, a fixed output + shape.""" + n = _given(kwargs.get("n")) + if n is not None and n != 1: + return "n" + if _given(kwargs.get("functions")) is not None: + return "functions" + shape = _given(kwargs.get("response_format")) + if shape is not None and not (isinstance(shape, dict) and shape.get("type") == "text"): + return "response_format" + choice = _given(kwargs.get("tool_choice")) + auto = choice in (None, "auto") or ( + isinstance(choice, dict) and choice.get("type") == "auto" + ) + if not auto: + return "tool_choice" + return None + + +def offered(kwargs: dict[str, Any], tool: str | None) -> bool: + """Whether the request offers `tool` (OpenAI's {"function": {"name"}}, or a top-level + name).""" + for spec in _given(kwargs.get("tools")) or []: + spec = spec.model_dump() if hasattr(spec, "model_dump") else spec + if not isinstance(spec, dict): + continue + function = spec.get("function") + name = function.get("name") if isinstance(function, dict) else spec.get("name") + if name == tool: + return True + return False + + class Router: """Holds the decision client, the trail and the mode for one wrapped client.""" @@ -41,18 +81,27 @@ def prepare(self, kwargs: dict[str, Any]) -> tuple[str | None, dict[str, Any]]: return conversation_id, kwargs def plan( - self, conversation_id: str | None, decision: Decision | None, error: str | None - ) -> tuple[str, Decision | None]: - """(route, the decision to answer with or None to forward).""" + self, + conversation_id: str | None, + decision: Decision | None, + error: str | None, + kwargs: dict[str, Any] | None = None, + ) -> tuple[str, Decision | None, str | None]: + """(route, the decision to answer with or None to forward, error text).""" if conversation_id is None: - return "no-conversation", None + return "no-conversation", None, error + if kwargs is not None and (why := unsupported(kwargs)) is not None: + return "unsupported", None, why if decision is None: - return "fail-open", None + return "fail-open", None, error if decision.action == "forward": - return "forwarded", None + return "forwarded", None, error if self.mode == SHADOW: - return "shadow", None - return "compiled", decision + return "shadow", None, error + if decision.action == "tool_call" and not offered(kwargs or {}, decision.tool): + # A tool the agent didn't offer this turn: its loop couldn't run it. + return "fail-open", None, "tool not offered" + return "compiled", decision, error def capture( self, @@ -112,11 +161,11 @@ def create(*args: Any, **kwargs: Any) -> Any: started = time.perf_counter() conversation_id, real_kwargs = router.prepare(kwargs) decision = decide_ms = error = None - if conversation_id is not None: + if conversation_id is not None and unsupported(real_kwargs) is None: decision, decide_ms, error = router.decider.decide( router.provider, conversation_id, payload(real_kwargs) ) - route, answer = router.plan(conversation_id, decision, error) + route, answer, error = router.plan(conversation_id, decision, error, real_kwargs) stream = bool(real_kwargs.get("stream")) if answer is not None: try: @@ -161,11 +210,11 @@ async def create(*args: Any, **kwargs: Any) -> Any: started = time.perf_counter() conversation_id, real_kwargs = router.prepare(kwargs) decision = decide_ms = error = None - if conversation_id is not None: + if conversation_id is not None and unsupported(real_kwargs) is None: decision, decide_ms, error = await router.decider.decide( router.provider, conversation_id, payload(real_kwargs) ) - route, answer = router.plan(conversation_id, decision, error) + route, answer, error = router.plan(conversation_id, decision, error, real_kwargs) stream = bool(real_kwargs.get("stream")) if answer is not None: try: diff --git a/src/agentcompile/_decide.py b/src/agentcompile/_decide.py index 063c4b3..3561a66 100644 --- a/src/agentcompile/_decide.py +++ b/src/agentcompile/_decide.py @@ -1,17 +1,30 @@ """The decision call to AgentCompile. Never raises: any problem is a `None` decision, which the -wrapper treats as forward (fail open).""" +wrapper treats as forward (fail open). + +Bounded: the whole call (name lookup, connect, send, read) takes at most `timeout`, and the +server is told how long that is so it answers in time. After a few failures in a row the +breaker opens: calls go straight to the model, without asking, until it tries again. +""" from __future__ import annotations +import asyncio +import concurrent.futures +import threading import time from dataclasses import dataclass, field -from typing import Any +from typing import Any, Callable import httpx COMPANY_HEADER = "x-agentcompiler-company" CONVERSATION_HEADER = "x-agentcompiler-conversation" KEY_HEADER = "x-agentcompiler-key" +MODE_HEADER = "x-agentcompiler-mode" +DEADLINE_HEADER = "x-agentcompiler-deadline-ms" + +FAILURES_TO_OPEN = 5 +OPEN_FOR_S = 30.0 @dataclass(frozen=True) @@ -58,6 +71,7 @@ class Settings: key: str | None company: str | None timeout: float + mode: str = "live" def _request( @@ -75,27 +89,113 @@ def _request( ) +def _decision_request( + settings: Settings, provider: str, conversation_id: str, payload: dict[str, Any] +) -> tuple[str, dict[str, str], dict[str, Any]]: + url, headers, body = _request(settings, provider, conversation_id, payload) + headers[DEADLINE_HEADER] = str(int(settings.timeout * 1000)) + if settings.mode == "shadow": # the server decides but keeps no state + headers[MODE_HEADER] = "shadow" + return url, headers, body + + +class Breaker: + """After `threshold` failures in a row (no answer, or a 5xx or 429), stop asking for + `open_for` seconds; then let one call through to see whether AgentCompile is back.""" + + def __init__( + self, + threshold: int = FAILURES_TO_OPEN, + open_for: float = OPEN_FOR_S, + clock: Callable[[], float] = time.monotonic, + ) -> None: + self.threshold, self.open_for, self._clock = threshold, open_for, clock + self._failures = 0 + self._open_until = 0.0 + self._trying = False + self._lock = threading.Lock() + + def allow(self) -> bool: + with self._lock: + if self._failures < self.threshold: + return True + if self._clock() < self._open_until or self._trying: + return False + self._trying = True # the one call that finds out + return True + + def record(self, ok: bool) -> None: + with self._lock: + self._trying = False + if ok: + self._failures = 0 + return + self._failures += 1 + if self._failures >= self.threshold: + self._open_until = self._clock() + self.open_for + + +def _answered(status: int) -> bool: + """The service is up: anything but a 5xx or being told to back off.""" + return status < 500 and status != 429 + + +def _read(response: Any, ms: float) -> tuple[Decision | None, float, str | None]: + if response.status_code != 200: + return None, ms, f"HTTP {response.status_code}" + decision = Decision.parse(response.json()) + return decision, ms, None if decision else "malformed decision" + + +_POOL: concurrent.futures.ThreadPoolExecutor | None = None +_POOL_LOCK = threading.Lock() + + +def _pool() -> concurrent.futures.ThreadPoolExecutor: + """Where sync decision calls run, so the caller waits at most `timeout` in all (httpx's + timeouts are per phase, and a name lookup has none). A call past its time finishes there, + bounded by httpx's own timeouts, while the caller has moved on.""" + global _POOL + with _POOL_LOCK: + if _POOL is None: + _POOL = concurrent.futures.ThreadPoolExecutor(8, "agentcompile-decide") + return _POOL + + class Decider: def __init__(self, settings: Settings, http: httpx.Client | None = None) -> None: self.settings = settings self._http = http or httpx.Client(timeout=settings.timeout) + self.breaker = Breaker() def decide( self, provider: str, conversation_id: str, payload: dict[str, Any] ) -> tuple[Decision | None, float, str | None]: """(decision or None, milliseconds, error text).""" - url, headers, body = _request(self.settings, provider, conversation_id, payload) + if not self.breaker.allow(): + return None, 0.0, "circuit open" + url, headers, body = _decision_request( + self.settings, provider, conversation_id, payload + ) start = time.perf_counter() + future = None try: - response = self._http.post( - url, headers=headers, json=body, timeout=self.settings.timeout + future = _pool().submit( + self._http.post, url, headers=headers, json=body, timeout=self.settings.timeout ) - ms = (time.perf_counter() - start) * 1000 - if response.status_code != 200: - return None, ms, f"HTTP {response.status_code}" - decision = Decision.parse(response.json()) - return decision, ms, None if decision else "malformed decision" + response = future.result(timeout=self.settings.timeout) + except concurrent.futures.TimeoutError: + if future is not None: + future.cancel() + self.breaker.record(False) + return None, (time.perf_counter() - start) * 1000, "deadline" except Exception as exc: # fail open + self.breaker.record(False) + return None, (time.perf_counter() - start) * 1000, type(exc).__name__ + self.breaker.record(_answered(response.status_code)) + try: + return _read(response, (time.perf_counter() - start) * 1000) + except Exception as exc: return None, (time.perf_counter() - start) * 1000, type(exc).__name__ @@ -103,20 +203,30 @@ class AsyncDecider: def __init__(self, settings: Settings, http: httpx.AsyncClient | None = None) -> None: self.settings = settings self._http = http or httpx.AsyncClient(timeout=settings.timeout) + self.breaker = Breaker() async def decide( self, provider: str, conversation_id: str, payload: dict[str, Any] ) -> tuple[Decision | None, float, str | None]: - url, headers, body = _request(self.settings, provider, conversation_id, payload) + if not self.breaker.allow(): + return None, 0.0, "circuit open" + url, headers, body = _decision_request( + self.settings, provider, conversation_id, payload + ) start = time.perf_counter() try: - response = await self._http.post( - url, headers=headers, json=body, timeout=self.settings.timeout + response = await asyncio.wait_for( + self._http.post(url, headers=headers, json=body, timeout=self.settings.timeout), + timeout=self.settings.timeout, ) - ms = (time.perf_counter() - start) * 1000 - if response.status_code != 200: - return None, ms, f"HTTP {response.status_code}" - decision = Decision.parse(response.json()) - return decision, ms, None if decision else "malformed decision" + except asyncio.TimeoutError: + self.breaker.record(False) + return None, (time.perf_counter() - start) * 1000, "deadline" except Exception as exc: # fail open + self.breaker.record(False) + return None, (time.perf_counter() - start) * 1000, type(exc).__name__ + self.breaker.record(_answered(response.status_code)) + try: + return _read(response, (time.perf_counter() - start) * 1000) + except Exception as exc: return None, (time.perf_counter() - start) * 1000, type(exc).__name__ diff --git a/src/agentcompile/_outcome.py b/src/agentcompile/_outcome.py index 39415b9..499d84c 100644 --- a/src/agentcompile/_outcome.py +++ b/src/agentcompile/_outcome.py @@ -17,14 +17,14 @@ OUTCOMES = ("resolved", "escalated", "unresolved", "abandoned", "reopened", "complaint") -_settings: list[Settings] = [] # the most recently wrapped client's, last -_threads: list[threading.Thread] = [] +_settings: list[Settings] = [] # the most recently wrapped client's (one, not every one) +_threads: list[threading.Thread] = [] # still sending (finished ones are let go) sent = 0 failed = 0 def remember(settings: Settings) -> None: - _settings.append(settings) + _settings[:] = [settings] def outcome(conversation_id: str, label: str, note: str | None = None) -> None: @@ -42,6 +42,7 @@ def outcome(conversation_id: str, label: str, note: str | None = None) -> None: if note: record["note"] = note thread = threading.Thread(target=_send, args=(_settings[-1], record), daemon=True) + _threads[:] = [t for t in _threads if t.is_alive()] _threads.append(thread) thread.start() diff --git a/src/agentcompile/_streams.py b/src/agentcompile/_streams.py index bc90b05..c637e94 100644 --- a/src/agentcompile/_streams.py +++ b/src/agentcompile/_streams.py @@ -1,4 +1,5 @@ -"""Compiled answers as streams: iterators shaped like the SDKs' Stream objects.""" +"""Compiled answers as streams: iterators shaped like the SDKs' Stream objects (one pass: +`for`, `next()` and `__anext__` all read from the same position, as theirs do).""" from __future__ import annotations @@ -8,10 +9,13 @@ class CompiledStream: def __init__(self, items: list[Any]) -> None: - self._items = items + self._it = iter(items) def __iter__(self) -> Iterator[Any]: - return iter(self._items) + return self._it + + def __next__(self) -> Any: + return next(self._it) def __enter__(self) -> CompiledStream: return self @@ -25,10 +29,9 @@ def close(self) -> None: class AsyncCompiledStream: def __init__(self, items: list[Any]) -> None: - self._items = items + self._it = iter(items) def __aiter__(self) -> AsyncCompiledStream: - self._it = iter(self._items) return self async def __anext__(self) -> Any: diff --git a/src/agentcompile/_trail.py b/src/agentcompile/_trail.py index a62f50d..3189a80 100644 --- a/src/agentcompile/_trail.py +++ b/src/agentcompile/_trail.py @@ -1,8 +1,13 @@ """The trail: one JSON line per model call, so you can see what AgentCompile did with each one. Routes: `compiled` (answered without your model), `forwarded` (sent to your model), -`fail-open` (AgentCompile failed or was slow; sent to your model), `shadow` (decided but always -forwarded), `no-conversation` (no conversation id; sent to your model).""" +`fail-open` (AgentCompile failed or was slow, or answered with a tool you didn't offer; sent to +your model), `shadow` (decided but always forwarded), `no-conversation` (no conversation id; +sent to your model), `unsupported` (a request a compiled answer couldn't honour: `n` > 1, a +forced `tool_choice`, a `response_format`; sent to your model without asking). + +The file is kept under 10 MB: past that it moves to `trail.jsonl.1` (the one before that is +dropped) and a new one starts.""" from __future__ import annotations @@ -14,6 +19,7 @@ from pathlib import Path from typing import Any, Callable +MAX_BYTES = 10 * 1024 * 1024 DEFAULT_PATH = Path( os.environ.get("AGENTCOMPILE_TRAIL", Path.home() / ".agentcompile" / "trail.jsonl") ) @@ -29,6 +35,7 @@ def __init__( DEFAULT_PATH if path is True else (Path(path) if path else None) ) self.on_event = on_event + self.max_bytes = MAX_BYTES self._lock = threading.Lock() def record(self, **event: Any) -> dict[str, Any]: @@ -42,6 +49,9 @@ def record(self, **event: Any) -> dict[str, Any]: self.path.parent.mkdir(parents=True, exist_ok=True) with self.path.open("a", encoding="utf-8") as fh: fh.write(json.dumps(event, default=str) + "\n") + full = fh.tell() > self.max_bytes + if full: + os.replace(self.path, self.path.with_name(self.path.name + ".1")) except OSError: pass # the trail never breaks the agent if self.on_event is not None: diff --git a/tests/test_capture.py b/tests/test_capture.py index 893f9d0..a5b3fcd 100644 --- a/tests/test_capture.py +++ b/tests/test_capture.py @@ -114,7 +114,10 @@ def test_a_failing_service_never_breaks_the_call() -> None: ) assert resp.choices[0].message.content == "from the model" agentcompile.flush(5) - assert _capture._ALL[-1].failed == 1 and _capture._ALL[-1].sent == 0 + assert ( + list(_capture._ALL.values())[-1].failed == 1 + and list(_capture._ALL.values())[-1].sent == 0 + ) def test_a_full_queue_drops_the_oldest(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/tests/test_capture_complete.py b/tests/test_capture_complete.py index 7cbed40..853a308 100644 --- a/tests/test_capture_complete.py +++ b/tests/test_capture_complete.py @@ -258,7 +258,9 @@ async def run() -> list[Any]: http_client=httpx.AsyncClient(transport=_provider(OPENAI_STREAM)), ) client = agentcompile.wrap(real, base_url="http://ac.test", trail=False, capture=True) - _capture._ALL[-1]._http = httpx.Client(transport=httpx.MockTransport(service)) + list(_capture._ALL.values())[-1]._http = httpx.Client( + transport=httpx.MockTransport(service) + ) stream = await client.chat.completions.create( model="m", messages=MESSAGES, stream=True, conversation_id="c1" ) diff --git a/tests/test_hardening.py b/tests/test_hardening.py new file mode 100644 index 0000000..7b38712 --- /dev/null +++ b/tests/test_hardening.py @@ -0,0 +1,314 @@ +"""What keeps the wrapper safe in production: a decision is bounded in time and tells the +server how long it waits, a sick service is skipped, a request a compiled answer couldn't honour +goes to the model, streams read like the SDKs' own, and nothing grows without bound.""" + +from __future__ import annotations + +import asyncio +import json +import time +from pathlib import Path +from typing import Any + +import httpx +import openai +import pytest + +import agentcompile +from agentcompile import _capture, _outcome +from agentcompile._assemble import CapturingStream +from agentcompile._decide import AsyncDecider, Breaker, Settings +from agentcompile._streams import AsyncCompiledStream, CompiledStream +from agentcompile._trail import Trail + +CALL = { + "action": "tool_call", + "tool": "get_order_details", + "args": {"order_id": "#W1"}, + "call_id": "call_1", +} +MESSAGES = [{"role": "user", "content": "cancel order #W1"}] +TOOLS = [ + { + "type": "function", + "function": {"name": "get_order_details", "parameters": {"type": "object"}}, + } +] +COMPLETION = { + "id": "c1", + "object": "chat.completion", + "created": 0, + "model": "gpt-x", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "from the model"}, + "finish_reason": "stop", + } + ], +} + + +class Service: + def __init__(self, decision: Any = CALL, status: int = 200, sleep: float = 0.0) -> None: + self.decision, self.status, self.sleep = decision, status, sleep + self.decides: list[httpx.Request] = [] + self.captures: list[dict[str, Any]] = [] + + def __call__(self, request: httpx.Request) -> httpx.Response: + if request.url.path == "/v1/capture": + self.captures.append(json.loads(request.content)) + return httpx.Response(200, json={"kept": 1}) + self.decides.append(request) + if self.sleep: + time.sleep(self.sleep) + return httpx.Response(self.status, json=self.decision) + + +def _client(service: Service, tmp_path: Path, **kw: Any) -> tuple[Any, list[Any]]: + provider_calls: list[Any] = [] + + def provider(request: httpx.Request) -> httpx.Response: + provider_calls.append(json.loads(request.content)) + return httpx.Response(200, json=COMPLETION) + + real = openai.OpenAI( + api_key="sk-test", + base_url="http://provider.test/v1", + http_client=httpx.Client(transport=httpx.MockTransport(provider)), + ) + client = 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)), + **kw, + ) + return client, provider_calls + + +def _trail(tmp_path: Path) -> list[dict[str, Any]]: + return [json.loads(line) for line in (tmp_path / "trail.jsonl").read_text().splitlines()] + + +@pytest.fixture(autouse=True) +def _no_leftover_senders() -> Any: + yield + _capture._ALL.clear() + + +# -- what the server is told --------------------------------------------------------------- + + +def test_the_server_hears_how_long_we_wait_and_whether_it_is_shadow(tmp_path: Path) -> None: + service = Service() + live, _ = _client(service, tmp_path, timeout=1.5) + live.chat.completions.create(model="m", messages=MESSAGES, tools=TOOLS, conversation_id="c") + shadow, _ = _client(service, tmp_path, mode="shadow") + shadow.chat.completions.create( + model="m", messages=MESSAGES, tools=TOOLS, conversation_id="c" + ) + first, second = service.decides + assert first.headers["x-agentcompiler-deadline-ms"] == "1500" + assert "x-agentcompiler-mode" not in first.headers + assert second.headers["x-agentcompiler-mode"] == "shadow" + + +# -- requests a compiled answer can't honour -------------------------------------------------- + + +@pytest.mark.parametrize( + ("extra", "why"), + [ + ({"n": 2}, "n"), + ({"tool_choice": "required"}, "tool_choice"), + ({"tool_choice": {"type": "function", "function": {"name": "x"}}}, "tool_choice"), + ({"response_format": {"type": "json_object"}}, "response_format"), + ], +) +def test_a_request_we_could_not_honour_goes_to_the_model_unasked( + tmp_path: Path, extra: dict[str, Any], why: str +) -> None: + service = Service() + client, provider = _client(service, tmp_path) + resp = client.chat.completions.create( + model="m", messages=MESSAGES, tools=TOOLS, conversation_id="c", **extra + ) + assert resp.choices[0].message.content == "from the model" + assert service.decides == [] and len(provider) == 1 + event = _trail(tmp_path)[0] + assert event["route"] == "unsupported" and event["reason"] == why + + +def test_auto_tool_choice_and_text_format_are_still_answered(tmp_path: Path) -> None: + service = Service() + client, provider = _client(service, tmp_path) + resp = client.chat.completions.create( + model="m", + messages=MESSAGES, + tools=TOOLS, + tool_choice="auto", + response_format={"type": "text"}, + n=1, + conversation_id="c", + ) + assert resp.choices[0].message.tool_calls and provider == [] + + +def test_a_tool_the_agent_did_not_offer_is_never_answered(tmp_path: Path) -> None: + service = Service() + client, provider = _client(service, tmp_path) + other = [{"type": "function", "function": {"name": "cancel_order"}}] + resp = client.chat.completions.create( + model="m", messages=MESSAGES, tools=other, conversation_id="c" + ) + assert resp.choices[0].message.content == "from the model" and len(provider) == 1 + event = _trail(tmp_path)[0] + assert event["route"] == "fail-open" and event["reason"] == "tool not offered" + + +# -- bounded in time --------------------------------------------------------------------------- + + +def test_a_slow_service_costs_at_most_the_timeout(tmp_path: Path) -> None: + service = Service(sleep=1.0) + client, provider = _client(service, tmp_path, timeout=0.2) + started = time.perf_counter() + resp = client.chat.completions.create( + model="m", messages=MESSAGES, tools=TOOLS, conversation_id="c" + ) + assert time.perf_counter() - started < 0.6 + assert resp.choices[0].message.content == "from the model" and len(provider) == 1 + event = _trail(tmp_path)[0] + assert event["route"] == "fail-open" and event["reason"] == "deadline" + + +async def test_the_async_client_is_bounded_too() -> None: + async def slow(request: httpx.Request) -> httpx.Response: + await asyncio.sleep(1.0) + return httpx.Response(200, json=CALL) + + decider = AsyncDecider( + Settings("http://ac.test", "k", None, 0.2), + httpx.AsyncClient(transport=httpx.MockTransport(slow)), + ) + started = time.perf_counter() + decision, _, error = await decider.decide("openai", "c", {"messages": MESSAGES}) + assert decision is None and error == "deadline" + assert time.perf_counter() - started < 0.6 + + +def test_a_sick_service_is_skipped_then_tried_again(tmp_path: Path) -> None: + service = Service(status=503) + client, provider = _client(service, tmp_path) + + def call() -> None: + client.chat.completions.create( + model="m", messages=MESSAGES, tools=TOOLS, conversation_id="c" + ) + + for _ in range(5): + call() + assert len(service.decides) == 5 + call() # open: straight to the model, not asked + assert len(service.decides) == 5 and len(provider) == 6 + assert _trail(tmp_path)[-1]["reason"] == "circuit open" + + +def test_the_breaker_lets_one_call_through_after_a_while() -> None: + now = [0.0] + breaker = Breaker(threshold=2, open_for=10.0, clock=lambda: now[0]) + breaker.record(False) + assert breaker.allow() + breaker.record(False) + assert not breaker.allow() + now[0] = 11.0 + assert breaker.allow() and not breaker.allow() # one trial at a time + breaker.record(True) + assert breaker.allow() and breaker.allow() # closed again + breaker.record(False) + assert breaker.allow() # one failure isn't a pattern + + +# -- streams read like the SDKs' own ---------------------------------------------------------- + + +def test_next_reads_a_compiled_stream_from_where_it_is() -> None: + stream = CompiledStream([1, 2, 3]) + assert next(stream) == 1 + assert list(stream) == [2, 3] + + +async def test_anext_works_without_iterating_first() -> None: + stream = AsyncCompiledStream([1, 2]) + assert await stream.__anext__() == 1 + assert [x async for x in stream] == [2] + + +def test_next_reads_a_captured_stream_and_capture_still_sees_it_whole() -> None: + seen: list[Any] = [] + stream = CapturingStream(iter([{"a": 1}, {"a": 2}]), lambda c, done: seen.append((c, done))) + assert next(stream) == {"a": 1} + assert list(stream) == [{"a": 2}] + assert seen == [([{"a": 1}, {"a": 2}], True)] + + +def test_a_compiled_stream_answers_next_from_the_wrapped_client(tmp_path: Path) -> None: + client, _ = _client(Service(), tmp_path) + stream = client.chat.completions.create( + model="m", messages=MESSAGES, tools=TOOLS, conversation_id="c", stream=True + ) + first = next(stream) + assert first.choices[0].delta.tool_calls[0].function.name == "get_order_details" + + +# -- nothing grows without bound --------------------------------------------------------- + + +def test_the_trail_moves_aside_when_full(tmp_path: Path) -> None: + trail = Trail(tmp_path / "trail.jsonl") + trail.max_bytes = 200 + for n in range(20): + trail.record(route="forwarded", n=n) + moved = tmp_path / "trail.jsonl.1" + assert moved.exists() and moved.stat().st_size <= 200 + 100 # one line past the limit + current = tmp_path / "trail.jsonl" + assert not current.exists() or current.stat().st_size <= 200 + 100 + + +def test_capture_batches_stay_under_the_byte_limit_in_order( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(_capture, "MAX_BATCH_BYTES", 1100) + monkeypatch.setattr(_capture, "MAX_RECORD_BYTES", 2000) + service = Service() + capturer = _capture.Capturer( + Settings("http://ac.test", "ack_acme.k", None, 1.0), + httpx.Client(transport=httpx.MockTransport(service)), + ) + capturer._start = lambda: None # type: ignore[method-assign] + for n in range(4): + capturer.add("openai", f"c{n}", {"messages": [{"content": "x" * 300}]}, {"n": n}) + capturer.add("openai", "huge", {"messages": [{"content": "x" * 3000}]}, {}) + capturer.flush() + sent = [r["conversation_id"] for batch in service.captures for r in batch["exchanges"]] + assert sent == ["c0", "c1", "c2", "c3"] and len(service.captures) == 2 + assert capturer.dropped == 1 + + +def test_wrapping_again_reuses_the_sender_and_keeps_one_destination(tmp_path: Path) -> None: + service = Service() + http = httpx.Client(transport=httpx.MockTransport(service)) + for _ in range(3): + agentcompile.wrap( + openai.OpenAI(api_key="sk-test"), + key="ack_acme.k", + base_url="http://ac.test", + trail=False, + http_client=http, + capture=True, + scrub=False, + ) + assert len(_capture._ALL) == 1 + assert len(_outcome._settings) == 1 diff --git a/tests/test_wrap.py b/tests/test_wrap.py index 4903d3e..25f960f 100644 --- a/tests/test_wrap.py +++ b/tests/test_wrap.py @@ -134,6 +134,13 @@ def _trail(tmp_path: Path) -> list[dict[str, Any]]: MESSAGES = [{"role": "user", "content": "cancel order #W1"}] +TOOLS = [ + { + "type": "function", + "function": {"name": "get_order_details", "parameters": {"type": "object"}}, + } +] +ANTHROPIC_TOOLS = [{"name": "get_order_details", "input_schema": {"type": "object"}}] def test_a_compiled_tool_call_is_a_real_chat_completion_and_the_model_is_not_called( @@ -142,7 +149,7 @@ def test_a_compiled_tool_call_is_a_real_chat_completion_and_the_model_is_not_cal fake = Fake(CALL) client = _openai(fake, tmp_path) resp = client.chat.completions.create( - model="gpt-x", messages=MESSAGES, conversation_id="c1" + model="gpt-x", messages=MESSAGES, tools=TOOLS, conversation_id="c1" ) call = resp.choices[0].message.tool_calls[0] assert resp.choices[0].finish_reason == "tool_calls" @@ -161,7 +168,7 @@ def test_a_compiled_tool_call_is_a_real_chat_completion_and_the_model_is_not_cal def test_forward_calls_the_model_with_our_keyword_removed(tmp_path: Path) -> None: fake = Fake(FORWARD) resp = _openai(fake, tmp_path).chat.completions.create( - model="gpt-x", messages=MESSAGES, conversation_id="c1" + model="gpt-x", messages=MESSAGES, tools=TOOLS, conversation_id="c1" ) assert resp.choices[0].message.content == "from the model" assert "conversation_id" not in fake.provider_requests[0] @@ -176,7 +183,7 @@ def test_the_trail_keeps_what_agentcompile_decided_along_the_way(tmp_path: Path) ] fake = Fake({**FORWARD, "events": events}) _openai(fake, tmp_path).chat.completions.create( - model="gpt-x", messages=MESSAGES, conversation_id="c1" + model="gpt-x", messages=MESSAGES, tools=TOOLS, conversation_id="c1" ) assert _trail(tmp_path)[0]["events"] == events @@ -184,7 +191,7 @@ def test_the_trail_keeps_what_agentcompile_decided_along_the_way(tmp_path: Path) def test_malformed_events_are_dropped_not_fatal(tmp_path: Path) -> None: fake = Fake({**CALL, "events": ["not a dict", {"decision": "emit_read"}]}) _openai(fake, tmp_path).chat.completions.create( - model="gpt-x", messages=MESSAGES, conversation_id="c1" + model="gpt-x", messages=MESSAGES, tools=TOOLS, conversation_id="c1" ) event = _trail(tmp_path)[0] assert event["route"] == "compiled" and event["events"] == [{"decision": "emit_read"}] @@ -201,7 +208,7 @@ def test_malformed_events_are_dropped_not_fatal(tmp_path: Path) -> None: ) def test_any_decision_problem_fails_open_to_the_model(fake: Fake, tmp_path: Path) -> None: resp = _openai(fake, tmp_path).chat.completions.create( - model="gpt-x", messages=MESSAGES, conversation_id="c1" + model="gpt-x", messages=MESSAGES, tools=TOOLS, conversation_id="c1" ) assert resp.choices[0].message.content == "from the model" assert _trail(tmp_path)[0]["route"] == "fail-open" @@ -226,7 +233,7 @@ def test_shadow_mode_always_calls_the_model_and_records_what_it_would_have_done( ) -> None: fake = Fake(CALL) resp = _openai(fake, tmp_path, mode="shadow").chat.completions.create( - model="gpt-x", messages=MESSAGES, conversation_id="c1" + model="gpt-x", messages=MESSAGES, tools=TOOLS, conversation_id="c1" ) assert resp.choices[0].message.content == "from the model" event = _trail(tmp_path)[0] @@ -236,7 +243,7 @@ def test_shadow_mode_always_calls_the_model_and_records_what_it_would_have_done( def test_streaming_a_compiled_tool_call(tmp_path: Path) -> None: fake = Fake(CALL) stream = _openai(fake, tmp_path).chat.completions.create( - model="gpt-x", messages=MESSAGES, conversation_id="c1", stream=True + model="gpt-x", messages=MESSAGES, tools=TOOLS, conversation_id="c1", stream=True ) with stream: chunks = list(stream) @@ -249,14 +256,16 @@ def test_pydantic_messages_from_earlier_responses_are_sent_as_json(tmp_path: Pat fake = Fake(CALL) client = _openai(fake, tmp_path) first = client.chat.completions.create( - model="gpt-x", messages=MESSAGES, conversation_id="c1" + model="gpt-x", messages=MESSAGES, tools=TOOLS, conversation_id="c1" ) history = [ *MESSAGES, first.choices[0].message, {"role": "tool", "tool_call_id": "call_1", "content": "{}"}, ] - client.chat.completions.create(model="gpt-x", messages=history, conversation_id="c1") + client.chat.completions.create( + model="gpt-x", messages=history, tools=TOOLS, conversation_id="c1" + ) sent = json.loads(fake.decide_requests[1].content)["request"]["messages"] assert sent[1]["tool_calls"][0]["function"]["name"] == "get_order_details" @@ -279,6 +288,7 @@ def test_anthropic_compiled_reply_and_stream(tmp_path: Path) -> None: model="claude-x", max_tokens=100, messages=MESSAGES, + tools=ANTHROPIC_TOOLS, conversation_id="c1", stream=True, ) @@ -311,12 +321,12 @@ async def test_async_openai_compiled_and_forward(tmp_path: Path) -> None: http_client=httpx.AsyncClient(transport=httpx.MockTransport(fake.decide)), ) resp = await client.chat.completions.create( - model="gpt-x", messages=MESSAGES, conversation_id="c1" + model="gpt-x", messages=MESSAGES, tools=TOOLS, conversation_id="c1" ) assert resp.choices[0].message.tool_calls[0].function.name == "get_order_details" fake.decision = FORWARD resp = await client.chat.completions.create( - model="gpt-x", messages=MESSAGES, conversation_id="c1" + model="gpt-x", messages=MESSAGES, tools=TOOLS, conversation_id="c1" ) assert resp.choices[0].message.content == "from the model" @@ -334,9 +344,13 @@ def test_the_trail_command_prints_a_summary( fake = Fake(CALL) client = _openai(fake, tmp_path) - client.chat.completions.create(model="gpt-x", messages=MESSAGES, conversation_id="c1") + client.chat.completions.create( + model="gpt-x", messages=MESSAGES, tools=TOOLS, conversation_id="c1" + ) fake.decision = FORWARD - client.chat.completions.create(model="gpt-x", messages=MESSAGES, conversation_id="c1") + client.chat.completions.create( + model="gpt-x", messages=MESSAGES, tools=TOOLS, conversation_id="c1" + ) assert main(["trail", "--path", str(tmp_path / "trail.jsonl")]) == 0 out = capsys.readouterr().out assert (