arbiter-cli 0.1.0__py3-none-any.whl

arbiter/core/providers/anthropic_provider.py

@@ -0,0 +1,133 @@
+"""Anthropic provider - Claude models via the Anthropic API."""
+
+from __future__ import annotations
+
+import json
+from typing import AsyncIterator, Optional
+
+import httpx
+
+from arbiter.core.providers.base import LLMProvider, StreamChunk
+
+
+class AnthropicProvider(LLMProvider):
+    """Provider for Anthropic Claude models."""
+
+    provider_name = "anthropic"
+
+    def __init__(self, api_key: str):
+        self.api_key = api_key
+        self.base_url = "https://api.anthropic.com/v1"
+
+    async def stream_generate(
+        self,
+        model: str,
+        prompt: str,
+        system: Optional[str] = None,
+        image_path: Optional[str] = None,
+    ) -> AsyncIterator[StreamChunk]:
+        """Stream tokens from the Anthropic API."""
+        messages = []
+
+        if image_path:
+            import base64
+            from pathlib import Path
+
+            img_data = base64.b64encode(Path(image_path).read_bytes()).decode()
+            messages.append(
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "image",
+                            "source": {
+                                "type": "base64",
+                                "media_type": "image/jpeg",
+                                "data": img_data,
+                            },
+                        },
+                        {"type": "text", "text": prompt},
+                    ],
+                }
+            )
+        else:
+            messages.append({"role": "user", "content": prompt})
+
+        payload: dict = {
+            "model": model,
+            "messages": messages,
+            "max_tokens": 4096,
+            "stream": True,
+        }
+        if system:
+            payload["system"] = system
+
+        headers = {
+            "x-api-key": self.api_key,
+            "anthropic-version": "2023-06-01",
+            "Content-Type": "application/json",
+        }
+
+        async with httpx.AsyncClient(timeout=httpx.Timeout(300.0)) as client:
+            async with client.stream(
+                "POST",
+                f"{self.base_url}/messages",
+                json=payload,
+                headers=headers,
+            ) as response:
+                response.raise_for_status()
+                async for line in response.aiter_lines():
+                    line = line.strip()
+                    if not line or not line.startswith("data: "):
+                        continue
+                    data_str = line[6:]
+                    if not data_str:
+                        continue
+                    try:
+                        data = json.loads(data_str)
+                    except json.JSONDecodeError:
+                        continue
+
+                    event_type = data.get("type", "")
+
+                    if event_type == "content_block_delta":
+                        delta = data.get("delta", {})
+                        text = delta.get("text", "")
+                        yield StreamChunk(text=text, done=False)
+
+                    elif event_type == "message_delta":
+                        usage = data.get("usage", {})
+                        yield StreamChunk(
+                            text="",
+                            done=True,
+                            meta={
+                                "output_tokens": usage.get("output_tokens"),
+                                "stop_reason": data.get("delta", {}).get(
+                                    "stop_reason"
+                                ),
+                            },
+                        )
+
+                    elif event_type == "message_start":
+                        # Capture input token count
+                        msg = data.get("message", {})
+                        usage = msg.get("usage", {})
+                        if usage.get("input_tokens"):
+                            yield StreamChunk(
+                                text="",
+                                done=False,
+                                meta={"input_tokens": usage["input_tokens"]},
+                            )
+
+    async def list_models(self) -> list[dict]:
+        """List available Anthropic models."""
+        # Anthropic doesn't have a list models endpoint, return known models
+        return [
+            {"name": "claude-sonnet-4-20250514", "size": None},
+            {"name": "claude-haiku-4-20250414", "size": None},
+            {"name": "claude-opus-4-20250514", "size": None},
+        ]
+
+    async def check_connection(self) -> bool:
+        """Check if the Anthropic API key is configured."""
+        return bool(self.api_key)

arbiter/core/providers/base.py

