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

traccia/integrations/langchain/callback.py

@@ -0,0 +1,418 @@
+"""Traccia callback handler for LangChain."""
+
+from __future__ import annotations
+
+import json
+from typing import Any, Dict, List, Optional, Set
+from uuid import UUID
+
+try:
+    from langchain_core.callbacks import BaseCallbackHandler
+    from langchain_core.outputs import LLMResult
+    from langchain_core.messages import BaseMessage
+except ImportError as e:
+    raise ModuleNotFoundError(
+        "LangChain integration requires langchain-core. "
+        "Install with: pip install traccia[langchain]"
+    ) from e
+
+from traccia import get_tracer
+from traccia.tracer.span import SpanStatus
+
+
+class TracciaCallbackHandler(BaseCallbackHandler):
+    """
+    LangChain callback handler that creates Traccia spans for LLM and chain runs.
+    
+    This handler integrates LangChain with Traccia's tracing system, creating spans
+    for LLM calls with the same attributes used by Traccia's OpenAI instrumentation.
+    
+    Usage:
+        ```python
+        from traccia import init
+        from traccia.integrations.langchain import CallbackHandler  # or TracciaCallbackHandler
+        from langchain_openai import ChatOpenAI
+        
+        # Initialize Traccia
+        init()
+        
+        # Create Traccia handler (no args)
+        traccia_handler = CallbackHandler()
+        
+        # Use with any LangChain runnable
+        llm = ChatOpenAI(model="gpt-4o-mini")
+        result = llm.invoke(
+            "Tell me a joke",
+            config={"callbacks": [traccia_handler]}
+        )
+        ```
+    
+    Note:
+        Requires langchain-core to be installed:
+        ```bash
+        pip install traccia[langchain]
+        ```
+    """
+    
+    def __init__(self):
+        """Initialize the callback handler."""
+        super().__init__()
+        self.tracer = get_tracer("traccia.langchain")
+        
+        # Track active spans by run_id
+        self._spans: Dict[UUID, Any] = {}
+        self._context_tokens: Dict[UUID, Any] = {}
+        
+        # Track parent relationships
+        self._parent_map: Dict[UUID, Optional[UUID]] = {}
+    
+    def on_llm_start(
+        self,
+        serialized: Dict[str, Any],
+        prompts: List[str],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Handle LLM start event."""
+        self._parent_map[run_id] = parent_run_id
+        
+        try:
+            # Extract attributes
+            attributes = self._build_llm_attributes(
+                serialized, prompts, None, kwargs, metadata
+            )
+            
+            # Start span
+            span = self.tracer.start_as_current_span(
+                "llm.langchain.run",
+                attributes=attributes
+            )
+            
+            # Store span
+            self._spans[run_id] = span
+            self._context_tokens[run_id] = span
+            
+        except Exception as e:
+            # Don't break LangChain execution
+            import logging
+            logging.getLogger(__name__).exception(f"Error in on_llm_start: {e}")
+    
+    def on_chat_model_start(
+        self,
+        serialized: Dict[str, Any],
+        messages: List[List[BaseMessage]],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Handle chat model start event."""
+        self._parent_map[run_id] = parent_run_id
+        
+        try:
+            # Convert messages to prompt format
+            message_dicts = []
+            for msg_list in messages:
+                for msg in msg_list:
+                    message_dicts.append(self._convert_message_to_dict(msg))
+            
+            # Extract attributes
+            attributes = self._build_llm_attributes(
+                serialized, None, message_dicts, kwargs, metadata
+            )
+            
+            # Start span
+            span = self.tracer.start_as_current_span(
+                "llm.langchain.run",
+                attributes=attributes
+            )
+            
+            # Store span
+            self._spans[run_id] = span
+            self._context_tokens[run_id] = span
+            
+        except Exception as e:
+            import logging
+            logging.getLogger(__name__).exception(f"Error in on_chat_model_start: {e}")
+    
+    def on_llm_end(
+        self,
+        response: LLMResult,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Handle LLM end event."""
+        try:
+            span = self._spans.pop(run_id, None)
+            if span is None:
+                return
+            
+            # Extract usage and output
+            self._set_llm_response_attributes(span, response)
+            
+            # End span
+            span.__exit__(None, None, None)
+            
+            # Clean up context
+            self._context_tokens.pop(run_id, None)
+            self._parent_map.pop(run_id, None)
+            
+        except Exception as e:
+            import logging
+            logging.getLogger(__name__).exception(f"Error in on_llm_end: {e}")
+    
+    def on_llm_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Handle LLM error event."""
+        try:
+            span = self._spans.pop(run_id, None)
+            if span is None:
+                return
+            
+            # Record exception
+            span._otel_span.record_exception(error)
+            span.set_status(SpanStatus.ERROR, str(error))
+            
+            # End span
+            span.__exit__(type(error), error, None)
+            
+            # Clean up
+            self._context_tokens.pop(run_id, None)
+            self._parent_map.pop(run_id, None)
+            
+        except Exception as e:
+            import logging
+            logging.getLogger(__name__).exception(f"Error in on_llm_error: {e}")
+    
+    def on_chain_start(
+        self,
+        serialized: Dict[str, Any],
+        inputs: Dict[str, Any],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Handle chain start event (optional Phase 2)."""
+        self._parent_map[run_id] = parent_run_id
+        # Phase 2: Can add chain spans here
+    
+    def on_chain_end(
+        self,
+        outputs: Dict[str, Any],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Handle chain end event (optional Phase 2)."""
+        self._parent_map.pop(run_id, None)
+        # Phase 2: Can end chain spans here
+    
+    def on_chain_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Handle chain error event (optional Phase 2)."""
+        self._parent_map.pop(run_id, None)
+        # Phase 2: Can handle chain errors here
+    
+    def _build_llm_attributes(
+        self,
+        serialized: Dict[str, Any],
+        prompts: Optional[List[str]],
+        messages: Optional[List[Dict[str, Any]]],
+        kwargs: Dict[str, Any],
+        metadata: Optional[Dict[str, Any]],
+    ) -> Dict[str, Any]:
+        """Build LLM span attributes."""
+        from traccia.integrations.langchain.utils import extract_model_name
+        
+        attributes: Dict[str, Any] = {}
+        
+        # Extract vendor
+        vendor = self._extract_vendor(serialized)
+        if vendor:
+            attributes["llm.vendor"] = vendor
+        
+        # Extract model
+        model = extract_model_name(serialized, kwargs, metadata)
+        if model:
+            attributes["llm.model"] = model
+        
+        # Set prompt
+        if messages:
+            # Chat messages
+            prompt_text = self._format_messages_as_prompt(messages)
+            if prompt_text:
+                attributes["llm.prompt"] = prompt_text
+            
+            # Store messages as JSON (truncated)
+            try:
+                messages_json = json.dumps(messages)[:1000]
+                attributes["llm.openai.messages"] = messages_json
+            except Exception:
+                pass
+        elif prompts:
+            # Text prompts
+            if len(prompts) == 1:
+                attributes["llm.prompt"] = prompts[0]
+            else:
+                attributes["llm.prompt"] = json.dumps(prompts)[:1000]
+        
+        return attributes
+    
+    def _set_llm_response_attributes(
+        self,
+        span: Any,
+        response: LLMResult,
+    ) -> None:
+        """Set response attributes on span."""
+        usage = self._parse_usage(response)
+        if usage:
+            span.set_attribute("llm.usage.source", "provider_usage")
+            prompt_tokens = usage.get("prompt_tokens") or usage.get("input_tokens")
+            if prompt_tokens is not None:
+                span.set_attribute("llm.usage.prompt_tokens", int(prompt_tokens))
+                span.set_attribute("llm.usage.prompt_source", "provider_usage")
+            completion_tokens = usage.get("completion_tokens") or usage.get("output_tokens")
+            if completion_tokens is not None:
+                span.set_attribute("llm.usage.completion_tokens", int(completion_tokens))
+                span.set_attribute("llm.usage.completion_source", "provider_usage")
+            total = usage.get("total_tokens")
+            if total is not None:
+                span.set_attribute("llm.usage.total_tokens", int(total))
+
+        # Extract completion (last generation text or message content)
+        if response.generations and len(response.generations) > 0:
+            last_gen = response.generations[-1]
+            if last_gen and len(last_gen) > 0:
+                chunk = last_gen[-1]
+                completion = getattr(chunk, "text", None) or (
+                    getattr(getattr(chunk, "message", None), "content", None)
+                )
+                if completion:
+                    span.set_attribute("llm.completion", str(completion))
+    
+    def _extract_vendor(self, serialized: Dict[str, Any]) -> Optional[str]:
+        """Extract vendor from serialized LLM config."""
+        if not serialized or "id" not in serialized:
+            return None
+        
+        id_list = serialized["id"]
+        if not isinstance(id_list, list) or len(id_list) == 0:
+            return None
+        
+        # Get last component (class name)
+        class_name = id_list[-1].lower()
+        
+        # Map to vendor
+        if "openai" in class_name:
+            return "openai"
+        elif "anthropic" in class_name:
+            return "anthropic"
+        elif "cohere" in class_name:
+            return "cohere"
+        elif "huggingface" in class_name:
+            return "huggingface"
+        elif "vertexai" in class_name or "vertex" in class_name:
+            return "google"
+        elif "bedrock" in class_name:
+            return "aws"
+        
+        return "langchain"
+    
+    def _convert_message_to_dict(self, message: BaseMessage) -> Dict[str, Any]:
+        """Convert LangChain message to dict."""
+        return {
+            "role": getattr(message, "type", "unknown"),
+            "content": str(message.content) if hasattr(message, "content") else str(message),
+        }
+    
+    def _format_messages_as_prompt(self, messages: List[Dict[str, Any]]) -> str:
+        """Format messages as a prompt string."""
+        parts = []
+        for msg in messages:
+            role = msg.get("role", "unknown")
+            content = msg.get("content", "")
+            if content:
+                parts.append(f"{role}: {content}")
+        return "\n".join(parts) if parts else ""
+
+    def _parse_usage(self, response: LLMResult) -> Optional[Dict[str, Any]]:
+        """
+        Extract token usage from LLMResult.
+        Checks llm_output['token_usage'], llm_output['usage'], and generation_info.
+        """
+        usage = None
+        if response.llm_output:
+            for key in ("token_usage", "usage"):
+                if key in response.llm_output and response.llm_output[key]:
+                    raw = response.llm_output[key]
+                    usage = self._normalize_usage(raw)
+                    if usage:
+                        break
+        if not usage and response.generations:
+            for gen_list in response.generations:
+                for chunk in gen_list:
+                    if getattr(chunk, "generation_info", None) and isinstance(
+                        chunk.generation_info, dict
+                    ):
+                        raw = chunk.generation_info.get("usage_metadata")
+                        if raw:
+                            usage = self._normalize_usage(raw)
+                            break
+                    msg = getattr(chunk, "message", None)
+                    if msg is not None:
+                        meta = getattr(msg, "response_metadata", None) or {}
+                        raw = meta.get("usage") if isinstance(meta, dict) else None
+                        if raw:
+                            usage = self._normalize_usage(raw)
+                            break
+                if usage:
+                    break
+        return usage
+
+    @staticmethod
+    def _normalize_usage(raw: Any) -> Optional[Dict[str, Any]]:
+        """Normalize usage dict to prompt_tokens, completion_tokens, total_tokens."""
+        if raw is None:
+            return None
+        if hasattr(raw, "__dict__"):
+            raw = getattr(raw, "__dict__", raw)
+        if not isinstance(raw, dict):
+            return None
+        # Map common keys to Traccia/OpenAI style
+        prompt = raw.get("prompt_tokens") or raw.get("input_tokens") or raw.get("input")
+        completion = raw.get("completion_tokens") or raw.get("output_tokens") or raw.get("output")
+        total = raw.get("total_tokens") or raw.get("total")
+        if prompt is None and completion is None and total is None:
+            return None
+        out: Dict[str, Any] = {}
+        if prompt is not None:
+            out["prompt_tokens"] = int(prompt) if not isinstance(prompt, list) else sum(prompt)
+        if completion is not None:
+            out["completion_tokens"] = int(completion) if not isinstance(completion, list) else sum(completion)
+        if total is not None:
+            out["total_tokens"] = int(total) if not isinstance(total, list) else sum(total)
+        return out if out else None

traccia/integrations/langchain/utils.py

@@ -0,0 +1,129 @@
+"""Utility functions for LangChain integration."""
+
+from __future__ import annotations
+
+import re
+from typing import Any, Dict, List, Literal, Optional, cast
+
+
+def extract_model_name(
+    serialized: Optional[Dict[str, Any]],
+    kwargs: Dict[str, Any],
+    metadata: Optional[Dict[str, Any]] = None,
+) -> Optional[str]:
+    """
+    Extract model name from LangChain serialized config, invocation params, or metadata.
+    
+    Args:
+        serialized: LangChain's serialized component dict
+        kwargs: Keyword arguments from the callback (contains invocation_params)
+        metadata: Optional metadata dict
+        
+    Returns:
+        Model name string or None
+    """
+    # Check metadata first
+    if metadata:
+        model_from_meta = _parse_model_name_from_metadata(metadata)
+        if model_from_meta:
+            return model_from_meta
+    
+    # Try known model paths by ID
+    models_by_id = [
+        ("ChatOpenAI", ["invocation_params", "model_name"], "kwargs"),
+        ("ChatOpenAI", ["invocation_params", "model"], "kwargs"),
+        ("OpenAI", ["invocation_params", "model_name"], "kwargs"),
+        ("AzureChatOpenAI", ["invocation_params", "model"], "kwargs"),
+        ("AzureChatOpenAI", ["invocation_params", "model_name"], "kwargs"),
+        ("AzureChatOpenAI", ["invocation_params", "azure_deployment"], "kwargs"),
+        ("ChatAnthropic", ["invocation_params", "model"], "kwargs"),
+        ("ChatAnthropic", ["invocation_params", "model_name"], "kwargs"),
+        ("ChatGoogleGenerativeAI", ["kwargs", "model"], "serialized"),
+        ("ChatVertexAI", ["kwargs", "model_name"], "serialized"),
+        ("BedrockChat", ["kwargs", "model_id"], "serialized"),
+        ("ChatBedrock", ["kwargs", "model_id"], "serialized"),
+    ]
+    
+    for model_name, keys, select_from in models_by_id:
+        model = _extract_model_by_path_for_id(
+            model_name,
+            serialized,
+            kwargs,
+            keys,
+            cast(Literal["serialized", "kwargs"], select_from),
+        )
+        if model:
+            return model
+    
+    # Try common paths as catch-all
+    common_paths = [
+        ["invocation_params", "model_name"],
+        ["invocation_params", "model"],
+        ["kwargs", "model_name"],
+        ["kwargs", "model"],
+    ]
+    
+    for select in ["kwargs", "serialized"]:
+        for path in common_paths:
+            model = _extract_model_by_path(
+                serialized, kwargs, path, cast(Literal["serialized", "kwargs"], select)
+            )
+            if model:
+                return str(model)
+    
+    return None
+
+
+def _parse_model_name_from_metadata(metadata: Dict[str, Any]) -> Optional[str]:
+    """Extract model name from metadata if present."""
+    if not isinstance(metadata, dict):
+        return None
+    return metadata.get("ls_model_name", None)
+
+
+def _extract_model_by_path_for_id(
+    id: str,
+    serialized: Optional[Dict[str, Any]],
+    kwargs: Dict[str, Any],
+    keys: List[str],
+    select_from: Literal["serialized", "kwargs"],
+) -> Optional[str]:
+    """Extract model if the serialized ID matches."""
+    if serialized is None and select_from == "serialized":
+        return None
+    
+    if serialized:
+        serialized_id = serialized.get("id")
+        if (
+            serialized_id
+            and isinstance(serialized_id, list)
+            and len(serialized_id) > 0
+            and serialized_id[-1] == id
+        ):
+            result = _extract_model_by_path(serialized, kwargs, keys, select_from)
+            return str(result) if result is not None else None
+    
+    return None
+
+
+def _extract_model_by_path(
+    serialized: Optional[Dict[str, Any]],
+    kwargs: dict,
+    keys: List[str],
+    select_from: Literal["serialized", "kwargs"],
+) -> Optional[str]:
+    """Extract value by following a path in the dict."""
+    if serialized is None and select_from == "serialized":
+        return None
+    
+    current_obj = kwargs if select_from == "kwargs" else serialized
+    
+    for key in keys:
+        if current_obj and isinstance(current_obj, dict):
+            current_obj = current_obj.get(key)
+        else:
+            return None
+        if not current_obj:
+            return None
+    
+    return str(current_obj) if current_obj else None

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",
+]