hud-python 0.4.35__py3-none-any.whl → 0.4.37__py3-none-any.whl

hud/agents/__init__.py

@@ -2,12 +2,14 @@ from __future__ import annotations
 
 from .base import MCPAgent
 from .claude import ClaudeAgent
+from .lite_llm import LiteAgent
 from .openai import OperatorAgent
 from .openai_chat_generic import GenericOpenAIChatAgent
 
 __all__ = [
     "ClaudeAgent",
     "GenericOpenAIChatAgent",
+    "LiteAgent",
     "MCPAgent",
     "OperatorAgent",
 ]

hud/agents/lite_llm.py

@@ -0,0 +1,72 @@
+"""LiteLLM MCP Agent implementation.
+
+Same OpenAI chat-completions shape + MCP tool plumbing,
+but transport is LiteLLM and (optionally) tools are shaped by LiteLLM's MCP transformer.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, ClassVar
+
+import litellm
+
+from .openai_chat_generic import GenericOpenAIChatAgent
+
+logger = logging.getLogger(__name__)
+
+# Prefer LiteLLM's built-in MCP -> OpenAI tool transformer (handles Bedrock nuances)
+try:
+    from litellm.experimental_mcp_client.tools import (
+        transform_mcp_tool_to_openai_tool,
+    )
+except Exception:  # pragma: no cover - optional dependency
+    transform_mcp_tool_to_openai_tool = None  # type: ignore
+
+
+class LiteAgent(GenericOpenAIChatAgent):
+    """
+    Same OpenAI chat-completions shape + MCP tool plumbing,
+    but transport is LiteLLM and (optionally) tools are shaped by LiteLLM's MCP transformer.
+    """
+
+    metadata: ClassVar[dict[str, Any]] = {}
+
+    def __init__(
+        self,
+        *,
+        model_name: str = "gpt-4o-mini",
+        completion_kwargs: dict[str, Any] | None = None,
+        **agent_kwargs: Any,
+    ) -> None:
+        # We don't need an OpenAI client; pass None
+        super().__init__(
+            openai_client=None,
+            model_name=model_name,
+            completion_kwargs=completion_kwargs,
+            **agent_kwargs,
+        )
+
+    def get_tool_schemas(self) -> list[dict]:
+        # Prefer LiteLLM's stricter transformer (handles Bedrock & friends)
+        if transform_mcp_tool_to_openai_tool is not None:
+            return [
+                transform_mcp_tool_to_openai_tool(t)  # returns ChatCompletionToolParam-like dict
+                for t in self.get_available_tools()
+            ]
+        # Fallback to the generic OpenAI sanitizer
+        return GenericOpenAIChatAgent.get_tool_schemas(self)
+
+    async def _invoke_chat_completion(
+        self,
+        *,
+        messages: list[Any],
+        tools: list[dict] | None,
+        extra: dict[str, Any],
+    ):
+        return await litellm.acompletion(
+            model=self.model_name,
+            messages=messages,
+            tools=tools or None,  # LiteLLM tolerates None better than []
+            **extra,
+        )

hud/agents/openai_chat_generic.py

@@ -42,7 +42,7 @@ class GenericOpenAIChatAgent(MCPAgent):
     def __init__(
         self,
         *,
-        openai_client: AsyncOpenAI,
+        openai_client: AsyncOpenAI | None,
         model_name: str = "gpt-4o-mini",
         completion_kwargs: dict[str, Any] | None = None,
         **agent_kwargs: Any,
@@ -171,6 +171,23 @@ class GenericOpenAIChatAgent(MCPAgent):
             openai_tools.append(openai_tool)
         return openai_tools
 
+    async def _invoke_chat_completion(
+        self,
+        *,
+        messages: list[Any],
+        tools: list[dict] | None,
+        extra: dict[str, Any],
+    ):
+        if self.oai is None:
+            raise ValueError("openai_client is required for GenericOpenAIChatAgent")
+        # default transport = OpenAI SDK
+        return await self.oai.chat.completions.create(
+            model=self.model_name,
+            messages=messages,
+            tools=tools,  # already ChatCompletionToolParam-shaped
+            **extra,
+        )
+
     @instrument(
         span_type="agent",
         record_args=False,
@@ -180,17 +197,14 @@ class GenericOpenAIChatAgent(MCPAgent):
         """Send chat request to OpenAI and convert the response."""
 
         # Convert MCP tool schemas to OpenAI format
-        mcp_schemas = self.get_tool_schemas()
+        tools = cast("list[ChatCompletionToolParam]", self.get_tool_schemas())
 
         protected_keys = {"model", "messages", "tools"}
         extra = {k: v for k, v in (self.completion_kwargs or {}).items() if k not in protected_keys}
 
         try:
-            response = await self.oai.chat.completions.create(
-                model=self.model_name,
-                messages=messages,
-                tools=cast("list[ChatCompletionToolParam]", mcp_schemas),
-                **extra,
+            response = await self._invoke_chat_completion(
+                messages=messages, tools=tools, extra=extra
             )
         except Exception as e:
             error_content = f"Error getting response {e}"

hud/agents/tests/test_claude.py

@@ -86,6 +86,7 @@ class TestClaudeAgent:
             model_client=mock_model_client,
             model="claude-3-opus-20240229",
             max_tokens=1000,
+            validate_api_key=False,  # Skip validation in tests
         )
 
         assert agent.model_name == "claude-3-opus-20240229"
@@ -93,10 +94,14 @@ class TestClaudeAgent:
         assert agent.anthropic_client == mock_model_client
 
     @pytest.mark.asyncio
-    async def test_init_without_model_client(self, mock_mcp_client):
+    async def test_init_without_model_client(self, mock_mcp_client, mock_anthropic):
         """Test agent initialization without model client."""
         with patch("hud.settings.settings.anthropic_api_key", "test_key"):
-            agent = ClaudeAgent(mcp_client=mock_mcp_client, model="claude-3-opus-20240229")
+            agent = ClaudeAgent(
+                mcp_client=mock_mcp_client,
+                model="claude-3-opus-20240229",
+                validate_api_key=False,  # Skip validation in tests
+            )
 
             assert agent.model_name == "claude-3-opus-20240229"
             assert agent.anthropic_client is not None
@@ -105,7 +110,11 @@ class TestClaudeAgent:
     async def test_format_blocks(self, mock_mcp_client):
         """Test formatting content blocks into Claude messages."""
         mock_model_client = MagicMock()
-        agent = ClaudeAgent(mcp_client=mock_mcp_client, model_client=mock_model_client)
+        agent = ClaudeAgent(
+            mcp_client=mock_mcp_client,
+            model_client=mock_model_client,
+            validate_api_key=False,  # Skip validation in tests
+        )
 
         # Test with text only
         text_blocks: list[types.ContentBlock] = [
@@ -141,7 +150,11 @@ class TestClaudeAgent:
     async def test_format_tool_results_method(self, mock_mcp_client):
         """Test the agent's format_tool_results method."""
         mock_model_client = MagicMock()
