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

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
+

traccia/processors/rate_limiter.py

@@ -0,0 +1,223 @@
+"""Rate limiting processor for span export."""
+
+from __future__ import annotations
+
+import logging
+import threading
+import time
+from collections import deque
+from typing import Optional
+
+from opentelemetry.sdk.trace import ReadableSpan
+
+from traccia.errors import RateLimitError
+
+logger = logging.getLogger(__name__)
+
+
+class RateLimiter:
+    """
+    Token bucket rate limiter with hybrid blocking/dropping behavior.
+    
+    Features:
+    - Token bucket algorithm for smooth rate limiting
+    - Short blocking period before dropping spans
+    - Detailed logging of dropped spans
+    - Thread-safe implementation
+    """
+    
+    def __init__(
+        self,
+        max_spans_per_second: Optional[float] = None,
+        max_block_ms: int = 100,
+    ):
+        """
+        Initialize rate limiter.
+        
+        Args:
+            max_spans_per_second: Maximum spans per second (None = unlimited)
+            max_block_ms: Maximum milliseconds to block before dropping
+        """
+        self.max_spans_per_second = max_spans_per_second
+        self.max_block_ms = max_block_ms
+        self.enabled = max_spans_per_second is not None and max_spans_per_second > 0
+        
+        # Token bucket state
+        self._tokens: float = max_spans_per_second or 0
+        self._max_tokens: float = max_spans_per_second or 0
+        self._last_refill_time: float = time.time()
+        self._lock = threading.Lock()
+        
+        # Stats
+        self._total_spans = 0
+        self._dropped_spans = 0
+        self._blocked_spans = 0
+        
+        # Recent timestamps for sliding window (backup)
+        self._recent_timestamps: deque = deque()
+        self._window_seconds = 1.0
+    
+    def acquire(self, span: Optional[ReadableSpan] = None) -> bool:
+        """
+        Try to acquire permission to process a span.
+        
+        Returns True if span should be processed, False if it should be dropped.
+        
+        Behavior:
+        1. If unlimited (disabled), always return True
+        2. Try to acquire a token immediately
+        3. If no token, block for up to max_block_ms
+        4. If still no token after blocking, drop and return False
+        
+        Args:
+            span: Optional span for logging purposes
+            
+        Returns:
+            True if span should be processed, False if dropped
+        """
+        if not self.enabled:
+            return True
+        
+        self._total_spans += 1
+        
+        with self._lock:
+            # Refill tokens based on elapsed time
+            self._refill_tokens()
+            
+            # Try to acquire immediately
+            if self._tokens >= 1.0:
+                self._tokens -= 1.0
+                return True
+            
+            # No tokens available, try blocking
+            if self.max_block_ms > 0:
+                block_start = time.time()
+                blocked_ms = 0
+                
+                while blocked_ms < self.max_block_ms:
+                    # Release lock briefly to allow other threads
+                    self._lock.release()
+                    time.sleep(0.001)  # Sleep 1ms
+                    self._lock.acquire()
+                    
+                    # Refill and try again
+                    self._refill_tokens()
+                    if self._tokens >= 1.0:
+                        self._tokens -= 1.0
+                        self._blocked_spans += 1
+                        return True
+                    
+                    blocked_ms = (time.time() - block_start) * 1000
+            
+            # Still no tokens after blocking - drop the span
+            self._dropped_spans += 1
+            
+            # Log dropped span
+            span_name = span.name if span else "unknown"
+            logger.warning(
+                f"Rate limit exceeded - dropping span '{span_name}'. "
+                f"Total dropped: {self._dropped_spans}/{self._total_spans} "
+                f"({self._dropped_spans / self._total_spans * 100:.1f}%)"
+            )
+            
+            return False
+    
+    def _refill_tokens(self) -> None:
+        """Refill tokens based on elapsed time (token bucket algorithm)."""
+        now = time.time()
+        elapsed = now - self._last_refill_time
+        
+        if elapsed > 0:
+            # Add tokens based on rate and elapsed time
+            new_tokens = elapsed * self.max_spans_per_second
+            self._tokens = min(self._max_tokens, self._tokens + new_tokens)
+            self._last_refill_time = now
+    
+    def get_stats(self) -> dict:
+        """Get rate limiting statistics."""
+        with self._lock:
+            drop_rate = (self._dropped_spans / self._total_spans * 100) if self._total_spans > 0 else 0
+            return {
+                "enabled": self.enabled,
+                "max_spans_per_second": self.max_spans_per_second,
+                "total_spans": self._total_spans,
+                "dropped_spans": self._dropped_spans,
+                "blocked_spans": self._blocked_spans,
+                "drop_rate_percent": round(drop_rate, 2),
+                "current_tokens": round(self._tokens, 2),
+            }
+    
+    def reset_stats(self) -> None:
+        """Reset statistics counters."""
+        with self._lock:
+            self._total_spans = 0
+            self._dropped_spans = 0
+            self._blocked_spans = 0
+
+
+class RateLimitingSpanProcessor:
+    """
+    Span processor that enforces rate limiting before passing to next processor.
+    
+    This should be added early in the processor chain to drop spans before
+    they consume resources in downstream processors.
+    """
+    
+    def __init__(
+        self,
+        next_processor,
+        max_spans_per_second: Optional[float] = None,
+        max_block_ms: int = 100,
+    ):
+        """
+        Initialize rate limiting processor.
+        
+        Args:
+            next_processor: Next processor in the chain
+            max_spans_per_second: Maximum spans per second (None = unlimited)
+            max_block_ms: Maximum milliseconds to block before dropping
+        """
+        self.next_processor = next_processor
+        self.rate_limiter = RateLimiter(
+            max_spans_per_second=max_spans_per_second,
+            max_block_ms=max_block_ms,
+        )
+    
+    def on_start(self, span, parent_context=None):
+        """Called when span starts - pass through to next processor."""
+        if self.next_processor and hasattr(self.next_processor, 'on_start'):
+            self.next_processor.on_start(span, parent_context)
+    
+    def on_end(self, span):
+        """
+        Called when span ends - check rate limit before passing to next processor.
+        
+        If rate limit is exceeded, span is dropped and not passed to next processor.
+        """
+        # Check rate limit
+        if not self.rate_limiter.acquire(span):
+            # Span dropped - don't pass to next processor
+            return
+        
+        # Pass to next processor
+        if self.next_processor and hasattr(self.next_processor, 'on_end'):
+            self.next_processor.on_end(span)
+    
+    def shutdown(self):
+        """Shutdown processor and log final stats."""
+        stats = self.rate_limiter.get_stats()
+        if stats["enabled"] and stats["dropped_spans"] > 0:
+            logger.info(
+                f"Rate limiter shutdown. Final stats: "
+                f"{stats['dropped_spans']}/{stats['total_spans']} spans dropped "
+                f"({stats['drop_rate_percent']}%)"
+            )
+        
+        if self.next_processor and hasattr(self.next_processor, 'shutdown'):
+            self.next_processor.shutdown()
+    
+    def force_flush(self, timeout_millis: int = 30000):
+        """Force flush - pass through to next processor."""
+        if self.next_processor and hasattr(self.next_processor, 'force_flush'):
+            return self.next_processor.force_flush(timeout_millis)
+        return True

