Update core/llm_services.py

core/llm_services.py

@@ -2,160 +2,98 @@
 import os
 import google.generativeai as genai
 from huggingface_hub import InferenceClient
-# from dotenv import load_dotenv # Optional: for local .env file
-# load_dotenv() # Load environment variables from .env file if present
+# from dotenv import load_dotenv
+# load_dotenv()
 
-GOOGLE_API_KEY = os.getenv("STORYVERSE_GOOGLE_API_KEY") # Use specific env var names
-HF_TOKEN = os.getenv("STORYVERSE_HF_TOKEN")
+GOOGLE_API_KEY = os.getenv("STORYVERSE_GOOGLE_API_KEY")
+HF_TOKEN = os.getenv("STORYVERSE_HF_TOKEN") # For fallback
 
 GEMINI_TEXT_CONFIGURED = False
-HF_TEXT_CONFIGURED = False
+HF_TEXT_CONFIGURED = False # For fallback text model
 hf_inference_text_client = None
 
 class LLMTextResponse:
     def __init__(self, text=None, error=None, success=True, model_id_used="unknown_text_llm"):
         self.text, self.error, self.success, self.model_id_used = text, error, success, model_id_used
+    def __str__(self): return str(self.text) if self.success and self.text is not None else f"ERROR (Text Model: {self.model_id_used}): {self.error}"
 
-def initialize_text_llms():
-    global GOOGLE_API_KEY, HF_TOKEN, GEMINI_TEXT_CONFIGURED, HF_TEXT_CONFIGURED, hf_inference_text_client
-    print("INFO: llm_services.py - Initializing Text LLM clients...")
-    if GOOGLE_API_KEY and GOOGLE_API_KEY.strip():
-        try:
-            genai.configure(api_key=GOOGLE_API_KEY)
-            GEMINI_TEXT_CONFIGURED = True
-            print("SUCCESS: llm_services.py - Google Gemini API (for text) configured.")
-        except Exception as e:
-            print(f"ERROR: llm_services.py - Failed to configure Google Gemini API: {e}")
-            GEMINI_TEXT_CONFIGURED = False
-    else:
-        print("WARNING: llm_services.py - STORYVERSE_GOOGLE_API_KEY not found or empty.")
-        GEMINI_TEXT_CONFIGURED = False
-
-    if HF_TOKEN and HF_TOKEN.strip():
-        try:
-            hf_inference_text_client = InferenceClient(token=HF_TOKEN)
-            HF_TEXT_CONFIGURED = True
-            print("SUCCESS: llm_services.py - Hugging Face InferenceClient (for text) initialized.")
-        except Exception as e:
-            print(f"ERROR: llm_services.py - Failed to initialize HF InferenceClient: {e}")
-            HF_TEXT_CONFIGURED = False
-    else:
-        print("WARNING: llm_services.py - STORYVERSE_HF_TOKEN not found or empty.")
-        HF_TEXT_CONFIGURED = False
-    print(f"INFO: llm_services.py - Text LLM Init complete. Gemini Text: {GEMINI_TEXT_CONFIGURED}, HF Text: {HF_TEXT_CONFIGURED}")
-
-def is_gemini_text_ready(): return GEMINI_TEXT_CONFIGURED
-def is_hf_text_ready(): return HF_TEXT_CONFIGURED
-
-def generate_text_gemini(prompt: str, model_id: str = "gemini-1.5-flash-latest", system_prompt: str = None, temperature: float = 0.7, max_tokens: int = 512) -> LLMTextResponse:
-    if not is_gemini_text_ready():
-        return LLMTextResponse(error="Gemini text API not configured.", success=False, model_id_used=model_id)
-    try:
-        model = genai.GenerativeModel(model_name=model_id, system_instruction=system_prompt)
-        config = genai.types.GenerationConfig(temperature=temperature, max_output_tokens=max_tokens)
-        response = model.generate_content(prompt, generation_config=config)
-        # Add robust response checking as in AlgoForge's llm_clients.py
-        if response.prompt_feedback and response.prompt_feedback.block_reason:
-            return LLMTextResponse(error=f"Gemini: Prompt blocked ({response.prompt_feedback.block_reason})", success=False, model_id_used=model_id)
-        if not response.candidates or not response.candidates[0].content.parts:
-            return LLMTextResponse(error=f"Gemini: No content generated (Finish reason: {response.candidates[0].finish_reason if response.candidates else 'Unknown'})", success=False, model_id_used=model_id)
-        return LLMTextResponse(text=response.text, model_id_used=model_id)
-    except Exception as e:
-        return LLMTextResponse(error=f"Gemini API Error ({model_id}): {type(e).__name__} - {str(e)}", success=False, model_id_used=model_id)
-
-def generate_text_hf(prompt: str, model_id: str = "mistralai/Mistral-7B-Instruct-v0.2", system_prompt: str = None, temperature: float = 0.7, max_tokens: int = 512) -> LLMTextResponse:
-    if not is_hf_text_ready() or not hf_inference_text_client:
-        return LLMTextResponse(error="HF text API not configured.", success=False, model_id_used=model_id)
-    
-    full_prompt = f"<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n{prompt} [/INST]" if system_prompt else prompt
-    try:
-        use_sample = temperature > 0.001
-        response_text = hf_inference_text_client.text_generation(
-            full_prompt, model=model_id, max_new_tokens=max_tokens,
-            temperature=temperature if use_sample else None, do_sample=use_sample
-        )
-        return LLMTextResponse(text=response_text, model_id_used=model_id)
-    except Exception as e:
-        return LLMTextResponse(error=f"HF API Error ({model_id}): {type(e).__name__} - {str(e)}", success=False, model_id_used=model_id)
-
-print("DEBUG: core.llm_services (for StoryVerseWeaver) - Module defined.")# storyverse_weaver/core/llm_services.py
-import os
-import google.generativeai as genai
-from huggingface_hub import InferenceClient
-# from dotenv import load_dotenv # Optional: for local .env file
-# load_dotenv() # Load environment variables from .env file if present
-
-GOOGLE_API_KEY = os.getenv("STORYVERSE_GOOGLE_API_KEY") # Use specific env var names
-HF_TOKEN = os.getenv("STORYVERSE_HF_TOKEN")
-
-GEMINI_TEXT_CONFIGURED = False
-HF_TEXT_CONFIGURED = False
-hf_inference_text_client = None
-
-class LLMTextResponse:
-    def __init__(self, text=None, error=None, success=True, model_id_used="unknown_text_llm"):
-        self.text, self.error, self.success, self.model_id_used = text, error, success, model_id_used
 
 def initialize_text_llms():
     global GOOGLE_API_KEY, HF_TOKEN, GEMINI_TEXT_CONFIGURED, HF_TEXT_CONFIGURED, hf_inference_text_client
