Update llm_spark.py

llm_spark.py

@@ -1,23 +1,30 @@
 """
 Spark LLM Implementation
 """
+import os
 import httpx
-from typing import Dict, List, Any
+import json
+from typing import Dict, List, Any, AsyncIterator
 from llm_interface import LLMInterface
 from logger import log_info, log_error, log_warning, log_debug
 
+# Get timeout from environment
+DEFAULT_LLM_TIMEOUT = int(os.getenv("LLM_TIMEOUT_SECONDS", "60"))
+MAX_RESPONSE_LENGTH = int(os.getenv("LLM_MAX_RESPONSE_LENGTH", "4096"))
+
 class SparkLLM(LLMInterface):
-    """Spark LLM integration"""
+    """Spark LLM integration with improved error handling"""
     
     def __init__(self, spark_endpoint: str, spark_token: str, provider_variant: str = "cloud", settings: Dict[str, Any] = None):
         super().__init__(settings)
         self.spark_endpoint = spark_endpoint.rstrip("/")
         self.spark_token = spark_token
         self.provider_variant = provider_variant
-        log_info(f"🔌 SparkLLM initialized with endpoint: {self.spark_endpoint}")
+        self.timeout = self.settings.get("timeout", DEFAULT_LLM_TIMEOUT)
+        log_info(f"🔌 SparkLLM initialized", endpoint=self.spark_endpoint, timeout=self.timeout)
     
     async def generate(self, system_prompt: str, user_input: str, context: List[Dict]) -> str:
-        """Generate response from Spark LLM"""
+        """Generate response with improved error handling and streaming support"""
         headers = {
             "Authorization": f"Bearer {self.spark_token}",
             "Content-Type": "application/json"
@@ -25,85 +32,85 @@ class SparkLLM(LLMInterface):
         
         # Build context messages
         messages = []
-        for msg in context:
+        if system_prompt:
+            messages.append({
+                "role": "system",
+                "content": system_prompt
+            })
+        
+        for msg in context[-10:]:  # Last 10 messages for context
             messages.append({
                 "role": msg.get("role", "user"),
                 "content": msg.get("content", "")
             })
         
+        messages.append({
+            "role": "user",
+            "content": user_input
+        })
+        
         payload = {
-            "user_input": user_input,
-            "system_prompt": system_prompt,
-            "context": messages,
-            "mode": self.provider_variant
+            "messages": messages,
+            "mode": self.provider_variant,
+            "max_tokens": self.settings.get("max_tokens", 2048),
+            "temperature": self.settings.get("temperature", 0.7),
+            "stream": False  # For now, no streaming
         }
         
         try:
-            async with httpx.AsyncClient(timeout=60) as client:
-                response = await client.post(
-                    f"{self.spark_endpoint}/generate",
-                    json=payload,
-                    headers=headers
-                )
-                response.raise_for_status()
-                result = response.json()
-                return result.get("model_answer", "")
+            async with httpx.AsyncClient(timeout=self.timeout) as client:
+                with LogTimer(f"Spark LLM request"):
+                    response = await client.post(
+                        f"{self.spark_endpoint}/generate",
+                        json=payload,
+                        headers=headers
+                    )
+                    
+                    # Check for rate limiting
+                    if response.status_code == 429:
+                        retry_after = response.headers.get("Retry-After", "60")
+                        log_warning(f"Rate limited by Spark", retry_after=retry_after)
+                        raise httpx.HTTPStatusError(
+                            f"Rate limited. Retry after {retry_after}s",
+                            request=response.request,
+                            response=response
+                        )
+                    
+                    response.raise_for_status()
+                    result = response.json()
+                    
+                    # Extract response
+                    content = result.get("model_answer", "")
+                    
+                    # Check response length
+                    if len(content) > MAX_RESPONSE_LENGTH:
+                        log_warning(f"Response exceeded max length, truncating", 
+                                  original_length=len(content),
+                                  max_length=MAX_RESPONSE_LENGTH)
+                        content = content[:MAX_RESPONSE_LENGTH] + "..."
+                    
+                    return content
+                    
         except httpx.TimeoutException:
-            log_warning("⏱️ Spark request timed out")
+            log_error(f"Spark request timed out", timeout=self.timeout)
             raise
-        except Exception as e:
-            log_error("❌ Spark error", e)
+        except httpx.HTTPStatusError as e:
+            log_error(f"Spark HTTP error", 
+                     status_code=e.response.status_code,
+                     response=e.response.text[:500])
             raise
-    
-    async def startup(self, project_config: Dict) -> bool:
-        """Initialize Spark with project config"""
-        try:
-            headers = {
-                "Authorization": f"Bearer {self.spark_token}",
-                "Content-Type": "application/json"
-            }
-            
-            # Extract version config
-            version = None
-            for v in project_config.get("versions", []):
-                if v.get("published"):
-                    version = v
-                    break
-            
-            if not version:
-                log_info("❌ No published version found")
-                return False
-            
-            llm_config = version.get("llm", {})
-            payload = {
-                "project_name": project_config.get("name"),
-                "repo_id": llm_config.get("repo_id", ""),
-                "use_fine_tune": llm_config.get("use_fine_tune", False),
-                "fine_tune_zip": llm_config.get("fine_tune_zip", ""),
-                "generation_config": llm_config.get("generation_config", {})
-            }
-            
-            async with httpx.AsyncClient(timeout=30) as client:
-                response = await client.post(
-                    f"{self.spark_endpoint}/startup",
-                    json=payload,
-                    headers=headers
-                )
-                response.raise_for_status()
-                log_info("✅ Spark startup successful")
-                return True
         except Exception as e:
-            log_error("❌ Spark startup failed", e)
-            return False
+            log_error("Spark unexpected error", error=str(e))
+            raise
     
     def get_provider_name(self) -> str:
-        """Get provider name"""
         return f"spark-{self.provider_variant}"
     
     def get_model_info(self) -> Dict[str, Any]:
-        """Get model information"""
         return {
             "provider": "spark",
             "variant": self.provider_variant,
-            "endpoint": self.spark_endpoint
+            "endpoint": self.spark_endpoint,
+            "max_tokens": self.settings.get("max_tokens", 2048),
+            "temperature": self.settings.get("temperature", 0.7)
         }