control-zero 0.2.0__py3-none-any.whl

control_zero/integrations/langchain/callbacks.py

@@ -0,0 +1,441 @@
+"""
+Control Zero Callback Handler for LangChain.
+
+Provides comprehensive logging and monitoring for all LangChain operations
+including LLM calls, tool executions, chain runs, and agent actions.
+"""
+
+import time
+from typing import Any, Dict, List, Optional, Union
+from uuid import UUID
+
+try:
+    from langchain_core.callbacks import BaseCallbackHandler
+    from langchain_core.outputs import LLMResult
+    from langchain_core.agents import AgentAction, AgentFinish
+    from langchain_core.messages import BaseMessage
+    LANGCHAIN_AVAILABLE = True
+except ImportError:
+    LANGCHAIN_AVAILABLE = False
+    class BaseCallbackHandler:
+        pass
+    class LLMResult:
+        pass
+    class AgentAction:
+        pass
+    class AgentFinish:
+        pass
+    class BaseMessage:
+        pass
+
+from control_zero.client import ControlZeroClient
+
+
+class ControlZeroCallbackHandler(BaseCallbackHandler):
+    """
+    Callback handler that logs all LangChain operations to Control Zero.
+
+    Captures:
+    - LLM calls (start, end, error)
+    - Tool executions (start, end, error)
+    - Chain runs (start, end, error)
+    - Agent actions
+    - Token usage and costs
+
+    Usage:
+        from langchain_openai import ChatOpenAI
+        from control_zero.integrations.langchain import ControlZeroCallbackHandler
+
+        client = ControlZeroClient(api_key="...")
+        client.initialize()
+
+        callback = ControlZeroCallbackHandler(client)
+
+        llm = ChatOpenAI(callbacks=[callback])
+        # or
+        agent.run("query", callbacks=[callback])
+    """
+
+    def __init__(
+        self,
+        client: ControlZeroClient,
+        log_prompts: bool = False,
+        log_responses: bool = False,
+        track_costs: bool = True,
+    ):
+        """
+        Initialize callback handler.
+
+        Args:
+            client: Control Zero client
+            log_prompts: Whether to log prompt content (may contain sensitive data)
+            log_responses: Whether to log response content
+            track_costs: Whether to track token costs
+        """
+        if not LANGCHAIN_AVAILABLE:
+            raise ImportError(
+                "langchain is required. Install with: pip install langchain-core"
+            )
+
+        super().__init__()
+        self._client = client
+        self._log_prompts = log_prompts
+        self._log_responses = log_responses
+        self._track_costs = track_costs
+
+        # Tracking state
+        self._run_starts: Dict[str, float] = {}
+        self._total_tokens = 0
+        self._total_cost = 0.0
+
+    # ==================== LLM Callbacks ====================
+
+    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:
+        """Called when LLM starts."""
+        self._run_starts[str(run_id)] = time.perf_counter()
+
+        model_name = serialized.get("name", "unknown_llm")
+
+        log_data = {
+            "event": "llm_start",
+            "model": model_name,
+            "num_prompts": len(prompts),
+        }
+
+        if self._log_prompts:
+            log_data["prompts"] = [p[:500] for p in prompts]
+
+        if tags:
+            log_data["tags"] = tags
+
+        self._client._log(
+            tool=f"llm:{model_name}",
+            method="start",
+            status="running",
+            latency_ms=0,
+            **log_data
+        )
+
+    def on_llm_end(
+        self,
+        response: LLMResult,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when LLM ends."""
+        start_time = self._run_starts.pop(str(run_id), time.perf_counter())
+        latency_ms = int((time.perf_counter() - start_time) * 1000)
+
+        # Extract usage info
+        llm_output = response.llm_output or {}
+        token_usage = llm_output.get("token_usage", {})
+        model_name = llm_output.get("model_name", "unknown_llm")
+
+        prompt_tokens = token_usage.get("prompt_tokens", 0)
+        completion_tokens = token_usage.get("completion_tokens", 0)
+        total_tokens = token_usage.get("total_tokens", prompt_tokens + completion_tokens)
+
+        if self._track_costs:
+            self._total_tokens += total_tokens
+
+        log_data = {
+            "event": "llm_end",
+            "model": model_name,
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": total_tokens,
+        }
+
+        if self._log_responses and response.generations:
+            texts = [g[0].text[:500] if g else "" for g in response.generations]
+            log_data["responses"] = texts
+
+        self._client._log(
+            tool=f"llm:{model_name}",
+            method="complete",
+            status="success",
+            latency_ms=latency_ms,
+            **log_data
+        )
+
+    def on_llm_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when LLM errors."""
+        start_time = self._run_starts.pop(str(run_id), time.perf_counter())
+        latency_ms = int((time.perf_counter() - start_time) * 1000)
+
+        self._client._log(
+            tool="llm:unknown",
+            method="complete",
+            status="error",
+            latency_ms=latency_ms,
+            error_type=type(error).__name__,
+            error_message=str(error),
+        )
+
+    # ==================== Chat Model Callbacks ====================
+
+    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:
+        """Called when chat model starts."""
+        self._run_starts[str(run_id)] = time.perf_counter()
+
+        model_name = serialized.get("name", kwargs.get("invocation_params", {}).get("model", "unknown_chat"))
+
+        log_data = {
+            "event": "chat_start",
+            "model": model_name,
+            "num_messages": sum(len(m) for m in messages),
+        }
+
+        if tags:
+            log_data["tags"] = tags
+
+        self._client._log(
+            tool=f"chat:{model_name}",
+            method="start",
+            status="running",
+            latency_ms=0,
+            **log_data
+        )
+
+    # ==================== Tool Callbacks ====================
+
+    def on_tool_start(
+        self,
+        serialized: Dict[str, Any],
+        input_str: str,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        inputs: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when tool starts."""
+        self._run_starts[str(run_id)] = time.perf_counter()
+
+        tool_name = serialized.get("name", "unknown_tool")
+
+        log_data = {
+            "event": "tool_start",
+        }
+
+        if self._log_prompts:
+            log_data["input"] = input_str[:500] if input_str else None
+
+        self._client._log(
+            tool=tool_name,
+            method="start",
+            status="running",
+            latency_ms=0,
+            **log_data
+        )
+
+    def on_tool_end(
+        self,
+        output: Any,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when tool ends."""
+        start_time = self._run_starts.pop(str(run_id), time.perf_counter())
+        latency_ms = int((time.perf_counter() - start_time) * 1000)
+
+        log_data = {"event": "tool_end"}
+
+        if self._log_responses:
+            log_data["output"] = str(output)[:500]
+
+        self._client._log(
+            tool="tool:unknown",
+            method="complete",
+            status="success",
+            latency_ms=latency_ms,
+            **log_data
+        )
+
+    def on_tool_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when tool errors."""
+        start_time = self._run_starts.pop(str(run_id), time.perf_counter())
+        latency_ms = int((time.perf_counter() - start_time) * 1000)
+
+        self._client._log(
+            tool="tool:unknown",
+            method="complete",
+            status="error",
+            latency_ms=latency_ms,
+            error_type=type(error).__name__,
+            error_message=str(error),
+        )
+
+    # ==================== Chain Callbacks ====================
+
+    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:
+        """Called when chain starts."""
+        self._run_starts[str(run_id)] = time.perf_counter()
+
+        chain_name = serialized.get("name", "unknown_chain")
+
+        log_data = {"event": "chain_start"}
+
+        if self._log_prompts:
+            log_data["inputs"] = {k: str(v)[:200] for k, v in inputs.items()}
+
+        self._client._log(
+            tool=f"chain:{chain_name}",
+            method="start",
+            status="running",
+            latency_ms=0,
+            **log_data
+        )
+
+    def on_chain_end(
+        self,
+        outputs: Dict[str, Any],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when chain ends."""
+        start_time = self._run_starts.pop(str(run_id), time.perf_counter())
+        latency_ms = int((time.perf_counter() - start_time) * 1000)
+
+        log_data = {"event": "chain_end"}
+
+        if self._log_responses:
+            log_data["outputs"] = {k: str(v)[:200] for k, v in outputs.items()}
+
+        self._client._log(
+            tool="chain:unknown",
+            method="complete",
+            status="success",
+            latency_ms=latency_ms,
+            **log_data
+        )
+
+    def on_chain_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when chain errors."""
+        start_time = self._run_starts.pop(str(run_id), time.perf_counter())
+        latency_ms = int((time.perf_counter() - start_time) * 1000)
+
+        self._client._log(
+            tool="chain:unknown",
+            method="complete",
+            status="error",
+            latency_ms=latency_ms,
+            error_type=type(error).__name__,
+            error_message=str(error),
+        )
+
+    # ==================== Agent Callbacks ====================
+
+    def on_agent_action(
+        self,
+        action: AgentAction,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when agent takes an action."""
+        self._client._log(
+            tool=f"agent:action:{action.tool}",
+            method="action",
+            status="running",
+            latency_ms=0,
+            event="agent_action",
+            tool_name=action.tool,
+            tool_input=str(action.tool_input)[:500] if self._log_prompts else None,
+        )
+
+    def on_agent_finish(
+        self,
+        finish: AgentFinish,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when agent finishes."""
+        self._client._log(
+            tool="agent:finish",
+            method="finish",
+            status="success",
+            latency_ms=0,
+            event="agent_finish",
+            output=str(finish.return_values)[:500] if self._log_responses else None,
+        )
+
+    # ==================== Metrics ====================
+
+    @property
+    def total_tokens(self) -> int:
+        """Get total tokens used."""
+        return self._total_tokens
+
+    @property
+    def total_cost(self) -> float:
+        """Get estimated total cost."""
+        return self._total_cost
+
+    def reset_metrics(self) -> None:
+        """Reset accumulated metrics."""
+        self._total_tokens = 0
+        self._total_cost = 0.0
+        self._run_starts.clear()