@@ -0,0 +1,62 @@
+"""Base provider interface for all LLM backends."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import AsyncIterator, Optional
+
+
+@dataclass
+class StreamChunk:
+    """A single chunk from a streaming LLM response."""
+
+    text: str
+    done: bool = False
+    # Provider-specific metadata (eval_count, eval_duration, etc.)
+    meta: dict = field(default_factory=dict)
+
+
+@dataclass
+class GenerationResult:
+    """Complete result from a model generation."""
+
+    model: str
+    provider: str
+    output: str
+    total_tokens: int
+    eval_duration_ns: Optional[int] = None  # nanoseconds
+    total_duration_ns: Optional[int] = None
+    prompt_tokens: Optional[int] = None
+    # Raw provider response metadata
+    raw_meta: dict = field(default_factory=dict)
+
+
+class LLMProvider(ABC):
+    """Abstract base class for LLM providers."""
+
+    provider_name: str = "base"
+
+    @abstractmethod
+    async def stream_generate(
+        self,
+        model: str,
+        prompt: str,
+        system: Optional[str] = None,
+        image_path: Optional[str] = None,
+    ) -> AsyncIterator[StreamChunk]:
+        """Stream tokens from the model. Yields StreamChunk objects."""
+        ...
+
+    @abstractmethod
+    async def list_models(self) -> list[dict]:
+        """List available models for this provider.
+
+        Returns a list of dicts with at least {"name": str, "size": int|None}.
+        """
+        ...
+
+    @abstractmethod
+    async def check_connection(self) -> bool:
+        """Check if the provider is reachable and configured."""
+        ...

arbiter/core/providers/factory.py

@@ -0,0 +1,79 @@
+"""Factory for creating provider instances from config."""
+
+from __future__ import annotations
+
+from arbiter.core.config import (
+    PROVIDER_ANTHROPIC,
+    PROVIDER_GOOGLE,
+    PROVIDER_OLLAMA,
+    PROVIDER_OPENAI,
+    PROVIDER_OPENAI_COMPAT,
+    ProviderConfig,
+    resolve_model,
+)
+from arbiter.core.providers.anthropic_provider import AnthropicProvider
+from arbiter.core.providers.base import LLMProvider
+from arbiter.core.providers.google_provider import GoogleProvider
+from arbiter.core.providers.ollama import OllamaProvider
+from arbiter.core.providers.openai_provider import OpenAIProvider
+
+
+def create_provider(config: ProviderConfig) -> tuple[LLMProvider, str]:
+    """Create a provider instance and return (provider, clean_model_name).
+
+    Args:
+        config: ProviderConfig from resolve_model()
+
+    Returns:
+        Tuple of (provider_instance, model_name_to_use)
+
+    Raises:
+        ValueError: If API key is required but missing
+    """
+    model = config.extra["model"]
+
+    if config.provider == PROVIDER_OLLAMA:
+        provider = OllamaProvider(base_url=config.base_url or "http://localhost:11434")
+        return provider, model
+
+    if config.provider == PROVIDER_OPENAI:
+        if not config.api_key:
+            raise ValueError(
+                "OpenAI API key required. Set OPENAI_API_KEY environment variable."
+            )
+        provider = OpenAIProvider(api_key=config.api_key)
+        return provider, model
+
+    if config.provider == PROVIDER_OPENAI_COMPAT:
+        if not config.api_key:
+            raise ValueError(
+                "API key required for OpenAI-compatible endpoint. Set OPENAI_API_KEY."
+            )
+        provider = OpenAIProvider(
+            api_key=config.api_key, base_url=config.base_url or ""
+        )
+        return provider, model
+
+    if config.provider == PROVIDER_ANTHROPIC:
+        if not config.api_key:
+            raise ValueError(
+                "Anthropic API key required. Set ANTHROPIC_API_KEY environment variable."
+            )
+        provider = AnthropicProvider(api_key=config.api_key)
+        return provider, model
+
+    if config.provider == PROVIDER_GOOGLE:
+        if not config.api_key:
+            raise ValueError(
+                "Google API key required. Set GOOGLE_API_KEY environment variable."
+            )
+        provider = GoogleProvider(api_key=config.api_key)
+        return provider, model
+
+    raise ValueError(f"Unknown provider: {config.provider}")
+
+
+def provider_from_spec(model_spec: str) -> tuple[LLMProvider, str]:
+    """Convenience: resolve a model spec string directly to (provider, model_name)."""
+    config = resolve_model(model_spec)
+    return create_provider(config)

arbiter/core/providers/google_provider.py

@@ -0,0 +1,126 @@
+"""Google Gemini provider via the Generative Language API."""
+
+from __future__ import annotations
+
+import json
+from typing import AsyncIterator, Optional
+
+import httpx
+
+from arbiter.core.providers.base import LLMProvider, StreamChunk
+
+
+class GoogleProvider(LLMProvider):
+    """Provider for Google Gemini models."""
+
+    provider_name = "google"
+
+    def __init__(self, api_key: str):
+        self.api_key = api_key
+        self.base_url = "https://generativelanguage.googleapis.com/v1beta"
+
+    async def stream_generate(
+        self,
+        model: str,
+        prompt: str,
+        system: Optional[str] = None,
+        image_path: Optional[str] = None,
+    ) -> AsyncIterator[StreamChunk]:
+        """Stream tokens from the Google Gemini API."""
+        contents = []
+
+        if image_path:
+            import base64
+            from pathlib import Path
+
+            img_data = base64.b64encode(Path(image_path).read_bytes()).decode()
+            contents.append(
+                {
+                    "parts": [
+                        {"inline_data": {"mime_type": "image/jpeg", "data": img_data}},
+                        {"text": prompt},
+                    ]
+                }
+            )
+        else:
+            contents.append({"parts": [{"text": prompt}]})
+
+        payload: dict = {"contents": contents}
+        if system:
+            payload["system_instruction"] = {"parts": [{"text": system}]}
+
+        url = (
+            f"{self.base_url}/models/{model}:streamGenerateContent"
+            f"?key={self.api_key}&alt=sse"
+        )
+
+        async with httpx.AsyncClient(timeout=httpx.Timeout(300.0)) as client:
+            async with client.stream("POST", url, json=payload) as response:
+                response.raise_for_status()
+                async for line in response.aiter_lines():
+                    line = line.strip()
+                    if not line or not line.startswith("data: "):
+                        continue
+                    data_str = line[6:]
+                    try:
+                        data = json.loads(data_str)
+                    except json.JSONDecodeError:
+                        continue
+
+                    candidates = data.get("candidates", [])
+                    if not candidates:
+                        continue
+
+                    candidate = candidates[0]
+                    content = candidate.get("content", {})
+                    parts = content.get("parts", [])
+                    text = "".join(p.get("text", "") for p in parts)
+
+                    finish_reason = candidate.get("finishReason")
+                    done = finish_reason is not None and finish_reason != "STOP"
+
+                    chunk = StreamChunk(text=text, done=False)
+
+                    if finish_reason == "STOP":
+                        usage = data.get("usageMetadata", {})
+                        chunk.done = True
+                        chunk.meta = {
+                            "prompt_tokens": usage.get("promptTokenCount"),
+                            "output_tokens": usage.get("candidatesTokenCount"),
+                            "total_tokens": usage.get("totalTokenCount"),
+                            "finish_reason": finish_reason,
+                        }
+
+                    yield chunk
+
+    async def list_models(self) -> list[dict]:
+        """List available Gemini models."""
+        try:
+            url = f"{self.base_url}/models?key={self.api_key}"
+            async with httpx.AsyncClient(timeout=httpx.Timeout(10.0)) as client:
+                resp = await client.get(url)
+                resp.raise_for_status()
+                data = resp.json()
+            return [
+                {
+                    "name": m["name"].replace("models/", ""),
+                    "size": None,
+                    "display_name": m.get("displayName"),
+                }
+                for m in data.get("models", [])
+                if "generateContent" in m.get("supportedGenerationMethods", [])
+            ]
+        except (httpx.HTTPError, KeyError):
+            return []
+
+    async def check_connection(self) -> bool:
+        """Check if the Google API key is configured and valid."""
+        if not self.api_key:
+            return False
+        try:
+            url = f"{self.base_url}/models?key={self.api_key}"
+            async with httpx.AsyncClient(timeout=httpx.Timeout(10.0)) as client:
+                resp = await client.get(url)
+                return resp.status_code == 200
+        except (httpx.ConnectError, httpx.TimeoutException):
+            return False

arbiter/core/providers/ollama.py

@@ -0,0 +1,103 @@
+"""Ollama provider - local model execution."""
+
+from __future__ import annotations
+
+import base64
+import json
+from pathlib import Path
+from typing import AsyncIterator, Optional
+
+import httpx
+
+from arbiter.core.providers.base import GenerationResult, LLMProvider, StreamChunk
+
+
+class OllamaProvider(LLMProvider):
+    """Provider for locally running Ollama models."""
+
+    provider_name = "ollama"
+
+    def __init__(self, base_url: str = "http://localhost:11434"):
+        self.base_url = base_url.rstrip("/")
+
+    async def stream_generate(
+        self,
+        model: str,
+        prompt: str,
+        system: Optional[str] = None,
+        image_path: Optional[str] = None,
+    ) -> AsyncIterator[StreamChunk]:
+        """Stream tokens from an Ollama model."""
+        payload: dict = {"model": model, "prompt": prompt, "stream": True}
+
+        if system:
+            payload["system"] = system
+
+        if image_path:
+            img_data = Path(image_path).read_bytes()
+            payload["images"] = [base64.b64encode(img_data).decode()]
+
+        async with httpx.AsyncClient(timeout=httpx.Timeout(300.0)) as client:
+            async with client.stream(
+                "POST", f"{self.base_url}/api/generate", json=payload
+            ) as response:
+                response.raise_for_status()
+                async for line in response.aiter_lines():
+                    if not line.strip():
+                        continue
+                    data = json.loads(line)
+                    chunk = StreamChunk(
+                        text=data.get("response", ""),
+                        done=data.get("done", False),
+                    )
+                    if chunk.done:
+                        chunk.meta = {
+                            "total_duration": data.get("total_duration"),
+                            "eval_count": data.get("eval_count"),
+                            "eval_duration": data.get("eval_duration"),
+                            "prompt_eval_count": data.get("prompt_eval_count"),
+                            "prompt_eval_duration": data.get("prompt_eval_duration"),
+                        }
+                    yield chunk
+
+    async def list_models(self) -> list[dict]:
+        """List installed Ollama models."""
+        async with httpx.AsyncClient(timeout=httpx.Timeout(10.0)) as client:
+            resp = await client.get(f"{self.base_url}/api/tags")
+            resp.raise_for_status()
+            data = resp.json()
+
+        models = []
+        for m in data.get("models", []):
+            details = m.get("details", {})
+            models.append(
+                {
+                    "name": m["name"],
+                    "size": m.get("size"),
+                    "parameter_size": details.get("parameter_size"),
+                    "quantization": details.get("quantization_level"),
+                    "family": details.get("family"),
+                    "families": details.get("families"),
+                    "format": details.get("format"),
+                    "modified_at": m.get("modified_at"),
+                }
+            )
+        return models
+
+    async def get_model_info(self, model: str) -> dict:
+        """Get detailed info for a specific model."""
+        async with httpx.AsyncClient(timeout=httpx.Timeout(10.0)) as client:
+            resp = await client.post(
+                f"{self.base_url}/api/show", json={"name": model}
+            )
+            resp.raise_for_status()
+            return resp.json()
+
+    async def check_connection(self) -> bool:
+        """Check if Ollama is running."""
+        try:
+            async with httpx.AsyncClient(timeout=httpx.Timeout(5.0)) as client:
+                resp = await client.get(f"{self.base_url}/api/tags")
+                return resp.status_code == 200
+        except (httpx.ConnectError, httpx.TimeoutException):
+            return False

arbiter/core/providers/openai_provider.py

@@ -0,0 +1,120 @@
+"""OpenAI-compatible provider - works with OpenAI, Together, Groq, any OpenAI-compatible API."""
+
+from __future__ import annotations
+
+import json
+from typing import AsyncIterator, Optional
+
+import httpx
+
+from arbiter.core.providers.base import LLMProvider, StreamChunk
+
+
+class OpenAIProvider(LLMProvider):
+    """Provider for OpenAI and any OpenAI-compatible API."""
+
+    provider_name = "openai"
+
+    def __init__(self, api_key: str, base_url: str = "https://api.openai.com/v1"):
+        self.api_key = api_key
+        self.base_url = base_url.rstrip("/")
+
+    async def stream_generate(
+        self,
+        model: str,
+        prompt: str,
+        system: Optional[str] = None,
+        image_path: Optional[str] = None,
+    ) -> AsyncIterator[StreamChunk]:
+        """Stream tokens from an OpenAI-compatible API."""
+        messages = []
+        if system:
+            messages.append({"role": "system", "content": system})
+
+        if image_path:
+            import base64
+            from pathlib import Path
+
+            img_data = base64.b64encode(Path(image_path).read_bytes()).decode()
+            messages.append(
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": prompt},
+                        {
+                            "type": "image_url",
+                            "image_url": {"url": f"data:image/jpeg;base64,{img_data}"},
+                        },
+                    ],
+                }
+            )
+        else:
+            messages.append({"role": "user", "content": prompt})
+
+        payload = {"model": model, "messages": messages, "stream": True}
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+        }
+
+        async with httpx.AsyncClient(timeout=httpx.Timeout(300.0)) as client:
+            async with client.stream(
+                "POST",
+                f"{self.base_url}/chat/completions",
+                json=payload,
+                headers=headers,
+            ) as response:
+                response.raise_for_status()
+                async for line in response.aiter_lines():
+                    line = line.strip()
+                    if not line or not line.startswith("data: "):
+                        continue
+                    data_str = line[6:]
+                    if data_str == "[DONE]":
+                        yield StreamChunk(text="", done=True)
+                        return
+                    data = json.loads(data_str)
+                    choices = data.get("choices", [])
+                    if not choices:
+                        continue
+                    delta = choices[0].get("delta", {})
+                    text = delta.get("content", "")
+                    finish = choices[0].get("finish_reason")
+                    chunk = StreamChunk(text=text, done=finish is not None)
+                    if finish:
+                        chunk.meta = {
+                            "usage": data.get("usage", {}),
+                            "finish_reason": finish,
+                        }
+                    yield chunk
+
+    async def list_models(self) -> list[dict]:
+        """List available models from the API."""
+        headers = {"Authorization": f"Bearer {self.api_key}"}
+        try:
+            async with httpx.AsyncClient(timeout=httpx.Timeout(10.0)) as client:
+                resp = await client.get(
+                    f"{self.base_url}/models", headers=headers
+                )
+                resp.raise_for_status()
+                data = resp.json()
+            return [
+                {"name": m["id"], "size": None, "owned_by": m.get("owned_by")}
+                for m in data.get("data", [])
+            ]
+        except (httpx.HTTPError, KeyError):
+            return []
+
+    async def check_connection(self) -> bool:
+        """Check if the API is reachable and the key is valid."""
+        if not self.api_key:
+            return False
+        headers = {"Authorization": f"Bearer {self.api_key}"}
+        try:
+            async with httpx.AsyncClient(timeout=httpx.Timeout(10.0)) as client:
+                resp = await client.get(
+                    f"{self.base_url}/models", headers=headers
+                )
+                return resp.status_code == 200
+        except (httpx.ConnectError, httpx.TimeoutException):
+            return False