-    print("INFO: llm_services.py - Initializing Text LLM clients...")
+    print("INFO: llm_services.py - Initializing Text LLM clients (Gemini primary)...")
+    
+    # Google Gemini (Primary)
     if GOOGLE_API_KEY and GOOGLE_API_KEY.strip():
+        print("INFO: llm_services.py - STORYVERSE_GOOGLE_API_KEY found in environment.")
         try:
             genai.configure(api_key=GOOGLE_API_KEY)
+            # Simple test: list available models to confirm API key works and API is enabled
+            models = [m for m in genai.list_models() if 'generateContent' in m.supported_generation_methods and "gemini" in m.name]
+            if not models:
+                 raise Exception("No usable Gemini models found with this API key, or Generative Language API not fully enabled/propagated.")
             GEMINI_TEXT_CONFIGURED = True
-            print("SUCCESS: llm_services.py - Google Gemini API (for text) configured.")
+            print(f"SUCCESS: llm_services.py - Google Gemini API (for text) configured. Found models like: {models[0].name}")
         except Exception as e:
-            print(f"ERROR: llm_services.py - Failed to configure Google Gemini API: {e}")
             GEMINI_TEXT_CONFIGURED = False
+            print(f"ERROR: llm_services.py - Failed to configure/validate Google Gemini API.")
+            print(f"  Gemini Init Error Details: {type(e).__name__}: {e}")
     else:
-        print("WARNING: llm_services.py - STORYVERSE_GOOGLE_API_KEY not found or empty.")
         GEMINI_TEXT_CONFIGURED = False
+        print("WARNING: llm_services.py - STORYVERSE_GOOGLE_API_KEY not found or empty.")
 
+    # Hugging Face (Fallback)
     if HF_TOKEN and HF_TOKEN.strip():
+        print("INFO: llm_services.py - STORYVERSE_HF_TOKEN found (for fallback text model).")
         try:
             hf_inference_text_client = InferenceClient(token=HF_TOKEN)
             HF_TEXT_CONFIGURED = True
-            print("SUCCESS: llm_services.py - Hugging Face InferenceClient (for text) initialized.")
+            print("SUCCESS: llm_services.py - Hugging Face InferenceClient (for fallback text) initialized.")
         except Exception as e:
-            print(f"ERROR: llm_services.py - Failed to initialize HF InferenceClient: {e}")
             HF_TEXT_CONFIGURED = False
+            print(f"ERROR: llm_services.py - Failed to initialize HF InferenceClient for fallback text: {e}")
+            hf_inference_text_client = None # Ensure client is None on failure
     else:
-        print("WARNING: llm_services.py - STORYVERSE_HF_TOKEN not found or empty.")
         HF_TEXT_CONFIGURED = False