-        agent = ClaudeAgent(mcp_client=mock_mcp_client, model_client=mock_model_client)
+        agent = ClaudeAgent(
+            mcp_client=mock_mcp_client,
+            model_client=mock_model_client,
+            validate_api_key=False,  # Skip validation in tests
+        )
 
         tool_calls = [
             MCPToolCall(name="test_tool", arguments={}, id="id1"),
@@ -171,7 +184,11 @@ class TestClaudeAgent:
         """Test getting model response from Claude API."""
         # Disable telemetry for this test to avoid backend configuration issues
         with patch("hud.settings.settings.telemetry_enabled", False):
-            agent = ClaudeAgent(mcp_client=mock_mcp_client, model_client=mock_anthropic)
+            agent = ClaudeAgent(
+                mcp_client=mock_mcp_client,
+                model_client=mock_anthropic,
+                validate_api_key=False,  # Skip validation in tests
+            )
 
             # Mock the API response
             mock_response = MagicMock()
@@ -215,7 +232,11 @@ class TestClaudeAgent:
         """Test getting text-only response."""
         # Disable telemetry for this test to avoid backend configuration issues
         with patch("hud.settings.settings.telemetry_enabled", False):
-            agent = ClaudeAgent(mcp_client=mock_mcp_client, model_client=mock_anthropic)
+            agent = ClaudeAgent(
+                mcp_client=mock_mcp_client,
+                model_client=mock_anthropic,
+                validate_api_key=False,  # Skip validation in tests
+            )
 
             mock_response = MagicMock()
             # Create text block
@@ -242,7 +263,11 @@ class TestClaudeAgent:
         """Test handling API errors."""
         # Disable telemetry for this test to avoid backend configuration issues
         with patch("hud.settings.settings.telemetry_enabled", False):
-            agent = ClaudeAgent(mcp_client=mock_mcp_client, model_client=mock_anthropic)
+            agent = ClaudeAgent(
+                mcp_client=mock_mcp_client,
+                model_client=mock_anthropic,
+                validate_api_key=False,  # Skip validation in tests
+            )
 
             # Mock API error
             mock_anthropic.beta.messages.create = AsyncMock(

hud/agents/tests/test_openai.py

@@ -44,7 +44,10 @@ class TestOperatorAgent:
         """Test agent initialization."""
         mock_model_client = MagicMock()
         agent = OperatorAgent(
-            mcp_client=mock_mcp_client, model_client=mock_model_client, model="gpt-4"
+            mcp_client=mock_mcp_client,
+            model_client=mock_model_client,
+            model="gpt-4",
+            validate_api_key=False,  # Skip validation in tests
         )
 
         assert agent.model_name == "openai-gpt-4"
@@ -55,7 +58,11 @@ class TestOperatorAgent:
     async def test_format_blocks(self, mock_mcp_client):
         """Test formatting content blocks."""
         mock_model_client = MagicMock()
-        agent = OperatorAgent(mcp_client=mock_mcp_client, model_client=mock_model_client)
+        agent = OperatorAgent(
+            mcp_client=mock_mcp_client,
+            model_client=mock_model_client,
+            validate_api_key=False,  # Skip validation in tests
+        )
 
         # Test with text blocks
         blocks: list[types.ContentBlock] = [
@@ -85,7 +92,11 @@ class TestOperatorAgent:
     @pytest.mark.asyncio
     async def test_format_tool_results(self, mock_mcp_client, mock_openai):
         """Test formatting tool results."""
-        agent = OperatorAgent(mcp_client=mock_mcp_client, model_client=mock_openai)
+        agent = OperatorAgent(
+            mcp_client=mock_mcp_client,
+            model_client=mock_openai,
+            validate_api_key=False,  # Skip validation in tests
+        )
 
         tool_calls = [
             MCPToolCall(name="test_tool", arguments={}, id="call_123"),  # type: ignore
@@ -111,7 +122,11 @@ class TestOperatorAgent:
     @pytest.mark.asyncio
     async def test_format_tool_results_with_error(self, mock_mcp_client, mock_openai):
         """Test formatting tool results with errors."""
-        agent = OperatorAgent(mcp_client=mock_mcp_client, model_client=mock_openai)
+        agent = OperatorAgent(
+            mcp_client=mock_mcp_client,
+            model_client=mock_openai,
+            validate_api_key=False,  # Skip validation in tests
+        )
 
         tool_calls = [
             MCPToolCall(name="failing_tool", arguments={}, id="call_error"),  # type: ignore
@@ -131,7 +146,11 @@ class TestOperatorAgent:
     @pytest.mark.asyncio
     async def test_get_model_response(self, mock_mcp_client, mock_openai):
         """Test getting model response from OpenAI API."""
-        agent = OperatorAgent(mcp_client=mock_mcp_client, model_client=mock_openai)
+        agent = OperatorAgent(
+            mcp_client=mock_mcp_client,
+            model_client=mock_openai,
+            validate_api_key=False,  # Skip validation in tests
+        )
 
         # Set up available tools so agent doesn't return "No computer use tools available"
         agent._available_tools = [
@@ -162,7 +181,11 @@ class TestOperatorAgent:
     @pytest.mark.asyncio
     async def test_handle_empty_response(self, mock_mcp_client, mock_openai):
         """Test handling empty response from API."""
-        agent = OperatorAgent(mcp_client=mock_mcp_client, model_client=mock_openai)
+        agent = OperatorAgent(
+            mcp_client=mock_mcp_client,
+            model_client=mock_openai,
+            validate_api_key=False,  # Skip validation in tests
+        )
 
         # Set up available tools
         agent._available_tools = [