Implement HF endpoint monitoring and integration with wake-up functionality

src/llm/factory.py

@@ -1,9 +1,10 @@
 import logging
-from typing import Optional, List
+from typing import Optional
 from src.llm.base_provider import LLMProvider
 from src.llm.hf_provider import HuggingFaceProvider
 from src.llm.ollama_provider import OllamaProvider
 from utils.config import config
+from src.services.hf_monitor import hf_monitor
 
 logger = logging.getLogger(__name__)
 
@@ -24,87 +25,54 @@ class LLMFactory:
     def get_provider(self, preferred_provider: Optional[str] = None) -> LLMProvider:
         """
         Get an LLM provider based on preference and availability.
-        
-        Args:
-            preferred_provider: Preferred provider name ('huggingface', 'ollama')
-            
-        Returns:
-            LLMProvider instance
-            
-        Raises:
-            ProviderNotAvailableError: When no providers are available
         """
-        # Build provider chain based on configuration and preference
-        provider_chain = self._build_provider_chain(preferred_provider)
+        # Try preferred provider first
+        if preferred_provider:
+            provider = self._try_provider(preferred_provider)
+            if provider:
+                return provider
+        
+        # Try HF provider if configured
+        if config.hf_token:
+            provider = self._try_provider("huggingface")
+            if provider:
+                return provider
         
-        # Try providers in order
-        for provider_name, provider_class, model_name in provider_chain:
-            try:
-                logger.info(f"Attempting to initialize {provider_name} provider...")
-                provider = provider_class(model_name=model_name)
-                # Test that provider is working
-                if self._test_provider(provider):
-                    logger.info(f"Successfully initialized {provider_name} provider")
-                    return provider
-                else:
-                    logger.warning(f"{provider_name} provider failed validation test")
-            except Exception as e:
-                logger.warning(f"Failed to initialize {provider_name} provider: {e}")
-                continue
+        # Try Ollama as fallback
+        if config.ollama_host:
+            provider = self._try_provider("ollama")
+            if provider:
+                return provider
         
         raise ProviderNotAvailableError("No LLM providers are available or configured")
 
-    def _build_provider_chain(self, preferred_provider: Optional[str]) -> List[tuple]:
-        """Build provider chain based on preference and configuration"""
-        chain = []
-        
-        # Add preferred provider first if specified
-        if preferred_provider:
-            provider_info = self._get_provider_info(preferred_provider)
-            if provider_info:
-                chain.append(provider_info)
-        
-        # Add fallback providers based on configuration
-        if config.use_fallback:
-            # Add HF if configured
-            if config.hf_token:
-                chain.append((
-                    "huggingface",
-                    HuggingFaceProvider,
-                    "DavidAU/OpenAi-GPT-oss-20b-abliterated-uncensored-NEO-Imatrix-gguf"
-                ))
+    def _try_provider(self, provider_name: str) -> Optional[LLMProvider]:
+        """Try to initialize a specific provider"""
+        try:
+            if provider_name == "huggingface" and config.hf_token:
+                # Check if HF endpoint is available
+                status = hf_monitor.get_endpoint_status()
+                if status["available"] or status["initializing"]:
+                    return HuggingFaceProvider(
+                        model_name="DavidAU/OpenAi-GPT-oss-20b-abliterated-uncensored-NEO-Imatrix-gguf"
+                    )
+                elif status["status"] == "scaled_to_zero":
+                    # Attempt to wake up the endpoint
+                    logger.info("Attempting to wake up HF endpoint...")
+                    if hf_monitor.attempt_wake_up():
+                        return HuggingFaceProvider(
+                            model_name="DavidAU/OpenAi-GPT-oss-20b-abliterated-uncensored-NEO-Imatrix-gguf"
+                        )
+                        
+            elif provider_name == "ollama" and config.ollama_host:
+                return OllamaProvider(
+                    model_name=config.local_model_name
+                )
+                
+        except Exception as e:
+            logger.warning(f"Failed to initialize {provider_name} provider: {e}")
             
-            # Add Ollama if configured
-            if config.ollama_host:
-                chain.append((
-                    "ollama",
-                    OllamaProvider,
-                    config.local_model_name
-                ))
-        
-        return chain
-
-    def _get_provider_info(self, provider_name: str) -> Optional[tuple]:
-        """Get provider class and model info"""
-        provider_map = {
-            "huggingface": (
-                "huggingface",
-                HuggingFaceProvider,
-                "DavidAU/OpenAi-GPT-oss-20b-abliterated-uncensored-NEO-Imatrix-gguf"
-            ),
-            "ollama": (
-                "ollama",
-                OllamaProvider,
-                config.local_model_name
-            )
-        }
-        return provider_map.get(provider_name)
-
-    def _test_provider(self, provider: LLMProvider) -> bool:
-        """Test if provider is working (stub implementation)"""
-        # In a real implementation, you might want to do a lightweight test
-        # For now, we'll assume initialization success means it's working
-        return True
+        return None
 
 # Global factory instance
 llm_factory = LLMFactory()

