Feat/scenario improvements

hud/agents/__init__.py

@@ -1,19 +1,70 @@
 from __future__ import annotations
 
+from typing import Any
+
 from .base import MCPAgent
 from .openai import OpenAIAgent
 from .openai_chat import OpenAIChatAgent
 from .operator import OperatorAgent
 
-# Note: These agents are not exported here to avoid requiring optional dependencies.
-# Import directly if needed:
-#   from hud.agents.claude import ClaudeAgent  # requires anthropic
-#   from hud.agents.gemini import GeminiAgent  # requires google-genai
-#   from hud.agents.gemini_cua import GeminiCUAAgent  # requires google-genai
-
 __all__ = [
     "MCPAgent",
     "OpenAIAgent",
     "OpenAIChatAgent",
     "OperatorAgent",
+    "create_agent",
 ]
+
+
+def create_agent(model: str, **kwargs: Any) -> MCPAgent:
+    """Create an agent for a gateway model.
+
+    This routes ALL requests through the HUD gateway. For direct API access
+    (using your own API keys), use the agent classes directly.
+
+    Args:
+        model: Model name (e.g., "gpt-4o", "claude-sonnet-4-5").
+        **kwargs: Additional params passed to agent.create().
+
+    Returns:
+        Configured MCPAgent instance with gateway routing.
+
+    Example:
+        ```python
+        # Gateway routing (recommended)
+        agent = create_agent("gpt-4o")
+        agent = create_agent("claude-sonnet-4-5", temperature=0.7)
+
+        # Direct API access (use agent classes)
+        from hud.agents.claude import ClaudeAgent
+
+        agent = ClaudeAgent.create(model="claude-sonnet-4-5")
+        ```
+    """
+    from hud.agents.gateway import build_gateway_client
+    from hud.agents.resolver import resolve_cls
+
+    # Resolve class and gateway info
+    agent_cls, gateway_info = resolve_cls(model)
+
+    # Get model ID from gateway info or use input
+    model_id = model
+    if gateway_info:
+        model_id = gateway_info.get("model") or gateway_info.get("id") or model
+
+    # Build gateway client
+    provider = gateway_info.get("provider", "openai") if gateway_info else "openai"
+    client = build_gateway_client(provider)
+
+    # Set up kwargs
+    kwargs.setdefault("model", model_id)
+
+    # Use correct client key based on agent type
+    if agent_cls == OpenAIChatAgent:
+        kwargs.setdefault("openai_client", client)
+    else:
+        # Claude and other agents use model_client and validate_api_key
+        kwargs.setdefault("model_client", client)
+        kwargs.setdefault("validate_api_key", False)
+
+    return agent_cls.create(**kwargs)

hud/agents/gateway.py

@@ -0,0 +1,42 @@
+"""Gateway client utilities for HUD inference gateway."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def build_gateway_client(provider: str) -> Any:
+    """Build a client configured for HUD gateway routing.
+
+    Args:
+        provider: Provider name ("anthropic", "openai", "gemini", etc.)
+
+    Returns:
+        Configured async client for the provider.
+    """
+    from hud.settings import settings
+
+    provider = provider.lower()
+
+    if provider == "anthropic":
+        from anthropic import AsyncAnthropic
+
+        return AsyncAnthropic(api_key=settings.api_key, base_url=settings.hud_gateway_url)
+
+    if provider == "gemini":
+        from google import genai
+        from google.genai.types import HttpOptions
+
+        return genai.Client(
+            api_key="PLACEHOLDER",
+            http_options=HttpOptions(
+                api_version="v1beta",
+                base_url=settings.hud_gateway_url,
+                headers={"Authorization": f"Bearer {settings.api_key}"},
+            ),
+        )
+
+    # OpenAI-compatible (openai, azure, together, groq, fireworks, etc.)
+    from openai import AsyncOpenAI
+
+    return AsyncOpenAI(api_key=settings.api_key, base_url=settings.hud_gateway_url)

hud/agents/resolver.py

@@ -0,0 +1,70 @@
+"""Model resolution - maps model strings to agent classes."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from hud.agents.base import MCPAgent
+
+__all__ = ["resolve_cls"]
+
+_models_cache: list[dict[str, Any]] | None = None
+
+# Provider name → AgentType value (only anthropic differs)
+_PROVIDER_TO_AGENT = {"anthropic": "claude"}
+
+
+def _fetch_gateway_models() -> list[dict[str, Any]]:
+    """Fetch available models from HUD gateway (cached)."""
+    global _models_cache
+    if _models_cache is not None:
+        return _models_cache
+
+    import httpx
+
+    from hud.settings import settings
+
+    if not settings.api_key:
+        return []
+
+    try:
+        resp = httpx.get(
+            f"{settings.hud_gateway_url}/models",
+            headers={"Authorization": f"Bearer {settings.api_key}"},
+            timeout=10.0,
+        )
+        resp.raise_for_status()
+        data = resp.json()
+        _models_cache = data.get("data", data) if isinstance(data, dict) else data
+        return _models_cache or []
+    except Exception:
+        return []
+
+
+def resolve_cls(model: str) -> tuple[type[MCPAgent], dict[str, Any] | None]:
+    """Resolve model string to (agent_class, gateway_info).
+
+    Returns:
+        (agent_class, None) for known AgentTypes
+        (agent_class, gateway_model_info) for gateway models
+    """
+    from hud.types import AgentType
+
+    # Known AgentType → no gateway info
+    try:
+        return AgentType(model).cls, None
+    except ValueError:
+        pass
+
+    # Gateway lookup
+    for m in _fetch_gateway_models():
+        if model in (m.get("id"), m.get("name"), m.get("model")):
+            provider = m.get("provider", "openai_compatible").lower()
+            agent_str = _PROVIDER_TO_AGENT.get(provider, provider)
+            try:
+                return AgentType(agent_str).cls, m
+            except ValueError:
+                return AgentType.OPENAI_COMPATIBLE.cls, m
+
+    raise ValueError(f"Model '{model}' not found")

