224 lines
9.2 KiB
Python
224 lines
9.2 KiB
Python
"""Trace state and observation lifecycle for the Langfuse plugin.
|
|
|
|
Holds no process-global state: the live ``TraceState`` registry, its lock and
|
|
the client cache live in the package ``__init__`` (tests patch them there).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Dict, Optional
|
|
|
|
from ._capture import _capture_mode, _debug, _extract_last_user_message
|
|
|
|
try:
|
|
from langfuse import propagate_attributes
|
|
except Exception: # pragma: no cover - fail-open when optional dep is missing
|
|
propagate_attributes = None
|
|
|
|
|
|
@dataclass
|
|
class TraceState:
|
|
trace_id: str
|
|
root_ctx: Any
|
|
root_span: Any
|
|
generations: Dict[str, Any] = field(default_factory=dict)
|
|
tools: Dict[str, Any] = field(default_factory=dict)
|
|
pending_tools_by_name: Dict[str, list] = field(default_factory=dict)
|
|
turn_tool_calls: list[dict[str, Any]] = field(default_factory=list)
|
|
# Keyed by child_session_id: subagent_stop carries no child_subagent_id.
|
|
subagents: Dict[str, Any] = field(default_factory=dict)
|
|
# Fingerprints of MoA fan-outs already recorded: the client holds its last
|
|
# fan-out until the next one, so tool-loop turns would re-emit advisors.
|
|
moa_emitted: set = field(default_factory=set)
|
|
last_updated_at: float = field(default_factory=time.time)
|
|
|
|
def touch(self) -> None:
|
|
self.last_updated_at = time.time()
|
|
|
|
def pop_tool(self, tool_call_id: str, tool_name: str) -> Any:
|
|
"""Detach the open tool observation: by id, else FIFO by name for id-less callers."""
|
|
observation = self.tools.pop(tool_call_id, None) if tool_call_id else None
|
|
if observation is None:
|
|
queue = self.pending_tools_by_name.get(tool_name)
|
|
if queue:
|
|
observation = queue.pop(0)
|
|
if not queue:
|
|
self.pending_tools_by_name.pop(tool_name, None)
|
|
return observation
|
|
|
|
def record_tool_output(self, tool_call_id: str, output: Any) -> None:
|
|
"""Backfill the generation's tool_call record so it carries the result alongside arguments."""
|
|
for tool_call in reversed(self.turn_tool_calls):
|
|
if tool_call.get("id") == tool_call_id:
|
|
tool_call["output"] = output
|
|
function_payload = tool_call.get("function")
|
|
if isinstance(function_payload, dict):
|
|
function_payload["output"] = output
|
|
return
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Trace keys
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _scope_prefix(task_id: str, session_id: str) -> str:
|
|
if task_id:
|
|
return f"task:{task_id}"
|
|
if session_id:
|
|
return f"session:{session_id}"
|
|
return f"thread:{threading.get_ident()}"
|
|
|
|
|
|
def _trace_key(task_id: str, session_id: str, *, turn_id: str = "", api_request_id: str = "") -> str:
|
|
"""Stable in-process trace scope key for one agent turn.
|
|
|
|
``turn_id``/``api_request_id`` scope state so concurrent requests sharing a
|
|
task/session never collide. ``turn_id`` wins over ``api_request_id`` so the
|
|
turn-level post_llm_call hook (no api_request_id) resolves to the same key
|
|
as request-level hooks. Legacy shape: bare ``task_id`` (no ``task:`` prefix),
|
|
kept for keys minted before turn/request scoping existed.
|
|
"""
|
|
if turn_id:
|
|
return f"{_scope_prefix(task_id, session_id)}:turn:{turn_id}"
|
|
if api_request_id:
|
|
return f"{_scope_prefix(task_id, session_id)}:api:{api_request_id}"
|
|
if task_id:
|
|
return task_id
|
|
return _scope_prefix(task_id, session_id)
|
|
|
|
|
|
def _request_key(api_call_count: Any) -> str:
|
|
return str(api_call_count or 0)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Observations
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _start_root_trace(task_key: str, *, task_id: str, session_id: str, platform: str, provider: str, model: str,
|
|
api_mode: str, messages: Any, client: Any,
|
|
turn_id: str = "", api_request_id: str = "") -> TraceState:
|
|
trace_id = client.create_trace_id(seed=f"{session_id or 'sessionless'}::{task_id or task_key}")
|
|
trace_input = _extract_last_user_message(messages)
|
|
metadata = {
|
|
"source": "hermes", "task_id": task_id, "turn_id": turn_id, "api_request_id": api_request_id,
|
|
"platform": platform, "provider": provider, "model": model, "api_mode": api_mode,
|
|
"capture_mode": _capture_mode(),
|
|
}
|
|
# session_id must be in trace_context for Langfuse session grouping.
|
|
trace_ctx: Dict[str, Any] = {"trace_id": trace_id}
|
|
if session_id:
|
|
trace_ctx["session_id"] = session_id
|
|
|
|
def open_root():
|
|
ctx = client.start_as_current_observation(
|
|
trace_context=trace_ctx, name="Hermes turn", as_type="chain",
|
|
input=trace_input, metadata=metadata, end_on_exit=False,
|
|
)
|
|
return ctx, ctx.__enter__()
|
|
|
|
root_ctx = root_span = None
|
|
if propagate_attributes is not None:
|
|
try:
|
|
with propagate_attributes(session_id=session_id or task_key, trace_name="Hermes turn",
|
|
tags=["hermes", "langfuse"]):
|
|
root_ctx, root_span = open_root()
|
|
except Exception:
|
|
root_ctx = None
|
|
if root_ctx is None:
|
|
root_ctx, root_span = open_root()
|
|
|
|
# SDK v3 uses update_trace(); failures must never block the turn.
|
|
try:
|
|
root_span.update_trace(input=trace_input)
|
|
except Exception as exc:
|
|
_debug(f"update_trace(input) failed: {exc}")
|
|
|
|
_debug(f"started trace {trace_id} for {task_key}")
|
|
return TraceState(trace_id=trace_id, root_ctx=root_ctx, root_span=root_span)
|
|
|
|
|
|
def _start_child_observation(state: TraceState, *, name: str, as_type: str, input_value: Any,
|
|
metadata: Optional[dict] = None, model: Optional[str] = None,
|
|
model_parameters: Optional[dict] = None) -> Any:
|
|
return state.root_span.start_observation(
|
|
name=name, as_type=as_type, input=input_value, metadata=metadata or {},
|
|
model=model, model_parameters=model_parameters,
|
|
)
|
|
|
|
|
|
def _end_observation(observation: Any, *, output: Any = None, metadata: Optional[dict] = None,
|
|
usage_details: Optional[dict] = None, cost_details: Optional[dict] = None) -> None:
|
|
if observation is None:
|
|
return
|
|
try:
|
|
update_kwargs: Dict[str, Any] = {}
|
|
if output is not None:
|
|
update_kwargs["output"] = output
|
|
for key, val in (("metadata", metadata), ("usage_details", usage_details), ("cost_details", cost_details)):
|
|
if val:
|
|
update_kwargs[key] = val
|
|
if update_kwargs:
|
|
observation.update(**update_kwargs)
|
|
observation.end()
|
|
except Exception as exc: # pragma: no cover - fail-open
|
|
_debug(f"end observation failed: {exc}")
|
|
|
|
|
|
def _end_children(state: TraceState, *, include_subagents: bool = False) -> None:
|
|
for observation in (*state.generations.values(), *state.tools.values()):
|
|
_end_observation(observation)
|
|
for queue in state.pending_tools_by_name.values():
|
|
for observation in queue:
|
|
_end_observation(observation)
|
|
if include_subagents:
|
|
for observation in state.subagents.values():
|
|
_end_observation(observation)
|
|
|
|
|
|
def _close_root(state: TraceState, *, label: str, output: Any = None,
|
|
children: bool = True, include_subagents: bool = False) -> None:
|
|
"""End ``state``'s root: children first, then trace output, root ``end()``, context unwind.
|
|
|
|
Never raises. Neither output update may prevent ``end()``, else children
|
|
export without a root. The root context manager is unwound now, while
|
|
``opentelemetry.trace.Span`` is still a real type: a GC-driven close at
|
|
interpreter teardown raises TypeError inside ``use_span``'s isinstance check.
|
|
"""
|
|
try:
|
|
if children:
|
|
_end_children(state, include_subagents=include_subagents)
|
|
if output is not None:
|
|
# update_trace sets TRACE-level I/O (SDK v3); root I/O via update().
|
|
for method, what in (("update_trace", "update_trace(output)"), ("update", "root update(output)")):
|
|
try:
|
|
getattr(state.root_span, method)(output=output)
|
|
except Exception as exc:
|
|
_debug(f"{what} failed: {exc}")
|
|
try:
|
|
state.root_span.end()
|
|
except Exception as exc:
|
|
_debug(f"root end() failed: {exc}")
|
|
if state.root_ctx is not None:
|
|
try:
|
|
state.root_ctx.__exit__(None, None, None)
|
|
except Exception: # pragma: no cover - fail-open
|
|
pass
|
|
except Exception as exc: # pragma: no cover - fail-open
|
|
_debug(f"{label} failed: {exc}")
|
|
# Last-chance end so an unexpected error still exports the root.
|
|
try:
|
|
state.root_span.end()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _merge_trace_output(output: Any, state: TraceState) -> Any:
|
|
if not state.turn_tool_calls:
|
|
return output
|
|
merged = dict(output) if isinstance(output, dict) else {"content": output}
|
|
merged["tool_calls"] = list(state.turn_tool_calls)
|
|
return merged
|