traccia 0.1.2__py3-none-any.whl → 0.1.6__py3-none-any.whl

traccia/integrations/openai_agents/processor.py

@@ -0,0 +1,262 @@
+"""Traccia processor for OpenAI Agents SDK tracing."""
+
+from __future__ import annotations
+
+import json
+from typing import Any, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    try:
+        from agents.tracing import Span, Trace
+        from agents.tracing.span_data import (
+            AgentSpanData,
+            FunctionSpanData,
+            GenerationSpanData,
+            HandoffSpanData,
+            GuardrailSpanData,
+            ResponseSpanData,
+            CustomSpanData,
+        )
+    except ImportError:
+        pass
+
+
+class TracciaAgentsTracingProcessor:
+    """
+    Traccia processor for OpenAI Agents SDK.
+    
+    Implements the Agents SDK TracingProcessor interface to capture agent
+    runs, tool calls, handoffs, and generations into Traccia spans.
+    """
+
+    def __init__(self):
+        """Initialize the processor."""
+        self._trace_map = {}  # Map Agents trace_id -> Traccia trace context
+        self._span_map = {}   # Map Agents span_id -> Traccia span
+        self._tracer = None
+
+    def _get_tracer(self):
+        """Get or create the Traccia tracer."""
+        if self._tracer is None:
+            import traccia
+            self._tracer = traccia.get_tracer("openai.agents")
+        return self._tracer
+
+    def on_trace_start(self, trace: Trace) -> None:
+        """Called when an Agents trace starts."""
+        try:
+            # Store trace context for span correlation
+            self._trace_map[trace.trace_id] = {
+                "trace_id": trace.trace_id,
+                "started_at": getattr(trace, "started_at", None),
+            }
+        except Exception:
+            # Don't break agent execution on tracing errors
+            pass
+
+    def on_trace_end(self, trace: Trace) -> None:
+        """Called when an Agents trace ends."""
+        try:
+            # Clean up trace mapping
+            self._trace_map.pop(trace.trace_id, None)
+        except Exception:
+            pass
+
+    def on_span_start(self, span: Span[Any]) -> None:
+        """Called when an Agents span starts."""
+        try:
+            tracer = self._get_tracer()
+            span_data = span.span_data
+            
+            # Determine span name based on span type
+            span_name = self._get_span_name(span_data)
+            
+            # Start Traccia span
+            attributes = self._extract_attributes(span_data)
+            traccia_span = tracer.start_span(span_name, attributes=attributes)
+            
+            # Store mapping
+            self._span_map[span.span_id] = traccia_span
+        except Exception:
+            # Don't break agent execution
+            pass
+
+    def on_span_end(self, span: Span[Any]) -> None:
+        """Called when an Agents span ends."""
+        try:
+            traccia_span = self._span_map.pop(span.span_id, None)
+            if not traccia_span:
+                return
+            
+            # Update attributes with final data
+            span_data = span.span_data
+            self._update_span_attributes(traccia_span, span_data)
+            
+            # Record error if present
+            error = getattr(span, "error", None)
+            if error:
+                from traccia.tracer.span import SpanStatus
+                error_msg = str(error.get("message", "Unknown error") if isinstance(error, dict) else error)
+                traccia_span.set_status(SpanStatus.ERROR, error_msg)
+            
+            # End the span
+            traccia_span.end()
+        except Exception:
+            # Ensure span ends even if there's an error
+            try:
+                if traccia_span:
+                    traccia_span.end()
+            except:
+                pass
+
+    def _get_span_name(self, span_data: Any) -> str:
+        """Determine Traccia span name from Agents span data."""
+        span_type = getattr(span_data, "type", "unknown")
+        
+        if span_type == "agent":
+            agent_name = getattr(span_data, "name", "unknown")
+            return f"agent.{agent_name}"
+        elif span_type == "generation":
+            return "llm.agents.generation"
+        elif span_type == "function":
+            func_name = getattr(span_data, "name", "unknown")
+            return f"agent.tool.{func_name}"
+        elif span_type == "handoff":
+            return "agent.handoff"
+        elif span_type == "guardrail":
+            guardrail_name = getattr(span_data, "name", "unknown")
+            return f"agent.guardrail.{guardrail_name}"
+        elif span_type == "response":
+            return "agent.response"
+        elif span_type == "custom":
+            custom_name = getattr(span_data, "name", "unknown")
+            return f"agent.custom.{custom_name}"
+        else:
+            return f"agent.{span_type}"
+
+    def _extract_attributes(self, span_data: Any) -> dict[str, Any]:
+        """Extract initial attributes from Agents span data."""
+        attrs = {
+            "agent.span.type": getattr(span_data, "type", "unknown"),
+        }
+        
+        span_type = getattr(span_data, "type", None)
+        
+        if span_type == "agent":
+            attrs["agent.name"] = getattr(span_data, "name", None)
+            tools = getattr(span_data, "tools", None)
+            if tools:
+                attrs["agent.tools"] = json.dumps(tools)[:500]
+            handoffs = getattr(span_data, "handoffs", None)
+            if handoffs:
+                attrs["agent.handoffs"] = json.dumps(handoffs)[:500]
+            output_type = getattr(span_data, "output_type", None)
+            if output_type:
+                attrs["agent.output_type"] = str(output_type)
+                
+        elif span_type == "generation":
+            model = getattr(span_data, "model", None)
+            if model:
+                attrs["llm.model"] = str(model)
+            model_config = getattr(span_data, "model_config", None)
+            if model_config:
+                attrs["llm.model_config"] = json.dumps(model_config)[:500]
+                
+        elif span_type == "function":
+            func_name = getattr(span_data, "name", None)
+            if func_name:
+                attrs["agent.tool.name"] = func_name
+                
+        elif span_type == "handoff":
+            from_agent = getattr(span_data, "from_agent", None)
+            to_agent = getattr(span_data, "to_agent", None)
+            if from_agent:
+                attrs["agent.handoff.from"] = from_agent
+            if to_agent:
+                attrs["agent.handoff.to"] = to_agent
+                
+        elif span_type == "guardrail":
+            guardrail_name = getattr(span_data, "name", None)
+            if guardrail_name:
+                attrs["agent.guardrail.name"] = guardrail_name
+        
+        return attrs
+
+    def _update_span_attributes(self, traccia_span: Any, span_data: Any) -> None:
+        """Update Traccia span with final attributes from completed Agents span."""
+        try:
+            span_type = getattr(span_data, "type", None)
+            
+            if span_type == "generation":
+                # Add usage info
+                usage = getattr(span_data, "usage", None)
+                if usage and isinstance(usage, dict):
+                    input_tokens = usage.get("input_tokens")
+                    output_tokens = usage.get("output_tokens")
+                    if input_tokens is not None:
+                        traccia_span.set_attribute("llm.usage.input_tokens", input_tokens)
+                        traccia_span.set_attribute("llm.usage.prompt_tokens", input_tokens)
+                    if output_tokens is not None:
+                        traccia_span.set_attribute("llm.usage.output_tokens", output_tokens)
+                        traccia_span.set_attribute("llm.usage.completion_tokens", output_tokens)
+                    if input_tokens is not None and output_tokens is not None:
+                        traccia_span.set_attribute("llm.usage.total_tokens", input_tokens + output_tokens)
+                
+                # Add truncated input/output for observability
+                input_data = getattr(span_data, "input", None)
+                if input_data:
+                    try:
+                        input_str = json.dumps(input_data)[:1000]
+                        traccia_span.set_attribute("llm.input", input_str)
+                    except:
+                        traccia_span.set_attribute("llm.input", str(input_data)[:1000])
+                
+                output_data = getattr(span_data, "output", None)
+                if output_data:
+                    try:
+                        output_str = json.dumps(output_data)[:1000]
+                        traccia_span.set_attribute("llm.output", output_str)
+                    except:
+                        traccia_span.set_attribute("llm.output", str(output_data)[:1000])
+            
+            elif span_type == "function":
+                # Add function input/output
+                func_input = getattr(span_data, "input", None)
+                if func_input:
+                    traccia_span.set_attribute("agent.tool.input", str(func_input)[:500])
+                
+                func_output = getattr(span_data, "output", None)
+                if func_output:
+                    traccia_span.set_attribute("agent.tool.output", str(func_output)[:500])
+                
+                mcp_data = getattr(span_data, "mcp_data", None)
+                if mcp_data:
+                    traccia_span.set_attribute("agent.tool.mcp", json.dumps(mcp_data)[:500])
+            
+            elif span_type == "guardrail":
+                triggered = getattr(span_data, "triggered", False)
+                traccia_span.set_attribute("agent.guardrail.triggered", triggered)
+            
+            elif span_type == "response":
+                response = getattr(span_data, "response", None)
+                if response:
+                    response_id = getattr(response, "id", None)
+                    if response_id:
+                        traccia_span.set_attribute("agent.response.id", response_id)
+        
+        except Exception:
+            # Don't break tracing on attribute errors
+            pass
+
+    def shutdown(self) -> None:
+        """Shutdown the processor."""
+        try:
+            self._trace_map.clear()
+            self._span_map.clear()
+        except Exception:
+            pass
+
+    def force_flush(self) -> None:
+        """Force flush any queued spans."""
+        # Traccia handles flushing at the provider level
+        pass

