arbiter-cli 0.1.0__py3-none-any.whl

arbiter/core/config.py

@@ -0,0 +1,137 @@
+"""Configuration and defaults for Arbiter."""
+
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Optional
+
+
+ARBITER_DIR = Path.home() / ".arbiter"
+LEADERBOARD_FILE = ARBITER_DIR / "leaderboard.json"
+CONFIG_FILE = ARBITER_DIR / "config.json"
+
+# Ollama defaults
+DEFAULT_OLLAMA_HOST = "http://localhost:11434"
+
+# Provider identifiers
+PROVIDER_OLLAMA = "ollama"
+PROVIDER_OPENAI = "openai"
+PROVIDER_ANTHROPIC = "anthropic"
+PROVIDER_GOOGLE = "google"
+PROVIDER_OPENAI_COMPAT = "openai-compatible"
+
+
+@dataclass
+class ProviderConfig:
+    """Configuration for a single LLM provider."""
+
+    provider: str
+    api_key: Optional[str] = None
+    base_url: Optional[str] = None
+    extra: dict = field(default_factory=dict)
+
+
+def ensure_arbiter_dir() -> Path:
+    """Create ~/.arbiter/ if it doesn't exist."""
+    ARBITER_DIR.mkdir(parents=True, exist_ok=True)
+    return ARBITER_DIR
+
+
+def get_ollama_host() -> str:
+    """Get the Ollama API host from env or default."""
+    return os.environ.get("OLLAMA_HOST", DEFAULT_OLLAMA_HOST)
+
+
+def detect_provider(model_name: str) -> tuple[str, str]:
+    """Detect the provider and clean model name from a model specifier.
+
+    Formats:
+        "gemma4"                     -> (ollama, gemma4)
+        "ollama:gemma4"              -> (ollama, gemma4)
+        "openai:gpt-4o"              -> (openai, gpt-4o)
+        "anthropic:claude-sonnet-4-20250514" -> (anthropic, claude-sonnet-4-20250514)
+        "google:gemini-2.0-flash"    -> (google, gemini-2.0-flash)
+        "http://host:port/v1:model"  -> (openai-compatible, model) with base_url
+    """
+    if ":" in model_name:
+        prefix, _, rest = model_name.partition(":")
+
+        # Check for known providers
+        if prefix in (PROVIDER_OLLAMA, PROVIDER_OPENAI, PROVIDER_ANTHROPIC, PROVIDER_GOOGLE):
+            return prefix, rest
+
+        # Check for URL-based custom endpoint (openai-compatible)
+        if prefix in ("http", "https"):
+            # Format: http://host:port/v1:model_name
+            # Find the last colon that separates URL from model
+            full = model_name
+            last_colon = full.rfind(":")
+            # Walk backwards to find the actual model separator
+            # URLs have colons in protocol and port, so we look for the pattern
+            # where what's after the colon looks like a model name (no slashes)
+            parts = full.rsplit(":", 1)
+            if len(parts) == 2 and "/" not in parts[1]:
+                base_url = parts[0]
+                model = parts[1]
+                return PROVIDER_OPENAI_COMPAT, f"{base_url}||{model}"
+
+        # Ollama model with tag (e.g. "gemma4:latest" or "qwen2.5:7b")
+        return PROVIDER_OLLAMA, model_name
+
+    return PROVIDER_OLLAMA, model_name
+
+
+def get_api_key(provider: str) -> Optional[str]:
+    """Get API key for a provider from environment variables."""
+    env_map = {
+        PROVIDER_OPENAI: "OPENAI_API_KEY",
+        PROVIDER_ANTHROPIC: "ANTHROPIC_API_KEY",
+        PROVIDER_GOOGLE: "GOOGLE_API_KEY",
+        PROVIDER_OPENAI_COMPAT: "OPENAI_API_KEY",
+    }
+    env_var = env_map.get(provider)
+    if env_var:
+        return os.environ.get(env_var)
+    return None
+
+
+def resolve_model(model_spec: str) -> ProviderConfig:
+    """Resolve a model specifier into a full provider config.
+
+    Returns a ProviderConfig with provider name, API key, base URL, and
+    the clean model name stored in extra["model"].
+    """
+    provider, model_name = detect_provider(model_spec)
+
+    base_url = None
+    if provider == PROVIDER_OPENAI_COMPAT and "||" in model_name:
+        base_url, model_name = model_name.split("||", 1)
+    elif provider == PROVIDER_OLLAMA:
+        base_url = get_ollama_host()
+
+    api_key = get_api_key(provider)
+
+    return ProviderConfig(
+        provider=provider,
+        api_key=api_key,
+        base_url=base_url,
+        extra={"model": model_name, "original_spec": model_spec},
+    )
+
+
+def load_config() -> dict:
+    """Load user config from ~/.arbiter/config.json."""
+    if CONFIG_FILE.exists():
+        with open(CONFIG_FILE) as f:
+            return json.load(f)
+    return {}
+
+
+def save_config(config: dict) -> None:
+    """Save user config to ~/.arbiter/config.json."""
+    ensure_arbiter_dir()
+    with open(CONFIG_FILE, "w") as f:
+        json.dump(config, f, indent=2)

