Spaces:
Sleeping
Sleeping
Add multi-LLM provider support (OpenAI, Ollama, HuggingFace) with unified interface for seamless switching
93eb4a4 verified | """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 | |
| 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 | |
| def generate(self, prompt: str, **kwargs) -> LLMResponse: | |
| """Generate text from prompt""" | |
| pass | |
| 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 | |