src/core/gemini_llm.py

import google.generativeai as genai
import os
import logging

logger = logging.getLogger(__name__)

# --- Gemini Client Initialization ---

def get_gemini_client():
    """Initializes and returns the Gemini client."""
    api_key = os.environ.get("GEMINI_API_KEY")
    if not api_key:
        raise ValueError("GEMINI_API_KEY environment variable is required")
    
    genai.configure(api_key=api_key)
    
    # List available models for debugging
    try:
        available_models = genai.list_models()
        logger.info("Available Gemini models:")
        for model in available_models:
            if 'generateContent' in model.supported_generation_methods:
                logger.info(f"  - {model.name}")
    except Exception as e:
        logger.warning(f"Could not list available models: {str(e)}")
    
    # Try different model names in order of preference
    # The model names change based on the google-generativeai version
    model_names = [
        'gemini-1.5-flash',      # Current free tier model
        'gemini-1.5-pro',        # Alternative model
        'models/gemini-1.5-flash', # With models/ prefix
        'models/gemini-1.5-pro',   # With models/ prefix
    ]
    
    for model_name in model_names:
        try:
            model = genai.GenerativeModel(model_name)
            logger.info(f"Successfully initialized Gemini model: {model_name}")
            return model
        except Exception as e:
            logger.warning(f"Failed to initialize model {model_name}: {str(e)}")
            continue
    
    # If all fail, raise an error
    raise ValueError("Could not initialize any Gemini model. Please check your API key and model availability.")

# --- Prompt Generation ---

def format_prompt(query: str, context: list[dict]) -> str:
    """Formats the prompt for the LLM with the retrieved context."""
    context_str = "\n".join([item.payload.get('text') for item in context])
    prompt = f"""**Instruction**:
Answer the user's query based *only* on the provided context.
If the context does not contain the answer, state that you cannot answer the question with the given information.
Do not use any prior knowledge.

**Context**:
{context_str}

**Query**:
{query}

**Answer**:
"""
    return prompt

# --- LLM Interaction ---

def generate_response(client, prompt: str):
    """Generates a response from the Gemini LLM."""
    try:
        response = client.generate_content(prompt)
        return response.text
    except Exception as e:
        logger.error(f"Gemini API error: {str(e)}")
        raise Exception(f"Failed to generate response: {str(e)}")

# --- Embedding Functions ---

def get_embeddings(texts: list[str]) -> list[list[float]]:
    """Generate embeddings for a list of texts using Gemini."""
    api_key = os.environ.get("GEMINI_API_KEY")
    if not api_key:
        raise ValueError("GEMINI_API_KEY environment variable is required")
    
    genai.configure(api_key=api_key)
    
    embeddings = []
    for text in texts:
        try:
            # Use Gemini's embedding model
            result = genai.embed_content(
                model="models/embedding-001",
                content=text,
                task_type="retrieval_document"
            )
            embeddings.append(result['embedding'])
        except Exception as e:
            logger.error(f"Failed to generate embedding for text: {str(e)}")
            # Return a zero vector as fallback
            embeddings.append([0.0] * 768)  # Gemini embeddings are 768-dimensional
    
    return embeddings