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

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
+

traccia/runtime_config.py

@@ -0,0 +1,106 @@
+"""Runtime configuration state management."""
+
+from typing import Optional, List
+
+# Global runtime configuration state
+_config = {
+    "auto_instrument_tools": False,
+    "tool_include": [],
+    "max_tool_spans": 1000,
+    "max_span_depth": 100,
+    "session_id": None,
+    "user_id": None,
+    "tenant_id": None,
+    "project_id": None,
+    "agent_id": None,
+    "debug": False,
+    "attr_truncation_limit": 1000,
+}
+
+
+def set_auto_instrument_tools(value: bool) -> None:
+    _config["auto_instrument_tools"] = value
+
+
+def get_auto_instrument_tools() -> bool:
+    return _config["auto_instrument_tools"]
+
+
+def set_tool_include(value: List[str]) -> None:
+    _config["tool_include"] = value
+
+
+def get_tool_include() -> List[str]:
+    return _config["tool_include"]
+
+
+def set_max_tool_spans(value: int) -> None:
+    _config["max_tool_spans"] = value
+
+
+def get_max_tool_spans() -> int:
+    return _config["max_tool_spans"]
+
+
+def set_max_span_depth(value: int) -> None:
+    _config["max_span_depth"] = value
+
+
+def get_max_span_depth() -> int:
+    return _config["max_span_depth"]
+
+
+def set_session_id(value: Optional[str]) -> None:
+    _config["session_id"] = value
+
+
+def get_session_id() -> Optional[str]:
+    return _config["session_id"]
+
+
+def set_user_id(value: Optional[str]) -> None:
+    _config["user_id"] = value
+
+
+def get_user_id() -> Optional[str]:
+    return _config["user_id"]
+
+
+def set_tenant_id(value: Optional[str]) -> None:
+    _config["tenant_id"] = value
+
+
+def get_tenant_id() -> Optional[str]:
+    return _config["tenant_id"]
+
+
+def set_project_id(value: Optional[str]) -> None:
+    _config["project_id"] = value
+
+
+def get_project_id() -> Optional[str]:
+    return _config["project_id"]
+
+
+def set_agent_id(value: Optional[str]) -> None:
+    _config["agent_id"] = value
+
+
+def get_agent_id() -> Optional[str]:
+    return _config["agent_id"]
+
+
+def set_debug(value: bool) -> None:
+    _config["debug"] = value
+
+
+def get_debug() -> bool:
+    return _config["debug"]
+
+
+def set_attr_truncation_limit(value: int) -> None:
+    _config["attr_truncation_limit"] = value
+
+
+def get_attr_truncation_limit() -> int:
+    return _config["attr_truncation_limit"]

traccia/tracer/__init__.py

@@ -0,0 +1,15 @@
+"""Tracer components for the tracing SDK."""
+
+from traccia.tracer.provider import SpanProcessor, TracerProvider
+from traccia.tracer.span import Span, SpanStatus
+from traccia.tracer.span_context import SpanContext
+from traccia.tracer.tracer import Tracer
+
+__all__ = [
+    "Span",
+    "SpanStatus",
+    "SpanContext",
+    "Tracer",
+    "TracerProvider",
+    "SpanProcessor",
+]