src/llm/hf_provider.py

@@ -1,8 +1,9 @@
 import time
 import logging
 from typing import List, Dict, Optional, Union
-from src.llm.base_provider import LLMProvider
+from src.llm.enhanced_provider import EnhancedLLMProvider
 from utils.config import config
+from src.services.context_enrichment import context_service
 
 logger = logging.getLogger(__name__)
 
@@ -13,10 +14,10 @@ except ImportError:
     HF_SDK_AVAILABLE = False
     OpenAI = None
 
-class HuggingFaceProvider(LLMProvider):
+class HuggingFaceProvider(EnhancedLLMProvider):
     """Hugging Face LLM provider for your custom endpoint"""
 
-    def __init__(self, model_name: str, timeout: int = 60, max_retries: int = 3):
+    def __init__(self, model_name: str, timeout: int = 120, max_retries: int = 2):
         super().__init__(model_name, timeout, max_retries)
         
         if not HF_SDK_AVAILABLE:
@@ -34,24 +35,20 @@ class HuggingFaceProvider(LLMProvider):
 
     def generate(self, prompt: str, conversation_history: List[Dict]) -> Optional[str]:
         """Generate a response synchronously"""
-        return self._retry_with_backoff(self._generate_impl, prompt, conversation_history)
-
-    def stream_generate(self, prompt: str, conversation_history: List[Dict]) -> Optional[Union[str, List[str]]]:
-        """Generate a response with streaming support"""
-        return self._retry_with_backoff(self._stream_generate_impl, prompt, conversation_history)
-
-    def _generate_impl(self, prompt: str, conversation_history: List[Dict]) -> str:
-        """Implementation of synchronous generation"""
         try:
