Keep each model turn's raw trace in the stream and the record

CHANGELOG.md

@@ -8,6 +8,13 @@ Notable changes to Cloudfall. The format follows
 
 ### Added
 
+- Every turn of `agent investigate` keeps its raw trace, in the stream's
+  `turn` event (`trace`) and in the investigation record (`spec.trace`): the
+  model's reasoning when it returns one (Nemotron's `reasoning_content`),
+  each tool call exactly as written (name and raw arguments), the
+  endpoint's response id and `x-request-id`, the milliseconds the turn took,
+  and its prompt, completion and reasoning tokens. The record can now say
+  what the model thought when it proposed a change
 - `agent investigate` streams: one bare JSON line with `_seq` per model turn
   (its words and how many tools it called), one per tool call (the step and
   what the tool handed back, `effect: created` when it wrote a decision),

config/schemas/v1/agent-investigation.schema.json

@@ -117,6 +117,44 @@
         "answer": {
           "type": "string"
         },
+        "trace": {
+          "type": "array",
+          "items": {
+            "type": "object",
+            "properties": {
+              "turn": {"type": "integer", "minimum": 1},
+              "text": {"type": "string"},
+              "reasoning": {"type": "string"},
+              "toolCalls": {
+                "type": "array",
+                "items": {
+                  "type": "object",
+                  "properties": {
+                    "name": {"type": "string"},
+                    "arguments": {"type": "string"}
+                  },
+                  "required": ["name", "arguments"],
+                  "additionalProperties": false
+                }
+              },
+              "responseId": {"type": "string", "maxLength": 256},
+              "requestId": {"type": "string", "maxLength": 256},
+              "elapsedMs": {"type": "integer", "minimum": 0},
+              "tokens": {
+                "type": "object",
+                "properties": {
+                  "prompt": {"type": "integer", "minimum": 0},
+                  "completion": {"type": "integer", "minimum": 0},
+                  "reasoning": {"type": "integer", "minimum": 0}
+                },
+                "required": ["prompt", "completion", "reasoning"],
+                "additionalProperties": false
+              }
+            },
+            "required": ["turn", "text", "reasoning", "toolCalls", "elapsedMs", "tokens"],
+            "additionalProperties": false
+          }
+        },
         "replay": {
           "type": "string",
           "minLength": 1,

sdk/src/cloudfall/app.py

@@ -140,6 +140,7 @@
     InvestigationStatus,
     InvestigationStore,
     ModelEndpoint,
+    ModelReply,
     ModelUnavailableError,
     Progress,
     StepOutcome,
@@ -2231,16 +2232,18 @@ def endpoint(self) -> ModelEndpoint:
 class InvestigationEvent:
     """``agent investigate``: one event of the stream.
 
-    ``turn`` is the model's words and how many tools it called; ``call`` is one
-    tool call and what the tool handed back; the last event, ``investigation``,
-    is the record as it was saved.
+    ``turn`` is one model turn as the endpoint answered it (``trace``: its
+    reasoning, the tool calls as written, the response and request ids, the
+    time and the tokens); ``call`` is one tool call and what the tool handed
+    back; the last event, ``investigation``, is the record as it was saved.
     """
 
     effect: str
     kind: str
     turn: int | None = None
     text: str | None = None
     calls: int | None = None
+    trace: dict[str, object] = Out(default_factory=dict, ordered=True)
     step: dict[str, object] = Out(default_factory=dict, ordered=True)
     result: dict[str, object] = Out(default_factory=dict, ordered=True)
     status: str | None = None
@@ -2345,6 +2348,7 @@ def _progress_event(progress: Progress) -> InvestigationEvent:
             turn=progress.turn,
             text=progress.text,
             calls=progress.calls,
+            trace=progress.as_dict(),
         )
     wrote = (
         progress.step.decision is not None
@@ -2444,7 +2448,7 @@ def _tool_result(
 def _chat_transport(ctx: Ctx, api_key: str) -> ChatTransport:
     """Send one chat request through ``ctx.http``."""
 
-    def send(url: str, body: Mapping[str, object]) -> Mapping[str, object]:
+    def send(url: str, body: Mapping[str, object]) -> ModelReply:
         try:
             response = ctx.http.post(
                 url,
@@ -2473,7 +2477,10 @@ def send(url: str, body: Mapping[str, object]) -> Mapping[str, object]:
         if not isinstance(answer, Mapping):
             message = "the model endpoint answered JSON that is not an object"
             raise InvestigationError(ERROR_MODEL_ANSWER_MALFORMED, message)
-        return cast("Mapping[str, object]", answer)
+        return ModelReply(
+            body=cast("Mapping[str, object]", answer),
+            request_id=response.headers.get("x-request-id"),
+        )
 
     return send
 

sdk/src/cloudfall/investigation.py

@@ -16,6 +16,7 @@
 
 import json
 import re
+import time
 from collections.abc import Callable, Generator, Mapping, Sequence
 from dataclasses import dataclass, replace
 from datetime import UTC, datetime
@@ -39,6 +40,8 @@
 """Most characters of host output one tool result hands the model."""
 TOOL_NAME_LIMIT = 128
 """Most characters of a tool name the record keeps, as its schema allows."""
+REPLY_ID_LIMIT = 256
+"""Most characters of an endpoint's response or request id the trace keeps."""
 SAVE_ATTEMPTS = 100
 """Most investigations recorded under one second's id, with their suffixes."""
 ENDPOINT_LIMIT = 512
@@ -164,11 +167,18 @@ class AgentTool:
     call: Callable[[Mapping[str, object]], ToolResult]
 
 
-ChatTransport = Callable[[str, Mapping[str, object]], Mapping[str, object]]
-"""Send one chat request body to a URL and return the answer's JSON.
+@dataclass(frozen=True, slots=True)
+class ModelReply:
+    """One answer of the endpoint: its JSON, and the request id it was served under."""
 
-A failed request raises ``InvestigationError``.
-"""
+    body: Mapping[str, object]
+    request_id: str | None = None
+
+
+ChatTransport = Callable[[str, Mapping[str, object]], ModelReply]
+"""Send one chat request body to a URL and return the endpoint's reply.
+
+A failed request raises ``InvestigationError``."""
 
 
 @dataclass(frozen=True, slots=True)
@@ -247,6 +257,8 @@ class Investigation:
     finding: Finding | None
     replay: str | None = None
     """The recording played back instead of the hosts, when there was one."""
+    trace: tuple[Turn, ...] = ()
+    """Every turn as the endpoint answered it, reasoning included."""
 
     def as_document(self) -> dict[str, object]:
         """Serialize the investigation as a schema-valid record."""
@@ -268,6 +280,8 @@ def as_document(self) -> dict[str, object]:
             spec["finding"] = self.finding.as_dict()
         if self.replay is not None:
             spec["replay"] = self.replay
+        if self.trace:
+            spec["trace"] = [turn.as_dict() for turn in self.trace]
         return {
             "apiVersion": "cloudfall/v1",
             "kind": "AgentInvestigation",
@@ -322,13 +336,62 @@ def save(self, investigation: Investigation) -> Investigation:
         raise InvestigationError(ERROR_INVESTIGATION_EXISTS, message)
 
 
+@dataclass(frozen=True, slots=True)
+class ToolCallText:
+    """One tool call exactly as the model wrote it: the name and the raw arguments."""
+
+    name: str
+    arguments: str
+
+    def as_dict(self) -> dict[str, object]:
+        """Serialize the call for the record."""
+        return {"name": self.name, "arguments": self.arguments}
+
+
 @dataclass(frozen=True, slots=True)
 class Turn:
-    """The model answered one turn: its words, and how many tools it called."""
+    """One turn of the model, as the endpoint answered it.
+
+    Its words, its reasoning when the model returns it (Nemotron does, as
+    ``reasoning_content``), the tool calls as written, the ids the endpoint
+    served it under, how long it took and the tokens it cost.
+    """
 
     turn: int
     text: str
-    calls: int
+    reasoning: str
+    tool_calls: tuple[ToolCallText, ...]
+    response_id: str | None
+    request_id: str | None
+    elapsed_ms: int
+    prompt_tokens: int
+    completion_tokens: int
+    reasoning_tokens: int
+
+    @property
+    def calls(self) -> int:
+        """How many tools the model called this turn."""
+        return len(self.tool_calls)
+
+    def as_dict(self) -> dict[str, object]:
+        """Serialize the turn for the record's trace."""
+        document: dict[str, object] = {
+            "turn": self.turn,
+            "text": self.text,
+            "reasoning": self.reasoning,
+            "toolCalls": [call.as_dict() for call in self.tool_calls],
+            "elapsedMs": self.elapsed_ms,
+            "tokens": {
+                "prompt": self.prompt_tokens,
+                "completion": self.completion_tokens,
+                "reasoning": self.reasoning_tokens,
+            },
+        }
+        if self.response_id is not None:
+            document["responseId"] = self.response_id
+        if self.request_id is not None:
+            document["requestId"] = self.request_id
+        return document
 
 
 @dataclass(frozen=True, slots=True)
@@ -392,6 +455,7 @@ def investigation_events(  # noqa: PLR0913 - as investigate.
     steps: list[Step] = []
     seen: dict[str, ToolResult] = {}
     usage = Usage()
+    trace: list[Turn] = []
     answer = ""
     status = InvestigationStatus.TURN_LIMIT
     turns = 0
@@ -412,11 +476,13 @@ def record(ending: InvestigationStatus, finding: Finding | None) -> Investigatio
             answer=answer,
             finding=finding,
             replay=replay,
+            trace=tuple(trace),
         )
 
     try:
         while turns < max_turns:
             turns += 1
+            asked = time.monotonic()
             reply = transport(
                 endpoint.chat_url,
                 {
@@ -426,14 +492,17 @@ def record(ending: InvestigationStatus, finding: Finding | None) -> Investigatio
                     "temperature": 0.2,
                 },
             )
-            usage = usage.plus(reply)
-            message = _first_message(reply)
+            elapsed_ms = int((time.monotonic() - asked) * 1000)
+            usage = usage.plus(reply.body)
+            message = _first_message(reply.body)
             messages.append(message)
             calls = cast(
                 "Sequence[Mapping[str, object]]", message.get("tool_calls") or ()
             )
             text = str(message.get("content") or "").strip()
-            yield Turn(turn=turns, text=text, calls=len(calls))
+            turn = _turn(turns, reply, message, calls, elapsed_ms)
+            trace.append(turn)
+            yield turn
             if not calls:
                 answer = text
                 status = InvestigationStatus.UNSTRUCTURED
@@ -482,6 +551,49 @@ def _tool_definition(tool: AgentTool) -> dict[str, object]:
     }
 
 
+def _turn(
+    number: int,
+    reply: ModelReply,
+    message: Mapping[str, object],
+    calls: Sequence[Mapping[str, object]],
+    elapsed_ms: int,
+) -> Turn:
+    """Read one turn's raw trace out of the endpoint's answer."""
+    usage = reply.body.get("usage")
+    usage = usage if isinstance(usage, Mapping) else {}
+    details = usage.get("completion_tokens_details")
+    details = details if isinstance(details, Mapping) else {}
+    reasoning = message.get("reasoning_content") or message.get("reasoning") or ""
+    response_id = reply.body.get("id")
+    # The ids are the endpoint's text; the record keeps what its schema takes.
+    return Turn(
+        turn=number,
+        text=str(message.get("content") or "").strip(),
+        reasoning=str(reasoning).strip(),
+        tool_calls=tuple(_call_text(call) for call in calls),
+        response_id=(
+            response_id[:REPLY_ID_LIMIT] if isinstance(response_id, str) else None
+        ),
+        request_id=(
+            None if reply.request_id is None else reply.request_id[:REPLY_ID_LIMIT]
+        ),
+        elapsed_ms=elapsed_ms,
+        prompt_tokens=_count(usage, "prompt_tokens"),
+        completion_tokens=_count(usage, "completion_tokens"),
+        reasoning_tokens=_count(details, "reasoning_tokens"),
+    )
+
+
+def _call_text(call: Mapping[str, object]) -> ToolCallText:
+    function = call.get("function")
+    function = function if isinstance(function, Mapping) else {}
+    arguments = function.get("arguments", "")
+    return ToolCallText(
+        name=str(function.get("name", "")),
+        arguments=arguments if isinstance(arguments, str) else json.dumps(arguments),
+    )
+
+
 def _first_message(reply: Mapping[str, object]) -> dict[str, object]:
     choices = reply.get("choices")
     if not isinstance(choices, list) or not choices: