Update llm_factory.py

llm_factory.py

@@ -5,7 +5,9 @@ import os
 from typing import Optional
 from dotenv import load_dotenv
 
-from llm_interface import LLMInterface, SparkLLM, GPT4oLLM
+from llm_interface import LLMInterface
+from llm_spark import SparkLLM
+from llm_openai import OpenAILLM
 from config_provider import ConfigProvider
 from utils import log
 
@@ -39,63 +41,44 @@ class LLMFactory:
             raise ValueError(f"Unsupported LLM provider: {provider_name}")
     
     @staticmethod
-    def _create_spark_provider(llm_config, api_key: str, provider_def) -> SparkLLM:
+    def _create_spark_provider(llm_config, api_key, provider_def):
         """Create Spark LLM provider"""
-        if not llm_config.endpoint:
-            raise ValueError("Spark endpoint is required")
+        endpoint = llm_config.endpoint
+        if not endpoint:
+            raise ValueError("Spark endpoint not configured")
         
-        if not api_key:
-            raise ValueError("Spark API token is required")
-        
-        # Extract work mode variant (for backward compatibility)
-        provider_variant = "cloud"  # Default
-        if os.getenv("SPACE_ID"):  # HuggingFace Space
-            provider_variant = "hfcloud"
-        
-        log(f"🚀 Initializing SparkLLM: {llm_config.endpoint}")
-        log(f"🔧 Provider variant: {provider_variant}")
+        # Determine variant based on environment
+        is_cloud = bool(os.environ.get("SPACE_ID"))
+        variant = "hfcloud" if is_cloud else "on-premise"
         
         return SparkLLM(
-            spark_endpoint=llm_config.endpoint,
+            spark_endpoint=endpoint,
             spark_token=api_key,
-            provider_variant=provider_variant,
+            provider_variant=variant,
             settings=llm_config.settings
         )
     
     @staticmethod
-    def _create_gpt_provider(llm_config, api_key: str, provider_def) -> GPT4oLLM:
-        """Create GPT-4 provider"""
-        if not api_key:
-            raise ValueError("OpenAI API key is required")
-        
-        # Get model-specific settings
-        settings = llm_config.settings or {}
-        model = provider_def.name  # gpt4o or gpt4o-mini
-        
-        log(f"🤖 Initializing GPT4oLLM with model: {model}")
-        
-        return GPT4oLLM(
+    def _create_gpt_provider(llm_config, api_key, provider_def):
+        """Create OpenAI GPT provider"""
+        return OpenAILLM(
             api_key=api_key,
-            model=model,
-            settings=settings
+            model=llm_config.name,
+            settings=llm_config.settings
         )
     
     @staticmethod
-    def _get_api_key(provider_name: str, config_key: Optional[str]) -> Optional[str]:
+    def _get_api_key(provider_name: str, configured_key: Optional[str]) -> str:
         """Get API key from config or environment"""
-        # First check config
-        if config_key:
-            if config_key.startswith("enc:"):
-                # Decrypt if encrypted
+        # First try configured key
+        if configured_key:
+            # Handle encrypted keys
+            if configured_key.startswith("enc:"):
                 from encryption_utils import decrypt
-                decrypted = decrypt(config_key)
-                log(f"🔓 Using encrypted API key from config: ***{decrypted[-4:]}")
-                return decrypted
-            else:
-                log(f"🔑 Using plain API key from config: ***{config_key[-4:]}")
-                return config_key
+                return decrypt(configured_key)
+            return configured_key
         
-        # Then check environment
+        # Then try environment variables
         env_mappings = {
             "spark": "SPARK_TOKEN",
             "gpt4o": "OPENAI_API_KEY",
@@ -104,19 +87,9 @@ class LLMFactory:
         
         env_var = env_mappings.get(provider_name)
         if env_var:
-            # Check if we're in HuggingFace Space
-            if os.getenv("SPACE_ID"):
-                # HuggingFace mode - direct environment
-                api_key = os.environ.get(env_var)
-                if api_key:
-                    log(f"🔑 Using API key from HuggingFace secrets: {env_var}")
-                    return api_key
-            else:
-                # Local mode - use dotenv
-                load_dotenv()
-                api_key = os.getenv(env_var)
-                if api_key:
-                    log(f"🔑 Using API key from .env: {env_var}")
-                    return api_key
+            key = os.environ.get(env_var)
+            if key:
+                log(f"📌 Using API key from environment: {env_var}")
+                return key
         
-        return None
+        raise ValueError(f"No API key found for provider: {provider_name}")