+            # Enrich context
+            enriched_history = self._enrich_context(conversation_history)
+            
             response = self.client.chat.completions.create(
                 model=self.model_name,
-                messages=conversation_history,
+                messages=enriched_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 self._is_scale_to_zero_error(e):
                 logger.info("HF endpoint is scaling up, waiting...")
@@ -67,12 +64,15 @@ class HuggingFaceProvider(LLMProvider):
                 return response.choices[0].message.content
             raise
 
-    def _stream_generate_impl(self, prompt: str, conversation_history: List[Dict]) -> List[str]:
-        """Implementation of streaming generation"""
+    def stream_generate(self, prompt: str, conversation_history: List[Dict]) -> Optional[Union[str, List[str]]]:
+        """Generate a response with streaming support"""
         try:
+            # Enrich context
+            enriched_history = self._enrich_context(conversation_history)
+            
             response = self.client.chat.completions.create(
                 model=self.model_name,
-                messages=conversation_history,
+                messages=enriched_history,
                 max_tokens=8192,
                 temperature=0.7,
                 stream=True
@@ -85,6 +85,7 @@ class HuggingFaceProvider(LLMProvider):
                     chunks.append(content)
             return chunks
         except Exception as e:
+            logger.error(f"HF stream generation failed: {e}")
             # Handle scale-to-zero behavior
             if self._is_scale_to_zero_error(e):
                 logger.info("HF endpoint is scaling up, waiting...")
@@ -106,6 +107,28 @@ class HuggingFaceProvider(LLMProvider):
                 return chunks
             raise
 
+    def _enrich_context(self, conversation_history: List[Dict]) -> List[Dict]:
+        """Add current context to conversation"""
+        # Get the last user message to determine context needs
+        last_user_message = ""
+        for msg in reversed(conversation_history):
+            if msg["role"] == "user":
+                last_user_message = msg["content"]
+                break
+                
+        # Get current context
+        context = context_service.get_current_context(last_user_message)
+        
+        # Add context as system message at the beginning
+        context_message = {
+            "role": "system",
+            "content": f"[Current Context: {context['current_time']} | Weather: {context['weather']}]"
+        }
+        
+        # Insert context at the beginning
+        enriched_history = [context_message] + conversation_history
+        return enriched_history
+
     def _is_scale_to_zero_error(self, error: Exception) -> bool:
         """Check if the error is related to scale-to-zero initialization"""
         error_str = str(error).lower()
@@ -113,6 +136,7 @@ class HuggingFaceProvider(LLMProvider):
             "503",
             "service unavailable",
             "initializing",
-            "cold start"
+            "cold start",
+            "timeout"
         ]
         return any(indicator in error_str for indicator in scale_to_zero_indicators)

src/services/hf_monitor.py

@@ -0,0 +1,143 @@
+import requests
+import time
+import logging
+from typing import Dict
+from utils.config import config
+
+logger = logging.getLogger(__name__)
+
+class HFEndpointMonitor:
+    """Monitor Hugging Face endpoint status and health"""
+    
+    def __init__(self):
+        self.endpoint_url = config.hf_api_url.rstrip('/') if config.hf_api_url else ""
+        self.hf_token = config.hf_token
+        self.last_check = 0
+        self.check_interval = 300  # 5 minutes
+        self._cached_status = None
+        
+    def get_endpoint_status(self) -> Dict:
+        """Get current HF endpoint status"""
+        current_time = time.time()
+        
+        # Return cached status if checked recently
+        if (self._cached_status and 
+            current_time - self.last_check < 60):
+            return self._cached_status
+            
+        self.last_check = current_time
+        
+        # Check if configured
+        if not self.endpoint_url or not self.hf_token:
+            status = {
+                "status": "not_configured",
+                "message": "HF endpoint not configured",
+                "available": False,
+                "initializing": False
+            }
+            self._cached_status = status
+            return status
+            
+        try:
+            # Check endpoint status
+            headers = {"Authorization": f"Bearer {self.hf_token}"}
+            models_url = f"{self.endpoint_url}/models"
+            
+            response = requests.get(
+                models_url, 
+                headers=headers, 
+                timeout=15
+            )
+            
+            if response.status_code in [200, 201]:
+                status = {
+                    "status": "available",
+                    "message": "HF endpoint is ready",
+                    "available": True,
+                    "initializing": False,
+                    "status_code": response.status_code
+                }
+            elif response.status_code == 503:
+                status = {
+                    "status": "scaled_to_zero",
+                    "message": "HF endpoint is scaled to zero",
+                    "available": False,
+                    "initializing": False,
+                    "status_code": 503
+                }
+            else:
+                status = {
+                    "status": "error",
+                    "message": f"HF endpoint error: {response.status_code}",
+                    "available": False,
+                    "initializing": False,
+                    "status_code": response.status_code
+                }
+                
+        except requests.exceptions.Timeout:
+            status = {
+                "status": "timeout",
+                "message": "HF endpoint timeout (may be initializing)",
+                "available": False,
+                "initializing": True
+            }
+        except Exception as e:
+            status = {
+                "status": "error",
+                "message": f"HF endpoint error: {str(e)}",
+                "available": False,
+                "initializing": False
+            }
+            
+        self._cached_status = status
+        return status
+    
+    def get_human_readable_status(self) -> str:
+        """Get human-readable status message"""
+        status = self.get_endpoint_status()
+        
+        status_messages = {
+            "not_configured": "🟡 HF Endpoint: Not configured",
+            "available": "🟢 HF Endpoint: Available and ready",
+            "scaled_to_zero": "🔴 HF Endpoint: Scaled to zero (send message to wake up)",
+            "timeout": "⏳ HF Endpoint: Initializing (may take 4 minutes)",
+            "error": f"❌ HF Endpoint: Error - {status.get('message', 'Unknown error')}"
+        }
+        
+        return status_messages.get(status["status"], "⚪ HF Endpoint: Unknown status")
+    
+    def attempt_wake_up(self) -> bool:
+        """Attempt to wake up the HF endpoint"""
+        if not self.endpoint_url or not self.hf_token:
+            return False
+            
+        try:
+            headers = {
+                "Authorization": f"Bearer {self.hf_token}",
+                "Content-Type": "application/json"
+            }
+            
+            # Send a minimal request to wake up the endpoint
+            payload = {
+                "model": "DavidAU/OpenAi-GPT-oss-20b-abliterated-uncensored-NEO-Imatrix-gguf",
+                "messages": [{"role": "user", "content": "Hello"}],
+                "max_tokens": 10,
+                "stream": False
+            }
+            
+            chat_url = f"{self.endpoint_url}/chat/completions"
+            response = requests.post(
+                chat_url,
+                headers=headers,
+                json=payload,
+                timeout=45
+            )
+            
+            return response.status_code in [200, 201]
+            
+        except Exception as e:
+            logger.warning(f"Failed to wake up HF endpoint: {e}")
+            return False
+
+# Global instance
+hf_monitor = HFEndpointMonitor()