mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-03 18:46:54 +00:00
Alex
cf9a5e9cb8
Agents store their model only in model_id, so neither log_activity nor agent_log_context bound a model key. start_agent_span already reads it.
407 lines
15 KiB
Python
407 lines
15 KiB
Python
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)
|