From a74420ed2c1a3fe95e1ccd3cebdbef1f9757da66 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:42:05 +0100 Subject: [PATCH] Trace unexecuted tool calls in one place; tidy span lifecycles Paused, denied, skipped and client-run tool calls all go through trace_unexecuted_tool_call. The research synthesis, continuation and workflow step spans use with-blocks instead of hand-rolled GeneratorExit handling. Fixes blank lines and adds a missing hint and docstrings. --- docsgpt/agents/base.py | 21 ++++-------------- docsgpt/agents/research_agent.py | 22 ++++++------------- docsgpt/agents/tool_executor.py | 16 ++++++++++++++ docsgpt/agents/workflows/workflow_engine.py | 7 +----- docsgpt/guardrails/engine.py | 6 +++-- docsgpt/llm/handlers/base.py | 18 ++++----------- .../storage/db/repositories/request_traces.py | 2 ++ 7 files changed, 38 insertions(+), 54 deletions(-) diff --git a/docsgpt/agents/base.py b/docsgpt/agents/base.py index 3dea7254..4f620236 100644 --- a/docsgpt/agents/base.py +++ b/docsgpt/agents/base.py @@ -6,11 +6,9 @@ from abc import ABC, abstractmethod from datetime import datetime, timezone from typing import Any, Dict, Generator, List, Optional -from docsgpt import tracing from docsgpt.agents.tool_executor import ( ToolExecutor, - finish_tool_span, - record_tool_span_start, + trace_unexecuted_tool_call, result_status, truncate_tool_result, ) @@ -504,18 +502,10 @@ class BaseAgent(ABC): pending_tool_calls: The pending tool call descriptors from the pause. tool_actions: Client-provided actions resolving the pending calls. """ - span = start_agent_span(self, continuation=True) - completed = False - try: + with start_agent_span(self, continuation=True): yield from self._gen_continuation_inner( messages, tools_dict, pending_tool_calls, tool_actions, reasoning_content ) - completed = True - except Exception as exc: - span.end(error=exc) - raise - finally: - span.end(None if completed else tracing.STATUS_CANCELLED) def _gen_continuation_inner( self, @@ -617,7 +607,7 @@ class BaseAgent(ABC): "arguments": args, "status": "denied", } - finish_tool_span(record_tool_span_start(tc), {**denied_data, "error": comment or None}) + trace_unexecuted_tool_call(tc, {**denied_data, "error": comment or None}) yield {"type": "tool_call", "data": denied_data} elif "result" in action: @@ -650,10 +640,7 @@ class BaseAgent(ABC): "result": truncate_tool_result(result_str), "status": result_status(result), } - finish_tool_span( - record_tool_span_start(tc, **{"docsgpt.client_executed": True}), - client_data, - ) + trace_unexecuted_tool_call(tc, client_data, **{"docsgpt.client_executed": True}) yield {"type": "tool_call", "data": client_data} # Resume the LLM loop with the updated messages diff --git a/docsgpt/agents/research_agent.py b/docsgpt/agents/research_agent.py index c4858ec0..859bc872 100644 --- a/docsgpt/agents/research_agent.py +++ b/docsgpt/agents/research_agent.py @@ -47,12 +47,6 @@ STEP_PROMPT = _load_prompt("step.txt") SYNTHESIS_PROMPT = _load_prompt("synthesis.txt") -# --------------------------------------------------------------------------- -# CitationManager -# --------------------------------------------------------------------------- - - - def _phase_span(phase: str, **attributes: Any) -> Any: """Open a trace span for one research phase (clarify, plan, a step, synthesis).""" return tracing.start_span( @@ -61,6 +55,12 @@ def _phase_span(phase: str, **attributes: Any) -> Any: attributes={"docsgpt.research.phase": phase.split(" ")[0], **attributes}, ) + +# --------------------------------------------------------------------------- +# CitationManager +# --------------------------------------------------------------------------- + + class CitationManager: """Tracks and deduplicates citations across research steps.""" @@ -266,18 +266,10 @@ class ResearchAgent(BaseAgent): "tokens_used": self._tokens_used, }, } - synthesis_span = _phase_span("synthesis") - try: + with _phase_span("synthesis"): yield from self._synthesis_phase( query, plan, intermediate_reports, tools_dict, log_context ) - except GeneratorExit: - synthesis_span.end(tracing.STATUS_CANCELLED) - raise - except Exception as exc: - synthesis_span.end(error=exc) - raise - synthesis_span.end() # Sources and tool calls self.retrieved_docs = self.citations.get_all_docs() diff --git a/docsgpt/agents/tool_executor.py b/docsgpt/agents/tool_executor.py index 8e2b65e1..bc7f809a 100644 --- a/docsgpt/agents/tool_executor.py +++ b/docsgpt/agents/tool_executor.py @@ -57,6 +57,22 @@ _TOOL_SPAN_STATUS = { } +def trace_unexecuted_tool_call(call: Any, data: Dict[str, Any], **attributes: Any) -> None: + """Record a tool call that this process did not execute. + + The single entry point for calls that never reach :meth:`ToolExecutor.execute` + (which traces executed calls itself): calls paused for approval or for the + client, denied ones, ones skipped at the context limit, and results a client + sends back on resume. A new pause or refusal path records its call here. + + Args: + call: The tool call (anything with ``name`` and ``id``). + data: The tool-call record emitted for it (``status``, ``arguments``, ...). + **attributes: Extra span attributes. + """ + finish_tool_span(record_tool_span_start(call, **attributes), data) + + def finish_tool_span(span: Any, data: Dict[str, Any]) -> None: """Close a tool span from the ``tool_calls`` entry the call produced. diff --git a/docsgpt/agents/workflows/workflow_engine.py b/docsgpt/agents/workflows/workflow_engine.py index f296fc05..3b80dc81 100644 --- a/docsgpt/agents/workflows/workflow_engine.py +++ b/docsgpt/agents/workflows/workflow_engine.py @@ -164,12 +164,8 @@ class WorkflowEngine: } try: - try: + with step_span: yield from self._execute_node(node) - except GeneratorExit: - step_span.end(tracing.STATUS_CANCELLED) - raise - step_span.end() log_entry["status"] = ExecutionStatus.COMPLETED.value self._finalize_log_entry(log_entry, pre_state) @@ -186,7 +182,6 @@ class WorkflowEngine: "output": node_output, } except Exception as e: - step_span.end(error=e) logger.error(f"Error executing node {node.id}: {e}", exc_info=True) log_entry["status"] = ExecutionStatus.FAILED.value log_entry["error"] = str(e) diff --git a/docsgpt/guardrails/engine.py b/docsgpt/guardrails/engine.py index 834edc36..5f42cf02 100644 --- a/docsgpt/guardrails/engine.py +++ b/docsgpt/guardrails/engine.py @@ -5,7 +5,7 @@ from __future__ import annotations import logging from threading import Thread from time import monotonic -from typing import List, Optional +from typing import Any, List, Optional, Sequence from docsgpt import tracing from docsgpt.guardrails.base import ScanContext @@ -27,7 +27,8 @@ logger = logging.getLogger(__name__) _MAX_WORKERS = 8 -def _start_guardrail_span(stage: Stage, controls) -> "tracing.Span": +def _start_guardrail_span(stage: Stage, controls: Sequence[Any]) -> "tracing.Span": + """Open a ``guardrail`` span for evaluating ``controls`` at ``stage``.""" return tracing.start_span( tracing.KIND_GUARDRAIL, f"guardrail {stage.value}", @@ -39,6 +40,7 @@ def _start_guardrail_span(stage: Stage, controls) -> "tracing.Span": def _content_fired(decision: StageDecision) -> bool: + """True when a check matched, or the text was blocked or redacted.""" return bool(decision.triggered or decision.blocked or decision.redacted) diff --git a/docsgpt/llm/handlers/base.py b/docsgpt/llm/handlers/base.py index c52f1823..0954dd53 100644 --- a/docsgpt/llm/handlers/base.py +++ b/docsgpt/llm/handlers/base.py @@ -5,6 +5,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass from typing import Any, Dict, Generator, List, Optional, Union +from docsgpt.agents.tool_executor import trace_unexecuted_tool_call from docsgpt.logging import build_stack_data logger = logging.getLogger(__name__) @@ -97,17 +98,6 @@ class LLMResponse: return bool(self.tool_calls) and self.finish_reason == "tool_calls" -def _trace_unexecuted_tool(call: Any, data: Dict[str, Any]) -> None: - """Record a tool call that was paused, denied or skipped instead of run. - - Those calls never reach ``ToolExecutor.execute`` (which traces executed - calls), so without this the trace would not show them at all. - """ - from docsgpt.agents.tool_executor import finish_tool_span, record_tool_span_start - - finish_tool_span(record_tool_span_start(call), data) - - def _is_restated_payload(existing: str, incoming: str) -> bool: """True when ``incoming`` restates an already-complete ``existing`` payload. @@ -1138,7 +1128,7 @@ class LLMHandler(ABC): # Mark remaining tools as skipped for remaining_call in tool_calls[i:]: - _trace_unexecuted_tool( + trace_unexecuted_tool_call( remaining_call, {"tool_name": "system", "action_name": remaining_call.name, "status": "skipped"}, ) @@ -1235,7 +1225,7 @@ class LLMHandler(ABC): "error_type", "tool_not_allowed" ), } - _trace_unexecuted_tool(call, denied_data) + trace_unexecuted_tool_call(call, denied_data) yield {"type": "tool_call", "data": denied_data} continue # Yield pause event so the client knows this tool is waiting @@ -1250,7 +1240,7 @@ class LLMHandler(ABC): # can wire the sticky "don't ask again" button. if pause_info.get("device_id"): pause_data["device_id"] = pause_info["device_id"] - _trace_unexecuted_tool(call, pause_data) + trace_unexecuted_tool_call(call, pause_data) yield {"type": "tool_call", "data": pause_data} pending_actions.append(pause_info) # Do NOT add messages for pending tools here. diff --git a/docsgpt/storage/db/repositories/request_traces.py b/docsgpt/storage/db/repositories/request_traces.py index 1fb5c594..79a7e3cb 100644 --- a/docsgpt/storage/db/repositories/request_traces.py +++ b/docsgpt/storage/db/repositories/request_traces.py @@ -33,10 +33,12 @@ _SUMMARY_COLUMNS = ( def _dump_jsonb(value: Any) -> str: + """Serialize ``value`` for a JSONB parameter, without NUL bytes Postgres rejects.""" return json.dumps(strip_null_bytes(value), cls=PGNativeJSONEncoder) def _uuid_or_none(value: Optional[str]) -> Optional[str]: + """``value`` as a string when it is a UUID, else ``None`` (the column is typed).""" return str(value) if value and looks_like_uuid(str(value)) else None