Add multi-LLM provider support (OpenAI, Ollama, HuggingFace) with unified interface for seamless switching

smolagents/llm_provider.py

@@ -0,0 +1,334 @@
+"""Multi-LLM Provider Support
+
+Unified interface for OpenAI, Ollama, and HuggingFace Inference API.
+Enables seamless switching between LLM backends with parameterized configuration.
+"""
+
+import os
+from abc import ABC, abstractmethod
+from typing import Optional, List, Dict, Any
+from dataclasses import dataclass
+
+
+@dataclass
+class LLMResponse:
+    """Standardized LLM response"""
+    text: str
+    model: str
+    provider: str
+    tokens_used: Optional[int] = None
+    metadata: Optional[Dict[str, Any]] = None
+
+
+class LLMProvider(ABC):
+    """Abstract base class for LLM providers"""
+    
+    def __init__(self, model_name: str, **kwargs):
+        self.model_name = model_name
+        self.config = kwargs
+    
+    @abstractmethod
+    def generate(self, prompt: str, **kwargs) -> LLMResponse:
+        """Generate text from prompt"""
+        pass
+    
+    @abstractmethod
+    def score_response(self, prompt: str, response: str) -> float:
+        """Score quality of a response (0-1)"""
+        pass
+
+
+class OpenAIProvider(LLMProvider):
+    """OpenAI API provider"""
+    
+    def __init__(self, model_name: str = "gpt-4", api_key: Optional[str] = None, **kwargs):
+        super().__init__(model_name, **kwargs)
+        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
+        
+        try:
+            from openai import OpenAI
+            self.client = OpenAI(api_key=self.api_key)
+            self.available = True
+        except ImportError:
+            print("Warning: openai package not installed. Install with: pip install openai")
+            self.available = False
+        except Exception as e:
+            print(f"Warning: Could not initialize OpenAI client: {e}")
+            self.available = False
+    
+    def generate(self, prompt: str, temperature: float = 0.7, max_tokens: int = 1000, **kwargs) -> LLMResponse:
+        if not self.available:
+            return LLMResponse(
+                text="OpenAI provider not available. Please install openai package and set API key.",
+                model=self.model_name,
+                provider="openai"
+            )
+        
+        try:
+            response = self.client.chat.completions.create(
+                model=self.model_name,
+                messages=[{"role": "user", "content": prompt}],
+                temperature=temperature,
+                max_tokens=max_tokens,
+                **kwargs
+            )
+            
+            return LLMResponse(
+                text=response.choices[0].message.content,
+                model=self.model_name,
+                provider="openai",
+                tokens_used=response.usage.total_tokens if hasattr(response, 'usage') else None,
+                metadata={"finish_reason": response.choices[0].finish_reason}
+            )
+        except Exception as e:
+            return LLMResponse(
+                text=f"Error generating response: {str(e)}",
+                model=self.model_name,
+                provider="openai"
+            )
+    
+    def score_response(self, prompt: str, response: str) -> float:
+        """Score response quality using a simple prompt"""
+        if not self.available:
+            return 0.5
+        
+        try:
+            score_prompt = f"""Rate the quality of this response on a scale of 0.0 to 1.0.
+
+Prompt: {prompt}
+
+Response: {response}
+
+Provide only a number between 0.0 and 1.0:"""
+            
+            result = self.generate(score_prompt, temperature=0.1, max_tokens=10)
+            
+            # Extract number from response
+            try:
+                score = float(result.text.strip())
+                return max(0.0, min(1.0, score))
+            except ValueError:
+                return 0.5
+        except Exception:
+            return 0.5
+
+
+class OllamaProvider(LLMProvider):
+    """Ollama local LLM provider"""
+    
+    def __init__(self, model_name: str = "llama2", base_url: str = "http://localhost:11434", **kwargs):
+        super().__init__(model_name, **kwargs)
+        self.base_url = base_url
+        
+        try:
+            import requests
+            self.requests = requests
+            # Test connection
+            response = requests.get(f"{base_url}/api/tags", timeout=2)
+            self.available = response.status_code == 200
+        except ImportError:
+            print("Warning: requests package not installed. Install with: pip install requests")
+            self.available = False
+        except Exception as e:
+            print(f"Warning: Could not connect to Ollama at {base_url}: {e}")
+            self.available = False
+    
+    def generate(self, prompt: str, temperature: float = 0.7, max_tokens: int = 1000, **kwargs) -> LLMResponse:
+        if not self.available:
+            return LLMResponse(
+                text="Ollama provider not available. Please ensure Ollama is running at " + self.base_url,
+                model=self.model_name,
+                provider="ollama"
+            )
+        
+        try:
+            response = self.requests.post(
+                f"{self.base_url}/api/generate",
+                json={
+                    "model": self.model_name,
+                    "prompt": prompt,
+                    "temperature": temperature,
+                    "num_predict": max_tokens,
+                    "stream": False
+                },
+                timeout=60
+            )
+            
+            if response.status_code == 200:
+                data = response.json()
+                return LLMResponse(
+                    text=data.get("response", ""),
+                    model=self.model_name,
+                    provider="ollama",
+                    metadata=data
+                )
+            else:
+                return LLMResponse(
+                    text=f"Ollama error: {response.status_code}",
+                    model=self.model_name,
+                    provider="ollama"
+                )
+        except Exception as e:
+            return LLMResponse(
+                text=f"Error generating response: {str(e)}",
+                model=self.model_name,
+                provider="ollama"
+            )
+    
+    def score_response(self, prompt: str, response: str) -> float:
+        """Simple heuristic scoring for Ollama"""
+        # Length-based heuristic (can be improved)
+        if len(response) < 10:
+            return 0.2
+        elif len(response) < 50:
+            return 0.5
+        elif len(response) < 200:
+            return 0.7
+        else:
+            return 0.9
+
+
+class HuggingFaceProvider(LLMProvider):
+    """HuggingFace Inference API provider"""
+    
+    def __init__(self, model_name: str = "HuggingFaceH4/zephyr-7b-beta", api_token: Optional[str] = None, **kwargs):
+        super().__init__(model_name, **kwargs)
+        self.api_token = api_token or os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN")
+        
+        try:
+            from huggingface_hub import InferenceClient
+            self.client = InferenceClient(token=self.api_token)
+            self.available = True
+        except ImportError:
+            print("Warning: huggingface_hub not installed. Install with: pip install huggingface_hub")
+            self.available = False
+        except Exception as e:
+            print(f"Warning: Could not initialize HuggingFace client: {e}")
+            self.available = False
+    
+    def generate(self, prompt: str, temperature: float = 0.7, max_tokens: int = 1000, **kwargs) -> LLMResponse:
+        if not self.available:
+            return LLMResponse(
+                text="HuggingFace provider not available. Please install huggingface_hub.",
+                model=self.model_name,
+                provider="huggingface"
+            )
+        
+        try:
+            response = self.client.text_generation(
+                prompt,
+                model=self.model_name,
+                max_new_tokens=max_tokens,
+                temperature=temperature,
+                **kwargs
+            )
+            
+            return LLMResponse(
+                text=response,
+                model=self.model_name,
+                provider="huggingface"
+            )
+        except Exception as e:
+            return LLMResponse(
+                text=f"Error generating response: {str(e)}",
+                model=self.model_name,
+                provider="huggingface"
+            )
+    
+    def score_response(self, prompt: str, response: str) -> float:
+        """Simple heuristic scoring"""
+        # Basic quality metrics
+        score = 0.5
+        
+        # Length check
+        if 50 <= len(response) <= 2000:
+            score += 0.2
+        
+        # Contains relevant terms from prompt
+        prompt_words = set(prompt.lower().split())
+        response_words = set(response.lower().split())
+        overlap = len(prompt_words & response_words) / max(len(prompt_words), 1)
+        score += overlap * 0.3
+        
+        return min(1.0, score)
+
+
+class LLMManager:
+    """Unified LLM manager for seamless provider switching"""
+    
+    def __init__(self, default_provider: str = "huggingface"):
+        self.providers: Dict[str, LLMProvider] = {}
+        self.default_provider = default_provider
+        self._initialize_providers()
+    
+    def _initialize_providers(self):
+        """Initialize all available providers"""
+        # OpenAI
+        openai_model = os.getenv("OPENAI_MODEL", "gpt-4")
+        self.providers["openai"] = OpenAIProvider(model_name=openai_model)
+        
+        # Ollama
+        ollama_model = os.getenv("OLLAMA_MODEL", "llama2")
+        ollama_url = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
+        self.providers["ollama"] = OllamaProvider(model_name=ollama_model, base_url=ollama_url)
+        
+        # HuggingFace
+        hf_model = os.getenv("HF_MODEL", "HuggingFaceH4/zephyr-7b-beta")
+        self.providers["huggingface"] = HuggingFaceProvider(model_name=hf_model)
+    
+    def generate(
+        self,
+        prompt: str,
+        provider: Optional[str] = None,
+        **kwargs
+    ) -> LLMResponse:
+        """Generate text using specified or default provider"""
+        provider_name = provider or self.default_provider
+        
+        if provider_name not in self.providers:
+            return LLMResponse(
+                text=f"Provider '{provider_name}' not found. Available: {list(self.providers.keys())}",
+                model="unknown",
+                provider="unknown"
+            )
+        
+        return self.providers[provider_name].generate(prompt, **kwargs)
+    
+    def score_response(self, prompt: str, response: str, provider: Optional[str] = None) -> float:
+        """Score response quality"""
+        provider_name = provider or self.default_provider
+        
+        if provider_name not in self.providers:
+            return 0.5
+        
+        return self.providers[provider_name].score_response(prompt, response)
+    
+    def get_available_providers(self) -> List[str]:
+        """Get list of available providers"""
+        return [
+            name for name, provider in self.providers.items()
+            if getattr(provider, 'available', True)
+        ]
+    
+    def set_default_provider(self, provider: str):
+        """Set default provider"""
+        if provider in self.providers:
+            self.default_provider = provider
+        else:
+            raise ValueError(f"Provider '{provider}' not found")
+
+
+# Singleton instance
+_llm_manager = None
+
+def get_llm_manager(default_provider: Optional[str] = None) -> LLMManager:
+    """Get or create LLM manager singleton"""
+    global _llm_manager
+    
+    if _llm_manager is None:
+        provider = default_provider or os.getenv("LLM_PROVIDER", "huggingface")
+        _llm_manager = LLMManager(default_provider=provider)
+    elif default_provider:
+        _llm_manager.set_default_provider(default_provider)
+    
+    return _llm_manager