Enhance modular LLM provider interface with HF endpoint integration

src/config/llm_config.py

@@ -4,23 +4,23 @@ from typing import Optional
 class LLMConfig:
     """Configuration loader for LLM providers"""
     
-    @staticmethod
-    def get_active_provider() -> Optional[str]:
-        """Get the name of the active provider based on environment variables"""
-        if os.getenv("HF_TOKEN"):
-            return "huggingface"
-        # elif os.getenv("OLLAMA_HOST"):
-        #     return "ollama"
-        # elif os.getenv("OPENAI_API_KEY"):
-        #     return "openai"
-        return None
+    def __init__(self):
+        # Load all environment variables
+        self.hf_token = os.getenv("HF_TOKEN")
+        self.ollama_host = os.getenv("OLLAMA_HOST")
+        self.local_model_name = os.getenv("LOCAL_MODEL_NAME", "mistral:latest")
+        self.hf_api_url = os.getenv("HF_API_ENDPOINT_URL", "https://zxzbfrlg3ssrk7d9.us-east-1.aws.endpoints.huggingface.cloud/v1/")
+        self.use_fallback = os.getenv("USE_FALLBACK", "true").lower() == "true"
+        self.openweather_api_key = os.getenv("OPENWEATHER_API_KEY")
+        self.nasa_api_key = os.getenv("NASA_API_KEY")
+        self.tavily_api_key = os.getenv("TAVILY_API_KEY")
+        self.redis_host = os.getenv("REDIS_HOST")
+        self.redis_port = os.getenv("REDIS_PORT")
+        self.redis_username = os.getenv("REDIS_USERNAME")
+        self.redis_password = os.getenv("REDIS_PASSWORD")
+        
+        # Detect if running on HF Spaces
+        self.is_hf_space = bool(os.getenv("SPACE_ID"))
 
-    @staticmethod
-    def get_provider_model(provider: str) -> str:
-        """Get the model name for a given provider"""
-        model_map = {
-            "huggingface": os.getenv("HF_MODEL_NAME", "meta-llama/Llama-2-7b-chat-hf"),
-            # "ollama": os.getenv("LOCAL_MODEL_NAME", "mistral:latest"),
-            # "openai": "gpt-3.5-turbo"
-        }
-        return model_map.get(provider, "unknown-model")
+# Global config instance
+config = LLMConfig()

src/llm/factory.py

@@ -1,17 +1,18 @@
-import os
+import logging
 from typing import Optional
 from src.llm.base_provider import LLMProvider
 from src.llm.hf_provider import HuggingFaceProvider
-# Import other providers as they are implemented
-# from src.llm.ollama_provider import OllamaProvider
-# from src.llm.openai_provider import OpenAIProvider
+from src.llm.ollama_provider import OllamaProvider
+from utils.config import config
+
+logger = logging.getLogger(__name__)
 
 class ProviderNotAvailableError(Exception):
     """Raised when no provider is available"""
     pass
 
 class LLMFactory:
-    """Factory for creating LLM providers"""
+    """Factory for creating LLM providers with fallback support"""
     
     _instance = None
 
@@ -25,7 +26,7 @@ class LLMFactory:
         Get an LLM provider based on preference and availability.
         
         Args:
-            preferred_provider: Preferred provider name ('huggingface', 'ollama', 'openai')
+            preferred_provider: Preferred provider name ('huggingface', 'ollama')
             
         Returns:
             LLMProvider instance
