backspin 0.5.1__py3-none-any.whl

backspin/integrations/openai.py

@@ -0,0 +1,298 @@
+"""Capture OpenAI-SDK-shaped chat completion calls into a Recorder.
+
+The wrapper duck-types: it works with the real ``openai.OpenAI`` /
+``openai.AsyncOpenAI`` clients and with any object shaped like them (see
+:mod:`backspin.testing`), sync or async, streaming or not.
+"""
+from __future__ import annotations
+
+import functools
+import inspect
+import time
+from typing import Any, Callable, Dict, List, Optional
+
+from ..fakes import stream_chunks  # noqa: F401  (re-exported for replay/tests)
+from ..runfile import jsonable
+
+# Request fields worth keeping in a recording. Everything else (clients,
+# callables, ...) is dropped rather than repr-ed to keep run files clean.
+_KEEP = {
+    "model",
+    "messages",
+    "tools",
+    "tool_choice",
+    "temperature",
+    "top_p",
+    "max_tokens",
+    "max_completion_tokens",
+    "stop",
+    "n",
+    "stream",
+    "stream_options",
+    "response_format",
+    "seed",
+    "user",
+    "metadata",
+    "frequency_penalty",
+    "presence_penalty",
+    "parallel_tool_calls",
+    "logprobs",
+    "top_logprobs",
+}
+
+
+def _clean_kwargs(kwargs: Dict[str, Any]) -> Dict[str, Any]:
+    return {k: jsonable(v) for k, v in kwargs.items() if k in _KEEP}
+
+
+def _usage_of(usage: Any) -> Optional[Dict[str, Any]]:
+    """Extract a dict from a usage object/dict of any shape."""
+    if usage is None:
+        return None
+    if isinstance(usage, dict):
+        return usage
+    if hasattr(usage, "model_dump"):
+        try:
+            return usage.model_dump(exclude_none=True)
+        except TypeError:
+            return usage.model_dump()
+    return {
+        k: getattr(usage, k)
+        for k in ("prompt_tokens", "completion_tokens", "total_tokens")
+        if getattr(usage, k, None) is not None
+    }
+
+
+def _usage(resp: Any) -> Optional[Dict[str, Any]]:
+    return _usage_of(getattr(resp, "usage", None))
+
+
+def _response_payload(resp: Any) -> Dict[str, Any]:
+    if hasattr(resp, "model_dump"):
+        try:
+            return resp.model_dump(exclude_none=True)
+        except TypeError:
+            return resp.model_dump()
+    return {"repr": repr(resp)}
+
+
+def _ms(t0: float) -> float:
+    return (time.perf_counter() - t0) * 1000
+
+
+class _StreamState:
+    """Accumulated view of a streamed completion, built from its chunks."""
+
+    def __init__(self) -> None:
+        self.content: List[str] = []
+        self.tool_calls: Dict[int, Dict[str, Any]] = {}
+        self.model: Optional[str] = None
+        self.usage: Optional[Dict[str, Any]] = None
+        self.finish_reason: Optional[str] = None
+
+    def absorb(self, chunk: Any) -> None:
+        model = getattr(chunk, "model", None)
+        if model:
+            self.model = model
+        usage = getattr(chunk, "usage", None)
+        if usage is not None:
+            self.usage = _usage_of(usage)
+        for choice in getattr(chunk, "choices", None) or []:
+            delta = getattr(choice, "delta", None)
+            if delta is None:
+                continue
+            content = getattr(delta, "content", None)
+            if content:
+                self.content.append(content)
+            for tc in getattr(delta, "tool_calls", None) or []:
+                idx = getattr(tc, "index", 0)
+                slot = self.tool_calls.setdefault(
+                    idx,
+                    {"id": "", "type": "function",
+                     "function": {"name": "", "arguments": ""}},
+                )
+                if getattr(tc, "id", None):
+                    slot["id"] = tc.id
+                fn = getattr(tc, "function", None)
+                if fn is not None:
+                    if getattr(fn, "name", None):
+                        slot["function"]["name"] = fn.name
+                    if getattr(fn, "arguments", None):
+                        slot["function"]["arguments"] += fn.arguments
+            finish = getattr(choice, "finish_reason", None)
+            if finish:
+                self.finish_reason = finish
+
+    def payload(self) -> Dict[str, Any]:
+        message: Dict[str, Any] = {"role": "assistant"}
+        if self.content:
+            message["content"] = "".join(self.content)
+        if self.tool_calls:
+            message["tool_calls"] = [self.tool_calls[i] for i in sorted(self.tool_calls)]
+        return {
+            "object": "chat.completion",
+            "model": self.model,
+            "reconstructed_from_stream": True,
+            "choices": [
+                {
+                    "index": 0,
+                    "finish_reason": self.finish_reason or "stop",
+                    "message": message,
+                }
+            ],
+            "usage": self.usage,
+        }
+
+
+def capture_openai(recorder: Any, client: Any) -> Any:
+    """Patch ``client.chat.completions.create`` to record every call."""
+    completions = client.chat.completions
+    original = completions.create
+    # Newer openai SDKs ship async methods that are not coroutine functions
+    # (they return awaitables), so the resource type name is the reliable tell.
+    is_async = inspect.iscoroutinefunction(original) or type(
+        completions
+    ).__name__.lower().startswith("async")
+    if is_async:
+        completions.create = _async_wrapper(recorder, original)
+    else:
+        completions.create = _sync_wrapper(recorder, original)
+    return client
+
+
+def _sync_wrapper(recorder: Any, original: Callable) -> Callable:
+    @functools.wraps(original)
+    def create(*args: Any, **kwargs: Any):
+        t0 = time.perf_counter()
+        try:
+            resp = original(*args, **kwargs)
+        except Exception as exc:
+            recorder.record_llm(
+                request=_clean_kwargs(kwargs),
+                model=kwargs.get("model"),
+                error=exc,
+                duration_ms=_ms(t0),
+            )
+            raise
+        if kwargs.get("stream"):
+            return _StreamRecorder(recorder, kwargs, resp, t0)
+        recorder.record_llm(
+            request=_clean_kwargs(kwargs),
+            response=_response_payload(resp),
+            usage=_usage(resp),
+            model=kwargs.get("model") or getattr(resp, "model", None),
+            duration_ms=_ms(t0),
+        )
+        return resp
+
+    return create
+
+
+def _async_wrapper(recorder: Any, original: Callable) -> Callable:
+    @functools.wraps(original)
+    async def create(*args: Any, **kwargs: Any):
+        t0 = time.perf_counter()
+        try:
+            resp = await original(*args, **kwargs)
+        except Exception as exc:
+            recorder.record_llm(
+                request=_clean_kwargs(kwargs),
+                model=kwargs.get("model"),
+                error=exc,
+                duration_ms=_ms(t0),
+            )
+            raise
+        if kwargs.get("stream"):
+            return _AsyncStreamRecorder(recorder, kwargs, resp, t0)
+        recorder.record_llm(
+            request=_clean_kwargs(kwargs),
+            response=_response_payload(resp),
+            usage=_usage(resp),
+            model=kwargs.get("model") or getattr(resp, "model", None),
+            duration_ms=_ms(t0),
+        )
+        return resp
+
+    return create
+
+
+class _StreamRecorder:
+    """Transparent iterator wrapper: passthrough chunks, record on end."""
+
+    def __init__(self, recorder: Any, kwargs: Dict[str, Any], stream: Any, t0: float):
+        self._recorder = recorder
+        self._kwargs = kwargs
+        self._stream = stream
+        self._t0 = t0
+        self._state = _StreamState()
+        self._done = False
+
+    def __iter__(self) -> "_StreamRecorder":
+        return self
+
+    def __next__(self) -> Any:
+        try:
+            chunk = next(self._stream)
+        except StopIteration:
+            self._finalize()
+            raise
+        except Exception as exc:
+            self._finalize(error=exc)
+            raise
+        self._state.absorb(chunk)
+        return chunk
+
+    def close(self) -> None:
+        closer = getattr(self._stream, "close", None)
+        if closer is not None:
+            closer()
+        self._finalize()
+
+    def __enter__(self) -> "_StreamRecorder":
+        enter = getattr(self._stream, "__enter__", None)
+        if enter is not None:
+            enter()
+        return self
+
+    def __exit__(self, *exc_info: Any):
+        self._finalize()
+        exit_ = getattr(self._stream, "__exit__", None)
+        if exit_ is not None:
+            return exit_(*exc_info)
+        return None
+
+    def _finalize(self, error: Optional[BaseException] = None) -> None:
+        if self._done:
+            return
+        self._done = True
+        self._recorder.record_llm(
+            request=_clean_kwargs(self._kwargs),
+            response=self._state.payload(),
+            usage=self._state.usage,
+            model=self._kwargs.get("model") or self._state.model,
+            duration_ms=_ms(self._t0),
+            error=error,
+        )
+
+
+class _AsyncStreamRecorder(_StreamRecorder):
+    def __aiter__(self) -> "_AsyncStreamRecorder":
+        return self
+
+    async def __anext__(self) -> Any:
+        try:
+            chunk = await self._stream.__anext__()
+        except StopAsyncIteration:
+            self._finalize()
+            raise
+        except Exception as exc:
+            self._finalize(error=exc)
+            raise
+        self._state.absorb(chunk)
+        return chunk
+
+    async def aclose(self) -> None:
+        closer = getattr(self._stream, "aclose", None)
+        if closer is not None:
+            await closer()
+        self._finalize()

