headroom-ai 0.2.13__py3-none-any.whl

headroom/integrations/langchain/langsmith.py

@@ -0,0 +1,324 @@
+"""LangSmith integration for Headroom compression metrics.
+
+This module provides HeadroomLangSmithCallbackHandler, a LangChain callback
+handler that adds Headroom compression metrics to LangSmith traces.
+
+When used with HeadroomChatModel, it automatically captures:
+- Tokens before/after optimization
+- Savings percentage
+- Transforms applied
+- Per-request compression details
+
+Example:
+    import os
+    from langchain_openai import ChatOpenAI
+    from headroom.integrations import (
+        HeadroomChatModel,
+        HeadroomLangSmithCallbackHandler,
+    )
+
+    # Enable LangSmith tracing
+    os.environ["LANGCHAIN_TRACING_V2"] = "true"
+    os.environ["LANGCHAIN_API_KEY"] = "..."
+
+    # Create handler
+    handler = HeadroomLangSmithCallbackHandler()
+
+    # Use with HeadroomChatModel
+    llm = HeadroomChatModel(
+        ChatOpenAI(model="gpt-4o"),
+        callbacks=[handler],
+    )
+
+    # Traces will include headroom.* metadata
+    response = llm.invoke("Hello!")
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any
+from uuid import UUID
+
+# LangChain imports - these are optional dependencies
+try:
+    from langchain_core.callbacks import BaseCallbackHandler
+    from langchain_core.messages import BaseMessage
+    from langchain_core.outputs import LLMResult
+
+    LANGCHAIN_AVAILABLE = True
+except ImportError:
+    LANGCHAIN_AVAILABLE = False
+    BaseCallbackHandler = object  # type: ignore[misc,assignment]
+    LLMResult = object  # type: ignore[misc,assignment]
+
+# LangSmith imports - optional
+try:
+    from langsmith import Client as LangSmithClient
+
+    LANGSMITH_AVAILABLE = True
+except ImportError:
+    LANGSMITH_AVAILABLE = False
+    LangSmithClient = None  # type: ignore[misc,assignment]
+
+logger = logging.getLogger(__name__)
+
+
+def _check_langchain_available() -> None:
+    """Raise ImportError if LangChain is not installed."""
+    if not LANGCHAIN_AVAILABLE:
+        raise ImportError(
+            "LangChain is required for this integration. "
+            "Install with: pip install headroom[langchain] "
+            "or: pip install langchain-core"
+        )
+
+
+@dataclass
+class PendingMetrics:
+    """Metrics pending attachment to a LangSmith run."""
+
+    tokens_before: int
+    tokens_after: int
+    tokens_saved: int
+    savings_percent: float
+    transforms_applied: list[str]
+    timestamp: datetime = field(default_factory=datetime.now)
+
+
+class HeadroomLangSmithCallbackHandler(BaseCallbackHandler):
+    """Callback handler that adds Headroom metrics to LangSmith traces.
+
+    Integrates with LangSmith to provide visibility into context
+    optimization within traces. Metrics appear as metadata with
+    the `headroom.` prefix.
+
+    Works automatically when:
+    1. LANGCHAIN_TRACING_V2=true is set
+    2. Used as a callback with HeadroomChatModel
+    3. LangSmith API key is configured
+
+    Example:
+        from headroom.integrations import (
+            HeadroomChatModel,
+            HeadroomLangSmithCallbackHandler,
+        )
+
+        handler = HeadroomLangSmithCallbackHandler()
+        llm = HeadroomChatModel(
+            ChatOpenAI(model="gpt-4o"),
+            callbacks=[handler],
+        )
+
+        response = llm.invoke("Hello!")
+        # LangSmith trace now includes:
+        # - headroom.tokens_before
+        # - headroom.tokens_after
+        # - headroom.tokens_saved
+        # - headroom.savings_percent
+        # - headroom.transforms_applied
+
+    Attributes:
+        langsmith_client: LangSmith client for updating runs.
+        pending_metrics: Metrics waiting to be attached to runs.
+    """
+
+    def __init__(
+        self,
+        langsmith_client: Any = None,
+        auto_update_runs: bool = True,
+    ):
+        """Initialize HeadroomLangSmithCallbackHandler.
+
+        Args:
+            langsmith_client: LangSmith client instance. Auto-creates
+                one if not provided and LangSmith is available.
+            auto_update_runs: If True, automatically updates LangSmith
+                runs with Headroom metadata. Default True.
+        """
+        _check_langchain_available()
+
+        self._client = langsmith_client
+        self._auto_update = auto_update_runs
+        self._pending_metrics: dict[str, PendingMetrics] = {}
+        self._run_metrics: dict[str, dict[str, Any]] = {}
+
+        # Initialize LangSmith client if available and not provided
+        if self._client is None and LANGSMITH_AVAILABLE and auto_update_runs:
+            try:
+                if os.environ.get("LANGCHAIN_API_KEY"):
+                    self._client = LangSmithClient()
+            except Exception as e:
+                logger.debug(f"Could not initialize LangSmith client: {e}")
+
+    def set_headroom_metrics(
+        self,
+        run_id: str | UUID,
+        tokens_before: int,
+        tokens_after: int,
+        transforms_applied: list[str] | None = None,
+    ) -> None:
+        """Set Headroom metrics for a run.
+
+        Call this from HeadroomChatModel after optimization to attach
+        metrics to the current run.
+
+        Args:
+            run_id: The LangSmith run ID.
+            tokens_before: Token count before optimization.
+            tokens_after: Token count after optimization.
+            transforms_applied: List of transforms that were applied.
+        """
+        run_id_str = str(run_id)
+        tokens_saved = tokens_before - tokens_after
+        savings_percent = (tokens_saved / tokens_before * 100) if tokens_before > 0 else 0.0
+
+        metrics = PendingMetrics(
+            tokens_before=tokens_before,
+            tokens_after=tokens_after,
+            tokens_saved=tokens_saved,
+            savings_percent=savings_percent,
+            transforms_applied=transforms_applied or [],
+        )
+
+        self._pending_metrics[run_id_str] = metrics
+
+        logger.debug(
+            f"Headroom metrics set for run {run_id_str}: "
+            f"{tokens_before} -> {tokens_after} tokens ({savings_percent:.1f}% saved)"
+        )
+
+    def on_chat_model_start(
+        self,
+        serialized: dict[str, Any],
+        messages: list[list[BaseMessage]],
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> None:
+        """Called when chat model starts.
+
+        Records the run ID for later metric attachment.
+        """
+        run_id_str = str(run_id)
+        # Initialize empty metrics for this run
+        self._run_metrics[run_id_str] = {}
+
+    def on_llm_end(
+        self,
+        response: LLMResult,
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> None:
+        """Called when LLM completes.
+
+        Attaches pending Headroom metrics to the LangSmith run.
+        """
+        run_id_str = str(run_id)
+
+        # Check for pending metrics
+        if run_id_str in self._pending_metrics:
+            metrics = self._pending_metrics.pop(run_id_str)
+            self._attach_metrics_to_run(run_id_str, metrics)
+
+    def _attach_metrics_to_run(self, run_id: str, metrics: PendingMetrics) -> None:
+        """Attach Headroom metrics to a LangSmith run.
+
+        Args:
+            run_id: The run ID.
+            metrics: Metrics to attach.
+        """
+        metadata = {
+            "headroom.tokens_before": metrics.tokens_before,
+            "headroom.tokens_after": metrics.tokens_after,
+            "headroom.tokens_saved": metrics.tokens_saved,
+            "headroom.savings_percent": round(metrics.savings_percent, 2),
+            "headroom.transforms_applied": metrics.transforms_applied,
+            "headroom.optimization_timestamp": metrics.timestamp.isoformat(),
+        }
+
+        # Store in run metrics
+        self._run_metrics[run_id] = metadata
+
+        # Update LangSmith run if client available
+        if self._client and self._auto_update:
+            try:
+                self._client.update_run(
+                    run_id=run_id,
+                    extra={"metadata": metadata},
+                )
+                logger.debug(f"Updated LangSmith run {run_id} with Headroom metrics")
+            except Exception as e:
+                logger.debug(f"Could not update LangSmith run: {e}")
+
+    def get_run_metrics(self, run_id: str | UUID) -> dict[str, Any]:
+        """Get Headroom metrics for a specific run.
+
+        Args:
+            run_id: The run ID.
+
+        Returns:
+            Dictionary of headroom.* metrics for the run.
+        """
+        return self._run_metrics.get(str(run_id), {})
+
+    def get_all_metrics(self) -> dict[str, dict[str, Any]]:
+        """Get all recorded run metrics.
+
+        Returns:
+            Dictionary mapping run IDs to their metrics.
+        """
+        return self._run_metrics.copy()
+
+    def get_summary(self) -> dict[str, Any]:
+        """Get summary statistics across all runs.
+
+        Returns:
+            Summary with total runs, tokens saved, etc.
+        """
+        if not self._run_metrics:
+            return {
+                "total_runs": 0,
+                "total_tokens_saved": 0,
+                "average_savings_percent": 0,
+            }
+
+        total_saved = sum(m.get("headroom.tokens_saved", 0) for m in self._run_metrics.values())
+        savings_percents = [
+            m.get("headroom.savings_percent", 0) for m in self._run_metrics.values()
+        ]
+
+        return {
+            "total_runs": len(self._run_metrics),
+            "total_tokens_saved": total_saved,
+            "average_savings_percent": (
+                sum(savings_percents) / len(savings_percents) if savings_percents else 0
+            ),
+        }
+
+    def reset(self) -> None:
+        """Clear all recorded metrics."""
+        self._pending_metrics.clear()
+        self._run_metrics.clear()
+
+
+def is_langsmith_available() -> bool:
+    """Check if LangSmith is available and configured.
+
+    Returns:
+        True if LangSmith is installed and API key is set.
+    """
+    return LANGSMITH_AVAILABLE and bool(os.environ.get("LANGCHAIN_API_KEY"))
+
+
+def is_langsmith_tracing_enabled() -> bool:
+    """Check if LangSmith tracing is enabled.
+
+    Returns:
+        True if LANGCHAIN_TRACING_V2 is set to "true".
+    """
+    return os.environ.get("LANGCHAIN_TRACING_V2", "").lower() == "true"

headroom/integrations/langchain/memory.py

@@ -0,0 +1,319 @@
+"""Memory integration for LangChain with automatic compression.
+
+This module provides HeadroomChatMessageHistory, a wrapper for any LangChain
+chat message history that automatically compresses conversation history
+when it exceeds a token threshold.
+
+Example:
+    from langchain.memory import ConversationBufferMemory
+    from langchain_community.chat_message_histories import ChatMessageHistory
+    from headroom.integrations import HeadroomChatMessageHistory
+
+    # Wrap any chat message history
+    base_history = ChatMessageHistory()
+    compressed_history = HeadroomChatMessageHistory(base_history)
+
+    # Use with ConversationBufferMemory (zero code changes to chain)
+    memory = ConversationBufferMemory(chat_memory=compressed_history)
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from headroom.providers.base import Provider
+
+# LangChain imports - these are optional dependencies
+try:
+    from langchain_core.chat_history import BaseChatMessageHistory
+    from langchain_core.messages import (
+        AIMessage,
+        BaseMessage,
+        HumanMessage,
+        SystemMessage,
+        ToolMessage,
+    )
+
+    LANGCHAIN_AVAILABLE = True
+except ImportError:
+    LANGCHAIN_AVAILABLE = False
+    BaseChatMessageHistory = object  # type: ignore[misc,assignment]
+
+from headroom import HeadroomConfig
+from headroom.config import RollingWindowConfig
+from headroom.providers import OpenAIProvider
+from headroom.transforms import TransformPipeline
+
+logger = logging.getLogger(__name__)
+
+
+def _check_langchain_available() -> None:
+    """Raise ImportError if LangChain is not installed."""
+    if not LANGCHAIN_AVAILABLE:
+        raise ImportError(
+            "LangChain is required for this integration. "
+            "Install with: pip install headroom[langchain] "
+            "or: pip install langchain-core"
+        )
+
+
+class HeadroomChatMessageHistory(BaseChatMessageHistory):
+    """Wraps any LangChain chat message history with automatic compression.
+
+    When conversation history exceeds the token threshold, automatically
+    applies RollingWindow compression to keep recent turns while fitting
+    within the limit.
+
+    This works with ANY memory type because it wraps at the storage layer:
+    - ConversationBufferMemory
+    - ConversationSummaryMemory
+    - ConversationBufferWindowMemory
+    - Redis, PostgreSQL, or any custom history
+
+    Example:
+        from langchain.memory import ConversationBufferMemory
+        from langchain_community.chat_message_histories import ChatMessageHistory
+        from headroom.integrations import HeadroomChatMessageHistory
+
+        # Wrap base history
+        base = ChatMessageHistory()
+        compressed = HeadroomChatMessageHistory(
+            base,
+            compress_threshold_tokens=4000,
+            keep_recent_turns=5,
+        )
+
+        # Use with any memory class
+        memory = ConversationBufferMemory(chat_memory=compressed)
+
+        # Messages are compressed automatically when accessed
+        chain = ConversationChain(llm=llm, memory=memory)
+        chain.invoke({"input": "Hello!"})
+
+    Attributes:
+        base_history: The underlying chat message history
+        compress_threshold_tokens: Token count that triggers compression
+        keep_recent_turns: Minimum recent turns to always preserve
+        model: Model name for token counting (default: "gpt-4o")
+    """
+
+    def __init__(
+        self,
+        base_history: BaseChatMessageHistory,
+        compress_threshold_tokens: int = 4000,
+        keep_recent_turns: int = 5,
+        model: str = "gpt-4o",
+        provider: Provider | None = None,
+    ):
+        """Initialize HeadroomChatMessageHistory.
+
+        Args:
+            base_history: Any LangChain BaseChatMessageHistory to wrap
+            compress_threshold_tokens: Apply compression when history exceeds
+                this many tokens. Default 4000.
+            keep_recent_turns: Minimum number of recent user/assistant turns
+                to always preserve during compression. Default 5.
+            model: Model name for token counting. Default "gpt-4o".
+            provider: Headroom provider for token counting. Auto-uses
+                OpenAIProvider if not specified.
+        """
+        _check_langchain_available()
+
+        self._base = base_history
+        self._threshold = compress_threshold_tokens
+        self._keep_recent_turns = keep_recent_turns
+        self._model = model
+        self._provider: Provider = provider or OpenAIProvider()
+
+        # Track compression stats
+        self._compression_count = 0
+        self._total_tokens_saved = 0
+
+    @property
+    def messages(self) -> list[BaseMessage]:  # type: ignore[override]
+        """Get messages, applying compression if over threshold.
+
+        Returns:
+            List of messages, potentially compressed to fit within threshold.
+        """
+        raw_messages = self._base.messages
+
+        if not raw_messages:
+            return []
+
+        # Count tokens
+        token_count = self._count_tokens(raw_messages)
+
+        if token_count <= self._threshold:
+            return list(raw_messages)
+
+        # Apply compression
+        compressed = self._apply_rolling_window(raw_messages)
+        tokens_after = self._count_tokens(compressed)
+
+        self._compression_count += 1
+        self._total_tokens_saved += token_count - tokens_after
+
+        logger.info(
+            f"HeadroomChatMessageHistory compressed: {token_count} -> {tokens_after} tokens "
+            f"({len(raw_messages)} -> {len(compressed)} messages)"
+        )
+
+        return compressed
+
+    def add_message(self, message: BaseMessage) -> None:
+        """Add a message to the underlying history.
+
+        Args:
+            message: The message to add.
+        """
+        self._base.add_message(message)
+
+    def add_user_message(self, message: HumanMessage | str) -> None:
+        """Add a user message to the history.
+
+        Args:
+            message: The user message (string or HumanMessage).
+        """
+        self._base.add_user_message(message)
+
+    def add_ai_message(self, message: AIMessage | str) -> None:
+        """Add an AI message to the history.
+
+        Args:
+            message: The AI message (string or AIMessage).
+        """
+        self._base.add_ai_message(message)
+
+    def clear(self) -> None:
+        """Clear all messages from history."""
+        self._base.clear()
+
+    def _count_tokens(self, messages: list[BaseMessage]) -> int:
+        """Count tokens in messages using provider's tokenizer.
+
+        Args:
+            messages: List of messages to count.
+
+        Returns:
+            Total token count.
+        """
+        token_counter = self._provider.get_token_counter(self._model)
+        total = 0
+        for msg in messages:
+            content = msg.content if isinstance(msg.content, str) else str(msg.content)
+            total += token_counter.count_text(content)
+        return total
+
+    def _apply_rolling_window(self, messages: list[BaseMessage]) -> list[BaseMessage]:
+        """Apply RollingWindow compression to messages.
+
+        Args:
+            messages: Messages to compress.
+
+        Returns:
+            Compressed messages fitting within threshold.
+        """
+        # Convert to OpenAI format for Headroom transforms
+        openai_messages = self._convert_to_openai(messages)
+
+        # Use TransformPipeline which handles tokenizer setup
+        config = HeadroomConfig(
+            rolling_window=RollingWindowConfig(keep_last_turns=self._keep_recent_turns),
+        )
+        pipeline = TransformPipeline(config=config, provider=self._provider)
+
+        # Apply compression via pipeline
+        result = pipeline.apply(
+            messages=openai_messages,
+            model=self._model,
+            model_limit=self._threshold,
+        )
+
+        # Convert back to LangChain format
+        return self._convert_from_openai(result.messages)
+
+    def _convert_to_openai(self, messages: list[BaseMessage]) -> list[dict[str, Any]]:
+        """Convert LangChain messages to OpenAI format.
+
+        Args:
+            messages: LangChain messages.
+
+        Returns:
+            OpenAI format messages.
+        """
+        result = []
+        for msg in messages:
+            content = msg.content if isinstance(msg.content, str) else str(msg.content)
+
+            if isinstance(msg, SystemMessage):
+                result.append({"role": "system", "content": content})
+            elif isinstance(msg, HumanMessage):
+                result.append({"role": "user", "content": content})
+            elif isinstance(msg, AIMessage):
+                entry: dict[str, Any] = {"role": "assistant", "content": content}
+                if hasattr(msg, "tool_calls") and msg.tool_calls:
+                    entry["tool_calls"] = msg.tool_calls
+                result.append(entry)
+            elif isinstance(msg, ToolMessage):
+                result.append(
+                    {
+                        "role": "tool",
+                        "tool_call_id": getattr(msg, "tool_call_id", ""),
+                        "content": content,
+                    }
+                )
+            else:
+                # Generic fallback
+                result.append(
+                    {
+                        "role": getattr(msg, "type", "user"),
+                        "content": content,
+                    }
+                )
+        return result
+
+    def _convert_from_openai(self, messages: list[dict[str, Any]]) -> list[BaseMessage]:
+        """Convert OpenAI format back to LangChain messages.
+
+        Args:
+            messages: OpenAI format messages.
+
+        Returns:
+            LangChain messages.
+        """
+        result: list[BaseMessage] = []
+        for msg in messages:
+            role = msg.get("role", "user")
+            content = msg.get("content", "")
+
+            if role == "system":
+                result.append(SystemMessage(content=content))
+            elif role == "user":
+                result.append(HumanMessage(content=content))
+            elif role == "assistant":
+                tool_calls = msg.get("tool_calls", [])
+                result.append(AIMessage(content=content, tool_calls=tool_calls))
+            elif role == "tool":
+                result.append(
+                    ToolMessage(
+                        content=content,
+                        tool_call_id=msg.get("tool_call_id", ""),
+                    )
+                )
+        return result
+
+    def get_compression_stats(self) -> dict[str, Any]:
+        """Get statistics about compression operations.
+
+        Returns:
+            Dictionary with compression_count, total_tokens_saved.
+        """
+        return {
+            "compression_count": self._compression_count,
+            "total_tokens_saved": self._total_tokens_saved,
+            "threshold_tokens": self._threshold,
+            "keep_recent_turns": self._keep_recent_turns,
+        }