diff --git a/CHANGELOG.md b/CHANGELOG.md index df046f8..c6fb8be 100644 --- a/CHANGELOG.md +++ b/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), diff --git a/config/schemas/v1/agent-investigation.schema.json b/config/schemas/v1/agent-investigation.schema.json index 133da8d..f691a7b 100644 --- a/config/schemas/v1/agent-investigation.schema.json +++ b/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, diff --git a/sdk/src/cloudfall/app.py b/sdk/src/cloudfall/app.py index 70d2b9b..1c5f070 100644 --- a/sdk/src/cloudfall/app.py +++ b/sdk/src/cloudfall/app.py @@ -140,6 +140,7 @@ InvestigationStatus, InvestigationStore, ModelEndpoint, + ModelReply, ModelUnavailableError, Progress, StepOutcome, @@ -2231,9 +2232,10 @@ 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 @@ -2241,6 +2243,7 @@ class InvestigationEvent: 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 diff --git a/sdk/src/cloudfall/investigation.py b/sdk/src/cloudfall/investigation.py index 6c19f25..5edb975 100644 --- a/sdk/src/cloudfall/investigation.py +++ b/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: diff --git a/sdk/tests/test_investigation.py b/sdk/tests/test_investigation.py index 1f48fbc..e0b4786 100644 --- a/sdk/tests/test_investigation.py +++ b/sdk/tests/test_investigation.py @@ -17,6 +17,7 @@ InvestigationStatus, InvestigationStore, ModelEndpoint, + ModelReply, ModelUnavailableError, StepOutcome, ToolResult, @@ -74,10 +75,10 @@ def __init__(self, *replies: Mapping[str, object]) -> None: self.replies = list(replies) self.requests: list[Mapping[str, object]] = [] - def __call__(self, url: str, body: Mapping[str, object]) -> Mapping[str, object]: + def __call__(self, url: str, body: Mapping[str, object]) -> ModelReply: assert url == "https://models.example/v1/chat/completions" self.requests.append(json.loads(json.dumps(body))) - return self.replies.pop(0) + return ModelReply(body=self.replies.pop(0), request_id="req-test") class _Tool: @@ -270,7 +271,7 @@ def test_a_malformed_model_answer_is_refused() -> None: class _Failing(_Script): """A model that answers from a list, then its endpoint refuses.""" - def __call__(self, url: str, body: Mapping[str, object]) -> Mapping[str, object]: + def __call__(self, url: str, body: Mapping[str, object]) -> ModelReply: if not self.replies: message = "the model endpoint answered 500: overloaded" raise InvestigationError(ERROR_MODEL_REFUSED, message) @@ -802,6 +803,11 @@ def test_the_cli_streams_each_turn_and_call_then_the_record( assert events[1]["effect"] == "created" assert "/var/log/shop holds 36G" in events[1]["result"]["output"] assert events[2]["text"].startswith("ROOT CAUSE") + # A turn carries its raw trace; this endpoint sends no x-request-id. + trace = events[0]["trace"] + assert trace["toolCalls"][0]["name"] == "operation_facts" + assert "requestId" not in trace + assert "sk-test" not in json.dumps(events) # The decision the call wrote, and the saved record. assert stream["end"]["_summary"] is True assert stream["end"]["effects"] == {"created": 2, "noop": 2} @@ -828,3 +834,66 @@ def test_an_incomplete_investigation_ends_with_the_record_then_93( # The error line holds no data: the record is the event before it. assert "data.investigation" not in error["suggestion"] assert "kind investigation" in error["suggestion"] + + +def test_each_turn_keeps_the_raw_trace() -> None: + report, rotate = _tools() + thinking = _call("disk-report", {}) + thinking["id"] = "chatcmpl-1" + first = thinking["choices"] + assert isinstance(first, list) + first[0]["message"]["reasoning_content"] = "The disk may be full; read it first." + thinking["usage"] = { + "prompt_tokens": 100, + "completion_tokens": 40, + "completion_tokens_details": {"reasoning_tokens": 30}, + } + script = _Script( + thinking, _say("ROOT CAUSE: disk\nPROPOSED: none\nWHY: read only.") + ) + + done = investigate("alert", _tool_list(report, rotate), ENDPOINT, script) + + turn = done.trace[0] + assert turn.reasoning == "The disk may be full; read it first." + assert turn.tool_calls[0].name == "disk-report" + assert turn.tool_calls[0].arguments == "{}" + assert (turn.response_id, turn.request_id) == ("chatcmpl-1", "req-test") + assert (turn.prompt_tokens, turn.completion_tokens, turn.reasoning_tokens) == ( + 100, + 40, + 30, + ) + recorded = done.as_document()["spec"] + assert isinstance(recorded, dict) + assert recorded["trace"][0]["reasoning"] == turn.reasoning + assert recorded["trace"][1]["toolCalls"] == [] + + +class _LongIds(_Script): + """An endpoint that serves its answers under over-long ids.""" + + def __call__(self, url: str, body: Mapping[str, object]) -> ModelReply: + reply = super().__call__(url, body) + return ModelReply(body=reply.body, request_id="r" * 300) + + +def test_an_over_long_endpoint_id_still_leaves_a_valid_record( + tmp_path: Path, +) -> None: + report, rotate = _tools() + answer = _say("ROOT CAUSE: disk\nPROPOSED: none\nWHY: read only.") + answer["id"] = "chatcmpl-" + "x" * 300 + script = _LongIds(answer) + done = investigate("alert", _tool_list(report, rotate), ENDPOINT, script) + + saved = InvestigationStore( + directory=tmp_path / "investigations", + catalog=SchemaCatalog(default_schema_directory()), + ).save(done) + + turn = saved.trace[0] + assert turn.response_id is not None + assert turn.request_id is not None + assert turn.response_id.startswith("chatcmpl-") + assert (len(turn.response_id), len(turn.request_id)) == (256, 256)