feat: capture streamed answers whole; customer ids; outcome()

README.md

@@ -67,6 +67,25 @@ pip install -e ".[dev]"
 pytest tests
 ```
 
+## Outcomes and customers
+
+Tell AgentCompile how a conversation ended, so it learns only from the ones that went well:
+
+```python
+agentcompile.outcome(ticket.id, "resolved")  # or escalated, unresolved, abandoned, reopened, complaint
+```
+
+Pass a stable customer id so repeat jobs are counted per customer (scrubbed like everything
+else when capture is on):
+
+```python
+with agentcompile.conversation(ticket.id, customer=ticket.customer_id):
+    run_agent()
+```
+
+Streamed answers are captured too: your agent reads the stream as always, and the whole answer
+is captured once it ends.
+
 ## Privacy
 
 With `capture=True`, personal data is scrubbed on your machine before anything is sent:

src/agentcompile/__init__.py

@@ -19,10 +19,13 @@
 from ._conversation import conversation
 from ._core import LIVE, SHADOW, Proxy, Router, async_create, sync_create
 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", "SHADOW", "Settings", "conversation", "flush", "wrap"]
+__all__ = ["LIVE", "OUTCOMES", "SHADOW", "Settings", "conversation", "flush", "outcome", "wrap"]
 try:
     from ._version import __version__
 except ImportError:  # running from a source checkout that was never built
@@ -64,6 +67,7 @@ def wrap(
         timeout=timeout,
     )
     the_trail = Trail(trail, on_event)
+    _remember(settings)  # outcome() sends with the client wrapped last
     if capture is None:
         capture = os.environ.get("AGENTCOMPILE_CAPTURE", "") in ("1", "true", "yes")
     # One sender per wrapped client; it always sends from its own thread (sync client).
@@ -103,8 +107,10 @@ def wrap(
 
 
 def flush(timeout: float = 5.0) -> None:
-    """Send every captured call still queued (short scripts and tests; also runs at exit)."""
+    """Send every captured call and outcome still queued (short scripts and tests; also runs
+    at exit)."""
     flush_all(timeout)
+    _flush_outcomes(timeout)
 
 
 def _router(

src/agentcompile/_assemble.py

@@ -0,0 +1,188 @@
+"""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."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import AsyncIterator, Iterator
+from typing import Any, Callable
+
+from ._payload import jsonable
+
+OnDone = Callable[[list[Any], bool], None]  # (chunks as JSON, the stream ran to its end)
+
+
+def assemble(provider: str, chunks: list[Any]) -> dict[str, Any] | None:
+    try:
+        return _openai(chunks) if provider == "openai" else _anthropic(chunks)
+    except Exception:  # never break the agent over a shape we didn't expect
+        return None
+
+
+def _openai(chunks: list[dict[str, Any]]) -> dict[str, Any]:
+    first = chunks[0] if chunks else {}
+    content: list[str] = []
+    calls: dict[int, dict[str, Any]] = {}
+    finish = None
+    usage = None
+    for chunk in chunks:
+        usage = chunk.get("usage") or usage
+        for choice in chunk.get("choices") or []:
+            if choice.get("index", 0) != 0:
+                continue
+            delta = choice.get("delta") or {}
+            if delta.get("content"):
+                content.append(delta["content"])
+            for call in delta.get("tool_calls") or []:
+                slot = calls.setdefault(
+                    call.get("index", 0),
+                    {"id": None, "type": "function", "function": {"name": "", "arguments": ""}},
+                )
+                slot["id"] = call.get("id") or slot["id"]
+                fn = call.get("function") or {}
+                slot["function"]["name"] += fn.get("name") or ""
+                slot["function"]["arguments"] += fn.get("arguments") or ""
+            finish = choice.get("finish_reason") or finish
+    message: dict[str, Any] = {"role": "assistant", "content": "".join(content) or None}
+    if calls:
+        message["tool_calls"] = [calls[i] for i in sorted(calls)]
+    out: dict[str, Any] = {
+        "id": first.get("id"),
+        "object": "chat.completion",
+        "created": first.get("created"),
+        "model": first.get("model"),
+        "choices": [{"index": 0, "message": message, "finish_reason": finish}],
+    }
+    if usage:
+        out["usage"] = usage
+    return out
+
+
+def _anthropic(events: list[dict[str, Any]]) -> dict[str, Any]:
+    message: dict[str, Any] = {}
+    blocks: dict[int, dict[str, Any]] = {}
+    partial: dict[int, list[str]] = {}
+    for event in events:
+        kind = event.get("type")
+        if kind == "message_start":
+            message = dict(event.get("message") or {})
+        elif kind == "content_block_start":
+            blocks[event["index"]] = dict(event.get("content_block") or {})
+        elif kind == "content_block_delta":
+            block = blocks.setdefault(event["index"], {"type": "text", "text": ""})
+            delta = event.get("delta") or {}
+            if delta.get("type") == "text_delta":
+                block["text"] = block.get("text", "") + delta.get("text", "")
+            elif delta.get("type") == "input_json_delta":
+                partial.setdefault(event["index"], []).append(delta.get("partial_json", ""))
+            elif delta.get("type") == "thinking_delta":
+                block["thinking"] = block.get("thinking", "") + delta.get("thinking", "")
+        elif kind == "message_delta":
+            delta = event.get("delta") or {}
+            message["stop_reason"] = delta.get("stop_reason", message.get("stop_reason"))
+            message["stop_sequence"] = delta.get("stop_sequence")
+            usage = dict(message.get("usage") or {})
+            usage.update(event.get("usage") or {})
+            message["usage"] = usage
+    for index, pieces in partial.items():
+        text = "".join(pieces)
+        blocks[index]["input"] = json.loads(text) if text else {}
+    message["content"] = [blocks[i] for i in sorted(blocks)]
+    return message
+
+
+class CapturingStream:
+    """Your stream, unchanged, with a copy of each chunk kept for capture."""
+
+    def __init__(self, stream: Any, on_done: OnDone) -> None:
+        self._stream = stream
+        self._on_done = on_done
+        self._chunks: list[Any] = []
+        self._finished = False
+
+    def __iter__(self) -> Iterator[Any]:
+        complete = False
+        try:
+            for item in self._stream:
+                self._chunks.append(jsonable(item))
+                yield item
+            complete = True
+        finally:
+            self._finish(complete=complete)
+
+    def __enter__(self) -> CapturingStream:
+        enter = getattr(self._stream, "__enter__", None)
+        if enter is not None:
+            enter()
+        return self
+
+    def __exit__(self, *exc: Any) -> Any:
+        self._finish(complete=False)
+        leave = getattr(self._stream, "__exit__", None)
+        return leave(*exc) if leave is not None else None
+
+    def close(self) -> None:
+        self._finish(complete=False)
+        close = getattr(self._stream, "close", None)
+        if close is not None:
+            close()
+
+    def _finish(self, *, complete: bool) -> None:
+        if not self._finished:
+            self._finished = True
+            self._on_done(self._chunks, complete)
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._stream, name)
+
+
+class AsyncCapturingStream:
+    """The async version."""
+
+    def __init__(self, stream: Any, on_done: OnDone) -> None:
+        self._stream = stream
+        self._on_done = on_done
+        self._chunks: list[Any] = []
+        self._finished = False
+
+    async def _iterate(self) -> AsyncIterator[Any]:
+        complete = False
+        try:
+            async for item in self._stream:
+                self._chunks.append(jsonable(item))
+                yield item
+            complete = True
+        finally:
+            self._finish(complete=complete)
+
+    def __aiter__(self) -> AsyncIterator[Any]:
+        return self._iterate()
+
+    async def __aenter__(self) -> AsyncCapturingStream:
+        enter = getattr(self._stream, "__aenter__", None)
+        if enter is not None:
+            await enter()
+        return self
+
+    async def __aexit__(self, *exc: Any) -> Any:
+        self._finish(complete=False)
+        leave = getattr(self._stream, "__aexit__", None)
+        return (await leave(*exc)) if leave is not None else None
+
+    async def close(self) -> None:
+        self._finish(complete=False)
+        close = getattr(self._stream, "close", None)
+        if close is not None:
+            result = close()
+            if hasattr(result, "__await__"):
+                await result
+
+    def _finish(self, *, complete: bool) -> None:
+        if not self._finished:
+            self._finished = True
+            self._on_done(self._chunks, complete)
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._stream, name)

src/agentcompile/_capture.py

@@ -18,9 +18,10 @@
 
 import httpx
 
+from ._assemble import assemble
 from ._decide import CONVERSATION_HEADER, Settings, _request
 from ._payload import jsonable, payload
-from ._scrub import scrub_call
+from ._scrub import scrub_call, scrub_value
 
 MAX_QUEUE = 2000
 BATCH = 50
@@ -63,9 +64,11 @@ def add(
         kwargs: dict[str, Any],
         response: Any,
         stream: bool = False,
+        streamed: dict[str, Any] | None = None,
+        customer: str | None = None,
     ) -> None:
-        """Queue one call. Never raises. A streamed answer is not captured yet: the record
-        keeps the request and says so."""
+        """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)
             if self.scrub_key is not None:
@@ -78,10 +81,20 @@ def add(
                 "request": request,
                 "response": answer,
             }
+            if customer:
+                # The customer's id, scrubbed like everything else (an email becomes a token).
+                record["end_user"] = (
+                    scrub_value(customer, self.scrub_key, "customer")
+                    if self.scrub_key is not None
+                    else customer
+                )
             if self.scrub_key is not None:
                 record["scrubbed"] = True
             if stream:
                 record["stream"] = True
+            if streamed is not None:
+                record["stream"] = True
+                record["stream_complete"] = streamed["complete"]
             with self._lock:
                 if len(self._queue) == self._queue.maxlen:
                     self.dropped += 1  # the deque drops the oldest
@@ -92,6 +105,26 @@ def add(
         except Exception:
             self.dropped += 1
 
+    def add_stream(
+        self,
+        provider: str,
+        conversation_id: str | None,
+        kwargs: dict[str, Any],
+        chunks: list[Any],
+        complete: bool,
+        customer: str | None = None,
+    ) -> None:
+        """Queue a streamed call once its stream is done: the chunks assembled into the answer
+        a non-streamed call would have had. A stream stopped early is kept and marked."""
+        self.add(
+            provider,
+            conversation_id,
+            kwargs,
+            assemble(provider, chunks),
+            streamed={"complete": complete, "chunks": len(chunks)},
+            customer=customer,
+        )
+
     def flush(self, timeout: float = 5.0) -> None:
         """Send everything queued now (tests, and at exit)."""
         deadline = time.monotonic() + timeout

src/agentcompile/_conversation.py

@@ -9,21 +9,31 @@
 _current: contextvars.ContextVar[str | None] = contextvars.ContextVar(
     "agentcompile_conversation", default=None
 )
+_customer: contextvars.ContextVar[str | None] = contextvars.ContextVar(
+    "agentcompile_customer", default=None
+)
 
 
 @contextlib.contextmanager
-def conversation(conversation_id: str) -> Iterator[None]:
-    """Every model call inside the block belongs to this conversation:
+def conversation(conversation_id: str, customer: str | None = None) -> Iterator[None]:
+    """Every model call inside the block belongs to this conversation (and, optionally, this
+    customer: a stable id for the person, so repeat jobs can be counted per customer):
 
-    with agentcompile.conversation(ticket.id):
+    with agentcompile.conversation(ticket.id, customer=ticket.customer_id):
         run_agent(...)
     """
     token = _current.set(conversation_id)
+    customer_token = _customer.set(customer)
     try:
         yield
     finally:
+        _customer.reset(customer_token)
         _current.reset(token)
 
 
 def current() -> str | None:
     return _current.get()
+
+
+def current_customer() -> str | None:
+    return _customer.get()