Update chat_handler.py

chat_handler.py

@@ -80,59 +80,67 @@ def setup_llm_provider():
             log("❌ SPARK_TOKEN not found")
             raise ValueError("SPARK_TOKEN not configured")
         
-        spark_endpoint = str(cfg.global_config.spark_endpoint)
-        llm_provider = SparkLLM(spark_endpoint, spark_token)
-        return
-    
-    # Parse internal prompt format: "provider:model"
-    parts = internal_prompt.split(":", 1)
-    if len(parts) != 2:
-        log(f"⚠️ Invalid internal_prompt format: {internal_prompt}, using Spark")
-        spark_token = _get_spark_token()
-        if not spark_token:
-            log("❌ SPARK_TOKEN not found")
-            raise ValueError("SPARK_TOKEN not configured")
+        spark_endpoint = str(cfg.global_config.spark_endpoint).rstrip("/")
+        work_mode = cfg.global_config.work_mode
+        
+        log(f"🔌 Initializing SparkLLM: {spark_endpoint}")
+        log(f"🔧 Work mode: {work_mode}")
+        
+        llm_provider = SparkLLM(
+            spark_endpoint=spark_endpoint,
+            spark_token=spark_token,
+            work_mode=work_mode
+        )
         
-        spark_endpoint = str(cfg.global_config.spark_endpoint)
-        llm_provider = SparkLLM(spark_endpoint, spark_token)
+        log("✅ SparkLLM initialized")
         return
     
-    provider, model = parts[0].lower(), parts[1]
-    
-    if provider == "openai":
-        # Get API key from environment
-        api_key = os.getenv("OPENAI_API_KEY")
+    # Check if it's a GPT-4o config
+    if internal_prompt.get("provider") == "gpt-4o":
+        api_key = internal_prompt.get("api_key")
         if not api_key:
-            log("❌ OPENAI_API_KEY not found in environment")
-            # Fallback to Spark
-            spark_token = _get_spark_token()
-            if not spark_token:
-                raise ValueError("Neither OPENAI_API_KEY nor SPARK_TOKEN configured")
-            
-            spark_endpoint = str(cfg.global_config.spark_endpoint)
-            llm_provider = SparkLLM(spark_endpoint, spark_token)
-            return
+            if cfg.global_config.is_cloud_mode():
+                api_key = os.environ.get("OPENAI_API_KEY")
+            else:
+                from dotenv import load_dotenv
+                load_dotenv()
+                api_key = os.getenv("OPENAI_API_KEY")
+        
+        if not api_key:
+            raise ValueError("OpenAI API key not configured")
         
-        log(f"🤖 Using OpenAI with model: {model}")
-        llm_provider = GPT4oLLM(api_key, model)
+        llm_provider = GPT4oLLM(
+            api_key=api_key,
+            model=internal_prompt.get("model", "gpt-4o"),
+            max_tokens=internal_prompt.get("max_tokens", 4096),
+            temperature=internal_prompt.get("temperature", 0.7)
+        )
+        log("✅ GPT-4o LLM initialized")
     else:
-        log(f"⚠️ Unknown provider: {provider}, using Spark")
+        # Default to Spark
         spark_token = _get_spark_token()
         if not spark_token:
             raise ValueError("SPARK_TOKEN not configured")
         
-        spark_endpoint = str(cfg.global_config.spark_endpoint)
-        llm_provider = SparkLLM(spark_endpoint, spark_token)
+        spark_endpoint = str(cfg.global_config.spark_endpoint).rstrip("/")
+        work_mode = cfg.global_config.work_mode
+        
+        llm_provider = SparkLLM(
+            spark_endpoint=spark_endpoint,
+            spark_token=spark_token,
+            work_mode=work_mode
+        )
+        log("✅ SparkLLM initialized (via internal_prompt)")
 
 def _get_spark_token() -> Optional[str]:
-    """Get Spark token based on work_mode"""
+    """Get Spark token based on work mode"""
     cfg = ConfigProvider.get()
     
     if cfg.global_config.is_cloud_mode():
-        # Cloud mode - use HuggingFace Secrets
-        token = os.getenv("SPARK_TOKEN")
-        if not token:
-            log("❌ SPARK_TOKEN not found in HuggingFace Secrets!")
+        # Cloud mode - use HuggingFace secrets
+        token = os.environ.get("SPARK_TOKEN")
+        if token:
+            log("🔑 Using SPARK_TOKEN from environment")
         return token
     else:
         # On-premise mode - use .env file