mirror of
https://github.com/dograh-hq/dograh.git
synced 2026-06-07 07:55:16 +02:00
00a1a22b74
* refactor: carve out extraction panel * refactor: create spec versions for node types * refactor: create a GenericNode and remove custom nodes * feat: add python and typescript sdk * add dograh sdk * fix: fetch draft workflow definition over published one * fix: fix routes of SDKs to use code gen * chore: remove doclink dependency to reduce image size * chore: format files * chore: bump pipecat * feat: let mcp fetch archived workflows on demand * chore: fix tests * feat: add sdk documentation * chore: change banner and add badge
87 lines
3.2 KiB
Python
87 lines
3.2 KiB
Python
"""OTel tracing for MCP tool invocations.
|
|
|
|
The project-wide tracing setup in
|
|
`api/services/pipecat/tracing_config.py` already routes spans to
|
|
per-organization Langfuse projects based on the `dograh.org_id` span
|
|
attribute. This module plugs MCP tool calls into that pipeline:
|
|
|
|
@mcp.tool
|
|
@traced_tool
|
|
async def my_tool(...): ...
|
|
|
|
Each decorated invocation produces one span named `mcp.<tool_name>` with
|
|
Langfuse-rendered input/output. Organization and user attributes are
|
|
stamped separately by `authenticate_mcp_request` when it runs inside
|
|
the tool body — the decorator's span is the `current_span` at that
|
|
point, so the attributes land on the right span and the router export
|
|
dispatches to the correct Langfuse project.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from functools import wraps
|
|
from typing import Any, Awaitable, Callable, TypeVar
|
|
|
|
from opentelemetry import trace
|
|
from opentelemetry.context import Context
|
|
from opentelemetry.trace import Status, StatusCode
|
|
|
|
R = TypeVar("R")
|
|
|
|
_TRACER = trace.get_tracer("dograh.mcp")
|
|
# Langfuse truncates long payloads anyway; cap here to keep span size
|
|
# bounded. Tune up if you find tool outputs consistently clipped.
|
|
_MAX_ATTR_LEN = 8000
|
|
|
|
|
|
def _safe_json(value: Any) -> str:
|
|
try:
|
|
return json.dumps(value, default=str, ensure_ascii=False)
|
|
except Exception: # noqa: BLE001
|
|
return str(value)
|
|
|
|
|
|
def traced_tool(fn: Callable[..., Awaitable[R]]) -> Callable[..., Awaitable[R]]:
|
|
"""Wrap an MCP tool so each invocation produces a span.
|
|
|
|
Captures tool name, input kwargs, output, and exceptions. Stacks
|
|
below `@mcp.tool` so FastMCP sees the wrapped function when
|
|
introspecting the tool schema (`functools.wraps` preserves the
|
|
signature the framework reads).
|
|
"""
|
|
|
|
@wraps(fn)
|
|
async def wrapper(*args: Any, **kwargs: Any) -> R:
|
|
# Each MCP tool call is its own root trace. Passing an empty
|
|
# `Context()` severs the inherited parent so the span doesn't
|
|
# graft onto whatever other trace happens to be active (e.g.
|
|
# the FastAPI request span, or a client-propagated context).
|
|
# One trace per tool invocation makes Langfuse diffing and
|
|
# per-org filtering clean.
|
|
with _TRACER.start_as_current_span(
|
|
f"mcp.{fn.__name__}",
|
|
context=Context(),
|
|
) as span:
|
|
span.set_attribute("mcp.tool.name", fn.__name__)
|
|
# Explicit trace-name override so the Langfuse UI shows
|
|
# `mcp.<tool>` at the top of the trace instead of whatever
|
|
# the framework happens to name the root span.
|
|
span.set_attribute("langfuse.trace.name", f"mcp.{fn.__name__}")
|
|
span.set_attribute(
|
|
"langfuse.observation.input",
|
|
_safe_json(kwargs)[:_MAX_ATTR_LEN],
|
|
)
|
|
try:
|
|
result = await fn(*args, **kwargs)
|
|
except Exception as e:
|
|
span.record_exception(e)
|
|
span.set_status(Status(StatusCode.ERROR, str(e)))
|
|
raise
|
|
span.set_attribute(
|
|
"langfuse.observation.output",
|
|
_safe_json(result)[:_MAX_ATTR_LEN],
|
|
)
|
|
return result
|
|
|
|
return wrapper
|