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

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

traccia/processors/cost_engine.py

@@ -0,0 +1,71 @@
+"""Cost calculation based on model pricing and token usage."""
+
+from __future__ import annotations
+
+from typing import Dict, Optional, Tuple
+
+DEFAULT_PRICING: Dict[str, Dict[str, float]] = {
+    # prices per 1k tokens
+    "gpt-4": {"prompt": 0.03, "completion": 0.06},
+    "gpt-4o": {"prompt": 0.005, "completion": 0.015},
+    "gpt-3.5-turbo": {"prompt": 0.0015, "completion": 0.002},
+    "claude-3-opus": {"prompt": 0.015, "completion": 0.075},
+    "claude-3-sonnet": {"prompt": 0.003, "completion": 0.015},
+    # Add other models via pricing overrides:
+    #   - env: AGENT_DASHBOARD_PRICING_JSON='{"model": {"prompt": x, "completion": y}}'
+    #   - start_tracing(pricing_override={...})
+}
+
+def _lookup_price(model: str, table: Dict[str, Dict[str, float]]) -> Optional[Tuple[str, Dict[str, float]]]:
+    """
+    Return (matched_key, price_dict) for a given model name.
+
+    Supports exact matches and prefix matches for version-suffixed model names:
+    e.g. "claude-3-opus-20240229" -> "claude-3-opus"
+         "gpt-4o-2024-08-06" -> "gpt-4o"
+    """
+    if not model:
+        return None
+    m = str(model).strip()
+    if not m:
+        return None
+    # exact (case sensitive + lower)
+    if m in table:
+        return m, table[m]
+    ml = m.lower()
+    if ml in table:
+        return ml, table[ml]
+    # prefix match (longest key wins)
+    for key in sorted(table.keys(), key=len, reverse=True):
+        if ml.startswith(key.lower()):
+            return key, table[key]
+    return None
+
+
+def match_pricing_model_key(
+    model: str, pricing_table: Optional[Dict[str, Dict[str, float]]] = None
+) -> Optional[str]:
+    """Return the pricing table key that would be used for `model`, if any."""
+    table = pricing_table or DEFAULT_PRICING
+    matched = _lookup_price(model, table)
+    if not matched:
+        return None
+    key, _ = matched
+    return key
+
+
+def compute_cost(
+    model: str,
+    prompt_tokens: int,
+    completion_tokens: int,
+    pricing_table: Optional[Dict[str, Dict[str, float]]] = None,
+) -> Optional[float]:
+    table = pricing_table or DEFAULT_PRICING
+    matched = _lookup_price(model, table)
+    if not matched:
+        return None
+    _, price = matched
+    prompt_cost = (prompt_tokens / 1000.0) * price.get("prompt", 0.0)
+    completion_cost = (completion_tokens / 1000.0) * price.get("completion", 0.0)
+    return round(prompt_cost + completion_cost, 6)
+

traccia/processors/cost_processor.py

@@ -0,0 +1,70 @@
+"""Span processor that annotates spans with cost based on token usage and pricing."""
+
+from __future__ import annotations
+
+from typing import Dict, Optional
+
+from traccia.processors.cost_engine import compute_cost, match_pricing_model_key
+from traccia.tracer.provider import SpanProcessor
+
+
+class CostAnnotatingProcessor(SpanProcessor):
+    """
+    Adds `llm.cost.usd` to spans when token usage and model info are available.
+
+    Expects spans to carry:
+      - llm.model (model name)
+      - llm.usage.prompt_tokens
+      - llm.usage.completion_tokens
+    """
+
+    def __init__(
+        self,
+        pricing_table: Optional[Dict[str, Dict[str, float]]] = None,
+        *,
+        pricing_source: str = "default",
+    ) -> None:
+        self.pricing_table = pricing_table or {}
+        self.pricing_source = pricing_source
+
+    def on_end(self, span) -> None:
+        if "llm.cost.usd" in (span.attributes or {}):
+            return
+        model = span.attributes.get("llm.model")
+        prompt = span.attributes.get("llm.usage.prompt_tokens")
+        completion = span.attributes.get("llm.usage.completion_tokens")
+        # Anthropic-style names (also supported)
+        if prompt is None:
+            prompt = span.attributes.get("llm.usage.input_tokens")
+        if completion is None:
+            completion = span.attributes.get("llm.usage.output_tokens")
+        if not model or prompt is None or completion is None:
+            return
+        cost = compute_cost(
+            model,
+            int(prompt),
+            int(completion),
+            pricing_table=self.pricing_table,
+        )
+        if cost is not None:
+            span.set_attribute("llm.cost.usd", cost)
+            # Provenance for downstream analysis.
+            span.set_attribute("llm.cost.source", span.attributes.get("llm.usage.source", "unknown"))
+            span.set_attribute("llm.pricing.source", self.pricing_source)
+            key = match_pricing_model_key(model, self.pricing_table)
+            if key:
+                span.set_attribute("llm.pricing.model_key", key)
+
+    def shutdown(self) -> None:
+        return None
+
+    def force_flush(self, timeout: Optional[float] = None) -> None:
+        return None
+
+    def update_pricing_table(
+        self, pricing_table: Dict[str, Dict[str, float]], pricing_source: Optional[str] = None
+    ) -> None:
+        self.pricing_table = pricing_table
+        if pricing_source:
+            self.pricing_source = pricing_source
+

traccia/processors/drop_policy.py

@@ -0,0 +1,44 @@
+"""Queue overflow handling strategies for span buffering."""
+
+from collections import deque
+from typing import Deque
+
+from traccia.tracer.span import Span
+
+
+class DropPolicy:
+    """Base policy deciding how to handle span queue overflow."""
+
+    def handle(self, queue: Deque[Span], span: Span, max_size: int) -> bool:
+        """
+        Apply the drop policy.
+
+        Returns True if the span was enqueued, False if it was dropped.
+        """
+        raise NotImplementedError
+
+
+class DropOldestPolicy(DropPolicy):
+    """Drop the oldest span to make room for a new one."""
+
+    def handle(self, queue: Deque[Span], span: Span, max_size: int) -> bool:
+        if len(queue) >= max_size and queue:
+            queue.popleft()
+        if len(queue) < max_size:
+            queue.append(span)
+            return True
+        return False
+
+
+class DropNewestPolicy(DropPolicy):
+    """Drop the incoming span if the queue is full."""
+
+    def handle(self, queue: Deque[Span], span: Span, max_size: int) -> bool:
+        if len(queue) < max_size:
+            queue.append(span)
+            return True
+        return False
+
+
+DEFAULT_DROP_POLICY = DropOldestPolicy()
+

traccia/processors/logging_processor.py

@@ -0,0 +1,31 @@
+"""Span processor that logs spans when they end."""
+
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+from traccia.tracer.provider import SpanProcessor
+
+
+class LoggingSpanProcessor(SpanProcessor):
+    """Logs span summary on end using the standard logging module."""
+
+    def __init__(self, logger: Optional[logging.Logger] = None) -> None:
+        self.logger = logger or logging.getLogger("traccia.traces")
+
+    def on_end(self, span) -> None:
+        attrs = span.attributes or {}
+        msg = (
+            f"[trace] name={span.name} trace_id={span.context.trace_id} "
+            f"span_id={span.context.span_id} status={span.status.name} "
+            f"duration_ns={span.duration_ns} attrs={attrs}"
+        )
+        self.logger.info(msg)
+
+    def shutdown(self) -> None:
+        return None
+
+    def force_flush(self, timeout: Optional[float] = None) -> None:
+        return None
+