diff --git a/docs/content/Deploying/Observability.mdx b/docs/content/Deploying/Observability.mdx index a0d8afaf..4643fcd0 100644 --- a/docs/content/Deploying/Observability.mdx +++ b/docs/content/Deploying/Observability.mdx @@ -14,8 +14,9 @@ launch command with `opentelemetry-instrument` and setting OTLP env vars. Auto-instrumentation covers Flask, Starlette, Celery, SQLAlchemy, -psycopg, Redis, requests, and Python logging. LLM/retriever calls are -not captured at this layer — see *Going further* below. +psycopg, Redis, requests, and Python logging. Agent runs, LLM calls, +tool calls and retrieval are recorded by DocsGPT itself and exported as +OpenTelemetry GenAI spans — see [Execution traces](#execution-traces). ## Enabling @@ -67,6 +68,88 @@ dotenv run -- opentelemetry-instrument celery -A docsgpt.app.celery worker -l IN the OTEL log handler. Without it, `logging` writes only to stdout. +## Execution traces + +Every request records an **execution trace**: a timed tree of the steps +behind it. Traces are recorded for chat turns (`/stream`, `/api/answer`, +`/v1/chat/completions`, including each round of a tool-approval pause), +scheduled and webhook runs, workflows, the research agent, `/api/search`, +the MCP `search_docs` tool, and graph builds. + +| Step | Recorded when | +| --- | --- | +| `invoke_agent` | An agent (or a workflow node's agent) runs | +| `chat` | An LLM call made during the request, including retries, fallbacks, query rephrasing, prescreening, history compression and guardrail judges | +| `execute_tool` | A tool is executed, paused for approval, denied or skipped | +| `retrieval` | A retriever or the multi-source dispatcher searches | +| `embeddings` | The query is embedded | +| `search` | One source is searched | +| `rerank` | Prescreening filters retrieved chunks | +| `guardrail` | A guardrail calls a remote check or fires | +| `step` | A workflow node or research phase runs | + +Traces are stored in the `request_traces` table and shown in the app: open +**Settings → Logs** (or an agent's **Logs** tab), expand an entry and choose +**View trace** to see a waterfall of every step with its timing, tokens, +cost and details. + +Stored traces keep short previews — tool arguments and results, retrieved +chunk titles and snippets, rephrased queries, answer excerpts — truncated +and with secret-named fields redacted. Full prompts are never stored. When a +guardrail fires during a request, every preview is dropped from its trace. + +```bash +TRACES_ENABLED=true # record traces at all +TRACES_CAPTURE_CONTENT=true # keep previews in stored traces +TRACES_PREVIEW_CHARS=2000 # characters kept per preview +TRACES_MAX_SPANS=500 # steps kept per trace; the rest are counted +TRACES_RETENTION_DAYS=30 # a daily task deletes older traces +TRACES_OTEL_EXPORT=true # also export traces as OTel GenAI spans +``` + +### GenAI spans and metrics + +When DocsGPT runs under `opentelemetry-instrument`, each finished trace is +also exported as spans that follow the +[OpenTelemetry GenAI semantic conventions](https://opentelemetry.io/docs/specs/semconv/gen-ai/): +`invoke_agent {agent}`, `chat {model}`, `execute_tool {tool}`, +`embeddings {model}` and `retrieval`, with attributes such as +`gen_ai.provider.name`, `gen_ai.request.model`, +`gen_ai.usage.input_tokens`, `gen_ai.usage.output_tokens`, +`gen_ai.usage.cache_read.input_tokens`, `gen_ai.conversation.id`, +`gen_ai.agent.id` and `gen_ai.tool.name`. DocsGPT-specific details use the +`docsgpt.*` prefix (`docsgpt.request_id`, `docsgpt.token_source`, +`docsgpt.cache_hit`, `docsgpt.ttft_ms`, ...). The trace's root span is a +child of the request's HTTP server span, and the stored trace keeps the +OTel trace id so you can move between the two. + +Two metrics are recorded for every model call: +`gen_ai.client.token.usage` and `gen_ai.client.operation.duration`. + +Backends that understand the GenAI conventions — Langfuse +(`/api/public/otel`), Arize Phoenix, Datadog LLM Observability, Grafana — +render these as LLM traces with token and cost views. + +Prompt and tool content is **not** exported by default, because the OTLP +backend may be a third party. Opt in with the standard variable: + +```bash +OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT=SPAN_ONLY +``` + +This adds the same redacted previews the app stores (for example +`gen_ai.tool.call.arguments` and `gen_ai.tool.call.result`). + + + GenAI spans are exported when the request finishes, with their original + timestamps. Consequences: a long research run appears only when it ends; + HTTP and database spans made during a step sit beside the step rather than + under it; and log records carry the request's span ids, not the step's. + + +The GenAI conventions are still in development upstream, so attribute names +may change in later releases. + ## Backend examples ### Axiom diff --git a/docs/content/Deploying/Settings-Reference.mdx b/docs/content/Deploying/Settings-Reference.mdx index d9f59ec7..b08af68a 100644 --- a/docs/content/Deploying/Settings-Reference.mdx +++ b/docs/content/Deploying/Settings-Reference.mdx @@ -1451,6 +1451,47 @@ Type `int`, default `30`, must be `>= 1`. Days guardrail events are kept before the cleanup task removes them. +## Execution traces + +Recording of agent, LLM, tool and retrieval steps per request. + +### `TRACES_ENABLED` + +Type `bool`, default `true`. + +Record an execution trace (agent runs, LLM calls, tool calls, retrieval, embeddings) for every request, store it in request_traces and show it in the Logs UI. False records nothing. + +### `TRACES_CAPTURE_CONTENT` + +Type `bool`, default `true`. + +Store short, secret-redacted previews (tool arguments and results, retrieved chunk titles, rephrased queries, answer excerpts) with each stored trace. Full prompts are never stored. OTel export follows OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT instead. + +### `TRACES_PREVIEW_CHARS` + +Type `int`, default `2000`, must be `>= 100`. + +Maximum characters kept per stored trace preview. + +### `TRACES_MAX_SPANS` + +Type `int`, default `500`, must be `>= 10`. + +Maximum spans recorded per trace; further spans are counted as dropped, not stored. + +### `TRACES_RETENTION_DAYS` + +Type `int`, default `30`, must be `>= 1`. + +Days stored traces are kept before the cleanup task removes them. + +### `TRACES_OTEL_EXPORT` + +Type `bool`, default `true`. + +Also emit each finished trace as OpenTelemetry GenAI spans (gen_ai.*) and metrics. Has no effect unless an OTel SDK is configured, e.g. by launching under opentelemetry-instrument. + + ## Quotas Quota window and the treatment of unpriced models. diff --git a/docsgpt/agents/base.py b/docsgpt/agents/base.py index 030d76fc..4f620236 100644 --- a/docsgpt/agents/base.py +++ b/docsgpt/agents/base.py @@ -8,6 +8,7 @@ from typing import Any, Dict, Generator, List, Optional from docsgpt.agents.tool_executor import ( ToolExecutor, + trace_unexecuted_tool_call, result_status, truncate_tool_result, ) @@ -29,7 +30,7 @@ from docsgpt.guardrails.stream import StreamingOutputGuard from docsgpt.guardrails.types import Action, Stage, resolve_tool_result from docsgpt.llm.handlers.handler_creator import LLMHandlerCreator from docsgpt.llm.llm_creator import LLMCreator -from docsgpt.logging import build_stack_data, log_activity, LogContext +from docsgpt.logging import build_stack_data, log_activity, LogContext, start_agent_span logger = logging.getLogger(__name__) @@ -492,12 +493,28 @@ class BaseAgent(ABC): or client-side results), appends the resulting messages, then hands back to the LLM to continue the conversation. + Unlike :meth:`gen` this is not wrapped by ``@log_activity``, so the + continuation's ``invoke_agent`` trace span is opened here. + Args: messages: The saved messages array from the pause point. tools_dict: The saved tools dictionary. pending_tool_calls: The pending tool call descriptors from the pause. tool_actions: Client-provided actions resolving the pending calls. """ + with start_agent_span(self, continuation=True): + yield from self._gen_continuation_inner( + messages, tools_dict, pending_tool_calls, tool_actions, reasoning_content + ) + + def _gen_continuation_inner( + self, + messages: List[Dict], + tools_dict: Dict, + pending_tool_calls: List[Dict], + tool_actions: List[Dict], + reasoning_content: str = "", + ) -> Generator[Dict, None, None]: self._prepare_tools(tools_dict) actions_by_id = {a["call_id"]: a for a in tool_actions} @@ -583,16 +600,15 @@ class BaseAgent(ABC): messages.append( self.llm_handler.create_tool_message(tc, denial) ) - yield { - "type": "tool_call", - "data": { - "tool_name": pending.get("tool_name", "unknown"), - "call_id": call_id, - "action_name": pending.get("llm_name", pending["name"]), - "arguments": args, - "status": "denied", - }, + denied_data = { + "tool_name": pending.get("tool_name", "unknown"), + "call_id": call_id, + "action_name": pending.get("llm_name", pending["name"]), + "arguments": args, + "status": "denied", } + trace_unexecuted_tool_call(tc, {**denied_data, "error": comment or None}) + yield {"type": "tool_call", "data": denied_data} elif "result" in action: result = action["result"] @@ -616,17 +632,16 @@ class BaseAgent(ABC): tc, _bound_tool_response_for_llm(result_str) ) ) - yield { - "type": "tool_call", - "data": { - "tool_name": pending.get("tool_name", "unknown"), - "call_id": call_id, - "action_name": pending.get("llm_name", pending["name"]), - "arguments": args, - "result": truncate_tool_result(result_str), - "status": result_status(result), - }, + client_data = { + "tool_name": pending.get("tool_name", "unknown"), + "call_id": call_id, + "action_name": pending.get("llm_name", pending["name"]), + "arguments": args, + "result": truncate_tool_result(result_str), + "status": result_status(result), } + 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 llm_response = self._llm_gen(messages, preserve_responses_state=True) diff --git a/docsgpt/agents/headless_runner.py b/docsgpt/agents/headless_runner.py index c2a4d713..d97e6407 100644 --- a/docsgpt/agents/headless_runner.py +++ b/docsgpt/agents/headless_runner.py @@ -5,6 +5,7 @@ from __future__ import annotations import logging from typing import Any, Dict, Iterable, List, Optional +from docsgpt import tracing from docsgpt.agents.agent_creator import AgentCreator from docsgpt.agents.tool_executor import ToolExecutor from docsgpt.api.answer.services.prompt_renderer import ( @@ -69,12 +70,61 @@ def run_agent_headless( endpoint: str = "headless", chat_history: Optional[List[Dict[str, Any]]] = None, conversation_id: Optional[str] = None, + request_id: Optional[str] = None, + trace_user_id: Optional[str] = None, ) -> Dict[str, Any]: """Run an agent with no live client; returns a structured outcome dict. + The run is recorded as one execution trace under ``endpoint`` as its + source. ``request_id`` links that trace to the caller's own record (the + scheduler passes its run id, the webhook worker its task id); it is kept + off the LLM's token-usage rows, whose request ids drive request counts. + ``trace_user_id`` owns the trace when the run belongs to someone other + than the agent's owner (a schedule a user set on a shared agent), so the + trace is visible wherever that user sees the run; it defaults to the owner. + Raises: QuotaExceededError: If the agent owner's usage quota is exhausted. """ + trace = tracing.start_trace( + source=endpoint, + request_id=request_id, + user_id=trace_user_id or _resolve_owner(agent_config), + agent_id=_resolve_agent_id(agent_config), + conversation_id=conversation_id, + ) + status = None + with tracing.activate(trace): + try: + outcome = _run_agent_headless( + agent_config, + query, + tool_allowlist=tool_allowlist, + model_id_override=model_id_override, + endpoint=endpoint, + chat_history=chat_history, + conversation_id=conversation_id, + ) + if outcome.get("error"): + status = tracing.STATUS_ERROR + return outcome + except BaseException: + status = tracing.STATUS_ERROR + raise + finally: + tracing.flush(trace, status) + + +def _run_agent_headless( + agent_config: Dict[str, Any], + query: str, + *, + tool_allowlist: Optional[Iterable[str]] = None, + model_id_override: Optional[str] = None, + endpoint: str = "headless", + chat_history: Optional[List[Dict[str, Any]]] = None, + conversation_id: Optional[str] = None, +) -> Dict[str, Any]: from docsgpt.core.model_utils import ( get_api_key_for_provider, get_default_model_id, diff --git a/docsgpt/agents/research_agent.py b/docsgpt/agents/research_agent.py index de96cf35..859bc872 100644 --- a/docsgpt/agents/research_agent.py +++ b/docsgpt/agents/research_agent.py @@ -2,8 +2,9 @@ import json import logging import os import time -from typing import Dict, Generator, List, Optional +from typing import Any, Dict, Generator, List, Optional +from docsgpt import tracing from docsgpt.agents.base import BaseAgent from docsgpt.agents.tool_executor import ToolExecutor from docsgpt.agents.tools.graph_search import add_graph_search_tool @@ -46,6 +47,15 @@ STEP_PROMPT = _load_prompt("step.txt") SYNTHESIS_PROMPT = _load_prompt("synthesis.txt") +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( + tracing.KIND_STEP, + f"research {phase}", + attributes={"docsgpt.research.phase": phase.split(" ")[0], **attributes}, + ) + + # --------------------------------------------------------------------------- # CitationManager # --------------------------------------------------------------------------- @@ -169,7 +179,8 @@ class ResearchAgent(BaseAgent): # Phase 0: Clarification (skip if user is responding to a prior clarification) if not self._is_follow_up(): - clarification = self._clarification_phase(query) + with _phase_span("clarify"): + clarification = self._clarification_phase(query) if clarification: yield {"metadata": {"is_clarification": True}} yield {"answer": clarification} @@ -182,7 +193,11 @@ class ResearchAgent(BaseAgent): # Phase 1: Planning (with adaptive depth) yield {"type": "research_progress", "data": {"status": "planning"}} - plan, complexity = self._planning_phase(query) + with _phase_span("plan") as plan_span: + plan, complexity = self._planning_phase(query) + plan_span.set( + **{"docsgpt.research.steps": len(plan or []), "docsgpt.research.complexity": complexity} + ) if not plan: logger.warning("ResearchAgent: Planning produced no steps, falling back") @@ -222,7 +237,9 @@ class ResearchAgent(BaseAgent): }, } - report = self._research_step(step_query, tools_dict) + with _phase_span(f"step {step_num}", **{"docsgpt.research.step": step_num}) as step_span: + step_span.preview("query", step_query) + report = self._research_step(step_query, tools_dict) intermediate_reports.append({"step": step, "content": report}) yield { @@ -249,9 +266,10 @@ class ResearchAgent(BaseAgent): "tokens_used": self._tokens_used, }, } - yield from self._synthesis_phase( - query, plan, intermediate_reports, tools_dict, log_context - ) + with _phase_span("synthesis"): + yield from self._synthesis_phase( + query, plan, intermediate_reports, tools_dict, log_context + ) # 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 9952e3d2..bc7f809a 100644 --- a/docsgpt/agents/tool_executor.py +++ b/docsgpt/agents/tool_executor.py @@ -13,6 +13,7 @@ from docsgpt.agents.default_tools import ( resolve_tool_by_id, synthesized_default_tools, ) +from docsgpt import tracing from docsgpt.agents.tools.tool_action_parser import ToolActionParser from docsgpt.agents.tools.tool_manager import ToolManager from docsgpt.guardrails.types import Stage as GuardrailStage, resolve_tool_result @@ -29,6 +30,85 @@ from docsgpt.storage.db.session import db_readonly, db_session logger = logging.getLogger(__name__) +def record_tool_span_start(call: Any, **attributes: Any) -> Any: + """Open an ``execute_tool`` span for ``call`` (no-op without an active trace).""" + llm_name = getattr(call, "name", None) or "unknown" + base = { + "gen_ai.operation.name": "execute_tool", + "gen_ai.tool.name": llm_name, + "gen_ai.tool.call.id": getattr(call, "id", None), + "gen_ai.tool.type": "function", + } + base.update(attributes) + return tracing.start_span( + tracing.KIND_TOOL, + f"execute_tool {llm_name}", + attributes={k: v for k, v in base.items() if v is not None}, + ) + + +#: ``tool_calls`` statuses mapped to trace span statuses. +_TOOL_SPAN_STATUS = { + "error": tracing.STATUS_ERROR, + "denied": tracing.STATUS_DENIED, + "skipped": tracing.STATUS_SKIPPED, + "awaiting_approval": tracing.STATUS_PENDING, + "requires_client_execution": tracing.STATUS_PENDING, +} + + +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. + + Args: + span: The span from :func:`record_tool_span_start`. + data: The tool-call record (``tool_name``, ``arguments``, ``result``, + ``status``, ...); results in it are already guardrail-scanned. + """ + if not span: + return + status = str(data.get("status") or "completed") + span.set( + **{ + "docsgpt.tool": data.get("tool_name"), + "docsgpt.action": data.get("action_name"), + "docsgpt.tool_status": status, + "docsgpt.artifact_id": data.get("artifact_id"), + } + ) + if data.get("arguments") is not None: + span.preview("arguments", data.get("arguments")) + if data.get("result") is not None: + span.preview("result", data.get("result")) + span_status = _TOOL_SPAN_STATUS.get(status, tracing.STATUS_OK) + # ``error`` and ``result`` are tool output or user text (a denial + # comment), so they travel only as previews, which honour the content + # capture settings; ``span.error`` is exported and stored unconditionally. + if data.get("error"): + span.preview("error", data["error"]) + if span_status == tracing.STATUS_ERROR: + error_type = str(data.get("error_type") or "ToolError") + span.error = f"Tool call failed ({error_type})" + span.set(**{"error.type": error_type}) + span.end(span_status) + + def _is_foreign_key_violation(exc: BaseException) -> bool: """Whether ``exc`` is a Postgres FK violation (SQLSTATE 23503).""" if not isinstance(exc, IntegrityError): @@ -950,7 +1030,31 @@ class ToolExecutor: MAX_ADVERTISED_TOOL_NAMES = 30 def execute(self, tools_dict: Dict, call, llm_class_name: str): - """Execute a tool call. Yields status events, returns (result, call_id).""" + """Execute a tool call. Yields status events, returns (result, call_id). + + Every call is recorded as an ``execute_tool`` trace span; its outcome + is read from the ``tool_calls`` entry each branch of :meth:`_execute` + appends, so the span matches what the conversation stores. + """ + span = record_tool_span_start(call) + recorded = len(self.tool_calls) + try: + outcome = yield from self._execute(tools_dict, call, llm_class_name) + except Exception as exc: + # A tool's exception text can quote its response; like tool + # output it goes only into the capture-gated preview. + span.preview("error", str(exc)) + span.error = f"Tool call failed ({type(exc).__name__})" + span.end(tracing.STATUS_ERROR, attributes={"error.type": type(exc).__name__}) + raise + except GeneratorExit: + span.end(tracing.STATUS_CANCELLED) + raise + data = self.tool_calls[-1] if len(self.tool_calls) > recorded else {} + finish_tool_span(span, data) + return outcome + + def _execute(self, tools_dict: Dict, call, llm_class_name: str): parser = ToolActionParser(llm_class_name, name_mapping=self._name_to_tool) tool_id, action_name, call_args = parser.parse_args(call) llm_name = getattr(call, "name", "unknown") diff --git a/docsgpt/agents/workflows/workflow_engine.py b/docsgpt/agents/workflows/workflow_engine.py index 15876c82..3b80dc81 100644 --- a/docsgpt/agents/workflows/workflow_engine.py +++ b/docsgpt/agents/workflows/workflow_engine.py @@ -6,6 +6,7 @@ import uuid from datetime import datetime, timezone from typing import Any, Dict, Generator, List, Optional, TYPE_CHECKING +from docsgpt import tracing from docsgpt.agents.workflows.cel_evaluator import CelEvaluationError, evaluate_cel from docsgpt.agents.workflows.node_agent import WorkflowNodeAgentFactory from docsgpt.agents.workflows.schemas import ( @@ -114,6 +115,9 @@ class WorkflowEngine: self, initial_inputs: WorkflowState, query: str ) -> Generator[Dict[str, str], None, None]: self._initialize_state(initial_inputs, query) + # A workflow run's Logs row finds its trace by this id; a nested + # engine (a workflow inside a workflow) keeps the outermost run's. + tracing.bind_if_unset(workflow_run_id=self.workflow_run_id) # Surface the run id up front so the client can list this run's # artifacts (GET /api/artifacts?workflow_run_id=) once it has been @@ -140,6 +144,16 @@ class WorkflowEngine: break log_entry = self._create_log_entry(node) self._last_node_tool_calls = [] + step_span = tracing.start_span( + tracing.KIND_STEP, + f"workflow_step {node.title or node.type.value}", + attributes={ + "docsgpt.workflow.node_id": node.id, + "docsgpt.workflow.node_type": node.type.value, + "docsgpt.workflow.node_title": node.title or None, + "docsgpt.workflow_run_id": self.workflow_run_id, + }, + ) yield { "type": "workflow_step", @@ -150,7 +164,8 @@ class WorkflowEngine: } try: - yield from self._execute_node(node) + with step_span: + yield from self._execute_node(node) log_entry["status"] = ExecutionStatus.COMPLETED.value self._finalize_log_entry(log_entry, pre_state) diff --git a/docsgpt/alembic/versions/0037_request_traces.py b/docsgpt/alembic/versions/0037_request_traces.py new file mode 100644 index 00000000..0f57f343 --- /dev/null +++ b/docsgpt/alembic/versions/0037_request_traces.py @@ -0,0 +1,91 @@ +"""0037 request_traces — one execution trace per request for the Logs UI. + +Each row is the full span tree of one execution (a chat turn, a tool-approval +continuation, a scheduled or webhook run, a search, a graph extraction): +agent runs, LLM calls, tool calls, retrieval and embeddings with their +timings, stored as a ``spans`` JSONB array. The Logs UI loads a trace whole, +so one row per trace keeps the write to a single INSERT and lets retention +and conversation deletion remove a trace in one step. + +``message_id`` cascades: truncating a conversation when a turn is +superseded removes that turn's traces with it. Deleting a conversation also +deletes, by ``conversation_id``, the traces that have no message (scheduled +runs, stateless ``/v1`` rounds, turns whose message was never reserved); the +conversation index serves that. The link ids (``request_id``, +``activity_id``, ``workflow_run_id``) carry partial indexes because only the +Logs rows that have them look traces up by them. Traces with no log row of +their own (searches, graph builds) are listed per user or agent and source, +hence the source-leading indexes. + +``status`` is the only CHECK; span kinds and sources are validated in code so +new ones need no migration. + +Idempotent both ways. + +Revision ID: 0037_request_traces +Revises: 0036_device_audit_created_idx +""" + +from typing import Sequence, Union + +from alembic import op + + +revision: str = "0037_request_traces" +down_revision: Union[str, None] = "0036_device_audit_created_idx" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.execute( + """ + CREATE TABLE IF NOT EXISTS request_traces ( + id UUID PRIMARY KEY, + request_id TEXT, + message_id UUID REFERENCES conversation_messages(id) ON DELETE CASCADE, + conversation_id UUID, + activity_id TEXT, + workflow_run_id UUID, + user_id TEXT, + agent_id UUID, + source TEXT NOT NULL, + name TEXT, + status TEXT NOT NULL + CONSTRAINT request_traces_status_chk + CHECK (status IN ('ok', 'error', 'paused', 'cancelled')), + started_at TIMESTAMPTZ NOT NULL, + duration_ms INTEGER NOT NULL DEFAULT 0, + span_count INTEGER NOT NULL DEFAULT 0, + dropped_spans INTEGER NOT NULL DEFAULT 0, + summary JSONB NOT NULL DEFAULT '{}'::jsonb, + spans JSONB NOT NULL DEFAULT '[]'::jsonb, + otel_trace_id TEXT, + created_at TIMESTAMPTZ NOT NULL DEFAULT now() + ); + """ + ) + op.execute( + """ + CREATE INDEX IF NOT EXISTS request_traces_request_idx + ON request_traces (request_id) WHERE request_id IS NOT NULL; + CREATE INDEX IF NOT EXISTS request_traces_message_idx + ON request_traces (message_id) WHERE message_id IS NOT NULL; + CREATE INDEX IF NOT EXISTS request_traces_activity_idx + ON request_traces (activity_id) WHERE activity_id IS NOT NULL; + CREATE INDEX IF NOT EXISTS request_traces_workflow_run_idx + ON request_traces (workflow_run_id) WHERE workflow_run_id IS NOT NULL; + CREATE INDEX IF NOT EXISTS request_traces_conversation_idx + ON request_traces (conversation_id) WHERE conversation_id IS NOT NULL; + CREATE INDEX IF NOT EXISTS request_traces_user_source_started_idx + ON request_traces (user_id, source, started_at DESC); + CREATE INDEX IF NOT EXISTS request_traces_agent_source_started_idx + ON request_traces (agent_id, source, started_at DESC) WHERE agent_id IS NOT NULL; + CREATE INDEX IF NOT EXISTS request_traces_created_idx + ON request_traces (created_at); + """ + ) + + +def downgrade() -> None: + op.execute("DROP TABLE IF EXISTS request_traces;") diff --git a/docsgpt/api/answer/routes/answer.py b/docsgpt/api/answer/routes/answer.py index 7cdcd722..4d7f8dd2 100644 --- a/docsgpt/api/answer/routes/answer.py +++ b/docsgpt/api/answer/routes/answer.py @@ -13,7 +13,10 @@ from docsgpt.api.answer.services.continuation_service import ( ResumeInProgressError, ) from docsgpt.api.answer.services.persistence_policy import resolve_persistence -from docsgpt.api.answer.services.stream_processor import StreamProcessor +from docsgpt.api.answer.services.stream_processor import ( + StreamProcessor, + flush_trace_after_request, +) logger = logging.getLogger(__name__) @@ -87,7 +90,10 @@ class AnswerResource(Resource, BaseAnswerResource): if error := self.validate_request(data): return error decoded_token = getattr(request, "decoded_token", None) - processor = StreamProcessor(data, decoded_token) + processor = StreamProcessor(data, decoded_token, trace_source="answer") + # ``complete_stream`` is consumed below and writes the trace itself; + # this covers requests refused before it runs. + flush_trace_after_request(processor) try: # ---- Continuation mode ---- if data.get("tool_actions"): @@ -115,6 +121,7 @@ class AnswerResource(Resource, BaseAnswerResource): decoded_token=processor.decoded_token, agent_id=processor.agent_id, model_id=processor.model_id, + trace=processor.trace, _continuation={ "messages": messages, "tools_dict": tools_dict, @@ -156,6 +163,8 @@ class AnswerResource(Resource, BaseAnswerResource): is_shared_usage=processor.is_shared_usage, shared_token=processor.shared_token, model_id=processor.model_id, + request_id=processor.request_id, + trace=processor.trace, ) stream_result = self.process_response_stream(stream) diff --git a/docsgpt/api/answer/routes/base.py b/docsgpt/api/answer/routes/base.py index e0cfe4cd..ca92d3cf 100644 --- a/docsgpt/api/answer/routes/base.py +++ b/docsgpt/api/answer/routes/base.py @@ -1,14 +1,17 @@ import datetime +import functools +import inspect import json import logging import threading import time import uuid -from typing import Any, Dict, Generator, List, Optional +from typing import Any, Callable, Dict, Generator, List, Optional from flask import jsonify, make_response, Response from flask_restx import Namespace +from docsgpt import tracing from docsgpt.api.answer.services.continuation_service import ContinuationService from docsgpt.api.answer.services.conversation_service import ( ConversationService, @@ -66,6 +69,54 @@ class StreamSuperseded(Exception): answer_ns = Namespace("answer", description="Answer related operations", path="/") +def _traced_stream( + method: Callable[..., Generator[str, None, None]], +) -> Callable[..., Generator[str, None, None]]: + """Record ``complete_stream`` as the request's execution trace. + + The stream runs in the SSE pump thread, not the request thread, so the + trace the route started is activated here for the whole stream and written + exactly once when the stream ends -- normally, paused, failed or + abandoned. The request id is resolved here too so the trace and the + stream agree on it. + """ + signature = inspect.signature(method) + + @functools.wraps(method) + def wrapper(*args: Any, **kwargs: Any) -> Generator[str, None, None]: + bound = signature.bind(*args, **kwargs) + arguments = bound.arguments + continuation = arguments.get("_continuation") + request_id = ( + continuation.get("request_id") if continuation else None + ) or arguments.get("request_id") or str(uuid.uuid4()) + arguments["request_id"] = request_id + trace = arguments.get("trace") + if not isinstance(trace, tracing.Trace): + agent = arguments.get("agent") + trace = tracing.start_trace( + source=str(getattr(agent, "endpoint", None) or "stream"), + capture_otel_context=False, + ) + arguments["trace"] = trace + if trace is not None: + decoded_token = arguments.get("decoded_token") + trace.bind( + request_id=request_id, + user_id=decoded_token.get("sub") if isinstance(decoded_token, dict) else None, + agent_id=arguments.get("agent_id"), + ) + with tracing.activate(trace): + try: + yield from method(*bound.args, **bound.kwargs) + finally: + # Written on a writer thread so the stream's connection closes + # without waiting on the OTel replay and the INSERT. + tracing.flush(trace, background=True) + + return wrapper + + class BaseAnswerResource: """Shared base class for answer endpoints""" @@ -93,6 +144,79 @@ class BaseAnswerResource: return missing_fields return None + def _persist_turn_log( + self, + *, + decoded_token: Dict[str, Any], + user_api_key: Optional[str], + agent_id: Optional[str], + question: str, + response: str, + sources: List[Dict[str, Any]], + tool_calls: Any, + attachment_ids: Optional[List[str]], + request_id: Optional[str], + message_id: Optional[str], + error: Optional[str] = None, + extra: Optional[Dict[str, Any]] = None, + ) -> None: + """Write the turn's ``user_logs`` row: its Logs entry, linked to its trace. + + Written for every finished turn, failed ones included (``level`` is + ``error`` then), so a failed chat shows up as a chat entry. A failure + to write is logged, never raised. + + Args: + decoded_token: The caller's token. + user_api_key: The agent API key the request used, if any. + agent_id: The agent that answered. + question: The question as stored (after input redaction). + response: The answer, or what streamed before a failure. + sources: Retrieved sources. + tool_calls: The turn's tool calls, before log truncation. + attachment_ids: Attached file ids. + request_id: The turn's request id (links its trace). + message_id: The reserved message id. + error: What failed the turn, if it failed. + extra: More fields for the row (structured-output details). + """ + log_data: Dict[str, Any] = { + "action": "stream_answer", + "level": "error" if error else "info", + "user": decoded_token.get("sub"), + "api_key": user_api_key, + "agent_id": agent_id, + "question": question, + "response": response, + "sources": sources, + "tool_calls": self._prepare_tool_calls_for_logging(tool_calls), + "attachments": attachment_ids, + "timestamp": datetime.datetime.now(datetime.timezone.utc), + # Links the Logs row to this turn's execution trace. + "request_id": request_id, + "message_id": message_id, + } + if error: + log_data["error"] = error + if extra: + log_data.update(extra) + # Clean up text fields to be no longer than 10000 characters. + for key, value in log_data.items(): + if isinstance(value, str) and len(value) > 10000: + log_data[key] = value[:10000] + try: + with db_session() as conn: + UserLogsRepository(conn).insert( + user_id=log_data.get("user"), + endpoint="stream_answer", + data=log_data, + ) + except Exception as log_err: + logger.error( + f"Failed to persist stream_answer user log: {log_err}", + exc_info=True, + ) + @staticmethod def _prepare_tool_calls_for_logging( tool_calls: Optional[List[Dict[str, Any]]], max_chars: int = 10000 @@ -244,6 +368,7 @@ class BaseAnswerResource: logger.exception("Failed to release resume claim after a usage refusal") return error + @_traced_stream def complete_stream( self, question: str, @@ -263,6 +388,8 @@ class BaseAnswerResource: model_user_id: Optional[str] = None, _continuation: Optional[Dict] = None, finalize_tool_pause_as_complete: bool = False, + request_id: Optional[str] = None, + trace: Optional["tracing.Trace"] = None, ) -> Generator[str, None, None]: """ Generator function that streams the complete conversation response. @@ -286,6 +413,14 @@ class BaseAnswerResource: shared_token: Token for shared agent model_id: Model ID used for the request retrieved_docs: Pre-fetched documents for sources (optional) + request_id: The request's id, minted by the route before the + agent was built so pre-fetch retrieval shares it. A + continuation's saved id takes precedence; absent both, a new + one is minted. + trace: The execution trace the route started (it already holds + the pre-fetch retrieval); :func:`_traced_stream` activates it + and writes it when the stream ends. A new one is started when + omitted. finalize_tool_pause_as_complete: Stateless-tool-round mode for the OpenAI-compatible ``/v1/chat/completions`` endpoint. OpenAI clients resume a tool call by re-POSTing the full @@ -336,10 +471,10 @@ class BaseAnswerResource: pending_pause_event: Optional[dict] = None # One id shared across the WAL row, primary LLM (token_usage - # attribution), the SSE event, and resumed continuations. + # attribution), the SSE event, resumed continuations and the trace. request_id = ( _continuation.get("request_id") if _continuation else None - ) or str(uuid.uuid4()) + ) or request_id or str(uuid.uuid4()) # Reserve the placeholder row before the LLM call so a crash # mid-stream still leaves the question queryable. Continuations @@ -405,6 +540,17 @@ class BaseAnswerResource: primary_llm = getattr(agent, "llm", None) if primary_llm is not None: primary_llm._request_id = request_id + # Side-channel LLMs built later (guardrail judge, retrievers the agent + # creates for its search tools) read the id off the agent. + if getattr(agent, "request_id", None) is None: + try: + agent.request_id = request_id + except Exception: + logger.debug("Could not stamp request_id on the agent") + tracing.bind( + message_id=reserved_message_id, + conversation_id=str(conversation_id) if conversation_id else None, + ) # Flipped to ``streaming`` on the first ``answer``/``sources`` chunk; # the reconciler reads ``status`` to tell "never started" from "in @@ -791,10 +937,20 @@ class BaseAnswerResource: # error silently — the exact shape of the bug being fixed here. if stream_error: query_metadata.setdefault("error", stream_error) + # A yielded error (e.g. a failed workflow node) ends the + # generator normally, so no span raised; the user still saw + # the turn fail, and its trace should say so. A pause below + # overrides this, as it does for the message row. + trace = tracing.current_trace() + if trace is not None: + trace.outcome = tracing.STATUS_ERROR # ---- Paused: save continuation state and end stream early ---- if paused: continuation = getattr(agent, "_pending_continuation", None) + trace = tracing.current_trace() + if trace is not None: + trace.outcome = tracing.STATUS_PAUSED # ---- Stateless-tool-round mode (OpenAI-compatible /v1) ---- # OpenAI clients resume by re-POSTing the whole message @@ -1216,44 +1372,27 @@ class BaseAnswerResource: ) yield _emit({"type": "id", "id": str(conversation_id)}) - tool_calls_for_logging = self._prepare_tool_calls_for_logging( - getattr(agent, "tool_calls", tool_calls) or tool_calls - ) - - log_data = { - "action": "stream_answer", - "level": "info", - "user": decoded_token.get("sub"), - "api_key": user_api_key, - "agent_id": agent_id, - "question": question, - "response": response_full, - "sources": source_log_docs, - "tool_calls": tool_calls_for_logging, - "attachments": attachment_ids, - "timestamp": datetime.datetime.now(datetime.timezone.utc), - } + extra: Dict[str, Any] = {} if is_structured: - log_data["structured_output"] = True + extra["structured_output"] = True if schema_info: - log_data["schema"] = schema_info - # Clean up text fields to be no longer than 10000 characters - - for key, value in log_data.items(): - if isinstance(value, str) and len(value) > 10000: - log_data[key] = value[:10000] - try: - with db_session() as conn: - UserLogsRepository(conn).insert( - user_id=log_data.get("user"), - endpoint="stream_answer", - data=log_data, - ) - except Exception as log_err: - logger.error( - f"Failed to persist stream_answer user log: {log_err}", - exc_info=True, - ) + extra["schema"] = schema_info + self._persist_turn_log( + decoded_token=decoded_token, + user_api_key=user_api_key, + agent_id=agent_id, + question=question, + response=response_full, + sources=source_log_docs, + tool_calls=getattr(agent, "tool_calls", tool_calls) or tool_calls, + attachment_ids=attachment_ids, + request_id=request_id, + message_id=reserved_message_id, + # A yielded error (a failed workflow node) ends the turn + # normally but still failed it. + error=stream_error, + extra=extra, + ) yield _emit({"type": "end"}) # Drain the journal buffer so the terminal ``end`` event is @@ -1469,9 +1608,14 @@ class BaseAnswerResource: ) if journal_writer is not None: journal_writer.close() + # The user replaced this turn; its trace describes nothing kept. + tracing.discard(tracing.current_trace()) return except Exception as e: logger.error(f"Error in stream: {str(e)}", exc_info=True) + trace = tracing.current_trace() + if trace is not None: + trace.outcome = tracing.STATUS_ERROR # This process took the resume claim, so it owns releasing it. The # only other way back is ``revert_stale_resuming``'s 600 s grace, # which leaves the user locked out of their own conversation for @@ -1526,6 +1670,22 @@ class BaseAnswerResource: f"Failed to finalize errored message: {fin_err}", exc_info=True, ) + # A failed turn is still a chat turn: log it as one (level + # ``error``), with its trace link, instead of leaving only the + # agent's system error row. + self._persist_turn_log( + decoded_token=decoded_token, + user_api_key=user_api_key, + agent_id=agent_id, + question=question, + response=response_full, + sources=source_log_docs, + tool_calls=getattr(agent, "tool_calls", tool_calls) or tool_calls, + attachment_ids=attachment_ids, + request_id=request_id, + message_id=reserved_message_id, + error=f"{type(e).__name__}: {e}", + ) yield _emit( { "type": "error", diff --git a/docsgpt/api/answer/routes/stream.py b/docsgpt/api/answer/routes/stream.py index cdee274b..21dbcbff 100644 --- a/docsgpt/api/answer/routes/stream.py +++ b/docsgpt/api/answer/routes/stream.py @@ -13,7 +13,10 @@ from docsgpt.api.answer.services.continuation_service import ( ResumeInProgressError, ) from docsgpt.api.answer.services.persistence_policy import resolve_persistence -from docsgpt.api.answer.services.stream_processor import StreamProcessor +from docsgpt.api.answer.services.stream_processor import ( + StreamProcessor, + flush_trace_after_request, +) from docsgpt.streaming.sse_keepalive import with_sse_keepalive logger = logging.getLogger(__name__) @@ -94,7 +97,8 @@ class StreamResource(Resource, BaseAnswerResource): if error := self.validate_request(data, "index" in data): return error decoded_token = getattr(request, "decoded_token", None) - processor = StreamProcessor(data, decoded_token) + processor = StreamProcessor(data, decoded_token, trace_source="stream") + flush_trace_after_request(processor) try: # ---- Continuation mode ---- @@ -130,6 +134,7 @@ class StreamResource(Resource, BaseAnswerResource): agent_id=processor.agent_id, model_id=processor.model_id, model_user_id=processor.model_user_id, + trace=processor.handoff_trace(), _continuation={ "messages": messages, "tools_dict": tools_dict, @@ -181,6 +186,8 @@ class StreamResource(Resource, BaseAnswerResource): shared_token=processor.shared_token, model_id=processor.model_id, model_user_id=processor.model_user_id, + request_id=processor.request_id, + trace=processor.handoff_trace(), ), ), mimetype="text/event-stream", diff --git a/docsgpt/api/answer/services/stream_processor.py b/docsgpt/api/answer/services/stream_processor.py index 004b616c..27be1028 100644 --- a/docsgpt/api/answer/services/stream_processor.py +++ b/docsgpt/api/answer/services/stream_processor.py @@ -1,9 +1,14 @@ import datetime +import functools import json import logging +import uuid from pathlib import Path -from typing import Any, Dict, Optional, Set +from typing import Any, Callable, Dict, Optional, Set, TypeVar +from flask import after_this_request + +from docsgpt import tracing from docsgpt.agents.agent_creator import AgentCreator from docsgpt.agents.default_tools import synthesized_default_tools from docsgpt.api.answer.services.compression import CompressionOrchestrator @@ -104,9 +109,63 @@ def get_prompt(prompt_id: str, prompts_collection=None) -> str: raise ValueError(f"Invalid prompt ID: {prompt_id}") from e +T = TypeVar("T") + + +def _traced_setup(method: Callable[..., T]) -> Callable[..., T]: + """Run a request-setup method inside the request's execution trace. + + Agent setup does real work worth seeing in the trace -- pre-fetch + retrieval, history compression -- before ``complete_stream`` runs, so + the trace is started here, in the request thread, and handed on. + """ + + @functools.wraps(method) + def wrapper(self: "StreamProcessor", *args: Any, **kwargs: Any) -> T: + trace = getattr(self, "trace", None) + if trace is None: + trace = tracing.start_trace(source=getattr(self, "trace_source", "stream")) + self.trace = trace + with tracing.activate(trace): + try: + return method(self, *args, **kwargs) + finally: + if trace is not None: + decoded = getattr(self, "decoded_token", None) + trace.bind( + request_id=getattr(self, "request_id", None), + user_id=decoded.get("sub") if isinstance(decoded, dict) else None, + agent_id=getattr(self, "agent_id", None), + ) + + return wrapper + + +def flush_trace_after_request(processor: "StreamProcessor") -> None: + """Write ``processor``'s setup trace when the request ends, unless it was claimed. + + Registered on the current request with ``after_this_request``; the hook + always hands the response back unchanged and never raises. + + Args: + processor: The request's processor. + """ + + @after_this_request + def _flush(response: Any) -> Any: + try: + processor.flush_unclaimed_trace() + except Exception: + logger.warning("Could not write an unclaimed request trace", exc_info=True) + return response + + class StreamProcessor: def __init__( - self, request_data: Dict[str, Any], decoded_token: Optional[Dict[str, Any]] + self, + request_data: Dict[str, Any], + decoded_token: Optional[Dict[str, Any]], + trace_source: str = "stream", ): # Legacy attribute retained as None for any external callers that # introspect the processor; all DB access uses per-op connections. @@ -137,6 +196,11 @@ class StreamProcessor: self.reserved_message_id: Optional[str] = None # Carried through resumes so multi-pause runs keep one request_id. self.request_id: Optional[str] = None + # The request's execution trace, started by the first traced setup + # step and handed to ``complete_stream``; ``trace_source`` names the + # entry point it is stored under. + self.trace: Optional[tracing.Trace] = None + self.trace_source = trace_source self.conversation_service = ConversationService() self.compression_orchestrator = CompressionOrchestrator( self.conversation_service @@ -161,12 +225,44 @@ class StreamProcessor: self._load_conversation_history() self._process_attachments() + def handoff_trace(self) -> Optional[tracing.Trace]: + """Hand the setup trace to a streaming ``complete_stream``. + + The stream writes the trace when it ends, after the view has + returned; marking the hand-off stops :meth:`flush_unclaimed_trace` + from writing it first. + + Returns: + The trace to pass as ``complete_stream(trace=...)``. + """ + self._trace_handed_off = True + return getattr(self, "trace", None) + + def flush_unclaimed_trace(self) -> None: + """Write the setup trace of a request that ended before streaming. + + A request refused after setup started (unauthorized, over its usage + limit, a resume conflict, a setup error) still records what ran, + marked ``error``. A trace the request already wrote, or handed to a + stream, is left alone. Routes arrange this with + :func:`flush_trace_after_request`. + """ + if not getattr(self, "_trace_handed_off", False): + tracing.flush(getattr(self, "trace", None), tracing.STATUS_ERROR) + + @_traced_setup def build_agent(self, question: str): """One call to go from request data to a ready-to-run agent. Combines initialize(), pre_fetch_docs(), pre_fetch_tools(), and - create_agent() into a single convenience method. + create_agent() into a single convenience method. The request id is + minted first so pre-fetch retrieval and its side-channel LLM calls + share it with the rest of the turn. It is always generated here, never + taken from the request body: request quotas count distinct request + ids, so a client-chosen id would let every call count as one. """ + if not getattr(self, "request_id", None): + self.request_id = str(uuid.uuid4()) self.initialize() agent_type = self.agent_config.get("agent_type", "classic") @@ -221,6 +317,7 @@ class StreamProcessor: tools_data=tools_data, ) + @_traced_setup def build_continuation_from_messages(self, messages, tool_actions): """Rebuild a tool continuation from the request messages (STATELESS). @@ -1118,7 +1215,7 @@ class StreamProcessor: user_api_key=self.agent_config["user_api_key"], agent_id=self.agent_id, decoded_token=self.decoded_token, - request_id=self.data.get("request_id"), + request_id=self.request_id or self.data.get("request_id"), ) def _legacy_classic(): @@ -1440,6 +1537,7 @@ class StreamProcessor: logger.warning(f"Failed to fetch memory tool data: {str(e)}") return None + @_traced_setup def resume_from_tool_actions( self, tool_actions: list, @@ -1813,7 +1911,7 @@ class StreamProcessor: "llm_name": provider or settings.LLM_PROVIDER, "api_key": system_api_key, "decoded_token": self.decoded_token, - "request_id": self.data.get("request_id"), + "request_id": self.request_id or self.data.get("request_id"), } elif agent_type == "workflow": diff --git a/docsgpt/api/pat/rules.py b/docsgpt/api/pat/rules.py index 227fe11a..c1d7683d 100644 --- a/docsgpt/api/pat/rules.py +++ b/docsgpt/api/pat/rules.py @@ -298,6 +298,7 @@ RULES: dict[tuple[str, str], Rule] = { ("/api/get_tool_analytics", "POST"): _rule("analytics:read", blocked_by=_ALL_FAMILIES), ("/api/get_schedule_analytics", "POST"): _rule("analytics:read", blocked_by=_ALL_FAMILIES), ("/api/get_user_logs", "POST"): _rule("analytics:read", blocked_by=_ALL_FAMILIES), + ("/api/traces", "GET"): _rule("analytics:read", blocked_by=_ALL_FAMILIES), # Teams (read only) ("/api/teams", "GET"): _rule("teams:read"), ("/api/teams/", "GET"): _rule("teams:read"), diff --git a/docsgpt/api/user/analytics/routes.py b/docsgpt/api/user/analytics/routes.py index 129f7e97..3642ed3a 100644 --- a/docsgpt/api/user/analytics/routes.py +++ b/docsgpt/api/user/analytics/routes.py @@ -1,10 +1,11 @@ """Analytics and reporting routes.""" import datetime +from typing import Optional, Tuple from flask import current_app, jsonify, make_response, request from flask_restx import fields, Namespace, Resource -from sqlalchemy import text as _sql_text +from sqlalchemy import Connection, text as _sql_text from docsgpt.api import api from docsgpt.api.user.base import ( @@ -14,6 +15,10 @@ from docsgpt.api.user.base import ( ) from docsgpt.storage.db.redaction import redact_secrets from docsgpt.storage.db.repositories.agents import AgentsRepository +from docsgpt.storage.db.repositories.request_traces import ( + REF_FIELDS as TRACE_REF_FIELDS, + RequestTracesRepository, +) from docsgpt.storage.db.repositories.token_usage import TokenUsageRepository from docsgpt.storage.db.session import db_readonly @@ -90,6 +95,157 @@ def _resolve_agent(conn, api_key_id, user_id): return agent, api_key, agent_pg_id +def _trace_branch(name: str, sources_sql: str, scope: str) -> dict: + """A ``get_user_logs`` branch listing stored traces of the given sources.""" + return { + "name": name, + "level": "CASE WHEN t.status = 'error' THEN 'error' ELSE 'info' END", + # A search is listed by its query (copied into the small ``summary`` + # at flush so this never detoasts ``spans``); graph builds are named + # after their source. + "summary": "COALESCE(t.summary->>'query', t.name, t.source)", + "where": [f"t.source IN {sources_sql}", scope], + "sql": f""" + SELECT '{name}' AS event_type, + CAST(t.id AS text) AS id, + t.user_id AS user_id, + t.started_at AS timestamp, + {{level}} AS level, + t.source AS action, + {{summary}} AS summary, + jsonb_build_object( + 'status', t.status, + 'source', t.source, + 'duration_ms', t.duration_ms + ) AS payload + FROM request_traces t + WHERE {{where}} + """, + } + + +def _trace_ref(item: dict) -> Optional[Tuple[str, str]]: + """The ``(field, value)`` that finds a Logs row's trace, if it has one.""" + event_type = item["event_type"] + row_id = item["id"].split("-", 1)[1] + if event_type == "chat": + return ("request_id", item.get("request_id")) if item.get("request_id") else None + if event_type in ("system", "webhook"): + return ("activity_id", item.get("activity_id")) if item.get("activity_id") else None + if event_type == "workflow": + return ("workflow_run_id", row_id) + if event_type == "schedule": + # The scheduler records the run under its run id. + return ("request_id", row_id) + if event_type in ("search", "graph"): + return ("id", row_id) + return None + + +def _merge_trace_summaries(traces: list) -> dict: + """One Logs-row summary over every trace for it (a turn plus its resumes).""" + totals: dict = {} + for trace in traces: + for key, value in (trace.get("summary") or {}).items(): + if isinstance(value, (int, float)) and not isinstance(value, bool): + totals[key] = totals.get(key, 0) + value + if "retrieval_ms" in totals: + totals["retrieval_ms"] = round(totals["retrieval_ms"], 1) + return { + "count": len(traces), + "duration_ms": sum(int(t.get("duration_ms") or 0) for t in traces), + "status": traces[-1].get("status"), + "started_at": traces[0].get("started_at"), + "summary": totals, + } + + +def _attach_trace_summaries( + conn: Connection, items: list, *, user_id: Optional[str], agent_id: Optional[str] +) -> None: + """Add a ``trace`` summary to each Logs row that has a stored trace. + + One batched lookup per link field for the whole page, rather than a join + inside the UNION, so the timeline query is unchanged. Rows without a + trace (older than the feature, or tracing disabled) get no key. + """ + refs: dict = {} + for item in items: + ref = _trace_ref(item) + if ref: + refs.setdefault(ref[0], set()).add(str(ref[1])) + if not refs: + return + found = RequestTracesRepository(conn).summaries_for_refs( + refs, user_id=user_id, agent_id=agent_id + ) + for item in items: + ref = _trace_ref(item) + if not ref: + continue + traces = found.get(ref[0], {}).get(str(ref[1])) + if traces: + item["trace"] = { + "ref": {"field": ref[0], "value": str(ref[1])}, + **_merge_trace_summaries(traces), + } + + +@analytics_ns.route("/traces") +class GetTraces(Resource): + @api.doc( + description=( + "Stored execution traces (span timelines) for one Logs row. Pass exactly " + "one of message_id, request_id, activity_id, workflow_run_id or id; " + "api_key_id scopes to an agent you own." + ), + params={ + "message_id": "Assistant message id", + "request_id": "Request id (chat turn, scheduled run)", + "activity_id": "Agent activity id (webhook and system rows)", + "workflow_run_id": "Workflow run id", + "id": "Trace id", + "api_key_id": "Agent id to scope to", + }, + ) + def get(self): + decoded_token = request.decoded_token + if not decoded_token: + return make_response(jsonify({"success": False}), 401) + user = decoded_token.get("sub") + refs = [ + (field, request.args.get(field)) + for field in TRACE_REF_FIELDS + if request.args.get(field) + ] + if len(refs) != 1: + return make_response( + jsonify( + { + "success": False, + "message": "Pass exactly one of: " + ", ".join(TRACE_REF_FIELDS), + } + ), + 400, + ) + field, value = refs[0] + api_key_id = request.args.get("api_key_id") + try: + with db_readonly() as conn: + agent, _api_key, agent_pg_id = _resolve_agent(conn, api_key_id, user) + if api_key_id and agent is None: + return make_response(jsonify({"success": True, "traces": []}), 200) + traces = RequestTracesRepository(conn).list_by_ref( + field, value, user_id=user, agent_id=agent_pg_id + ) + except Exception as err: + current_app.logger.error(f"Error getting traces: {err}", exc_info=True) + return make_response(jsonify({"success": False}), 400) + for trace in traces: + trace.pop("_id", None) + return make_response(jsonify({"success": True, "traces": traces}), 200) + + @analytics_ns.route("/get_message_analytics") class GetMessageAnalytics(Resource): get_message_analytics_model = api.model( @@ -701,7 +857,9 @@ class GetUserLogs(Resource): "event_type": fields.String( required=False, description="Filter by event source", - enum=["chat", "schedule", "webhook", "workflow", "system"], + enum=[ + "chat", "schedule", "webhook", "workflow", "system", "search", "graph", + ], ), "search": fields.String( required=False, description="Substring filter on the summary" @@ -741,6 +899,8 @@ class GetUserLogs(Resource): "webhook", "workflow", "system", + "search", + "graph", ): return make_response( jsonify({"success": False, "message": "Invalid option"}), 400 @@ -783,6 +943,16 @@ class GetUserLogs(Resource): "WHERE e.activity_id = s.activity_id " "AND e.level = 'error'))" ) + # A failed chat turn now writes its own chat row (level + # ``error``) linked to its trace; the agent's error row for + # the same activity would list the failure twice. Rows from + # before execution traces, or with tracing off, have no such + # trace and still show here. + chat_failure_dedupe = ( + "NOT EXISTS (SELECT 1 FROM request_traces t " + "WHERE t.activity_id = s.activity_id " + "AND t.source IN ('stream', 'answer', 'v1'))" + ) if api_key_id: # The owner-scoped lookup gates access, so the # chat/webhook/system branches match on the agent @@ -818,6 +988,7 @@ class GetUserLogs(Resource): "s.level = 'error'", "COALESCE(s.endpoint, '') NOT IN ('webhook', 'schedule')", stack_agent_match, + chat_failure_dedupe, ] # Owner-gated agent match: drop the user clause so a # shared agent's runs (stamped with the caller's @@ -831,6 +1002,7 @@ class GetUserLogs(Resource): "wr.user_id = :user_id", "wr.workflow_id = CAST(:agent_workflow_id AS uuid)", ] + trace_scope = "t.agent_id = CAST(:agent_pg_id AS uuid)" else: chat_where = ["l.user_id = :user_id"] webhook_where = [ @@ -842,6 +1014,7 @@ class GetUserLogs(Resource): "s.user_id = :user_id", "s.level = 'error'", "COALESCE(s.endpoint, '') NOT IN ('webhook', 'schedule')", + chat_failure_dedupe, ] # Terminal statuses only (worker writes ``success`` / # ``failed`` / ``timeout`` / ``skipped``; ``completed`` @@ -852,6 +1025,7 @@ class GetUserLogs(Resource): "r.status IN ('success', 'completed', 'failed', 'timeout', 'skipped')", ] workflow_where = ["wr.user_id = :user_id"] + trace_scope = "t.user_id = :user_id" # One normalized timeline over five event sources. # ``payload`` carries the per-type detail; the outer query @@ -892,7 +1066,8 @@ class GetUserLogs(Resource): {summary} AS summary, jsonb_build_object( 'endpoint', s.endpoint, - 'stacks', s.stacks + 'stacks', s.stacks, + 'activity_id', s.activity_id ) AS payload FROM stack_logs s WHERE {where} @@ -913,7 +1088,8 @@ class GetUserLogs(Resource): {summary} AS summary, jsonb_build_object( 'endpoint', s.endpoint, - 'stacks', s.stacks + 'stacks', s.stacks, + 'activity_id', s.activity_id ) AS payload FROM stack_logs s WHERE {where} @@ -989,6 +1165,11 @@ class GetUserLogs(Resource): WHERE {where} """, }, + # Runs with no log row of their own are listed from their + # stored trace: searches (/api/search and MCP) and graph + # builds. + _trace_branch("search", "('search', 'mcp')", trace_scope), + _trace_branch("graph", "('graph_extraction')", trace_scope), ] if level: @@ -1062,6 +1243,9 @@ class GetUserLogs(Resource): "tool_calls": payload.get("tool_calls"), "agent_id": payload.get("agent_id"), "attachments": payload.get("attachments"), + "request_id": payload.get("request_id"), + "message_id": payload.get("message_id"), + "error": payload.get("error"), } ) elif m["event_type"] in ("system", "webhook"): @@ -1072,6 +1256,15 @@ class GetUserLogs(Resource): # before write-time redaction still carry the # reflected provider/user secrets in ``stacks``. "stacks": redact_secrets(payload.get("stacks")), + "activity_id": payload.get("activity_id"), + } + ) + elif m["event_type"] in ("search", "graph"): + item.update( + { + "status": payload.get("status"), + "source": payload.get("source"), + "duration_ms": payload.get("duration_ms"), } ) elif m["event_type"] == "workflow": @@ -1104,6 +1297,22 @@ class GetUserLogs(Resource): } ) results.append(item) + if results: + # Trace chips are an extra: a failed lookup must leave the + # page intact, just without them. + try: + with db_readonly() as conn: + _attach_trace_summaries( + conn, + results, + user_id=user, + agent_id=agent_pg_id if api_key_id else None, + ) + except Exception: + current_app.logger.warning( + "Could not attach trace summaries to the logs page", + exc_info=True, + ) except Exception as err: current_app.logger.error( f"Error getting user logs: {err}", exc_info=True diff --git a/docsgpt/api/user/scheduler_worker.py b/docsgpt/api/user/scheduler_worker.py index 3e085f0e..be05c101 100644 --- a/docsgpt/api/user/scheduler_worker.py +++ b/docsgpt/api/user/scheduler_worker.py @@ -277,6 +277,10 @@ def execute_scheduled_run_body(run_id: str, celery_task_id: Optional[str]) -> Di endpoint="schedule", conversation_id=schedule.get("origin_conversation_id"), chat_history=chat_history, + # Links the run's execution trace to its Logs row, which belongs + # to the user who scheduled it (not always the agent's owner). + request_id=str(run_id), + trace_user_id=run.get("user_id") or schedule.get("user_id"), ) except SoftTimeLimitExceeded: timed_out = True diff --git a/docsgpt/api/user/tasks.py b/docsgpt/api/user/tasks.py index 82d247e7..b34945fe 100644 --- a/docsgpt/api/user/tasks.py +++ b/docsgpt/api/user/tasks.py @@ -608,6 +608,11 @@ def setup_periodic_tasks(sender, **kwargs): cleanup_guardrail_events.s(), name="cleanup-guardrail-events", ) + sender.add_periodic_task( + timedelta(hours=24), + cleanup_traces.s(), + name="cleanup-traces", + ) sender.add_periodic_task( timedelta(hours=24), cleanup_orphan_memories.s(), @@ -827,6 +832,30 @@ def cleanup_guardrail_events(self): return {"deleted": deleted, "ttl_days": ttl_days} +@celery.task(bind=True, acks_late=False) +def cleanup_traces(self): + """Delete ``request_traces`` rows older than ``TRACES_RETENTION_DAYS``. + + Every chat turn, scheduled run and search writes a trace, and each one + carries content previews, so the table is bounded by a retention window + like the other per-request journals. + """ + from docsgpt.core.settings import settings + if not settings.POSTGRES_URI: + return {"deleted": 0, "skipped": "POSTGRES_URI not set"} + + from docsgpt.storage.db.engine import get_engine + from docsgpt.storage.db.repositories.request_traces import ( + RequestTracesRepository, + ) + + ttl_days = settings.TRACES_RETENTION_DAYS + engine = get_engine() + with engine.begin() as conn: + deleted = RequestTracesRepository(conn).purge_older_than(ttl_days) + return {"deleted": deleted, "ttl_days": ttl_days} + + @celery.task(bind=True, acks_late=False) def cleanup_orphan_memories(self): """Sweep orphan memories left by the 0009 FK-to-trigger orphan window. diff --git a/docsgpt/api/v1/routes.py b/docsgpt/api/v1/routes.py index fa915d24..2b943088 100644 --- a/docsgpt/api/v1/routes.py +++ b/docsgpt/api/v1/routes.py @@ -21,7 +21,11 @@ from docsgpt.api.answer.services.continuation_service import ( RESUME_IN_PROGRESS_MESSAGE, ResumeInProgressError, ) -from docsgpt.api.answer.services.stream_processor import StreamProcessor +from docsgpt import tracing +from docsgpt.api.answer.services.stream_processor import ( + StreamProcessor, + flush_trace_after_request, +) from docsgpt.api.v1 import idempotency as v1_idempotency from docsgpt.api.v1.session_store import ( V1Session, @@ -257,7 +261,8 @@ def chat_completions(): internal_data["persist"] = True try: - processor = StreamProcessor(internal_data, decoded_token) + processor = StreamProcessor(internal_data, decoded_token, trace_source="v1") + flush_trace_after_request(processor) # Set when this request took the resume claim, so a refusal can release it. claimed_conversation_id = None @@ -371,6 +376,8 @@ def chat_completions(): # safe way to re-emit a recorded SSE stream (and the regression / # b2b client is non-streaming), so a streaming request never # claims a key. This is a known, accepted limitation. + # The stream writes the trace once it runs, after this returns. + processor.handoff_trace() return Response( with_sse_keepalive( _stream_response( @@ -404,6 +411,9 @@ def chat_completions(): if not claimed: # ``completed`` cache hit, or a 409 for an in-flight same-key # request — either way return without re-running the agent. + # The original request already has its trace; this retry's + # setup (pre-fetch retrieval) is not a failed run to record. + tracing.discard(processor.trace) return replay # An exception from the agent run propagates to the ``except`` handlers @@ -495,6 +505,8 @@ def _stream_response( visibility=visibility, _continuation=continuation, finalize_tool_pause_as_complete=finalize_stateless_tool_pause, + request_id=processor.request_id, + trace=getattr(processor, "trace", None), ) translation_state = StreamTranslationState() @@ -572,6 +584,8 @@ def _non_stream_response( visibility=visibility, _continuation=continuation, finalize_tool_pause_as_complete=finalize_stateless_tool_pause, + request_id=processor.request_id, + trace=getattr(processor, "trace", None), ) result = helper.process_response_stream(stream) diff --git a/docsgpt/cache.py b/docsgpt/cache.py index 4ee069de..820b20c2 100644 --- a/docsgpt/cache.py +++ b/docsgpt/cache.py @@ -8,6 +8,7 @@ from threading import Lock import redis from docsgpt.core.settings import settings +from docsgpt.tracing.llm import CACHE_HIT_ATTR, record_cached_gen from docsgpt.utils import get_hash logger = logging.getLogger(__name__) @@ -240,6 +241,7 @@ def gen_cache(func): if cached_response: decoded = cached_response.decode("utf-8") if not _is_stream_payload(decoded): + record_cached_gen(self, model, decoded) return decoded except Exception as e: logger.error(f"Error getting cached response: {e}", exc_info=True) @@ -295,6 +297,12 @@ def stream_cache(func): if cached_chunks is not None: logger.info(f"Cache hit for stream key: {cache_key}") + # ``stream_token_usage`` wraps this cache and owns + # the call's span; flag it as served from cache. + try: + setattr(self, CACHE_HIT_ATTR, True) + except AttributeError: + pass for chunk in cached_chunks: yield chunk time.sleep(0.03) # Simulate streaming delay diff --git a/docsgpt/core/settings/__init__.py b/docsgpt/core/settings/__init__.py index 817281c5..aab95515 100644 --- a/docsgpt/core/settings/__init__.py +++ b/docsgpt/core/settings/__init__.py @@ -35,6 +35,7 @@ from docsgpt.core.settings.scheduler import SchedulerSettings from docsgpt.core.settings.server import ServerSettings from docsgpt.core.settings.speech import SpeechSettings from docsgpt.core.settings.storage import StorageSettings +from docsgpt.core.settings.tracing import TracingSettings from docsgpt.core.settings.vectorstores import VectorStoreSettings from docsgpt.core.settings.workers import WorkerSettings @@ -55,6 +56,7 @@ SETTINGS_GROUPS: tuple[tuple[str, type[SettingsGroup]], ...] = ( ("Events and devices", EventsSettings), ("Agents", AgentSettings), ("Guardrails", GuardrailSettings), + ("Execution traces", TracingSettings), ("Quotas", QuotaSettings), ("Scheduler", SchedulerSettings), ("Sandbox", SandboxSettings), diff --git a/docsgpt/core/settings/tracing.py b/docsgpt/core/settings/tracing.py new file mode 100644 index 00000000..19a26341 --- /dev/null +++ b/docsgpt/core/settings/tracing.py @@ -0,0 +1,45 @@ +"""Execution traces: the per-request timeline shown in the Logs UI and exported as OTel GenAI spans.""" + +from __future__ import annotations + +from pydantic import Field + +from docsgpt.core.settings._shared import SettingsGroup + + +class TracingSettings(SettingsGroup): + """Recording of agent, LLM, tool and retrieval steps per request.""" + + TRACES_ENABLED: bool = Field( + default=True, + description=( + "Record an execution trace (agent runs, LLM calls, tool calls, retrieval, embeddings) for every " + "request, store it in request_traces and show it in the Logs UI. False records nothing." + ), + ) + TRACES_CAPTURE_CONTENT: bool = Field( + default=True, + description=( + "Store short, secret-redacted previews (tool arguments and results, retrieved chunk titles, " + "rephrased queries, answer excerpts) with each stored trace. Full prompts are never stored. OTel " + "export follows OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT instead." + ), + ) + TRACES_PREVIEW_CHARS: int = Field( + default=2000, ge=100, description="Maximum characters kept per stored trace preview." + ) + TRACES_MAX_SPANS: int = Field( + default=500, + ge=10, + description="Maximum spans recorded per trace; further spans are counted as dropped, not stored.", + ) + TRACES_RETENTION_DAYS: int = Field( + default=30, ge=1, description="Days stored traces are kept before the cleanup task removes them." + ) + TRACES_OTEL_EXPORT: bool = Field( + default=True, + description=( + "Also emit each finished trace as OpenTelemetry GenAI spans (gen_ai.*) and metrics. Has no effect " + "unless an OTel SDK is configured, e.g. by launching under opentelemetry-instrument." + ), + ) diff --git a/docsgpt/graphrag/extraction.py b/docsgpt/graphrag/extraction.py index daa8f910..80efd1fa 100644 --- a/docsgpt/graphrag/extraction.py +++ b/docsgpt/graphrag/extraction.py @@ -23,6 +23,7 @@ import logging import re from typing import Any, Callable, Dict, List, Optional +from docsgpt import tracing from docsgpt.core.model_utils import ( get_api_key_for_provider, get_provider_from_model_id, @@ -353,7 +354,9 @@ def extract_graph_for_source( if pool is not None: # ``map`` yields in submission order, so chunks are still applied in # the order they were given and a run stays reproducible. - prepared = pool.map(_prepare, items) + # Pool threads don't inherit context; carry the trace in so each + # chunk's extraction LLM call is recorded. + prepared = pool.map(tracing.wrap(_prepare), items) else: prepared = (_prepare(item) for item in items) missed = [] diff --git a/docsgpt/guardrails/engine.py b/docsgpt/guardrails/engine.py index 97a359b7..5f42cf02 100644 --- a/docsgpt/guardrails/engine.py +++ b/docsgpt/guardrails/engine.py @@ -5,8 +5,9 @@ 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 from docsgpt.guardrails.config import GuardrailsConfig from docsgpt.guardrails.guardrail_creator import GuardrailCreator @@ -26,6 +27,53 @@ logger = logging.getLogger(__name__) _MAX_WORKERS = 8 +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}", + attributes={ + "docsgpt.guardrail.stage": stage.value, + "docsgpt.guardrail.checks": [c.check for c in controls], + }, + ) + + +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) + + +def _decision_key(decision: StageDecision) -> tuple: + """Identity of a guardrail outcome, for recording each distinct one once.""" + return ( + "guardrail", + decision.stage.value, + tuple(sorted(v.check for v in decision.triggered)), + tuple(sorted(v.check for v in decision.unevaluated)), + decision.blocked, + decision.redacted, + ) + + +def _describe_decision(span: "tracing.Span", decision: StageDecision) -> None: + """Record a stage decision on its span; a firing guardrail drops trace previews.""" + triggered = [v.check for v in decision.triggered] + span.set( + **{ + "docsgpt.guardrail.triggered": triggered or None, + "docsgpt.guardrail.blocked": decision.blocked, + "docsgpt.guardrail.redacted": decision.redacted, + "docsgpt.guardrail.categories": decision.categories() or None, + "docsgpt.guardrail.unevaluated": [v.check for v in decision.unevaluated] or None, + } + ) + if _content_fired(decision): + # The scanned text (or text near it) sits in other spans' previews: + # the retrieval query, tool results, the answer. Keep none of it. + tracing.mark_content_blocked() + + class GuardrailEngine: """Stateless evaluator bound to one agent's resolved config.""" @@ -93,14 +141,29 @@ class GuardrailEngine: return decision if any(self._needs_deadline(c) for c in controls): - decision.verdicts = self._run_concurrent(controls, text, stage) + # Remote checks (e.g. an LLM judge) cost real time, so they are + # always traced; the judge's own LLM call nests under this span. + with _start_guardrail_span(stage, controls) as span: + decision.verdicts = self._run_concurrent(controls, text, stage) + self._reduce(decision) + _describe_decision(span, decision) else: # Bounded local checks are pattern matches measured in # microseconds. Running them inline keeps the streaming hot loop # free of thread churn. decision.verdicts = [self._run_control(c, text, stage) for c in controls] + self._reduce(decision) + # The output guard evaluates every streamed segment; tracing each + # clean local scan would bury the trace, so only a firing is kept, + # and a firing that repeats unchanged segment after segment is + # recorded once so it cannot use up the trace's span cap. + if not decision.clean or decision.unevaluated: + if tracing.first_occurrence(_decision_key(decision)): + with _start_guardrail_span(stage, controls) as span: + _describe_decision(span, decision) + elif _content_fired(decision): + tracing.mark_content_blocked() - self._reduce(decision) self._record(decision) return decision @@ -117,7 +180,7 @@ class GuardrailEngine: for control in controls[:_MAX_WORKERS]: slot: dict = {"control": control, "verdict": None} thread = Thread( - target=self._fill_slot, + target=tracing.wrap(self._fill_slot), args=(slot, control, text, stage), daemon=True, name=f"guardrail-{control.check}", diff --git a/docsgpt/llm/handlers/base.py b/docsgpt/llm/handlers/base.py index 039d7007..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__) @@ -1127,6 +1128,10 @@ class LLMHandler(ABC): # Mark remaining tools as skipped for remaining_call in tool_calls[i:]: + trace_unexecuted_tool_call( + remaining_call, + {"tool_name": "system", "action_name": remaining_call.name, "status": "skipped"}, + ) skip_message = { "type": "tool_call", "data": { @@ -1207,22 +1212,21 @@ class LLMHandler(ABC): message_id=agent.tool_executor.message_id, user_id=agent.tool_executor.user, ) - yield { - "type": "tool_call", - "data": { - "tool_name": pause_info["tool_name"], - "call_id": pause_info["call_id"], - "action_name": pause_info.get( - "llm_name", pause_info["name"] - ), - "arguments": pause_info["arguments"], - "status": "denied", - "error": deny_reason, - "error_type": pause_info.get( - "error_type", "tool_not_allowed" - ), - }, + denied_data = { + "tool_name": pause_info["tool_name"], + "call_id": pause_info["call_id"], + "action_name": pause_info.get( + "llm_name", pause_info["name"] + ), + "arguments": pause_info["arguments"], + "status": "denied", + "error": deny_reason, + "error_type": pause_info.get( + "error_type", "tool_not_allowed" + ), } + 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 pause_data = { @@ -1236,6 +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(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/llm/llm_creator.py b/docsgpt/llm/llm_creator.py index 995600cb..d64a37a2 100644 --- a/docsgpt/llm/llm_creator.py +++ b/docsgpt/llm/llm_creator.py @@ -128,6 +128,10 @@ class LLMCreator: # llm.model_id is the upstream name (BYOM resolves it above); stamp # the canonical id (UUID for BYOM) separately for token_usage. llm._canonical_model_id = model_id + # The provider plugin that built it: ``openai_compatible`` endpoints + # run through the OpenAI client, so the class alone cannot name them + # in traces (see ``docsgpt.tracing.llm.llm_provider``). + llm._provider_plugin = plugin.name # Calls to a user's own model are recorded at $0 (see ``docsgpt/usage.py``). llm._is_byom = model is not None and getattr(model, "source", "builtin") == "user" return llm diff --git a/docsgpt/logging.py b/docsgpt/logging.py index c05265e8..5440631f 100644 --- a/docsgpt/logging.py +++ b/docsgpt/logging.py @@ -7,6 +7,7 @@ import logging import uuid from typing import Any, Callable, Dict, Generator, List, Optional +from docsgpt import tracing from docsgpt.core import log_context from docsgpt.storage.db.repositories.stack_logs import StackLogsRepository from docsgpt.storage.db.session import db_session @@ -136,10 +137,17 @@ def log_activity() -> Callable: }, ) + # The outermost agent run names the trace's activity, which is how + # a webhook/system Logs row (keyed by activity_id) finds its trace. + tracing.bind_if_unset(activity_id=activity_id) + span = start_agent_span(args[0], agent_id=agent_id, model=model, endpoint=endpoint) + error: BaseException | None = None + completed = False try: generator = func(*args, **kwargs) yield from _consume_and_log(generator, context) + completed = True except Exception as exc: # Only ``Exception`` counts as an activity error; ``GeneratorExit`` # (consumer disconnected mid-stream) and ``KeyboardInterrupt`` @@ -154,6 +162,7 @@ def log_activity() -> Callable: started_at=started_at, error=error, ) + _end_agent_span(span, context, error=error, completed=completed) log_context.reset(ctx_token) return wrapper @@ -161,6 +170,78 @@ def log_activity() -> Callable: return decorator +def start_agent_span( + agent: Any, + *, + agent_id: Any = None, + model: Any = None, + endpoint: Optional[str] = None, + continuation: bool = False, +) -> Any: + """Open the ``invoke_agent`` span for one agent run (no-op without a trace). + + The one builder for agent spans: ``@log_activity`` uses it for every run + and ``gen_continuation`` for a resumed one. Values not passed are read off + the agent. + + Args: + agent: The agent instance. + agent_id: Overrides ``agent.agent_id``. + model: Overrides the agent's model (``gpt_model``, ``model`` or ``model_id``). + endpoint: Overrides ``agent.endpoint``. + continuation: True for a run resumed after tool approval. + + Returns: + The span, or a no-op span without an active trace. + """ + label = type(agent).__name__ + agent_id = agent_id or getattr(agent, "agent_id", None) + model = ( + model + or getattr(agent, "gpt_model", None) + or getattr(agent, "model", None) + or getattr(agent, "model_id", None) + ) + endpoint = endpoint or getattr(agent, "endpoint", None) + return tracing.start_span( + tracing.KIND_AGENT, + f"invoke_agent {label}", + attributes={ + k: v + for k, v in { + "gen_ai.operation.name": "invoke_agent", + "gen_ai.agent.id": str(agent_id) if agent_id else None, + "gen_ai.request.model": str(model) if model else None, + "docsgpt.agent_type": label, + "docsgpt.endpoint": str(endpoint) if endpoint else None, + "docsgpt.continuation": True if continuation else None, + }.items() + if v is not None + }, + ) + + +def _end_agent_span(span: Any, context: "LogContext", *, error: BaseException | None, completed: bool) -> None: + """Close the agent span with the run's response aggregates and outcome.""" + if not span: + return + attributes = { + "docsgpt.answer_chars": context.answer_length, + "docsgpt.source_count": context.source_count, + "docsgpt.tool_call_count": context.tool_call_count, + } + if error is not None: + span.end(error=error, attributes=attributes) + elif context.stream_error: + span.set(**attributes, **{"error.type": "StreamError"}) + # The error event's text is gated like any other content. + span.preview("error", context.stream_error) + span.error = "StreamError" + span.end(tracing.STATUS_ERROR) + else: + span.end(None if completed else tracing.STATUS_CANCELLED, attributes=attributes) + + def _emit_activity_finished( *, context: "LogContext", diff --git a/docsgpt/mcp_server.py b/docsgpt/mcp_server.py index f67120b2..9e8c7fbf 100644 --- a/docsgpt/mcp_server.py +++ b/docsgpt/mcp_server.py @@ -53,7 +53,7 @@ async def search_docs(query: str, chunks: int = 5) -> list[dict]: if not api_key: raise PermissionError("Missing Bearer token") try: - return await asyncio.to_thread(search, api_key, query, chunks) + return await asyncio.to_thread(search, api_key, query, chunks, source="mcp") except InvalidAPIKey as exc: raise PermissionError("Invalid API key") from exc except SearchFailed: diff --git a/docsgpt/requirements-docling.txt b/docsgpt/requirements-docling.txt index 2e272a2e..a22caa19 100644 --- a/docsgpt/requirements-docling.txt +++ b/docsgpt/requirements-docling.txt @@ -519,6 +519,7 @@ openpyxl==3.1.5 opentelemetry-api==1.44.0 # via # daytona + # docsgpt # fastmcp-slim # google-api-core # mcp diff --git a/docsgpt/requirements-milvus.txt b/docsgpt/requirements-milvus.txt index 46d3d2b9..713ded5c 100644 --- a/docsgpt/requirements-milvus.txt +++ b/docsgpt/requirements-milvus.txt @@ -438,6 +438,7 @@ openpyxl==3.1.5 opentelemetry-api==1.44.0 # via # daytona + # docsgpt # fastmcp-slim # google-api-core # mcp diff --git a/docsgpt/requirements.txt b/docsgpt/requirements.txt index 17db5b1d..b4b45c8f 100644 --- a/docsgpt/requirements.txt +++ b/docsgpt/requirements.txt @@ -429,6 +429,7 @@ openpyxl==3.1.5 opentelemetry-api==1.44.0 # via # daytona + # docsgpt # fastmcp-slim # google-api-core # mcp diff --git a/docsgpt/retriever/classic_rag.py b/docsgpt/retriever/classic_rag.py index db5e09b9..bcf49038 100644 --- a/docsgpt/retriever/classic_rag.py +++ b/docsgpt/retriever/classic_rag.py @@ -2,6 +2,11 @@ import logging from typing import Any, Dict, List, Optional, Tuple from docsgpt.core.settings import settings +from docsgpt.tracing.retrieval import ( + describe_documents, + start_retrieval_span, + start_source_search_span, +) from docsgpt.llm.llm_creator import LLMCreator from docsgpt.retriever.base import BaseRetriever from docsgpt.retriever.fanout import fetch_per_source, max_parallel_sources @@ -266,6 +271,15 @@ class ClassicRAG(BaseRetriever): logged and reported as ``None`` so one bad source cannot take the rest of the retrieval down with it. """ + span = start_source_search_span( + plan["id"], + top_k=plan["src_k"], + **{ + "docsgpt.vector_store": settings.VECTOR_STORE, + "docsgpt.retriever": type(self).__name__, + "docsgpt.shared_query_vector": query_vector is not None, + }, + ) try: if docsearch is None: docsearch = VectorCreator.create_vectorstore( @@ -279,8 +293,14 @@ class ClassicRAG(BaseRetriever): query_vector=query_vector, ) score_kind = self._score_kind(docsearch) if self.include_scores else None + try: + candidates = len(docs_temp) + except TypeError: + candidates = None + span.end(attributes={"docsgpt.candidate_count": candidates}) return docs_temp, score_kind except Exception as e: + span.end(error=e) logger.error( f"Error searching vectorstore {plan['id']}: {e}", exc_info=True ) @@ -409,11 +429,20 @@ class ClassicRAG(BaseRetriever): def search(self, query: str = ""): """Search for documents using optional query override""" - if query: - self.original_question = query - # Invalidate the cached rephrase so a per-source path that opts in - # rephrases against the new query, not a stale one. - self._rephrased_question = None - self.question = self._rephrase_query() - self._rephrased_question = self.question - return self._get_data() + with start_retrieval_span( + f"retrieval {type(self).__name__}", + sources=self.vectorstores, + **{"docsgpt.retriever": type(self).__name__, "docsgpt.top_k": self.chunks}, + ) as span: + if query: + self.original_question = query + # Invalidate the cached rephrase so a per-source path that opts in + # rephrases against the new query, not a stale one. + self._rephrased_question = None + self.question = self._rephrase_query() + self._rephrased_question = self.question + docs = self._get_data() + rephrased = self._rephrased_question + span.set(**{"docsgpt.rephrased": bool(rephrased and rephrased != self.original_question)}) + describe_documents(span, docs, query=rephrased or self.original_question) + return docs diff --git a/docsgpt/retriever/dispatcher.py b/docsgpt/retriever/dispatcher.py index 31cc75e0..7ec23c5c 100644 --- a/docsgpt/retriever/dispatcher.py +++ b/docsgpt/retriever/dispatcher.py @@ -23,6 +23,7 @@ from docsgpt.retriever.stages.prescreen import ( max_candidate_k, ) from docsgpt.storage.db.source_config import RetrievalConfig +from docsgpt.tracing.retrieval import describe_documents, start_retrieval_span from docsgpt.utils import num_tokens_from_string logger = logging.getLogger(__name__) @@ -294,6 +295,17 @@ class Dispatcher(BaseRetriever): def search(self, query: str = "") -> List[Dict[str, Any]]: """Run every group under the shared budget and merge the results.""" + sources = [s for group in self._groups for s in group.get("doc_ids", [])] + with start_retrieval_span( + "retrieval", + sources=sources, + **{"docsgpt.retriever": "Dispatcher", "docsgpt.group_count": len(self._groups)}, + ) as span: + docs = self._search_groups(query) + describe_documents(span, docs) + return docs + + def _search_groups(self, query: str) -> List[Dict[str, Any]]: groups = self._groups n_groups = len(groups) diff --git a/docsgpt/retriever/fanout.py b/docsgpt/retriever/fanout.py index cb1ad3a6..fcca623d 100644 --- a/docsgpt/retriever/fanout.py +++ b/docsgpt/retriever/fanout.py @@ -12,7 +12,9 @@ import logging from concurrent.futures import ThreadPoolExecutor from typing import Any, Callable, Dict, Iterable, List, Optional, TypeVar +from docsgpt import tracing from docsgpt.core.settings import settings +from docsgpt.tracing.retrieval import start_embedding_span logger = logging.getLogger(__name__) @@ -85,16 +87,22 @@ def embed_questions(docsearch, questions: Iterable[str]) -> Dict[str, List[float embedder = store_embeddings(docsearch) if embedder is None: return {} + distinct = list(dict.fromkeys(questions)) + span = start_embedding_span( + settings.EMBEDDINGS_NAME, + inputs=len(distinct), + **{"docsgpt.embeddings_backend": type(embedder).__name__}, + ) try: - return { - question: embedder.embed_query(question) - for question in dict.fromkeys(questions) - } + vectors = {question: embedder.embed_query(question) for question in distinct} except Exception as e: + span.end(error=e) logger.warning( "Query embedding failed (%s); each store will embed its own.", e ) return {} + span.end() + return vectors def run_source_jobs( @@ -121,10 +129,13 @@ def run_source_jobs( workers = max_parallel_sources(len(jobs)) if workers == 1: return [fn(job) for job in jobs] + # Pool threads don't inherit context: carry the trace and its current + # retrieval span in so per-source spans nest under it. + traced_fn = tracing.wrap(fn) with ThreadPoolExecutor( max_workers=workers, thread_name_prefix="rag-source" ) as pool: - return list(pool.map(fn, jobs)) + return list(pool.map(traced_fn, jobs)) def fetch_per_source( diff --git a/docsgpt/retriever/graph_rag.py b/docsgpt/retriever/graph_rag.py index 8d191f7b..59ee6163 100644 --- a/docsgpt/retriever/graph_rag.py +++ b/docsgpt/retriever/graph_rag.py @@ -33,6 +33,11 @@ from docsgpt.retriever.base import BaseRetriever from docsgpt.retriever.classic_rag import ClassicRAG from docsgpt.retriever.labels import labels_from_metadata from docsgpt.storage.db.source_config import GraphRetrievalConfig +from docsgpt.tracing.retrieval import ( + describe_documents, + start_embedding_span, + start_retrieval_span, +) from docsgpt.utils import num_tokens_from_string from docsgpt.vectorstore.base import get_embeddings @@ -231,7 +236,11 @@ class GraphRAGRetriever(BaseRetriever): def _embed_query(self, question: str) -> List[float]: embedding = get_embeddings() - return embedding.embed_query(question) + with start_embedding_span( + settings.EMBEDDINGS_NAME, + **{"docsgpt.embeddings_backend": type(embedding).__name__}, + ): + return embedding.embed_query(question) def _ppr_scores(self, subgraph, seeds) -> Dict[str, float]: """Run Personalized PageRank, then down-weight hub nodes by IDF. @@ -730,10 +739,19 @@ class GraphRAGRetriever(BaseRetriever): logging.debug("Error closing GraphRAG store: %s", e) def search(self, query: str = "") -> List[Dict[str, Any]]: - if query: - self.original_question = query - self._classic.original_question = query - self._classic._rephrased_question = None - self._classic.question = self._classic._rephrase_query() - self._classic._rephrased_question = self._classic.question - return self._get_data() + with start_retrieval_span( + f"retrieval {type(self).__name__}", + sources=self.vectorstores, + **{"docsgpt.retriever": type(self).__name__, "docsgpt.top_k": self.chunks}, + ) as span: + if query: + self.original_question = query + self._classic.original_question = query + self._classic._rephrased_question = None + self._classic.question = self._classic._rephrase_query() + self._classic._rephrased_question = self._classic.question + docs = self._get_data() + describe_documents( + span, docs, query=self._classic._rephrased_question or self.original_question + ) + return docs diff --git a/docsgpt/retriever/stages/prescreen.py b/docsgpt/retriever/stages/prescreen.py index cae50dea..2dfa5140 100644 --- a/docsgpt/retriever/stages/prescreen.py +++ b/docsgpt/retriever/stages/prescreen.py @@ -18,6 +18,7 @@ import re from concurrent.futures import ThreadPoolExecutor from typing import Any, Callable, Dict, List, Optional +from docsgpt import tracing from docsgpt.llm.llm_creator import LLMCreator from docsgpt.storage.db.source_config import PreScreenConfig @@ -164,20 +165,28 @@ class PreScreenStage: ] llm = self._build_llm() - max_workers = min(_MAX_WORKERS, len(batches)) - kept: List[Dict[str, Any]] = [] - if max_workers <= 1: - for batch in batches: - kept.extend(self._screen_batch(llm, query, batch)) - else: - with ThreadPoolExecutor(max_workers=max_workers) as pool: - results = pool.map( - lambda b: self._screen_batch(llm, query, b), batches - ) - for batch_result in results: - kept.extend(batch_result) - - return kept[: self.config.max_keep] + with tracing.span( + tracing.KIND_RERANK, + "rerank prescreen", + attributes={ + "gen_ai.operation.name": "rerank", + "docsgpt.candidate_count": len(docs), + "docsgpt.batch_count": len(batches), + }, + ) as span: + max_workers = min(_MAX_WORKERS, len(batches)) + kept: List[Dict[str, Any]] = [] + if max_workers <= 1: + for batch in batches: + kept.extend(self._screen_batch(llm, query, batch)) + else: + screen = tracing.wrap(lambda b: self._screen_batch(llm, query, b)) + with ThreadPoolExecutor(max_workers=max_workers) as pool: + for batch_result in pool.map(screen, batches): + kept.extend(batch_result) + survivors = kept[: self.config.max_keep] + span.set(**{"docsgpt.kept_count": len(survivors)}) + return survivors def build_prescreen_stages( diff --git a/docsgpt/services/search_service.py b/docsgpt/services/search_service.py index b67a1aca..68756c91 100644 --- a/docsgpt/services/search_service.py +++ b/docsgpt/services/search_service.py @@ -10,10 +10,12 @@ from __future__ import annotations import logging from typing import Any, Dict, List, Optional +from docsgpt import tracing from docsgpt.core.settings import settings from docsgpt.retriever.fanout import fetch_per_source from docsgpt.storage.db.repositories.agents import AgentsRepository from docsgpt.storage.db.session import db_readonly +from docsgpt.tracing.retrieval import describe_documents, start_retrieval_span from docsgpt.vectorstore.vector_creator import VectorCreator logger = logging.getLogger(__name__) @@ -219,14 +221,21 @@ def _search_sources( return results[:chunks] -def search(api_key: str, query: str, chunks: int = 5) -> List[Dict[str, Any]]: +def search( + api_key: str, query: str, chunks: int = 5, *, source: str = "search" +) -> List[Dict[str, Any]]: """Resolve an agent by API key and search its sources. + Every search that reaches the sources is recorded as an execution trace + owned by the agent's owner, under ``source``. + Args: api_key: Agent API key (the opaque string stored on ``agents.key`` in Postgres). query: Free-text search query. chunks: Max number of hits to return. + source: Trace source name: ``search`` for ``/api/search``, ``mcp`` + for the MCP ``search_docs`` tool. Returns: List of hit dicts with ``text``, ``title``, ``source`` keys. @@ -256,4 +265,20 @@ def search(api_key: str, query: str, chunks: int = 5) -> List[Dict[str, Any]]: if not source_ids: return [] - return _search_sources(query, source_ids, chunks) + trace = tracing.start_trace( + source=source, + user_id=agent.get("user_id"), + agent_id=str(agent.get("id")) if agent.get("id") else None, + ) + with tracing.activate(trace): + try: + with start_retrieval_span( + f"retrieval {source}", + sources=source_ids, + **{"docsgpt.top_k": chunks}, + ) as span: + results = _search_sources(query, source_ids, chunks) + describe_documents(span, results, query=query) + return results + finally: + tracing.flush(trace) diff --git a/docsgpt/storage/db/models.py b/docsgpt/storage/db/models.py index 79773b6e..ebd25567 100644 --- a/docsgpt/storage/db/models.py +++ b/docsgpt/storage/db/models.py @@ -825,6 +825,47 @@ Index( Index("ix_guardrail_events_message", guardrail_events_table.c.message_id) Index("ix_guardrail_events_created", guardrail_events_table.c.created_at) +# One execution trace per request (chat turn, continuation, scheduled or +# webhook run, search, graph extraction): the span tree as a JSONB array, +# rendered as a waterfall in the Logs UI. ``message_id`` cascades; deleting a +# conversation also deletes its traces by ``conversation_id``. Migration 0037. +request_traces_table = Table( + "request_traces", + metadata, + Column("id", UUID(as_uuid=True), primary_key=True), + Column("request_id", Text), + Column( + "message_id", + UUID(as_uuid=True), + ForeignKey("conversation_messages.id", ondelete="CASCADE"), + ), + Column("conversation_id", UUID(as_uuid=True)), + Column("activity_id", Text), + Column("workflow_run_id", UUID(as_uuid=True)), + Column("user_id", Text), + Column("agent_id", UUID(as_uuid=True)), + Column("source", Text, nullable=False), + Column("name", Text), + # ok | error | paused | cancelled + Column("status", Text, nullable=False), + Column("started_at", DateTime(timezone=True), nullable=False), + Column("duration_ms", Integer, nullable=False, server_default="0"), + Column("span_count", Integer, nullable=False, server_default="0"), + Column("dropped_spans", Integer, nullable=False, server_default="0"), + Column("summary", JSONB, nullable=False, server_default="{}"), + Column("spans", JSONB, nullable=False, server_default="[]"), + Column("otel_trace_id", Text), + Column("created_at", DateTime(timezone=True), nullable=False, server_default=func.now()), +) + +Index( + "request_traces_user_source_started_idx", + request_traces_table.c.user_id, + request_traces_table.c.source, + request_traces_table.c.started_at, +) +Index("request_traces_created_idx", request_traces_table.c.created_at) + tool_call_attempts_table = Table( "tool_call_attempts", metadata, diff --git a/docsgpt/storage/db/repositories/conversations.py b/docsgpt/storage/db/repositories/conversations.py index 95205350..7d27a3f5 100644 --- a/docsgpt/storage/db/repositories/conversations.py +++ b/docsgpt/storage/db/repositories/conversations.py @@ -576,6 +576,18 @@ class ConversationsRepository: artifacts = ArtifactsRepository(self._conn) paths = artifacts.storage_paths_for_conversation(conversation_id) + # Traces linked to a message go with it (ON DELETE CASCADE). This + # takes the rest -- scheduled runs, stateless /v1 rounds, turns whose + # message was never reserved -- whose tool results and chunk previews + # the user expects gone with the conversation. + self._conn.execute( + text( + "DELETE FROM request_traces WHERE conversation_id = CAST(:id AS uuid) " + "AND EXISTS (SELECT 1 FROM conversations " + "WHERE id = CAST(:id AS uuid) AND user_id = :user_id)" + ), + {"id": conversation_id, "user_id": user_id}, + ) result = self._conn.execute( text( "DELETE FROM conversations " @@ -597,6 +609,14 @@ class ConversationsRepository: artifacts = ArtifactsRepository(self._conn) paths = artifacts.storage_paths_for_user_conversations(user_id) artifacts.delete_for_user_conversations(user_id) + # See ``delete``: traces not linked to a message are removed explicitly. + self._conn.execute( + text( + "DELETE FROM request_traces WHERE conversation_id IN " + "(SELECT id FROM conversations WHERE user_id = :user_id)" + ), + {"user_id": user_id}, + ) result = self._conn.execute( text("DELETE FROM conversations WHERE user_id = :user_id"), {"user_id": user_id}, diff --git a/docsgpt/storage/db/repositories/request_traces.py b/docsgpt/storage/db/repositories/request_traces.py new file mode 100644 index 00000000..79a7e3cb --- /dev/null +++ b/docsgpt/storage/db/repositories/request_traces.py @@ -0,0 +1,221 @@ +"""Repository for ``request_traces``: one stored execution trace per request.""" + +from __future__ import annotations + +import datetime +import json +import logging +from typing import Any, Dict, Iterable, List, Optional + +from sqlalchemy import Connection, text +from sqlalchemy.exc import IntegrityError + +from docsgpt.storage.db.base_repository import looks_like_uuid, row_to_dict +from docsgpt.storage.db.serialization import PGNativeJSONEncoder +from docsgpt.utils import strip_null_bytes + +logger = logging.getLogger(__name__) + +#: Columns a trace can be looked up by, mapped to their SQL type. +REF_FIELDS: Dict[str, str] = { + "id": "uuid", + "request_id": "text", + "message_id": "uuid", + "activity_id": "text", + "workflow_run_id": "uuid", +} + +_SUMMARY_COLUMNS = ( + "id, request_id, message_id, conversation_id, activity_id, workflow_run_id, " + "user_id, agent_id, source, name, status, started_at, duration_ms, " + "span_count, dropped_spans, summary, otel_trace_id" +) + + +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 + + +def _started_at(record: Dict[str, Any]) -> datetime.datetime: + started_ns = record.get("started_at_ns") + if started_ns: + return datetime.datetime.fromtimestamp(started_ns / 1e9, tz=datetime.timezone.utc) + return datetime.datetime.now(datetime.timezone.utc) + + +def _scope_clause(user_id: Optional[str], agent_id: Optional[str]) -> tuple[str, Dict[str, Any]]: + """Owner scoping: an owned agent's traces, else the caller's own traces.""" + if agent_id: + return "agent_id = CAST(:scope_agent AS uuid)", {"scope_agent": str(agent_id)} + return "user_id = :scope_user", {"scope_user": user_id} + + +class RequestTracesRepository: + def __init__(self, conn: Connection) -> None: + self._conn = conn + + def insert(self, record: Dict[str, Any]) -> bool: + """Insert one finished trace (the dict from ``Trace.to_record``). + + A foreign-key failure means the message was deleted before the trace + flushed (e.g. the turn was superseded); the trace describes a turn the + user discarded, so it is dropped rather than raised. The insert runs + in a savepoint so that drop leaves the caller's transaction usable. + + Args: + record: The trace record. + + Returns: + True when the row was written. + """ + params = { + "id": record["id"], + "request_id": record.get("request_id"), + "message_id": _uuid_or_none(record.get("message_id")), + "conversation_id": _uuid_or_none(record.get("conversation_id")), + "activity_id": record.get("activity_id"), + "workflow_run_id": _uuid_or_none(record.get("workflow_run_id")), + "user_id": record.get("user_id"), + "agent_id": _uuid_or_none(record.get("agent_id")), + "source": record.get("source") or "unknown", + "name": strip_null_bytes(record.get("name")), + "status": record.get("status") or "ok", + "started_at": _started_at(record), + "duration_ms": int(record.get("duration_ms") or 0), + "span_count": int(record.get("span_count") or 0), + "dropped_spans": int(record.get("dropped_spans") or 0), + "summary": _dump_jsonb(record.get("summary") or {}), + "spans": _dump_jsonb(record.get("spans") or []), + "otel_trace_id": record.get("otel_trace_id"), + } + statement = text( + """ + INSERT INTO request_traces + (id, request_id, message_id, conversation_id, activity_id, + workflow_run_id, user_id, agent_id, source, name, status, + started_at, duration_ms, span_count, dropped_spans, summary, + spans, otel_trace_id) + VALUES + (CAST(:id AS uuid), :request_id, CAST(:message_id AS uuid), + CAST(:conversation_id AS uuid), :activity_id, + CAST(:workflow_run_id AS uuid), :user_id, CAST(:agent_id AS uuid), + :source, :name, :status, :started_at, :duration_ms, :span_count, + :dropped_spans, CAST(:summary AS jsonb), CAST(:spans AS jsonb), + :otel_trace_id) + ON CONFLICT (id) DO NOTHING + """ + ) + try: + with self._conn.begin_nested(): + result = self._conn.execute(statement, params) + except IntegrityError as exc: + sqlstate = getattr(getattr(exc, "orig", None), "sqlstate", None) + if sqlstate != "23503": + raise + logger.info("Dropped trace %s: its message no longer exists", record["id"]) + return False + return (result.rowcount or 0) > 0 + + def list_by_ref( + self, + field: str, + value: str, + *, + user_id: Optional[str], + agent_id: Optional[str] = None, + limit: int = 20, + ) -> List[dict]: + """Full traces matching ``field = value``, oldest first, owner-scoped. + + Args: + field: One of :data:`REF_FIELDS`. + value: The id to match. + user_id: The caller; used when ``agent_id`` is not given. + agent_id: An agent the caller owns; scopes to that agent's traces. + limit: Maximum traces returned. + + Returns: + Trace dicts including ``spans``; empty for an unknown field or id. + """ + sql_type = REF_FIELDS.get(field) + if sql_type is None or not value: + return [] + if sql_type == "uuid" and not looks_like_uuid(str(value)): + return [] + if not agent_id and not user_id: + return [] + scope, scope_params = _scope_clause(user_id, agent_id) + result = self._conn.execute( + text( + f""" + SELECT {_SUMMARY_COLUMNS}, spans FROM request_traces + WHERE {field} = CAST(:value AS {sql_type}) AND {scope} + ORDER BY started_at + LIMIT :limit + """ + ), + {"value": str(value), "limit": max(1, min(int(limit), 100)), **scope_params}, + ) + return [row_to_dict(row) for row in result.fetchall()] + + def summaries_for_refs( + self, + refs: Dict[str, Iterable[str]], + *, + user_id: Optional[str], + agent_id: Optional[str] = None, + ) -> Dict[str, Dict[str, List[dict]]]: + """Trace summaries (no spans) for a page of Logs rows. + + Args: + refs: ``{field: [ids...]}`` for fields in :data:`REF_FIELDS`. + user_id: The caller; used when ``agent_id`` is not given. + agent_id: An agent the caller owns. + + Returns: + ``{field: {id: [summary, ...]}}`` with summaries oldest first. + """ + out: Dict[str, Dict[str, List[dict]]] = {} + if not agent_id and not user_id: + return out + scope, scope_params = _scope_clause(user_id, agent_id) + for field, values in refs.items(): + sql_type = REF_FIELDS.get(field) + if sql_type is None: + continue + ids = sorted({str(v) for v in values if v}) + if sql_type == "uuid": + ids = [v for v in ids if looks_like_uuid(v)] + if not ids: + continue + result = self._conn.execute( + text( + f""" + SELECT {_SUMMARY_COLUMNS} FROM request_traces + WHERE {field} = ANY(CAST(:ids AS {sql_type}[])) AND {scope} + ORDER BY started_at + """ + ), + {"ids": ids, **scope_params}, + ) + for row in result.fetchall(): + data = row_to_dict(row) + out.setdefault(field, {}).setdefault(str(data[field]), []).append(data) + return out + + def purge_older_than(self, days: int) -> int: + """Delete traces older than ``days`` (the retention window).""" + result = self._conn.execute( + text( + "DELETE FROM request_traces " + "WHERE created_at < NOW() - CAST(:days || ' days' AS interval)" + ), + {"days": str(max(1, days))}, + ) + return result.rowcount or 0 diff --git a/docsgpt/tracing/__init__.py b/docsgpt/tracing/__init__.py new file mode 100644 index 00000000..d31ba387 --- /dev/null +++ b/docsgpt/tracing/__init__.py @@ -0,0 +1,89 @@ +"""Per-request execution traces. + +Records agent runs, LLM calls, tool calls, retrieval and embeddings as a tree +of timed spans. A finished trace is stored in ``request_traces`` (rendered as +a waterfall in the Logs UI) and replayed as OpenTelemetry GenAI spans when an +OTel SDK is configured. See ``docs/content/Deploying/Observability.mdx``. + +Typical use:: + + trace = tracing.start_trace(source="stream", request_id=request_id) + with tracing.activate(trace): + with tracing.span(tracing.KIND_TOOL, "execute_tool search") as s: + s.preview("arguments", args) + ... + tracing.flush(trace) + +Every call is a no-op when no trace is active. +""" + +from docsgpt.tracing.core import ( + BINDABLE_IDS, + CONTAINER_KINDS, + KIND_AGENT, + KIND_EMBEDDING, + KIND_GUARDRAIL, + KIND_LLM, + KIND_RERANK, + KIND_RETRIEVAL, + KIND_SEARCH, + KIND_STEP, + KIND_TOOL, + NOOP_SPAN, + STATUS_CANCELLED, + STATUS_DENIED, + STATUS_ERROR, + STATUS_OK, + STATUS_PAUSED, + STATUS_PENDING, + STATUS_SKIPPED, + Span, + Trace, + activate, + bind, + bind_if_unset, + current_trace, + first_occurrence, + mark_content_blocked, + span, + start_span, + start_trace, + wrap, +) +from docsgpt.tracing.sink import discard, flush + +__all__ = [ + "BINDABLE_IDS", + "CONTAINER_KINDS", + "KIND_AGENT", + "KIND_EMBEDDING", + "KIND_GUARDRAIL", + "KIND_LLM", + "KIND_RERANK", + "KIND_RETRIEVAL", + "KIND_SEARCH", + "KIND_STEP", + "KIND_TOOL", + "NOOP_SPAN", + "STATUS_CANCELLED", + "STATUS_DENIED", + "STATUS_ERROR", + "STATUS_OK", + "STATUS_PAUSED", + "STATUS_PENDING", + "STATUS_SKIPPED", + "Span", + "Trace", + "activate", + "bind", + "bind_if_unset", + "current_trace", + "discard", + "first_occurrence", + "flush", + "mark_content_blocked", + "span", + "start_span", + "start_trace", + "wrap", +] diff --git a/docsgpt/tracing/core.py b/docsgpt/tracing/core.py new file mode 100644 index 00000000..6e3aef52 --- /dev/null +++ b/docsgpt/tracing/core.py @@ -0,0 +1,640 @@ +"""In-memory recording of one execution trace. + +A :class:`Trace` collects spans for a single request (a chat turn, a +scheduled run, a search). Spans are recorded while the request runs and the +whole trace is written once, when its owner calls :func:`flush`. + +Nesting is tracked with one span stack per thread. Only *container* spans +(agent, tool, retrieval, rerank, guardrail, step) are pushed, so a leaf such as an LLM call can +never become the parent of a sibling that starts while it is still open -- +which matters because DocsGPT's agent loop is a chain of suspended +generators. Ending a span pops anything left above it (marked +``cancelled``), and :meth:`Trace.finish` closes whatever is still open, so a +generator finalized late can never corrupt the tree or the stored record. +""" + +from __future__ import annotations + +import contextlib +import functools +import logging +import threading +import time +import uuid +from contextvars import ContextVar +from typing import Any, Callable, Dict, Iterator, List, Optional + +from docsgpt.core.settings import settings +from docsgpt.tracing.preview import make_preview + +logger = logging.getLogger(__name__) + +KIND_AGENT = "agent" +KIND_LLM = "llm" +KIND_TOOL = "tool" +KIND_RETRIEVAL = "retrieval" +KIND_SEARCH = "search" +KIND_EMBEDDING = "embedding" +KIND_RERANK = "rerank" +KIND_GUARDRAIL = "guardrail" +KIND_STEP = "step" + +#: Kinds that become the implicit parent of spans started while they are open. +CONTAINER_KINDS = frozenset( + {KIND_AGENT, KIND_TOOL, KIND_RETRIEVAL, KIND_RERANK, KIND_GUARDRAIL, KIND_STEP} +) + +STATUS_OK = "ok" +STATUS_ERROR = "error" +STATUS_CANCELLED = "cancelled" +STATUS_PAUSED = "paused" +STATUS_PENDING = "pending" +STATUS_DENIED = "denied" +STATUS_SKIPPED = "skipped" + +#: Trace ids ``bind`` may set; anything else is ignored. +BINDABLE_IDS = frozenset( + { + "request_id", + "message_id", + "conversation_id", + "activity_id", + "workflow_run_id", + "user_id", + "agent_id", + "name", + } +) + +_current: ContextVar[Optional["Trace"]] = ContextVar("docsgpt_trace", default=None) + + +def _new_id() -> str: + return uuid.uuid4().hex[:16] + + +class Span: + """One timed step inside a trace. + + Attribute keys follow the OTel GenAI conventions (``gen_ai.*``) where one + exists and ``docsgpt.*`` otherwise, so the stored trace and the exported + spans share a vocabulary. + """ + + __slots__ = ( + "trace", + "id", + "parent_id", + "kind", + "name", + "attributes", + "previews", + "status", + "error", + "start_perf_ns", + "end_perf_ns", + "_thread", + "_pushed", + ) + + def __init__( + self, + trace: "Trace", + kind: str, + name: str, + parent_id: Optional[str], + attributes: Optional[Dict[str, Any]] = None, + ) -> None: + self.trace = trace + self.id = _new_id() + self.parent_id = parent_id + self.kind = kind + self.name = name + self.attributes: Dict[str, Any] = dict(attributes or {}) + self.previews: Dict[str, Any] = {} + self.status: Optional[str] = None + self.error: Optional[str] = None + self.start_perf_ns = time.perf_counter_ns() + self.end_perf_ns: Optional[int] = None + self._thread: Optional[int] = None + self._pushed = False + + @property + def ended(self) -> bool: + return self.end_perf_ns is not None + + @property + def duration_ms(self) -> Optional[float]: + if self.end_perf_ns is None: + return None + return (self.end_perf_ns - self.start_perf_ns) / 1e6 + + def set(self, **attributes: Any) -> "Span": + """Merge attributes; ``None`` values are ignored. Keys may contain dots via ``**{...}``.""" + if not self.ended: + self.attributes.update({k: v for k, v in attributes.items() if v is not None}) + return self + + def preview(self, key: str, value: Any) -> "Span": + """Attach a bounded, redacted content preview (skipped when capture is off).""" + if self.ended or value is None or not settings.TRACES_CAPTURE_CONTENT: + return self + try: + self.previews[key] = make_preview(value) + except Exception: # noqa: BLE001 - a preview must never break a request + logger.debug("trace preview failed for %s", key, exc_info=True) + return self + + def fail(self, exc: BaseException) -> "Span": + """Record ``exc`` on the span without ending it. + + ``error`` and ``error.type`` hold only the exception type: they are + stored and exported whatever the content settings, and exception text + can quote the prompt (a provider's content-filter error) or a tool + result. The message itself becomes the capture-gated ``error`` + preview, dropped with every other preview when capture is off or a + guardrail fired. + """ + self.status = STATUS_ERROR + self.error = type(exc).__name__ + self.attributes["error.type"] = type(exc).__name__ + message = str(exc) + if message: + self.preview("error", message) + return self + + def end( + self, + status: Optional[str] = None, + *, + error: Optional[BaseException] = None, + attributes: Optional[Dict[str, Any]] = None, + ) -> None: + """Close the span. A second end, or an end after the trace finished, is ignored.""" + if self.ended: + return + if attributes: + self.set(**attributes) + if error is not None: + self.fail(error) + if status is not None: + self.status = status + elif self.status is None: + self.status = STATUS_OK + self.end_perf_ns = time.perf_counter_ns() + self.trace._on_end(self) + + def __enter__(self) -> "Span": + return self + + def __exit__(self, exc_type, exc, tb) -> bool: + if exc is None: + self.end() + elif isinstance(exc, GeneratorExit): + self.end(STATUS_CANCELLED) + else: + self.end(error=exc) + return False + + +class _NoopSpan: + """Stand-in returned when no trace is active or the span cap is reached.""" + + id = None + parent_id = None + kind = None + name = "" + status = None + ended = True + duration_ms = None + + @property + def attributes(self) -> Dict[str, Any]: + return {} + + @property + def previews(self) -> Dict[str, Any]: + return {} + + def set(self, **attributes: Any) -> "_NoopSpan": + return self + + def preview(self, key: str, value: Any) -> "_NoopSpan": + return self + + def fail(self, exc: BaseException) -> "_NoopSpan": + return self + + def end(self, status=None, *, error=None, attributes=None) -> None: + return None + + def __enter__(self) -> "_NoopSpan": + return self + + def __exit__(self, exc_type, exc, tb) -> bool: + return False + + def __bool__(self) -> bool: + return False + + +NOOP_SPAN = _NoopSpan() + + +class Trace: + """All spans recorded for one execution, plus the ids that link it to logs.""" + + def __init__( + self, + *, + source: str, + name: Optional[str] = None, + request_id: Optional[str] = None, + message_id: Optional[str] = None, + conversation_id: Optional[str] = None, + activity_id: Optional[str] = None, + workflow_run_id: Optional[str] = None, + user_id: Optional[str] = None, + agent_id: Optional[str] = None, + otel_context: Any = None, + ) -> None: + self.id = str(uuid.uuid4()) + self.source = source + self.name = name or source + self.request_id = request_id + self.message_id = message_id + self.conversation_id = conversation_id + self.activity_id = activity_id + self.workflow_run_id = workflow_run_id + self.user_id = user_id + self.agent_id = agent_id + self.otel_context = otel_context + self.otel_trace_id: Optional[str] = None + self.start_ns = time.time_ns() + self.start_perf_ns = time.perf_counter_ns() + self.end_perf_ns: Optional[int] = None + self.status: Optional[str] = None + # Status the owner determined mid-run (e.g. ``paused`` at a tool + # approval); used by :meth:`finish` when no explicit status is given. + self.outcome: Optional[str] = None + self.spans: List[Span] = [] + self.dropped_spans = 0 + self.content_blocked = False + self.attributes: Dict[str, Any] = {} + self.finished = False + self.flushed = False + self._seen_keys: set = set() + self._lock = threading.Lock() + self._stacks: Dict[int, List[Any]] = {} + + # -- span lifecycle ------------------------------------------------- + + def _parent_for_current_thread(self) -> Optional[str]: + stack = self._stacks.get(threading.get_ident()) + return stack[-1].id if stack else None + + def start_span( + self, + kind: str, + name: str, + *, + parent: Any = None, + attributes: Optional[Dict[str, Any]] = None, + ) -> Any: + """Record a new span; returns :data:`NOOP_SPAN` once finished or over the cap.""" + with self._lock: + if self.finished: + return NOOP_SPAN + if len(self.spans) >= settings.TRACES_MAX_SPANS: + self.dropped_spans += 1 + return NOOP_SPAN + if parent is not None: + parent_id = getattr(parent, "id", None) + else: + parent_id = self._parent_for_current_thread() + span = Span(self, kind, name, parent_id, attributes) + self.spans.append(span) + if kind in CONTAINER_KINDS: + ident = threading.get_ident() + self._stacks.setdefault(ident, []).append(span) + span._thread = ident + span._pushed = True + return span + + def _on_end(self, span: Span) -> None: + if not span._pushed: + return + with self._lock: + stack = self._stacks.get(span._thread) + if not stack or span not in stack: + return + index = stack.index(span) + abandoned = stack[index + 1:] + del stack[index:] + if not stack: + self._stacks.pop(span._thread, None) + for child in reversed(abandoned): + if isinstance(child, Span) and not child.ended: + child.end(STATUS_CANCELLED) + + def _seed_thread(self, parent: Optional[Span]) -> Callable[[], None]: + """Make ``parent`` the implicit parent in this thread; returns an undo callable.""" + if parent is None: + return lambda: None + ident = threading.get_ident() + marker = _Seed(parent.id) + with self._lock: + self._stacks.setdefault(ident, []).append(marker) + + def _undo() -> None: + with self._lock: + stack = self._stacks.get(ident) + if stack and marker in stack: + del stack[stack.index(marker):] + if not stack: + self._stacks.pop(ident, None) + + return _undo + + def current_parent(self) -> Optional[Any]: + stack = self._stacks.get(threading.get_ident()) + return stack[-1] if stack else None + + def first_occurrence(self, key: Any) -> bool: + """True the first time ``key`` is seen in this trace, then False.""" + with self._lock: + if key in self._seen_keys: + return False + self._seen_keys.add(key) + return True + + # -- ids -------------------------------------------------------------- + + def bind(self, *, only_if_unset: bool = False, **ids: Any) -> None: + for key, value in ids.items(): + if key not in BINDABLE_IDS or value is None: + continue + if only_if_unset and getattr(self, key, None): + continue + setattr(self, key, str(value)) + + # -- completion --------------------------------------------------------- + + def finish(self, status: Optional[str] = None) -> None: + """Freeze the trace: close open spans as ``cancelled`` and set the status.""" + with self._lock: + if self.finished: + return + self.finished = True + open_spans = [s for s in self.spans if not s.ended] + self._stacks.clear() + now = time.perf_counter_ns() + for span in open_spans: + span.status = STATUS_CANCELLED + span.end_perf_ns = now + self.end_perf_ns = now + if status is None: + status = self.outcome + if status is not None: + self.status = status + else: + failed = any( + s.parent_id is None and s.status == STATUS_ERROR for s in self.spans + ) + self.status = STATUS_ERROR if failed else STATUS_OK + + @property + def duration_ms(self) -> Optional[float]: + if self.end_perf_ns is None: + return None + return (self.end_perf_ns - self.start_perf_ns) / 1e6 + + def span_start_ns(self, span: Span) -> int: + """Wall-clock start of ``span`` in ns, derived from the monotonic offset.""" + return self.start_ns + (span.start_perf_ns - self.start_perf_ns) + + def span_end_ns(self, span: Span) -> int: + end = span.end_perf_ns if span.end_perf_ns is not None else self.end_perf_ns + return self.start_ns + ((end or span.start_perf_ns) - self.start_perf_ns) + + def summary(self) -> Dict[str, Any]: + """Aggregate counts shown as chips in the Logs UI.""" + by_id = {s.id: s for s in self.spans} + llm = [s for s in self.spans if s.kind == KIND_LLM] + # Outermost retrieval spans only, so nested dispatcher/retriever + # spans are not double-counted. + retrieval = [ + s + for s in self.spans + if s.kind == KIND_RETRIEVAL + and not (s.parent_id in by_id and by_id[s.parent_id].kind == KIND_RETRIEVAL) + ] + # Only calls that ran count: a call paused for approval is recorded + # again when it runs in the next round, and denied or skipped calls + # never ran at all. + tools = [ + s for s in self.spans if s.kind == KIND_TOOL and s.status in (STATUS_OK, STATUS_ERROR) + ] + + def _tokens(key: str) -> int: + total = 0 + for s in llm: + value = s.attributes.get(key) + if isinstance(value, (int, float)): + total += int(value) + return total + + return { + "llm_calls": len(llm), + "tool_calls": len(tools), + "retrieval_calls": len(retrieval), + "retrieval_ms": round(sum(s.duration_ms or 0 for s in retrieval), 1), + "input_tokens": _tokens("gen_ai.usage.input_tokens"), + "output_tokens": _tokens("gen_ai.usage.output_tokens"), + "errors": sum(1 for s in self.spans if s.status == STATUS_ERROR), + } + + def to_record(self) -> Dict[str, Any]: + """The ``request_traces`` row for this (finished) trace.""" + summary = self.summary() + # Traces without a log row of their own (searches) are listed by + # their query; keep it in the small summary so listing never reads + # ``spans``. Dropped with the other previews when content is blocked. + if not self.content_blocked and self.spans: + query = self.spans[0].previews.get("query") + if isinstance(query, str): + summary["query"] = query + spans = [] + for s in self.spans: + entry: Dict[str, Any] = { + "id": s.id, + "parent_id": s.parent_id, + "kind": s.kind, + "name": s.name, + "status": s.status or STATUS_CANCELLED, + "offset_ms": round((s.start_perf_ns - self.start_perf_ns) / 1e6, 2), + "duration_ms": round(s.duration_ms or 0.0, 2), + "attributes": s.attributes, + } + if s.error: + entry["error"] = s.error + if s.previews and not self.content_blocked: + entry["preview"] = s.previews + spans.append(entry) + return { + "id": self.id, + "request_id": self.request_id, + "message_id": self.message_id, + "conversation_id": self.conversation_id, + "activity_id": self.activity_id, + "workflow_run_id": self.workflow_run_id, + "user_id": self.user_id, + "agent_id": self.agent_id, + "source": self.source, + "name": self.name, + "status": self.status or STATUS_OK, + "started_at_ns": self.start_ns, + "duration_ms": int(round(self.duration_ms or 0)), + "span_count": len(spans), + "dropped_spans": self.dropped_spans, + "summary": summary, + "spans": spans, + "otel_trace_id": self.otel_trace_id, + } + + +class _Seed: + """Placeholder stack entry naming a parent span owned by another thread.""" + + __slots__ = ("id",) + + def __init__(self, span_id: Optional[str]) -> None: + self.id = span_id + + +# -- module-level API -------------------------------------------------------- + + +def current_trace() -> Optional[Trace]: + """The trace active in this context, or ``None``.""" + return _current.get() + + +def start_trace(*, source: str, capture_otel_context: bool = True, **ids: Any) -> Optional[Trace]: + """Create a trace, or return ``None`` when ``TRACES_ENABLED`` is off. + + The OTel context current at this point (normally the HTTP server span) is + captured so the exported GenAI spans hang off the request's own trace. + """ + if not settings.TRACES_ENABLED: + return None + otel_context = None + if capture_otel_context: + try: + from opentelemetry import context as otel_ctx + + otel_context = otel_ctx.get_current() + except Exception: # noqa: BLE001 + otel_context = None + known = {k: v for k, v in ids.items() if k in BINDABLE_IDS} + known = {k: (str(v) if v is not None else None) for k, v in known.items()} + return Trace(source=source, otel_context=otel_context, **known) + + +@contextlib.contextmanager +def activate(trace: Optional[Trace]) -> Iterator[Optional[Trace]]: + """Make ``trace`` current for the enclosed block (a no-op for ``None``).""" + if trace is None: + yield None + return + token = _current.set(trace) + try: + yield trace + finally: + try: + _current.reset(token) + except ValueError: + # Reset from a different context (a generator finalized + # elsewhere); clearing is the closest safe equivalent. + _current.set(None) + + +def start_span( + kind: str, name: str, *, parent: Any = None, attributes: Optional[Dict[str, Any]] = None +) -> Any: + """Start a span in the current trace; returns :data:`NOOP_SPAN` without one.""" + trace = _current.get() + if trace is None: + return NOOP_SPAN + try: + return trace.start_span(kind, name, parent=parent, attributes=attributes) + except Exception: # noqa: BLE001 - tracing must never break a request + logger.debug("trace start_span failed", exc_info=True) + return NOOP_SPAN + + +def span( + kind: str, name: str, *, parent: Any = None, attributes: Optional[Dict[str, Any]] = None +) -> Any: + """Context-manager form of :func:`start_span` (errors mark the span and re-raise).""" + return start_span(kind, name, parent=parent, attributes=attributes) + + +def bind(**ids: Any) -> None: + """Set link ids (``message_id``, ``conversation_id``, ...) on the current trace.""" + trace = _current.get() + if trace is not None: + trace.bind(**ids) + + +def bind_if_unset(**ids: Any) -> None: + """Like :func:`bind` but keeps any value already set.""" + trace = _current.get() + if trace is not None: + trace.bind(only_if_unset=True, **ids) + + +def first_occurrence(key: Any) -> bool: + """True the first time ``key`` is seen in the current trace. + + For steps that repeat identically many times in one request (a guardrail + re-firing on every streamed segment): record the first, skip the rest, so + they cannot use up the span cap. False without an active trace. + """ + trace = _current.get() + return trace.first_occurrence(key) if trace is not None else False + + +def mark_content_blocked() -> None: + """A guardrail blocked or retracted content: drop every preview from the stored trace.""" + trace = _current.get() + if trace is not None: + trace.content_blocked = True + + +def wrap(fn: Callable[..., Any]) -> Callable[..., Any]: + """Bind the current trace and parent span into ``fn`` for another thread. + + Use around work handed to a thread pool or ``threading.Thread``; those do + not inherit context variables. Without an active trace ``fn`` is returned + unchanged. + """ + trace = _current.get() + if trace is None: + return fn + parent = trace.current_parent() + + @functools.wraps(fn) + def _run(*args: Any, **kwargs: Any) -> Any: + token = _current.set(trace) + undo = trace._seed_thread(parent) + try: + return fn(*args, **kwargs) + finally: + undo() + try: + _current.reset(token) + except ValueError: + _current.set(None) + + return _run diff --git a/docsgpt/tracing/llm.py b/docsgpt/tracing/llm.py new file mode 100644 index 00000000..ef169186 --- /dev/null +++ b/docsgpt/tracing/llm.py @@ -0,0 +1,187 @@ +"""``chat`` spans and GenAI metrics for LLM calls. + +Called from the token-usage wrappers in ``docsgpt/usage.py``: one span per +decorated invocation, so a primary attempt, its same-provider retry and a +fallback each get their own span with the provider that actually ran. +""" + +from __future__ import annotations + +from typing import Any, Dict, Iterable, Optional, Union +from urllib.parse import urlparse + +from docsgpt.core.settings import settings +from docsgpt.tracing import core +from docsgpt.tracing.otel import provider_name, record_llm_metrics + +#: LLM attribute a cache wrapper sets when it served the call from Redis. +CACHE_HIT_ATTR = "_trace_cache_hit" + +#: API hosts whose provider differs from the client class that calls them +#: (every OpenAI-compatible API runs through ``OpenAILLM``), mapped to their +#: ``gen_ai.provider.name``. Matched on the host or any subdomain of it. +_PROVIDER_HOSTS = ( + ("api.openai.com", "openai"), + ("openai.azure.com", "azure.ai.openai"), + ("deepseek.com", "deepseek"), + ("mistral.ai", "mistral_ai"), + ("x.ai", "x_ai"), + ("perplexity.ai", "perplexity"), + ("groq.com", "groq"), + ("openrouter.ai", "openrouter"), + ("novita.ai", "novita"), + ("cohere.com", "cohere"), + ("cohere.ai", "cohere"), + ("together.xyz", "together_ai"), + ("together.ai", "together_ai"), + ("fireworks.ai", "fireworks"), + ("anthropic.com", "anthropic"), + ("generativelanguage.googleapis.com", "gcp.gemini"), +) + + +def _endpoint_host(llm: Any) -> Optional[str]: + url = getattr(llm, "_effective_base_url", None) or getattr(llm, "base_url", None) + if not isinstance(url, str) or not url: + return None + try: + return urlparse(url).hostname + except ValueError: + return None + + +def llm_provider(llm: Any) -> str: + """The ``gen_ai.provider.name`` of the API ``llm`` actually calls. + + The client class alone is not enough: DeepSeek, a self-hosted vLLM or any + other OpenAI-compatible endpoint runs through ``OpenAILLM``, whose + ``provider_name`` is ``openai``. A known API host wins; an OpenAI client + pointed at any other host is reported as ``openai_compatible``. + + Args: + llm: The LLM instance making the call. + + Returns: + The provider name for spans and metrics. + """ + base = provider_name(getattr(llm, "provider_name", None)) + host = _endpoint_host(llm) + if host: + for suffix, name in _PROVIDER_HOSTS: + if host == suffix or host.endswith("." + suffix): + return name + if getattr(llm, "_provider_plugin", None) == "openai_compatible": + return "openai_compatible" + if base == "openai" and host: + # An OpenAI client aimed somewhere other than OpenAI's API. + return "openai_compatible" + return base + + +def start_llm_span( + llm: Any, model: Optional[str], *, stream: bool, tools: Any = None +) -> Union[core.Span, Any]: + """Open a ``chat {model}`` span for a call on ``llm`` (no-op without a trace).""" + if core.current_trace() is None: + return core.NOOP_SPAN + attributes = { + "gen_ai.operation.name": "chat", + "gen_ai.provider.name": llm_provider(llm), + "server.address": _endpoint_host(llm), + "gen_ai.request.model": str(model) if model else None, + "docsgpt.token_source": getattr(llm, "_token_usage_source", None) or "agent_stream", + "docsgpt.stream": bool(stream), + "docsgpt.tool_count": len(tools) if tools else None, + } + return core.start_span( + core.KIND_LLM, + f"chat {model}" if model else "chat", + attributes={k: v for k, v in attributes.items() if v is not None}, + ) + + +def output_text(chunks: Iterable[Any]) -> str: + """Concatenate the text deltas of a streamed or returned response.""" + if isinstance(chunks, str): + return chunks + return "".join(chunk for chunk in chunks if isinstance(chunk, str)) + + +def finish_llm_call( + span: Any, + llm: Any, + model: Optional[str], + call_usage: Dict[str, Any], + *, + duration_ms: int, + error: Optional[BaseException], + completed: bool = True, + ttft_ms: Optional[int] = None, + cost_usd: Optional[float] = None, + estimated: bool = True, + output: Union[str, Iterable[Any], None] = None, +) -> None: + """Close the call's span and record the GenAI client metrics. + + Args: + span: The span from :func:`start_llm_span`. + llm: The LLM instance that ran the call. + model: Model the call was made with. + call_usage: Final token counts (provider-reported when available). + duration_ms: Provider time for the call. + error: The exception the call raised, if any. + completed: False when a stream was abandoned before it finished. + ttft_ms: Time to first streamed chunk. + cost_usd: Cost priced for the call, when known. + estimated: True when token counts are local estimates. + output: The response text, or the streamed chunks to join into it; + joined only when the span keeps a preview. + """ + cache_hit = bool(getattr(llm, CACHE_HIT_ATTR, False)) + try: + setattr(llm, CACHE_HIT_ATTR, False) + except AttributeError: + pass + if not settings.TRACES_ENABLED: + return + record_llm_metrics( + provider=llm_provider(llm), + model=str(model) if model else None, + input_tokens=call_usage.get("prompt_tokens", 0), + output_tokens=call_usage.get("generated_tokens", 0), + duration_s=max(duration_ms, 0) / 1000.0, + error_type=type(error).__name__ if error is not None else None, + ) + if not span: + return + if output and settings.TRACES_CAPTURE_CONTENT: + span.preview("output", output_text(output)) + span.end( + None if completed or error is not None else core.STATUS_CANCELLED, + error=error, + attributes={ + "gen_ai.usage.input_tokens": int(call_usage.get("prompt_tokens") or 0), + "gen_ai.usage.output_tokens": int(call_usage.get("generated_tokens") or 0), + "gen_ai.usage.cache_read.input_tokens": call_usage.get("cached_tokens"), + "gen_ai.usage.cache_creation.input_tokens": call_usage.get("cache_write_tokens"), + "docsgpt.usage_estimated": estimated, + "docsgpt.provider_ms": duration_ms, + "docsgpt.ttft_ms": ttft_ms, + "docsgpt.cost_usd": cost_usd if isinstance(cost_usd, (int, float)) else None, + "docsgpt.cache_hit": True if cache_hit else None, + }, + ) + + +def record_cached_gen(llm: Any, model: Optional[str], output: Optional[str]) -> None: + """Record a non-streaming call answered from the response cache. + + The gen cache wraps the usage wrapper, so a hit never reaches it; this + records the zero-cost call so the trace still shows it happened. + """ + span = start_llm_span(llm, model, stream=False) + if not span: + return + if output: + span.preview("output", output) + span.end(attributes={"docsgpt.cache_hit": True}) diff --git a/docsgpt/tracing/otel.py b/docsgpt/tracing/otel.py new file mode 100644 index 00000000..40a8b5ef --- /dev/null +++ b/docsgpt/tracing/otel.py @@ -0,0 +1,274 @@ +"""OpenTelemetry export for execution traces (GenAI semantic conventions). + +A finished :class:`~docsgpt.tracing.core.Trace` is *replayed*: every recorded +span is re-created with its original start and end timestamps and explicit +parents, under the OTel context captured when the trace started (normally the +HTTP server span). Replaying after the fact means no OTel context is ever +attached across the agent loop's generator yields. The cost is that GenAI +spans reach the backend only when the request ends, and auto-instrumented +HTTP/DB spans made during a step do not nest under it. + +Attribute names follow the (still incubating) ``gen_ai.*`` conventions and are +written as literals rather than imported from the semconv package. Prompt and +tool content is exported only when +``OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT`` opts in. + +Nothing here runs unless an OTel SDK is installed as the global provider (for +example by launching under ``opentelemetry-instrument``). +""" + +from __future__ import annotations + +import json +import logging +import os +import threading +from typing import Any, Dict, Optional + +from docsgpt.core.settings import settings +from docsgpt.tracing.core import ( + KIND_EMBEDDING, + KIND_LLM, + KIND_TOOL, + STATUS_ERROR, + STATUS_OK, + Trace, +) + +logger = logging.getLogger(__name__) + +_INSTRUMENTATION_NAME = "docsgpt.tracing" + +#: DocsGPT provider names that differ from the ``gen_ai.provider.name`` values. +_PROVIDER_NAMES = { + "google": "gcp.gen_ai", + "azure_openai": "azure.ai.openai", + "aws_bedrock": "aws.bedrock", + "mistral": "mistral_ai", +} + +#: Preview keys with a dedicated GenAI attribute; others export as ``docsgpt.preview.``. +_TOOL_PREVIEW_ATTRIBUTES = { + "arguments": "gen_ai.tool.call.arguments", + "result": "gen_ai.tool.call.result", +} + +_CAPTURE_VALUES = {"span_only", "span_and_event", "true"} + +# Advisory buckets from the GenAI metrics conventions. +_DURATION_BUCKETS = [0.01, 0.02, 0.04, 0.08, 0.16, 0.32, 0.64, 1.28, 2.56, 5.12, 10.24, 20.48, 40.96, 81.92] +_TOKEN_BUCKETS = [1, 4, 16, 64, 256, 1024, 4096, 16384, 65536, 262144, 1048576, 4194304, 16777216, 67108864] + + +def provider_name(internal: Optional[str]) -> str: + """Map a DocsGPT ``provider_name`` to the ``gen_ai.provider.name`` value.""" + if not internal: + return "unknown" + return _PROVIDER_NAMES.get(internal, internal) + + +def content_capture_enabled() -> bool: + """Whether the standard GenAI content-capture variable opts into span content.""" + value = os.getenv("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", "") + return value.strip().lower() in _CAPTURE_VALUES + + +def _is_noop(provider: Any) -> bool: + from opentelemetry import trace as ot_trace + + return isinstance(provider, (ot_trace.ProxyTracerProvider, ot_trace.NoOpTracerProvider)) + + +def _attr_value(value: Any) -> Any: + """Coerce a recorded attribute into an OTel-legal value (or ``None`` to skip).""" + if value is None: + return None + if isinstance(value, (str, bool, int, float)): + return value + if isinstance(value, (list, tuple)): + items = [v for v in value if v is not None] + if items and all(isinstance(v, str) for v in items): + return list(items) + if items and all(isinstance(v, bool) for v in items): + return list(items) + if items and all(isinstance(v, (int, float)) and not isinstance(v, bool) for v in items): + return list(items) + if not items: + return None + try: + return json.dumps(value, ensure_ascii=False, default=str) + except (TypeError, ValueError): + return str(value) + + +def _attributes(raw: Dict[str, Any]) -> Dict[str, Any]: + out = {} + for key, value in raw.items(): + coerced = _attr_value(value) + if coerced is not None: + out[key] = coerced + return out + + +def _preview_attributes(kind: str, previews: Dict[str, Any]) -> Dict[str, Any]: + out = {} + for key, value in previews.items(): + name = _TOOL_PREVIEW_ATTRIBUTES.get(key) if kind == KIND_TOOL else None + out[name or f"docsgpt.preview.{key}"] = value if isinstance(value, str) else _attr_value(value) + return out + + +def _root_attributes(trace: Trace) -> Dict[str, Any]: + return _attributes( + { + "docsgpt.trace.id": trace.id, + "docsgpt.source": trace.source, + "docsgpt.request_id": trace.request_id, + "docsgpt.message_id": trace.message_id, + "docsgpt.activity_id": trace.activity_id, + "docsgpt.workflow_run_id": trace.workflow_run_id, + "gen_ai.conversation.id": trace.conversation_id, + "gen_ai.agent.id": trace.agent_id, + "docsgpt.status": trace.status, + "docsgpt.dropped_spans": trace.dropped_spans or None, + } + ) + + +def export_trace(trace: Trace, tracer_provider: Any = None) -> Optional[str]: + """Replay ``trace`` as OTel spans; returns the OTel trace id (hex) or ``None``. + + Args: + trace: A finished trace. + tracer_provider: Override the global provider (tests). + + Returns: + The 32-char hex OTel trace id, or ``None`` when export is disabled, + no SDK is configured, or the trace is empty. + """ + if not settings.TRACES_OTEL_EXPORT or not trace.spans: + return None + from opentelemetry import trace as ot_trace + from opentelemetry.trace import SpanKind, Status, StatusCode + + provider = tracer_provider or ot_trace.get_tracer_provider() + if _is_noop(provider): + return None + tracer = provider.get_tracer(_INSTRUMENTATION_NAME) + capture = content_capture_enabled() and not trace.content_blocked + + root = tracer.start_span( + f"docsgpt {trace.source}", + context=trace.otel_context, + kind=SpanKind.INTERNAL, + attributes=_root_attributes(trace), + start_time=trace.start_ns, + ) + if trace.status == STATUS_ERROR: + root.set_status(Status(StatusCode.ERROR)) + exported: Dict[str, Any] = {} + for span in trace.spans: + parent = exported.get(span.parent_id, root) + attributes = dict(span.attributes) + if trace.conversation_id and span.kind in (KIND_LLM, KIND_TOOL): + attributes.setdefault("gen_ai.conversation.id", trace.conversation_id) + if span.status and span.status not in (STATUS_OK, STATUS_ERROR): + attributes["docsgpt.status"] = span.status + attributes = _attributes(attributes) + if capture and span.previews: + attributes.update(_preview_attributes(span.kind, span.previews)) + otel_span = tracer.start_span( + span.name, + context=ot_trace.set_span_in_context(parent), + kind=SpanKind.CLIENT if span.kind in (KIND_LLM, KIND_EMBEDDING) else SpanKind.INTERNAL, + attributes=attributes, + start_time=trace.span_start_ns(span), + ) + if span.status == STATUS_ERROR: + otel_span.set_status(Status(StatusCode.ERROR, span.error)) + exported[span.id] = otel_span + for span in trace.spans: + exported[span.id].end(end_time=trace.span_end_ns(span)) + root.end(end_time=trace.start_ns + int((trace.duration_ms or 0) * 1e6)) + context = root.get_span_context() + return format(context.trace_id, "032x") if context.trace_id else None + + +# -- metrics -------------------------------------------------------------------- + +_instruments_lock = threading.Lock() +_global_instruments: Optional[tuple] = None + + +def _make_instruments(meter_provider: Any = None) -> tuple: + from opentelemetry import metrics + + meter = ( + meter_provider.get_meter(_INSTRUMENTATION_NAME) + if meter_provider is not None + else metrics.get_meter(_INSTRUMENTATION_NAME) + ) + + def _histogram(name: str, unit: str, description: str, buckets: list): + try: + return meter.create_histogram( + name, unit=unit, description=description, + explicit_bucket_boundaries_advisory=buckets, + ) + except TypeError: # API older than bucket advisories + return meter.create_histogram(name, unit=unit, description=description) + + duration = _histogram( + "gen_ai.client.operation.duration", "s", "GenAI operation duration.", _DURATION_BUCKETS + ) + tokens = _histogram( + "gen_ai.client.token.usage", "{token}", "Number of input and output tokens used.", _TOKEN_BUCKETS + ) + return duration, tokens + + +def _instruments(meter_provider: Any = None) -> tuple: + global _global_instruments + if meter_provider is not None: + return _make_instruments(meter_provider) + with _instruments_lock: + if _global_instruments is None: + _global_instruments = _make_instruments() + return _global_instruments + + +def record_llm_metrics( + *, + provider: Optional[str], + model: Optional[str], + input_tokens: int, + output_tokens: int, + duration_s: float, + error_type: Optional[str], + operation: str = "chat", + meter_provider: Any = None, +) -> None: + """Record the two GenAI client metrics for one model call. + + Token usage is recorded only for successful calls, as the conventions + require; duration is recorded for every call, with ``error.type`` on + failures. Without a configured SDK the meter is a no-op. + """ + if not settings.TRACES_OTEL_EXPORT: + return + try: + duration, tokens = _instruments(meter_provider) + attributes = { + "gen_ai.operation.name": operation, + "gen_ai.provider.name": provider_name(provider), + } + if model: + attributes["gen_ai.request.model"] = str(model) + if error_type: + duration.record(duration_s, {**attributes, "error.type": error_type}) + return + duration.record(duration_s, attributes) + tokens.record(int(input_tokens or 0), {**attributes, "gen_ai.token.type": "input"}) + tokens.record(int(output_tokens or 0), {**attributes, "gen_ai.token.type": "output"}) + except Exception: # noqa: BLE001 - metrics must never break a call + logger.debug("Failed to record GenAI metrics", exc_info=True) diff --git a/docsgpt/tracing/preview.py b/docsgpt/tracing/preview.py new file mode 100644 index 00000000..892e00f2 --- /dev/null +++ b/docsgpt/tracing/preview.py @@ -0,0 +1,58 @@ +"""Bounded, secret-redacted content previews stored with a trace span.""" + +from __future__ import annotations + +import json +from typing import Any + +from docsgpt.core.settings import settings +from docsgpt.storage.db.redaction import redact_secrets +from docsgpt.utils import strip_null_bytes + +_ELLIPSIS = "…" +# A list longer than this keeps its head only; the preview is for reading, +# not for reconstructing the payload. +_MAX_LIST_ITEMS = 50 + + +def _bound(value: Any, limit: int) -> Any: + """Truncate strings and long lists inside ``value``; stringify unknown objects.""" + if isinstance(value, str): + return value if len(value) <= limit else value[:limit] + _ELLIPSIS + if isinstance(value, dict): + return {str(k): _bound(v, limit) for k, v in value.items()} + if isinstance(value, (list, tuple)): + items = [_bound(v, limit) for v in list(value)[:_MAX_LIST_ITEMS]] + if len(value) > _MAX_LIST_ITEMS: + items.append(f"{_ELLIPSIS} {len(value) - _MAX_LIST_ITEMS} more") + return items + if value is None or isinstance(value, (bool, int, float)): + return value + return _bound(str(value), limit) + + +def make_preview(value: Any, limit: int | None = None) -> Any: + """Return a storable preview of ``value``. + + Secret-keyed fields are redacted, NUL bytes stripped and every string + truncated to ``limit`` (``TRACES_PREVIEW_CHARS`` by default). A structure + whose JSON form is still far larger than ``limit`` collapses to a + truncated JSON string so one huge tool result cannot bloat the row. + + Args: + value: Any JSON-like value (tool arguments, a result, a query string). + limit: Maximum characters per string; defaults to the setting. + + Returns: + A JSON-serializable preview. + """ + limit = limit or settings.TRACES_PREVIEW_CHARS + bounded = strip_null_bytes(_bound(redact_secrets(value), limit)) + if isinstance(bounded, (dict, list)): + try: + encoded = json.dumps(bounded, ensure_ascii=False, default=str) + except (TypeError, ValueError): + encoded = str(bounded) + if len(encoded) > limit * 4: + return encoded[:limit] + _ELLIPSIS + return bounded diff --git a/docsgpt/tracing/retrieval.py b/docsgpt/tracing/retrieval.py new file mode 100644 index 00000000..374fd207 --- /dev/null +++ b/docsgpt/tracing/retrieval.py @@ -0,0 +1,98 @@ +"""Spans for RAG retrieval: the search itself, the query embedding, each source.""" + +from __future__ import annotations + +from typing import Any, Dict, Iterable, List, Optional + +from docsgpt.tracing import core + +#: Characters of each retrieved chunk kept in the stored preview. +_SNIPPET_CHARS = 240 +#: Chunks listed in the preview; the count attribute still reports all of them. +_MAX_PREVIEW_CHUNKS = 20 + + +def start_retrieval_span( + name: str, *, sources: Optional[Iterable[Any]] = None, **attributes: Any +) -> Any: + """Open a ``retrieval`` span (a container: embeddings and searches nest under it).""" + source_ids = [str(s) for s in (sources or []) if s] + base = { + "gen_ai.operation.name": "retrieval", + "docsgpt.source_ids": source_ids or None, + "gen_ai.data_source.id": source_ids[0] if len(source_ids) == 1 else None, + } + base.update(attributes) + return core.start_span( + core.KIND_RETRIEVAL, + name, + attributes={k: v for k, v in base.items() if v is not None}, + ) + + +def describe_documents(span: Any, docs: Optional[List[Dict[str, Any]]], *, query: Optional[str] = None) -> None: + """Record what a retrieval returned: counts, sources, top scores and a chunk preview.""" + if not span: + return + docs = docs or [] + scores = [d.get("score") for d in docs if isinstance(d, dict) and isinstance(d.get("score"), (int, float))] + score_kinds = {d.get("score_kind") for d in docs if isinstance(d, dict) and d.get("score_kind")} + span.set( + **{ + "docsgpt.chunk_count": len(docs), + "docsgpt.top_score": max(scores) if scores else None, + "docsgpt.score_kind": score_kinds.pop() if len(score_kinds) == 1 else None, + } + ) + if query: + span.preview("query", query) + if docs: + span.preview( + "chunks", + [ + { + k: v + for k, v in { + "title": d.get("title") or d.get("filename"), + "source": d.get("source"), + "score": d.get("score"), + "text": str(d.get("text") or "")[:_SNIPPET_CHARS], + }.items() + if v not in (None, "") + } + for d in docs[:_MAX_PREVIEW_CHUNKS] + if isinstance(d, dict) + ], + ) + + +def start_embedding_span(model: Optional[str], *, inputs: int = 1, **attributes: Any) -> Any: + """Open an ``embeddings`` span for a query embedding (leaf).""" + base = { + "gen_ai.operation.name": "embeddings", + "gen_ai.request.model": model, + "docsgpt.input_count": inputs, + } + base.update(attributes) + return core.start_span( + core.KIND_EMBEDDING, + f"embeddings {model}" if model else "embeddings", + attributes={k: v for k, v in base.items() if v is not None}, + ) + + +def start_source_search_span( + source_id: Any, *, top_k: Optional[int] = None, **attributes: Any +) -> Any: + """Open a per-source vector ``search`` span (leaf).""" + base = { + "gen_ai.operation.name": "search", + "gen_ai.data_source.id": str(source_id) if source_id else None, + "docsgpt.top_k": top_k, + } + base.update(attributes) + return core.start_span( + core.KIND_SEARCH, + f"search {source_id}" if source_id else "search", + attributes={k: v for k, v in base.items() if v is not None}, + ) diff --git a/docsgpt/tracing/sink.py b/docsgpt/tracing/sink.py new file mode 100644 index 00000000..49a516a8 --- /dev/null +++ b/docsgpt/tracing/sink.py @@ -0,0 +1,93 @@ +"""Write a finished trace to its two sinks: Postgres and OpenTelemetry.""" + +from __future__ import annotations + +import logging +import threading +from concurrent.futures import Future, ThreadPoolExecutor +from typing import Optional + +from docsgpt.tracing.core import Trace + +logger = logging.getLogger(__name__) + +# Writes handed off by ``flush(background=True)``. Two workers keep up with +# chat traffic (each write is one INSERT); the pool's threads are joined at +# interpreter exit, so queued traces are still written on a clean shutdown. +_writer: Optional[ThreadPoolExecutor] = None +_writer_lock = threading.Lock() + + +def _executor() -> ThreadPoolExecutor: + global _writer + with _writer_lock: + if _writer is None: + _writer = ThreadPoolExecutor(max_workers=2, thread_name_prefix="trace-writer") + return _writer + + +def _write(trace: Trace) -> None: + """Export ``trace`` to OTel, then store it; each sink swallows its own failure. + + OTel replay runs first so the exported trace id can be stored with the + row. A trace is diagnostic data and must never fail the request that + produced it. + """ + try: + from docsgpt.tracing.otel import export_trace + + trace.otel_trace_id = export_trace(trace) + except Exception: # noqa: BLE001 + logger.warning("Failed to export trace %s to OpenTelemetry", trace.id, exc_info=True) + if not trace.spans: + # Nothing ran worth showing (e.g. an early validation error). + return + try: + from docsgpt.storage.db.repositories.request_traces import RequestTracesRepository + from docsgpt.storage.db.session import db_session + + with db_session() as conn: + RequestTracesRepository(conn).insert(trace.to_record()) + except Exception: # noqa: BLE001 + logger.warning("Failed to store trace %s", trace.id, exc_info=True) + + +def flush( + trace: Optional[Trace], status: Optional[str] = None, *, background: bool = False +) -> Optional[Future]: + """Finish ``trace`` and persist it once; later calls are no-ops. + + The trace is frozen immediately, so nothing recorded afterwards is + included, whichever way it is written. + + Args: + trace: The trace to write; ``None`` is accepted and ignored. + status: Final status; defaults to ``error`` when a top-level span + failed, else ``ok``. + background: Write on a writer thread instead of the caller's. A chat + stream flushes this way so its connection closes without waiting + on the OTel replay and the INSERT. + + Returns: + The pending write when ``background`` is set, else ``None``. + """ + if trace is None or trace.flushed: + return None + trace.flushed = True + trace.finish(status) + if background: + try: + return _executor().submit(_write, trace) + except RuntimeError: + # The pool is shut down (interpreter exit): write inline instead. + pass + _write(trace) + return None + + +def discard(trace: Optional[Trace]) -> None: + """Drop ``trace`` without writing it (the request was rejected or superseded).""" + if trace is None or trace.flushed: + return + trace.flushed = True + trace.finish() diff --git a/docsgpt/usage.py b/docsgpt/usage.py index 8ac6c436..0f48dfbc 100644 --- a/docsgpt/usage.py +++ b/docsgpt/usage.py @@ -3,6 +3,7 @@ import time from typing import Any, Dict from docsgpt.pricing import compute_cost_usd +from docsgpt.tracing.llm import finish_llm_call, start_llm_span from docsgpt.storage.db.repositories.token_usage import TokenUsageRepository from docsgpt.storage.db.session import db_session from docsgpt.utils import num_tokens_from_object_or_list, num_tokens_from_string @@ -108,9 +109,12 @@ def _persist_call_usage(llm, call_usage, *, duration_ms=None, ttft_ms=None): duration_ms: Wall-clock for the call, measured by the wrapper. ttft_ms: Time to the first streamed chunk; None for a non-streaming call and for a stream that failed before yielding anything. + + Returns: + The call's priced cost in USD, or None when no row was written. """ if call_usage["prompt_tokens"] == 0 and call_usage["generated_tokens"] == 0: - return + return None decoded_token = getattr(llm, "decoded_token", None) user_id = ( decoded_token.get("sub") if isinstance(decoded_token, dict) else None @@ -126,7 +130,7 @@ def _persist_call_usage(llm, call_usage, *, duration_ms=None, ttft_ms=None): "source": getattr(llm, "_token_usage_source", "agent_stream"), }, ) - return + return None model_id = getattr(llm, "_canonical_model_id", None) # Bring-your-own models run on the user's own provider key: recorded, never priced. if getattr(llm, "_is_byom", False): @@ -161,6 +165,7 @@ def _persist_call_usage(llm, call_usage, *, duration_ms=None, ttft_ms=None): ) except Exception: logger.exception("token_usage persist failed") + return cost def _call_cost_usd(model_id, call_usage) -> float: @@ -257,8 +262,10 @@ def gen_token_usage(func): usage_attachments=usage_attachments, **kwargs, ) + span = start_llm_span(self, model, stream=False, tools=tools) started_at = time.monotonic() error: BaseException | None = None + result = None try: result = func(self, model, messages, stream, tools, **kwargs) call_usage["generated_tokens"] += _count_tokens(result) @@ -268,11 +275,23 @@ def gen_token_usage(func): raise finally: duration_ms = int((time.monotonic() - started_at) * 1000) + estimated_usage = call_usage call_usage = _prefer_provider_usage(self, call_usage) self.token_usage["prompt_tokens"] += call_usage["prompt_tokens"] self.token_usage["generated_tokens"] += call_usage["generated_tokens"] # A non-streaming call has no first-token moment; ttft stays NULL. - _persist_call_usage(self, call_usage, duration_ms=duration_ms) + cost = _persist_call_usage(self, call_usage, duration_ms=duration_ms) + finish_llm_call( + span, + self, + model, + call_usage, + duration_ms=duration_ms, + error=error, + cost_usd=cost, + estimated=call_usage is estimated_usage, + output=result if isinstance(result, str) else None, + ) emit = getattr(self, "_emit_gen_finished_log", None) if callable(emit): try: @@ -314,6 +333,10 @@ def stream_token_usage(func): # non-streaming durations under one p50. provider_seconds = 0.0 error: BaseException | None = None + completed = False + # This body runs on the first ``next()``, not at ``gen_stream()`` + # time, so the span starts when the provider call really does. + span = start_llm_span(self, model, stream=True, tools=tools) try: result = func(self, model, messages, stream, tools, **kwargs) stream_iter = iter(result) @@ -323,6 +346,7 @@ def stream_token_usage(func): r = next(stream_iter) except StopIteration: provider_seconds += time.monotonic() - pull_started + completed = True break provider_seconds += time.monotonic() - pull_started if first_chunk_at is None: @@ -346,12 +370,26 @@ def stream_token_usage(func): ) for line in batch: call_usage["generated_tokens"] += _count_tokens(line) + estimated_usage = call_usage call_usage = _prefer_provider_usage(self, call_usage) self.token_usage["prompt_tokens"] += call_usage["prompt_tokens"] self.token_usage["generated_tokens"] += call_usage["generated_tokens"] - _persist_call_usage( + cost = _persist_call_usage( self, call_usage, duration_ms=duration_ms, ttft_ms=ttft_ms ) + finish_llm_call( + span, + self, + model, + call_usage, + duration_ms=duration_ms, + error=error, + completed=completed, + ttft_ms=ttft_ms, + cost_usd=cost, + estimated=call_usage is estimated_usage, + output=batch, + ) emit = getattr(self, "_emit_stream_finished_log", None) if callable(emit): try: diff --git a/docsgpt/worker.py b/docsgpt/worker.py index 9fbf57f3..f3311859 100755 --- a/docsgpt/worker.py +++ b/docsgpt/worker.py @@ -16,6 +16,7 @@ from urllib.parse import urljoin, urlsplit import requests +from docsgpt import tracing from docsgpt.core.settings import settings from docsgpt.events.publisher import publish_user_event from docsgpt.parser.chunking_creator import ChunkerCreator @@ -2129,6 +2130,7 @@ def agent_webhook_worker(self, agent_id, payload): input_data, tool_allowlist=_webhook_tool_allowlist(agent_config), endpoint="webhook", + request_id=getattr(getattr(self, "request", None), "id", None), ) result = { "answer": outcome.get("answer", ""), @@ -2942,16 +2944,36 @@ def extract_graph_worker(self, source_id, user): }, ) + trace = tracing.start_trace( + source="graph_extraction", + name=f"graph_extraction {source.get('name') or source_id}", + request_id=getattr(self.request, "id", None), + user_id=user, + ) try: - summary = extract_graph_for_source( - source_id, - user, - chunks, - config=cfg, - request_id=getattr(self.request, "id", None), - progress_cb=_progress, - ) + with tracing.activate(trace), tracing.span( + tracing.KIND_STEP, + "graph_extraction", + attributes={"docsgpt.source_id": source_id, "docsgpt.chunk_count": total}, + ) as span: + summary = extract_graph_for_source( + source_id, + user, + chunks, + config=cfg, + request_id=getattr(self.request, "id", None), + progress_cb=_progress, + ) + if isinstance(summary, dict): + span.set( + **{ + "docsgpt.graph.nodes": summary.get("nodes"), + "docsgpt.graph.edges": summary.get("edges"), + "docsgpt.graph.chunks_processed": summary.get("chunks_processed"), + } + ) except Exception as e: + tracing.flush(trace, tracing.STATUS_ERROR) _publish_graph_event( user, source_id, @@ -2960,6 +2982,7 @@ def extract_graph_worker(self, source_id, user): ) raise + tracing.flush(trace) _publish_graph_event( user, source_id, diff --git a/frontend/src/api/endpoints.ts b/frontend/src/api/endpoints.ts index 28a95dbc..7989be60 100644 --- a/frontend/src/api/endpoints.ts +++ b/frontend/src/api/endpoints.ts @@ -57,6 +57,7 @@ const endpoints = { TOOL_ANALYTICS: '/api/get_tool_analytics', SCHEDULE_ANALYTICS: '/api/get_schedule_analytics', LOGS: `/api/get_user_logs`, + TRACES: (params: string) => `/api/traces?${params}`, MANAGE_SYNC: '/api/manage_sync', SYNC_SOURCE: '/api/sync_source', REINGEST_SOURCE: '/api/sources/reingest', diff --git a/frontend/src/api/services/userService.ts b/frontend/src/api/services/userService.ts index fd9d81f6..85253b01 100644 --- a/frontend/src/api/services/userService.ts +++ b/frontend/src/api/services/userService.ts @@ -109,6 +109,17 @@ const userService = { apiClient.post(endpoints.USER.SCHEDULE_ANALYTICS, data, token), getLogs: (data: any, token: string | null): Promise => apiClient.post(endpoints.USER.LOGS, data, token), + getTraces: ( + params: Record, + token: string | null, + signal?: AbortSignal, + ): Promise => + apiClient.get( + endpoints.USER.TRACES(new URLSearchParams(params).toString()), + token, + {}, + signal, + ), manageSync: (data: any, token: string | null): Promise => apiClient.post(endpoints.USER.MANAGE_SYNC, data, token), syncSource: (data: any, token: string | null): Promise => diff --git a/frontend/src/conversation/AnswerFlow.tsx b/frontend/src/conversation/AnswerFlow.tsx index c9f32892..6fa071fa 100644 --- a/frontend/src/conversation/AnswerFlow.tsx +++ b/frontend/src/conversation/AnswerFlow.tsx @@ -333,7 +333,7 @@ function InlineToolCallChip({ ); } -function ToolCallPanel({ +export function ToolCallPanel({ title, copyText, children, diff --git a/frontend/src/locale/de.json b/frontend/src/locale/de.json index 061b5480..05e2770e 100644 --- a/frontend/src/locale/de.json +++ b/frontend/src/locale/de.json @@ -411,7 +411,9 @@ "schedule": "Geplant", "webhook": "Webhook", "workflow": "Workflow", - "system": "System" + "system": "System", + "search": "Suche", + "graph": "Graph-Aufbau" }, "detail": { "agent": "Agent", @@ -431,7 +433,112 @@ "workflow": "Workflow", "activity": "Aktivität", "steps": "Schritte", - "result": "Ergebnis" + "result": "Ergebnis", + "source": "Quelle" + }, + "trace": { + "view": "Trace anzeigen", + "viewRounds_one": "Trace anzeigen ({{count}} Runde)", + "viewRounds_other": "Trace anzeigen ({{count}} Runden)", + "title": "Ausführungs-Trace", + "empty": "Für diesen Eintrag wurde kein Trace aufgezeichnet.", + "failed": "Der Trace konnte nicht geladen werden.", + "noSpans": "Für diese Ausführung wurde nichts aufgezeichnet.", + "round": "Runde {{n}}", + "step": "Schritt", + "toggleChildren": "Verschachtelte Schritte ein- oder ausblenden", + "droppedSpans_one": "{{count}} weiterer Schritt wurde nicht aufgezeichnet (Größenlimit des Traces).", + "droppedSpans_other": "{{count}} weitere Schritte wurden nicht aufgezeichnet (Größenlimit des Traces).", + "chips": { + "llmCalls_one": "{{count}} LLM-Aufruf", + "llmCalls_other": "{{count}} LLM-Aufrufe", + "tokens": "{{input}} ein / {{output}} aus", + "toolCalls_one": "{{count}} Tool-Aufruf", + "toolCalls_other": "{{count}} Tool-Aufrufe", + "retrieval": "RAG {{duration}}", + "errors_one": "{{count}} Fehler", + "errors_other": "{{count}} Fehler" + }, + "headline": { + "tokens": "{{input}} → {{output}} Tok.", + "chunks_one": "{{count}} Abschnitt", + "chunks_other": "{{count}} Abschnitte", + "hits_one": "{{count}} Treffer", + "hits_other": "{{count}} Treffer", + "cached": "aus Cache" + }, + "status": { + "ok": "OK", + "error": "Fehler", + "cancelled": "Abgebrochen", + "paused": "Pausiert", + "pending": "Wartet auf Freigabe", + "denied": "Abgelehnt", + "skipped": "Übersprungen" + }, + "kinds": { + "agent": "Agent", + "llm": "LLM-Aufruf", + "tool": "Tool-Aufruf", + "retrieval": "Abruf", + "search": "Vektorsuche", + "embedding": "Embedding", + "rerank": "Neusortierung", + "guardrail": "Guardrail", + "step": "Schritt" + }, + "sources": { + "stream": "Chat", + "answer": "Antwort-API", + "v1": "OpenAI-kompatible API", + "webhook": "Webhook", + "schedule": "Geplante Ausführung", + "search": "Such-API", + "mcp": "MCP-Suche", + "graph_extraction": "Graph-Aufbau", + "headless": "Hintergrundausführung" + }, + "fields": { + "status": "Status", + "duration": "Dauer", + "modelTime": "Modellzeit", + "provider": "Anbieter", + "model": "Modell", + "inputTokens": "Eingabe-Tokens", + "outputTokens": "Ausgabe-Tokens", + "cachedTokens": "Gecachte Tokens", + "timeToFirstToken": "Zeit bis zum ersten Token", + "cost": "Kosten", + "tokenSource": "Zweck", + "cacheHit": "Aus dem Cache", + "estimatedUsage": "Geschätzte Tokens", + "yes": "Ja", + "tool": "Tool", + "action": "Aktion", + "callId": "Aufruf-ID", + "retriever": "Retriever", + "sources": "Quellen", + "topK": "Top K", + "chunks": "Abschnitte", + "candidates": "Kandidaten", + "topScore": "Beste Bewertung", + "vectorStore": "Vektorspeicher", + "agentType": "Agententyp", + "toolCalls": "Tool-Aufrufe", + "stage": "Phase", + "checks": "Prüfungen", + "triggered": "Ausgelöst", + "blocked": "Blockiert", + "redacted": "Geschwärzt", + "nodeType": "Knotentyp", + "phase": "Abschnitt", + "query": "Anfrage", + "arguments": "Argumente", + "result": "Ergebnis", + "output": "Ausgabe", + "retrievedChunks": "Abgerufene Abschnitte", + "allAttributes": "Alle Attribute" + } } }, "teams": { diff --git a/frontend/src/locale/en.json b/frontend/src/locale/en.json index 665e746c..9bb4018f 100644 --- a/frontend/src/locale/en.json +++ b/frontend/src/locale/en.json @@ -416,7 +416,9 @@ "schedule": "Scheduled", "webhook": "Webhook", "workflow": "Workflow", - "system": "System" + "system": "System", + "search": "Search", + "graph": "Graph build" }, "detail": { "agent": "Agent", @@ -436,7 +438,112 @@ "workflow": "Workflow", "activity": "Activity", "steps": "Steps", - "result": "Result" + "result": "Result", + "source": "Source" + }, + "trace": { + "view": "View trace", + "viewRounds_one": "View trace ({{count}} round)", + "viewRounds_other": "View trace ({{count}} rounds)", + "title": "Execution trace", + "empty": "No trace was recorded for this entry.", + "failed": "Could not load the trace.", + "noSpans": "Nothing was recorded for this run.", + "round": "Round {{n}}", + "step": "Step", + "toggleChildren": "Show or hide nested steps", + "droppedSpans_one": "{{count}} more step was not recorded (trace size limit).", + "droppedSpans_other": "{{count}} more steps were not recorded (trace size limit).", + "chips": { + "llmCalls_one": "{{count}} LLM call", + "llmCalls_other": "{{count}} LLM calls", + "tokens": "{{input}} in / {{output}} out", + "toolCalls_one": "{{count}} tool call", + "toolCalls_other": "{{count}} tool calls", + "retrieval": "RAG {{duration}}", + "errors_one": "{{count}} error", + "errors_other": "{{count}} errors" + }, + "headline": { + "tokens": "{{input}} → {{output}} tok", + "chunks_one": "{{count}} chunk", + "chunks_other": "{{count}} chunks", + "hits_one": "{{count}} hit", + "hits_other": "{{count}} hits", + "cached": "cached" + }, + "status": { + "ok": "OK", + "error": "Error", + "cancelled": "Cancelled", + "paused": "Paused", + "pending": "Awaiting approval", + "denied": "Denied", + "skipped": "Skipped" + }, + "kinds": { + "agent": "Agent", + "llm": "LLM call", + "tool": "Tool call", + "retrieval": "Retrieval", + "search": "Vector search", + "embedding": "Embedding", + "rerank": "Rerank", + "guardrail": "Guardrail", + "step": "Step" + }, + "sources": { + "stream": "Chat", + "answer": "Answer API", + "v1": "OpenAI-compatible API", + "webhook": "Webhook", + "schedule": "Scheduled run", + "search": "Search API", + "mcp": "MCP search", + "graph_extraction": "Graph build", + "headless": "Background run" + }, + "fields": { + "status": "Status", + "duration": "Duration", + "modelTime": "Model time", + "provider": "Provider", + "model": "Model", + "inputTokens": "Input tokens", + "outputTokens": "Output tokens", + "cachedTokens": "Cached tokens", + "timeToFirstToken": "Time to first token", + "cost": "Cost", + "tokenSource": "Purpose", + "cacheHit": "Served from cache", + "estimatedUsage": "Estimated tokens", + "yes": "Yes", + "tool": "Tool", + "action": "Action", + "callId": "Call ID", + "retriever": "Retriever", + "sources": "Sources", + "topK": "Top K", + "chunks": "Chunks", + "candidates": "Candidates", + "topScore": "Top score", + "vectorStore": "Vector store", + "agentType": "Agent type", + "toolCalls": "Tool calls", + "stage": "Stage", + "checks": "Checks", + "triggered": "Triggered", + "blocked": "Blocked", + "redacted": "Redacted", + "nodeType": "Node type", + "phase": "Phase", + "query": "Query", + "arguments": "Arguments", + "result": "Result", + "output": "Output", + "retrievedChunks": "Retrieved chunks", + "allAttributes": "All attributes" + } } }, "teams": { diff --git a/frontend/src/locale/es.json b/frontend/src/locale/es.json index 1c27a945..e5d9617c 100644 --- a/frontend/src/locale/es.json +++ b/frontend/src/locale/es.json @@ -411,7 +411,9 @@ "schedule": "Programado", "webhook": "Webhook", "workflow": "Flujo de trabajo", - "system": "Sistema" + "system": "Sistema", + "search": "Búsqueda", + "graph": "Creación de grafo" }, "detail": { "agent": "Agente", @@ -431,7 +433,112 @@ "workflow": "Flujo de trabajo", "activity": "Actividad", "steps": "Pasos", - "result": "Resultado" + "result": "Resultado", + "source": "Origen" + }, + "trace": { + "view": "Ver traza", + "viewRounds_one": "Ver traza ({{count}} ronda)", + "viewRounds_other": "Ver traza ({{count}} rondas)", + "title": "Traza de ejecución", + "empty": "No se registró ninguna traza para esta entrada.", + "failed": "No se pudo cargar la traza.", + "noSpans": "No se registró nada en esta ejecución.", + "round": "Ronda {{n}}", + "step": "Paso", + "toggleChildren": "Mostrar u ocultar pasos anidados", + "droppedSpans_one": "{{count}} paso más no se registró (límite de tamaño de la traza).", + "droppedSpans_other": "{{count}} pasos más no se registraron (límite de tamaño de la traza).", + "chips": { + "llmCalls_one": "{{count}} llamada al LLM", + "llmCalls_other": "{{count}} llamadas al LLM", + "tokens": "{{input}} entrada / {{output}} salida", + "toolCalls_one": "{{count}} llamada a herramienta", + "toolCalls_other": "{{count}} llamadas a herramientas", + "retrieval": "RAG {{duration}}", + "errors_one": "{{count}} error", + "errors_other": "{{count}} errores" + }, + "headline": { + "tokens": "{{input}} → {{output}} tok", + "chunks_one": "{{count}} fragmento", + "chunks_other": "{{count}} fragmentos", + "hits_one": "{{count}} resultado", + "hits_other": "{{count}} resultados", + "cached": "en caché" + }, + "status": { + "ok": "OK", + "error": "Error", + "cancelled": "Cancelado", + "paused": "En pausa", + "pending": "Esperando aprobación", + "denied": "Denegado", + "skipped": "Omitido" + }, + "kinds": { + "agent": "Agente", + "llm": "Llamada al LLM", + "tool": "Llamada a herramienta", + "retrieval": "Recuperación", + "search": "Búsqueda vectorial", + "embedding": "Embedding", + "rerank": "Reordenación", + "guardrail": "Guardrail", + "step": "Paso" + }, + "sources": { + "stream": "Chat", + "answer": "API de respuestas", + "v1": "API compatible con OpenAI", + "webhook": "Webhook", + "schedule": "Ejecución programada", + "search": "API de búsqueda", + "mcp": "Búsqueda MCP", + "graph_extraction": "Creación de grafo", + "headless": "Ejecución en segundo plano" + }, + "fields": { + "status": "Estado", + "duration": "Duración", + "modelTime": "Tiempo del modelo", + "provider": "Proveedor", + "model": "Modelo", + "inputTokens": "Tokens de entrada", + "outputTokens": "Tokens de salida", + "cachedTokens": "Tokens en caché", + "timeToFirstToken": "Tiempo hasta el primer token", + "cost": "Coste", + "tokenSource": "Propósito", + "cacheHit": "Servido desde caché", + "estimatedUsage": "Tokens estimados", + "yes": "Sí", + "tool": "Herramienta", + "action": "Acción", + "callId": "ID de llamada", + "retriever": "Recuperador", + "sources": "Fuentes", + "topK": "Top K", + "chunks": "Fragmentos", + "candidates": "Candidatos", + "topScore": "Mejor puntuación", + "vectorStore": "Almacén vectorial", + "agentType": "Tipo de agente", + "toolCalls": "Llamadas a herramientas", + "stage": "Etapa", + "checks": "Comprobaciones", + "triggered": "Activadas", + "blocked": "Bloqueado", + "redacted": "Censurado", + "nodeType": "Tipo de nodo", + "phase": "Fase", + "query": "Consulta", + "arguments": "Argumentos", + "result": "Resultado", + "output": "Salida", + "retrievedChunks": "Fragmentos recuperados", + "allAttributes": "Todos los atributos" + } } }, "teams": { diff --git a/frontend/src/locale/jp.json b/frontend/src/locale/jp.json index 537a6373..e0f4b035 100644 --- a/frontend/src/locale/jp.json +++ b/frontend/src/locale/jp.json @@ -410,7 +410,9 @@ "schedule": "スケジュール", "webhook": "Webhook", "workflow": "ワークフロー", - "system": "システム" + "system": "システム", + "search": "検索", + "graph": "グラフ構築" }, "detail": { "agent": "エージェント", @@ -430,7 +432,112 @@ "workflow": "ワークフロー", "activity": "アクティビティ", "steps": "ステップ", - "result": "結果" + "result": "結果", + "source": "ソース" + }, + "trace": { + "view": "トレースを表示", + "viewRounds_one": "トレースを表示（{{count}} ラウンド）", + "viewRounds_other": "トレースを表示（{{count}} ラウンド）", + "title": "実行トレース", + "empty": "このエントリのトレースは記録されていません。", + "failed": "トレースを読み込めませんでした。", + "noSpans": "この実行では何も記録されませんでした。", + "round": "ラウンド {{n}}", + "step": "ステップ", + "toggleChildren": "ネストされたステップの表示を切り替え", + "droppedSpans_one": "さらに {{count}} 個のステップは記録されていません（トレースのサイズ上限）。", + "droppedSpans_other": "さらに {{count}} 個のステップは記録されていません（トレースのサイズ上限）。", + "chips": { + "llmCalls_one": "LLM 呼び出し {{count}} 回", + "llmCalls_other": "LLM 呼び出し {{count}} 回", + "tokens": "入力 {{input}} / 出力 {{output}}", + "toolCalls_one": "ツール呼び出し {{count}} 回", + "toolCalls_other": "ツール呼び出し {{count}} 回", + "retrieval": "RAG {{duration}}", + "errors_one": "エラー {{count}} 件", + "errors_other": "エラー {{count}} 件" + }, + "headline": { + "tokens": "{{input}} → {{output}} トークン", + "chunks_one": "{{count}} チャンク", + "chunks_other": "{{count}} チャンク", + "hits_one": "{{count}} 件ヒット", + "hits_other": "{{count}} 件ヒット", + "cached": "キャッシュ" + }, + "status": { + "ok": "OK", + "error": "エラー", + "cancelled": "キャンセル", + "paused": "一時停止", + "pending": "承認待ち", + "denied": "拒否", + "skipped": "スキップ" + }, + "kinds": { + "agent": "エージェント", + "llm": "LLM 呼び出し", + "tool": "ツール呼び出し", + "retrieval": "検索", + "search": "ベクトル検索", + "embedding": "埋め込み", + "rerank": "再ランク付け", + "guardrail": "ガードレール", + "step": "ステップ" + }, + "sources": { + "stream": "チャット", + "answer": "回答 API", + "v1": "OpenAI 互換 API", + "webhook": "Webhook", + "schedule": "スケジュール実行", + "search": "検索 API", + "mcp": "MCP 検索", + "graph_extraction": "グラフ構築", + "headless": "バックグラウンド実行" + }, + "fields": { + "status": "ステータス", + "duration": "所要時間", + "modelTime": "モデル処理時間", + "provider": "プロバイダー", + "model": "モデル", + "inputTokens": "入力トークン", + "outputTokens": "出力トークン", + "cachedTokens": "キャッシュされたトークン", + "timeToFirstToken": "最初のトークンまでの時間", + "cost": "コスト", + "tokenSource": "用途", + "cacheHit": "キャッシュから提供", + "estimatedUsage": "推定トークン", + "yes": "はい", + "tool": "ツール", + "action": "アクション", + "callId": "呼び出し ID", + "retriever": "リトリーバー", + "sources": "ソース", + "topK": "Top K", + "chunks": "チャンク", + "candidates": "候補", + "topScore": "最高スコア", + "vectorStore": "ベクトルストア", + "agentType": "エージェントの種類", + "toolCalls": "ツール呼び出し", + "stage": "ステージ", + "checks": "チェック", + "triggered": "発動", + "blocked": "ブロック", + "redacted": "マスク済み", + "nodeType": "ノードの種類", + "phase": "フェーズ", + "query": "クエリ", + "arguments": "引数", + "result": "結果", + "output": "出力", + "retrievedChunks": "取得したチャンク", + "allAttributes": "すべての属性" + } } }, "teams": { diff --git a/frontend/src/locale/logsTrace.test.ts b/frontend/src/locale/logsTrace.test.ts new file mode 100644 index 00000000..98f36f1a --- /dev/null +++ b/frontend/src/locale/logsTrace.test.ts @@ -0,0 +1,63 @@ +import { describe, expect, it } from 'vitest'; + +import de from './de.json'; +import en from './en.json'; +import es from './es.json'; +import jp from './jp.json'; +import ru from './ru.json'; +import zhTW from './zh-TW.json'; +import zh from './zh.json'; + +type Tree = { [key: string]: string | Tree }; + +const PLURAL_SUFFIX = /_(zero|one|two|few|many|other)$/; + +const flatten = (tree: Tree, prefix = ''): string[] => + Object.entries(tree).flatMap(([key, value]) => + typeof value === 'string' + ? [prefix + key.replace(PLURAL_SUFFIX, '')] + : flatten(value, `${prefix}${key}.`), + ); + +const traceBlock = (locale: { settings: object }): Tree => + ((locale.settings as Tree).logs as Tree).trace as Tree; + +const keysOf = (locale: { settings: object }): string[] => + Array.from(new Set(flatten(traceBlock(locale)))).sort(); + +const LOCALES = { es, de, jp, ru, zh, zhTW }; + +describe('settings.logs.trace locale block', () => { + it.each(Object.entries(LOCALES))( + '%s has the same keys as en', + (_name, locale) => { + expect(keysOf(locale)).toEqual(keysOf(en)); + }, + ); + + it.each(Object.entries(LOCALES))( + '%s is translated, not an English copy', + (_name, locale) => { + expect(traceBlock(locale).title).not.toBe(traceBlock(en).title); + expect((traceBlock(locale).fields as Tree).duration).not.toBe( + (traceBlock(en).fields as Tree).duration, + ); + }, + ); + + it('every locale provides both plural forms i18next falls back through', () => { + [en, ...Object.values(LOCALES)].forEach((locale) => { + const chips = traceBlock(locale).chips as Tree; + expect(chips.llmCalls_one).toBeTruthy(); + expect(chips.llmCalls_other).toBeTruthy(); + }); + }); + + it('every locale labels the new log event types', () => { + [en, ...Object.values(LOCALES)].forEach((locale) => { + const types = ((locale.settings as Tree).logs as Tree).types as Tree; + expect(types.search).toBeTruthy(); + expect(types.graph).toBeTruthy(); + }); + }); +}); diff --git a/frontend/src/locale/ru.json b/frontend/src/locale/ru.json index 243954cb..ac6233ce 100644 --- a/frontend/src/locale/ru.json +++ b/frontend/src/locale/ru.json @@ -423,7 +423,9 @@ "schedule": "По расписанию", "webhook": "Вебхук", "workflow": "Рабочий процесс", - "system": "Система" + "system": "Система", + "search": "Поиск", + "graph": "Построение графа" }, "detail": { "agent": "Агент", @@ -443,7 +445,126 @@ "workflow": "Рабочий процесс", "activity": "Активность", "steps": "Шаги", - "result": "Результат" + "result": "Результат", + "source": "Источник" + }, + "trace": { + "view": "Открыть трассировку", + "viewRounds_one": "Открыть трассировку ({{count}} раунд)", + "viewRounds_few": "Открыть трассировку ({{count}} раунда)", + "viewRounds_many": "Открыть трассировку ({{count}} раундов)", + "viewRounds_other": "Открыть трассировку ({{count}} раунда)", + "title": "Трассировка выполнения", + "empty": "Для этой записи трассировка не сохранена.", + "failed": "Не удалось загрузить трассировку.", + "noSpans": "Для этого запуска ничего не записано.", + "round": "Раунд {{n}}", + "step": "Шаг", + "toggleChildren": "Показать или скрыть вложенные шаги", + "droppedSpans_one": "Ещё {{count}} шаг не записан (лимит размера трассировки).", + "droppedSpans_few": "Ещё {{count}} шага не записаны (лимит размера трассировки).", + "droppedSpans_many": "Ещё {{count}} шагов не записаны (лимит размера трассировки).", + "droppedSpans_other": "Ещё {{count}} шага не записаны (лимит размера трассировки).", + "chips": { + "llmCalls_one": "{{count}} вызов LLM", + "llmCalls_few": "{{count}} вызова LLM", + "llmCalls_many": "{{count}} вызовов LLM", + "llmCalls_other": "{{count}} вызова LLM", + "tokens": "{{input}} вход / {{output}} выход", + "toolCalls_one": "{{count}} вызов инструмента", + "toolCalls_few": "{{count}} вызова инструментов", + "toolCalls_many": "{{count}} вызовов инструментов", + "toolCalls_other": "{{count}} вызова инструментов", + "retrieval": "RAG {{duration}}", + "errors_one": "{{count}} ошибка", + "errors_few": "{{count}} ошибки", + "errors_many": "{{count}} ошибок", + "errors_other": "{{count}} ошибки" + }, + "headline": { + "tokens": "{{input}} → {{output}} ток.", + "chunks_one": "{{count}} фрагмент", + "chunks_few": "{{count}} фрагмента", + "chunks_many": "{{count}} фрагментов", + "chunks_other": "{{count}} фрагмента", + "hits_one": "{{count}} совпадение", + "hits_few": "{{count}} совпадения", + "hits_many": "{{count}} совпадений", + "hits_other": "{{count}} совпадения", + "cached": "из кэша" + }, + "status": { + "ok": "ОК", + "error": "Ошибка", + "cancelled": "Отменено", + "paused": "Приостановлено", + "pending": "Ожидает подтверждения", + "denied": "Отклонено", + "skipped": "Пропущено" + }, + "kinds": { + "agent": "Агент", + "llm": "Вызов LLM", + "tool": "Вызов инструмента", + "retrieval": "Поиск", + "search": "Векторный поиск", + "embedding": "Эмбеддинг", + "rerank": "Переранжирование", + "guardrail": "Ограничитель", + "step": "Шаг" + }, + "sources": { + "stream": "Чат", + "answer": "API ответов", + "v1": "OpenAI-совместимый API", + "webhook": "Вебхук", + "schedule": "Запуск по расписанию", + "search": "API поиска", + "mcp": "Поиск MCP", + "graph_extraction": "Построение графа", + "headless": "Фоновый запуск" + }, + "fields": { + "status": "Статус", + "duration": "Длительность", + "modelTime": "Время модели", + "provider": "Провайдер", + "model": "Модель", + "inputTokens": "Входные токены", + "outputTokens": "Выходные токены", + "cachedTokens": "Токены из кэша", + "timeToFirstToken": "Время до первого токена", + "cost": "Стоимость", + "tokenSource": "Назначение", + "cacheHit": "Из кэша", + "estimatedUsage": "Оценочные токены", + "yes": "Да", + "tool": "Инструмент", + "action": "Действие", + "callId": "ID вызова", + "retriever": "Ретривер", + "sources": "Источники", + "topK": "Top K", + "chunks": "Фрагменты", + "candidates": "Кандидаты", + "topScore": "Лучшая оценка", + "vectorStore": "Векторное хранилище", + "agentType": "Тип агента", + "toolCalls": "Вызовы инструментов", + "stage": "Этап", + "checks": "Проверки", + "triggered": "Сработали", + "blocked": "Заблокировано", + "redacted": "Скрыто", + "nodeType": "Тип узла", + "phase": "Фаза", + "query": "Запрос", + "arguments": "Аргументы", + "result": "Результат", + "output": "Вывод", + "retrievedChunks": "Найденные фрагменты", + "allAttributes": "Все атрибуты" + } } }, "teams": { diff --git a/frontend/src/locale/zh-TW.json b/frontend/src/locale/zh-TW.json index 3e10c564..beb8acf7 100644 --- a/frontend/src/locale/zh-TW.json +++ b/frontend/src/locale/zh-TW.json @@ -410,7 +410,9 @@ "schedule": "排程", "webhook": "Webhook", "workflow": "工作流程", - "system": "系統" + "system": "系統", + "search": "搜尋", + "graph": "圖譜建置" }, "detail": { "agent": "代理", @@ -430,7 +432,112 @@ "workflow": "工作流程", "activity": "活動", "steps": "步驟", - "result": "結果" + "result": "結果", + "source": "來源" + }, + "trace": { + "view": "檢視追蹤", + "viewRounds_one": "檢視追蹤（{{count}} 輪）", + "viewRounds_other": "檢視追蹤（{{count}} 輪）", + "title": "執行追蹤", + "empty": "此項目沒有記錄追蹤。", + "failed": "無法載入追蹤。", + "noSpans": "此次執行沒有記錄任何內容。", + "round": "第 {{n}} 輪", + "step": "步驟", + "toggleChildren": "顯示或隱藏巢狀步驟", + "droppedSpans_one": "另有 {{count}} 個步驟未被記錄（追蹤大小限制）。", + "droppedSpans_other": "另有 {{count}} 個步驟未被記錄（追蹤大小限制）。", + "chips": { + "llmCalls_one": "{{count}} 次 LLM 呼叫", + "llmCalls_other": "{{count}} 次 LLM 呼叫", + "tokens": "輸入 {{input}} / 輸出 {{output}}", + "toolCalls_one": "{{count}} 次工具呼叫", + "toolCalls_other": "{{count}} 次工具呼叫", + "retrieval": "RAG {{duration}}", + "errors_one": "{{count}} 個錯誤", + "errors_other": "{{count}} 個錯誤" + }, + "headline": { + "tokens": "{{input}} → {{output}} 詞元", + "chunks_one": "{{count}} 個片段", + "chunks_other": "{{count}} 個片段", + "hits_one": "{{count}} 個命中", + "hits_other": "{{count}} 個命中", + "cached": "快取" + }, + "status": { + "ok": "正常", + "error": "錯誤", + "cancelled": "已取消", + "paused": "已暫停", + "pending": "等待核准", + "denied": "已拒絕", + "skipped": "已略過" + }, + "kinds": { + "agent": "代理", + "llm": "LLM 呼叫", + "tool": "工具呼叫", + "retrieval": "檢索", + "search": "向量搜尋", + "embedding": "嵌入", + "rerank": "重新排序", + "guardrail": "防護機制", + "step": "步驟" + }, + "sources": { + "stream": "聊天", + "answer": "回答 API", + "v1": "OpenAI 相容 API", + "webhook": "Webhook", + "schedule": "排程執行", + "search": "搜尋 API", + "mcp": "MCP 搜尋", + "graph_extraction": "圖譜建置", + "headless": "背景執行" + }, + "fields": { + "status": "狀態", + "duration": "耗時", + "modelTime": "模型耗時", + "provider": "供應商", + "model": "模型", + "inputTokens": "輸入詞元", + "outputTokens": "輸出詞元", + "cachedTokens": "快取詞元", + "timeToFirstToken": "首個詞元耗時", + "cost": "費用", + "tokenSource": "用途", + "cacheHit": "來自快取", + "estimatedUsage": "估算詞元", + "yes": "是", + "tool": "工具", + "action": "動作", + "callId": "呼叫 ID", + "retriever": "檢索器", + "sources": "來源", + "topK": "Top K", + "chunks": "片段", + "candidates": "候選", + "topScore": "最高分", + "vectorStore": "向量儲存", + "agentType": "代理類型", + "toolCalls": "工具呼叫", + "stage": "階段", + "checks": "檢查項目", + "triggered": "已觸發", + "blocked": "已封鎖", + "redacted": "已遮蔽", + "nodeType": "節點類型", + "phase": "階段", + "query": "查詢", + "arguments": "參數", + "result": "結果", + "output": "輸出", + "retrievedChunks": "檢索到的片段", + "allAttributes": "所有屬性" + } } }, "teams": { diff --git a/frontend/src/locale/zh.json b/frontend/src/locale/zh.json index 201b6c4c..540c40a4 100644 --- a/frontend/src/locale/zh.json +++ b/frontend/src/locale/zh.json @@ -410,7 +410,9 @@ "schedule": "计划", "webhook": "Webhook", "workflow": "工作流", - "system": "系统" + "system": "系统", + "search": "搜索", + "graph": "图谱构建" }, "detail": { "agent": "代理", @@ -430,7 +432,112 @@ "workflow": "工作流", "activity": "活动", "steps": "步骤", - "result": "结果" + "result": "结果", + "source": "来源" + }, + "trace": { + "view": "查看追踪", + "viewRounds_one": "查看追踪（{{count}} 轮）", + "viewRounds_other": "查看追踪（{{count}} 轮）", + "title": "执行追踪", + "empty": "此条目没有记录追踪。", + "failed": "无法加载追踪。", + "noSpans": "此次运行没有记录任何内容。", + "round": "第 {{n}} 轮", + "step": "步骤", + "toggleChildren": "显示或隐藏嵌套步骤", + "droppedSpans_one": "另有 {{count}} 个步骤未被记录（追踪大小限制）。", + "droppedSpans_other": "另有 {{count}} 个步骤未被记录（追踪大小限制）。", + "chips": { + "llmCalls_one": "{{count}} 次 LLM 调用", + "llmCalls_other": "{{count}} 次 LLM 调用", + "tokens": "输入 {{input}} / 输出 {{output}}", + "toolCalls_one": "{{count}} 次工具调用", + "toolCalls_other": "{{count}} 次工具调用", + "retrieval": "RAG {{duration}}", + "errors_one": "{{count}} 个错误", + "errors_other": "{{count}} 个错误" + }, + "headline": { + "tokens": "{{input}} → {{output}} 词元", + "chunks_one": "{{count}} 个片段", + "chunks_other": "{{count}} 个片段", + "hits_one": "{{count}} 个命中", + "hits_other": "{{count}} 个命中", + "cached": "缓存" + }, + "status": { + "ok": "正常", + "error": "错误", + "cancelled": "已取消", + "paused": "已暂停", + "pending": "等待批准", + "denied": "已拒绝", + "skipped": "已跳过" + }, + "kinds": { + "agent": "智能体", + "llm": "LLM 调用", + "tool": "工具调用", + "retrieval": "检索", + "search": "向量搜索", + "embedding": "嵌入", + "rerank": "重排序", + "guardrail": "防护栏", + "step": "步骤" + }, + "sources": { + "stream": "聊天", + "answer": "回答 API", + "v1": "OpenAI 兼容 API", + "webhook": "Webhook", + "schedule": "定时运行", + "search": "搜索 API", + "mcp": "MCP 搜索", + "graph_extraction": "图谱构建", + "headless": "后台运行" + }, + "fields": { + "status": "状态", + "duration": "耗时", + "modelTime": "模型耗时", + "provider": "提供商", + "model": "模型", + "inputTokens": "输入词元", + "outputTokens": "输出词元", + "cachedTokens": "缓存词元", + "timeToFirstToken": "首个词元耗时", + "cost": "费用", + "tokenSource": "用途", + "cacheHit": "来自缓存", + "estimatedUsage": "估算词元", + "yes": "是", + "tool": "工具", + "action": "操作", + "callId": "调用 ID", + "retriever": "检索器", + "sources": "来源", + "topK": "Top K", + "chunks": "片段", + "candidates": "候选", + "topScore": "最高分", + "vectorStore": "向量存储", + "agentType": "智能体类型", + "toolCalls": "工具调用", + "stage": "阶段", + "checks": "检查项", + "triggered": "已触发", + "blocked": "已拦截", + "redacted": "已脱敏", + "nodeType": "节点类型", + "phase": "阶段", + "query": "查询", + "arguments": "参数", + "result": "结果", + "output": "输出", + "retrievedChunks": "检索到的片段", + "allAttributes": "全部属性" + } } }, "teams": { diff --git a/frontend/src/settings/Logs.tsx b/frontend/src/settings/Logs.tsx index b5b80a9b..16a0b247 100644 --- a/frontend/src/settings/Logs.tsx +++ b/frontend/src/settings/Logs.tsx @@ -1,3 +1,4 @@ +import { Activity } from 'lucide-react'; import React, { useCallback, useEffect, useRef, useState } from 'react'; import { useTranslation } from 'react-i18next'; import { useSelector } from 'react-redux'; @@ -6,6 +7,7 @@ import userService from '../api/services/userService'; import ChevronRight from '../assets/chevron-right.svg'; import CopyButton from '../components/CopyButton'; import SkeletonLoader from '../components/SkeletonLoader'; +import { Button } from '../components/ui/button'; import { Input } from '../components/ui/input'; import { Select, @@ -16,7 +18,10 @@ import { } from '../components/ui/select'; import { useLoaderState } from '../hooks'; import { selectToken } from '../preferences/preferenceSlice'; -import { LogData } from './types'; +import TraceChips from './traces/TraceChips'; +import TraceSheet from './traces/TraceSheet'; +import { formatDurationMs } from './traces/traceUtils'; +import { LogData, TraceRef } from './types'; type LogsProps = { agentId?: string; @@ -35,6 +40,7 @@ export default function Logs({ agentId, tableHeader }: LogsProps) { const [typeFilter, setTypeFilter] = useState('all'); const [searchInput, setSearchInput] = useState(''); const [search, setSearch] = useState(''); + const [openTrace, setOpenTrace] = useState(null); const logs = Object.values(logsByPage).flat(); @@ -127,6 +133,8 @@ export default function Logs({ agentId, tableHeader }: LogsProps) { { label: t('settings.logs.types.webhook'), value: 'webhook' }, { label: t('settings.logs.types.workflow'), value: 'workflow' }, { label: t('settings.logs.types.system'), value: 'system' }, + { label: t('settings.logs.types.search'), value: 'search' }, + { label: t('settings.logs.types.graph'), value: 'graph' }, ]; return ( @@ -178,8 +186,14 @@ export default function Logs({ agentId, tableHeader }: LogsProps) { setPage={setPage} loading={loadingLogs} tableHeader={tableHeader} + onViewTrace={setOpenTrace} /> + setOpenTrace(null)} + /> ); } @@ -189,8 +203,15 @@ type LogsTableProps = { setPage: React.Dispatch>; loading: boolean; tableHeader?: string; + onViewTrace: (ref: TraceRef) => void; }; -function LogsTable({ logs, setPage, loading, tableHeader }: LogsTableProps) { +function LogsTable({ + logs, + setPage, + loading, + tableHeader, + onViewTrace, +}: LogsTableProps) { const { t } = useTranslation(); const observerRef = useRef(null); const [openLogId, setOpenLogId] = useState(null); @@ -248,6 +269,7 @@ function LogsTable({ logs, setPage, loading, tableHeader }: LogsTableProps) { log={log} isOpen={openLogId === log.id} onToggle={handleLogToggle} + onViewTrace={onViewTrace} /> ); @@ -258,6 +280,7 @@ function LogsTable({ logs, setPage, loading, tableHeader }: LogsTableProps) { log={log} isOpen={openLogId === log.id} onToggle={handleLogToggle} + onViewTrace={onViewTrace} /> ); })} @@ -271,17 +294,19 @@ function formatDuration(start?: string, end?: string): string | null { if (!start || !end) return null; const ms = new Date(end).getTime() - new Date(start).getTime(); if (isNaN(ms) || ms < 0) return null; - return ms < 1000 ? `${ms}ms` : `${(ms / 1000).toFixed(1)}s`; + return formatDurationMs(ms); } function Log({ log, isOpen, onToggle, + onViewTrace, }: { log: LogData; isOpen: boolean; onToggle: (id: string) => void; + onViewTrace: (ref: TraceRef) => void; }) { const { t } = useTranslation(); const logLevelColor = { @@ -336,6 +361,17 @@ function Log({ } else if (log.event_type === 'system' || log.event_type === 'webhook') { if (log.endpoint) detailRows.push([t('settings.logs.detail.endpoint'), log.endpoint]); + } else if (log.event_type === 'search' || log.event_type === 'graph') { + if (log.source) + detailRows.push([ + t('settings.logs.detail.source'), + t(`settings.logs.trace.sources.${log.source}`, log.source), + ]); + if (log.status) + detailRows.push([ + t('settings.logs.detail.status'), + t(`settings.logs.trace.status.${log.status}`, log.status), + ]); } const textBlocks: { label: string; text: string; isError?: boolean }[] = []; @@ -394,8 +430,17 @@ function Log({ return (
onToggle(log.id)} - className={`text-foreground flex cursor-pointer flex-row items-start gap-2 p-2 px-4 py-3 ${ + onKeyDown={(e) => { + if (e.key === 'Enter' || e.key === ' ') { + e.preventDefault(); + onToggle(log.id); + } + }} + className={`text-foreground focus-visible:ring-ring/50 flex cursor-pointer flex-row items-start gap-2 p-2 px-4 py-3 outline-none focus-visible:ring-3 focus-visible:ring-inset ${ isOpen ? 'dark:bg-background rounded-t-xl bg-[#F1F1F1]' : '' }`} > @@ -412,6 +457,11 @@ function Log({ )}

{`[${log.action}]`}

+ {log.trace && ( +

+ {formatDurationMs(log.trace.duration_ms)} +

+ )}

@@ -423,6 +473,26 @@ function Log({

{isOpen && (
+ {log.trace && ( +
+ + +
+ )} {detailRows.length > 0 && (
{detailRows.map(([label, value]) => ( diff --git a/frontend/src/settings/traces/TraceChips.tsx b/frontend/src/settings/traces/TraceChips.tsx new file mode 100644 index 00000000..d1240a57 --- /dev/null +++ b/frontend/src/settings/traces/TraceChips.tsx @@ -0,0 +1,93 @@ +import { + Brain, + Coins, + Search, + Timer, + TriangleAlert, + Wrench, +} from 'lucide-react'; +import React from 'react'; +import { useTranslation } from 'react-i18next'; + +import { TraceCounts } from '../types'; +import { formatDurationMs, formatTokens } from './traceUtils'; + +type TraceChipsProps = { + durationMs?: number; + counts: TraceCounts; +}; + +/** Compact stat chips for a trace: duration, LLM calls, tokens, tools, RAG time, errors. */ +export default function TraceChips({ durationMs, counts }: TraceChipsProps) { + const { t } = useTranslation(); + const chips: { + key: string; + icon: React.ReactNode; + label: string; + tone?: 'danger'; + }[] = []; + if (durationMs !== undefined) + chips.push({ + key: 'duration', + icon: , + label: formatDurationMs(durationMs), + }); + if (counts.llm_calls) + chips.push({ + key: 'llm', + icon: , + label: t('settings.logs.trace.chips.llmCalls', { + count: counts.llm_calls, + }), + }); + if (counts.input_tokens || counts.output_tokens) + chips.push({ + key: 'tokens', + icon: , + label: t('settings.logs.trace.chips.tokens', { + input: formatTokens(counts.input_tokens), + output: formatTokens(counts.output_tokens), + }), + }); + if (counts.tool_calls) + chips.push({ + key: 'tools', + icon: , + label: t('settings.logs.trace.chips.toolCalls', { + count: counts.tool_calls, + }), + }); + if (counts.retrieval_calls) + chips.push({ + key: 'retrieval', + icon: , + label: t('settings.logs.trace.chips.retrieval', { + duration: formatDurationMs(counts.retrieval_ms), + }), + }); + if (counts.errors) + chips.push({ + key: 'errors', + icon: , + label: t('settings.logs.trace.chips.errors', { count: counts.errors }), + tone: 'danger', + }); + if (!chips.length) return null; + return ( +
+ {chips.map((chip) => ( + + {chip.icon} + {chip.label} + + ))} +
+ ); +} diff --git a/frontend/src/settings/traces/TraceSheet.tsx b/frontend/src/settings/traces/TraceSheet.tsx new file mode 100644 index 00000000..eb7e2385 --- /dev/null +++ b/frontend/src/settings/traces/TraceSheet.tsx @@ -0,0 +1,161 @@ +import React, { useEffect, useRef, useState } from 'react'; +import { useTranslation } from 'react-i18next'; +import { useSelector } from 'react-redux'; + +import userService from '../../api/services/userService'; +import Spinner from '../../components/Spinner'; +import { Sheet, SheetContent } from '../../components/ui/sheet'; +import { selectToken } from '../../preferences/preferenceSlice'; +import { formatDateTime } from '../../utils/dateTimeUtils'; +import { Trace, TraceRef } from '../types'; +import TraceChips from './TraceChips'; +import TraceWaterfall from './TraceWaterfall'; + +type TraceSheetProps = { + traceRef: TraceRef | null; + agentId?: string; + onClose: () => void; +}; + +/** Status pill tones. The theme has no success or warning token yet. */ +const STATUS_TONE: Record = { + ok: 'bg-muted text-foreground', + error: 'bg-destructive/10 text-destructive', + paused: 'bg-primary/10 text-primary', + cancelled: 'bg-muted text-muted-foreground', +}; + +/** + * Side panel showing the execution trace(s) behind one Logs row. A chat turn + * paused for tool approval has one trace per round, shown in order. + */ +export default function TraceSheet({ + traceRef, + agentId, + onClose, +}: TraceSheetProps) { + const { t } = useTranslation(); + const token = useSelector(selectToken); + const [traces, setTraces] = useState([]); + const [loading, setLoading] = useState(false); + const [failed, setFailed] = useState(false); + // The sheet is opened from a Logs row, not a SheetTrigger, so Radix has no + // trigger to return focus to on close; remember what had focus instead. + const returnFocusRef = useRef(null); + + useEffect(() => { + if (!traceRef) return; + const controller = new AbortController(); + setLoading(true); + setFailed(false); + setTraces([]); + const params: Record = { + [traceRef.field]: traceRef.value, + }; + if (agentId) params.api_key_id = agentId; + userService + .getTraces(params, token, controller.signal) + .then(async (response: Response) => { + if (!response.ok) throw new Error('Failed to load traces'); + const data = await response.json(); + setTraces(data.traces ?? []); + }) + .catch((error: unknown) => { + if ((error as Error)?.name === 'AbortError') return; + console.error(error); + setFailed(true); + }) + .finally(() => { + if (!controller.signal.aborted) setLoading(false); + }); + return () => controller.abort(); + }, [traceRef?.field, traceRef?.value, agentId, token]); + + return ( + !open && onClose()}> + {/* The waterfall speaks for itself: the title is for screen readers + only, and there is no description to announce. */} + { + returnFocusRef.current = + document.activeElement instanceof HTMLElement + ? document.activeElement + : null; + }} + onCloseAutoFocus={(event) => { + const target = returnFocusRef.current; + returnFocusRef.current = null; + if (target?.isConnected) { + event.preventDefault(); + target.focus(); + } + }} + className="w-full overflow-y-auto sm:max-w-3xl" + > +
+ {loading && ( +
+ +
+ )} + {!loading && failed && ( +

+ {t('settings.logs.trace.failed')} +

+ )} + {!loading && !failed && traces.length === 0 && ( +

+ {t('settings.logs.trace.empty')} +

+ )} + {traces.map((trace, index) => ( +
+ {/* Right padding keeps the first row clear of the close button. */} +
+ {traces.length > 1 && ( + + {t('settings.logs.trace.round', { n: index + 1 })} + + )} + + {t( + `settings.logs.trace.sources.${trace.source}`, + trace.source, + )} + + + {formatDateTime(trace.started_at)} + + + {t( + `settings.logs.trace.status.${trace.status}`, + trace.status, + )} + +
+ + {trace.dropped_spans > 0 && ( +

+ {t('settings.logs.trace.droppedSpans', { + count: trace.dropped_spans, + })} +

+ )} + +
+ ))} +
+ + + ); +} diff --git a/frontend/src/settings/traces/TraceSpanDetails.tsx b/frontend/src/settings/traces/TraceSpanDetails.tsx new file mode 100644 index 00000000..146acef1 --- /dev/null +++ b/frontend/src/settings/traces/TraceSpanDetails.tsx @@ -0,0 +1,209 @@ +import React from 'react'; +import { useTranslation } from 'react-i18next'; + +import { Alert, AlertDescription } from '../../components/ui/alert'; +import { ToolCallPanel } from '../../conversation/AnswerFlow'; +import { TraceSpan } from '../types'; +import { formatDurationMs, formatTokens } from './traceUtils'; + +type Row = [string, string]; + +function str(value: unknown): string | undefined { + if (value === null || value === undefined || value === '') return undefined; + if (Array.isArray(value)) return value.map(String).join(', '); + return String(value); +} + +function jsonText(value: unknown): string { + return typeof value === 'string' ? value : JSON.stringify(value, null, 2); +} + +type ChunkPreview = { + title?: string; + source?: string; + score?: number; + text?: string; +}; + +/** Everything recorded for one span: its key facts, previews and raw attributes. */ +export default function TraceSpanDetails({ span }: { span: TraceSpan }) { + const { t } = useTranslation(); + const a = span.attributes || {}; + const f = (key: string) => t(`settings.logs.trace.fields.${key}`); + const yes = t('settings.logs.trace.fields.yes'); + const rows: Row[] = []; + const push = (label: string, value: unknown) => { + const text = str(value); + if (text !== undefined) rows.push([label, text]); + }; + + push( + f('status'), + t(`settings.logs.trace.status.${span.status}`, span.status), + ); + push(f('duration'), formatDurationMs(span.duration_ms)); + switch (span.kind) { + case 'llm': { + push(f('provider'), a['gen_ai.provider.name']); + push(f('model'), a['gen_ai.request.model']); + const input = a['gen_ai.usage.input_tokens']; + const output = a['gen_ai.usage.output_tokens']; + if (typeof input === 'number') + push(f('inputTokens'), formatTokens(input)); + if (typeof output === 'number') + push(f('outputTokens'), formatTokens(output)); + push(f('cachedTokens'), a['gen_ai.usage.cache_read.input_tokens']); + // A stream's span covers the whole time it was open, including the + // app's own work between chunks; model time is the provider's share. + if (a['docsgpt.stream'] && typeof a['docsgpt.provider_ms'] === 'number') + push(f('modelTime'), formatDurationMs(a['docsgpt.provider_ms'])); + if (typeof a['docsgpt.ttft_ms'] === 'number') + push(f('timeToFirstToken'), formatDurationMs(a['docsgpt.ttft_ms'])); + if ( + typeof a['docsgpt.cost_usd'] === 'number' && + a['docsgpt.cost_usd'] > 0 + ) + push(f('cost'), `$${a['docsgpt.cost_usd'].toFixed(6)}`); + push(f('tokenSource'), a['docsgpt.token_source']); + if (a['docsgpt.cache_hit']) push(f('cacheHit'), yes); + if (a['docsgpt.usage_estimated']) push(f('estimatedUsage'), yes); + break; + } + case 'tool': + push(f('tool'), a['docsgpt.tool']); + push(f('action'), a['docsgpt.action'] ?? a['gen_ai.tool.name']); + push(f('callId'), a['gen_ai.tool.call.id']); + break; + case 'retrieval': + case 'search': + case 'embedding': + case 'rerank': + push(f('retriever'), a['docsgpt.retriever']); + push(f('sources'), a['docsgpt.source_ids'] ?? a['gen_ai.data_source.id']); + push(f('model'), a['gen_ai.request.model']); + push(f('topK'), a['docsgpt.top_k']); + push(f('chunks'), a['docsgpt.chunk_count'] ?? a['docsgpt.kept_count']); + push(f('candidates'), a['docsgpt.candidate_count']); + push(f('topScore'), a['docsgpt.top_score']); + push(f('vectorStore'), a['docsgpt.vector_store']); + break; + case 'agent': + push(f('agentType'), a['docsgpt.agent_type']); + push(f('model'), a['gen_ai.request.model']); + push(f('sources'), a['docsgpt.source_count']); + push(f('toolCalls'), a['docsgpt.tool_call_count']); + break; + case 'guardrail': + push(f('stage'), a['docsgpt.guardrail.stage']); + push(f('checks'), a['docsgpt.guardrail.checks']); + push(f('triggered'), a['docsgpt.guardrail.triggered']); + if (a['docsgpt.guardrail.blocked']) push(f('blocked'), yes); + if (a['docsgpt.guardrail.redacted']) push(f('redacted'), yes); + break; + case 'step': + push(f('nodeType'), a['docsgpt.workflow.node_type']); + push(f('phase'), a['docsgpt.research.phase']); + break; + } + + const preview = span.preview || {}; + const chunks = Array.isArray(preview.chunks) + ? (preview.chunks as ChunkPreview[]) + : []; + const otherPreviews = Object.entries(preview).filter( + ([key]) => + !['arguments', 'result', 'query', 'chunks', 'output'].includes(key), + ); + + return ( +
+
+ {rows.map(([label, value]) => ( + + {label} + {value} + + ))} +
+ {span.error && ( + + + {span.error} + + + )} + {preview.query !== undefined && ( + +

+ {jsonText(preview.query)} +

+ + )} + {preview.arguments !== undefined && ( + +

+ {jsonText(preview.arguments)} +

+ + )} + {preview.result !== undefined && ( + +

+ {jsonText(preview.result)} +

+ + )} + {preview.output !== undefined && ( + +

+ {jsonText(preview.output)} +

+ + )} + {chunks.length > 0 && ( + +
    + {chunks.map((chunk, index) => ( +
  1. + + {index + 1}. {chunk.title || chunk.source || '—'} + {typeof chunk.score === 'number' && ( + + {chunk.score.toFixed(3)} + + )} + + {chunk.text && ( + + {chunk.text} + + )} +
    2. + ))} +
+ + )} + {otherPreviews.map(([key, value]) => ( + +

+ {jsonText(value)} +

+ + ))} +
+ + {f('allAttributes')} + + 
+          {JSON.stringify(a, null, 2)}
+
+
+
+ ); +} diff --git a/frontend/src/settings/traces/TraceWaterfall.tsx b/frontend/src/settings/traces/TraceWaterfall.tsx new file mode 100644 index 00000000..30957719 --- /dev/null +++ b/frontend/src/settings/traces/TraceWaterfall.tsx @@ -0,0 +1,208 @@ +import { + Binary, + Bot, + Brain, + ChevronRight, + Database, + ListFilter, + ListTree, + Search, + ShieldCheck, + Wrench, +} from 'lucide-react'; +import React, { useMemo, useState } from 'react'; +import { useTranslation } from 'react-i18next'; + +import { Button } from '../../components/ui/button'; +import { Trace, TraceSpan } from '../types'; +import TraceSpanDetails from './TraceSpanDetails'; +import { + barGeometry, + buildSpanRows, + formatDurationMs, + spanHeadline, + traceTotalMs, +} from './traceUtils'; + +/** + * Icon per span kind, and one of the five chart series per group of kinds: + * the agent and its steps, model calls, tool calls, retrieval (with its + * embedding, per-source search and rerank), and guardrails. + */ +const KIND_STYLE: Record = { + agent: { icon: Bot, bar: 'bg-chart-1' }, + step: { icon: ListTree, bar: 'bg-chart-1' }, + llm: { icon: Brain, bar: 'bg-chart-2' }, + tool: { icon: Wrench, bar: 'bg-chart-3' }, + retrieval: { icon: Search, bar: 'bg-chart-4' }, + search: { icon: Database, bar: 'bg-chart-4' }, + embedding: { icon: Binary, bar: 'bg-chart-4' }, + rerank: { icon: ListFilter, bar: 'bg-chart-4' }, + guardrail: { icon: ShieldCheck, bar: 'bg-chart-5' }, +}; + +const INACTIVE_STATUSES = ['cancelled', 'pending', 'denied', 'skipped']; + +function kindStyle(span: TraceSpan) { + return KIND_STYLE[span.kind] ?? KIND_STYLE.step; +} + +function barClass(span: TraceSpan): string { + if (span.status === 'error') return 'bg-destructive'; + if (INACTIVE_STATUSES.includes(span.status)) return 'bg-muted-foreground/40'; + return kindStyle(span).bar; +} + +/** Scale ticks: 0, a quarter, half, three quarters and the full duration. */ +const SCALE_TICKS = [ + { at: 0, className: 'left-0' }, + { at: 0.25, className: 'left-1/4 -translate-x-1/2' }, + { at: 0.5, className: 'left-1/2 -translate-x-1/2' }, + { at: 0.75, className: 'left-3/4 -translate-x-1/2' }, + { at: 1, className: 'right-0' }, +]; + +/** Indent per nesting level, in rem. */ +const INDENT_REM = 0.875; + +/** A waterfall of one trace's spans; select a row for its details. */ +export default function TraceWaterfall({ trace }: { trace: Trace }) { + const { t } = useTranslation(); + const rows = useMemo(() => buildSpanRows(trace.spans), [trace.spans]); + const totalMs = traceTotalMs(trace.duration_ms, trace.spans); + const [selectedId, setSelectedId] = useState(null); + const [collapsed, setCollapsed] = useState>(new Set()); + + const visibleRows = useMemo(() => { + const out: typeof rows = []; + let hiddenBelowDepth: number | null = null; + for (const row of rows) { + if (hiddenBelowDepth !== null && row.depth > hiddenBelowDepth) continue; + hiddenBelowDepth = collapsed.has(row.span.id) ? row.depth : null; + out.push(row); + } + return out; + }, [rows, collapsed]); + + const toggleCollapsed = (id: string) => + setCollapsed((prev) => { + const next = new Set(prev); + if (next.has(id)) next.delete(id); + else next.add(id); + return next; + }); + + const toggleSelected = (id: string) => + setSelectedId((current) => (current === id ? null : id)); + + if (!rows.length) { + return ( +

+ {t('settings.logs.trace.noSpans')} +

+ ); + } + + return ( +
+
+ {t('settings.logs.trace.step')} +
+ {SCALE_TICKS.map((tick) => ( + + {tick.at === 0 ? '0' : formatDurationMs(totalMs * tick.at)} + + ))} +
+
+ {visibleRows.map(({ span, depth, hasChildren }) => { + const Icon = kindStyle(span).icon; + const { left, width } = barGeometry(span, totalMs); + const selected = selectedId === span.id; + const headline = spanHeadline(span, t); + const isCollapsed = collapsed.has(span.id); + const geometry = { + '--trace-indent': `${depth * INDENT_REM}rem`, + '--trace-bar-left': `${left}%`, + '--trace-bar-width': `${width}%`, + } as React.CSSProperties; + return ( +
+
+
+ {hasChildren ? ( + + ) : ( + + )} + +
+ {/* Mouse shortcut to the row's button; keyboard users use the button. */} + +
+ {selected && ( +
+ +
+ )} +
+ ); + })} +
+ ); +} diff --git a/frontend/src/settings/traces/traceUtils.test.ts b/frontend/src/settings/traces/traceUtils.test.ts new file mode 100644 index 00000000..106c3dc8 --- /dev/null +++ b/frontend/src/settings/traces/traceUtils.test.ts @@ -0,0 +1,160 @@ +import { describe, expect, it } from 'vitest'; + +import { TraceSpan } from '../types'; +import { + barGeometry, + buildSpanRows, + formatDurationMs, + formatTokens, + spanHeadline, + traceTotalMs, +} from './traceUtils'; + +const span = (overrides: Partial): TraceSpan => ({ + id: 's', + parent_id: null, + kind: 'llm', + name: 'chat m', + status: 'ok', + offset_ms: 0, + duration_ms: 10, + attributes: {}, + ...overrides, +}); + +describe('buildSpanRows', () => { + it('orders depth-first by start offset with depths', () => { + const rows = buildSpanRows([ + span({ id: 'llm2', parent_id: 'agent', offset_ms: 50 }), + span({ id: 'agent', kind: 'agent', offset_ms: 5, duration_ms: 100 }), + span({ id: 'retrieval', kind: 'retrieval', offset_ms: 0 }), + span({ id: 'tool', kind: 'tool', parent_id: 'agent', offset_ms: 20 }), + span({ id: 'inner', parent_id: 'tool', offset_ms: 21 }), + ]); + expect(rows.map((r) => [r.span.id, r.depth])).toEqual([ + ['retrieval', 0], + ['agent', 0], + ['tool', 1], + ['inner', 2], + ['llm2', 1], + ]); + }); + + it('treats spans whose parent is missing as roots', () => { + const rows = buildSpanRows([span({ id: 'a', parent_id: 'gone' })]); + expect(rows).toHaveLength(1); + expect(rows[0].depth).toBe(0); + }); + + it('marks rows that have children', () => { + const rows = buildSpanRows([ + span({ id: 'p', kind: 'agent' }), + span({ id: 'c', parent_id: 'p' }), + ]); + expect(rows[0].hasChildren).toBe(true); + expect(rows[1].hasChildren).toBe(false); + }); +}); + +describe('traceTotalMs', () => { + it('uses the larger of the stored duration and the last span end', () => { + expect(traceTotalMs(100, [span({ offset_ms: 90, duration_ms: 30 })])).toBe( + 120, + ); + expect(traceTotalMs(500, [span({ offset_ms: 0, duration_ms: 30 })])).toBe( + 500, + ); + }); + + it('never returns zero', () => { + expect(traceTotalMs(0, [])).toBeGreaterThan(0); + }); +}); + +describe('barGeometry', () => { + it('positions a span as percentages of the total', () => { + expect(barGeometry(span({ offset_ms: 25, duration_ms: 50 }), 100)).toEqual({ + left: 25, + width: 50, + }); + }); + + it('keeps tiny spans visible and inside the track', () => { + const tiny = barGeometry(span({ offset_ms: 100, duration_ms: 0 }), 100); + expect(tiny.width).toBeGreaterThan(0); + expect(tiny.left + tiny.width).toBeLessThanOrEqual(100); + }); +}); + +describe('formatDurationMs', () => { + it.each([ + [0.4, '<1 ms'], + [12.4, '12 ms'], + [999, '999 ms'], + [1234, '1.23 s'], + [15432, '15.4 s'], + [65000, '1m 05s'], + [119_600, '2m 00s'], + ])('%s -> %s', (ms, expected) => { + expect(formatDurationMs(ms)).toBe(expected); + }); + + it('shows a dash for missing values', () => { + expect(formatDurationMs(undefined)).toBe('—'); + }); +}); + +describe('formatTokens', () => { + it('abbreviates large counts', () => { + expect(formatTokens(950)).toBe('950'); + expect(formatTokens(12345)).toBe('12.3k'); + expect(formatTokens(2_500_000)).toBe('2.5M'); + }); +}); + +const HEADLINES: Record = { + 'settings.logs.trace.headline.tokens': '{{input}} → {{output}} tok', + 'settings.logs.trace.headline.chunks': '{{count}} chunks', + 'settings.logs.trace.headline.hits': '{{count}} hits', + 'settings.logs.trace.headline.cached': 'cached', +}; + +const t = (key: string, options: Record = {}) => + (HEADLINES[key] ?? key).replace(/\{\{(\w+)\}\}/g, (_m, name) => + String(options[name]), + ); + +describe('spanHeadline', () => { + it('summarises an llm span by tokens', () => { + expect( + spanHeadline( + span({ + attributes: { + 'gen_ai.usage.input_tokens': 1200, + 'gen_ai.usage.output_tokens': 80, + }, + }), + t, + ), + ).toBe('1.2k → 80 tok'); + }); + + it('summarises a retrieval span by chunk count', () => { + expect( + spanHeadline( + span({ kind: 'retrieval', attributes: { 'docsgpt.chunk_count': 4 } }), + t, + ), + ).toBe('4 chunks'); + }); + + it('is empty when nothing is known', () => { + expect(spanHeadline(span({ kind: 'step' }), t)).toBe(''); + }); + + it('labels a cached llm call', () => { + expect( + spanHeadline(span({ attributes: { 'docsgpt.cache_hit': true } }), t), + ).toBe('cached'); + }); +}); diff --git a/frontend/src/settings/traces/traceUtils.ts b/frontend/src/settings/traces/traceUtils.ts new file mode 100644 index 00000000..e077299c --- /dev/null +++ b/frontend/src/settings/traces/traceUtils.ts @@ -0,0 +1,142 @@ +import { TraceSpan } from '../types'; + +export type SpanRow = { + span: TraceSpan; + depth: number; + hasChildren: boolean; +}; + +/** + * Flatten a trace's spans into display rows: depth-first, children in start + * order under their parent. A span whose parent is missing (dropped by the + * span cap) is shown at the top level rather than hidden. + */ +export function buildSpanRows(spans: TraceSpan[]): SpanRow[] { + const ids = new Set(spans.map((s) => s.id)); + const children = new Map(); + for (const span of spans) { + const parent = + span.parent_id && ids.has(span.parent_id) ? span.parent_id : null; + const list = children.get(parent) ?? []; + list.push(span); + children.set(parent, list); + } + for (const list of children.values()) { + list.sort((a, b) => a.offset_ms - b.offset_ms); + } + const rows: SpanRow[] = []; + const visit = (parent: string | null, depth: number) => { + for (const span of children.get(parent) ?? []) { + const kids = children.get(span.id); + rows.push({ span, depth, hasChildren: Boolean(kids?.length) }); + visit(span.id, depth + 1); + } + }; + visit(null, 0); + return rows; +} + +/** Width of the timeline: the stored duration, or the last span end if later. */ +export function traceTotalMs( + durationMs: number | undefined, + spans: TraceSpan[], +): number { + const lastEnd = spans.reduce( + (max, s) => Math.max(max, (s.offset_ms || 0) + (s.duration_ms || 0)), + 0, + ); + return Math.max(durationMs || 0, lastEnd, 1); +} + +const MIN_BAR_PERCENT = 0.6; + +/** A span's bar as left offset and width, in percent of the timeline. */ +export function barGeometry( + span: TraceSpan, + totalMs: number, +): { left: number; width: number } { + const total = totalMs > 0 ? totalMs : 1; + const width = Math.max( + ((span.duration_ms || 0) / total) * 100, + MIN_BAR_PERCENT, + ); + const left = Math.min( + Math.max(((span.offset_ms || 0) / total) * 100, 0), + 100 - width, + ); + return { left: round(left), width: round(width) }; +} + +function round(n: number): number { + return Math.round(n * 100) / 100; +} + +/** Human duration: `<1 ms`, `12 ms`, `1.23 s`, `15.4 s`, `1m 05s`. */ +export function formatDurationMs(ms: number | null | undefined): string { + if (ms === null || ms === undefined || Number.isNaN(ms)) return '—'; + if (ms < 1) return '<1 ms'; + if (ms < 1000) return `${Math.round(ms)} ms`; + if (ms < 10_000) return `${(ms / 1000).toFixed(2)} s`; + if (ms < 60_000) return `${(ms / 1000).toFixed(1)} s`; + // Round to whole seconds before splitting, so 119.6 s reads 2m 00s, not 1m 60s. + const totalSeconds = Math.round(ms / 1000); + const minutes = Math.floor(totalSeconds / 60); + const seconds = totalSeconds % 60; + return `${minutes}m ${String(seconds).padStart(2, '0')}s`; +} + +/** Token counts abbreviated: `950`, `12.3k`, `2.5M`. */ +export function formatTokens(n: number | null | undefined): string { + const value = n ?? 0; + if (value >= 1_000_000) return `${trimZero((value / 1_000_000).toFixed(1))}M`; + if (value >= 1000) return `${trimZero((value / 1000).toFixed(1))}k`; + return String(Math.round(value)); +} + +function trimZero(s: string): string { + return s.endsWith('.0') ? s.slice(0, -2) : s; +} + +function num(value: unknown): number | undefined { + return typeof value === 'number' && Number.isFinite(value) + ? value + : undefined; +} + +type Translate = (key: string, options?: Record) => string; + +/** + * A short, kind-specific summary shown next to a span's name, e.g. + * `1.2k → 80 tok` for an LLM call or `4 chunks` for a retrieval. + */ +export function spanHeadline(span: TraceSpan, t: Translate): string { + const a = span.attributes || {}; + switch (span.kind) { + case 'llm': { + if (a['docsgpt.cache_hit']) + return t('settings.logs.trace.headline.cached'); + const input = num(a['gen_ai.usage.input_tokens']); + const output = num(a['gen_ai.usage.output_tokens']); + if (input === undefined && output === undefined) return ''; + return t('settings.logs.trace.headline.tokens', { + input: formatTokens(input), + output: formatTokens(output), + }); + } + case 'retrieval': + case 'rerank': { + const chunks = num(a['docsgpt.chunk_count'] ?? a['docsgpt.kept_count']); + return chunks === undefined + ? '' + : t('settings.logs.trace.headline.chunks', { count: chunks }); + } + case 'search': { + const hits = num(a['docsgpt.candidate_count']); + return hits === undefined + ? '' + : t('settings.logs.trace.headline.hits', { count: hits }); + } + default: + return ''; + } +} diff --git a/frontend/src/settings/types/index.ts b/frontend/src/settings/types/index.ts index 81e557a8..76006b84 100644 --- a/frontend/src/settings/types/index.ts +++ b/frontend/src/settings/types/index.ts @@ -17,7 +17,74 @@ export type ChunkType = { }; export type LogEventType = - 'chat' | 'schedule' | 'webhook' | 'workflow' | 'system'; + 'chat' | 'schedule' | 'webhook' | 'workflow' | 'system' | 'search' | 'graph'; + +/** Counts rolled up from a trace's spans; every field may be absent. */ +export type TraceCounts = { + llm_calls?: number; + tool_calls?: number; + retrieval_calls?: number; + retrieval_ms?: number; + input_tokens?: number; + output_tokens?: number; + errors?: number; +}; + +/** The id a Logs row's traces are looked up by. */ +export type TraceRef = { + field: 'id' | 'request_id' | 'message_id' | 'activity_id' | 'workflow_run_id'; + value: string; +}; + +/** Trace summary attached to a Logs row (merged over its rounds). */ +export type LogTraceSummary = { + ref: TraceRef; + count: number; + duration_ms: number; + status?: string; + started_at?: string; + summary: TraceCounts; +}; + +export type TraceSpanKind = + | 'agent' + | 'llm' + | 'tool' + | 'retrieval' + | 'search' + | 'embedding' + | 'rerank' + | 'guardrail' + | 'step'; + +export type TraceSpan = { + id: string; + parent_id: string | null; + kind: TraceSpanKind | string; + name: string; + status: string; + offset_ms: number; + duration_ms: number; + attributes: Record; + preview?: Record; + error?: string; +}; + +export type Trace = { + id: string; + request_id?: string | null; + message_id?: string | null; + source: string; + name?: string | null; + status: string; + started_at: string; + duration_ms: number; + span_count: number; + dropped_spans: number; + summary: TraceCounts; + spans: TraceSpan[]; + otel_trace_id?: string | null; +}; export type LogData = { id: string; @@ -33,9 +100,17 @@ export type LogData = { tool_calls?: Record[]; agent_id?: string; attachments?: string[]; + request_id?: string; + message_id?: string; // system + webhook events (stack_logs) endpoint?: string; stacks?: Record[]; + activity_id?: string; + // search + graph events (request_traces) + source?: string; + duration_ms?: number; + // Present when the row has a stored execution trace. + trace?: LogTraceSummary; // workflow events (workflow_runs) workflow_name?: string; result?: Record; diff --git a/pyproject.toml b/pyproject.toml index 5e24bbaa..dbbfc373 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -78,6 +78,8 @@ dependencies = [ "openapi3-parser>=2.0.0,<3", # pandas reads .xlsx through openpyxl but does not depend on it. "openpyxl>=3.1.5,<4", + # Imported directly by docsgpt.tracing to replay traces as GenAI spans. + "opentelemetry-api>=1.29.0,<2", "opentelemetry-distro>=0.50b0,<1", "opentelemetry-exporter-otlp>=1.29.0,<2", "opentelemetry-instrumentation-celery>=0.50b0,<1", diff --git a/tests/agents/test_research_agent.py b/tests/agents/test_research_agent.py index cd9aabd9..79b899e7 100644 --- a/tests/agents/test_research_agent.py +++ b/tests/agents/test_research_agent.py @@ -1626,3 +1626,53 @@ class TestResearchAgentClarificationCoverage: text = 'Here is my response: {"needs_clarification": false} end.' result = agent._parse_clarification_json(text) assert result == {"needs_clarification": False} + + +@pytest.mark.unit +class TestGenInnerTraceSpans: + """Each research phase is a step span in the execution trace.""" + + def test_phases_become_step_spans( + self, + agent_base_params, + mock_llm, + mock_llm_creator, + mock_llm_handler_creator, + log_context, + monkeypatch, + ): + from docsgpt import tracing + from docsgpt.core.settings import settings + + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + agent = ResearchAgent(**agent_base_params) + plan_steps = [{"query": "s1", "rationale": "r"}, {"query": "s2", "rationale": "r"}] + + def fake_step(query, tools_dict): + with tracing.span(tracing.KIND_LLM, "chat m"): + return f"report {query}" + + trace = tracing.start_trace(source="stream", capture_otel_context=False) + with tracing.activate(trace), \ + patch.object(agent, "_setup_tools", return_value={}), \ + patch.object(agent, "_is_follow_up", return_value=False), \ + patch.object(agent, "_clarification_phase", return_value=None), \ + patch.object(agent, "_planning_phase", return_value=(plan_steps, "moderate")), \ + patch.object(agent, "_research_step", side_effect=fake_step), \ + patch.object(agent, "_synthesis_phase", return_value=iter([{"answer": "final"}])), \ + patch.object(agent, "_get_truncated_tool_calls", return_value=[]): + list(agent._gen_inner("Compare A and B", log_context)) + + names = [s.name for s in trace.spans if s.kind == tracing.KIND_STEP] + assert names == [ + "research clarify", + "research plan", + "research step 1", + "research step 2", + "research synthesis", + ] + by_name = {s.name: s for s in trace.spans} + assert by_name["research plan"].attributes["docsgpt.research.steps"] == 2 + llm_spans = [s for s in trace.spans if s.kind == tracing.KIND_LLM] + assert llm_spans[0].parent_id == by_name["research step 1"].id + assert by_name["research step 1"].previews["query"] == "s1" diff --git a/tests/agents/test_workflow_engine_coverage.py b/tests/agents/test_workflow_engine_coverage.py index 5382523c..1b8c6707 100644 --- a/tests/agents/test_workflow_engine_coverage.py +++ b/tests/agents/test_workflow_engine_coverage.py @@ -1081,3 +1081,51 @@ class TestNodeDocumentManifest: monkeypatch.setattr("docsgpt.storage.db.session.db_readonly", broken_readonly) config = AgentNodeConfig(agent_type="classic", input_documents=["*"]) assert engine._node_document_manifest(config) == "" + + +class TestTraceSpans: + """Every executed node is a ``workflow_step`` span; the run id links the trace.""" + + @pytest.fixture(autouse=True) + def _trace(self, monkeypatch): + from docsgpt import tracing + from docsgpt.core.settings import settings + + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + self.trace = tracing.start_trace(source="stream", capture_otel_context=False) + with tracing.activate(self.trace): + yield + + @pytest.mark.unit + def test_step_spans_and_run_binding(self): + nodes = [ + _make_node("n1", NodeType.START, "Start"), + _make_node("n2", NodeType.END, "End", config={"config": {}}), + ] + engine = WorkflowEngine( + _make_graph(nodes, [_make_edge("e1", "n1", "n2")]), _make_agent() + ) + list(engine.execute({}, "hello")) + steps = [s for s in self.trace.spans if s.kind == "step"] + assert [s.name for s in steps] == ["workflow_step Start", "workflow_step End"] + assert all(s.status == "ok" for s in steps) + assert steps[0].attributes["docsgpt.workflow.node_type"] == "start" + assert self.trace.workflow_run_id == str(engine.workflow_run_id) + + @pytest.mark.unit + def test_failed_node_is_an_error_span(self): + nodes = [ + _make_node("n1", NodeType.START), + _make_node( + "n2", + NodeType.STATE, + "State", + config={"config": {"operations": [{"expression": "bad!!!", "target_variable": "x"}]}}, + ), + ] + engine = WorkflowEngine( + _make_graph(nodes, [_make_edge("e1", "n1", "n2")]), _make_agent() + ) + list(engine.execute({}, "q")) + state_span = [s for s in self.trace.spans if s.name == "workflow_step State"][0] + assert state_span.status == "error" diff --git a/tests/api/answer/test_trace_wiring.py b/tests/api/answer/test_trace_wiring.py new file mode 100644 index 00000000..01be1475 --- /dev/null +++ b/tests/api/answer/test_trace_wiring.py @@ -0,0 +1,368 @@ +"""``complete_stream`` owns the request's execution trace and writes it once.""" + +from __future__ import annotations + +from contextlib import contextmanager +from unittest.mock import MagicMock, patch + +import pytest + +from docsgpt import tracing +from docsgpt.core.settings import settings + + +@pytest.fixture(autouse=True) +def _tracing_on(monkeypatch): + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + monkeypatch.setattr(settings, "TRACES_OTEL_EXPORT", False) + + +@contextmanager +def _captured_flushes(): + """Record every flushed trace instead of writing it.""" + flushed = [] + + def _fake_flush(trace, status=None, **_kwargs): + if trace is None or trace.flushed: + return + trace.flushed = True + trace.finish(status) + flushed.append(trace) + + with patch("docsgpt.tracing.flush", side_effect=_fake_flush): + yield flushed + + +def _agent(events): + agent = MagicMock() + + def _gen(query): + with tracing.span(tracing.KIND_AGENT, "invoke_agent Fake"): + with tracing.span(tracing.KIND_LLM, "chat m"): + pass + yield from events + + agent.gen.side_effect = _gen + agent.tool_calls = [] + agent.compression_metadata = None + agent.compression_saved = False + return agent + + +def _run(resource, agent, **kwargs): + base = dict( + question="q", + agent=agent, + conversation_id=None, + user_api_key=None, + decoded_token={"sub": "u-trace"}, + should_persist=False, + ) + base.update(kwargs) + return list(resource.complete_stream(**base)) + + +@pytest.mark.unit +class TestTraceLifecycle: + def test_normal_turn_flushes_once_with_ids(self, flask_app, mock_mongo_db): + from docsgpt.api.answer.routes.base import BaseAnswerResource + + with flask_app.app_context(), _captured_flushes() as flushed: + _run(BaseAnswerResource(), _agent([{"answer": "hi"}]), request_id="req-1") + (trace,) = flushed + assert trace.request_id == "req-1" + assert trace.user_id == "u-trace" + assert trace.status == "ok" + assert [s.name for s in trace.spans] == ["invoke_agent Fake", "chat m"] + + def test_route_trace_is_reused(self, flask_app, mock_mongo_db): + from docsgpt.api.answer.routes.base import BaseAnswerResource + + route_trace = tracing.start_trace(source="answer", capture_otel_context=False) + with tracing.activate(route_trace): + with tracing.span(tracing.KIND_RETRIEVAL, "retrieval"): + pass + with flask_app.app_context(), _captured_flushes() as flushed: + _run( + BaseAnswerResource(), + _agent([{"answer": "hi"}]), + request_id="req-2", + trace=route_trace, + ) + assert flushed == [route_trace] + assert route_trace.source == "answer" + assert [s.kind for s in route_trace.spans] == ["retrieval", "agent", "llm"] + + def test_mock_trace_is_ignored(self, flask_app, mock_mongo_db): + from docsgpt.api.answer.routes.base import BaseAnswerResource + + with flask_app.app_context(), _captured_flushes() as flushed: + _run(BaseAnswerResource(), _agent([{"answer": "hi"}]), trace=MagicMock()) + assert len(flushed) == 1 + assert isinstance(flushed[0], tracing.Trace) + + def test_continuation_keeps_saved_request_id(self, flask_app, mock_mongo_db): + from docsgpt.api.answer.routes.base import BaseAnswerResource + + agent = _agent([]) + agent.gen_continuation.side_effect = lambda **_kw: iter([{"answer": "done"}]) + with flask_app.app_context(), _captured_flushes() as flushed: + _run( + BaseAnswerResource(), + agent, + question="", + request_id="fresh", + _continuation={ + "messages": [], + "tools_dict": {}, + "pending_tool_calls": [], + "tool_actions": [], + "request_id": "saved-req", + }, + ) + assert flushed[0].request_id == "saved-req" + + def test_agent_error_marks_trace_error(self, flask_app, mock_mongo_db): + from docsgpt.api.answer.routes.base import BaseAnswerResource + + agent = MagicMock() + agent.gen.side_effect = RuntimeError("upstream down") + with flask_app.app_context(), _captured_flushes() as flushed: + stream = _run(BaseAnswerResource(), agent) + assert any('"type": "error"' in s for s in stream) + assert flushed[0].status == "error" + + def test_yielded_error_marks_trace_error(self, flask_app, mock_mongo_db): + """A failed workflow node yields an error event instead of raising.""" + from docsgpt.api.answer.routes.base import BaseAnswerResource + + with flask_app.app_context(), _captured_flushes() as flushed: + _run( + BaseAnswerResource(), + _agent([{"type": "error", "error": "node failed"}]), + ) + assert flushed[0].status == "error" + + def test_abandoned_stream_still_flushes(self, flask_app, mock_mongo_db): + from docsgpt.api.answer.routes.base import BaseAnswerResource + + with flask_app.app_context(), _captured_flushes() as flushed: + gen = BaseAnswerResource().complete_stream( + question="q", + agent=_agent([{"answer": "a"}, {"answer": "b"}]), + conversation_id=None, + user_api_key=None, + decoded_token={"sub": "u"}, + should_persist=False, + ) + next(gen) + gen.close() + assert len(flushed) == 1 + + +@pytest.mark.unit +class TestTraceWithPersistence: + def test_message_and_conversation_ids_bound(self, pg_conn, flask_app): + from docsgpt.api.answer.routes.base import BaseAnswerResource + from tests.api.answer.test_base_routes import _patch_db_session + + with flask_app.app_context(), _patch_db_session(pg_conn), _captured_flushes() as flushed: + _run( + BaseAnswerResource(), + _agent([{"answer": "persisted"}]), + should_persist=True, + model_id="gpt-4", + request_id="req-p", + ) + trace = flushed[0] + assert trace.message_id + assert trace.conversation_id + from sqlalchemy import text as sql_text + + row = pg_conn.execute( + sql_text("SELECT data FROM user_logs WHERE user_id = 'u-trace'") + ).fetchone() + assert row[0]["request_id"] == "req-p" + assert row[0]["message_id"] == trace.message_id + + def test_failed_turn_is_logged_as_a_chat_row(self, pg_conn, flask_app): + """A raised failure still writes the turn's chat row, at level error.""" + from docsgpt.api.answer.routes.base import BaseAnswerResource + from tests.api.answer.test_base_routes import _patch_db_session + + agent = MagicMock() + agent.gen.side_effect = RuntimeError("upstream down") + agent.tool_calls = [] + with flask_app.app_context(), _patch_db_session(pg_conn), _captured_flushes(): + _run( + BaseAnswerResource(), + agent, + should_persist=True, + model_id="gpt-4", + request_id="req-failed", + ) + from sqlalchemy import text as sql_text + + rows = pg_conn.execute( + sql_text("SELECT data FROM user_logs WHERE user_id = 'u-trace'") + ).fetchall() + assert len(rows) == 1 + data = rows[0][0] + assert data["level"] == "error" + assert data["request_id"] == "req-failed" + assert data["error"] == "RuntimeError: upstream down" + + def test_yielded_error_logs_the_chat_row_at_error_level(self, pg_conn, flask_app): + from docsgpt.api.answer.routes.base import BaseAnswerResource + from tests.api.answer.test_base_routes import _patch_db_session + + with flask_app.app_context(), _patch_db_session(pg_conn), _captured_flushes(): + _run( + BaseAnswerResource(), + _agent([{"type": "error", "error": "node failed"}]), + should_persist=True, + model_id="gpt-4", + ) + from sqlalchemy import text as sql_text + + data = pg_conn.execute( + sql_text("SELECT data FROM user_logs WHERE user_id = 'u-trace'") + ).fetchone()[0] + assert data["level"] == "error" + assert data["error"] == "node failed" + + def test_paused_turn_is_flushed_paused(self, pg_conn, flask_app): + from docsgpt.api.answer.routes.base import BaseAnswerResource + from tests.api.answer.test_base_routes import _patch_db_session + + agent = _agent( + [ + { + "type": "tool_calls_pending", + "data": {"pending_tool_calls": [{"call_id": "c1"}]}, + } + ] + ) + agent._pending_continuation = { + "messages": [], + "tools_dict": {}, + "pending_tool_calls": [{"call_id": "c1"}], + } + with flask_app.app_context(), _patch_db_session(pg_conn), patch( + "docsgpt.api.answer.services.continuation_service.ContinuationService.save_state", + return_value=True, + ), _captured_flushes() as flushed: + _run(BaseAnswerResource(), agent, should_persist=True, model_id="gpt-4") + assert flushed[0].status == "paused" + + +@pytest.mark.unit +class TestProcessorTraceSetup: + def test_build_agent_mints_request_id_inside_the_trace(self): + from docsgpt.api.answer.services.stream_processor import StreamProcessor + + seen = {} + + class _Stop(Exception): + pass + + def _initialize(): + seen["trace"] = tracing.current_trace() + with tracing.span(tracing.KIND_RETRIEVAL, "retrieval"): + pass + raise _Stop() + + processor = StreamProcessor({"question": "q"}, {"sub": "u1"}, trace_source="answer") + with patch.object(processor, "initialize", side_effect=_initialize): + with pytest.raises(_Stop): + processor.build_agent("q") + trace = processor.trace + assert seen["trace"] is trace + assert trace.source == "answer" + assert processor.request_id and trace.request_id == processor.request_id + assert trace.user_id == "u1" + assert [s.kind for s in trace.spans] == ["retrieval"] + assert tracing.current_trace() is None + + def test_client_supplied_request_id_is_ignored(self): + """Quotas count distinct request ids; a client must not choose its own.""" + from docsgpt.api.answer.services.stream_processor import StreamProcessor + + processor = StreamProcessor({"request_id": "client-rid"}, {"sub": "u1"}) + with patch.object(processor, "initialize", side_effect=RuntimeError("stop")): + with pytest.raises(RuntimeError): + processor.build_agent("q") + assert processor.request_id and processor.request_id != "client-rid" + + def test_refused_request_still_writes_its_trace(self): + from docsgpt.api.answer.services.stream_processor import StreamProcessor + + processor = StreamProcessor({}, {"sub": "u1"}) + + def _initialize(): + with tracing.span(tracing.KIND_RETRIEVAL, "retrieval"): + pass + + with patch.object(processor, "initialize", side_effect=_initialize), patch.object( + processor, "pre_fetch_docs", return_value=(None, None) + ), patch.object(processor, "pre_fetch_tools", return_value=None), patch.object( + processor, "create_agent", return_value=MagicMock() + ), patch.object(processor, "_exposure_partition", return_value=([], [])): + processor.build_agent("q") + with _captured_flushes() as flushed: + processor.flush_unclaimed_trace() + (trace,) = flushed + assert trace.status == "error" + assert [s.kind for s in trace.spans] == ["retrieval"] + + def test_handed_off_trace_is_left_to_the_stream(self): + from docsgpt.api.answer.services.stream_processor import StreamProcessor + + processor = StreamProcessor({}, {"sub": "u1"}) + processor.trace = tracing.start_trace(source="stream", capture_otel_context=False) + assert processor.handoff_trace() is processor.trace + with _captured_flushes() as flushed: + processor.flush_unclaimed_trace() + assert flushed == [] + + def test_tracing_disabled_leaves_no_trace(self, monkeypatch): + from docsgpt.api.answer.services.stream_processor import StreamProcessor + + monkeypatch.setattr(settings, "TRACES_ENABLED", False) + processor = StreamProcessor({}, {"sub": "u1"}) + with patch.object(processor, "initialize", side_effect=RuntimeError("stop")): + with pytest.raises(RuntimeError): + processor.build_agent("q") + assert processor.trace is None + assert processor.request_id + + +@pytest.mark.unit +class TestRouteFlushesRefusedRequests: + def test_unauthorized_answer_request_writes_its_trace(self, mock_mongo_db, flask_app): + """The route registers the flush, and the hook never replaces the response.""" + import json + + from flask_restx import Api + + from docsgpt.api.answer.routes.answer import answer_ns + + api = Api(flask_app) + api.add_namespace(answer_ns) + client = flask_app.test_client() + processor = MagicMock() + processor.decoded_token = None + processor.flush_unclaimed_trace.return_value = "not a response" + with patch( + "docsgpt.api.answer.routes.answer.StreamProcessor", return_value=processor + ), patch( + "docsgpt.api.answer.routes.answer.AnswerResource.validate_request", + return_value=None, + ): + resp = client.post( + "/api/answer", + data=json.dumps({"question": "q"}), + content_type="application/json", + ) + assert resp.status_code == 401 + processor.flush_unclaimed_trace.assert_called_once_with() diff --git a/tests/api/user/test_tasks.py b/tests/api/user/test_tasks.py index a593bbb5..a3ff2a19 100644 --- a/tests/api/user/test_tasks.py +++ b/tests/api/user/test_tasks.py @@ -299,7 +299,7 @@ class TestSetupPeriodicTasks: setup_periodic_tasks(sender) - assert sender.add_periodic_task.call_count == 14 + assert sender.add_periodic_task.call_count == 15 calls = sender.add_periodic_task.call_args_list @@ -326,20 +326,23 @@ class TestSetupPeriodicTasks: # guardrail_events retention sweep (24h) assert calls[8][0][0] == timedelta(hours=24) assert calls[8][1].get("name") == "cleanup-guardrail-events" - # orphan memories sweep (24h) + # request_traces retention sweep (24h) assert calls[9][0][0] == timedelta(hours=24) - assert calls[9][1].get("name") == "cleanup-orphan-memories" + assert calls[9][1].get("name") == "cleanup-traces" + # orphan memories sweep (24h) + assert calls[10][0][0] == timedelta(hours=24) + assert calls[10][1].get("name") == "cleanup-orphan-memories" # scheduler dispatcher - assert calls[10][1].get("name") == "dispatch-scheduled-runs" + assert calls[11][1].get("name") == "dispatch-scheduled-runs" # schedule runs cleanup (24h) - assert calls[11][0][0] == timedelta(hours=24) - assert calls[11][1].get("name") == "cleanup-schedule-runs" + assert calls[12][0][0] == timedelta(hours=24) + assert calls[12][1].get("name") == "cleanup-schedule-runs" # sandbox session reaper (60s) - assert calls[12][0][0] == timedelta(seconds=60) - assert calls[12][1].get("name") == "reap-sandbox-sessions" + assert calls[13][0][0] == timedelta(seconds=60) + assert calls[13][1].get("name") == "reap-sandbox-sessions" # stale workflow-run reaper (5m) - assert calls[13][0][0] == timedelta(seconds=300) - assert calls[13][1].get("name") == "reap-stale-workflow-runs" + assert calls[14][0][0] == timedelta(seconds=300) + assert calls[14][1].get("name") == "reap-stale-workflow-runs" class TestMcpOauthTask: @@ -679,6 +682,67 @@ class TestCleanupMessageEventsTask: assert [r["sequence_no"] for r in rows] == [1] +class TestCleanupTracesTask: + """Retention janitor for ``request_traces``.""" + + @pytest.mark.unit + def test_skips_when_postgres_uri_missing(self, monkeypatch): + from docsgpt.api.user.tasks import cleanup_traces + from docsgpt.core.settings import settings + + monkeypatch.setattr(settings, "POSTGRES_URI", None, raising=False) + + assert cleanup_traces.run() == {"deleted": 0, "skipped": "POSTGRES_URI not set"} + + @pytest.mark.unit + def test_deletes_traces_past_retention_window(self, pg_conn, monkeypatch): + import time + import uuid + + from sqlalchemy import text as _text + + from docsgpt.api.user.tasks import cleanup_traces + from docsgpt.core.settings import settings + from docsgpt.storage.db.repositories.request_traces import ( + RequestTracesRepository, + ) + + repo = RequestTracesRepository(pg_conn) + for request_id in ("stale", "fresh"): + repo.insert( + { + "id": str(uuid.uuid4()), + "request_id": request_id, + "user_id": "u1", + "source": "stream", + "status": "ok", + "started_at_ns": time.time_ns(), + } + ) + pg_conn.execute( + _text( + "UPDATE request_traces SET created_at = now() - interval '45 days' " + "WHERE request_id = 'stale'" + ) + ) + monkeypatch.setattr(settings, "POSTGRES_URI", "postgresql://stub", raising=False) + monkeypatch.setattr(settings, "TRACES_RETENTION_DAYS", 30) + + @contextmanager + def _fake_begin(): + yield pg_conn + + fake_engine = MagicMock() + fake_engine.begin = _fake_begin + + with patch("docsgpt.storage.db.engine.get_engine", return_value=fake_engine): + result = cleanup_traces.run() + + assert result == {"deleted": 1, "ttl_days": 30} + remaining = pg_conn.execute(_text("SELECT request_id FROM request_traces")).scalars().all() + assert remaining == ["fresh"] + + class TestCleanupOrphanMemoriesTask: """Sweeps orphan memories from the FK-to-trigger orphan window.""" diff --git a/tests/api/user/test_traces_routes.py b/tests/api/user/test_traces_routes.py new file mode 100644 index 00000000..bbafdd79 --- /dev/null +++ b/tests/api/user/test_traces_routes.py @@ -0,0 +1,250 @@ +"""Tests for ``GET /api/traces`` and the trace summaries on ``get_user_logs`` rows.""" + +from __future__ import annotations + +import time +import uuid +from contextlib import contextmanager +from unittest.mock import patch + +import pytest +from flask import Flask + +from docsgpt.storage.db.repositories.request_traces import RequestTracesRepository + + +@pytest.fixture +def app(): + return Flask(__name__) + + +@contextmanager +def _patch_db(conn): + @contextmanager + def _yield_conn(): + yield conn + + with patch("docsgpt.api.user.analytics.routes.db_readonly", _yield_conn): + yield + + +def _trace(pg_conn, **overrides): + record = { + "id": str(uuid.uuid4()), + "request_id": "req-1", + "user_id": "owner", + "source": "stream", + "name": "stream", + "status": "ok", + "started_at_ns": time.time_ns(), + "duration_ms": 900, + "span_count": 2, + "summary": {"llm_calls": 1, "tool_calls": 1, "input_tokens": 50, "retrieval_ms": 12.5}, + "spans": [ + { + "id": "s1", + "parent_id": None, + "kind": "retrieval", + "name": "retrieval", + "status": "ok", + "offset_ms": 0, + "duration_ms": 12.5, + "attributes": {}, + "preview": {"query": "how do I deploy"}, + } + ], + } + record.update(overrides) + RequestTracesRepository(pg_conn).insert(record) + return record + + +def _get(app, pg_conn, user, query): + from docsgpt.api.user.analytics.routes import GetTraces + + with _patch_db(pg_conn), app.test_request_context("/api/traces", query_string=query): + from flask import request + + request.decoded_token = {"sub": user} if user else None + return GetTraces().get() + + +def _logs(app, pg_conn, user, body): + from docsgpt.api.user.analytics.routes import GetUserLogs + + with _patch_db(pg_conn), app.test_request_context( + "/api/get_user_logs", method="POST", json=body + ): + from flask import request + + request.decoded_token = {"sub": user} + return GetUserLogs().post() + + +class TestGetTraces: + def test_requires_auth(self, app, pg_conn): + assert _get(app, pg_conn, None, {"request_id": "x"}).status_code == 401 + + def test_requires_exactly_one_ref(self, app, pg_conn): + assert _get(app, pg_conn, "owner", {}).status_code == 400 + assert _get(app, pg_conn, "owner", {"request_id": "a", "id": "b"}).status_code == 400 + + def test_returns_own_trace_with_spans(self, app, pg_conn): + record = _trace(pg_conn) + response = _get(app, pg_conn, "owner", {"request_id": "req-1"}) + assert response.status_code == 200 + (trace,) = response.json["traces"] + assert trace["id"] == record["id"] + assert trace["spans"][0]["preview"]["query"] == "how do I deploy" + assert "_id" not in trace + + def test_other_users_get_nothing(self, app, pg_conn): + _trace(pg_conn) + response = _get(app, pg_conn, "intruder", {"request_id": "req-1"}) + assert response.json["traces"] == [] + + def test_agent_owner_sees_shared_callers_traces(self, app, pg_conn): + from docsgpt.storage.db.repositories.agents import AgentsRepository + + agent = AgentsRepository(pg_conn).create("owner", "a", "published", key="k1") + _trace(pg_conn, user_id="caller", agent_id=str(agent["id"])) + response = _get( + app, pg_conn, "owner", {"request_id": "req-1", "api_key_id": str(agent["id"])} + ) + assert len(response.json["traces"]) == 1 + + def test_unowned_agent_scope_is_empty(self, app, pg_conn): + from docsgpt.storage.db.repositories.agents import AgentsRepository + + agent = AgentsRepository(pg_conn).create("someone", "a", "published", key="k2") + _trace(pg_conn, user_id="someone", agent_id=str(agent["id"])) + response = _get( + app, pg_conn, "owner", {"request_id": "req-1", "api_key_id": str(agent["id"])} + ) + assert response.json["traces"] == [] + + +class TestLogsTraceSummaries: + def test_chat_row_gets_merged_summary_across_rounds(self, app, pg_conn): + from docsgpt.storage.db.repositories.user_logs import UserLogsRepository + + UserLogsRepository(pg_conn).insert( + user_id="owner", + endpoint="stream_answer", + data={"action": "stream_answer", "question": "q", "request_id": "req-1"}, + ) + _trace(pg_conn, status="paused") + _trace(pg_conn, started_at_ns=time.time_ns() + 1_000_000) + response = _logs(app, pg_conn, "owner", {"event_type": "chat"}) + (row,) = response.json["logs"] + assert row["request_id"] == "req-1" + trace = row["trace"] + assert trace["ref"] == {"field": "request_id", "value": "req-1"} + assert trace["count"] == 2 + assert trace["duration_ms"] == 1800 + assert trace["status"] == "ok" + assert trace["summary"]["llm_calls"] == 2 + assert trace["summary"]["retrieval_ms"] == 25.0 + + def test_row_without_trace_has_no_summary(self, app, pg_conn): + from docsgpt.storage.db.repositories.user_logs import UserLogsRepository + + UserLogsRepository(pg_conn).insert( + user_id="owner", endpoint="stream_answer", data={"question": "old row"} + ) + (row,) = _logs(app, pg_conn, "owner", {}).json["logs"] + assert "trace" not in row + + def test_webhook_row_links_by_activity_id(self, app, pg_conn): + from docsgpt.storage.db.repositories.stack_logs import StackLogsRepository + + StackLogsRepository(pg_conn).insert( + activity_id="act-1", endpoint="webhook", level="info", user_id="owner", query="{}" + ) + _trace(pg_conn, request_id="task-1", source="webhook", activity_id="act-1") + (row,) = _logs(app, pg_conn, "owner", {"event_type": "webhook"}).json["logs"] + assert row["activity_id"] == "act-1" + assert row["trace"]["ref"] == {"field": "activity_id", "value": "act-1"} + + def test_search_traces_are_listed(self, app, pg_conn): + record = _trace( + pg_conn, + request_id=None, + source="mcp", + name="mcp", + status="error", + summary={"llm_calls": 0, "query": "how do I deploy"}, + ) + _trace(pg_conn, request_id=None, source="graph_extraction", name="graph_extraction Docs") + search_rows = _logs(app, pg_conn, "owner", {"event_type": "search"}).json["logs"] + assert len(search_rows) == 1 + row = search_rows[0] + assert row["id"] == f"search-{record['id']}" + assert row["question"] == "how do I deploy" + assert row["level"] == "error" + assert row["action"] == "mcp" + assert row["trace"]["ref"] == {"field": "id", "value": record["id"]} + graph_rows = _logs(app, pg_conn, "owner", {"event_type": "graph"}).json["logs"] + assert len(graph_rows) == 1 + + def test_search_traces_of_other_users_are_hidden(self, app, pg_conn): + _trace(pg_conn, request_id=None, source="search", user_id="someone") + assert _logs(app, pg_conn, "owner", {"event_type": "search"}).json["logs"] == [] + + def test_unknown_event_type_is_rejected(self, app, pg_conn): + assert _logs(app, pg_conn, "owner", {"event_type": "bogus"}).status_code == 400 + + +class TestSummaryFailure: + def test_page_survives_a_failed_trace_lookup(self, app, pg_conn): + from docsgpt.storage.db.repositories.user_logs import UserLogsRepository + + UserLogsRepository(pg_conn).insert( + user_id="owner", + endpoint="stream_answer", + data={"question": "q", "request_id": "req-1"}, + ) + with patch( + "docsgpt.api.user.analytics.routes.RequestTracesRepository.summaries_for_refs", + side_effect=RuntimeError("statement timeout"), + ): + response = _logs(app, pg_conn, "owner", {}) + assert response.status_code == 200 + (row,) = response.json["logs"] + assert "trace" not in row + + +class TestFailedChatTurns: + def _failed_turn(self, pg_conn, *, with_trace): + from docsgpt.storage.db.repositories.stack_logs import StackLogsRepository + from docsgpt.storage.db.repositories.user_logs import UserLogsRepository + + StackLogsRepository(pg_conn).insert( + activity_id="act-chat", endpoint="stream", level="error", user_id="owner", query="q" + ) + UserLogsRepository(pg_conn).insert( + user_id="owner", + endpoint="stream_answer", + data={ + "question": "q", + "level": "error", + "error": "RuntimeError: upstream down", + "request_id": "req-f", + }, + ) + if with_trace: + _trace(pg_conn, request_id="req-f", activity_id="act-chat", status="error") + + def test_failed_turn_is_one_chat_row_not_also_a_system_row(self, app, pg_conn): + self._failed_turn(pg_conn, with_trace=True) + rows = _logs(app, pg_conn, "owner", {}).json["logs"] + assert [r["event_type"] for r in rows] == ["chat"] + assert rows[0]["level"] == "error" + assert rows[0]["error"] == "RuntimeError: upstream down" + assert rows[0]["trace"]["status"] == "error" + + def test_untraced_failures_keep_their_system_row(self, app, pg_conn): + """Rows from before traces (or with tracing off) have nothing to dedupe against.""" + self._failed_turn(pg_conn, with_trace=False) + rows = _logs(app, pg_conn, "owner", {}).json["logs"] + assert sorted(r["event_type"] for r in rows) == ["chat", "system"] diff --git a/tests/api/v1/test_v1_idempotency.py b/tests/api/v1/test_v1_idempotency.py index 08897b24..7715a886 100644 --- a/tests/api/v1/test_v1_idempotency.py +++ b/tests/api/v1/test_v1_idempotency.py @@ -19,6 +19,7 @@ These tests pin both layers against an ephemeral Postgres (``pg_engine``): from __future__ import annotations +import time import uuid from contextlib import contextmanager from typing import Any, Dict, List @@ -437,3 +438,53 @@ class TestV1IdempotencyRoute: {"n": TASK_NAME}, ).scalar() assert completed == 1 + + +@pytest.mark.integration +class TestV1ReplayTrace: + """A replayed request must not leave a failed trace behind.""" + + QUESTION = TestV1IdempotencyRoute.QUESTION + + def test_replay_discards_its_setup_trace(self, pg_engine, monkeypatch): + from docsgpt import tracing + from docsgpt.core.settings import settings + + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + monkeypatch.setattr(settings, "TRACES_OTEL_EXPORT", False) + user_id = f"user-{uuid.uuid4().hex[:8]}" + api_key = f"key-{uuid.uuid4().hex[:8]}" + with pg_engine.begin() as conn: + _seed_user(conn, user_id) + _seed_agent(conn, user_id, api_key) + + def _build_agent(processor, question): # noqa: ARG001 + # What the real setup does: a trace holding the pre-fetch retrieval. + processor.trace = tracing.start_trace(source="v1", capture_otel_context=False) + with tracing.activate(processor.trace): + with tracing.span(tracing.KIND_RETRIEVAL, "retrieval"): + pass + return _CountingAnswerAgent() + + written = [] + app = _build_app() + idem_key = f"idem-{uuid.uuid4().hex}" + with _wire_v1_route_db(pg_engine, monkeypatch), patch( + "docsgpt.api.answer.services.stream_processor.StreamProcessor.build_agent", + _build_agent, + ), patch("docsgpt.tracing.sink._write", side_effect=written.append): + _wire_idem_db(pg_engine, monkeypatch) + with app.test_client() as c: + resp1 = _post_chat(c, self.QUESTION, api_key, idem_key=idem_key) + resp2 = _post_chat(c, self.QUESTION, api_key, idem_key=idem_key) + + assert resp1.status_code == resp2.status_code == 200 + # The first request's trace is written on the background writer, so + # wait for it. A replay's trace is dropped synchronously, and an + # unclaimed-trace flush would be synchronous too, so any stray write + # is already in ``written`` by the time the responses return. + deadline = time.monotonic() + 5 + while not written and time.monotonic() < deadline: + time.sleep(0.02) + # Only the first request's trace is written, and it is not a failure. + assert [t.status for t in written] == ["ok"] diff --git a/tests/conftest.py b/tests/conftest.py index 876e6d89..f5bee6a6 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -44,6 +44,10 @@ import os os.environ.setdefault("AUTO_MIGRATE", "false") os.environ.setdefault("AUTO_CREATE_DB", "false") os.environ.setdefault("AUTO_VECTOR_SCHEMA", "false") +# Execution traces write to ``request_traces`` through their own DB session +# when a request finishes. Off by default so unrelated tests never touch a +# database through that path; tracing tests switch it on explicitly. +os.environ.setdefault("TRACES_ENABLED", "false") import subprocess import sys diff --git a/tests/e2e/specs/tier-a/execution-traces.spec.ts b/tests/e2e/specs/tier-a/execution-traces.spec.ts new file mode 100644 index 00000000..d0474961 --- /dev/null +++ b/tests/e2e/specs/tier-a/execution-traces.spec.ts @@ -0,0 +1,188 @@ +/** + * Execution traces — a chat turn's steps are stored and shown in the Logs UI. + * + * One `/stream` turn with a tool call writes a `request_traces` row holding + * the agent run, every LLM call and the tool call. The Logs row links to it + * through `request_id` (stamped into `user_logs.data`), carries a trace + * summary, and "View trace" opens the waterfall. + * + * // Silent-break covered: the Logs row loses its trace link. If the + * request id stops reaching `user_logs` (or the stored trace), the row + * silently renders without the "View trace" button and nothing errors. + */ + +import * as playwright from '@playwright/test'; +const { expect, test } = playwright; + +import { dirname, resolve } from 'node:path'; +import { fileURLToPath } from 'node:url'; + +import { authedRequest } from '../../helpers/api.js'; +import { newUserContext } from '../../helpers/auth.js'; +import { pg } from '../../helpers/db.js'; +import { resetDb } from '../../helpers/reset.js'; +import { + multipartContext, + postUpload, + waitForTask, +} from '../../helpers/uploads.js'; + +const HERE = dirname(fileURLToPath(import.meta.url)); +const README_MD = resolve(HERE, '..', '..', 'fixtures', 'docs', 'readme.md'); + +/** `memory_view` ships in DEFAULT_CHAT_TOOLS, so every user can run it. */ +const TOOL_DIRECTIVE = '[[MOCK_LLM_TOOLCALL:memory_view:once]]'; + +interface TraceRow { + source: string; + status: string; + request_id: string | null; + message_id: string | null; + // pg parses JSONB to JS values + // eslint-disable-next-line @typescript-eslint/no-explicit-any + spans: any[]; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + summary: any; +} + +async function tracesFor(userId: string): Promise { + const { rows } = await pg.query( + `SELECT source, status, request_id, message_id::text AS message_id, spans, summary + FROM request_traces WHERE user_id = $1 ORDER BY started_at`, + [userId], + ); + return rows; +} + +test.describe('tier-a · execution traces', () => { + test.beforeEach(async () => { + await resetDb(); + }); + + test('a chat turn with a tool call is traced and opens from the Logs page', async ({ + browser, + }) => { + const { context, sub, token } = await newUserContext(browser); + const api = await authedRequest(playwright, token); + try { + const res = await api.post('/stream', { + data: { + question: `what do you remember about me? ${TOOL_DIRECTIVE}`, + save_conversation: true, + }, + timeout: 180_000, + }); + expect(res.ok()).toBeTruthy(); + expect(await res.text()).toContain('"type": "end"'); + + // Stored trace: one row for the turn, linked to its message. It is + // written just after the stream closes, so wait for it. + await expect.poll(async () => (await tracesFor(sub)).length).toBe(1); + const traces = await tracesFor(sub); + const [trace] = traces; + expect(trace.source).toBe('stream'); + expect(trace.request_id).toBeTruthy(); + expect(trace.message_id).toBeTruthy(); + const kinds = new Set(trace.spans.map((s) => s.kind)); + expect(kinds.has('agent')).toBe(true); + expect(kinds.has('llm')).toBe(true); + expect(kinds.has('tool')).toBe(true); + expect(trace.summary.llm_calls).toBeGreaterThan(0); + const tool = trace.spans.find((s) => s.kind === 'tool'); + expect(tool.attributes['gen_ai.tool.name']).toBe('memory_view'); + + // The Logs API links the row to the trace. + const logsRes = await api.post('/api/get_user_logs', { + data: { page: 1, page_size: 10, event_type: 'chat' }, + }); + const logs = (await logsRes.json()).logs; + expect(logs[0].request_id).toBe(trace.request_id); + expect(logs[0].trace.ref).toEqual({ + field: 'request_id', + value: trace.request_id, + }); + + // And the UI opens it. + const page = await context.newPage(); + await page.goto('/settings/logs'); + await page.getByText('what do you remember about me?').first().click(); + await page.getByRole('button', { name: 'View trace' }).click(); + const sheet = page.getByRole('dialog', { name: 'Execution trace' }); + await expect(sheet).toBeVisible(); + await expect(sheet.getByText('execute_tool memory_view').first()).toBeVisible(); + await sheet.getByText('execute_tool memory_view').first().click(); + await expect(sheet.getByText('Arguments')).toBeVisible(); + if (process.env.TRACE_SCREENSHOT) { + await page.screenshot({ path: process.env.TRACE_SCREENSHOT, fullPage: false }); + } + + // Closing returns focus to the button that opened the drawer. + await page.keyboard.press('Escape'); + await expect(sheet).toBeHidden(); + await expect(page.getByRole('button', { name: 'View trace' })).toBeFocused(); + } finally { + await api.dispose(); + await context.close(); + } + }); + + test('a RAG turn records the retrieval, its query embedding and each source search', async ({ + browser, + }) => { + const { context, sub, token } = await newUserContext(browser); + const api = await authedRequest(playwright, token); + const multi = await multipartContext(token); + try { + const taskId = await postUpload(multi, README_MD, { + user: sub, + name: 'trace-docs', + mimeType: 'text/markdown', + }); + expect((await waitForTask(api, taskId, 120_000)).status).toBe('SUCCESS'); + const { rows } = await pg.query<{ id: string }>( + `SELECT id::text AS id FROM sources WHERE user_id = $1 ORDER BY date DESC LIMIT 1`, + [sub], + ); + const sourceId = rows[0].id; + + const res = await api.post('/stream', { + data: { + question: 'how do I get started?', + active_docs: sourceId, + save_conversation: true, + }, + timeout: 180_000, + }); + expect(await res.text()).toContain('"type": "end"'); + + await expect.poll(async () => (await tracesFor(sub)).length).toBe(1); + const [trace] = await tracesFor(sub); + const byKind = (kind: string) => trace.spans.filter((s) => s.kind === kind); + const retrieval = byKind('retrieval'); + expect(retrieval.length).toBeGreaterThan(0); + expect(byKind('embedding').length).toBeGreaterThan(0); + const searches = byKind('search'); + expect(searches.map((s) => s.attributes['gen_ai.data_source.id'])).toContain(sourceId); + // Retrieval ran before the agent answered, inside the same trace. + const agent = byKind('agent')[0]; + expect(retrieval[0].offset_ms).toBeLessThan(agent.offset_ms); + expect(trace.summary.retrieval_calls).toBeGreaterThan(0); + + const page = await context.newPage(); + await page.goto('/settings/logs'); + await page.getByText('how do I get started?').first().click(); + await page.getByRole('button', { name: 'View trace' }).click(); + const sheet = page.getByRole('dialog', { name: 'Execution trace' }); + await expect(sheet).toBeVisible(); + await sheet.getByText(/^retrieval/).first().click(); + await expect(sheet.getByText('Retrieved chunks')).toBeVisible(); + if (process.env.TRACE_SCREENSHOT_RAG) { + await page.screenshot({ path: process.env.TRACE_SCREENSHOT_RAG, fullPage: false }); + } + } finally { + await multi.dispose(); + await api.dispose(); + await context.close(); + } + }); +}); diff --git a/tests/guardrails/test_engine.py b/tests/guardrails/test_engine.py index b3316804..420192de 100644 --- a/tests/guardrails/test_engine.py +++ b/tests/guardrails/test_engine.py @@ -340,3 +340,71 @@ def test_action_enum_serialises_as_value(): dumped = config.model_dump(mode="json") assert dumped["controls"][0]["action"] == Action.REDACT.value assert dumped["controls"][0]["stage"] == Stage.INPUT.value + + +class QuickRemoteCheck(GuardrailCheck): + name = "_test_quick_remote" + label = "Quick remote" + supported_stages = {Stage.INPUT, Stage.OUTPUT} + remote = True + + def scan(self, text, stage, context): + from docsgpt import tracing + + with tracing.span(tracing.KIND_LLM, "chat judge"): + pass + return CheckOutcome.clean() + + +class TestTraceSpans: + """Guardrail evaluations appear in the execution trace only when they matter.""" + + @pytest.fixture(autouse=True) + def _trace(self, monkeypatch): + from docsgpt import tracing + from docsgpt.core.settings import settings + + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + GuardrailCreator.register(QuickRemoteCheck.name, QuickRemoteCheck) + self.trace = tracing.start_trace(source="stream", capture_otel_context=False) + with tracing.activate(self.trace): + yield + GuardrailCreator.checks.pop(QuickRemoteCheck.name, None) + + def test_clean_local_scan_is_not_traced(self): + engine = GuardrailEngine( + _config(controls=[{"check": "_test_never", "stage": "output", "action": "block"}]) + ) + engine.evaluate("text", Stage.OUTPUT) + assert self.trace.spans == [] + assert self.trace.content_blocked is False + + def test_firing_local_scan_is_traced_and_blocks_previews(self): + engine = GuardrailEngine(_config()) + engine.evaluate("some text", Stage.INPUT) + (span,) = self.trace.spans + assert span.kind == "guardrail" + assert span.attributes["docsgpt.guardrail.blocked"] is True + assert span.attributes["docsgpt.guardrail.triggered"] == ["_test_always"] + assert self.trace.content_blocked is True + + def test_a_repeating_firing_is_recorded_once(self): + """The output guard re-scans every segment; one firing must not fill the span cap.""" + engine = GuardrailEngine( + _config(controls=[{"check": "_test_always", "stage": "output", "action": "flag"}]) + ) + for _ in range(50): + engine.evaluate("segment", Stage.OUTPUT) + assert len(self.trace.spans) == 1 + assert self.trace.content_blocked is True + + def test_remote_scan_is_traced_with_judge_nested(self): + engine = GuardrailEngine( + _config(controls=[{"check": "_test_quick_remote", "stage": "input", "action": "block"}]) + ) + decision = engine.evaluate("text", Stage.INPUT) + assert decision.clean + guard, judge = self.trace.spans + assert guard.kind == "guardrail" + assert judge.parent_id == guard.id + assert self.trace.content_blocked is False diff --git a/tests/retriever/test_graph_rag.py b/tests/retriever/test_graph_rag.py index 5203eee1..802bfc99 100644 --- a/tests/retriever/test_graph_rag.py +++ b/tests/retriever/test_graph_rag.py @@ -1021,3 +1021,41 @@ class TestPersonalizedPageRankWithoutScipy: docs = rag._get_data() assert [doc["text"] for doc in docs] == ["near", "far"] + + +class TestTraceSpans: + """GraphRAG searches and query embeddings are recorded in the execution trace.""" + + @pytest.fixture(autouse=True) + def _trace(self, monkeypatch): + from docsgpt import tracing + from docsgpt.core.settings import settings + + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + monkeypatch.setattr(settings, "TRACES_CAPTURE_CONTENT", True) + self.trace = tracing.start_trace(source="stream", capture_otel_context=False) + with tracing.activate(self.trace): + yield + + def test_search_is_a_retrieval_span(self, _patch_llm_creator): + rag = _make_retriever() + docs = [{"text": "alpha", "title": "Doc A", "source": "a.md"}] + with patch.object(rag, "_get_data", return_value=docs): + assert rag.search("new question") == docs + (span,) = [s for s in self.trace.spans if s.kind == "retrieval"] + assert span.name == "retrieval GraphRAGRetriever" + assert span.attributes["docsgpt.retriever"] == "GraphRAGRetriever" + assert span.attributes["gen_ai.data_source.id"] == "src1" + assert span.attributes["docsgpt.chunk_count"] == 1 + assert span.previews["chunks"][0]["title"] == "Doc A" + + def test_embed_query_is_an_embedding_span(self): + embedder = Mock() + embedder.embed_query.return_value = [0.5] + with patch( + "docsgpt.retriever.graph_rag.get_embeddings", return_value=embedder + ): + assert GraphRAGRetriever._embed_query(object(), "q") == [0.5] + (span,) = self.trace.spans + assert span.kind == "embedding" + assert span.attributes["gen_ai.operation.name"] == "embeddings" diff --git a/tests/services/test_mcp_server.py b/tests/services/test_mcp_server.py index d8441db5..b048eb77 100644 --- a/tests/services/test_mcp_server.py +++ b/tests/services/test_mcp_server.py @@ -99,7 +99,7 @@ class TestSearchDocsTool: ): out = await search_docs(query="q", chunks=7) assert out == hits - mock_search.assert_called_once_with("the-key", "q", 7) + mock_search.assert_called_once_with("the-key", "q", 7, source="mcp") @pytest.mark.asyncio async def test_default_chunks_is_5(self): @@ -115,7 +115,7 @@ class TestSearchDocsTool: ) as mock_search, ): await search_docs(query="q") - mock_search.assert_called_once_with("k", "q", 5) + mock_search.assert_called_once_with("k", "q", 5, source="mcp") @pytest.mark.asyncio async def test_bearer_scheme_case_insensitive(self): @@ -131,4 +131,4 @@ class TestSearchDocsTool: ) as mock_search, ): await search_docs(query="q") - mock_search.assert_called_once_with("lowercase-scheme", "q", 5) + mock_search.assert_called_once_with("lowercase-scheme", "q", 5, source="mcp") diff --git a/tests/storage/db/repositories/test_request_traces.py b/tests/storage/db/repositories/test_request_traces.py new file mode 100644 index 00000000..7d380188 --- /dev/null +++ b/tests/storage/db/repositories/test_request_traces.py @@ -0,0 +1,200 @@ +"""Tests for RequestTracesRepository against a real Postgres instance.""" + +from __future__ import annotations + +import time +import uuid + +from sqlalchemy import text + +from docsgpt.storage.db.repositories.conversations import ConversationsRepository +from docsgpt.storage.db.repositories.request_traces import RequestTracesRepository + + +def _record(**overrides): + base = { + "id": str(uuid.uuid4()), + "request_id": "req-1", + "message_id": None, + "conversation_id": None, + "activity_id": None, + "workflow_run_id": None, + "user_id": "u1", + "agent_id": None, + "source": "stream", + "name": "stream", + "status": "ok", + "started_at_ns": time.time_ns(), + "duration_ms": 1234, + "span_count": 1, + "dropped_spans": 0, + "summary": {"llm_calls": 1}, + "spans": [ + { + "id": "s1", + "parent_id": None, + "kind": "llm", + "name": "chat m", + "status": "ok", + "offset_ms": 1.0, + "duration_ms": 10.0, + "attributes": {"gen_ai.request.model": "m"}, + } + ], + "otel_trace_id": None, + } + base.update(overrides) + return base + + +def _message(pg_conn, user_id="u1"): + convs = ConversationsRepository(pg_conn) + conv = convs.create(user_id, "t") + msg = convs.reserve_message( + str(conv["id"]), prompt="q", placeholder_response="..." + ) + return str(conv["id"]), str(msg["id"]) + + +class TestInsert: + def test_round_trip(self, pg_conn): + repo = RequestTracesRepository(pg_conn) + record = _record() + assert repo.insert(record) is True + rows = repo.list_by_ref("request_id", "req-1", user_id="u1") + assert len(rows) == 1 + row = rows[0] + assert row["id"] == record["id"] + assert row["duration_ms"] == 1234 + assert row["summary"] == {"llm_calls": 1} + assert row["spans"][0]["attributes"]["gen_ai.request.model"] == "m" + + def test_strips_nul_bytes(self, pg_conn): + repo = RequestTracesRepository(pg_conn) + record = _record( + spans=[{"id": "s", "name": "a\x00b", "preview": {"result": "x\x00y"}}] + ) + assert repo.insert(record) + row = repo.list_by_ref("id", record["id"], user_id="u1")[0] + assert row["spans"][0]["name"] == "ab" + assert row["spans"][0]["preview"]["result"] == "xy" + + def test_non_uuid_link_ids_become_null(self, pg_conn): + repo = RequestTracesRepository(pg_conn) + record = _record(agent_id="not-a-uuid", conversation_id="nope") + assert repo.insert(record) + row = repo.list_by_ref("id", record["id"], user_id="u1")[0] + assert row["agent_id"] is None + assert row["conversation_id"] is None + + def test_missing_message_drops_row_and_keeps_transaction_usable(self, pg_conn): + repo = RequestTracesRepository(pg_conn) + assert repo.insert(_record(message_id=str(uuid.uuid4()))) is False + # The savepoint rolled back; the outer transaction still works. + assert repo.insert(_record(request_id="req-after")) is True + + def test_duplicate_id_is_ignored(self, pg_conn): + repo = RequestTracesRepository(pg_conn) + record = _record() + assert repo.insert(record) is True + assert repo.insert(record) is False + + +class TestScoping: + def test_other_users_cannot_read(self, pg_conn): + repo = RequestTracesRepository(pg_conn) + repo.insert(_record()) + assert repo.list_by_ref("request_id", "req-1", user_id="u2") == [] + + def test_agent_scope_returns_agent_traces_from_any_user(self, pg_conn): + repo = RequestTracesRepository(pg_conn) + agent_id = str(uuid.uuid4()) + repo.insert(_record(user_id="someone-else", agent_id=agent_id)) + rows = repo.list_by_ref("request_id", "req-1", user_id="owner", agent_id=agent_id) + assert len(rows) == 1 + + def test_unknown_field_and_bad_uuid_return_empty(self, pg_conn): + repo = RequestTracesRepository(pg_conn) + repo.insert(_record()) + assert repo.list_by_ref("user_id", "u1", user_id="u1") == [] + assert repo.list_by_ref("message_id", "not-uuid", user_id="u1") == [] + assert repo.list_by_ref("request_id", "req-1", user_id=None) == [] + + +class TestMessageLink: + def test_rounds_for_a_message_are_ordered(self, pg_conn): + conv_id, msg_id = _message(pg_conn) + repo = RequestTracesRepository(pg_conn) + first = _record(message_id=msg_id, status="paused", started_at_ns=time.time_ns()) + second = _record(message_id=msg_id, started_at_ns=time.time_ns() + 1_000_000) + repo.insert(second) + repo.insert(first) + rows = repo.list_by_ref("message_id", msg_id, user_id="u1") + assert [r["id"] for r in rows] == [first["id"], second["id"]] + + def test_deleting_the_conversation_cascades(self, pg_conn): + conv_id, msg_id = _message(pg_conn) + repo = RequestTracesRepository(pg_conn) + repo.insert(_record(message_id=msg_id, conversation_id=conv_id)) + ConversationsRepository(pg_conn).delete(conv_id, "u1") + count = pg_conn.execute(text("SELECT count(*) FROM request_traces")).scalar() + assert count == 0 + + +class TestConversationDeletion: + """Every trace of a conversation goes with it, not only message-linked ones.""" + + def test_trace_without_message_is_deleted_with_its_conversation(self, pg_conn): + conv_id, _msg_id = _message(pg_conn) + repo = RequestTracesRepository(pg_conn) + # e.g. a scheduled run in this conversation, or a stateless /v1 round. + repo.insert(_record(conversation_id=conv_id, source="schedule")) + repo.insert(_record(request_id="other")) # another conversation's trace + ConversationsRepository(pg_conn).delete(conv_id, "u1") + remaining = pg_conn.execute(text("SELECT request_id FROM request_traces")).scalars().all() + assert remaining == ["other"] + + def test_delete_all_for_user_removes_their_conversation_traces(self, pg_conn): + conv_id, _msg_id = _message(pg_conn) + repo = RequestTracesRepository(pg_conn) + repo.insert(_record(conversation_id=conv_id)) + ConversationsRepository(pg_conn).delete_all_for_user("u1") + assert pg_conn.execute(text("SELECT count(*) FROM request_traces")).scalar() == 0 + + +class TestSummaries: + def test_summaries_grouped_by_field_and_id(self, pg_conn): + repo = RequestTracesRepository(pg_conn) + a = _record(request_id="r-a") + b = _record(request_id="r-b", activity_id="act-1") + repo.insert(a) + repo.insert(b) + out = repo.summaries_for_refs( + {"request_id": ["r-a", "r-b", "r-missing"], "activity_id": ["act-1"]}, + user_id="u1", + ) + assert set(out["request_id"]) == {"r-a", "r-b"} + assert out["activity_id"]["act-1"][0]["id"] == b["id"] + assert "spans" not in out["request_id"]["r-a"][0] + + def test_summaries_are_scoped(self, pg_conn): + repo = RequestTracesRepository(pg_conn) + repo.insert(_record(request_id="r-a")) + assert repo.summaries_for_refs({"request_id": ["r-a"]}, user_id="u2") == {} + + +class TestPurge: + def test_purge_older_than(self, pg_conn): + repo = RequestTracesRepository(pg_conn) + old = _record(request_id="old") + repo.insert(old) + repo.insert(_record(request_id="new")) + pg_conn.execute( + text( + "UPDATE request_traces SET created_at = now() - interval '40 days' " + "WHERE request_id = 'old'" + ) + ) + assert repo.purge_older_than(30) == 1 + assert repo.list_by_ref("request_id", "old", user_id="u1") == [] + assert len(repo.list_by_ref("request_id", "new", user_id="u1")) == 1 diff --git a/tests/storage/db/test_migration_0037.py b/tests/storage/db/test_migration_0037.py new file mode 100644 index 00000000..e8838a74 --- /dev/null +++ b/tests/storage/db/test_migration_0037.py @@ -0,0 +1,78 @@ +"""Migration round-trip test for 0037_request_traces.""" + +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path + +import pytest +from sqlalchemy import text + + +pytestmark = pytest.mark.integration + + +def _alembic_ini() -> Path: + return Path(__file__).resolve().parents[3] / "docsgpt" / "alembic.ini" + + +def _run_alembic(url: str, *args: str) -> None: + subprocess.check_call( + [sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), *args], + timeout=60, + env={**os.environ, "POSTGRES_URI": url}, + ) + + +def _alembic_version(conn) -> str: + return conn.execute(text("SELECT version_num FROM alembic_version")).scalar() + + +def _table_exists(conn, table: str) -> bool: + return conn.execute(text("SELECT to_regclass(:t)"), {"t": f"public.{table}"}).scalar() is not None + + +_0037 = "0037_request_traces" +_0036 = "0036_device_audit_created_idx" + + +class TestMigration0037RoundTrip: + def test_head_has_request_traces(self, pg_engine): + with pg_engine.connect() as conn: + assert _alembic_version(conn) >= _0037 + assert _table_exists(conn, "request_traces") + + def test_downgrade_drops_then_upgrade_restores(self, pg_engine): + url = pg_engine.url.render_as_string(hide_password=False) + _run_alembic(url, "downgrade", _0036) + with pg_engine.connect() as conn: + assert _alembic_version(conn) == _0036 + assert not _table_exists(conn, "request_traces") + _run_alembic(url, "upgrade", "head") + with pg_engine.connect() as conn: + assert _table_exists(conn, "request_traces") + + def test_status_check_rejects_unknown_status(self, pg_engine): + with pytest.raises(Exception): + with pg_engine.begin() as conn: + conn.execute( + text( + "INSERT INTO request_traces (id, source, status, started_at) " + "VALUES (gen_random_uuid(), 'stream', 'weird', now())" + ) + ) + + def test_listing_and_deletion_indexes_exist(self, pg_engine): + with pg_engine.connect() as conn: + names = set( + conn.execute( + text("SELECT indexname FROM pg_indexes WHERE tablename = 'request_traces'") + ).scalars() + ) + assert { + "request_traces_conversation_idx", + "request_traces_user_source_started_idx", + "request_traces_agent_source_started_idx", + } <= names diff --git a/tests/test_logging.py b/tests/test_logging.py index 2a18573a..07fba72a 100644 --- a/tests/test_logging.py +++ b/tests/test_logging.py @@ -398,3 +398,96 @@ class TestAccumulateResponseSummary: assert ctx.tool_call_count == 2 +class TestLogActivityTraceSpan: + """``@log_activity`` opens the ``invoke_agent`` span for every agent run.""" + + @pytest.fixture(autouse=True) + def _tracing_on(self, monkeypatch): + from docsgpt.core.settings import settings + + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + + class _Agent: + endpoint = "stream" + user = "user1" + user_api_key = "" + agent_id = "agent-1" + model_id = "gpt-4o" + + def test_agent_span_wraps_the_run_and_binds_activity_id(self): + from docsgpt import tracing + from docsgpt.logging import log_activity + + @log_activity() + def gen(agent, log_context=None): + with tracing.span(tracing.KIND_LLM, "chat m"): + pass + yield {"answer": "hi"} + yield {"sources": [{"title": "a"}]} + + trace = tracing.start_trace(source="stream", capture_otel_context=False) + with patch("docsgpt.logging._log_activity_to_db"), tracing.activate(trace): + list(gen(self._Agent())) + agent_span, llm_span = trace.spans + assert agent_span.kind == tracing.KIND_AGENT + assert agent_span.name == "invoke_agent _Agent" + assert agent_span.status == "ok" + assert agent_span.attributes["gen_ai.operation.name"] == "invoke_agent" + assert agent_span.attributes["gen_ai.agent.id"] == "agent-1" + assert agent_span.attributes["docsgpt.answer_chars"] == 2 + # BaseAgent keeps its model in ``model_id``; the span reads it too. + assert agent_span.attributes["gen_ai.request.model"] == "gpt-4o" + assert agent_span.attributes["docsgpt.source_count"] == 1 + assert llm_span.parent_id == agent_span.id + assert trace.activity_id is not None + + def test_nested_agent_keeps_outer_activity_id(self): + from docsgpt import tracing + from docsgpt.logging import log_activity + + @log_activity() + def inner(agent, log_context=None): + yield "x" + + @log_activity() + def outer(agent, log_context=None): + yield log_context.activity_id + yield from inner(agent) + + trace = tracing.start_trace(source="stream", capture_otel_context=False) + with patch("docsgpt.logging._log_activity_to_db"), tracing.activate(trace): + outer_activity = list(outer(self._Agent()))[0] + assert trace.activity_id == outer_activity + outer_span, inner_span = trace.spans + assert inner_span.parent_id == outer_span.id + + def test_yielded_error_marks_span_error(self): + from docsgpt import tracing + from docsgpt.logging import log_activity + + @log_activity() + def gen(agent, log_context=None): + yield {"type": "error", "error": "node failed"} + + trace = tracing.start_trace(source="stream", capture_otel_context=False) + with patch("docsgpt.logging._log_activity_to_db"), tracing.activate(trace): + list(gen(self._Agent())) + assert trace.spans[0].status == "error" + assert trace.spans[0].error == "StreamError" + assert trace.spans[0].previews["error"] == "node failed" + + def test_raised_error_marks_span_error(self): + from docsgpt import tracing + from docsgpt.logging import log_activity + + @log_activity() + def gen(agent, log_context=None): + yield "a" + raise RuntimeError("boom") + + trace = tracing.start_trace(source="stream", capture_otel_context=False) + with patch("docsgpt.logging._log_activity_to_db"), tracing.activate(trace), \ + pytest.raises(RuntimeError): + list(gen(self._Agent())) + assert trace.spans[0].status == "error" + assert trace.spans[0].attributes["error.type"] == "RuntimeError" diff --git a/tests/tracing/__init__.py b/tests/tracing/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/tracing/test_entry_points.py b/tests/tracing/test_entry_points.py new file mode 100644 index 00000000..5a613892 --- /dev/null +++ b/tests/tracing/test_entry_points.py @@ -0,0 +1,191 @@ +"""Non-chat entry points each record and write one execution trace.""" + +from __future__ import annotations + +from contextlib import contextmanager +from unittest.mock import MagicMock, patch + +import pytest + +from docsgpt import tracing +from docsgpt.core.settings import settings + + +@pytest.fixture(autouse=True) +def _tracing_on(monkeypatch): + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + monkeypatch.setattr(settings, "TRACES_OTEL_EXPORT", False) + + +@pytest.fixture() +def flushed(): + """Capture flushed traces instead of writing them.""" + captured = [] + + def _fake_flush(trace, status=None, **_kwargs): + if trace is None or trace.flushed: + return + trace.flushed = True + trace.finish(status) + captured.append(trace) + + with patch("docsgpt.tracing.flush", side_effect=_fake_flush): + yield captured + + +def _headless(events, monkeypatch, **kwargs): + from docsgpt.agents import headless_runner as hr + + agent = MagicMock(name="agent") + + def _gen(query): + with tracing.span(tracing.KIND_AGENT, "invoke_agent Fake"): + yield from events + + agent.gen.side_effect = _gen + agent.llm.token_usage = {"prompt_tokens": 1, "generated_tokens": 1} + retriever = MagicMock(name="retriever") + + def _search(query): + with tracing.span(tracing.KIND_RETRIEVAL, "retrieval"): + return [] + + retriever.search.side_effect = _search + tool_executor = MagicMock(name="tool_executor") + tool_executor.headless_denials = [] + monkeypatch.setattr(hr, "get_prompt", lambda _pid: "system prompt") + monkeypatch.setattr( + hr.RetrieverCreator, "create_retriever", classmethod(lambda cls, *a, **kw: retriever) + ) + monkeypatch.setattr(hr, "ToolExecutor", lambda *a, **kw: tool_executor) + monkeypatch.setattr( + hr.AgentCreator, "create_agent", classmethod(lambda cls, *a, **kw: agent) + ) + with patch("docsgpt.core.model_utils.validate_model_id", return_value=True), \ + patch("docsgpt.core.model_utils.get_default_model_id", return_value="m"), \ + patch("docsgpt.core.model_utils.get_provider_from_model_id", return_value="openai"), \ + patch("docsgpt.core.model_utils.get_api_key_for_provider", return_value="k"), \ + patch("docsgpt.utils.calculate_doc_token_budget", return_value=1000): + return hr.run_agent_headless( + {"user_id": "owner-1", "id": "11111111-1111-1111-1111-111111111111"}, + "do the thing", + **kwargs, + ) + + +@pytest.mark.unit +class TestHeadless: + def test_scheduled_run_is_traced_with_its_run_id(self, monkeypatch, flushed): + _headless([{"answer": "done"}], monkeypatch, endpoint="schedule", request_id="run-1") + (trace,) = flushed + assert trace.source == "schedule" + assert trace.request_id == "run-1" + assert trace.user_id == "owner-1" + assert trace.status == "ok" + assert [s.kind for s in trace.spans] == ["retrieval", "agent"] + + def test_trace_can_belong_to_the_scheduling_user(self, monkeypatch, flushed): + """A schedule on a shared agent is the scheduler's run, not the owner's.""" + _headless( + [{"answer": "done"}], + monkeypatch, + endpoint="schedule", + request_id="run-2", + trace_user_id="scheduler-user", + ) + assert flushed[0].user_id == "scheduler-user" + + def test_stream_error_marks_trace_error(self, monkeypatch, flushed): + outcome = _headless([{"type": "error", "error": "boom"}], monkeypatch, endpoint="webhook") + assert outcome["error_type"] == "stream_error" + assert flushed[0].status == "error" + + def test_raised_error_still_flushes(self, monkeypatch, flushed): + from docsgpt.agents import headless_runner as hr + + monkeypatch.setattr(hr, "_run_agent_headless", MagicMock(side_effect=RuntimeError("x"))) + with pytest.raises(RuntimeError): + hr.run_agent_headless({"user_id": "u"}, "q") + assert flushed[0].status == "error" + + def test_trace_request_id_stays_off_llm_usage_rows(self, monkeypatch, flushed): + """The headless LLM's own request id is untouched (quota counts depend on it).""" + from docsgpt.agents import headless_runner as hr + + created = {} + + def _create(cls, *a, **kw): + agent = MagicMock() + agent.gen.return_value = iter([{"answer": "x"}]) + agent.llm.token_usage = {} + agent.llm._request_id = None + created["agent"] = agent + return agent + + monkeypatch.setattr(hr.AgentCreator, "create_agent", classmethod(_create)) + monkeypatch.setattr(hr, "get_prompt", lambda _pid: "p") + monkeypatch.setattr( + hr.RetrieverCreator, "create_retriever", + classmethod(lambda cls, *a, **kw: MagicMock(search=MagicMock(return_value=[]))), + ) + monkeypatch.setattr(hr, "ToolExecutor", lambda *a, **kw: MagicMock(headless_denials=[])) + with patch("docsgpt.core.model_utils.validate_model_id", return_value=True), \ + patch("docsgpt.core.model_utils.get_default_model_id", return_value="m"), \ + patch("docsgpt.core.model_utils.get_provider_from_model_id", return_value="openai"), \ + patch("docsgpt.core.model_utils.get_api_key_for_provider", return_value="k"), \ + patch("docsgpt.utils.calculate_doc_token_budget", return_value=1000): + hr.run_agent_headless({"user_id": "u"}, "q", endpoint="schedule", request_id="run-9") + assert created["agent"].llm._request_id is None + + +@contextmanager +def _search_env(agent): + repo = MagicMock() + repo.find_by_key.return_value = agent + + @contextmanager + def _conn(): + yield MagicMock() + + store = MagicMock() + store.search.return_value = [{"text": "hit", "metadata": {"title": "T", "source": "s"}}] + with patch("docsgpt.api.user.team_sharing.can_access", return_value=True), \ + patch("docsgpt.services.search_service.db_readonly", _conn), \ + patch("docsgpt.services.search_service.AgentsRepository", return_value=repo), \ + patch( + "docsgpt.services.search_service.VectorCreator.create_vectorstore", + return_value=store, + ): + yield + + +@pytest.mark.unit +class TestSearch: + def test_search_is_traced(self, flushed): + from docsgpt.services.search_service import search + + agent = {"id": "a-1", "source_id": "src-1", "extra_source_ids": [], "user_id": "owner"} + with _search_env(agent): + results = search("k", "what is x", 3) + assert results + (trace,) = flushed + assert trace.source == "search" + assert trace.user_id == "owner" + retrieval = trace.spans[0] + assert retrieval.kind == tracing.KIND_RETRIEVAL + assert retrieval.attributes["docsgpt.chunk_count"] == 1 + + def test_mcp_source_name(self, flushed): + from docsgpt.services.search_service import search + + agent = {"id": "a-1", "source_id": "src-1", "extra_source_ids": [], "user_id": "owner"} + with _search_env(agent): + search("k", "q", 3, source="mcp") + assert flushed[0].source == "mcp" + + def test_no_sources_records_nothing(self, flushed): + from docsgpt.services.search_service import search + + with _search_env({"id": "a-1", "source_id": None, "user_id": "owner"}): + assert search("k", "q", 3) == [] + assert flushed == [] diff --git a/tests/tracing/test_llm_spans.py b/tests/tracing/test_llm_spans.py new file mode 100644 index 00000000..966719de --- /dev/null +++ b/tests/tracing/test_llm_spans.py @@ -0,0 +1,279 @@ +"""LLM calls become ``chat`` spans through the token-usage wrappers.""" + +from __future__ import annotations + +from unittest.mock import patch + +import pytest + +from docsgpt import tracing +from docsgpt.cache import gen_cache, stream_cache +from docsgpt.core.settings import settings +from docsgpt.usage import gen_token_usage, stream_token_usage + + +class _LLM: + provider_name = "openai" + + def __init__(self, source=None): + self.token_usage = {"prompt_tokens": 0, "generated_tokens": 0} + self.decoded_token = {"sub": "u1"} + self.user_api_key = None + self.agent_id = None + if source: + self._token_usage_source = source + + +@pytest.fixture(autouse=True) +def _env(monkeypatch): + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + monkeypatch.setattr(settings, "TRACES_CAPTURE_CONTENT", True) + with patch("docsgpt.usage._persist_call_usage", return_value=0.0012): + yield + + +@pytest.fixture() +def trace(): + t = tracing.start_trace(source="stream", capture_otel_context=False) + with tracing.activate(t): + yield t + + +@pytest.fixture() +def metrics(): + with patch("docsgpt.tracing.llm.record_llm_metrics") as rec: + yield rec + + +class TestNonStreaming: + def test_span_with_usage_and_preview(self, trace, metrics): + @gen_token_usage + def _gen(self, model, messages, stream, tools, **kwargs): + return "the answer" + + _gen(_LLM(), "gpt-4o", [{"role": "user", "content": "hi"}], False, None) + (span,) = trace.spans + assert span.kind == tracing.KIND_LLM + assert span.name == "chat gpt-4o" + assert span.status == "ok" + attrs = span.attributes + assert attrs["gen_ai.operation.name"] == "chat" + assert attrs["gen_ai.provider.name"] == "openai" + assert attrs["gen_ai.request.model"] == "gpt-4o" + assert attrs["gen_ai.usage.input_tokens"] > 0 + assert attrs["gen_ai.usage.output_tokens"] > 0 + assert attrs["docsgpt.token_source"] == "agent_stream" + assert attrs["docsgpt.cost_usd"] == 0.0012 + assert span.previews["output"] == "the answer" + metrics.assert_called_once() + assert metrics.call_args.kwargs["error_type"] is None + + def test_failure_marks_span_error(self, trace, metrics): + @gen_token_usage + def _gen(self, model, messages, stream, tools, **kwargs): + raise TimeoutError("slow") + + with pytest.raises(TimeoutError): + _gen(_LLM(source="fallback"), "m", [], False, None) + (span,) = trace.spans + assert span.status == "error" + assert span.attributes["error.type"] == "TimeoutError" + assert span.attributes["docsgpt.token_source"] == "fallback" + assert metrics.call_args.kwargs["error_type"] == "TimeoutError" + + def test_provider_reported_usage_is_flagged(self, trace, metrics): + llm = _LLM() + + @gen_token_usage + def _gen(self, model, messages, stream, tools, **kwargs): + self._last_usage = { + "prompt_tokens": 100, + "completion_tokens": 7, + "prompt_tokens_details": {"cached_tokens": 40}, + } + self._last_usage_claimed = False + return "x" + + _gen(llm, "m", [], False, None) + attrs = trace.spans[0].attributes + assert attrs["gen_ai.usage.input_tokens"] == 100 + assert attrs["gen_ai.usage.output_tokens"] == 7 + assert attrs["gen_ai.usage.cache_read.input_tokens"] == 40 + assert attrs["docsgpt.usage_estimated"] is False + + def test_no_trace_no_span(self, metrics): + @gen_token_usage + def _gen(self, model, messages, stream, tools, **kwargs): + return "x" + + assert _gen(_LLM(), "m", [], False, None) == "x" + metrics.assert_called_once() + + +class TestDisabled: + def test_no_metrics_when_tracing_is_off(self, metrics, monkeypatch): + monkeypatch.setattr(settings, "TRACES_ENABLED", False) + + @gen_token_usage + def _gen(self, model, messages, stream, tools, **kwargs): + return "x" + + _gen(_LLM(), "m", [], False, None) + metrics.assert_not_called() + + +class TestStreaming: + def test_span_starts_on_first_next_not_on_call(self, trace, metrics): + @stream_token_usage + def _stream(self, model, messages, stream, tools, **kwargs): + yield "a" + yield {"type": "thought", "thought": "hmm"} + yield "b" + + gen = _stream(_LLM(), "m", [], True, None) + assert trace.spans == [] + assert list(gen) == ["a", {"type": "thought", "thought": "hmm"}, "b"] + (span,) = trace.spans + assert span.status == "ok" + assert span.attributes["docsgpt.ttft_ms"] is not None + assert span.attributes["docsgpt.stream"] is True + assert span.previews["output"] == "ab" + + def test_abandoned_stream_is_cancelled(self, trace, metrics): + @stream_token_usage + def _stream(self, model, messages, stream, tools, **kwargs): + yield "a" + yield "b" + + gen = _stream(_LLM(), "m", [], True, None) + next(gen) + gen.close() + assert trace.spans[0].status == "cancelled" + + def test_failed_stream(self, trace, metrics): + @stream_token_usage + def _stream(self, model, messages, stream, tools, **kwargs): + yield "a" + raise ConnectionError("reset") + + with pytest.raises(ConnectionError): + list(_stream(_LLM(), "m", [], True, None)) + assert trace.spans[0].status == "error" + + def test_primary_and_fallback_are_siblings(self, trace, metrics): + @stream_token_usage + def _primary(self, model, messages, stream, tools, **kwargs): + raise ConnectionError("down") + yield # pragma: no cover + + @stream_token_usage + def _fallback(self, model, messages, stream, tools, **kwargs): + yield "ok" + + with tracing.span(tracing.KIND_AGENT, "agent"): + with pytest.raises(ConnectionError): + list(_primary(_LLM(), "m", [], True, None)) + list(_fallback(_LLM(source="fallback"), "m2", [], True, None)) + agent, primary, fallback = trace.spans + assert primary.parent_id == agent.id == fallback.parent_id + assert primary.status == "error" + assert fallback.attributes["docsgpt.token_source"] == "fallback" + + +class _FakeRedis: + def __init__(self): + self.store = {} + + def get(self, key): + return self.store.get(key) + + def set(self, key, value, ex=None): + self.store[key] = value.encode("utf-8") if isinstance(value, str) else value + + def delete(self, key): + self.store.pop(key, None) + + +class TestCacheHits: + def test_gen_cache_hit_records_a_cached_span(self, trace, metrics): + redis = _FakeRedis() + + @gen_cache + @gen_token_usage + def _gen(self, model, messages, stream, tools=None, **kwargs): + return "fresh" + + with patch("docsgpt.cache.get_redis_instance", return_value=redis): + _gen(_LLM(), "m", [{"role": "user", "content": "q"}], False) + _gen(_LLM(), "m", [{"role": "user", "content": "q"}], False) + first, second = trace.spans + assert first.attributes.get("docsgpt.cache_hit") is None + assert second.attributes["docsgpt.cache_hit"] is True + assert second.status == "ok" + + def test_stream_cache_hit_flags_the_open_span(self, trace, metrics, monkeypatch): + monkeypatch.setattr("docsgpt.cache.time.sleep", lambda _s: None) + redis = _FakeRedis() + + @stream_token_usage + @stream_cache + def _stream(self, model, messages, stream, tools=None, **kwargs): + yield "fresh" + + with patch("docsgpt.cache.get_redis_instance", return_value=redis): + list(_stream(_LLM(), "m", [{"role": "user", "content": "q"}], True, None)) + list(_stream(_LLM(), "m", [{"role": "user", "content": "q"}], True, None)) + first, second = trace.spans + assert first.attributes.get("docsgpt.cache_hit") is None + assert second.attributes["docsgpt.cache_hit"] is True + + +class TestProviderResolution: + """The span names the provider actually called, not the client class used.""" + + @staticmethod + def _llm(provider="openai", base_url=None, plugin=None): + llm = _LLM() + llm.provider_name = provider + if base_url is not None: + llm._effective_base_url = base_url + if plugin is not None: + llm._provider_plugin = plugin + return llm + + @pytest.mark.parametrize( + "base_url, plugin, expected", + [ + ("https://api.deepseek.com/v1", "openai_compatible", "deepseek"), + ("https://my-res.openai.azure.com/openai", "openai", "azure.ai.openai"), + ("https://api.mistral.ai/v1", "openai_compatible", "mistral_ai"), + ("https://api.x.ai/v1", "openai_compatible", "x_ai"), + ("https://api.openai.com/v1", "openai", "openai"), + ("http://10.0.0.5:8000/v1", "openai_compatible", "openai_compatible"), + ("http://127.0.0.1:7899/v1", "openai", "openai_compatible"), + (None, "openai", "openai"), + ], + ) + def test_provider_from_endpoint(self, base_url, plugin, expected): + from docsgpt.tracing.llm import llm_provider + + assert llm_provider(self._llm(base_url=base_url, plugin=plugin)) == expected + + def test_native_providers_are_unchanged(self): + from docsgpt.tracing.llm import llm_provider + + assert llm_provider(self._llm(provider="anthropic")) == "anthropic" + assert llm_provider(self._llm(provider="google")) == "gcp.gen_ai" + + def test_span_and_metrics_use_it_and_record_server_address(self, trace, metrics): + llm = self._llm(base_url="https://api.deepseek.com/v1", plugin="openai_compatible") + + @gen_token_usage + def _gen(self, model, messages, stream, tools, **kwargs): + return "x" + + _gen(llm, "deepseek-chat", [], False, None) + attrs = trace.spans[0].attributes + assert attrs["gen_ai.provider.name"] == "deepseek" + assert attrs["server.address"] == "api.deepseek.com" + assert metrics.call_args.kwargs["provider"] == "deepseek" diff --git a/tests/tracing/test_otel_export.py b/tests/tracing/test_otel_export.py new file mode 100644 index 00000000..c93ca935 --- /dev/null +++ b/tests/tracing/test_otel_export.py @@ -0,0 +1,235 @@ +"""Tests for replaying a finished trace as OpenTelemetry GenAI spans.""" + +from __future__ import annotations + +import pytest +from opentelemetry import trace as ot_trace +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import InMemoryMetricReader +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import SpanKind, StatusCode + +from docsgpt import tracing +from docsgpt.core.settings import settings +from docsgpt.tracing import otel as trace_otel + + +@pytest.fixture(autouse=True) +def _settings(monkeypatch): + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + monkeypatch.setattr(settings, "TRACES_OTEL_EXPORT", True) + monkeypatch.setattr(settings, "TRACES_CAPTURE_CONTENT", True) + monkeypatch.delenv("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", raising=False) + + +@pytest.fixture() +def provider(): + exporter = InMemorySpanExporter() + tp = TracerProvider() + tp.add_span_processor(SimpleSpanProcessor(exporter)) + tp.exporter = exporter + return tp + + +def _sample_trace(otel_context=None): + trace = tracing.start_trace( + source="stream", + request_id="req-1", + message_id="msg-1", + conversation_id="conv-1", + agent_id="agent-1", + capture_otel_context=False, + ) + trace.otel_context = otel_context + with tracing.activate(trace): + with tracing.span( + tracing.KIND_AGENT, + "invoke_agent Support", + attributes={"gen_ai.operation.name": "invoke_agent", "gen_ai.agent.name": "Support"}, + ): + with tracing.span( + tracing.KIND_LLM, + "chat gpt-4o", + attributes={ + "gen_ai.operation.name": "chat", + "gen_ai.provider.name": "openai", + "gen_ai.request.model": "gpt-4o", + "gen_ai.usage.input_tokens": 12, + "gen_ai.usage.output_tokens": 3, + "docsgpt.sources": ["a", "b"], + "docsgpt.meta": {"nested": True}, + }, + ) as llm: + llm.preview("output", "hello") + with tracing.span( + tracing.KIND_TOOL, + "execute_tool search", + attributes={"gen_ai.tool.name": "search"}, + ) as tool: + tool.preview("arguments", {"q": "x"}) + tool.preview("result", "found") + tool.end(error=RuntimeError("tool broke")) + trace.finish() + return trace + + +class TestReplay: + def test_spans_are_parented_and_timed(self, provider): + trace = _sample_trace() + otel_trace_id = trace_otel.export_trace(trace, tracer_provider=provider) + spans = {s.name: s for s in provider.exporter.get_finished_spans()} + root = spans["docsgpt stream"] + agent = spans["invoke_agent Support"] + llm = spans["chat gpt-4o"] + tool = spans["execute_tool search"] + assert agent.parent.span_id == root.context.span_id + assert llm.parent.span_id == agent.context.span_id + assert tool.parent.span_id == agent.context.span_id + assert otel_trace_id == format(root.context.trace_id, "032x") + recorded = {s.name: s for s in trace.spans} + assert llm.start_time == trace.span_start_ns(recorded["chat gpt-4o"]) + assert llm.end_time == trace.span_end_ns(recorded["chat gpt-4o"]) + assert root.start_time == trace.start_ns + assert root.end_time >= agent.end_time + + def test_semconv_attributes_and_kinds(self, provider): + trace_otel.export_trace(_sample_trace(), tracer_provider=provider) + spans = {s.name: s for s in provider.exporter.get_finished_spans()} + llm = spans["chat gpt-4o"] + assert llm.kind == SpanKind.CLIENT + assert llm.attributes["gen_ai.usage.input_tokens"] == 12 + assert llm.attributes["gen_ai.conversation.id"] == "conv-1" + assert tuple(llm.attributes["docsgpt.sources"]) == ("a", "b") + assert llm.attributes["docsgpt.meta"] == '{"nested": true}' + root = spans["docsgpt stream"] + assert root.attributes["docsgpt.request_id"] == "req-1" + assert root.attributes["docsgpt.message_id"] == "msg-1" + + def test_error_status(self, provider): + trace_otel.export_trace(_sample_trace(), tracer_provider=provider) + tool = {s.name: s for s in provider.exporter.get_finished_spans()}["execute_tool search"] + assert tool.status.status_code == StatusCode.ERROR + assert tool.attributes["error.type"] == "RuntimeError" + + def test_no_content_by_default(self, provider): + trace_otel.export_trace(_sample_trace(), tracer_provider=provider) + for span in provider.exporter.get_finished_spans(): + assert not any("preview" in k or "call.arguments" in k for k in span.attributes) + + def test_content_when_opted_in(self, provider, monkeypatch): + monkeypatch.setenv("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", "span_only") + trace_otel.export_trace(_sample_trace(), tracer_provider=provider) + spans = {s.name: s for s in provider.exporter.get_finished_spans()} + tool = spans["execute_tool search"] + assert tool.attributes["gen_ai.tool.call.arguments"] == '{"q": "x"}' + assert tool.attributes["gen_ai.tool.call.result"] == "found" + assert spans["chat gpt-4o"].attributes["docsgpt.preview.output"] == "hello" + + def test_content_blocked_trace_never_exports_content(self, provider, monkeypatch): + monkeypatch.setenv("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", "SPAN_ONLY") + trace = _sample_trace() + trace.content_blocked = True + trace_otel.export_trace(trace, tracer_provider=provider) + for span in provider.exporter.get_finished_spans(): + assert not any("preview" in k or "call.result" in k for k in span.attributes) + + def test_root_is_parented_to_captured_context(self, provider): + server = provider.get_tracer("test").start_span("GET /stream") + context = ot_trace.set_span_in_context(server) + server.end() + trace_otel.export_trace(_sample_trace(otel_context=context), tracer_provider=provider) + root = {s.name: s for s in provider.exporter.get_finished_spans()}["docsgpt stream"] + assert root.parent.span_id == server.get_span_context().span_id + assert root.context.trace_id == server.get_span_context().trace_id + + def test_skipped_without_sdk_provider(self): + assert trace_otel.export_trace(_sample_trace(), tracer_provider=ot_trace.NoOpTracerProvider()) is None + + def test_skipped_when_disabled(self, provider, monkeypatch): + monkeypatch.setattr(settings, "TRACES_OTEL_EXPORT", False) + assert trace_otel.export_trace(_sample_trace(), tracer_provider=provider) is None + assert provider.exporter.get_finished_spans() == () + + +class TestProviderName: + @pytest.mark.parametrize( + "internal, expected", + [ + ("openai", "openai"), + ("google", "gcp.gen_ai"), + ("anthropic", "anthropic"), + ("azure_openai", "azure.ai.openai"), + ("groq", "groq"), + (None, "unknown"), + ], + ) + def test_mapping(self, internal, expected): + assert trace_otel.provider_name(internal) == expected + + +class TestMetrics: + def test_llm_metrics_recorded(self): + reader = InMemoryMetricReader() + mp = MeterProvider(metric_readers=[reader]) + trace_otel.record_llm_metrics( + provider="openai", + model="gpt-4o", + input_tokens=10, + output_tokens=5, + duration_s=0.25, + error_type=None, + meter_provider=mp, + ) + data = reader.get_metrics_data() + metrics = { + m.name: m + for rm in data.resource_metrics + for sm in rm.scope_metrics + for m in sm.metrics + } + usage = metrics["gen_ai.client.token.usage"] + by_type = { + p.attributes["gen_ai.token.type"]: p.sum for p in usage.data.data_points + } + assert by_type == {"input": 10, "output": 5} + point = usage.data.data_points[0] + assert point.attributes["gen_ai.provider.name"] == "openai" + assert point.attributes["gen_ai.request.model"] == "gpt-4o" + assert point.attributes["gen_ai.operation.name"] == "chat" + duration = metrics["gen_ai.client.operation.duration"].data.data_points[0] + assert duration.sum == pytest.approx(0.25) + + def test_error_type_on_duration_only(self): + reader = InMemoryMetricReader() + mp = MeterProvider(metric_readers=[reader]) + trace_otel.record_llm_metrics( + provider="openai", + model="m", + input_tokens=0, + output_tokens=0, + duration_s=0.1, + error_type="Timeout", + meter_provider=mp, + ) + metrics = { + m.name: m + for rm in reader.get_metrics_data().resource_metrics + for sm in rm.scope_metrics + for m in sm.metrics + } + assert "gen_ai.client.token.usage" not in metrics + point = metrics["gen_ai.client.operation.duration"].data.data_points[0] + assert point.attributes["error.type"] == "Timeout" + + def test_metrics_disabled(self, monkeypatch): + monkeypatch.setattr(settings, "TRACES_OTEL_EXPORT", False) + reader = InMemoryMetricReader() + mp = MeterProvider(metric_readers=[reader]) + trace_otel.record_llm_metrics( + provider="openai", model="m", input_tokens=1, output_tokens=1, + duration_s=0.1, error_type=None, meter_provider=mp, + ) + data = reader.get_metrics_data() + assert data is None or not data.resource_metrics diff --git a/tests/tracing/test_retrieval_spans.py b/tests/tracing/test_retrieval_spans.py new file mode 100644 index 00000000..be45d98c --- /dev/null +++ b/tests/tracing/test_retrieval_spans.py @@ -0,0 +1,153 @@ +"""RAG retrieval becomes retrieval / embeddings / search / rerank spans.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + +from docsgpt import tracing +from docsgpt.core.settings import settings +from docsgpt.retriever import fanout + + +@pytest.fixture(autouse=True) +def _env(monkeypatch): + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + monkeypatch.setattr(settings, "TRACES_CAPTURE_CONTENT", True) + + +@pytest.fixture() +def trace(): + t = tracing.start_trace(source="stream", capture_otel_context=False) + with tracing.activate(t): + yield t + + +class TestFanout: + def test_embedding_span(self, trace): + embedder = MagicMock() + embedder.embed_query.side_effect = lambda q: [0.1, 0.2] + store = MagicMock(_embedding=embedder) + vectors = fanout.embed_questions(store, ["q", "q", "r"]) + assert set(vectors) == {"q", "r"} + (span,) = trace.spans + assert span.kind == tracing.KIND_EMBEDDING + assert span.attributes["gen_ai.operation.name"] == "embeddings" + assert span.attributes["docsgpt.input_count"] == 2 + + def test_embedding_failure_is_recorded(self, trace): + embedder = MagicMock() + embedder.embed_query.side_effect = RuntimeError("model down") + assert fanout.embed_questions(MagicMock(_embedding=embedder), ["q"]) == {} + assert trace.spans[0].status == "error" + + def test_pool_jobs_nest_under_the_open_retrieval(self, trace, monkeypatch): + monkeypatch.setattr(settings, "RETRIEVAL_MAX_PARALLEL_SOURCES", 4) + + def job(i): + with tracing.span(tracing.KIND_SEARCH, f"search {i}"): + return i + + with tracing.span(tracing.KIND_RETRIEVAL, "retrieval") as outer: + assert fanout.run_source_jobs(job, [1, 2, 3]) == [1, 2, 3] + searches = [s for s in trace.spans if s.kind == tracing.KIND_SEARCH] + assert len(searches) == 3 + assert all(s.parent_id == outer.id for s in searches) + + +class TestClassicRAG: + def _rag(self, stores): + from docsgpt.retriever.classic_rag import ClassicRAG + + with patch("docsgpt.retriever.classic_rag.LLMCreator.create_llm", return_value=MagicMock()): + rag = ClassicRAG( + source={"question": "what is x", "active_docs": list(stores)}, + chunks=4, + decoded_token={"sub": "u1"}, + ) + return rag + + def test_search_records_retrieval_tree(self, trace, monkeypatch): + monkeypatch.setattr(settings, "RETRIEVAL_MAX_PARALLEL_SOURCES", 4) + embedder = MagicMock() + embedder.embed_query.return_value = [0.1] + + def make_store(*_args, **_kwargs): + store = MagicMock(_embedding=embedder) + store.search.return_value = [ + {"text": "chunk body", "metadata": {"title": "Doc", "source": "doc.md"}} + ] + return store + + rag = self._rag(["s1", "s2"]) + with patch( + "docsgpt.retriever.classic_rag.VectorCreator.create_vectorstore", + side_effect=make_store, + ): + docs = rag.search() + assert len(docs) == 2 + retrieval = trace.spans[0] + assert retrieval.kind == tracing.KIND_RETRIEVAL + assert retrieval.attributes["docsgpt.source_ids"] == ["s1", "s2"] + assert retrieval.attributes["docsgpt.chunk_count"] == 2 + assert retrieval.previews["query"] == "what is x" + assert retrieval.previews["chunks"][0]["text"] == "chunk body" + children = [s for s in trace.spans[1:]] + assert {s.kind for s in children} == {tracing.KIND_EMBEDDING, tracing.KIND_SEARCH} + assert all(s.parent_id == retrieval.id for s in children) + searches = [s for s in children if s.kind == tracing.KIND_SEARCH] + assert sorted(s.attributes["gen_ai.data_source.id"] for s in searches) == ["s1", "s2"] + assert all(s.attributes["docsgpt.candidate_count"] == 1 for s in searches) + + def test_failed_source_search_is_an_error_span(self, trace, monkeypatch): + monkeypatch.setattr(settings, "RETRIEVAL_MAX_PARALLEL_SOURCES", 1) + rag = self._rag(["s1", "s2"]) + + def make_store(_kind, source_id, _key): + if source_id == "s2": + raise RuntimeError("no store") + store = MagicMock(_embedding=None, _embeddings=None, embeddings=None) + store._get_embeddings.side_effect = RuntimeError("no embedder") + store.search.return_value = [] + return store + + with patch( + "docsgpt.retriever.classic_rag.VectorCreator.create_vectorstore", + side_effect=make_store, + ): + assert rag.search() == [] + search = {s.attributes["gen_ai.data_source.id"]: s for s in trace.spans if s.kind == tracing.KIND_SEARCH} + assert search["s1"].status == "ok" + assert search["s2"].status == "error" + assert trace.spans[0].status == "ok" + + +class TestPrescreen: + def test_rerank_span_contains_screening_calls(self, trace): + from docsgpt.retriever.stages.prescreen import PreScreenStage + from docsgpt.storage.db.source_config import PreScreenConfig + + stage = PreScreenStage( + PreScreenConfig(candidate_k=10, max_keep=2, batch_size=1), + llm_name="openai", + api_key=None, + model_id="m", + ) + llm = MagicMock() + + def fake_gen(**_kwargs): + with tracing.span(tracing.KIND_LLM, "chat m"): + return '{"keep": [0]}' + + llm.gen.side_effect = fake_gen + stage._build_llm = lambda: llm + docs = [{"text": f"d{i}"} for i in range(3)] + kept = stage(docs, {"query": "q"}) + assert len(kept) == 2 + rerank = trace.spans[0] + assert rerank.kind == tracing.KIND_RERANK + assert rerank.attributes["docsgpt.kept_count"] == 2 + llm_spans = [s for s in trace.spans if s.kind == tracing.KIND_LLM] + assert len(llm_spans) == 3 + assert all(s.parent_id == rerank.id for s in llm_spans) diff --git a/tests/tracing/test_sink.py b/tests/tracing/test_sink.py new file mode 100644 index 00000000..4581434e --- /dev/null +++ b/tests/tracing/test_sink.py @@ -0,0 +1,128 @@ +"""Tests for writing a finished trace to Postgres and OTel.""" + +from __future__ import annotations + +from contextlib import contextmanager +from unittest.mock import MagicMock, patch + +import pytest + +from docsgpt import tracing +from docsgpt.core.settings import settings + + +@pytest.fixture(autouse=True) +def _settings(monkeypatch): + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + + +def _trace_with_span(): + trace = tracing.start_trace(source="stream", request_id="r1", user_id="u1") + with tracing.activate(trace): + with tracing.span(tracing.KIND_LLM, "chat m"): + pass + return trace + + +@contextmanager +def _patched_store(): + repo = MagicMock() + repo.insert.return_value = True + + @contextmanager + def _session(): + yield MagicMock() + + with patch("docsgpt.storage.db.session.db_session", _session), patch( + "docsgpt.storage.db.repositories.request_traces.RequestTracesRepository", + return_value=repo, + ): + yield repo + + +class TestFlush: + def test_flush_writes_once(self): + trace = _trace_with_span() + with _patched_store() as repo: + tracing.flush(trace) + tracing.flush(trace) + assert repo.insert.call_count == 1 + record = repo.insert.call_args[0][0] + assert record["request_id"] == "r1" + assert record["status"] == "ok" + assert trace.finished and trace.flushed + + def test_flush_status_override(self): + trace = _trace_with_span() + with _patched_store() as repo: + tracing.flush(trace, status="paused") + assert repo.insert.call_args[0][0]["status"] == "paused" + + def test_empty_trace_is_not_stored(self): + trace = tracing.start_trace(source="stream") + with _patched_store() as repo: + tracing.flush(trace) + repo.insert.assert_not_called() + + def test_store_failure_is_swallowed(self): + trace = _trace_with_span() + with _patched_store() as repo: + repo.insert.side_effect = RuntimeError("db down") + tracing.flush(trace) # must not raise + + def test_otel_failure_does_not_block_store(self): + trace = _trace_with_span() + with _patched_store() as repo, patch( + "docsgpt.tracing.otel.export_trace", side_effect=RuntimeError("otel") + ): + tracing.flush(trace) + repo.insert.assert_called_once() + + def test_otel_trace_id_is_stored(self): + trace = _trace_with_span() + with _patched_store() as repo, patch( + "docsgpt.tracing.otel.export_trace", return_value="ab" * 16 + ): + tracing.flush(trace) + assert repo.insert.call_args[0][0]["otel_trace_id"] == "ab" * 16 + + def test_none_is_ignored(self): + tracing.flush(None) + tracing.discard(None) + + +class TestDiscard: + def test_discard_never_writes(self): + trace = _trace_with_span() + with _patched_store() as repo: + tracing.discard(trace) + tracing.flush(trace) + repo.insert.assert_not_called() + + +class TestBackgroundFlush: + def test_background_flush_writes_off_the_calling_thread(self): + import threading + + trace = _trace_with_span() + seen = {} + + def _insert(record): + seen["thread"] = threading.current_thread().name + return True + + with _patched_store() as repo: + repo.insert.side_effect = _insert + future = tracing.flush(trace, background=True) + assert trace.finished # frozen at once, written later + future.result(timeout=5) + assert seen["thread"].startswith("trace-writer") + assert repo.insert.call_count == 1 + + def test_background_flush_is_still_once(self): + trace = _trace_with_span() + with _patched_store() as repo: + future = tracing.flush(trace, background=True) + assert tracing.flush(trace, background=True) is None + future.result(timeout=5) + assert repo.insert.call_count == 1 diff --git a/tests/tracing/test_tool_spans.py b/tests/tracing/test_tool_spans.py new file mode 100644 index 00000000..54ef8b40 --- /dev/null +++ b/tests/tracing/test_tool_spans.py @@ -0,0 +1,268 @@ +"""Tool calls become ``execute_tool`` spans: executed, paused, denied, skipped.""" + +from __future__ import annotations + +from unittest.mock import Mock + +import pytest + +from docsgpt import tracing +from docsgpt.agents.tool_executor import ToolExecutor +from docsgpt.core.settings import settings +from docsgpt.llm.handlers.base import LLMHandler, ToolCall + + +@pytest.fixture(autouse=True) +def _env(monkeypatch): + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + monkeypatch.setattr(settings, "TRACES_CAPTURE_CONTENT", True) + + +@pytest.fixture() +def trace(): + t = tracing.start_trace(source="stream", capture_otel_context=False) + with tracing.activate(t): + yield t + + +def _executor_with(inner): + executor = ToolExecutor.__new__(ToolExecutor) + executor.tool_calls = [] + executor._execute = lambda tools_dict, call, llm_class_name: inner(executor, call) + return executor + + +def _drain(gen): + events = [] + while True: + try: + events.append(next(gen)) + except StopIteration as stop: + return events, stop.value + + +class TestExecute: + def test_completed_call(self, trace): + def inner(executor, call): + yield {"type": "tool_call", "data": {"status": "pending"}} + executor.tool_calls.append( + { + "tool_name": "brave", + "action_name": "search", + "arguments": {"q": "x", "api_key": "k"}, + "result": "3 results", + "status": "completed", + } + ) + return "3 results", call.id + + executor = _executor_with(inner) + call = ToolCall(id="c1", name="search_1", arguments="{}") + events, outcome = _drain(executor.execute({}, call, "OpenAILLM")) + assert outcome == ("3 results", "c1") + assert len(events) == 1 + (span,) = trace.spans + assert span.kind == tracing.KIND_TOOL + assert span.name == "execute_tool search_1" + assert span.status == "ok" + assert span.attributes["gen_ai.tool.call.id"] == "c1" + assert span.attributes["docsgpt.tool"] == "brave" + assert span.previews["arguments"]["api_key"] == "[REDACTED]" + assert span.previews["result"] == "3 results" + + def test_in_band_error(self, trace): + def inner(executor, call): + executor.tool_calls.append( + {"tool_name": "unknown", "result": "no such tool", "status": "error"} + ) + return "no such tool", call.id + yield # pragma: no cover + + executor = _executor_with(inner) + _drain(executor.execute({}, ToolCall(id="c", name="x", arguments="{}"), "L")) + span = trace.spans[0] + assert span.status == "error" + # Tool output stays in the capture-gated preview, never in ``error``. + assert span.error == "Tool call failed (ToolError)" + assert span.previews["result"] == "no such tool" + + def test_raised_error(self, trace): + def inner(executor, call): + yield {"type": "tool_call", "data": {"status": "pending"}} + raise RuntimeError("api down") + + executor = _executor_with(inner) + with pytest.raises(RuntimeError): + _drain(executor.execute({}, ToolCall(id="c", name="x", arguments="{}"), "L")) + assert trace.spans[0].status == "error" + assert trace.spans[0].attributes["error.type"] == "RuntimeError" + assert trace.spans[0].error == "Tool call failed (RuntimeError)" + assert trace.spans[0].previews["error"] == "api down" + + def test_nested_retrieval_parents_to_tool(self, trace): + def inner(executor, call): + with tracing.span(tracing.KIND_RETRIEVAL, "retrieval"): + pass + executor.tool_calls.append({"status": "completed"}) + return "", call.id + yield # pragma: no cover + + executor = _executor_with(inner) + _drain(executor.execute({}, ToolCall(id="c", name="internal_search", arguments="{}"), "L")) + tool, retrieval = trace.spans + assert retrieval.parent_id == tool.id + + +def test_tool_content_is_not_kept_when_capture_is_off(trace, monkeypatch): + monkeypatch.setattr(settings, "TRACES_CAPTURE_CONTENT", False) + + def inner(executor, call): + executor.tool_calls.append( + {"tool_name": "t", "result": "secret tool output", "status": "error"} + ) + return "", call.id + yield # pragma: no cover + + _drain(_executor_with(inner).execute({}, ToolCall(id="c", name="x", arguments="{}"), "L")) + trace.finish() + record = trace.to_record() + assert "secret tool output" not in str(record) + + +class _Handler(LLMHandler): + def parse_response(self, response): # pragma: no cover - unused + raise NotImplementedError + + def create_tool_message(self, tool_call, result): + return {"role": "tool", "content": str(result)} + + def _iterate_stream(self, response): # pragma: no cover - unused + return iter(()) + + +class TestPausedCalls: + def _agent(self, pause): + agent = Mock() + agent._check_context_limit = Mock(return_value=False) + agent.context_limit_reached = False + agent.llm.__class__.__name__ = "MockLLM" + agent.tool_executor.check_pause = Mock(return_value=pause) + return agent + + def test_awaiting_approval_is_pending(self, trace): + agent = self._agent( + { + "call_id": "c1", + "name": "send_0", + "tool_name": "telegram", + "tool_id": "0", + "action_name": "send", + "arguments": {"text": "hi"}, + "pause_type": "awaiting_approval", + } + ) + call = ToolCall(id="c1", name="send_0", arguments='{"text": "hi"}') + _drain(_Handler().handle_tool_calls(agent, [call], {"0": {"name": "telegram"}}, [])) + (span,) = trace.spans + assert span.status == "pending" + assert span.attributes["docsgpt.tool_status"] == "awaiting_approval" + + def test_headless_denial_is_denied(self, trace): + agent = self._agent( + { + "call_id": "c1", + "name": "send_0", + "tool_name": "telegram", + "tool_id": "0", + "action_name": "send", + "arguments": {}, + "pause_type": "headless_denied", + "deny_reason": "not allowed", + } + ) + agent.tool_executor.message_id = None + agent.tool_executor.user = "u" + agent.tool_executor.agent_id = None + call = ToolCall(id="c1", name="send_0", arguments="{}") + with pytest.MonkeyPatch.context() as mp: + mp.setattr("docsgpt.agents.tool_executor._record_proposed", lambda *a, **k: False) + _drain(_Handler().handle_tool_calls(agent, [call], {"0": {"name": "telegram"}}, [])) + (span,) = trace.spans + assert span.status == "denied" + assert span.error is None + assert span.previews["error"] == "not allowed" + + +class TestContinuation: + def _agent(self): + from docsgpt.agents.classic_agent import ClassicAgent + + llm = Mock() + llm._supports_tools = True + llm.gen_stream = Mock(return_value=iter(["Answer"])) + llm._supports_structured_output = Mock(return_value=False) + handler = Mock() + handler.process_message_flow = Mock(return_value=iter([])) + handler.create_tool_message = Mock(return_value={"role": "tool", "content": "x"}) + executor = Mock() + executor.tool_calls = [] + executor.prepare_tools_for_llm = Mock(return_value=[]) + executor.get_truncated_tool_calls = Mock(return_value=[]) + return ClassicAgent( + endpoint="stream", + llm_name="openai", + model_id="gpt-4", + api_key="test", + llm=llm, + llm_handler=handler, + tool_executor=executor, + ) + + def _pending(self): + return [ + { + "call_id": "c1", + "name": "danger_0", + "tool_name": "danger", + "tool_id": "0", + "action_name": "danger", + "arguments": {}, + "pause_type": "awaiting_approval", + } + ] + + def test_continuation_opens_agent_span_with_denied_tool(self, trace): + agent = self._agent() + list( + agent.gen_continuation( + [{"role": "system", "content": "s"}], + {"0": {"name": "danger"}}, + self._pending(), + [{"call_id": "c1", "decision": "denied", "comment": "too risky"}], + ) + ) + agent_span = trace.spans[0] + tool_span = trace.spans[1] + assert agent_span.kind == tracing.KIND_AGENT + assert agent_span.attributes["docsgpt.continuation"] is True + assert agent_span.attributes["gen_ai.request.model"] == "gpt-4" + assert agent_span.name == "invoke_agent ClassicAgent" + assert agent_span.status == "ok" + assert tool_span.parent_id == agent_span.id + assert tool_span.status == "denied" + assert tool_span.error is None + assert tool_span.previews["error"] == "too risky" + + def test_client_result_is_recorded(self, trace): + agent = self._agent() + list( + agent.gen_continuation( + [{"role": "system", "content": "s"}], + {"0": {"name": "danger"}}, + self._pending(), + [{"call_id": "c1", "result": {"ok": True}}], + ) + ) + tool_span = trace.spans[1] + assert tool_span.status == "ok" + assert tool_span.attributes["docsgpt.client_executed"] is True diff --git a/tests/tracing/test_trace.py b/tests/tracing/test_trace.py new file mode 100644 index 00000000..b6e8c558 --- /dev/null +++ b/tests/tracing/test_trace.py @@ -0,0 +1,410 @@ +"""Tests for the execution-trace recorder in ``docsgpt.tracing``.""" + +from __future__ import annotations + +import threading +from concurrent.futures import ThreadPoolExecutor + +import pytest + +from docsgpt import tracing +from docsgpt.core.settings import settings + + +@pytest.fixture(autouse=True) +def _tracing_on(monkeypatch): + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + monkeypatch.setattr(settings, "TRACES_CAPTURE_CONTENT", True) + monkeypatch.setattr(settings, "TRACES_MAX_SPANS", 500) + monkeypatch.setattr(settings, "TRACES_PREVIEW_CHARS", 2000) + + +def _by_name(trace): + return {s.name: s for s in trace.spans} + + +class TestNoActiveTrace: + def test_span_calls_are_noops_without_a_trace(self): + assert tracing.current_trace() is None + with tracing.span(tracing.KIND_LLM, "chat gpt") as s: + s.set(foo=1) + s.preview("output", "hi") + handle = tracing.start_span(tracing.KIND_TOOL, "tool") + handle.end(status="error") + tracing.bind(message_id="m") + + def test_start_trace_returns_none_when_disabled(self, monkeypatch): + monkeypatch.setattr(settings, "TRACES_ENABLED", False) + assert tracing.start_trace(source="stream") is None + with tracing.activate(None): + assert tracing.current_trace() is None + + +class TestNesting: + def test_containers_push_and_leaves_record_parent(self): + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + with tracing.span(tracing.KIND_AGENT, "agent"): + with tracing.span(tracing.KIND_LLM, "llm-1"): + pass + with tracing.span(tracing.KIND_TOOL, "tool"): + with tracing.span(tracing.KIND_RETRIEVAL, "retrieval"): + with tracing.span(tracing.KIND_EMBEDDING, "embed"): + pass + with tracing.span(tracing.KIND_LLM, "llm-2"): + pass + spans = _by_name(trace) + assert spans["agent"].parent_id is None + assert spans["llm-1"].parent_id == spans["agent"].id + assert spans["tool"].parent_id == spans["agent"].id + assert spans["retrieval"].parent_id == spans["tool"].id + assert spans["embed"].parent_id == spans["retrieval"].id + assert spans["llm-2"].parent_id == spans["agent"].id + assert all(s.status == "ok" for s in trace.spans) + + def test_leaf_span_does_not_become_a_parent(self): + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + llm = tracing.start_span(tracing.KIND_LLM, "llm") + with tracing.span(tracing.KIND_TOOL, "tool"): + pass + llm.end() + spans = _by_name(trace) + assert spans["tool"].parent_id is None + + def test_explicit_parent_wins(self): + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + outer = tracing.start_span(tracing.KIND_RETRIEVAL, "outer") + with tracing.span(tracing.KIND_AGENT, "agent"): + child = tracing.start_span(tracing.KIND_SEARCH, "child", parent=outer) + child.end() + outer.end() + assert _by_name(trace)["child"].parent_id == _by_name(trace)["outer"].id + + def test_interleaved_generators_nest_by_start_order(self): + """An agent generator suspended at yield keeps its children nested.""" + trace = tracing.start_trace(source="stream") + + def llm_stream(): + handle = tracing.start_span(tracing.KIND_LLM, "llm") + try: + yield "a" + yield "b" + finally: + handle.end() + + def agent_gen(): + handle = tracing.start_span(tracing.KIND_AGENT, "agent") + try: + yield from llm_stream() + with tracing.span(tracing.KIND_TOOL, "tool"): + pass + yield "c" + finally: + handle.end() + + with tracing.activate(trace): + assert list(agent_gen()) == ["a", "b", "c"] + spans = _by_name(trace) + assert spans["llm"].parent_id == spans["agent"].id + assert spans["tool"].parent_id == spans["agent"].id + + def test_out_of_order_end_cancels_abandoned_children(self): + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + agent = tracing.start_span(tracing.KIND_AGENT, "agent") + tool = tracing.start_span(tracing.KIND_TOOL, "tool") + agent.end() + after = tracing.start_span(tracing.KIND_LLM, "after") + after.end() + tool.end() # late end is ignored + spans = _by_name(trace) + assert spans["tool"].status == "cancelled" + assert spans["agent"].status == "ok" + assert spans["after"].parent_id is None + + def test_exception_marks_span_error_and_propagates(self): + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + with pytest.raises(ValueError): + with tracing.span(tracing.KIND_TOOL, "tool"): + raise ValueError("boom") + span = trace.spans[0] + assert span.status == "error" + assert span.attributes["error.type"] == "ValueError" + # Exception text can quote a prompt: it is a capture-gated preview, + # never the always-stored/exported ``error``. + assert span.error == "ValueError" + assert span.previews["error"] == "boom" + + def test_exception_text_is_dropped_when_capture_is_off(self, monkeypatch): + monkeypatch.setattr(settings, "TRACES_CAPTURE_CONTENT", False) + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + with pytest.raises(ValueError): + with tracing.span(tracing.KIND_LLM, "chat"): + raise ValueError("content filter: ") + trace.finish() + assert "the user's prompt" not in str(trace.to_record()) + + def test_generator_exit_marks_span_cancelled(self): + trace = tracing.start_trace(source="stream") + + def gen(): + with tracing.span(tracing.KIND_AGENT, "agent"): + yield 1 + yield 2 + + with tracing.activate(trace): + g = gen() + next(g) + g.close() + assert trace.spans[0].status == "cancelled" + + +class TestFinish: + def test_finish_cancels_open_spans_and_ignores_late_ends(self): + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + open_span = tracing.start_span(tracing.KIND_AGENT, "agent") + trace.finish() + assert open_span.status == "cancelled" + duration = open_span.duration_ms + open_span.end(status="ok") + assert open_span.status == "cancelled" + assert open_span.duration_ms == duration + + def test_finish_status_defaults_to_error_when_a_top_level_span_failed(self): + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + tracing.start_span(tracing.KIND_AGENT, "agent").end(status="error") + trace.finish() + assert trace.status == "error" + + def test_explicit_finish_status(self): + trace = tracing.start_trace(source="stream") + trace.finish(status="paused") + assert trace.status == "paused" + + def test_spans_after_finish_are_dropped(self): + trace = tracing.start_trace(source="stream") + trace.finish() + with tracing.activate(trace): + tracing.start_span(tracing.KIND_LLM, "late").end() + assert trace.spans == [] + + def test_summary_rolls_up_llm_tool_retrieval(self): + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + with tracing.span(tracing.KIND_AGENT, "agent"): + with tracing.span( + tracing.KIND_LLM, + "chat", + attributes={ + "gen_ai.usage.input_tokens": 10, + "gen_ai.usage.output_tokens": 4, + }, + ): + pass + with tracing.span(tracing.KIND_RETRIEVAL, "r"): + with tracing.span(tracing.KIND_RETRIEVAL, "inner"): + pass + with tracing.span(tracing.KIND_TOOL, "t") as t: + t.end(status="error") + trace.finish() + summary = trace.summary() + assert summary["llm_calls"] == 1 + assert summary["tool_calls"] == 1 + assert summary["input_tokens"] == 10 + assert summary["output_tokens"] == 4 + assert summary["retrieval_calls"] == 1 + assert summary["errors"] == 1 + assert summary["retrieval_ms"] >= 0 + + +class TestCap: + def test_span_cap_counts_dropped(self, monkeypatch): + monkeypatch.setattr(settings, "TRACES_MAX_SPANS", 10) + trace = tracing.start_trace(source="graph_extraction") + with tracing.activate(trace): + for i in range(15): + with tracing.span(tracing.KIND_LLM, f"llm-{i}"): + pass + assert len(trace.spans) == 10 + assert trace.dropped_spans == 5 + + +class TestThreads: + def test_wrap_carries_trace_and_parent_into_pool_threads(self): + trace = tracing.start_trace(source="stream") + + def work(i): + with tracing.span(tracing.KIND_SEARCH, f"search-{i}"): + with tracing.span(tracing.KIND_EMBEDDING, f"embed-{i}"): + pass + return threading.get_ident() + + with tracing.activate(trace): + with tracing.span(tracing.KIND_RETRIEVAL, "retrieval"): + with ThreadPoolExecutor(max_workers=3) as pool: + list(pool.map(tracing.wrap(work), range(3))) + spans = _by_name(trace) + for i in range(3): + assert spans[f"search-{i}"].parent_id == spans["retrieval"].id + assert spans[f"embed-{i}"].parent_id == spans["retrieval"].id + + def test_pool_threads_without_wrap_see_no_trace(self): + trace = tracing.start_trace(source="stream") + seen = [] + with tracing.activate(trace): + t = threading.Thread(target=lambda: seen.append(tracing.current_trace())) + t.start() + t.join() + assert seen == [None] + + def test_wrap_without_trace_is_passthrough(self): + fn = tracing.wrap(lambda x: x + 1) + assert fn(1) == 2 + + +class TestBind: + def test_bind_sets_ids_on_active_trace(self): + trace = tracing.start_trace(source="stream", request_id="r1") + with tracing.activate(trace): + tracing.bind(message_id="m1", conversation_id="c1", unknown="x") + assert trace.request_id == "r1" + assert trace.message_id == "m1" + assert trace.conversation_id == "c1" + + def test_bind_if_unset_keeps_first_value(self): + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + tracing.bind_if_unset(activity_id="a1") + tracing.bind_if_unset(activity_id="a2") + assert trace.activity_id == "a1" + + +class TestPreviews: + def test_preview_redacts_secrets_and_bounds_strings(self, monkeypatch): + monkeypatch.setattr(settings, "TRACES_PREVIEW_CHARS", 100) + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + with tracing.span(tracing.KIND_TOOL, "tool") as s: + s.preview("arguments", {"api_key": "sk-123", "q": "x" * 500}) + s.preview("result", "y" * 500 + "\x00") + preview = trace.spans[0].previews + assert preview["arguments"]["api_key"] == "[REDACTED]" + assert len(preview["arguments"]["q"]) < 200 + assert preview["result"].endswith("…") + assert "\x00" not in preview["result"] + + def test_preview_skipped_when_capture_disabled(self, monkeypatch): + monkeypatch.setattr(settings, "TRACES_CAPTURE_CONTENT", False) + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + with tracing.span(tracing.KIND_TOOL, "tool") as s: + s.preview("result", "secret stuff") + assert trace.spans[0].previews == {} + + def test_guardrail_trigger_strips_all_previews(self): + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + with tracing.span(tracing.KIND_TOOL, "tool") as s: + s.preview("result", "leaked") + tracing.mark_content_blocked() + trace.finish() + record = trace.to_record() + assert all(not s.get("preview") for s in record["spans"]) + assert trace.content_blocked is True + + def test_huge_structures_collapse_to_a_string(self, monkeypatch): + monkeypatch.setattr(settings, "TRACES_PREVIEW_CHARS", 100) + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + with tracing.span(tracing.KIND_TOOL, "tool") as s: + s.preview("result", [{"k": "v" * 50} for _ in range(100)]) + value = trace.spans[0].previews["result"] + assert isinstance(value, str) + assert len(value) <= 101 + + +class TestRecord: + def test_to_record_shape(self): + trace = tracing.start_trace( + source="stream", request_id="r", user_id="u", agent_id="a" + ) + with tracing.activate(trace): + with tracing.span( + tracing.KIND_LLM, "chat m", attributes={"gen_ai.request.model": "m"} + ): + pass + trace.finish() + record = trace.to_record() + assert record["source"] == "stream" + assert record["request_id"] == "r" + assert record["status"] == "ok" + assert record["span_count"] == 1 + span = record["spans"][0] + assert set(span) >= { + "id", "parent_id", "kind", "name", "status", "offset_ms", + "duration_ms", "attributes", + } + assert span["attributes"]["gen_ai.request.model"] == "m" + assert span["offset_ms"] >= 0 + + +class TestOutcome: + def test_outcome_is_used_when_no_status_given(self): + trace = tracing.start_trace(source="stream") + trace.outcome = "paused" + trace.finish() + assert trace.status == "paused" + + def test_explicit_status_beats_outcome(self): + trace = tracing.start_trace(source="stream") + trace.outcome = "paused" + trace.finish(status="error") + assert trace.status == "error" + + +class TestRecordQuery: + def test_first_span_query_is_copied_into_summary(self): + trace = tracing.start_trace(source="search") + with tracing.activate(trace): + with tracing.span(tracing.KIND_RETRIEVAL, "retrieval") as s: + s.preview("query", "how do I deploy") + trace.finish() + assert trace.to_record()["summary"]["query"] == "how do I deploy" + + def test_query_is_dropped_when_content_is_blocked(self): + trace = tracing.start_trace(source="search") + with tracing.activate(trace): + with tracing.span(tracing.KIND_RETRIEVAL, "retrieval") as s: + s.preview("query", "leaky") + tracing.mark_content_blocked() + trace.finish() + assert "query" not in trace.to_record()["summary"] + + +class TestToolCount: + def test_only_tool_calls_that_ran_are_counted(self): + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + tracing.start_span(tracing.KIND_TOOL, "paused").end(tracing.STATUS_PENDING) + tracing.start_span(tracing.KIND_TOOL, "denied").end(tracing.STATUS_DENIED) + tracing.start_span(tracing.KIND_TOOL, "skipped").end(tracing.STATUS_SKIPPED) + tracing.start_span(tracing.KIND_TOOL, "ran").end() + tracing.start_span(tracing.KIND_TOOL, "failed").end(tracing.STATUS_ERROR) + trace.finish() + assert trace.summary()["tool_calls"] == 2 + + +class TestFirstOccurrence: + def test_key_is_new_once_per_trace(self): + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + assert tracing.first_occurrence(("a", 1)) is True + assert tracing.first_occurrence(("a", 1)) is False + assert tracing.first_occurrence(("a", 2)) is True + assert tracing.first_occurrence(("a", 3)) is False # no active trace diff --git a/tests/worker/test_extract_graph.py b/tests/worker/test_extract_graph.py index 7e040fbe..10283ad5 100644 --- a/tests/worker/test_extract_graph.py +++ b/tests/worker/test_extract_graph.py @@ -170,3 +170,53 @@ class TestExtractGraphWorker: worker.extract_graph_worker(task_self, source_id, "alice") assert "graph.extract.failed" in [e[0] for e in events] + + +@pytest.mark.unit +class TestExtractGraphTrace: + """A graph build is one execution trace holding every extraction call.""" + + def _run(self, pg_conn, monkeypatch, task_self, extract): + from docsgpt import tracing, worker + from docsgpt.core.settings import settings + + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + monkeypatch.setattr(settings, "TRACES_OTEL_EXPORT", False) + source_id = _seed_source(pg_conn) + _patch_store(monkeypatch, [{"doc_id": "c1", "text": "alpha"}]) + monkeypatch.setattr("docsgpt.graphrag.graphrag_available", lambda: True) + monkeypatch.setattr("docsgpt.graphrag.extraction.extract_graph_for_source", extract) + flushed = [] + + def _flush(trace, status=None, **_kwargs): + trace.flushed = True + trace.finish(status) + flushed.append(trace) + + monkeypatch.setattr(tracing, "flush", _flush) + return worker, source_id, flushed + + def test_successful_build_is_traced(self, pg_conn, patch_worker_db, task_self, monkeypatch): + from docsgpt import tracing + + def _extract(*_a, **_kw): + with tracing.span(tracing.KIND_LLM, "chat m"): + pass + return {"nodes": 3, "edges": 2, "chunks_processed": 1} + + worker, source_id, flushed = self._run(pg_conn, monkeypatch, task_self, _extract) + worker.extract_graph_worker(task_self, source_id, "alice") + (trace,) = flushed + assert trace.source == "graph_extraction" + assert trace.user_id == "alice" + step, llm = trace.spans + assert step.attributes["docsgpt.graph.nodes"] == 3 + assert llm.parent_id == step.id + + def test_failed_build_is_an_error_trace(self, pg_conn, patch_worker_db, task_self, monkeypatch): + worker, source_id, flushed = self._run( + pg_conn, monkeypatch, task_self, MagicMock(side_effect=RuntimeError("llm down")) + ) + with pytest.raises(RuntimeError): + worker.extract_graph_worker(task_self, source_id, "alice") + assert flushed[0].status == "error" diff --git a/uv.lock b/uv.lock index dba25614..13780640 100644 --- a/uv.lock +++ b/uv.lock @@ -1375,6 +1375,7 @@ dependencies = [ { name = "openai" }, { name = "openapi3-parser" }, { name = "openpyxl" }, + { name = "opentelemetry-api" }, { name = "opentelemetry-distro" }, { name = "opentelemetry-exporter-otlp" }, { name = "opentelemetry-instrumentation-celery" }, @@ -1488,6 +1489,7 @@ requires-dist = [ { name = "openai", specifier = ">=3.13.0,<4" }, { name = "openapi3-parser", specifier = ">=2.0.0,<3" }, { name = "openpyxl", specifier = ">=3.1.5,<4" }, + { name = "opentelemetry-api", specifier = ">=1.29.0,<2" }, { name = "opentelemetry-distro", specifier = ">=0.50b0,<1" }, { name = "opentelemetry-exporter-otlp", specifier = ">=1.29.0,<2" }, { name = "opentelemetry-instrumentation-celery", specifier = ">=0.50b0,<1" },