traccia/pricing_config.py

@@ -0,0 +1,58 @@
+"""Pricing configuration fetcher with optional env override.
+
+Pricing should be treated as configuration, not source code: vendors update
+prices and model versions frequently. The SDK therefore supports:
+- defaults (stub)
+- env override: AGENT_DASHBOARD_PRICING_JSON
+- direct override via start_tracing(pricing_override=...)
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from typing import Dict, Literal, Tuple
+
+from traccia.processors.cost_engine import DEFAULT_PRICING
+
+
+def fetch_remote_pricing() -> Dict[str, Dict[str, float]]:
+    """
+    Placeholder for remote pricing sync.
+    In production this would fetch from backend service; here we return defaults.
+    """
+    return DEFAULT_PRICING.copy()
+
+
+PricingSource = Literal["default", "env", "override"]
+
+
+def load_pricing_with_source(
+    override: Dict[str, Dict[str, float]] | None = None,
+) -> Tuple[Dict[str, Dict[str, float]], PricingSource]:
+    """
+    Return (pricing_table, source_of_latest_override).
+    """
+    pricing = fetch_remote_pricing()
+    source: PricingSource = "default"
+
+    env_override = os.getenv("AGENT_DASHBOARD_PRICING_JSON")
+    if env_override:
+        try:
+            env_pricing = json.loads(env_override)
+            if isinstance(env_pricing, dict):
+                pricing.update(env_pricing)
+                source = "env"
+        except Exception:
+            pass
+    if override:
+        pricing.update(override)
+        source = "override"
+    return pricing, source
+
+
+def load_pricing(override: Dict[str, Dict[str, float]] | None = None) -> Dict[str, Dict[str, float]]:
+    """Backward-compatible helper returning only the pricing table."""
+    pricing, _ = load_pricing_with_source(override)
+    return pricing
+

traccia/processors/__init__.py

@@ -0,0 +1,35 @@
+"""Span processors and supporting utilities."""
+
+from traccia.processors.batch_processor import BatchSpanProcessor
+from traccia.processors.drop_policy import (
+    DEFAULT_DROP_POLICY,
+    DropNewestPolicy,
+    DropOldestPolicy,
+    DropPolicy,
+)
+from traccia.processors.sampler import Sampler, SamplingResult
+from traccia.processors.token_counter import TokenCountingProcessor, estimate_tokens_from_text
+from traccia.processors.cost_engine import compute_cost, DEFAULT_PRICING
+from traccia.processors.cost_processor import CostAnnotatingProcessor
+from traccia.processors.logging_processor import LoggingSpanProcessor
+from traccia.processors.agent_enricher import AgentEnrichmentProcessor
+from traccia.processors.rate_limiter import RateLimiter, RateLimitingSpanProcessor
+
+__all__ = [
+    "BatchSpanProcessor",
+    "DropPolicy",
+    "DropOldestPolicy",
+    "DropNewestPolicy",
+    "DEFAULT_DROP_POLICY",
+    "Sampler",
+    "SamplingResult",
+    "TokenCountingProcessor",
+    "estimate_tokens_from_text",
+    "compute_cost",
+    "DEFAULT_PRICING",
+    "CostAnnotatingProcessor",
+    "LoggingSpanProcessor",
+    "AgentEnrichmentProcessor",
+    "RateLimiter",
+    "RateLimitingSpanProcessor",
+]

traccia/processors/agent_enricher.py

@@ -0,0 +1,159 @@
+"""Span processor that enriches spans with agent metadata and cost."""
+
+from __future__ import annotations
+
+import json
+import os
+from typing import Any, Dict, Optional
+
+from traccia.processors.cost_engine import compute_cost
+from traccia.tracer.provider import SpanProcessor
+
+
+def _load_agent_catalog(path: Optional[str]) -> Dict[str, Dict[str, Any]]:
+    """
+    Load agent metadata from a JSON file.
+    Supports:
+      { "agents": [ { "id": "...", "name": "...", ... } ] }
+      or { "agent-id": { "name": "...", ... }, ... }
+    """
+    if not path:
+        return {}
+    if not os.path.exists(path):
+        return {}
+    try:
+        with open(path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+    except Exception:
+        return {}
+    if isinstance(data, dict) and "agents" in data and isinstance(data["agents"], list):
+        out = {}
+        for agent in data["agents"]:
+            if not isinstance(agent, dict):
+                continue
+            aid = agent.get("id")
+            if aid:
+                out[str(aid)] = agent
+        return out
+    if isinstance(data, dict):
+        return {str(k): v for k, v in data.items() if isinstance(v, dict)}
+    return {}
+
+
+class AgentEnrichmentProcessor(SpanProcessor):
+    """
+    Enrich spans with agent metadata (id/name/env/owner/team/org) and compute llm.cost.usd if missing.
+    Static metadata can come from:
+      - span attributes (preferred)
+      - environment variables (AGENT_DASHBOARD_AGENT_ID/NAME/ENV/OWNER/TEAM/ORG_ID/SUB_ORG_ID/DESCRIPTION)
+      - JSON config file pointed by AGENT_DASHBOARD_AGENT_CONFIG
+    """
+
+    def __init__(
+        self,
+        *,
+        agent_config_path: Optional[str] = None,
+        default_agent_id: Optional[str] = None,
+        default_env: str = "production",
+    ) -> None:
+        self.default_agent_id = default_agent_id or os.getenv("AGENT_DASHBOARD_AGENT_ID")
+        self.default_env = os.getenv("AGENT_DASHBOARD_ENV") or default_env
+        self.default_name = os.getenv("AGENT_DASHBOARD_AGENT_NAME")
+        self.default_type = os.getenv("AGENT_DASHBOARD_AGENT_TYPE")
+        self.default_owner = os.getenv("AGENT_DASHBOARD_AGENT_OWNER")
+        self.default_team = os.getenv("AGENT_DASHBOARD_AGENT_TEAM")
+        self.default_org = os.getenv("AGENT_DASHBOARD_ORG_ID")
+        self.default_sub_org = os.getenv("AGENT_DASHBOARD_SUB_ORG_ID")
+        self.default_description = os.getenv("AGENT_DASHBOARD_AGENT_DESCRIPTION")
+        cfg_path = (
+            agent_config_path
+            or os.getenv("AGENT_DASHBOARD_AGENT_CONFIG")
+            or "agent_config.json"
+        )
+        self.catalog = _load_agent_catalog(cfg_path)
+        # If only one agent is declared, remember it for convenient fallback.
+        self.single_agent_id: Optional[str] = None
+        if len(self.catalog) == 1:
+            self.single_agent_id = next(iter(self.catalog.keys()))
+
+    def on_end(self, span) -> None:
+        attrs = span.attributes
+        # Resolve agent id
+        agent_id = (
+            attrs.get("agent.id")
+            or attrs.get("agent")
+            or self.default_agent_id
+        )
+        # Try using tracer instrumentation scope as a fallback id
+        if not agent_id and getattr(span, "tracer", None) is not None:
+            agent_id = getattr(span.tracer, "instrumentation_scope", None)
+        # If not found in attributes/env/scope, and only one agent exists in catalog, use it
+        if not agent_id and self.single_agent_id:
+            agent_id = self.single_agent_id
+        # If still missing, skip enrichment
+        if not agent_id:
+            return
+
+        # Look up static metadata
+        meta = self.catalog.get(agent_id, {})
+        # If the resolved id is not in catalog but we have a single agent defined, use that entry
+        if not meta and self.single_agent_id:
+            agent_id = self.single_agent_id
+            meta = self.catalog.get(agent_id, {})
+
+        def set_if_missing(key: str, value: Any) -> None:
+            if value is None:
+                return
+            if key not in attrs or attrs.get(key) in (None, ""):
+                attrs[key] = value
+
+        attrs["agent.id"] = agent_id
+        set_if_missing("agent.name", meta.get("name") or self.default_name or agent_id)
+        set_if_missing("agent.type", meta.get("type") or self.default_type or "workflow")
+        set_if_missing("agent.description", meta.get("description") or self.default_description or "")
+        set_if_missing("owner", meta.get("owner") or self.default_owner)
+        set_if_missing("team", meta.get("team") or self.default_team)
+        set_if_missing("org.id", meta.get("org_id") or self.default_org)
+        set_if_missing("sub_org.id", meta.get("sub_org_id") or self.default_sub_org)
+
+        # Environment
+        set_if_missing("env", meta.get("env") or self.default_env)
+        set_if_missing("environment", meta.get("env") or self.default_env)
+
+        # Consumers (store as list)
+        consumers = meta.get("consuming_teams")
+        if consumers and "agent.consuming_teams" not in attrs:
+            attrs["agent.consuming_teams"] = consumers
+
+        # Cost: fill llm.cost.usd if we have tokens + model
+        if "llm.cost.usd" not in attrs:
+            model = attrs.get("llm.model")
+            prompt_tokens = attrs.get("llm.usage.prompt_tokens") or 0
+            completion_tokens = attrs.get("llm.usage.completion_tokens") or 0
+            if model and (prompt_tokens or completion_tokens):
+                try:
+                    cost = compute_cost(
+                        model=model,
+                        prompt_tokens=int(prompt_tokens or 0),
+                        completion_tokens=int(completion_tokens or 0),
+                    )
+                    if cost is not None:
+                        attrs["llm.cost.usd"] = cost
+                except Exception:
+                    pass
+
+        # Span type inference if missing
+        if "span.type" not in attrs and "type" not in attrs:
+            span_type = None
+            if attrs.get("llm.model"):
+                span_type = "LLM"
+            elif attrs.get("tool.name") or attrs.get("tool") or attrs.get("http.url"):
+                span_type = "TOOL"
+            if span_type:
+                attrs["span.type"] = span_type
+
+    def shutdown(self) -> None:
+        return
+
+    def force_flush(self, timeout: Optional[float] = None) -> None:
+        return

traccia/processors/batch_processor.py

@@ -0,0 +1,140 @@
+"""Batching span processor with bounded queue and background flush."""
+
+from __future__ import annotations
+
+import threading
+import time
+from collections import deque
+from typing import Deque, Iterable, List, Optional
+
+from traccia.processors.drop_policy import DEFAULT_DROP_POLICY, DropPolicy
+from traccia.processors.sampler import Sampler
+from traccia.tracer.provider import SpanProcessor
+from traccia.tracer.span import Span
+
+
+class BatchSpanProcessor(SpanProcessor):
+    """
+    Batch span processor that queues spans for export.
+    
+    Note: This runs as an enrichment processor (before span.end()),
+    but it queues spans and exports them after they end.
+    When exporting, it extracts ReadableSpan from the OTel span.
+    """
+    
+    def __init__(
+        self,
+        exporter=None,
+        *,
+        max_queue_size: int = 5000,
+        max_export_batch_size: int = 512,
+        schedule_delay_millis: int = 5000,
+        drop_policy: Optional[DropPolicy] = None,
+        sampler: Optional[Sampler] = None,
+    ) -> None:
+        self.exporter = exporter
+        self.max_queue_size = max_queue_size
+        self.max_export_batch_size = max_export_batch_size
+        self.schedule_delay = schedule_delay_millis / 1000.0
+        self.drop_policy = drop_policy or DEFAULT_DROP_POLICY
+        self.sampler = sampler
+
+        self._queue: Deque[Span] = deque()
+        self._lock = threading.Lock()
+        self._event = threading.Event()
+        self._shutdown = False
+        self._worker = threading.Thread(target=self._worker_loop, daemon=True)
+        self._worker.start()
+
+    def on_end(self, span: Span) -> None:
+        """
+        Called when a span ends (BEFORE span.end() is called).
+        
+        We queue the span here, but it hasn't ended yet.
+        The span will end after enrichment processors run.
+        
+        Note: We mark the span as queued to prevent double-queuing.
+        """
+        if self._shutdown:
+            return
+
+        # Head-based sampling is recorded on SpanContext.trace_flags.
+        # If a sampler is configured, traces marked as not-sampled (0) are dropped.
+        if self.sampler and getattr(span.context, "trace_flags", 1) == 0:
+            return
+
+        # Prevent double-queuing (span might be queued multiple times if on_end is called multiple times)
+        if hasattr(span, '_batch_queued') and span._batch_queued:
+            return
+
+        with self._lock:
+            enqueued = self.drop_policy.handle(self._queue, span, self.max_queue_size)
+            if enqueued:
+                span._batch_queued = True  # Mark as queued
+                self._event.set()
+
+    def force_flush(self, timeout: Optional[float] = None) -> None:
+        """Force flush any pending spans."""
+        deadline = time.time() + timeout if timeout else None
+        while True:
+            flushed_any = self._flush_once()
+            if not flushed_any:
+                return
+            if deadline and time.time() >= deadline:
+                return
+
+    def shutdown(self) -> None:
+        """Shutdown the processor."""
+        self._shutdown = True
+        self._event.set()
+        self._worker.join(timeout=self.schedule_delay * 2)
+        self.force_flush()
+
+    # Internal
+    def _worker_loop(self) -> None:
+        """Background worker that periodically flushes spans."""
+        while not self._shutdown:
+            self._event.wait(timeout=self.schedule_delay)
+            self._event.clear()
+            self._flush_once()
+
+    def _flush_once(self) -> bool:
+        """Flush one batch of spans."""
+        spans = self._drain_queue(self.max_export_batch_size)
+        if not spans:
+            return False
+        self._export(spans)
+        return True
+
+    def _drain_queue(self, limit: int) -> List[Span]:
+        """Drain spans from queue up to limit."""
+        items: List[Span] = []
+        with self._lock:
+            while self._queue and len(items) < limit:
+                items.append(self._queue.popleft())
+        return items
+
+    def _export(self, spans: Iterable[Span]) -> None:
+        """
+        Export spans to exporter.
+        
+        Spans should be ended by the time they're flushed from the queue.
+        We filter out any spans that aren't ended yet.
+        """
+        if self.exporter is None:
+            return
+        
+        try:
+            # Filter to only ended spans
+            ended_spans = [span for span in spans if span._ended]
+            
+            if not ended_spans:
+                return
+            
+            # Export spans - exporter will handle conversion if needed
+            self.exporter.export(ended_spans)
+        except Exception as e:
+            # Export errors are swallowed; resilience over strictness.
+            import traceback
+            traceback.print_exc()
+            return