Update app/rag_setup.py

app/rag_setup.py

@@ -141,83 +141,58 @@ class OpenRouterLLM:
             self.client_ready = False
             return
         
-        # Don't test connection during initialization to avoid blocking startup
-        # Testing will happen when first needed
-        self.client_ready = True
-        logger.info("✅ OpenRouter LLM initialized (connection will be tested on first use)")
-        logger.info("=" * 60)
-
-    def test_connection(self) -> dict:
-        """Test the API connection with minimal request."""
-        logger.info("🔍 Testing OpenRouter API connection...")
-        
-        # Use a very simple test prompt and minimal tokens
-        test_prompt = "Hi"
-        response = self._make_api_request(
-            prompt=test_prompt, 
-            max_tokens=5,  # Minimal tokens for faster response
-            timeout=15,    # Short timeout for testing
-            is_test=True   # Flag to indicate this is a test
-        )
+        # Test the connection with minimal tokens
+        try:
+            logger.info("🔍 Testing OpenRouter connection...")
+            test_response = self._make_api_request("Hello", max_tokens=5)
+            if test_response and "error" not in test_response:
+                logger.info("✅ OpenRouter connection test successful")
+                self.client_ready = True
+            else:
+                logger.error(f"❌ OpenRouter connection test failed: {test_response}")
+                self.client_ready = False
+        except Exception as e:
+            logger.error(f"❌ OpenRouter connection test failed: {e}")
+            self.client_ready = False
         
-        if response and "error" not in response:
-            logger.info("✅ OpenRouter connection test successful")
-            return {"valid": True, "message": "API key is valid and working"}
-        else:
-            error_msg = response.get("error", "Unknown error") if response else "No response"
-            logger.error(f"❌ OpenRouter connection test failed: {error_msg}")
-            return {"valid": False, "message": f"API test failed: {error_msg}"}
+        logger.info("=" * 60)
 
-    def _make_api_request(self, prompt: str, max_tokens: int = 2000, timeout: int = None, is_test: bool = False) -> dict:
-        """Make a direct HTTP request to OpenRouter API with improved timeout handling."""
+    def _make_api_request(self, prompt: str, max_tokens: int = 2000, timeout: int = None) -> dict:
+        """Make a direct HTTP request to OpenRouter API with configurable token limits."""
         
-        # Calculate dynamic timeout based on context
+        # Calculate dynamic timeout based on max_tokens and prompt length
         if timeout is None:
