src/embeddings.py

from google import genai
import os
import logging
from typing import List, Optional
from src.config import Config

logger = logging.getLogger("Embeddings")


def generate_embedding(text: str) -> Optional[List[float]]:
    """
    Generates a text embedding using Gemini with the new google.genai SDK.
    """
    api_key = os.environ.get("GEMINI_API_KEY") or Config.GEMINI_API_KEY
    
    if not api_key:
        logger.warning("GEMINI_API_KEY not set. Skipping embedding.")
        return None
        
    try:
        client = genai.Client(api_key=api_key)
        
        # Truncate text if too long (Gemini limit is usually high, but let's be safe)
        text = text[:8000]
        
        result = client.models.embed_content(
            model="text-embedding-004",
            contents=text
        )
        
        # New SDK returns embeddings in a different structure
        if result.embeddings and len(result.embeddings) > 0:
            return result.embeddings[0].values
        return None
        
    except Exception as e:
        logger.warning(f"Embedding Generation Failed: {e}")
        return None