traccia/processors/sampler.py

@@ -0,0 +1,22 @@
+"""Sampling decisions for traces."""
+
+import random
+from dataclasses import dataclass
+
+
+@dataclass
+class SamplingResult:
+    sampled: bool
+
+
+class Sampler:
+    """Head-based sampler using a fixed probability."""
+
+    def __init__(self, sample_rate: float = 1.0) -> None:
+        if not 0.0 <= sample_rate <= 1.0:
+            raise ValueError("sample_rate must be between 0.0 and 1.0")
+        self.sample_rate = sample_rate
+
+    def should_sample(self) -> SamplingResult:
+        return SamplingResult(sampled=random.random() <= self.sample_rate)
+

traccia/processors/token_counter.py

@@ -0,0 +1,216 @@
+"""Token counting utilities and processor for spans with LLM usage.
+
+Best practice:
+- Prefer provider-reported usage tokens when available.
+- Otherwise, estimate with the vendor tokenizer when available (tiktoken for
+  OpenAI) and record the estimate source on the span.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional, Tuple
+
+from traccia.tracer.provider import SpanProcessor
+
+try:  # optional dependency for accurate counting
+    import tiktoken  # type: ignore
+except Exception:  # pragma: no cover
+    tiktoken = None  # fallback to heuristic
+
+
+MODEL_TO_ENCODING = {
+    # OpenAI mappings (approximate; kept current as of gpt-4o family)
+    "gpt-4o": "o200k_base",
+    "gpt-4o-mini": "o200k_base",
+    "gpt-4": "cl100k_base",
+    "gpt-3.5-turbo": "cl100k_base",
+}
+
+
+def _encoding_for_model(model: Optional[str]):
+    if tiktoken is None:
+        return None
+    if not model:
+        return None
+    m = str(model)
+    # First try tiktoken's model registry (best when available).
+    try:
+        return tiktoken.encoding_for_model(m)
+    except Exception:
+        pass
+    # Then try our explicit mapping, supporting version-suffixed models by prefix.
+    encoding_name = MODEL_TO_ENCODING.get(m)
+    if encoding_name is None:
+        for key in sorted(MODEL_TO_ENCODING.keys(), key=len, reverse=True):
+            if m.startswith(key):
+                encoding_name = MODEL_TO_ENCODING[key]
+                break
+    if encoding_name:
+        try:
+            return tiktoken.get_encoding(encoding_name)
+        except Exception:
+            return None
+    return None
+
+
+def _count_with_tiktoken(text: str, model: Optional[str]) -> Optional[int]:
+    if tiktoken is None or not text:
+        return None
+    encoding = _encoding_for_model(model)
+    if encoding is None:
+        return None
+    try:
+        return len(encoding.encode(text))
+    except Exception:
+        return None
+
+
+def estimate_tokens_from_text(text: str, model: Optional[str] = None) -> int:
+    """
+    Estimate tokens. Prefer model-accurate count via tiktoken when available,
+    otherwise fall back to a rough whitespace split.
+    """
+    if not text:
+        return 0
+    exact = _count_with_tiktoken(text, model)
+    if exact is not None:
+        return exact
+    return len(text.split())
+
+
+def estimate_tokens_from_text_with_source(
+    text: str, model: Optional[str] = None
+) -> Tuple[int, str]:
+    """
+    Return (token_count, source) where source is:
+      - "estimated.tiktoken"
+      - "estimated.heuristic"
+    """
+    if not text:
+        return 0, "estimated.heuristic"
+    exact = _count_with_tiktoken(text, model)
+    if exact is not None:
+        return exact, "estimated.tiktoken"
+    return len(text.split()), "estimated.heuristic"
+
+
+def _openai_chat_overhead(model: Optional[str]) -> Tuple[int, int, int]:
+    """
+    Return (tokens_per_message, tokens_per_name, tokens_for_reply).
+
+    These constants are model-dependent in OpenAI's chat format. For estimation
+    we use a reasonable default that is close for many modern chat models.
+    """
+    # Defaults (works reasonably for gpt-4/4o families as an estimate)
+    return 3, 1, 3
+
+
+def estimate_openai_chat_prompt_tokens_with_source(
+    messages: Any, model: Optional[str] = None
+) -> Optional[Tuple[int, str]]:
+    """
+    Estimate prompt tokens from a list of chat messages.
+
+    This is best-effort and should be treated as an estimate unless provider
+    usage is available.
+    """
+    if not isinstance(messages, (list, tuple)) or not messages:
+        return None
+    if tiktoken is None:
+        # Heuristic fallback: count whitespace tokens across role/content.
+        parts = []
+        for msg in list(messages)[:50]:
+            if not isinstance(msg, dict):
+                continue
+            role = msg.get("role") or ""
+            content = msg.get("content") or ""
+            if not isinstance(content, str):
+                content = str(content)
+            parts.append(f"{role} {content}".strip())
+        text = "\n".join([p for p in parts if p])
+        return (len(text.split()), "estimated.chat_heuristic")
+    try:
+        encoding = _encoding_for_model(model) or tiktoken.get_encoding("cl100k_base")
+    except Exception:
+        return None
+
+    tokens_per_message, tokens_per_name, tokens_for_reply = _openai_chat_overhead(model)
+    total = 0
+    for msg in list(messages)[:50]:
+        if not isinstance(msg, dict):
+            continue
+        total += tokens_per_message
+        role = msg.get("role") or ""
+        name = msg.get("name")
+        content = msg.get("content") or ""
+        if not isinstance(content, str):
+            content = str(content)
+        total += len(encoding.encode(str(role)))
+        total += len(encoding.encode(content))
+        if name:
+            total += tokens_per_name
+            total += len(encoding.encode(str(name)))
+    total += tokens_for_reply
+    return total, "estimated.tiktoken_chat"
+
+
+class TokenCountingProcessor(SpanProcessor):
+    """
+    A processor that infers token counts when not provided by the LLM response.
+    It prefers a model-specific tokenizer (tiktoken) when available.
+    """
+
+    def on_end(self, span) -> None:
+        prompt = span.attributes.get("llm.prompt")
+        completion = span.attributes.get("llm.completion")
+        model = span.attributes.get("llm.model")
+        openai_messages = span.attributes.get("llm.openai.messages")
+
+        wrote_any = False
+        wrote_prompt = False
+        wrote_completion = False
+
+        if "llm.usage.prompt_tokens" not in span.attributes:
+            # Prefer chat-structure estimation when available.
+            est = estimate_openai_chat_prompt_tokens_with_source(openai_messages, model)
+            if est is not None:
+                count, source = est
+                span.set_attribute("llm.usage.prompt_tokens", count)
+                span.set_attribute("llm.usage.prompt_source", source)
+                wrote_any = True
+                wrote_prompt = True
+            elif isinstance(prompt, str):
+                count, source = estimate_tokens_from_text_with_source(prompt, model)
+                span.set_attribute("llm.usage.prompt_tokens", count)
+                span.set_attribute("llm.usage.prompt_source", source)
+                wrote_any = True
+                wrote_prompt = True
+
+        if "llm.usage.completion_tokens" not in span.attributes and isinstance(completion, str):
+            count, source = estimate_tokens_from_text_with_source(completion, model)
+            span.set_attribute("llm.usage.completion_tokens", count)
+            span.set_attribute("llm.usage.completion_source", source)
+            wrote_any = True
+            wrote_completion = True
+
+        # Synthesize overall usage source if not provided by instrumentation.
+        if wrote_any and "llm.usage.source" not in span.attributes:
+            ps = span.attributes.get("llm.usage.prompt_source")
+            cs = span.attributes.get("llm.usage.completion_source")
+            if ps and cs and ps == cs:
+                span.set_attribute("llm.usage.source", ps)
+            elif ps or cs:
+                span.set_attribute("llm.usage.source", "mixed")
+
+        # If provider already marked usage as provider_usage, and we filled any missing
+        # fields, mark it as mixed.
+        if wrote_any and span.attributes.get("llm.usage.source") == "provider_usage":
+            if wrote_prompt or wrote_completion:
+                span.set_attribute("llm.usage.source", "mixed")
+
+    def shutdown(self) -> None:
+        return None
+
+    def force_flush(self, timeout: Optional[float] = None) -> None:
+        return None
+