aitracer 0.1.0__py3-none-any.whl

python/integrations/langchain.py

@@ -0,0 +1,452 @@
+"""
+LangChain integration for AITracer.
+
+This module provides a callback handler that automatically logs
+LangChain LLM calls to AITracer.
+
+Usage:
+    from aitracer.integrations.langchain import AITracerCallbackHandler
+    from langchain_openai import ChatOpenAI
+
+    handler = AITracerCallbackHandler(api_key="your-api-key", project="my-project")
+
+    llm = ChatOpenAI(callbacks=[handler])
+    response = llm.invoke("Hello!")
+"""
+
+from __future__ import annotations
+
+import time
+import uuid
+from decimal import Decimal
+from typing import Any, Dict, List, Optional, Union
+
+try:
+    from langchain_core.callbacks import BaseCallbackHandler
+    from langchain_core.outputs import LLMResult
+    from langchain_core.messages import BaseMessage
+    LANGCHAIN_AVAILABLE = True
+except ImportError:
+    LANGCHAIN_AVAILABLE = False
+    # Create a dummy base class for type hints
+    class BaseCallbackHandler:
+        pass
+
+
+class AITracerCallbackHandler(BaseCallbackHandler):
+    """
+    LangChain callback handler that logs LLM calls to AITracer.
+
+    This handler captures:
+    - LLM calls (start, end, error)
+    - Chain execution
+    - Tool usage
+    - Token usage and costs
+
+    Example:
+        >>> from aitracer.integrations.langchain import AITracerCallbackHandler
+        >>> from langchain_openai import ChatOpenAI
+        >>>
+        >>> handler = AITracerCallbackHandler(
+        ...     api_key="at-xxx",
+        ...     project="my-langchain-app"
+        ... )
+        >>> llm = ChatOpenAI(callbacks=[handler])
+        >>> response = llm.invoke("What is the capital of France?")
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        project: Optional[str] = None,
+        base_url: str = "https://api.aitracer.co",
+        session_id: Optional[str] = None,
+        user_id: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        flush_on_chain_end: bool = True,
+    ):
+        """
+        Initialize the AITracer callback handler.
+
+        Args:
+            api_key: AITracer API key. If not provided, uses AITRACER_API_KEY env var.
+            project: Project name for logging.
+            base_url: AITracer API base URL.
+            session_id: Optional session ID for grouping related calls.
+            user_id: Optional user ID for tracking user interactions.
+            metadata: Additional metadata to include with all logs.
+            flush_on_chain_end: Whether to flush logs when a chain completes.
+        """
+        if not LANGCHAIN_AVAILABLE:
+            raise ImportError(
+                "LangChain is not installed. Install it with: pip install langchain-core"
+            )
+
+        super().__init__()
+
+        # Import AITracer client
+        from aitracer import AITracer
+
+        self.client = AITracer(api_key=api_key, project=project, base_url=base_url)
+        self.project = project
+        self.session_id = session_id or str(uuid.uuid4())
+        self.user_id = user_id
+        self.default_metadata = metadata or {}
+        self.flush_on_chain_end = flush_on_chain_end
+
+        # Track active runs
+        self._runs: Dict[str, Dict[str, Any]] = {}
+        self._chain_stack: List[str] = []
+
+    def _get_trace_id(self) -> str:
+        """Get the current trace ID (from chain or generate new)."""
+        if self._chain_stack:
+            return self._chain_stack[0]
+        return str(uuid.uuid4())
+
+    def _get_parent_run_id(self) -> Optional[str]:
+        """Get the parent run ID if in a chain."""
+        if len(self._chain_stack) > 1:
+            return self._chain_stack[-1]
+        return None
+
+    def _extract_model_info(self, serialized: Dict[str, Any]) -> Dict[str, str]:
+        """Extract provider and model from serialized LLM info."""
+        # Default values
+        provider = "unknown"
+        model = "unknown"
+
+        # Try to extract from kwargs
+        kwargs = serialized.get("kwargs", {})
+
+        # Check for model name in various places
+        if "model" in kwargs:
+            model = kwargs["model"]
+        elif "model_name" in kwargs:
+            model = kwargs["model_name"]
+
+        # Determine provider from class name or model name
+        class_name = serialized.get("name", "").lower()
+        if "openai" in class_name:
+            provider = "openai"
+        elif "anthropic" in class_name or "claude" in class_name:
+            provider = "anthropic"
+        elif "google" in class_name or "gemini" in class_name:
+            provider = "google"
+        elif "mistral" in class_name:
+            provider = "mistral"
+        elif "cohere" in class_name:
+            provider = "cohere"
+
+        # Also check model name for provider hints
+        model_lower = model.lower()
+        if "gpt" in model_lower or model_lower.startswith("o1"):
+            provider = "openai"
+        elif "claude" in model_lower:
+            provider = "anthropic"
+        elif "gemini" in model_lower:
+            provider = "google"
+
+        return {"provider": provider, "model": model}
+
+    def _format_messages(self, messages: List[Any]) -> List[Dict[str, str]]:
+        """Format messages for logging."""
+        formatted = []
+        for msg in messages:
+            if isinstance(msg, dict):
+                formatted.append(msg)
+            elif hasattr(msg, "type") and hasattr(msg, "content"):
+                formatted.append({
+                    "role": msg.type,
+                    "content": msg.content,
+                })
+            elif hasattr(msg, "role") and hasattr(msg, "content"):
+                formatted.append({
+                    "role": msg.role,
+                    "content": msg.content,
+                })
+            else:
+                formatted.append({"role": "unknown", "content": str(msg)})
+        return formatted
+
+    # ========== LLM Callbacks ==========
+
+    def on_llm_start(
+        self,
+        serialized: Dict[str, Any],
+        prompts: List[str],
+        *,
+        run_id: uuid.UUID,
+        parent_run_id: Optional[uuid.UUID] = None,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when an LLM starts processing."""
+        run_id_str = str(run_id)
+        model_info = self._extract_model_info(serialized)
+
+        self._runs[run_id_str] = {
+            "type": "llm",
+            "start_time": time.time(),
+            "provider": model_info["provider"],
+            "model": model_info["model"],
+            "input": {"prompts": prompts},
+            "trace_id": self._get_trace_id(),
+            "parent_run_id": str(parent_run_id) if parent_run_id else self._get_parent_run_id(),
+            "tags": tags or [],
+            "metadata": {**self.default_metadata, **(metadata or {})},
+        }
+
+    def on_chat_model_start(
+        self,
+        serialized: Dict[str, Any],
+        messages: List[List[BaseMessage]],
+        *,
+        run_id: uuid.UUID,
+        parent_run_id: Optional[uuid.UUID] = None,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when a chat model starts processing."""
+        run_id_str = str(run_id)
+        model_info = self._extract_model_info(serialized)
+
+        # Format messages for logging
+        formatted_messages = []
+        for msg_list in messages:
+            formatted_messages.extend(self._format_messages(msg_list))
+
+        self._runs[run_id_str] = {
+            "type": "chat",
+            "start_time": time.time(),
+            "provider": model_info["provider"],
+            "model": model_info["model"],
+            "input": {"messages": formatted_messages},
+            "trace_id": self._get_trace_id(),
+            "parent_run_id": str(parent_run_id) if parent_run_id else self._get_parent_run_id(),
+            "tags": tags or [],
+            "metadata": {**self.default_metadata, **(metadata or {})},
+        }
+
+    def on_llm_end(
+        self,
+        response: LLMResult,
+        *,
+        run_id: uuid.UUID,
+        parent_run_id: Optional[uuid.UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when an LLM finishes processing."""
+        run_id_str = str(run_id)
+        run_data = self._runs.pop(run_id_str, None)
+
+        if not run_data:
+            return
+
+        latency_ms = int((time.time() - run_data["start_time"]) * 1000)
+
+        # Extract token usage
+        token_usage = {}
+        if response.llm_output:
+            token_usage = response.llm_output.get("token_usage", {})
+
+        # Extract output
+        output_text = ""
+        if response.generations:
+            for gen_list in response.generations:
+                for gen in gen_list:
+                    output_text += gen.text
+
+        # Log to AITracer
+        self.client.log(
+            trace_id=run_data["trace_id"],
+            span_id=run_id_str,
+            parent_span_id=run_data.get("parent_run_id"),
+            provider=run_data["provider"],
+            model=run_data["model"],
+            input=run_data["input"],
+            output={"content": output_text},
+            input_tokens=token_usage.get("prompt_tokens", 0),
+            output_tokens=token_usage.get("completion_tokens", 0),
+            latency_ms=latency_ms,
+            status="success",
+            metadata={
+                **run_data["metadata"],
+                "tags": run_data["tags"],
+                "session_id": self.session_id,
+                "user_id": self.user_id,
+                "source": "langchain",
+            },
+        )
+
+    def on_llm_error(
+        self,
+        error: Union[Exception, KeyboardInterrupt],
+        *,
+        run_id: uuid.UUID,
+        parent_run_id: Optional[uuid.UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when an LLM errors."""
+        run_id_str = str(run_id)
+        run_data = self._runs.pop(run_id_str, None)
+
+        if not run_data:
+            return
+
+        latency_ms = int((time.time() - run_data["start_time"]) * 1000)
+
+        # Log error to AITracer
+        self.client.log(
+            trace_id=run_data["trace_id"],
+            span_id=run_id_str,
+            parent_span_id=run_data.get("parent_run_id"),
+            provider=run_data["provider"],
+            model=run_data["model"],
+            input=run_data["input"],
+            output=None,
+            latency_ms=latency_ms,
+            status="error",
+            error_message=str(error),
+            metadata={
+                **run_data["metadata"],
+                "tags": run_data["tags"],
+                "session_id": self.session_id,
+                "user_id": self.user_id,
+                "source": "langchain",
+            },
+        )
+
+    # ========== Chain Callbacks ==========
+
+    def on_chain_start(
+        self,
+        serialized: Dict[str, Any],
+        inputs: Dict[str, Any],
+        *,
+        run_id: uuid.UUID,
+        parent_run_id: Optional[uuid.UUID] = None,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when a chain starts."""
+        run_id_str = str(run_id)
+        self._chain_stack.append(run_id_str)
+
+        self._runs[run_id_str] = {
+            "type": "chain",
+            "start_time": time.time(),
+            "name": serialized.get("name", "unknown"),
+            "inputs": inputs,
+            "tags": tags or [],
+            "metadata": {**self.default_metadata, **(metadata or {})},
+        }
+
+    def on_chain_end(
+        self,
+        outputs: Dict[str, Any],
+        *,
+        run_id: uuid.UUID,
+        parent_run_id: Optional[uuid.UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when a chain ends."""
+        run_id_str = str(run_id)
+
+        if run_id_str in self._chain_stack:
+            self._chain_stack.remove(run_id_str)
+
+        self._runs.pop(run_id_str, None)
+
+        # Flush logs if configured
+        if self.flush_on_chain_end and not self._chain_stack:
+            self.client.flush()
+
+    def on_chain_error(
+        self,
+        error: Union[Exception, KeyboardInterrupt],
+        *,
+        run_id: uuid.UUID,
+        parent_run_id: Optional[uuid.UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when a chain errors."""
+        run_id_str = str(run_id)
+
+        if run_id_str in self._chain_stack:
+            self._chain_stack.remove(run_id_str)
+
+        self._runs.pop(run_id_str, None)
+
+    # ========== Tool Callbacks ==========
+
+    def on_tool_start(
+        self,
+        serialized: Dict[str, Any],
+        input_str: str,
+        *,
+        run_id: uuid.UUID,
+        parent_run_id: Optional[uuid.UUID] = None,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when a tool starts."""
+        run_id_str = str(run_id)
+
+        self._runs[run_id_str] = {
+            "type": "tool",
+            "start_time": time.time(),
+            "name": serialized.get("name", "unknown"),
+            "input": input_str,
+            "tags": tags or [],
+            "metadata": {**self.default_metadata, **(metadata or {})},
+        }
+
+    def on_tool_end(
+        self,
+        output: str,
+        *,
+        run_id: uuid.UUID,
+        parent_run_id: Optional[uuid.UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when a tool ends."""
+        run_id_str = str(run_id)
+        run_data = self._runs.pop(run_id_str, None)
+
+        # Tools are tracked but not logged as separate entries
+        # They appear in the chain context
+
+    def on_tool_error(
+        self,
+        error: Union[Exception, KeyboardInterrupt],
+        *,
+        run_id: uuid.UUID,
+        parent_run_id: Optional[uuid.UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when a tool errors."""
+        run_id_str = str(run_id)
+        self._runs.pop(run_id_str, None)
+
+    # ========== Utility Methods ==========
+
+    def flush(self) -> None:
+        """Flush any pending logs to AITracer."""
+        self.client.flush()
+
+    def set_session_id(self, session_id: str) -> None:
+        """Set the session ID for subsequent calls."""
+        self.session_id = session_id
+
+    def set_user_id(self, user_id: str) -> None:
+        """Set the user ID for subsequent calls."""
+        self.user_id = user_id
+
+    def add_metadata(self, key: str, value: Any) -> None:
+        """Add metadata that will be included with all subsequent logs."""
+        self.default_metadata[key] = value

python/pii.py

@@ -0,0 +1,223 @@
+"""PII Detection and Masking utilities for SDK."""
+
+from __future__ import annotations
+
+import hashlib
+import re
+from dataclasses import dataclass
+from typing import Any, Optional
+
+
+@dataclass
+class PIIMatch:
+    """Represents a detected PII match."""
+
+    pii_type: str
+    value: str
+    start: int
+    end: int
+    confidence: float = 1.0
+
+
+# PII Detection Patterns
+PII_PATTERNS = {
+    "email": {
+        "pattern": r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b",
+        "mask": "[EMAIL]",
+        "confidence": 0.95,
+    },
+    "phone": {
+        # International and Japanese phone formats
+        "pattern": r"\b(?:\+?[0-9]{1,4}[-.\s]?)?(?:\(?[0-9]{2,4}\)?[-.\s]?)?[0-9]{2,4}[-.\s]?[0-9]{3,4}[-.\s]?[0-9]{3,4}\b",
+        "mask": "[PHONE]",
+        "confidence": 0.85,
+    },
+    "credit_card": {
+        # Major credit card formats (Visa, Mastercard, Amex, etc.)
+        "pattern": r"\b(?:4[0-9]{3}[-\s]?[0-9]{4}[-\s]?[0-9]{4}[-\s]?[0-9]{4}|5[1-5][0-9]{2}[-\s]?[0-9]{4}[-\s]?[0-9]{4}[-\s]?[0-9]{4}|3[47][0-9]{2}[-\s]?[0-9]{6}[-\s]?[0-9]{5}|6(?:011|5[0-9]{2})[-\s]?[0-9]{4}[-\s]?[0-9]{4}[-\s]?[0-9]{4})\b",
+        "mask": "[CREDIT_CARD]",
+        "confidence": 0.95,
+    },
+    "ssn": {
+        # US Social Security Number and Japanese My Number
+        "pattern": r"\b(?:[0-9]{3}[-\s]?[0-9]{2}[-\s]?[0-9]{4}|[0-9]{4}[-\s]?[0-9]{4}[-\s]?[0-9]{4})\b",
+        "mask": "[SSN]",
+        "confidence": 0.90,
+    },
+    "ip_address": {
+        # IPv4 and IPv6
+        "pattern": r"\b(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\b|(?:[A-Fa-f0-9]{1,4}:){7}[A-Fa-f0-9]{1,4}\b",
+        "mask": "[IP_ADDRESS]",
+        "confidence": 0.95,
+    },
+    "name": {
+        # Japanese names (simplified pattern)
+        "pattern": r"[一-龯]{1,4}[\s　][一-龯]{1,4}",
+        "mask": "[NAME]",
+        "confidence": 0.70,
+    },
+}
+
+
+class PIIDetector:
+    """PII Detection and Masking class."""
+
+    def __init__(
+        self,
+        enabled_types: Optional[list[str]] = None,
+        action: str = "mask",
+        custom_patterns: Optional[list[dict]] = None,
+    ):
+        """
+        Initialize PII Detector.
+
+        Args:
+            enabled_types: List of PII types to detect. None means all.
+            action: Action to take - "mask", "redact", "hash", "none"
+            custom_patterns: List of custom pattern dicts with 'name', 'pattern', 'mask' keys
+        """
+        self.enabled_types = enabled_types or list(PII_PATTERNS.keys())
+        self.action = action
+        self.custom_patterns = custom_patterns or []
+
+        # Compile patterns
+        self.compiled_patterns: dict[str, dict] = {}
+        for pii_type in self.enabled_types:
+            if pii_type in PII_PATTERNS:
+                pattern_info = PII_PATTERNS[pii_type]
+                self.compiled_patterns[pii_type] = {
+                    "regex": re.compile(pattern_info["pattern"], re.IGNORECASE),
+                    "mask": pattern_info["mask"],
+                    "confidence": pattern_info["confidence"],
+                }
+
+        # Add custom patterns
+        for custom in self.custom_patterns:
+            name = custom.get("name", f"custom_{len(self.compiled_patterns)}")
+            pattern = custom.get("pattern")
+            mask = custom.get("mask", f"[{name.upper()}]")
+            if pattern:
+                try:
+                    self.compiled_patterns[name] = {
+                        "regex": re.compile(pattern, re.IGNORECASE),
+                        "mask": mask,
+                        "confidence": custom.get("confidence", 0.8),
+                    }
+                except re.error:
+                    pass  # Skip invalid patterns
+
+    def detect(self, text: str) -> list[PIIMatch]:
+        """
+        Detect PII in text.
+
+        Args:
+            text: Text to scan for PII
+
+        Returns:
+            List of PIIMatch objects
+        """
+        if not text:
+            return []
+
+        matches: list[PIIMatch] = []
+        for pii_type, pattern_info in self.compiled_patterns.items():
+            for match in pattern_info["regex"].finditer(text):
+                matches.append(
+                    PIIMatch(
+                        pii_type=pii_type,
+                        value=match.group(),
+                        start=match.start(),
+                        end=match.end(),
+                        confidence=pattern_info["confidence"],
+                    )
+                )
+
+        # Sort by position and remove overlapping matches (keep higher confidence)
+        matches.sort(key=lambda m: (m.start, -m.confidence))
+        filtered_matches: list[PIIMatch] = []
+        last_end = -1
+        for match in matches:
+            if match.start >= last_end:
+                filtered_matches.append(match)
+                last_end = match.end
+
+        return filtered_matches
+
+    def mask_text(self, text: str) -> tuple[str, list[PIIMatch]]:
+        """
+        Detect and mask PII in text.
+
+        Args:
+            text: Text to process
+
+        Returns:
+            Tuple of (masked_text, list of matches)
+        """
+        matches = self.detect(text)
+        if not matches:
+            return text, []
+
+        # Apply masking based on action
+        result = text
+        offset = 0
+
+        for match in sorted(matches, key=lambda m: m.start):
+            start = match.start + offset
+            end = match.end + offset
+
+            if self.action == "mask":
+                replacement = self.compiled_patterns[match.pii_type]["mask"]
+            elif self.action == "redact":
+                replacement = "***"
+            elif self.action == "hash":
+                hash_value = hashlib.sha256(match.value.encode()).hexdigest()[:12]
+                replacement = f"[HASH:{hash_value}]"
+            else:  # none
+                continue
+
+            result = result[:start] + replacement + result[end:]
+            offset += len(replacement) - len(match.value)
+
+        return result, matches
+
+    def process_json(
+        self, data: Any, path: str = ""
+    ) -> tuple[Any, list[tuple[str, PIIMatch]]]:
+        """
+        Recursively process JSON data for PII.
+
+        Args:
+            data: JSON data (dict, list, or primitive)
+            path: Current path in the JSON structure
+
+        Returns:
+            Tuple of (processed_data, list of (path, match) tuples)
+        """
+        all_matches: list[tuple[str, PIIMatch]] = []
+
+        if isinstance(data, dict):
+            result = {}
+            for key, value in data.items():
+                new_path = f"{path}.{key}" if path else key
+                processed, matches = self.process_json(value, new_path)
+                result[key] = processed
+                all_matches.extend(matches)
+            return result, all_matches
+
+        elif isinstance(data, list):
+            result = []
+            for i, item in enumerate(data):
+                new_path = f"{path}[{i}]"
+                processed, matches = self.process_json(item, new_path)
+                result.append(processed)
+                all_matches.extend(matches)
+            return result, all_matches
+
+        elif isinstance(data, str):
+            masked_text, matches = self.mask_text(data)
+            for match in matches:
+                all_matches.append((path, match))
+            return masked_text, all_matches
+
+        else:
+            return data, []