hud/agents/tests/test_resolver.py

@@ -0,0 +1,192 @@
+"""Tests for model resolution and create_agent."""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from hud.agents import create_agent
+from hud.agents.resolver import resolve_cls
+
+
+class TestResolveCls:
+    """Tests for resolve_cls function."""
+
+    def test_resolves_known_agent_type(self) -> None:
+        """Known AgentType strings resolve to their class."""
+        from hud.agents.claude import ClaudeAgent
+
+        cls, gateway_info = resolve_cls("claude")
+        assert cls == ClaudeAgent
+        assert gateway_info is None
+
+    def test_resolves_openai(self) -> None:
+        """Resolves 'openai' to OpenAIAgent."""
+        from hud.agents import OpenAIAgent
+
+        cls, _gateway_info = resolve_cls("openai")
+        assert cls == OpenAIAgent
+
+    def test_resolves_gemini(self) -> None:
+        """Resolves 'gemini' to GeminiAgent."""
+        from hud.agents.gemini import GeminiAgent
+
+        cls, _gateway_info = resolve_cls("gemini")
+        assert cls == GeminiAgent
+
+    def test_unknown_model_without_gateway_raises(self) -> None:
+        """Unknown model with no gateway models raises ValueError."""
+        with (
+            patch("hud.agents.resolver._fetch_gateway_models", return_value=[]),
+            pytest.raises(ValueError, match="not found"),
+        ):
+            resolve_cls("unknown-model-xyz")
+
+    def test_resolves_gateway_model(self) -> None:
+        """Resolves model found in gateway."""
+        from hud.agents import OpenAIAgent
+
+        mock_models = [
+            {"id": "gpt-4o", "model": "gpt-4o", "provider": "openai"},
+        ]
+
+        with patch("hud.agents.resolver._fetch_gateway_models", return_value=mock_models):
+            cls, info = resolve_cls("gpt-4o")
+            assert cls == OpenAIAgent
+            assert info is not None
+            assert info["id"] == "gpt-4o"
+
+    def test_resolves_anthropic_provider_to_claude(self) -> None:
+        """Provider 'anthropic' maps to ClaudeAgent."""
+        from hud.agents.claude import ClaudeAgent
+
+        mock_models = [
+            {"id": "claude-sonnet", "model": "claude-3-sonnet", "provider": "anthropic"},
+        ]
+
+        with patch("hud.agents.resolver._fetch_gateway_models", return_value=mock_models):
+            cls, _info = resolve_cls("claude-sonnet")
+            assert cls == ClaudeAgent
+
+    def test_resolves_unknown_provider_to_openai_compatible(self) -> None:
+        """Unknown provider maps to OpenAIChatAgent."""
+        from hud.agents.openai_chat import OpenAIChatAgent
+
+        mock_models = [
+            {"id": "custom-model", "model": "custom", "provider": "custom-provider"},
+        ]
+
+        with patch("hud.agents.resolver._fetch_gateway_models", return_value=mock_models):
+            cls, _info = resolve_cls("custom-model")
+            assert cls == OpenAIChatAgent
+
+
+class TestCreateAgent:
+    """Tests for create_agent function - gateway-only."""
+
+    def test_creates_with_gateway_client(self) -> None:
+        """create_agent always uses gateway routing."""
+        from hud.agents import OpenAIAgent
+
+        mock_models = [
+            {"id": "gpt-4o", "model": "gpt-4o", "provider": "openai"},
+        ]
+
+        with (
+            patch("hud.agents.resolver._fetch_gateway_models", return_value=mock_models),
+            patch.object(OpenAIAgent, "create") as mock_create,
+            patch("hud.agents.gateway.build_gateway_client") as mock_build_client,
+        ):
+            mock_client = MagicMock()
+            mock_build_client.return_value = mock_client
+            mock_agent = MagicMock()
+            mock_create.return_value = mock_agent
+
+            agent = create_agent("gpt-4o")
+
+            # Should have set model and model_client
+            call_kwargs = mock_create.call_args.kwargs
+            assert call_kwargs["model"] == "gpt-4o"
+            assert "model_client" in call_kwargs
+            assert agent == mock_agent
+
+    def test_passes_kwargs_to_create(self) -> None:
+        """Extra kwargs are passed to agent.create()."""
+        from hud.agents import OpenAIAgent
+
+        mock_models = [
+            {"id": "gpt-4o", "model": "gpt-4o", "provider": "openai"},
+        ]
+
+        with (
+            patch("hud.agents.resolver._fetch_gateway_models", return_value=mock_models),
+            patch.object(OpenAIAgent, "create") as mock_create,
+            patch("hud.agents.gateway.build_gateway_client"),
+        ):
+            mock_create.return_value = MagicMock()
+
+            create_agent("gpt-4o", temperature=0.5, max_tokens=1000)
+
+            call_kwargs = mock_create.call_args.kwargs
+            assert call_kwargs["temperature"] == 0.5
+            assert call_kwargs["max_tokens"] == 1000
+
+    def test_known_agent_type_also_uses_gateway(self) -> None:
+        """Even 'claude' string uses gateway (it's a gateway shortcut)."""
+        from hud.agents.claude import ClaudeAgent
+
+        with (
+            patch.object(ClaudeAgent, "create") as mock_create,
+            patch("hud.agents.gateway.build_gateway_client") as mock_build_client,
+        ):
+            mock_client = MagicMock()
+            mock_build_client.return_value = mock_client
+            mock_create.return_value = MagicMock()
+
+            create_agent("claude")
+
+            # Should still build gateway client
+            mock_build_client.assert_called_once()
+            call_kwargs = mock_create.call_args.kwargs
+            assert "model_client" in call_kwargs
+
+
+class TestBuildGatewayClient:
+    """Tests for build_gateway_client function."""
+
+    def test_builds_anthropic_client(self) -> None:
+        """Builds AsyncAnthropic for anthropic provider."""
+        from hud.agents.gateway import build_gateway_client
+
+        with patch("hud.settings.settings") as mock_settings:
+            mock_settings.api_key = "test-key"
+            mock_settings.hud_gateway_url = "https://gateway.hud.ai"
+
+            with patch("anthropic.AsyncAnthropic") as mock_client_cls:
+                build_gateway_client("anthropic")
+                mock_client_cls.assert_called_once()
+
+    def test_builds_openai_client_for_openai(self) -> None:
+        """Builds AsyncOpenAI for openai provider."""
+        from hud.agents.gateway import build_gateway_client
+
+        with patch("hud.settings.settings") as mock_settings:
+            mock_settings.api_key = "test-key"
+            mock_settings.hud_gateway_url = "https://gateway.hud.ai"
+
+            with patch("openai.AsyncOpenAI") as mock_client_cls:
+                build_gateway_client("openai")
+                mock_client_cls.assert_called_once()
+
+    def test_builds_openai_client_for_unknown(self) -> None:
+        """Builds AsyncOpenAI for unknown providers (openai-compatible)."""
+        from hud.agents.gateway import build_gateway_client
+
+        with patch("hud.settings.settings") as mock_settings:
+            mock_settings.api_key = "test-key"
+            mock_settings.hud_gateway_url = "https://gateway.hud.ai"
+
+            with patch("openai.AsyncOpenAI") as mock_client_cls:
+                build_gateway_client("together")
+                mock_client_cls.assert_called_once()