import datetime import functools import inspect import time import logging import uuid from contextlib import contextmanager from typing import Any, Callable, Dict, Generator, Iterator, 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 logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" ) class LogContext: def __init__(self, endpoint, activity_id, user, api_key, query, agent_id=None): self.endpoint = endpoint self.activity_id = activity_id self.user = user self.api_key = api_key self.agent_id = agent_id self.query = query self.stacks = [] # Per-activity response aggregates populated by ``_consume_and_log`` # while it forwards stream items, then flushed onto the # ``activity_finished`` event so every Flask request gets the # same summary that ``run_agent_logic`` used to log only for the # Celery webhook path. self.answer_length = 0 self.thought_length = 0 self.source_count = 0 self.tool_call_count = 0 # Terminal ``error`` events are *yielded*, not raised (workflow node # failures, agent-reported errors), so they never reach the decorator's # ``except``. Recording one here keeps ``activity_finished`` from # reporting ``status="ok"`` on a turn the user saw fail — that gap is # why blank-answer incidents did not show up in error dashboards. self.stream_error: str | None = None def build_stack_data( obj: Any, include_attributes: List[str] = None, exclude_attributes: List[str] = None, custom_data: Dict = None, ) -> Dict: if obj is None: raise ValueError("The 'obj' parameter cannot be None") data = {} if include_attributes is None: include_attributes = [] for name, value in inspect.getmembers(obj): if ( not name.startswith("_") and not inspect.ismethod(value) and not inspect.isfunction(value) ): include_attributes.append(name) for attr_name in include_attributes: if exclude_attributes and attr_name in exclude_attributes: continue try: attr_value = getattr(obj, attr_name) if attr_value is not None: if isinstance(attr_value, (int, float, str, bool)): data[attr_name] = attr_value elif isinstance(attr_value, list): if all(isinstance(item, dict) for item in attr_value): data[attr_name] = attr_value elif all(hasattr(item, "__dict__") for item in attr_value): data[attr_name] = [item.__dict__ for item in attr_value] else: data[attr_name] = [str(item) for item in attr_value] elif isinstance(attr_value, dict): data[attr_name] = {k: str(v) for k, v in attr_value.items()} except AttributeError as e: logging.warning(f"AttributeError while accessing {attr_name}: {e}") except AttributeError: pass if custom_data: data.update(custom_data) return data def _agent_log_keys(agent: Any, data: Optional[Dict] = None) -> Dict[str, Any]: """Return the log-context keys identifying an agent run.""" if data is None: data = build_stack_data(agent) return { "user_id": data.get("user", "local"), "agent_id": getattr(agent, "agent_id", None), "conversation_id": getattr(agent, "conversation_id", None), "endpoint": data.get("endpoint", ""), "model": ( getattr(agent, "gpt_model", None) or getattr(agent, "model", None) or getattr(agent, "model_id", None) ), } @contextmanager def agent_log_context(agent: Any) -> Iterator[None]: """Bind an agent's identity to the log context without opening an activity. Any enclosing ``activity_id`` is kept. Args: agent: The agent whose identity the log lines should carry. Yields: None, with the context bound until the block exits. """ token = log_context.bind(**_agent_log_keys(agent)) try: yield finally: log_context.reset(token) def log_activity() -> Callable: def decorator(func: Callable) -> Callable: @functools.wraps(func) def wrapper(*args: Any, **kwargs: Any) -> Any: activity_id = str(uuid.uuid4()) data = build_stack_data(args[0]) keys = _agent_log_keys(args[0], data) endpoint = keys["endpoint"] user = keys["user_id"] api_key = data.get("user_api_key", "") query = kwargs.get("query", getattr(args[0], "query", "")) agent_id = keys["agent_id"] or kwargs.get("agent_id") conversation_id = kwargs.get("conversation_id") or keys["conversation_id"] model = keys["model"] # Capture the surrounding activity_id before overlaying ours, # so nested activities record the parent → child link. parent_activity_id = log_context.snapshot().get("activity_id") context = LogContext( endpoint, activity_id, user, api_key, query, agent_id=agent_id ) kwargs["log_context"] = context ctx_token = log_context.bind( activity_id=activity_id, parent_activity_id=parent_activity_id, user_id=user, agent_id=agent_id, conversation_id=conversation_id, endpoint=endpoint, model=model, ) started_at = time.monotonic() logging.info( "activity_started", extra={ "activity_id": activity_id, "parent_activity_id": parent_activity_id, "user_id": user, "agent_id": agent_id, "conversation_id": conversation_id, "endpoint": endpoint, "model": model, }, ) # 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`` # flow through the finally as ``status="ok"``, matching # ``_consume_and_log``. error = exc raise finally: _emit_activity_finished( context=context, parent_activity_id=parent_activity_id, started_at=started_at, error=error, ) _end_agent_span(span, context, error=error, completed=completed) log_context.reset(ctx_token) return wrapper 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", parent_activity_id: str | None, started_at: float, error: BaseException | None, ) -> None: """Emit the paired ``activity_finished`` event with duration, outcome, and per-activity response aggregates accumulated in ``_consume_and_log``. """ duration_ms = int((time.monotonic() - started_at) * 1000) logging.info( "activity_finished", extra={ "activity_id": context.activity_id, "parent_activity_id": parent_activity_id, "user_id": context.user, "endpoint": context.endpoint, "duration_ms": duration_ms, "status": "error" if (error is not None or context.stream_error) else "ok", "error_class": ( type(error).__name__ if error is not None else ("StreamError" if context.stream_error else None) ), "answer_length": context.answer_length, "thought_length": context.thought_length, "source_count": context.source_count, "tool_call_count": context.tool_call_count, }, ) def _accumulate_response_summary(item: Any, context: "LogContext") -> None: """Mirror the per-line aggregation that ``run_agent_logic`` did for the Celery webhook path, but at the generator-consumption layer so every ``Agent.gen`` activity (Flask streaming, sub-agents, workflow agents) gets the same summary. """ if not isinstance(item, dict): return if item.get("type") == "error": # Fall back to a sentinel: an error event carrying no message would # otherwise store "" and read as falsy, reporting the activity "ok". context.stream_error = str(item.get("error") or "")[:200] or "unspecified" return if "answer" in item: context.answer_length += len(str(item["answer"])) return if "thought" in item: context.thought_length += len(str(item["thought"])) return sources = item.get("sources") if "sources" in item else None if isinstance(sources, list): context.source_count += len(sources) return tool_calls = item.get("tool_calls") if "tool_calls" in item else None if isinstance(tool_calls, list): context.tool_call_count += len(tool_calls) def _consume_and_log(generator: Generator, context: "LogContext"): try: for item in generator: _accumulate_response_summary(item, context) yield item except Exception as e: logging.exception(f"Error in {context.endpoint} - {context.activity_id}: {e}") context.stacks.append({"component": "error", "data": {"message": str(e)}}) _log_activity_to_db( endpoint=context.endpoint, activity_id=context.activity_id, user=context.user, api_key=context.api_key, agent_id=context.agent_id, query=context.query, stacks=context.stacks, level="error", ) raise finally: _log_activity_to_db( endpoint=context.endpoint, activity_id=context.activity_id, user=context.user, api_key=context.api_key, agent_id=context.agent_id, query=context.query, stacks=context.stacks, level="info", ) def _log_activity_to_db( endpoint: str, activity_id: str, user: str, api_key: str, query: str, stacks: List[Dict], level: str, agent_id: Optional[str] = None, ) -> None: """Append a per-request activity log row to Postgres (``stack_logs``).""" try: # Clean up text fields to be no longer than 10000 characters so a # runaway payload can't blow up the insert. def _truncate(val): if isinstance(val, str) and len(val) > 10000: return val[:10000] return val with db_session() as conn: StackLogsRepository(conn).insert( activity_id=activity_id, endpoint=_truncate(endpoint), level=_truncate(level), user_id=_truncate(user), api_key=_truncate(api_key), agent_id=agent_id, query=_truncate(query), stacks=stacks, timestamp=datetime.datetime.now(datetime.timezone.utc), ) logging.debug(f"Logged activity to Postgres: {activity_id}") except Exception as e: logging.error(f"Failed to log activity to Postgres: {e}", exc_info=True)