webber/webber-cli/webber_cli/client.py

"""
Webber API client.

Communicates with the Webber API backend for agent execution.
"""
import json
import httpx
from collections.abc import AsyncIterator
from dataclasses import dataclass
from typing import Any


@dataclass
class AgentResponse:
    """Response from agent execution."""
    response: str
    agent_type: str
    success: bool
    error: str | None = None


@dataclass
class AgentInfo:
    """Information about an available agent."""
    name: str
    description: str


class WebberClient:
    """
    Client for the Webber API.

    Usage:
        client = WebberClient("http://localhost:8086")
        response = await client.run_agent("explore", "find python files", "/path/to/project")
    """

    def __init__(
        self,
        base_url: str = "http://localhost:8086",
        api_key: str | None = None,
        timeout: float = 120.0,
    ):
        """
        Initialize the Webber client.

        Args:
            base_url: Webber API URL
            api_key: Optional API key for authentication
            timeout: Request timeout in seconds
        """
        self.base_url = base_url.rstrip("/")
        self.api_key = api_key
        self.timeout = timeout
        self._client: httpx.AsyncClient | None = None

    async def _get_client(self) -> httpx.AsyncClient:
        """Get or create the HTTP client."""
        if self._client is None or self._client.is_closed:
            headers = {}
            if self.api_key:
                headers["X-API-Key"] = self.api_key
            self._client = httpx.AsyncClient(
                base_url=self.base_url,
                headers=headers,
                timeout=self.timeout,
            )
        return self._client

    async def close(self) -> None:
        """Close the HTTP client."""
        if self._client and not self._client.is_closed:
            await self._client.aclose()
            self._client = None

    async def health_check(self) -> bool:
        """Check if the API is healthy."""
        try:
            client = await self._get_client()
            response = await client.get("/health")
            return response.status_code == 200
        except httpx.RequestError:
            return False

    async def list_agents(self) -> list[AgentInfo]:
        """List available agents."""
        client = await self._get_client()
        response = await client.get("/agents/")
        response.raise_for_status()
        data = response.json()
        return [AgentInfo(**a) for a in data.get("agents", [])]

    async def get_agent(self, agent_type: str) -> AgentInfo | None:
        """Get information about a specific agent."""
        client = await self._get_client()
        response = await client.get(f"/agents/{agent_type}")
        if response.status_code == 404:
            return None
        response.raise_for_status()
        return AgentInfo(**response.json())

    async def run_agent(
        self,
        agent_type: str,
        prompt: str,
        working_dir: str = ".",
    ) -> AgentResponse:
        """
        Run an agent with the given prompt.

        Args:
            agent_type: Type of agent (e.g., "explore")
            prompt: User prompt/query
            working_dir: Working directory for the agent

        Returns:
            AgentResponse with the result
        """
        client = await self._get_client()
        response = await client.post(
            "/agents/run",
            json={
                "agent_type": agent_type,
                "prompt": prompt,
                "working_dir": working_dir,
            },
        )
        response.raise_for_status()
        data = response.json()
        return AgentResponse(
            response=data.get("response", ""),
            agent_type=data.get("agent_type", agent_type),
            success=data.get("success", True),
            error=data.get("error"),
        )

    async def run_agent_stream(
        self,
        agent_type: str,
        prompt: str,
        working_dir: str = ".",
    ) -> AsyncIterator[str]:
        """
        Run an agent with streaming response.

        Args:
            agent_type: Type of agent (e.g., "explore")
            prompt: User prompt/query
            working_dir: Working directory for the agent

        Yields:
            Text chunks as they arrive
        """
        # Use a fresh client for streaming with longer timeout
        async with httpx.AsyncClient(
            base_url=self.base_url,
            timeout=httpx.Timeout(300.0, connect=10.0),
        ) as client:
            async with client.stream(
                "POST",
                "/agents/stream",
                json={
                    "agent_type": agent_type,
                    "prompt": prompt,
                    "working_dir": working_dir,
                },
            ) as response:
                response.raise_for_status()
                async for line in response.aiter_lines():
                    if line.startswith("data: "):
                        try:
                            data = json.loads(line[6:])
                            event = data.get("event")
                            if event == "chunk":
                                yield data.get("data", "")
                            elif event == "error":
                                raise Exception(data.get("data", "Unknown error"))
                            elif event == "done":
                                break
                        except json.JSONDecodeError:
                            continue

    async def __aenter__(self) -> "WebberClient":
        """Async context manager entry."""
        return self

    async def __aexit__(self, *args: Any) -> None:
        """Async context manager exit."""
        await self.close()