models.py

"""
LLM Model interface abstraction for Gemini API.
Handles API calls with rate limiting, error handling, and logging.
"""

import logging
import time
from typing import Optional
from abc import ABC, abstractmethod
import google.generativeai as genai
from google.api_core import retry


logger = logging.getLogger(__name__)


class LLMModel(ABC):
    """Abstract base class for LLM models."""
    
    @abstractmethod
    def generate(self, prompt: str) -> str:
        """Generate text based on prompt."""
        pass
    
    @abstractmethod
    def generate_with_system(self, system_prompt: str, user_prompt: str) -> str:
        """Generate text with system and user prompts."""
        pass


class GeminiModel(LLMModel):
    """Interface for Google Gemini API."""
    
    def __init__(self, api_key: str, model_name: str = "gemini-1.5-pro",
                 temperature: float = 0.7, max_tokens: int = 4096,
                 timeout_seconds: int = 120):
        """
        Initialize Gemini model.
        
        Args:
            api_key: Gemini API key
            model_name: Model name (e.g., "gemini-1.5-pro")
            temperature: Generation temperature (0.0-1.0)
            max_tokens: Maximum tokens in response
            timeout_seconds: API call timeout
        """
        self.api_key = api_key
        self.model_name = model_name
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.timeout_seconds = timeout_seconds
        
        # Initialize Gemini
        genai.configure(api_key=self.api_key)
        self.model = genai.GenerativeModel(
            model_name=self.model_name,
            generation_config={
                "temperature": self.temperature,
                "max_output_tokens": self.max_tokens,
            }
        )
        
        logger.info(f"Initialized Gemini model: {self.model_name}")
    
    def generate(self, prompt: str) -> str:
        """
        Generate text from a single prompt.
        
        Args:
            prompt: Input prompt text
            
        Returns:
            Generated text response
            
        Raises:
            RuntimeError: If API call fails
        """
        try:
            logger.debug(f"Calling Gemini with prompt (length: {len(prompt)})")
            
            response = self.model.generate_content(
                prompt,
                safety_settings=[
                    {
                        "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
                        "threshold": "BLOCK_NONE"
                    },
                    {
                        "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
                        "threshold": "BLOCK_NONE"
                    },
                    {
                        "category": "HARM_CATEGORY_HARASSMENT",
                        "threshold": "BLOCK_NONE"
                    },
                    {
                        "category": "HARM_CATEGORY_HATE_SPEECH",
                        "threshold": "BLOCK_NONE"
                    },
                ]
            )
            
            text = response.text
            
            # Log if enabled
            if logger.isEnabledFor(logging.INFO):
                logger.info(f"Gemini prompt: {prompt[:200]}...")
                logger.info(f"Gemini response: {text[:200]}...")
            
            return text
            
        except Exception as e:
            logger.error(f"Gemini API error: {str(e)}")
            raise RuntimeError(f"Failed to generate text from Gemini: {str(e)}")
    
    def generate_with_system(self, system_prompt: str, user_prompt: str) -> str:
        """
        Generate text with system and user prompts.
        Gemini 1.5 supports system prompts natively via GenerativeModel initialization.
        
        Args:
            system_prompt: System context/instructions
            user_prompt: User's actual request
            
        Returns:
            Generated text response
            
        Raises:
            RuntimeError: If API call fails
        """
        try:
            logger.debug(
                f"Calling Gemini with system prompt (length: {len(system_prompt)}) "
                f"and user prompt (length: {len(user_prompt)})"
            )
            
            # Create a new GenerativeModel instance with system_instruction
            # system_instruction must be passed during model initialization, not generate_content()
            model_with_system = genai.GenerativeModel(
                model_name=self.model_name,
                generation_config={
                    "temperature": self.temperature,
                    "max_output_tokens": self.max_tokens,
                },
                system_instruction=system_prompt
            )
            
            response = model_with_system.generate_content(
                user_prompt,
                safety_settings=[
                    {
                        "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
                        "threshold": "BLOCK_NONE"
                    },
                    {
                        "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
                        "threshold": "BLOCK_NONE"
                    },
                    {
                        "category": "HARM_CATEGORY_HARASSMENT",
                        "threshold": "BLOCK_NONE"
                    },
                    {
                        "category": "HARM_CATEGORY_HATE_SPEECH",
                        "threshold": "BLOCK_NONE"
                    },
                ]
            )
            
            text = response.text
            
            # Log if enabled
            if logger.isEnabledFor(logging.INFO):
                logger.info(f"Gemini system prompt: {system_prompt[:200]}...")
                logger.info(f"Gemini user prompt: {user_prompt[:200]}...")
                logger.info(f"Gemini response: {text[:200]}...")
            
            return text
            
        except Exception as e:
            logger.error(f"Gemini API error with system prompt: {str(e)}")
            raise RuntimeError(
                f"Failed to generate text from Gemini with system prompt: {str(e)}"
            )


class RateLimitedGeminiModel(GeminiModel):
    """Gemini model with rate limiting support."""
    
    def __init__(self, *args, **kwargs):
        """Initialize with rate limiting."""
        super().__init__(*args, **kwargs)
        self.min_interval_between_calls = 1.0  # seconds
        self.last_call_time = 0
    
    def _wait_for_rate_limit(self) -> None:
        """Enforce minimum interval between API calls."""
        elapsed = time.time() - self.last_call_time
        if elapsed < self.min_interval_between_calls:
            wait_time = self.min_interval_between_calls - elapsed
            logger.debug(f"Rate limiting: waiting {wait_time:.2f}s")
            time.sleep(wait_time)
    
    def generate(self, prompt: str) -> str:
        """Generate with rate limiting."""
        self._wait_for_rate_limit()
        self.last_call_time = time.time()
        return super().generate(prompt)
    
    def generate_with_system(self, system_prompt: str, user_prompt: str) -> str:
        """Generate with system prompt and rate limiting."""
        self._wait_for_rate_limit()
        self.last_call_time = time.time()
        return super().generate_with_system(system_prompt, user_prompt)


# Factory function
def create_model(config: dict) -> LLMModel:
    """
    Create an LLM model based on configuration.
    
    Args:
        config: Configuration dictionary with API settings
        
    Returns:
        LLMModel instance
    """
    provider = config.get("provider", "gemini").lower()
    
    if provider == "gemini":
        api_key = config.get("api_key")
        if not api_key:
            raise ValueError("Gemini API key not provided in config")
        
        return RateLimitedGeminiModel(
            api_key=api_key,
            model_name=config.get("model", "gemini-1.5-pro"),
            temperature=config.get("temperature", 0.7),
            max_tokens=config.get("max_tokens", 4096),
            timeout_seconds=config.get("timeout_seconds", 120),
        )
    else:
        raise ValueError(f"Unsupported LLM provider: {provider}")