@@ -33,18 +34,43 @@ class LLMFactory:
         Raises:
             ProviderNotAvailableError: When no providers are available
         """
-        # For now, we only have HF provider implemented
-        if preferred_provider == "huggingface" or (preferred_provider is None and os.getenv("HF_TOKEN")):
-            return HuggingFaceProvider(
-                model_name=os.getenv("HF_MODEL_NAME", "meta-llama/Llama-2-7b-chat-hf")
-            )
-        
-        # Add other providers as they are implemented
-        # elif preferred_provider == "ollama" or (preferred_provider is None and os.getenv("OLLAMA_HOST")):
-        #     return OllamaProvider(model_name=os.getenv("LOCAL_MODEL_NAME", "mistral:latest"))
-        # elif preferred_provider == "openai" or (preferred_provider is None and os.getenv("OPENAI_API_KEY")):
-        #     return OpenAIProvider(model_name="gpt-3.5-turbo")
+        # Check preferred provider first
+        if preferred_provider == "huggingface" and config.hf_token:
+            try:
+                return HuggingFaceProvider(
+                    model_name="DavidAU/OpenAi-GPT-oss-20b-abliterated-uncensored-NEO-Imatrix-gguf"
+                )
+            except Exception as e:
+                logger.warning(f"Failed to initialize HF provider: {e}")
         
+        elif preferred_provider == "ollama" and config.ollama_host:
+            try:
+                return OllamaProvider(
+                    model_name=config.local_model_name
+                )
+            except Exception as e:
+                logger.warning(f"Failed to initialize Ollama provider: {e}")
+
+        # Fallback logic based on configuration
+        if config.use_fallback:
+            # Try HF first if configured
+            if config.hf_token:
+                try:
+                    return HuggingFaceProvider(
+                        model_name="DavidAU/OpenAi-GPT-oss-20b-abliterated-uncensored-NEO-Imatrix-gguf"
+                    )
+                except Exception as e:
+                    logger.warning(f"Failed to initialize HF provider: {e}")
+            
+            # Then try Ollama if configured
+            if config.ollama_host:
+                try:
+                    return OllamaProvider(
+                        model_name=config.local_model_name
+                    )
+                except Exception as e:
+                    logger.warning(f"Failed to initialize Ollama provider: {e}")
+
         raise ProviderNotAvailableError("No LLM providers are available or configured")
 
 # Global factory instance

src/llm/hf_provider.py

@@ -1,20 +1,101 @@
+import time
+import logging
 from typing import List, Dict, Optional, Union
 from src.llm.base_provider import LLMProvider
+from utils.config import config
+
+logger = logging.getLogger(__name__)
+
+try:
+    from openai import OpenAI
+    HF_SDK_AVAILABLE = True
+except ImportError:
+    HF_SDK_AVAILABLE = False
+    OpenAI = None
 
 class HuggingFaceProvider(LLMProvider):
-    """Hugging Face LLM provider stub implementation"""
+    """Hugging Face LLM provider for your custom endpoint"""
 
-    def __init__(self, model_name: str, timeout: int = 30, max_retries: int = 3):
+    def __init__(self, model_name: str, timeout: int = 60, max_retries: int = 3):
         super().__init__(model_name, timeout, max_retries)
-        # Placeholder for actual client initialization
-        self.client = None
+        
+        if not HF_SDK_AVAILABLE:
+            raise ImportError("Hugging Face provider requires 'openai' package")
+            
+        if not config.hf_token:
+            raise ValueError("HF_TOKEN not set - required for Hugging Face provider")
+
+        # Use your specific endpoint URL
+        self.client = OpenAI(
+            base_url=config.hf_api_url,
+            api_key=config.hf_token
+        )
+        logger.info(f"Initialized HF provider with endpoint: {config.hf_api_url}")
 
     def generate(self, prompt: str, conversation_history: List[Dict]) -> Optional[str]:
-        """Stub for synchronous generation"""
-        # In a real implementation, this would call the HF API
-        return f"[HuggingFace Stub] Response to: {prompt}"
+        """Generate a response synchronously"""
+        try:
+            response = self.client.chat.completions.create(
+                model=self.model_name,
+                messages=conversation_history,
+                max_tokens=8192,
+                temperature=0.7,
+                stream=False
+            )
+            return response.choices[0].message.content
+        except Exception as e:
+            logger.error(f"HF generation failed: {e}")
+            # Handle scale-to-zero behavior
+            if "503" in str(e) or "service unavailable" in str(e).lower():
+                logger.info("HF endpoint is scaling up, waiting...")
+                time.sleep(60)  # Wait for endpoint to initialize
+                # Retry once
+                response = self.client.chat.completions.create(
+                    model=self.model_name,
+                    messages=conversation_history,
+                    max_tokens=8192,
+                    temperature=0.7,
+                    stream=False
+                )
+                return response.choices[0].message.content
+            raise
 
     def stream_generate(self, prompt: str, conversation_history: List[Dict]) -> Optional[Union[str, List[str]]]:
-        """Stub for streaming generation"""
-        # In a real implementation, this would stream from the HF API
-        return ["[HuggingFace", " Streaming", " Stub]"]
+        """Generate a response with streaming support"""
+        try:
+            response = self.client.chat.completions.create(
+                model=self.model_name,
+                messages=conversation_history,
+                max_tokens=8192,
+                temperature=0.7,
+                stream=True
+            )
+            
+            chunks = []
+            for chunk in response:
+                content = chunk.choices[0].delta.content
+                if content:
+                    chunks.append(content)
+            return chunks
+        except Exception as e:
+            logger.error(f"HF stream generation failed: {e}")
+            # Handle scale-to-zero behavior
+            if "503" in str(e) or "service unavailable" in str(e).lower():
+                logger.info("HF endpoint is scaling up, waiting...")
+                time.sleep(60)  # Wait for endpoint to initialize
+                # Retry once
+                response = self.client.chat.completions.create(
+                    model=self.model_name,
+                    messages=conversation_history,
+                    max_tokens=8192,
+                    temperature=0.7,
+                    stream=True
+                )
+                
+                chunks = []
+                for chunk in response:
+                    content = chunk.choices[0].delta.content
+                    if content:
+                        chunks.append(content)
+                return chunks
+            raise

src/llm/ollama_provider.py

@@ -0,0 +1,87 @@
+import requests
+import logging
+import re
+from typing import List, Dict, Optional, Union
+from src.llm.base_provider import LLMProvider
+from utils.config import config
+
+logger = logging.getLogger(__name__)
+
+class OllamaProvider(LLMProvider):
+    """Ollama LLM provider implementation"""
+
+    def __init__(self, model_name: str, timeout: int = 60, max_retries: int = 3):
+        super().__init__(model_name, timeout, max_retries)
+        self.host = self._sanitize_host(config.ollama_host or "http://localhost:11434")
+        self.headers = {
+            "ngrok-skip-browser-warning": "true",
+            "User-Agent": "CosmicCat-AI-Assistant"
+        }
+
+    def _sanitize_host(self, host: str) -> str:
+        """Sanitize host URL by removing whitespace and control characters"""
+        if not host:
+            return "http://localhost:11434"
+        host = host.strip()
+        host = re.sub(r'[\r\n\t\0]+', '', host)
+        if not host.startswith(('http://', 'https://')):
+            host = 'http://' + host
+        return host
+
+    def generate(self, prompt: str, conversation_history: List[Dict]) -> Optional[str]:
+        """Generate a response synchronously"""
+        try:
+            url = f"{self.host}/api/chat"
+            payload = {
+                "model": self.model_name,
+                "messages": conversation_history,
+                "stream": False
+            }
+            
+            response = requests.post(
+                url,
+                json=payload,
+                headers=self.headers,
+                timeout=self.timeout
+            )
+            response.raise_for_status()
+            result = response.json()
+            return result["message"]["content"]
+        except Exception as e:
+            logger.error(f"Ollama generation failed: {e}")
+            raise
+
+    def stream_generate(self, prompt: str, conversation_history: List[Dict]) -> Optional[Union[str, List[str]]]:
+        """Generate a response with streaming support"""
+        try:
+            url = f"{self.host}/api/chat"
+            payload = {
+                "model": self.model_name,
+                "messages": conversation_history,
+                "stream": True
+            }
+            
+            response = requests.post(
+                url,
+                json=payload,
+                headers=self.headers,
+                timeout=self.timeout,
+                stream=True
+            )
+            response.raise_for_status()
+            
+            chunks = []
+            for line in response.iter_lines():
+                if line:
+                    chunk = line.decode('utf-8')
+                    try:
+                        data = eval(chunk)
+                        content = data.get("message", {}).get("content", "")
+                        if content:
+                            chunks.append(content)
+                    except:
+                        continue
+            return chunks
+        except Exception as e:
+            logger.error(f"Ollama stream generation failed: {e}")
+            raise