-    print(f"INFO: llm_services.py - Text LLM Init complete. Gemini Text: {GEMINI_TEXT_CONFIGURED}, HF Text: {HF_TEXT_CONFIGURED}")
+        print("WARNING: llm_services.py - STORYVERSE_HF_TOKEN not found or empty (for fallback text model).")
+    
+    print(f"INFO: llm_services.py - Text LLM Init complete. Gemini Text Ready: {GEMINI_TEXT_CONFIGURED}, HF Text (Fallback) Ready: {HF_TEXT_CONFIGURED}")
 
 def is_gemini_text_ready(): return GEMINI_TEXT_CONFIGURED
-def is_hf_text_ready(): return HF_TEXT_CONFIGURED
+def is_hf_text_ready(): return HF_TEXT_CONFIGURED # Still useful to know if fallback is available
 
 def generate_text_gemini(prompt: str, model_id: str = "gemini-1.5-flash-latest", system_prompt: str = None, temperature: float = 0.7, max_tokens: int = 512) -> LLMTextResponse:
     if not is_gemini_text_ready():
         return LLMTextResponse(error="Gemini text API not configured.", success=False, model_id_used=model_id)
+    print(f"DEBUG: llm_services.py - Calling Gemini ({model_id}) for text. System prompt: {'Yes' if system_prompt else 'No'}")
     try:
         model = genai.GenerativeModel(model_name=model_id, system_instruction=system_prompt)
         config = genai.types.GenerationConfig(temperature=temperature, max_output_tokens=max_tokens)
-        response = model.generate_content(prompt, generation_config=config)
-        # Add robust response checking as in AlgoForge's llm_clients.py
+        response = model.generate_content(prompt, generation_config=config) # Pass prompt directly
+
         if response.prompt_feedback and response.prompt_feedback.block_reason:
-            return LLMTextResponse(error=f"Gemini: Prompt blocked ({response.prompt_feedback.block_reason})", success=False, model_id_used=model_id)
+            return LLMTextResponse(error=f"Gemini: Prompt blocked ({response.prompt_feedback.block_reason_message or response.prompt_feedback.block_reason})", success=False, model_id_used=model_id)
         if not response.candidates or not response.candidates[0].content.parts:
             return LLMTextResponse(error=f"Gemini: No content generated (Finish reason: {response.candidates[0].finish_reason if response.candidates else 'Unknown'})", success=False, model_id_used=model_id)
-        return LLMTextResponse(text=response.text, model_id_used=model_id)
+        
+        generated_text = response.text # Simpler access for Gemini
+        print(f"DEBUG: llm_services.py - Gemini text generated successfully ({model_id}). Snippet: {generated_text[:50]}...")
+        return LLMTextResponse(text=generated_text, model_id_used=model_id)
     except Exception as e:
-        return LLMTextResponse(error=f"Gemini API Error ({model_id}): {type(e).__name__} - {str(e)}", success=False, model_id_used=model_id)
+        error_msg = f"Gemini API Error during text_generation ({model_id}): {type(e).__name__} - {str(e)}"
+        # Add specific checks for Google API errors
+        if "API key not valid" in str(e) or "PERMISSION_DENIED" in str(e):
+            error_msg += " Check your GOOGLE_API_KEY and ensure Generative Language API is enabled in Google Cloud."
+        print(f"ERROR: llm_services.py - {error_msg}")
+        return LLMTextResponse(error=error_msg, success=False, model_id_used=model_id)
 
 def generate_text_hf(prompt: str, model_id: str = "mistralai/Mistral-7B-Instruct-v0.2", system_prompt: str = None, temperature: float = 0.7, max_tokens: int = 512) -> LLMTextResponse:
-    if not is_hf_text_ready() or not hf_inference_text_client:
-        return LLMTextResponse(error="HF text API not configured.", success=False, model_id_used=model_id)
-    
+    # ... (This function remains the same as before, for fallback)
+    if not is_hf_text_ready() or not hf_inference_text_client: return LLMTextResponse(error="HF text API not configured.", success=False, model_id_used=model_id)
     full_prompt = f"<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n{prompt} [/INST]" if system_prompt else prompt
     try:
         use_sample = temperature > 0.001
-        response_text = hf_inference_text_client.text_generation(
-            full_prompt, model=model_id, max_new_tokens=max_tokens,
-            temperature=temperature if use_sample else None, do_sample=use_sample
-        )
+        response_text = hf_inference_text_client.text_generation(full_prompt, model=model_id, max_new_tokens=max_tokens, temperature=temperature if use_sample else None, do_sample=use_sample)
         return LLMTextResponse(text=response_text, model_id_used=model_id)
-    except Exception as e:
-        return LLMTextResponse(error=f"HF API Error ({model_id}): {type(e).__name__} - {str(e)}", success=False, model_id_used=model_id)
+    except Exception as e: return LLMTextResponse(error=f"HF API Error ({model_id}): {type(e).__name__} - {str(e)}", success=False, model_id_used=model_id)
 
-print("DEBUG: core.llm_services (for StoryVerseWeaver) - Module defined.")
+print("DEBUG: core.llm_services (Gemini Primary for StoryVerseWeaver) - Module defined.")