arbiter/core/discover.py

@@ -0,0 +1,184 @@
+"""Model discovery -- list all available models across providers."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+from arbiter.core.config import get_api_key, PROVIDER_OPENAI, PROVIDER_ANTHROPIC, PROVIDER_GOOGLE
+from arbiter.core.providers.ollama import OllamaProvider
+from arbiter.core.providers.openai_provider import OpenAIProvider
+from arbiter.core.providers.anthropic_provider import AnthropicProvider
+from arbiter.core.providers.google_provider import GoogleProvider
+
+
+def get_system_memory() -> dict:
+    """Get system memory info in GB."""
+    import psutil
+    mem = psutil.virtual_memory()
+    return {
+        "total_gb": round(mem.total / (1024**3), 1),
+        "available_gb": round(mem.available / (1024**3), 1),
+        "used_gb": round(mem.used / (1024**3), 1),
+        "percent": mem.percent,
+    }
+
+
+def get_max_model_size_gb() -> float:
+    """Get the max model size that can safely run on this machine.
+
+    Rule of thumb: model should use at most 70% of available RAM
+    to leave room for OS, Ollama overhead, and the app itself.
+    """
+    mem = get_system_memory()
+    return round(mem["available_gb"] * 0.70, 1)
+
+
+@dataclass
+class DiscoveredModel:
+    """A model discovered from a provider."""
+
+    name: str
+    provider: str
+    size: Optional[int] = None
+    parameter_size: Optional[str] = None
+    quantization: Optional[str] = None
+    family: Optional[str] = None
+    multimodal: bool = False
+    display_name: Optional[str] = None
+
+    @property
+    def size_gb(self) -> Optional[float]:
+        if self.size:
+            return round(self.size / (1024**3), 1)
+        return None
+
+    @property
+    def fits_in_memory(self) -> bool:
+        """Check if this model can run without swapping."""
+        if self.size_gb is None:
+            return True  # cloud models, no size info
+        max_size = get_max_model_size_gb()
+        return self.size_gb <= max_size
+
+    @property
+    def memory_warning(self) -> Optional[str]:
+        """Return a warning if the model is too large for this machine."""
+        if self.size_gb is None:
+            return None
+        mem = get_system_memory()
+        if self.size_gb > mem["available_gb"]:
+            return f"Model is {self.size_gb}GB but only {mem['available_gb']}GB available. Will be extremely slow."
+        if self.size_gb > mem["available_gb"] * 0.7:
+            return f"Model is {self.size_gb}GB which is tight for {mem['total_gb']}GB RAM. May be slow."
+        return None
+
+    @property
+    def spec(self) -> str:
+        """Return the full model specifier for use with arbiter."""
+        if self.provider == "ollama":
+            return self.name
+        return f"{self.provider}:{self.name}"
+
+
+async def discover_ollama() -> list[DiscoveredModel]:
+    """Discover models from local Ollama installation."""
+    ollama = OllamaProvider()
+    if not await ollama.check_connection():
+        return []
+
+    raw_models = await ollama.list_models()
+    models = []
+    for m in raw_models:
+        families = m.get("families") or []
+        multimodal = any(f in families for f in ["clip", "mllama"])
+        models.append(
+            DiscoveredModel(
+                name=m["name"],
+                provider="ollama",
+                size=m.get("size"),
+                parameter_size=m.get("parameter_size"),
+                quantization=m.get("quantization"),
+                family=m.get("family"),
+                multimodal=multimodal,
+            )
+        )
+    return models
+
+
+async def discover_openai() -> list[DiscoveredModel]:
+    """Discover models from OpenAI (if API key is set)."""
+    api_key = get_api_key(PROVIDER_OPENAI)
+    if not api_key:
+        return []
+
+    provider = OpenAIProvider(api_key=api_key)
+    raw_models = await provider.list_models()
+    return [
+        DiscoveredModel(
+            name=m["name"],
+            provider="openai",
+            display_name=m["name"],
+        )
+        for m in raw_models
+        if any(
+            prefix in m["name"]
+            for prefix in ("gpt-", "o1", "o3", "o4", "chatgpt")
+        )
+    ]
+
+
+async def discover_anthropic() -> list[DiscoveredModel]:
+    """Discover models from Anthropic (if API key is set)."""
+    api_key = get_api_key(PROVIDER_ANTHROPIC)
+    if not api_key:
+        return []
+
+    provider = AnthropicProvider(api_key=api_key)
+    raw_models = await provider.list_models()
+    return [
+        DiscoveredModel(
+            name=m["name"],
+            provider="anthropic",
+            display_name=m["name"],
+            multimodal=True,
+        )
+        for m in raw_models
+    ]
+
+
+async def discover_google() -> list[DiscoveredModel]:
+    """Discover models from Google (if API key is set)."""
+    api_key = get_api_key(PROVIDER_GOOGLE)
+    if not api_key:
+        return []
+
+    provider = GoogleProvider(api_key=api_key)
+    raw_models = await provider.list_models()
+    return [
+        DiscoveredModel(
+            name=m["name"],
+            provider="google",
+            display_name=m.get("display_name", m["name"]),
+        )
+        for m in raw_models
+    ]
+
+
+async def discover_all() -> list[DiscoveredModel]:
+    """Discover all available models across all providers."""
+    import asyncio
+
+    results = await asyncio.gather(
+        discover_ollama(),
+        discover_openai(),
+        discover_anthropic(),
+        discover_google(),
+        return_exceptions=True,
+    )
+
+    all_models = []
+    for result in results:
+        if isinstance(result, list):
+            all_models.extend(result)
+    return all_models

arbiter/core/judge.py

@@ -0,0 +1,193 @@
+"""Auto-judge system -- uses a model to score all outputs."""
+
+from __future__ import annotations
+
+import json
+import re
+from typing import Optional
+
+from arbiter.core.config import resolve_model
+from arbiter.core.metrics import ComparisonResult, ModelMetrics
+from arbiter.core.providers.factory import create_provider
+
+
+JUDGE_SYSTEM_PROMPT = """You are an expert judge evaluating AI model outputs. You will be given a prompt and multiple model responses. Score each response on these dimensions:
+
+1. **Correctness** (1-10): Is the answer factually accurate and free of errors?
+2. **Completeness** (1-10): Does it fully address the prompt?
+3. **Clarity** (1-10): Is it well-organized and easy to understand?
+4. **Code Quality** (1-10): If code is present, is it clean, efficient, and correct? If no code, score based on writing quality.
+
+Respond ONLY with valid JSON in this exact format:
+{
+  "results": [
+    {
+      "model": "<model_name>",
+      "correctness": <1-10>,
+      "completeness": <1-10>,
+      "clarity": <1-10>,
+      "code_quality": <1-10>,
+      "overall": <1-10>,
+      "reasoning": "<brief explanation>"
+    }
+  ],
+  "winner": "<model_name>"
+}"""
+
+
+def _build_judge_prompt(prompt: str, models: list[ModelMetrics]) -> str:
+    """Build the prompt for the judge model."""
+    parts = [f'Original prompt: """{prompt}"""\n']
+
+    for i, m in enumerate(models, 1):
+        output = m.output
+        if output.startswith("[ERROR]"):
+            output = "(Model failed to generate a response)"
+        # Truncate very long outputs for the judge
+        if len(output) > 4000:
+            output = output[:4000] + "\n... (truncated)"
+        parts.append(f'--- Response from {m.model} ---\n"""{output}"""\n')
+
+    parts.append(
+        "Score each response. Remember to respond ONLY with valid JSON."
+    )
+    return "\n".join(parts)
+
+
+def _parse_judge_response(text: str, models: list[ModelMetrics]) -> dict:
+    """Parse the judge's JSON response, handling common formatting issues."""
+    # Try to extract JSON from the response
+    # Sometimes models wrap JSON in markdown code blocks
+    json_match = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", text, re.DOTALL)
+    if json_match:
+        text = json_match.group(1)
+    else:
+        # Try to find raw JSON
+        brace_start = text.find("{")
+        brace_end = text.rfind("}") + 1
+        if brace_start >= 0 and brace_end > brace_start:
+            text = text[brace_start:brace_end]
+
+    try:
+        data = json.loads(text)
+    except json.JSONDecodeError:
+        # Fallback: return equal scores
+        model_names = [m.model for m in models]
+        return {
+            "results": [
+                {
+                    "model": name,
+                    "correctness": 5,
+                    "completeness": 5,
+                    "clarity": 5,
+                    "code_quality": 5,
+                    "overall": 5,
+                    "reasoning": "Judge response could not be parsed",
+                }
+                for name in model_names
+            ],
+            "winner": model_names[0] if model_names else None,
+        }
+
+    return data
+
+
+async def judge_comparison(
+    comparison: ComparisonResult,
+    judge_model: str = "auto",
+) -> ComparisonResult:
+    """Score all model outputs using a judge model.
+
+    Args:
+        comparison: The comparison result to judge
+        judge_model: Model spec for the judge. "auto" picks the best available.
+
+    Returns:
+        The same ComparisonResult with quality_scores and winner filled in.
+    """
+    # Filter out models that errored
+    valid_models = [m for m in comparison.models if not m.output.startswith("[ERROR]")]
+    if not valid_models:
+        return comparison
+
+    # Resolve judge model
+    if judge_model == "auto":
+        judge_model = await _pick_judge_model(comparison.models)
+
+    config = resolve_model(judge_model)
+    provider, model_name = create_provider(config)
+
+    judge_prompt = _build_judge_prompt(comparison.prompt, valid_models)
+
+    # Generate judge response with timeout and token limit
+    import time
+    output_parts = []
+    token_count = 0
+    start = time.perf_counter()
+    async for chunk in provider.stream_generate(
+        model=model_name,
+        prompt=judge_prompt,
+        system=JUDGE_SYSTEM_PROMPT,
+    ):
+        output_parts.append(chunk.text)
+        token_count += 1
+        if chunk.done:
+            break
+        # Hard limits to prevent hanging
+        if token_count >= 1500:
+            break
+        if time.perf_counter() - start > 90:
+            break
+
+    judge_text = "".join(output_parts)
+    parsed = _parse_judge_response(judge_text, valid_models)
+
+    # Apply scores to metrics
+    results_by_model = {}
+    for r in parsed.get("results", []):
+        results_by_model[r.get("model", "")] = r
+
+    for m in comparison.models:
+        scores = results_by_model.get(m.model, {})
+        if scores:
+            m.quality_scores = {
+                "correctness": scores.get("correctness", 5),
+                "completeness": scores.get("completeness", 5),
+                "clarity": scores.get("clarity", 5),
+                "code_quality": scores.get("code_quality", 5),
+            }
+            m.overall_score = scores.get("overall", 5)
+
+    # Judge no longer picks the winner -- composite scoring does that
+    comparison.judge_model = judge_model
+
+    return comparison
+
+
+async def _pick_judge_model(models: list[ModelMetrics]) -> str:
+    """Auto-select the best available judge model.
+
+    Prefers large local models, falls back to the first available model
+    that isn't in the comparison set.
+    """
+    from arbiter.core.providers.ollama import OllamaProvider
+
+    compared_models = {m.model for m in models}
+
+    # Try to find a local model not in the comparison
+    try:
+        ollama = OllamaProvider()
+        if await ollama.check_connection():
+            available = await ollama.list_models()
+            # Sort by size (largest first) to pick the best judge
+            available.sort(key=lambda x: x.get("size", 0) or 0, reverse=True)
+            for m in available:
+                name = m["name"]
+                # Skip models that are being compared
+                if name not in compared_models and name.split(":")[0] not in compared_models:
+                    return name
+    except Exception:
+        pass
+
+    # Fallback: use the first compared model as judge (not ideal but works)
+    return models[0].model if models else "llama3.2"