import datetime import functools import inspect import time import logging import uuid from typing import Any, Callable, Dict, Generator, List from application.core import log_context from application.storage.db.repositories.stack_logs import StackLogsRepository from application.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): self.endpoint = endpoint self.activity_id = activity_id self.user = user self.api_key = api_key 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 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 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]) endpoint = data.get("endpoint", "") user = data.get("user", "local") api_key = data.get("user_api_key", "") query = kwargs.get("query", getattr(args[0], "query", "")) agent_id = getattr(args[0], "agent_id", None) or kwargs.get("agent_id") conversation_id = ( kwargs.get("conversation_id") or getattr(args[0], "conversation_id", None) ) model = getattr(args[0], "gpt_model", None) or getattr(args[0], "model", None) # 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) 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, }, ) error: BaseException | None = None try: generator = func(*args, **kwargs) yield from _consume_and_log(generator, context) 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, ) log_context.reset(ctx_token) return wrapper return decorator 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 else "ok", "error_class": type(error).__name__ if error is not None 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 "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, 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, 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, ) -> 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), 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)