"""Trace/span instrumentation for agent runtimes.
Trace and span IDs follow W3C Trace Context format (32-hex / 16-hex,
no prefix) so they round-trip through any OTel-aware collector. When
``opentelemetry`` is installed, every Taimoe span is mirrored as an OTel
span via :mod:`._otel` so downstream APMs pick the trace up automatically.
"""
from __future__ import annotations
import sys
from contextvars import ContextVar, Token
from dataclasses import dataclass
from datetime import datetime, timezone
from time import perf_counter
from types import TracebackType
from typing import Any
from taimoe.platform._ids import new_span_id, new_trace_id, normalize_trace_id
from taimoe.platform.client import TaimoeClient
from taimoe.platform.errors import ConfigurationError
from taimoe.platform.types.observability import SpanBatch, SpanEvent, SpanType
from ._otel import start_bridged_span
_current_trace: ContextVar[Any] = ContextVar("taimoe_current_trace", default=None)
_current_span: ContextVar[Any] = ContextVar("taimoe_current_span", default=None)
[docs]
@dataclass(frozen=True)
class ObservabilityConfig:
"""Runtime observability configuration."""
runtime_id: str
base_url: str
api_key: str | None = None
enabled: bool = True
raise_on_export_error: bool = False
timeout: float = 10.0
[docs]
class ObservabilityState:
"""Holds the active exporter client for process-wide instrumentation."""
def __init__(self) -> None:
self.config: ObservabilityConfig | None = None
self.client: TaimoeClient | None = None
[docs]
def require_config(self) -> ObservabilityConfig:
if self.config is None:
raise ConfigurationError("taimoe.init() must be called before taimoe.trace().")
return self.config
[docs]
def export(self, batch: SpanBatch) -> None:
config = self.require_config()
if not config.enabled:
return
if self.client is None:
raise ConfigurationError("taimoe observability client is not configured.")
self.client.observability.submit_spans(batch)
_state = ObservabilityState()
[docs]
class TaimoeTrace:
"""Context manager representing one agent run or conversation trace."""
def __init__(
self,
name: str,
*,
trace_id: str | None = None,
session_id: str | None = None,
agent_id: str | None = None,
team_id: str | None = None,
attributes: dict[str, Any] | None = None,
state: ObservabilityState = _state,
) -> None:
config = state.require_config()
self.name = name
self.trace_id = normalize_trace_id(trace_id) if trace_id else new_trace_id()
self.session_id = session_id
self.runtime_id = config.runtime_id
self.agent_id = agent_id
self.team_id = team_id
self.attributes = attributes or {}
self._state = state
self._spans: list[SpanEvent] = []
self._token: Token[TaimoeTrace | None] | None = None
self.export_error: Exception | None = None
def __enter__(self) -> TaimoeTrace:
self._token = _current_trace.set(self)
return self
def __exit__(
self,
exc_type: type[BaseException] | None,
exc: BaseException | None,
traceback: TracebackType | None,
) -> bool:
self.flush()
if self._token is not None:
_current_trace.reset(self._token)
return False
@property
def spans(self) -> tuple[SpanEvent, ...]:
return tuple(self._spans)
[docs]
def span(
self,
span_type: SpanType,
*,
name: str,
input: Any = None,
attributes: dict[str, Any] | None = None,
agent_id: str | None = None,
) -> TaimoeSpan:
parent = _current_span.get()
parent_span_id = parent.span_id if parent is not None else None
return TaimoeSpan(
trace=self,
span_type=span_type,
name=name,
parent_span_id=parent_span_id,
input=input,
attributes=attributes,
agent_id=agent_id or self.agent_id,
)
[docs]
def record_span(self, span: SpanEvent) -> None:
self._spans.append(span)
[docs]
def flush(self) -> None:
if not self._spans:
return
batch = SpanBatch(
runtime_id=self.runtime_id,
trace_id=self.trace_id,
session_id=self.session_id,
spans=tuple(self._spans),
)
try:
self._state.export(batch)
except Exception as error:
self.export_error = error
print(f"FAILED TO EXPORT TRACE: {error}", file=sys.stderr)
config = self._state.require_config()
if config.raise_on_export_error:
raise
[docs]
class TaimoeSpan:
"""Context manager for one agent, LLM, tool, app, or policy operation."""
def __init__(
self,
*,
trace: TaimoeTrace,
span_type: SpanType,
name: str,
parent_span_id: str | None = None,
input: Any = None,
attributes: dict[str, Any] | None = None,
agent_id: str | None = None,
) -> None:
self.trace = trace
self.span_type = span_type
self.name = name
self.span_id = new_span_id()
self.parent_span_id = parent_span_id
self.input = input
self.output: Any = None
self.attributes = attributes or {}
self.agent_id = agent_id
self.prompt_tokens: int | None = None
self.completion_tokens: int | None = None
self.cost_usd: float | None = None
self.error: str | None = None
self._start_time: datetime | None = None
self._start_counter: float | None = None
self._token: Token[TaimoeSpan | None] | None = None
self._finished = False
self._otel_handle: Any = None
def __enter__(self) -> TaimoeSpan:
self._start_time = datetime.now(timezone.utc)
self._start_counter = perf_counter()
self._token = _current_span.set(self)
self._otel_handle = start_bridged_span(
name=self.name,
trace_id_hex=self.trace.trace_id,
span_id_hex=self.span_id,
parent_span_id_hex=self.parent_span_id,
attributes={
"taimoe.span_type": str(self.span_type),
"taimoe.runtime_id": self.trace.runtime_id,
**({"taimoe.agent_id": self.agent_id} if self.agent_id else {}),
},
)
return self
def __exit__(
self,
exc_type: type[BaseException] | None,
exc: BaseException | None,
traceback: TracebackType | None,
) -> bool:
if exc is not None:
self.set_error(exc)
self.finish()
if self._token is not None:
_current_span.reset(self._token)
return False
[docs]
def set_output(self, value: Any) -> None:
self.output = value
[docs]
def set_attribute(self, key: str, value: Any) -> None:
self.attributes[key] = value
[docs]
def set_usage(
self,
*,
prompt_tokens: int | None = None,
completion_tokens: int | None = None,
cost_usd: float | None = None,
) -> None:
self.prompt_tokens = prompt_tokens
self.completion_tokens = completion_tokens
self.cost_usd = cost_usd
[docs]
def set_error(self, error: BaseException | str) -> None:
self.error = str(error)
if self._otel_handle is not None:
self._otel_handle.record_error(error)
[docs]
def finish(self) -> SpanEvent:
if self._start_time is None or self._start_counter is None:
raise RuntimeError("span.finish() called before entering the span context.")
if self._finished:
raise RuntimeError("span.finish() called more than once.")
end_time = datetime.now(timezone.utc)
event = SpanEvent(
trace_id=self.trace.trace_id,
span_id=self.span_id,
parent_span_id=self.parent_span_id,
session_id=self.trace.session_id,
runtime_id=self.trace.runtime_id,
agent_id=self.agent_id,
team_id=self.trace.team_id,
span_type=self.span_type,
name=self.name,
status="error" if self.error else "ok",
start_time=self._start_time,
end_time=end_time,
latency_ms=(perf_counter() - self._start_counter) * 1000,
input=self.input,
output=self.output,
error=self.error,
prompt_tokens=self.prompt_tokens,
completion_tokens=self.completion_tokens,
cost_usd=self.cost_usd,
attributes=self.attributes,
)
self.trace.record_span(event)
if self._otel_handle is not None:
self._otel_handle.end()
self._finished = True
return event
[docs]
def init(
*,
runtime_id: str,
base_url: str | None = None,
platform_url: str | None = None,
api_key: str | None = None,
enabled: bool = True,
raise_on_export_error: bool = False,
timeout: float = 10.0,
) -> None:
"""Configure process-wide observability export to the Taimoe Platform."""
resolved_base_url = base_url or platform_url
if not resolved_base_url:
raise ConfigurationError("taimoe.init() requires base_url or platform_url.")
_state.configure(
runtime_id=runtime_id,
base_url=resolved_base_url,
api_key=api_key,
enabled=enabled,
raise_on_export_error=raise_on_export_error,
timeout=timeout,
)
[docs]
def trace(
name: str,
*,
trace_id: str | None = None,
session_id: str | None = None,
agent_id: str | None = None,
team_id: str | None = None,
attributes: dict[str, Any] | None = None,
) -> TaimoeTrace:
"""Start a trace for one agent run or conversation."""
return TaimoeTrace(
name,
trace_id=trace_id,
session_id=session_id,
agent_id=agent_id,
team_id=team_id,
attributes=attributes,
)
[docs]
def current_trace() -> TaimoeTrace | None:
"""Return the active trace for the current context."""
return _current_trace.get()
[docs]
def current_span() -> TaimoeSpan | None:
"""Return the active span for the current context."""
return _current_span.get()