dispatch_agents 0.9.0__py3-none-any.whl

dispatch_agents/llm_langchain.py

@@ -0,0 +1,394 @@
+"""LangChain adapter for Dispatch LLM Gateway.
+
+Provides a LangChain-compatible ChatModel that uses the Dispatch LLM proxy
+for automatic trace correlation and cost tracking.
+
+Example:
+    from dispatch_agents.llm_langchain import ChatDispatch
+
+    # Use as a drop-in replacement for ChatOpenAI
+    llm = ChatDispatch(model="gpt-4o")
+
+    # Works with LangGraph
+    agent = create_react_agent(llm, tools=tools)
+
+    # Works with structured output
+    structured_llm = llm.with_structured_output(MySchema)
+"""
+
+from collections.abc import Sequence
+from typing import Any
+
+from langchain_core.callbacks import (
+    AsyncCallbackManagerForLLMRun,
+    CallbackManagerForLLMRun,
+)
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import (
+    AIMessage,
+    BaseMessage,
+    HumanMessage,
+    SystemMessage,
+    ToolMessage,
+)
+from langchain_core.outputs import ChatGeneration, ChatResult
+from pydantic import Field
+
+from .llm import LLMClient, LLMResponse
+
+
+def _get_tool_schema(tool: Any) -> dict[str, Any]:
+    """Extract JSON schema from a tool's args_schema, handling both Pydantic models and dicts."""
+    if not hasattr(tool, "args_schema") or not tool.args_schema:
+        return {}
+
+    args_schema = tool.args_schema
+
+    # If it's already a dict, return it directly
+    if isinstance(args_schema, dict):
+        return args_schema
+
+    # If it's a Pydantic model class, call model_json_schema()
+    if hasattr(args_schema, "model_json_schema"):
+        return args_schema.model_json_schema()
+
+    # Fallback: try to convert to dict
+    return {}
+
+
+def _convert_tools_to_openai_format(tools: Sequence[Any]) -> list[dict[str, Any]]:
+    """Convert various tool formats to OpenAI function calling format.
+
+    Handles:
+    - LangChain BaseTool objects
+    - MCP tools (with args_schema as dict)
+    - Raw dicts already in OpenAI format
+    - Callable functions with docstrings
+    """
+    formatted_tools = []
+    for tool in tools:
+        if hasattr(tool, "name") and hasattr(tool, "description"):
+            # LangChain BaseTool or MCP tool
+            tool_schema: dict[str, Any] = {
+                "type": "function",
+                "function": {
+                    "name": tool.name,
+                    "description": tool.description or "",
+                },
+            }
+            # Get parameters schema, handling both Pydantic models and dicts
+            params = _get_tool_schema(tool)
+            if params:
+                tool_schema["function"]["parameters"] = params
+            formatted_tools.append(tool_schema)
+        elif isinstance(tool, dict):
+            # Already in OpenAI format
+            formatted_tools.append(tool)
+        elif callable(tool):
+            # Function with docstring
+            import inspect
+
+            sig = inspect.signature(tool)
+            formatted_tools.append(
+                {
+                    "type": "function",
+                    "function": {
+                        "name": tool.__name__,
+                        "description": tool.__doc__ or "",
+                        "parameters": {
+                            "type": "object",
+                            "properties": {
+                                name: {"type": "string"} for name in sig.parameters
+                            },
+                        },
+                    },
+                }
+            )
+    return formatted_tools
+
+
+def _serialize_tool_arguments(args: Any) -> str:
+    """Serialize tool call arguments to a JSON string.
+
+    OpenAI-compatible APIs expect tool call arguments as a JSON string,
+    not a Python dict string representation.
+    """
+    import json
+
+    if isinstance(args, str):
+        return args
+    if isinstance(args, dict):
+        return json.dumps(args)
+    # Fallback: try to convert to dict then serialize
+    try:
+        return json.dumps(dict(args))
+    except (TypeError, ValueError):
+        return json.dumps({})
+
+
+def _convert_message_to_dict(message: BaseMessage) -> dict[str, Any]:
+    """Convert a LangChain message to a dict for the LLM API."""
+    if isinstance(message, SystemMessage):
+        return {"role": "system", "content": message.content}
+    elif isinstance(message, HumanMessage):
+        return {"role": "user", "content": message.content}
+    elif isinstance(message, AIMessage):
+        msg_dict: dict[str, Any] = {"role": "assistant", "content": message.content}
+        # Include tool calls if present
+        if message.tool_calls:
+            msg_dict["tool_calls"] = [
+                {
+                    "id": tc["id"],
+                    "type": "function",
+                    "function": {
+                        "name": tc["name"],
+                        "arguments": _serialize_tool_arguments(tc["args"]),
+                    },
+                }
+                for tc in message.tool_calls
+            ]
+        return msg_dict
+    elif isinstance(message, ToolMessage):
+        # tool_call_id is required by OpenAI-compatible APIs - use message id as fallback
+        tool_call_id = message.tool_call_id
+        if not tool_call_id:
+            # Fallback: use message id or generate one
+            tool_call_id = getattr(message, "id", None) or f"call_{id(message)}"
+        return {
+            "role": "tool",
+            "content": str(message.content) if message.content else "",
+            "tool_call_id": tool_call_id,
+        }
+    else:
+        # Fallback for other message types
+        return {"role": "user", "content": str(message.content)}
+
+
+def _parse_tool_arguments(arguments: Any) -> dict[str, Any]:
+    """Parse tool call arguments to a dict.
+
+    OpenAI returns arguments as a JSON string, but LangChain expects a dict.
+    Handles various input formats gracefully.
+    """
+    import json
+
+    if isinstance(arguments, dict):
+        return arguments
+    if isinstance(arguments, str):
+        if not arguments or arguments.strip() == "":
+            return {}
+        try:
+            parsed = json.loads(arguments)
+            return parsed if isinstance(parsed, dict) else {}
+        except json.JSONDecodeError:
+            return {}
+    return {}
+
+
+def _convert_response_to_message(response: LLMResponse) -> AIMessage:
+    """Convert an LLM response to a LangChain AIMessage."""
+    # Build tool_calls list if present
+    tool_calls = []
+    if response.tool_calls:
+        for tc in response.tool_calls:
+            tool_calls.append(
+                {
+                    "id": tc.id,
+                    "name": tc.function.name,
+                    "args": _parse_tool_arguments(tc.function.arguments),
+                }
+            )
+
+    # Build response metadata
+    response_metadata = {
+        "llm_call_id": response.llm_call_id,
+        "model": response.model,
+        "provider": response.provider,
+        "finish_reason": response.finish_reason,
+        "input_tokens": response.input_tokens,
+        "output_tokens": response.output_tokens,
+        "cost_usd": response.cost_usd,
+        "latency_ms": response.latency_ms,
+    }
+    if response.variant_name:
+        response_metadata["variant_name"] = response.variant_name
+
+    # Build usage metadata
+    usage_metadata = {
+        "input_tokens": response.input_tokens,
+        "output_tokens": response.output_tokens,
+        "total_tokens": response.input_tokens + response.output_tokens,
+    }
+
+    return AIMessage(
+        content=response.content or "",
+        tool_calls=tool_calls if tool_calls else [],
+        response_metadata=response_metadata,
+        usage_metadata=usage_metadata,
+    )
+
+
+class ChatDispatch(BaseChatModel):
+    """LangChain ChatModel that uses Dispatch LLM Gateway.
+
+    Provides automatic trace correlation and cost tracking for all LLM calls
+    when used within a Dispatch agent handler.
+
+    Example:
+        from dispatch_agents.llm_langchain import ChatDispatch
+
+        # Basic usage
+        llm = ChatDispatch(model="gpt-4o")
+        response = await llm.ainvoke([HumanMessage(content="Hello!")])
+
+        # With LangGraph ReAct agent
+        agent = create_react_agent(llm, tools=my_tools)
+
+        # With structured output
+        structured_llm = llm.with_structured_output(MyPydanticModel)
+        result = await structured_llm.ainvoke([...])
+
+    Args:
+        model: Model name (e.g., "gpt-4o", "claude-3-5-sonnet").
+               Uses org default if not specified.
+        provider: Provider name (e.g., "openai", "anthropic").
+                  Uses org default if not specified.
+        temperature: Sampling temperature (0-2), default 1.0
+        max_tokens: Maximum tokens in response (optional)
+    """
+
+    model: str | None = Field(default=None, description="Model name to use")
+    provider: str | None = Field(default=None, description="Provider name")
+    temperature: float = Field(default=1.0, description="Sampling temperature")
+    max_tokens: int | None = Field(default=None, description="Max tokens in response")
+
+    _client: LLMClient | None = None
+
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self._client = LLMClient()
+
+    @property
+    def _llm_type(self) -> str:
+        """Return the type of LLM."""
+        return "dispatch"
+
+    @property
+    def _identifying_params(self) -> dict[str, Any]:
+        """Return identifying parameters for caching."""
+        return {
+            "model": self.model,
+            "provider": self.provider,
+            "temperature": self.temperature,
+            "max_tokens": self.max_tokens,
+        }
+
+    def _generate(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> ChatResult:
+        """Synchronous generation - wraps async implementation."""
+        import asyncio
+
+        # Get or create event loop
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            loop = None
+
+        if loop and loop.is_running():
+            # We're in an async context, need to run in a new thread
+            import concurrent.futures
+
+            with concurrent.futures.ThreadPoolExecutor() as executor:
+                future = executor.submit(
+                    asyncio.run,
+                    self._agenerate(messages, stop, run_manager, **kwargs),  # type: ignore[arg-type]
+                )
+                return future.result()
+        else:
+            # No running loop, we can use asyncio.run
+            return asyncio.run(
+                self._agenerate(messages, stop, run_manager, **kwargs)  # type: ignore[arg-type]
+            )
+
+    async def _agenerate(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: AsyncCallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> ChatResult:
+        """Async generation using Dispatch LLM Gateway."""
+        if self._client is None:
+            self._client = LLMClient()
+
+        # Convert messages to API format
+        message_dicts = [_convert_message_to_dict(m) for m in messages]
+
+        # Extract tools from kwargs if provided (for function calling)
+        tools = kwargs.get("tools")
+        if tools:
+            # Convert tools to OpenAI format, handling various tool types
+            formatted_tools = _convert_tools_to_openai_format(tools)
+            tools = formatted_tools if formatted_tools else None
+
+        # Call the LLM Gateway
+        response = await self._client.inference(
+            messages=message_dicts,
+            model=self.model,
+            provider=self.provider,
+            tools=tools,
+            temperature=self.temperature,
+            max_tokens=self.max_tokens,
+        )
+
+        # Convert response to LangChain format
+        ai_message = _convert_response_to_message(response)
+        generation = ChatGeneration(message=ai_message)
+
+        return ChatResult(
+            generations=[generation],
+            llm_output={
+                "model": response.model,
+                "provider": response.provider,
+                "llm_call_id": response.llm_call_id,
+                "cost_usd": response.cost_usd,
+                "latency_ms": response.latency_ms,
+            },
+        )
+
+    def bind_tools(
+        self,
+        tools: Sequence[Any],
+        **kwargs: Any,
+    ) -> "ChatDispatch":
+        """Bind tools to this model for function calling.
+
+        Args:
+            tools: List of tools (LangChain tools, MCP tools, functions, or tool dicts)
+
+        Returns:
+            A new ChatDispatch instance with tools bound
+        """
+        # Convert tools to OpenAI format, handling various tool types
+        formatted_tools = _convert_tools_to_openai_format(tools)
+
+        # Return a new instance that will pass tools in kwargs.
+        # .bind() returns Runnable, not ChatDispatch, but callers expect ChatDispatch.
+        bound = (
+            self.__class__(
+                model=self.model,
+                provider=self.provider,
+                temperature=self.temperature,
+                max_tokens=self.max_tokens,
+            )
+            .configurable_fields(
+                # Store tools in config to be passed through
+            )
+            .bind(tools=formatted_tools, **kwargs)
+        )
+        return bound  # type: ignore[return-value]

dispatch_agents/logging_config.py

@@ -0,0 +1,133 @@
+"""Logging configuration for the Dispatch Agents SDK.
+
+Controls the verbosity of SDK logging output via environment variables:
+- DISPATCH_VERBOSE=1 or DISPATCH_VERBOSE=true: Enable verbose/debug logging
+- DISPATCH_LOG_LEVEL=DEBUG/INFO/WARNING/ERROR: Set specific log level
+
+By default, the SDK logs at INFO level, hiding debug messages like
+re-subscription events. Enable verbose mode for debugging.
+"""
+
+import logging
+import os
+import sys
+
+# SDK logger namespace
+SDK_LOGGER_NAME = "dispatch_agents"
+
+# Check if we've already configured logging
+_logging_configured = False
+
+
+def _parse_bool_env(value: str | None) -> bool:
+    """Parse boolean from environment variable value."""
+    if value is None:
+        return False
+    return value.lower() in ("1", "true", "yes", "on")
+
+
+def _get_log_level() -> int:
+    """Determine the log level from environment variables.
+
+    Priority:
+    1. DISPATCH_LOG_LEVEL (explicit level)
+    2. DISPATCH_VERBOSE (boolean toggle for DEBUG)
+    3. Default: WARNING (minimal output)
+    """
+    # Check explicit log level first
+    log_level_str = os.environ.get("DISPATCH_LOG_LEVEL", "").upper()
+    if log_level_str:
+        level_map = {
+            "DEBUG": logging.DEBUG,
+            "INFO": logging.INFO,
+            "WARNING": logging.WARNING,
+            "WARN": logging.WARNING,
+            "ERROR": logging.ERROR,
+            "CRITICAL": logging.CRITICAL,
+        }
+        if log_level_str in level_map:
+            return level_map[log_level_str]
+
+    # Check verbose flag
+    if _parse_bool_env(os.environ.get("DISPATCH_VERBOSE")):
+        return logging.DEBUG
+
+    # Default: WARNING to minimize noise
+    # Users see errors and warnings, but not routine info/debug
+    return logging.WARNING
+
+
+def configure_logging(force: bool = False) -> None:
+    """Configure logging for the Dispatch Agents SDK.
+
+    This sets up the SDK logger with appropriate level and format.
+    Called automatically when the SDK is imported, but can be called
+    again with force=True to reconfigure.
+
+    Args:
+        force: If True, reconfigure even if already configured
+    """
+    global _logging_configured
+
+    if _logging_configured and not force:
+        return
+
+    # Get the SDK root logger
+    sdk_logger = logging.getLogger(SDK_LOGGER_NAME)
+
+    # Set the log level
+    level = _get_log_level()
+    sdk_logger.setLevel(level)
+
+    # Only add handler if none exist (avoid duplicate handlers on force)
+    if not sdk_logger.handlers:
+        # Create a handler that writes to stderr
+        handler = logging.StreamHandler(sys.stderr)
+        handler.setLevel(level)
+
+        # Format: simple for normal use, more detail for debug
+        if level == logging.DEBUG:
+            formatter = logging.Formatter(
+                "%(asctime)s [%(name)s] %(levelname)s: %(message)s",
+                datefmt="%H:%M:%S",
+            )
+        else:
+            formatter = logging.Formatter("%(levelname)s: %(message)s")
+
+        handler.setFormatter(formatter)
+        sdk_logger.addHandler(handler)
+
+    # Prevent propagation to root logger (avoid duplicate messages)
+    sdk_logger.propagate = False
+
+    _logging_configured = True
+
+
+def get_logger(name: str | None = None) -> logging.Logger:
+    """Get a logger for SDK modules.
+
+    Args:
+        name: Optional module name (e.g., __name__). If provided,
+              creates a child logger under SDK_LOGGER_NAME.
+
+    Returns:
+        Configured logger instance
+    """
+    # Ensure logging is configured
+    configure_logging()
+
+    if name:
+        # Create child logger: dispatch_agents.grpc_server, etc.
+        if name.startswith(SDK_LOGGER_NAME):
+            return logging.getLogger(name)
+        return logging.getLogger(f"{SDK_LOGGER_NAME}.{name}")
+    return logging.getLogger(SDK_LOGGER_NAME)
+
+
+def is_verbose() -> bool:
+    """Check if verbose logging is enabled.
+
+    Returns:
+        True if DISPATCH_VERBOSE is set or log level is DEBUG
+    """
+    return _get_log_level() == logging.DEBUG