core/llm_services.py

# storyverse_weaver/core/llm_services.py
import os
import google.generativeai as genai
from huggingface_hub import InferenceClient
# from dotenv import load_dotenv
# load_dotenv()

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 # 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 (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(f"SUCCESS: llm_services.py - Google Gemini API (for text) configured. Found models like: {models[0].name}")
        except Exception as 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:
        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 fallback text) initialized.")
        except Exception as 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:
        HF_TEXT_CONFIGURED = False
        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 # 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) # 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_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)
        
        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:
        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:
    # ... (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)
        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 (Gemini Primary for StoryVerseWeaver) - Module defined.")