Refactor LLM module to support multiple providers with unified interface

core/llm.py

@@ -1,57 +1,175 @@
+import json
+import time
 import requests
+from abc import ABC, abstractmethod
+from typing import Union, Generator
 import openai
 from utils.config import config
 
-class LLMClient:
-    def __init__(self, provider="ollama", model_name=None):
-        self.provider = provider
-        self.model_name = model_name or config.local_model_name
+class LLMProvider(ABC):
+    """Abstract base class for all LLM providers"""
+    
+    def __init__(self, model_name: str, timeout: int = 30, retries: int = 3):
+        self.model_name = model_name
+        self.timeout = timeout
+        self.retries = retries
+
+    @abstractmethod
+    def generate(self, prompt: str, max_tokens: int = 500, stream: bool = False) -> Union[str, Generator[str, None, None]]:
+        """Generate text completion - must be implemented by subclasses"""
+        pass
+
+    def _retry_request(self, func, *args, **kwargs):
+        """Generic retry wrapper with exponential backoff"""
+        last_exception = None
+        for attempt in range(self.retries + 1):
+            try:
+                return func(*args, **kwargs)
+            except Exception as e:
+                last_exception = e
+                if attempt < self.retries:
+                    time.sleep(1 * (2 ** attempt))  # Exponential backoff
+                continue
+        raise last_exception
+
 
-        # Set up OpenAI client for Hugging Face endpoint
-        self.hf_client = openai.OpenAI(
+class OllamaProvider(LLMProvider):
+    def __init__(self, model_name: str, host: str = None, timeout: int = 30, retries: int = 3):
+        super().__init__(model_name, timeout, retries)
+        self.host = host or config.ollama_host
+        self.headers = {
+            "ngrok-skip-browser-warning": "true",
+            "User-Agent": "AI-Life-Coach"
+        }
+
+    def generate(self, prompt: str, max_tokens: int = 500, stream: bool = False) -> Union[str, Generator[str, None, None]]:
+        def _make_request():
+            url = f"{self.host}/api/generate"
+            payload = {
+                "model": self.model_name,
+                "prompt": prompt,
+                "stream": stream,
+                "options": {
+                    "num_predict": max_tokens
+                }
+            }
+            
+            response = requests.post(
+                url,
+                json=payload,
+                headers=self.headers,
+                timeout=self.timeout,
+                stream=stream
+            )
+            
+            if response.status_code != 200:
+                raise Exception(f"Ollama API error: {response.text}")
+                
+            if stream:
+                def stream_response():
+                    for line in response.iter_lines():
+                        if line:
+                            try:
+                                data = json.loads(line.decode('utf-8'))
+                                if 'response' in data:
+                                    yield data['response']
+                            except:
+                                continue
+                return stream_response()
+            else:
+                return response.json()["response"]
+                
+        return self._retry_request(_make_request)
+
+
+class HuggingFaceProvider(LLMProvider):
+    def __init__(self, model_name: str, timeout: int = 30, retries: int = 3):
+        super().__init__(model_name, timeout, retries)
+        self.client = openai.OpenAI(
             base_url=config.hf_api_url,
             api_key=config.hf_token
         )
 
-    def generate(self, prompt, max_tokens=8192, stream=True):
-        if self.provider == "ollama":
-            return self._generate_ollama(prompt, max_tokens, stream)
-        elif self.provider == "huggingface":
-            return self._generate_hf(prompt, max_tokens, stream)
-        else:
-            raise ValueError(f"Unsupported provider: {self.provider}")
-
-    def _generate_ollama(self, prompt, max_tokens, stream):
-        url = f"{config.ollama_host}/api/generate"
-        payload = {
-            "model": self.model_name,
-            "prompt": prompt,
-            "stream": stream
-        }
+    def generate(self, prompt: str, max_tokens: int = 500, stream: bool = False) -> Union[str, Generator[str, None, None]]:
+        def _make_request():
+            response = self.client.chat.completions.create(
+                model=self.model_name,
+                messages=[{"role": "user", "content": prompt}],
+                max_tokens=max_tokens,
+                stream=stream
+            )
+            
+            if stream:
+                def stream_response():
+                    for chunk in response:
+                        content = chunk.choices[0].delta.content
+                        if content:
+                            yield content
+                return stream_response()
+            else:
+                return response.choices[0].message.content
+                
+        return self._retry_request(_make_request)
 
-        try:
-            with requests.post(url, json=payload, stream=stream) as response:
-                if response.status_code != 200:
-                    raise Exception(f"Ollama API error: {response.text}")
-
-                if stream:
-                    return (chunk.decode("utf-8") for chunk in response.iter_content())
-                else:
-                    return response.json()["response"]
-        except Exception as e:
-            raise Exception(f"Ollama request failed: {e}")
-
-    def _generate_hf(self, prompt, max_tokens, stream):
-        try:
-            response = self.hf_client.chat.completions.create(
+
+class OpenAIProvider(LLMProvider):
+    def __init__(self, model_name: str, api_key: str = None, timeout: int = 30, retries: int = 3):
+        super().__init__(model_name, timeout, retries)
+        self.client = openai.OpenAI(api_key=api_key or config.openai_api_key)
+
+    def generate(self, prompt: str, max_tokens: int = 500, stream: bool = False) -> Union[str, Generator[str, None, None]]:
+        def _make_request():
+            response = self.client.chat.completions.create(
                 model=self.model_name,
                 messages=[{"role": "user", "content": prompt}],
                 max_tokens=max_tokens,
                 stream=stream
             )
+            
             if stream:
-                return (chunk.choices[0].delta.content or "" for chunk in response)
+                def stream_response():
+                    for chunk in response:
+                        content = chunk.choices[0].delta.content
+                        if content:
+                            yield content
+                return stream_response()
             else:
-                return response.choices[0].text
-        except Exception as e:
-            raise Exception(f"Hugging Face API error: {e}")
+                return response.choices[0].message.content
+                
+        return self._retry_request(_make_request)
+
+
+class LLMClient:
+    PROVIDER_MAP = {
+        "ollama": OllamaProvider,
+        "huggingface": HuggingFaceProvider,
+        "openai": OpenAIProvider
+    }
+    
+    def __init__(self, provider: str = "ollama", model_name: str = None, **provider_kwargs):
+        self.provider_name = provider.lower()
+        self.model_name = model_name or self._get_default_model()
+        
+        if self.provider_name not in self.PROVIDER_MAP:
+            raise ValueError(f"Unsupported provider: {provider}")
+            
+        provider_class = self.PROVIDER_MAP[self.provider_name]
+        self.provider = provider_class(self.model_name, **provider_kwargs)
+    
+    def _get_default_model(self) -> str:
+        """Get default model based on provider"""
+        defaults = {
+            "ollama": config.local_model_name,
+            "huggingface": "meta-llama/Meta-Llama-3-8B-Instruct",
+            "openai": "gpt-3.5-turbo"
+        }
+        return defaults.get(self.provider_name, "mistral")
+    
+    def generate(self, prompt: str, max_tokens: int = 500, stream: bool = False) -> Union[str, Generator[str, None, None]]:
+        """Unified generate method that delegates to provider"""
+        return self.provider.generate(prompt, max_tokens, stream)
+    
+    @classmethod
+    def get_available_providers(cls) -> list:
+        """Return list of supported providers"""
+        return list(cls.PROVIDER_MAP.keys())