Add llm_client.py — unified multi-provider LLM client

multeclaw/llm_client.py

@@ -0,0 +1,381 @@
+"""
+Multeclaw LLM Client — unified interface across OpenAI, Anthropic, HuggingFace, Groq, Ollama.
+Uses native SDKs for maximum control, with LiteLLM as fallback router.
+"""
+
+import os
+import json
+import time
+import traceback
+from typing import Generator, Optional, Any
+from dataclasses import dataclass
+
+from multeclaw.config import Provider, ModelDef, MODEL_REGISTRY
+
+
+@dataclass
+class LLMResponse:
+    """Standardized response across all providers."""
+    content: str
+    model: str
+    provider: str
+    finish_reason: str = "stop"
+    input_tokens: int = 0
+    output_tokens: int = 0
+    latency_ms: float = 0.0
+    error: Optional[str] = None
+
+    @property
+    def total_tokens(self) -> int:
+        return self.input_tokens + self.output_tokens
+
+
+class MultiModelClient:
+    """
+    Unified LLM client supporting multiple providers.
+    Handles streaming, error recovery, and provider-specific API differences.
+    """
+
+    def __init__(self):
+        self._clients: dict[str, Any] = {}
+        self._api_keys: dict[str, str] = {}
+        self._ollama_url: str = "http://localhost:11434"
+        self._load_env_keys()
+
+    # ─── Key Management ────────────────────────────────────────────────────
+    def _load_env_keys(self):
+        """Load API keys from environment variables."""
+        mappings = {
+            "openai": "OPENAI_API_KEY",
+            "anthropic": "ANTHROPIC_API_KEY",
+            "huggingface": "HF_TOKEN",
+            "groq": "GROQ_API_KEY",
+        }
+        for provider, env_var in mappings.items():
+            key = os.environ.get(env_var, "")
+            if key:
+                self._api_keys[provider] = key
+
+    def set_api_key(self, provider: str, key: str):
+        """Set an API key for a provider, re-initializing its client."""
+        self._api_keys[provider] = key
+        self._clients.pop(provider, None)  # Force re-init
+
+    def set_ollama_url(self, url: str):
+        self._ollama_url = url
+        self._clients.pop("ollama", None)
+
+    def get_available_models(self) -> list[str]:
+        """Return model names that have valid API keys configured."""
+        available = []
+        for name, model_def in MODEL_REGISTRY.items():
+            provider = model_def.provider.value
+            if provider == "ollama":
+                available.append(name)  # Always show local models
+            elif provider in self._api_keys and self._api_keys[provider]:
+                available.append(name)
+        return available
+
+    def check_connections(self) -> dict[str, dict]:
+        """Test connectivity for all configured providers."""
+        results = {}
+        for provider_name, key in self._api_keys.items():
+            try:
+                if provider_name == "openai":
+                    import openai
+                    c = openai.OpenAI(api_key=key, timeout=10)
+                    c.models.list()
+                    results[provider_name] = {"status": "✅ Connected", "models": "Available"}
+                elif provider_name == "anthropic":
+                    results[provider_name] = {"status": "✅ Key Set", "models": "Available"}
+                elif provider_name == "huggingface":
+                    from huggingface_hub import InferenceClient
+                    c = InferenceClient(api_key=key, timeout=10)
+                    results[provider_name] = {"status": "✅ Key Set", "models": "Available"}
+                elif provider_name == "groq":
+                    results[provider_name] = {"status": "✅ Key Set", "models": "Available"}
+            except Exception as e:
+                results[provider_name] = {"status": f"❌ Error: {str(e)[:80]}", "models": "Unavailable"}
+
+        # Check Ollama
+        try:
+            import httpx
+            r = httpx.get(f"{self._ollama_url}/api/tags", timeout=5)
+            if r.status_code == 200:
+                models = [m["name"] for m in r.json().get("models", [])]
+                results["ollama"] = {"status": "✅ Running", "models": ", ".join(models[:5]) or "None"}
+            else:
+                results["ollama"] = {"status": "⚠️ Responded but error", "models": "Unknown"}
+        except Exception:
+            results["ollama"] = {"status": "⚪ Not running (optional)", "models": "N/A"}
+
+        return results
+
+    # ─── Client Initialization ─────────────────────────────────────────────
+    def _get_openai_client(self):
+        if "openai" not in self._clients:
+            import openai
+            self._clients["openai"] = openai.OpenAI(
+                api_key=self._api_keys.get("openai", ""),
+                timeout=120,
+            )
+        return self._clients["openai"]
+
+    def _get_anthropic_client(self):
+        if "anthropic" not in self._clients:
+            import anthropic
+            self._clients["anthropic"] = anthropic.Anthropic(
+                api_key=self._api_keys.get("anthropic", ""),
+                timeout=120,
+            )
+        return self._clients["anthropic"]
+
+    def _get_hf_client(self):
+        if "huggingface" not in self._clients:
+            from huggingface_hub import InferenceClient
+            self._clients["huggingface"] = InferenceClient(
+                provider="novita",
+                api_key=self._api_keys.get("huggingface", ""),
+                timeout=120,
+            )
+        return self._clients["huggingface"]
+
+    # ─── Completion (Non-streaming) ────────────────────────────────────────
+    def complete(
+        self,
+        model_name: str,
+        messages: list[dict],
+        system_prompt: str = "",
+        temperature: float = 0.7,
+        max_tokens: int = 4096,
+        tools: Optional[list] = None,
+    ) -> LLMResponse:
+        """
+        Send a completion request to the appropriate provider.
+        Returns a standardized LLMResponse.
+        """
+        if model_name not in MODEL_REGISTRY:
+            return LLMResponse(content="", model=model_name, provider="unknown",
+                               error=f"Unknown model: {model_name}")
+
+        model_def = MODEL_REGISTRY[model_name]
+        provider = model_def.provider
+        start = time.time()
+
+        try:
+            if provider == Provider.OPENAI:
+                return self._complete_openai(model_def, messages, system_prompt, temperature, max_tokens, tools, start)
+            elif provider == Provider.ANTHROPIC:
+                return self._complete_anthropic(model_def, messages, system_prompt, temperature, max_tokens, tools, start)
+            elif provider == Provider.HUGGINGFACE:
+                return self._complete_hf(model_def, messages, system_prompt, temperature, max_tokens, start)
+            elif provider == Provider.GROQ:
+                return self._complete_groq(model_def, messages, system_prompt, temperature, max_tokens, start)
+            elif provider == Provider.OLLAMA:
+                return self._complete_ollama(model_def, messages, system_prompt, temperature, max_tokens, start)
+            else:
+                return LLMResponse(content="", model=model_name, provider=provider.value,
+                                   error=f"Unsupported provider: {provider}")
+        except Exception as e:
+            return LLMResponse(
+                content="", model=model_name, provider=provider.value,
+                error=f"{type(e).__name__}: {str(e)}",
+                latency_ms=(time.time() - start) * 1000,
+            )
+
+    def _complete_openai(self, model_def, messages, system_prompt, temperature, max_tokens, tools, start):
+        client = self._get_openai_client()
+        msgs = self._build_openai_messages(messages, system_prompt)
+        kwargs = dict(model=model_def.model_id, messages=msgs, temperature=temperature, max_tokens=max_tokens)
+        if tools:
+            kwargs["tools"] = [{"type": "function", "function": t} for t in tools]
+            kwargs["tool_choice"] = "auto"
+        resp = client.chat.completions.create(**kwargs)
+        choice = resp.choices[0]
+        content = choice.message.content or ""
+        # Handle tool calls
+        if choice.message.tool_calls:
+            tool_calls = [{"name": tc.function.name, "arguments": tc.function.arguments} for tc in choice.message.tool_calls]
+            content = json.dumps({"tool_calls": tool_calls}, indent=2)
+        return LLMResponse(
+            content=content, model=model_def.model_id, provider="openai",
+            finish_reason=choice.finish_reason or "stop",
+            input_tokens=resp.usage.prompt_tokens if resp.usage else 0,
+            output_tokens=resp.usage.completion_tokens if resp.usage else 0,
+            latency_ms=(time.time() - start) * 1000,
+        )
+
+    def _complete_anthropic(self, model_def, messages, system_prompt, temperature, max_tokens, tools, start):
+        client = self._get_anthropic_client()
+        # Anthropic: system is a top-level param, NOT in messages
+        filtered = [m for m in messages if m.get("role") != "system"]
+        kwargs = dict(model=model_def.model_id, messages=filtered, max_tokens=max_tokens, temperature=temperature)
+        if system_prompt:
+            kwargs["system"] = system_prompt
+        if tools:
+            kwargs["tools"] = [{"name": t["name"], "description": t["description"], "input_schema": t["parameters"]} for t in tools]
+        resp = client.messages.create(**kwargs)
+        content = ""
+        for block in resp.content:
+            if hasattr(block, "text"):
+                content += block.text
+            elif block.type == "tool_use":
+                content += json.dumps({"tool_use": {"name": block.name, "input": block.input, "id": block.id}}, indent=2)
+        return LLMResponse(
+            content=content, model=model_def.model_id, provider="anthropic",
+            finish_reason=resp.stop_reason or "end_turn",
+            input_tokens=resp.usage.input_tokens if resp.usage else 0,
+            output_tokens=resp.usage.output_tokens if resp.usage else 0,
+            latency_ms=(time.time() - start) * 1000,
+        )
+
+    def _complete_hf(self, model_def, messages, system_prompt, temperature, max_tokens, start):
+        client = self._get_hf_client()
+        msgs = self._build_openai_messages(messages, system_prompt)
+        resp = client.chat_completion(model=model_def.model_id, messages=msgs, max_tokens=max_tokens, temperature=max(temperature, 0.01))
+        content = resp.choices[0].message.content or ""
+        return LLMResponse(
+            content=content, model=model_def.model_id, provider="huggingface",
+            finish_reason=resp.choices[0].finish_reason or "stop",
+            input_tokens=resp.usage.prompt_tokens if resp.usage else 0,
+            output_tokens=resp.usage.completion_tokens if resp.usage else 0,
+            latency_ms=(time.time() - start) * 1000,
+        )
+
+    def _complete_groq(self, model_def, messages, system_prompt, temperature, max_tokens, start):
+        """Groq uses OpenAI-compatible API."""
+        import openai
+        client = openai.OpenAI(
+            api_key=self._api_keys.get("groq", ""),
+            base_url="https://api.groq.com/openai/v1",
+            timeout=60,
+        )
+        msgs = self._build_openai_messages(messages, system_prompt)
+        resp = client.chat.completions.create(model=model_def.model_id, messages=msgs, temperature=temperature, max_tokens=max_tokens)
+        choice = resp.choices[0]
+        return LLMResponse(
+            content=choice.message.content or "", model=model_def.model_id, provider="groq",
+            finish_reason=choice.finish_reason or "stop",
+            input_tokens=resp.usage.prompt_tokens if resp.usage else 0,
+            output_tokens=resp.usage.completion_tokens if resp.usage else 0,
+            latency_ms=(time.time() - start) * 1000,
+        )
+
+    def _complete_ollama(self, model_def, messages, system_prompt, temperature, max_tokens, start):
+        """Ollama uses OpenAI-compatible API."""
+        import openai
+        client = openai.OpenAI(
+            api_key="ollama",
+            base_url=f"{self._ollama_url}/v1",
+            timeout=120,
+        )
+        msgs = self._build_openai_messages(messages, system_prompt)
+        resp = client.chat.completions.create(model=model_def.model_id, messages=msgs, temperature=temperature, max_tokens=max_tokens)
+        choice = resp.choices[0]
+        return LLMResponse(
+            content=choice.message.content or "", model=model_def.model_id, provider="ollama",
+            finish_reason=choice.finish_reason or "stop",
+            latency_ms=(time.time() - start) * 1000,
+        )
+
+    # ─── Streaming Completion ──────────────────────────────────────────────
+    def stream(
+        self,
+        model_name: str,
+        messages: list[dict],
+        system_prompt: str = "",
+        temperature: float = 0.7,
+        max_tokens: int = 4096,
+    ) -> Generator[str, None, None]:
+        """
+        Stream a completion. Yields partial text chunks.
+        Handles provider-specific streaming differences.
+        """
+        if model_name not in MODEL_REGISTRY:
+            yield f"❌ Unknown model: {model_name}"
+            return
+
+        model_def = MODEL_REGISTRY[model_name]
+        provider = model_def.provider
+
+        try:
+            if provider == Provider.OPENAI:
+                yield from self._stream_openai(model_def, messages, system_prompt, temperature, max_tokens)
+            elif provider == Provider.ANTHROPIC:
+                yield from self._stream_anthropic(model_def, messages, system_prompt, temperature, max_tokens)
+            elif provider == Provider.HUGGINGFACE:
+                yield from self._stream_hf(model_def, messages, system_prompt, temperature, max_tokens)
+            elif provider == Provider.GROQ:
+                yield from self._stream_groq(model_def, messages, system_prompt, temperature, max_tokens)
+            elif provider == Provider.OLLAMA:
+                yield from self._stream_ollama(model_def, messages, system_prompt, temperature, max_tokens)
+            else:
+                yield f"❌ Unsupported provider for streaming: {provider}"
+        except Exception as e:
+            yield f"\n\n❌ **Streaming Error** ({type(e).__name__}): {str(e)}"
+
+    def _stream_openai(self, model_def, messages, system_prompt, temperature, max_tokens):
+        client = self._get_openai_client()
+        msgs = self._build_openai_messages(messages, system_prompt)
+        stream = client.chat.completions.create(
+            model=model_def.model_id, messages=msgs, temperature=temperature,
+            max_tokens=max_tokens, stream=True,
+        )
+        for chunk in stream:
+            delta = chunk.choices[0].delta.content
+            if delta:
+                yield delta
+
+    def _stream_anthropic(self, model_def, messages, system_prompt, temperature, max_tokens):
+        client = self._get_anthropic_client()
+        filtered = [m for m in messages if m.get("role") != "system"]
+        kwargs = dict(model=model_def.model_id, messages=filtered, max_tokens=max_tokens, temperature=temperature)
+        if system_prompt:
+            kwargs["system"] = system_prompt
+        with client.messages.stream(**kwargs) as stream:
+            for text in stream.text_stream:
+                yield text
+
+    def _stream_hf(self, model_def, messages, system_prompt, temperature, max_tokens):
+        client = self._get_hf_client()
+        msgs = self._build_openai_messages(messages, system_prompt)
+        stream = client.chat_completion(
+            model=model_def.model_id, messages=msgs, max_tokens=max_tokens,
+            temperature=max(temperature, 0.01), stream=True,
+        )
+        for chunk in stream:
+            delta = chunk.choices[0].delta.content
+            if delta:
+                yield delta
+
+    def _stream_groq(self, model_def, messages, system_prompt, temperature, max_tokens):
+        import openai
+        client = openai.OpenAI(api_key=self._api_keys.get("groq", ""), base_url="https://api.groq.com/openai/v1", timeout=60)
+        msgs = self._build_openai_messages(messages, system_prompt)
+        stream = client.chat.completions.create(model=model_def.model_id, messages=msgs, temperature=temperature, max_tokens=max_tokens, stream=True)
+        for chunk in stream:
+            delta = chunk.choices[0].delta.content
+            if delta:
+                yield delta
+
+    def _stream_ollama(self, model_def, messages, system_prompt, temperature, max_tokens):
+        import openai
+        client = openai.OpenAI(api_key="ollama", base_url=f"{self._ollama_url}/v1", timeout=120)
+        msgs = self._build_openai_messages(messages, system_prompt)
+        stream = client.chat.completions.create(model=model_def.model_id, messages=msgs, temperature=temperature, max_tokens=max_tokens, stream=True)
+        for chunk in stream:
+            delta = chunk.choices[0].delta.content
+            if delta:
+                yield delta
+
+    # ─── Helpers ───────────────────────────────────────────────────────────
+    @staticmethod
+    def _build_openai_messages(messages: list[dict], system_prompt: str = "") -> list[dict]:
+        """Build OpenAI-format message list with system prompt prepended."""
+        msgs = []
+        if system_prompt:
+            msgs.append({"role": "system", "content": system_prompt})
+        for m in messages:
+            if m.get("role") != "system":  # Avoid duplicate system messages
+                msgs.append({"role": m["role"], "content": m["content"]})
+        return msgs