backspin/proxy.py

@@ -0,0 +1,367 @@
+"""OpenAI-compatible local proxy: record or replay, framework-agnostic.
+
+Record mode — point any agent at the proxy instead of the vendor endpoint;
+every request is forwarded upstream and captured into a run file::
+
+    backspin proxy --upstream https://api.openai.com --port 8840
+    # client: base_url = http://127.0.0.1:8840/v1
+
+Replay mode — serve a recorded run back as an API, no upstream and no
+network, from any language that speaks the OpenAI protocol::
+
+    backspin proxy --replay runs/live.backspin.jsonl --port 8840
+
+The proxy speaks streaming SSE on both paths: streamed upstream answers are
+reconstructed into a recorded completion, and replays are streamed back as
+chunks. Authorization headers are forwarded to upstream but never recorded.
+"""
+from __future__ import annotations
+
+import json
+import time
+from contextlib import suppress
+from typing import Any, Dict, Optional
+
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.responses import JSONResponse, StreamingResponse
+
+from .fakes import stream_chunks
+from .integrations.anthropic import _Acc as _AnthropicAcc
+from .recorder import Recorder
+from .replay import Cassette
+from .runfile import fingerprint_request
+
+
+def _norm_anthropic_usage(payload: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
+    usage = (payload or {}).get("usage") or {}
+    if not usage:
+        return None
+    return {
+        "prompt_tokens": usage.get("input_tokens") or 0,
+        "completion_tokens": usage.get("output_tokens") or 0,
+    }
+
+
+def _ms_since(t0: float) -> float:
+    return (time.perf_counter() - t0) * 1000
+
+
+def _sse(obj: Any) -> str:
+    return "data: " + json.dumps(obj, ensure_ascii=False) + "\n\n"
+
+
+async def _json_body(request: "Request") -> Dict[str, Any]:
+    """Parse the request body as a JSON object; 400 on anything else."""
+    try:
+        body = await request.json()
+    except Exception as exc:
+        raise HTTPException(status_code=400, detail=f"invalid JSON body: {exc}") from exc
+    if not isinstance(body, dict):
+        raise HTTPException(status_code=400, detail="request body must be a JSON object")
+    return body
+
+
+def _acc_payload(acc: Dict[str, Any]) -> Dict[str, Any]:
+    """Rebuild a completion-shaped dict from absorbed stream chunks."""
+    message: Dict[str, Any] = {"role": "assistant"}
+    if acc["content"]:
+        message["content"] = "".join(acc["content"])
+    if acc["tools"]:
+        message["tool_calls"] = [acc["tools"][i] for i in sorted(acc["tools"])]
+    return {
+        "object": "chat.completion",
+        "model": acc["model"],
+        "reconstructed_from_stream": True,
+        "choices": [{
+            "index": 0,
+            "finish_reason": acc["finish"] or "stop",
+            "message": message,
+        }],
+        "usage": acc["usage"],
+    }
+
+
+def _absorb_chunk(acc: Dict[str, Any], chunk: Dict[str, Any]) -> None:
+    """Accumulate one upstream SSE chunk (dict) into ``acc``."""
+    if chunk.get("model"):
+        acc["model"] = chunk["model"]
+    if chunk.get("usage"):
+        acc["usage"] = chunk["usage"]
+    for choice in chunk.get("choices") or []:
+        delta = choice.get("delta") or {}
+        if delta.get("content"):
+            acc["content"].append(delta["content"])
+        for tc in delta.get("tool_calls") or []:
+            idx = tc.get("index", 0)
+            slot = acc["tools"].setdefault(
+                idx, {"id": "", "type": "function",
+                      "function": {"name": "", "arguments": ""}}
+            )
+            if tc.get("id"):
+                slot["id"] = tc["id"]
+            fn = tc.get("function") or {}
+            if fn.get("name"):
+                slot["function"]["name"] = fn["name"]
+            if fn.get("arguments"):
+                slot["function"]["arguments"] += fn["arguments"]
+        if choice.get("finish_reason"):
+            acc["finish"] = choice["finish_reason"]
+
+
+def create_proxy_app(
+    upstream: Optional[str] = None,
+    cassette: Optional[Cassette] = None,
+    runs_dir: str = "runs",
+    agent: str = "proxy",
+    metadata: Optional[Dict[str, Any]] = None,
+) -> FastAPI:
+    """Build the proxy app. Set ``upstream`` (record) or ``cassette``
+    (replay) — exactly one.
+
+    The recorder is available as ``app.state.recorder`` after creation.
+    """
+    try:
+        import httpx
+    except ImportError as exc:  # pragma: no cover
+        raise ImportError(
+            "proxy mode requires httpx: pip install 'backspin[proxy]'"
+        ) from exc
+
+    if (upstream is None) == (cassette is None):
+        raise ValueError(
+            "create_proxy_app(): set exactly one of upstream= (record) or cassette= (replay)"
+        )
+
+    app = FastAPI(title="backspin proxy", docs_url=None, redoc_url=None)
+    if upstream is not None:
+        recorder = Recorder(
+            dir=runs_dir, agent=agent,
+            metadata={**(metadata or {}), "mode": "record", "upstream": upstream},
+        )
+        client = httpx.AsyncClient(
+            base_url=upstream,
+            timeout=httpx.Timeout(connect=15.0, read=600.0, write=600.0, pool=15.0),
+        )
+    else:
+        recorder = Recorder(
+            dir=runs_dir, agent=agent,
+            metadata={**(metadata or {}), "mode": "replay"},
+        )
+        client = None
+    app.state.recorder = recorder
+
+    @app.get("/v1/models")
+    async def models():
+        if client is None:
+            return JSONResponse({"object": "list", "data": []})
+        resp = await client.get("/v1/models")
+        return JSONResponse(resp.json(), status_code=resp.status_code)
+
+    @app.post("/v1/chat/completions")
+    async def chat(request: Request):
+        body = await _json_body(request)
+        stream = body.get("stream", False)
+        t0 = time.perf_counter()
+
+        # ---- replay mode: answer from the cassette, no upstream ----------
+        if client is None:
+            assert cassette is not None  # exactly one of upstream/cassette is set
+            fp = fingerprint_request(body.get("model"), body.get("messages"))
+            entry, _exact = cassette.match(fp)
+            if entry is None:
+                raise HTTPException(
+                    status_code=503,
+                    detail="backspin replay proxy: cassette exhausted",
+                )
+            data = dict(entry["response"])
+            if stream:
+                async def replay_gen():
+                    for chunk in stream_chunks(data):
+                        yield _sse(chunk.model_dump())
+                    yield "data: [DONE]\n\n"
+
+                return StreamingResponse(replay_gen(), media_type="text/event-stream")
+            return JSONResponse(data)
+
+        # ---- record mode: forward upstream, capture the exchange ---------
+        fwd_headers = {
+            k: v for k, v in request.headers.items()
+            if k.lower() in ("authorization", "content-type", "accept")
+        }
+
+        if not stream:
+            resp = await client.post("/v1/chat/completions", json=body, headers=fwd_headers)
+            try:
+                payload = resp.json()
+            except Exception:
+                payload = None
+            if resp.status_code >= 400:
+                recorder.record_llm(
+                    request=body, model=body.get("model"), duration_ms=_ms_since(t0),
+                    error=RuntimeError(f"upstream HTTP {resp.status_code}: {resp.text[:300]}"),
+                )
+            else:
+                recorder.record_llm(
+                    request=body, response=payload,
+                    usage=payload.get("usage") if payload else None,
+                    model=(payload or {}).get("model") or body.get("model"),
+                    duration_ms=_ms_since(t0),
+                )
+            return JSONResponse(
+                payload if payload is not None else {}, status_code=resp.status_code
+            )
+
+        upstream_resp = await client.send(
+            client.build_request("POST", "/v1/chat/completions", json=body, headers=fwd_headers),
+            stream=True,
+        )
+        if upstream_resp.status_code >= 400:
+            raw = (await upstream_resp.aread()).decode("utf-8", "replace")
+            await upstream_resp.aclose()
+            recorder.record_llm(
+                request=body, model=body.get("model"), duration_ms=_ms_since(t0),
+                error=RuntimeError(f"upstream HTTP {upstream_resp.status_code}: {raw[:300]}"),
+            )
+            return JSONResponse(
+                {"error": {"message": raw[:500]}}, status_code=upstream_resp.status_code
+            )
+
+        acc: Dict[str, Any] = {
+            "content": [], "tools": {}, "usage": None,
+            "finish": None, "model": body.get("model"),
+        }
+
+        async def passthrough_gen():
+            try:
+                async for line in upstream_resp.aiter_lines():
+                    yield line + "\n"
+                    stripped = line.strip()
+                    if stripped.startswith("data:"):
+                        data = stripped[5:].strip()
+                        if data and data != "[DONE]":
+                            with suppress(json.JSONDecodeError):
+                                _absorb_chunk(acc, json.loads(data))
+            finally:
+                await upstream_resp.aclose()
+                recorder.record_llm(
+                    request=body, response=_acc_payload(acc), usage=acc["usage"],
+                    model=acc["model"], duration_ms=_ms_since(t0),
+                )
+
+        return StreamingResponse(passthrough_gen(), media_type="text/event-stream")
+
+    @app.post("/v1/messages")
+    async def anthropic_messages(request: Request):
+        """Anthropic Messages protocol: same record/replay behavior."""
+        body = await _json_body(request)
+        stream = body.get("stream", False)
+        t0 = time.perf_counter()
+
+        # ---- replay -------------------------------------------------------
+        if client is None:
+            assert cassette is not None  # exactly one of upstream/cassette is set
+            fp = fingerprint_request(body.get("model"), body.get("messages"))
+            entry, _exact = cassette.match(fp)
+            if entry is None:
+                raise HTTPException(
+                    status_code=503,
+                    detail="backspin replay proxy: cassette exhausted",
+                )
+            data = dict(entry["response"])
+            if stream:
+                from .fakes import anthropic_stream_events
+
+                async def anthropic_replay_gen():
+                    for event in anthropic_stream_events(data):
+                        yield "event: " + event.get("type", "message") + "\n" + _sse(event)
+                    yield "event: message_stop\n" + _sse({"type": "message_stop"})
+
+                return StreamingResponse(anthropic_replay_gen(), media_type="text/event-stream")
+            return JSONResponse(data)
+
+        # ---- record -------------------------------------------------------
+        fwd_headers = {
+            k: v for k, v in request.headers.items()
+            if k.lower() in ("x-api-key", "anthropic-version", "content-type", "accept")
+        }
+
+        if not stream:
+            resp = await client.post("/v1/messages", json=body, headers=fwd_headers)
+            try:
+                payload = resp.json()
+            except Exception:
+                payload = None
+            if resp.status_code >= 400:
+                recorder.record_llm(
+                    request=body, model=body.get("model"), duration_ms=_ms_since(t0),
+                    error=RuntimeError(f"upstream HTTP {resp.status_code}: {resp.text[:300]}"),
+                    provider="anthropic",
+                )
+            else:
+                recorder.record_llm(
+                    request=body, response=payload,
+                    usage=_norm_anthropic_usage(payload),
+                    model=(payload or {}).get("model") or body.get("model"),
+                    duration_ms=_ms_since(t0), provider="anthropic",
+                )
+            return JSONResponse(
+                payload if payload is not None else {}, status_code=resp.status_code
+            )
+
+        upstream_resp = await client.send(
+            client.build_request("POST", "/v1/messages", json=body, headers=fwd_headers),
+            stream=True,
+        )
+        if upstream_resp.status_code >= 400:
+            raw = (await upstream_resp.aread()).decode("utf-8", "replace")
+            await upstream_resp.aclose()
+            recorder.record_llm(
+                request=body, model=body.get("model"), duration_ms=_ms_since(t0),
+                error=RuntimeError(f"upstream HTTP {upstream_resp.status_code}: {raw[:300]}"),
+                provider="anthropic",
+            )
+            return JSONResponse(
+                {"error": {"message": raw[:500]}}, status_code=upstream_resp.status_code
+            )
+
+        acc = _AnthropicAcc(body.get("model"))
+
+        async def anthropic_passthrough_gen():
+            try:
+                async for line in upstream_resp.aiter_lines():
+                    yield line + "\n"
+                    stripped = line.strip()
+                    if stripped.startswith("data:"):
+                        data = stripped[5:].strip()
+                        if data:
+                            with suppress(json.JSONDecodeError):
+                                acc.absorb(json.loads(data))
+            finally:
+                await upstream_resp.aclose()
+                usage = None
+                if acc.input_tokens is not None or acc.output_tokens is not None:
+                    usage = {"prompt_tokens": acc.input_tokens or 0,
+                             "completion_tokens": acc.output_tokens or 0}
+                recorder.record_llm(
+                    request=body, response=acc.payload(), usage=usage,
+                    model=acc.model, duration_ms=_ms_since(t0),
+                    meta={"provider": "anthropic"},
+                )
+
+        return StreamingResponse(anthropic_passthrough_gen(), media_type="text/event-stream")
+
+    return app
+
+
+def serve_proxy(
+    upstream: Optional[str] = None,
+    cassette: Optional[Cassette] = None,
+    runs_dir: str = "runs",
+    host: str = "127.0.0.1",
+    port: int = 8840,
+) -> None:
+    import uvicorn
+
+    app = create_proxy_app(upstream=upstream, cassette=cassette, runs_dir=runs_dir)
+    print(f"recording to: {app.state.recorder.path}")
+    uvicorn.run(app, host=host, port=port, log_level="warning")