llm_interface.py

"""
LLM interface for Objective 5 — supports Gemini (cloud) and Ollama (local).

Gemini: uses Google's API with caching and rate limiting.
        Free-tier safety filters block most political content in this corpus.

Ollama: runs a local open-weight model (e.g. Llama 3.2) on the server GPU.
        No API key, no rate limits, no safety filters.
        Switch provider via config.LLM_PROVIDER = 'ollama'.
"""
import os
import time
from typing import List, Dict, Optional
import google.generativeai as genai
from diskcache import Cache
import config


class LLMInterface:
    """
    Unified LLM interface supporting Gemini (cloud) and Ollama (local).
    Provider is selected via config.LLM_PROVIDER.
    """

    def __init__(self, api_key: Optional[str] = None,
                 model_name: str = None,
                 temperature: float = None,
                 enable_cache: bool = None):
        """
        Initialize Gemini interface.

        Parameters:
        -----------
        api_key : Optional[str]
            Gemini API key (uses config if not provided)
        model_name : str
            Model name (uses config if not provided)
        temperature : float
            Generation temperature (uses config if not provided)
        enable_cache : bool
            Enable response caching (uses config if not provided)
        """
        self.api_key = api_key or config.GEMINI_API_KEY
        self.model_name = model_name or config.GEMINI_MODEL
        self.temperature = temperature if temperature is not None else config.LLM_TEMPERATURE
        self.enable_cache = enable_cache if enable_cache is not None else config.LLM_ENABLE_CACHE

        self.provider = config.LLM_PROVIDER  # 'gemini' or 'ollama'

        if self.provider == "ollama":
            # Ollama runs locally — no API key or SDK needed
            import requests as _requests
            self._requests = _requests
            self.ollama_model = config.OLLAMA_MODEL
            self.ollama_host  = config.OLLAMA_HOST
            print(f"Ollama initialized: {self.ollama_model} @ {self.ollama_host}")
        else:
            # Configure Gemini
            genai.configure(api_key=self.api_key)
            self.model = genai.GenerativeModel(self.model_name)
            print(f"Gemini initialized: {self.model_name}")

        # Initialize cache (shared across providers)
        if self.enable_cache:
            os.makedirs(config.LLM_CACHE_DIR, exist_ok=True)
            self.cache = Cache(config.LLM_CACHE_DIR)
        else:
            self.cache = None

    def _generate_ollama(self, prompt: str, max_tokens: Optional[int] = None) -> str:
        """Generate text using a local Ollama model (no filters, no rate limits)."""
        try:
            response = self._requests.post(
                f"{self.ollama_host}/api/generate",
                json={
                    "model":  self.ollama_model,
                    "prompt": prompt,
                    "stream": False,
                    "options": {
                        "temperature": self.temperature,
                        "num_predict": max_tokens or config.LLM_MAX_TOKENS,
                    },
                },
                timeout=120,
            )
            response.raise_for_status()
            return response.json().get("response", "").strip()
        except Exception as e:
            print(f"  Ollama error: {e}")
            return ""

    def generate(self, prompt: str, max_tokens: Optional[int] = None) -> str:
        """
        Generate text using the configured LLM provider (Gemini or Ollama).

        Parameters:
        -----------
        prompt : str
            Input prompt
        max_tokens : Optional[int]
            Maximum tokens to generate

        Returns:
        --------
        str
            Generated text
        """
        # Route to Ollama if configured
        if self.provider == "ollama":
            cache_key = f"ollama:{self.ollama_model}:{hash(prompt)}"
            if self.cache is not None and cache_key in self.cache:
                print("  [Cache hit]")
                return self.cache[cache_key]
            result = self._generate_ollama(prompt, max_tokens)
            if result and self.cache is not None:
                self.cache[cache_key] = result
            return result

        # ---- Gemini path (unchanged) ----
        # Check cache first
        if self.cache is not None:
            cache_key = f"{self.model_name}:{hash(prompt)}"
            if cache_key in self.cache:
                print("  [Cache hit]")
                return self.cache[cache_key]

        # Generate with retry logic
        max_retries = 3
        for attempt in range(max_retries):
            try:
                generation_config = {
                    'temperature': self.temperature,
                    'max_output_tokens': max_tokens or config.LLM_MAX_TOKENS,
                }

                response = self.model.generate_content(
                    prompt,
                    generation_config=generation_config
                )

                # Check if response has valid content
                if response.candidates and len(response.candidates) > 0:
                    candidate = response.candidates[0]

                    # Check finish reason
                    if candidate.finish_reason == 1:  # STOP - successful completion
                        if hasattr(response, 'text') and response.text:
                            result = response.text.strip()

                            # Cache result
                            if self.cache is not None:
                                self.cache[cache_key] = result

                            return result
                    elif candidate.finish_reason == 2:  # RECITATION
                        # Content was blocked due to recitation
                        print(f"  ⚠️ Warning: Content blocked (RECITATION)")
                        return "BLOCKED_RECITATION"
                    elif candidate.finish_reason == 3:  # SAFETY
                        # Content was blocked due to safety filters
                        print(f"  ⚠️ Warning: Content blocked (SAFETY)")
                        return "BLOCKED_SAFETY"
                    elif candidate.finish_reason == 4:  # MAX_TOKENS
                        # Response was truncated
                        if hasattr(response, 'text') and response.text:
                            print(f"  ⚠️ Warning: Response truncated (MAX_TOKENS)")
                            return response.text.strip()
                        return "BLOCKED_MAX_TOKENS"
                    else:
                        print(f"  ⚠️ Warning: Unknown finish reason: {candidate.finish_reason}")
                        return "BLOCKED_UNKNOWN"

                print(f"  ⚠️ Warning: No valid candidates in response")
                return "BLOCKED_NO_CANDIDATES"

            except Exception as e:
                error_str = str(e)

                # Check if it's a rate limit error
                if "429" in error_str or "quota" in error_str.lower():
                    if attempt < max_retries - 1:
                        # Extract retry delay from error message if available
                        import re
                        retry_match = re.search(r'retry in ([\d.]+)s', error_str)
                        if retry_match:
                            wait_time = float(retry_match.group(1)) + 1  # Add 1 second buffer
                        else:
                            wait_time = 10  # Default wait time for rate limits

                        print(f"  Rate limit hit, waiting {wait_time:.0f}s...")
                        time.sleep(wait_time)
                        continue
                    else:
                        return ""  # Give up after retries
                else:
                    # Other errors, don't retry
                    return ""

        return ""

    def generate_batch(self, prompts: List[str],
                      max_tokens: Optional[int] = None,
                      delay: float = 0.5) -> List[str]:
        """
        Generate multiple responses with rate limiting.

        Parameters:
        -----------
        prompts : List[str]
            List of prompts
        max_tokens : Optional[int]
            Maximum tokens per response
        delay : float
            Delay between requests (seconds)

        Returns:
        --------
        List[str]
            Generated responses
        """
        responses = []

        for i, prompt in enumerate(prompts):
            if i > 0:
                time.sleep(delay)  # Rate limiting

            print(f"Generating {i+1}/{len(prompts)}...")
            response = self.generate(prompt, max_tokens)
            responses.append(response)

        return responses

    def test_connection(self) -> bool:
        """
        Test Gemini API connection.

        Returns:
        --------
        bool
            True if connection successful
        """
        try:
            # Use a simple original prompt to avoid recitation filters
            response = self.generate("What is 2 plus 2?", max_tokens=20)
            # Just check if we got any response, even empty is OK (means API is working)
            return True
        except Exception as e:
            print(f"Connection test failed: {e}")
            return False

    def clear_cache(self):
        """Clear the response cache."""
        if self.cache is not None:
            self.cache.clear()
            print("Cache cleared")

    def get_cache_stats(self) -> Dict:
        """
        Get cache statistics.

        Returns:
        --------
        Dict
            Cache statistics
        """
        if self.cache is None:
            return {'enabled': False}

        return {
            'enabled': True,
            'size': len(self.cache),
            'directory': config.LLM_CACHE_DIR
        }