-            if is_test:
-                timeout = 15  # Short timeout for tests
-            else:
-                base_timeout = 60  # Reduced base timeout
-                # More conservative timeout calculation
-                token_timeout = max(10, max_tokens // 200)  # ~1 second per 200 tokens
-                prompt_timeout = max(5, len(prompt) // 2000)  # ~1 second per 2000 characters
-                timeout = min(base_timeout + token_timeout + prompt_timeout, 300)  # Cap at 5 minutes
+            base_timeout = 120
+            # More tokens = longer generation time
+            token_timeout = max(20, max_tokens // 100)  # ~1 second per 100 tokens
+            prompt_timeout = max(10, len(prompt) // 1000)  # ~1 second per 2000 characters
+            timeout = min(base_timeout + token_timeout + prompt_timeout, 600)  # Cap at 5 minutes
         
         logger.info(f"🌐 Making API request to OpenRouter")
         logger.info(f"📏 Prompt length: {len(prompt)} characters")
         logger.info(f"🎯 Max tokens: {max_tokens}")
         logger.info(f"⏱️  Timeout: {timeout}s")
-        logger.info(f"🧪 Is test: {is_test}")
         
         headers = {
             "Authorization": f"Bearer {self.api_key}",
             "Content-Type": "application/json",
             "HTTP-Referer": "https://github.com/Ab-Romia/ContextIQ-RAG",
-            "X-Title": "Context Aware AI",
-            "User-Agent": "ContextIQ/1.0"  # Add user agent
+            "X-Title": "Context Aware AI"
         }
         
-        # Optimize payload for faster responses, especially for tests
+        # Optimize payload for longer responses
         payload = {
             "model": self.model,
             "messages": [{"role": "user", "content": prompt}],
             "max_tokens": max_tokens,
+            "temperature": 0.7,
+            "top_p": 0.9,
             "stream": False,
+            # Add parameters to encourage complete responses
+            "presence_penalty": 0.1,  # Slight penalty for repetition
+            "frequency_penalty": 0.1,  # Slight penalty for frequency
         }
         
-        # Add performance optimizations for tests
-        if is_test:
-            payload.update({
-                "temperature": 0.1,  # Lower temperature for faster, more deterministic responses
-                "top_p": 0.5,        # Lower top_p for faster generation
-            })
-        else:
-            payload.update({
-                "temperature": 0.7,
-                "top_p": 0.9,
-                "presence_penalty": 0.1,
-                "frequency_penalty": 0.1,
-            })
-        
         # Log the request payload (without sensitive data)
         safe_payload = payload.copy()
         safe_payload["messages"] = [{"role": "user", "content": f"[CONTENT: {len(prompt)} chars]"}]
@@ -226,35 +201,12 @@ class OpenRouterLLM:
         try:
             start_time = time.time()
             
-            # Use session with connection pooling and improved settings
             with requests.Session() as session:
-                # Configure session for better performance
-                session.headers.update(headers)
-                
-                # Configure adapters for retry and connection pooling
-                from requests.adapters import HTTPAdapter
-                from urllib3.util.retry import Retry
-                
-                # Define retry strategy for transient errors
-                retry_strategy = Retry(
-                    total=2,  # Reduced retries for faster failure
-                    backoff_factor=0.5,
-                    status_forcelist=[429, 500, 502, 503, 504, 408],  # Include 408 timeout
-                    allowed_methods=["POST"]
-                )
-                
-                adapter = HTTPAdapter(
-                    max_retries=retry_strategy,
-                    pool_connections=10,
-                    pool_maxsize=10
-                )
-                session.mount("http://", adapter)
-                session.mount("https://", adapter)
-                
                 response = session.post(
                     self.api_url,
+                    headers=headers,
                     json=payload,
-                    timeout=(10, timeout),  # (connection timeout, read timeout)
+                    timeout=timeout
                 )
             
             request_time = time.time() - start_time
@@ -280,47 +232,27 @@ class OpenRouterLLM:
                         if completion_tokens >= max_tokens * 0.95:  # If we used 95% of max tokens
                             logger.warning(f"⚠️  Response may be truncated (used {completion_tokens}/{max_tokens} tokens)")
                     
-                    if not is_test:  # Don't log content for tests
-                        content_preview = content[:300] + "..." if len(content) > 300 else content
-                        logger.info(f"📄 Response preview: {content_preview}")
+                    content_preview = content[:300] + "..." if len(content) > 300 else content
+                    logger.info(f"📄 Response preview: {content_preview}")
                 
                 return response_data
-            
-            elif response.status_code == 408:
-                logger.error(f"⏱️  Request timed out (408) after {request_time:.2f}s")
-                return {"error": f"Request timed out. Try again or reduce the request complexity."}
-            
-            elif response.status_code == 429:
-                logger.error(f"🚦 Rate limited (429)")
-                return {"error": "Rate limited. Please wait a moment and try again."}
-            
-            elif response.status_code in [500, 502, 503, 504]:
-                logger.error(f"🏥 Server error ({response.status_code})")
-                return {"error": f"Server error ({response.status_code}). Please try again later."}
-            
             else:
                 logger.error(f"❌ API request failed with status {response.status_code}")
                 logger.error(f"📄 Response text: {response.text}")
                 return {"error": f"HTTP {response.status_code}: {response.text}"}
                 
-        except requests.exceptions.Timeout as e:
-            logger.error(f"⏱️  Request timed out after {timeout}s: {e}")
-            return {"error": f"Request timed out after {timeout}s. Please try again or reduce the request size."}
-        
+        except requests.exceptions.Timeout:
+            logger.error(f"⏱️  API request timed out after {timeout}s")
+            return {"error": f"Request timed out after {timeout}s. Try reducing the context length or max tokens."}
         except requests.exceptions.ConnectionError as e:
             logger.error(f"🌐 Connection error: {e}")
-            return {"error": f"Connection error. Please check your internet connection."}
-        
-        except requests.exceptions.HTTPError as e:
-            logger.error(f"🌐 HTTP error: {e}")
-            return {"error": f"HTTP error: {str(e)}"}
-        
+            return {"error": f"Connection error: {str(e)}"}
         except Exception as e:
-            logger.error(f"❌ Unexpected error in API request: {e}")
-            return {"error": f"Unexpected error: {str(e)}"}
+            logger.error(f"❌ API request failed: {e}")
+            return {"error": str(e)}
 
     def generate_content(self, prompt: str, max_tokens: int = 2000) -> str:
-        """Generate content with improved error handling and timeout management."""
+        """Generate content with configurable token limits."""
         logger.info("=" * 80)
         logger.info("🧠 LLM CONTENT GENERATION STARTED")
         logger.info("=" * 80)
@@ -371,9 +303,9 @@ class OpenRouterLLM:
             logger.error(error_msg)
             return error_msg
         
-        max_retries = 2  # Reduced retries for faster failure
+        max_retries = 3
         retry_count = 0
-        base_wait_time = 1  # Reduced wait time
+        base_wait_time = 2
 
         while retry_count <= max_retries:
             try:
@@ -385,7 +317,7 @@ class OpenRouterLLM:
                 
                 if retry_count > 0:
                     # Reduce max_tokens on retries for faster responses
-                    current_max_tokens = max(500, max_tokens - (retry_count * 500))
+                    current_max_tokens = max(1000, max_tokens - (retry_count * 500))
                     logger.info(f"🔧 Retry attempt - reducing max_tokens to {current_max_tokens}")
                 
                 response = self._make_api_request(prompt, max_tokens=current_max_tokens, timeout=timeout)
@@ -397,9 +329,8 @@ class OpenRouterLLM:
                     if "timeout" in error_msg.lower() or "408" in error_msg:
                         logger.warning(f"⏱️  Timeout error on attempt {retry_count + 1}")
                         if retry_count < max_retries:
-                            retry_count += 1
                             continue
-                    elif "429" in error_msg or "rate limit" in error_msg.lower():
+                    elif "429" in error_msg:
                         logger.warning(f"🚦 Rate limit error on attempt {retry_count + 1}")
                         wait_time = base_wait_time * (2 ** retry_count)
                         logger.info(f"⏳ Waiting {wait_time}s for rate limit cooldown...")
@@ -459,7 +390,7 @@ class OpenRouterLLM:
                     logger.info("=" * 80)
                     return final_error
                 
-                wait_time = base_wait_time * retry_count
+                wait_time = base_wait_time * retry_count + (retry_count * 0.5)
                 logger.info(f"⏳ Waiting {wait_time:.1f}s before retry...")
                 time.sleep(wait_time)
 
@@ -483,9 +414,6 @@ except Exception as e:
     class DummyLLM:
         def generate_content(self, prompt: str) -> str:
             return f"❌ AI model is not available. Initialization error: {str(e)}"
-        
-        def test_connection(self) -> dict:
-            return {"valid": False, "message": f"Model not available: {str(e)}"}
     
     generation_model = DummyLLM()
     logger.warning("⚠️  Using dummy LLM due to initialization failure")