gemini_client.py

"""
Gemini API Client for Math Tutor Bot
Integrates Google's Gemini API with fallback to local Phi-3 Mini
"""

import google.generativeai as genai
from google import genai as genai_new
import os
from typing import Optional, List, Dict
import traceback

class GeminiClient:
    def __init__(self, api_key: Optional[str] = None):
        """
        Initialize Gemini API client

        Args:
            api_key: Google AI Studio API key (if None, loads from environment)
        """
        self.api_key = api_key or os.getenv('GEMINI_API_KEY')
        self.model = None
        self.enabled = False

        if self.api_key:
            try:
                genai.configure(api_key=self.api_key)
                # Use Gemini 3 Flash Preview for latest, most powerful responses
                self.model = genai.GenerativeModel('gemini-3-flash-preview')
                self.enabled = True
                print("[Gemini Client] ✓ Initialized successfully with Gemini 3 Flash Preview")
            except Exception as e:
                print(f"[Gemini Client] ✗ Failed to initialize: {e}")
                self.enabled = False
        else:
            print("[Gemini Client] ⚠ No API key provided, Gemini disabled")

    def generate_response(self,
                         messages: List[Dict[str, str]],
                         max_tokens: int = 1024,
                         temperature: float = 0.7) -> Optional[str]:
        """
        Generate a response using Gemini API

        Args:
            messages: List of message dicts with 'role' and 'content' keys
            max_tokens: Maximum tokens in response
            temperature: Sampling temperature (0.0 to 1.0)

        Returns:
            Generated text response, or None if failed
        """
        if not self.enabled:
            return None

        try:
            # Convert messages to Gemini format
            # Gemini uses 'user' and 'model' roles
            gemini_messages = []
            system_prompt = None

            for msg in messages:
                role = msg['role']
                content = msg['content']

                if role == 'system':
                    # Gemini doesn't have a system role, so we'll prepend it to the first user message
                    system_prompt = content
                elif role == 'user':
                    if system_prompt:
                        # Prepend system prompt to first user message
                        content = f"{system_prompt}\n\n{content}"
                        system_prompt = None
                    gemini_messages.append({'role': 'user', 'parts': [content]})
                elif role == 'assistant':
                    gemini_messages.append({'role': 'model', 'parts': [content]})

            # Generate response
            generation_config = genai.types.GenerationConfig(
                max_output_tokens=max_tokens,
                temperature=temperature,
            )

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

            # Debug: Check finish reason
            if hasattr(response, 'candidates') and response.candidates:
                finish_reason = response.candidates[0].finish_reason
                print(f"[Gemini Client] Finish reason: {finish_reason}")

            return response.text

        except Exception as e:
            print(f"[Gemini Client] Error generating response: {e}")
            traceback.print_exc()
            return None

    def generate_simple(self, prompt: str, max_tokens: int = 1024) -> Optional[str]:
        """
        Generate a simple response from a single prompt (for diagram generation)

        Args:
            prompt: Single prompt string
            max_tokens: Maximum tokens in response

        Returns:
            Generated text response, or None if failed
        """
        if not self.enabled:
            return None

        try:
            generation_config = genai.types.GenerationConfig(
                max_output_tokens=max_tokens,
                temperature=0.3,  # Lower temperature for code generation
            )

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

            return response.text

        except Exception as e:
            print(f"[Gemini Client] Error generating simple response: {e}")
            traceback.print_exc()
            return None

    def is_available(self) -> bool:
        """Check if Gemini API is available"""
        return self.enabled

    def generate_image(self, prompt: str) -> Optional[bytes]:
        """
        Generate an image using Gemini 2.5 Flash Image model

        Args:
            prompt: Description of the image to generate

        Returns:
            Image bytes (PNG format), or None if failed
        """
        if not self.enabled:
            return None

        try:
            # Create new client for image generation
            client = genai_new.Client(api_key=self.api_key)

            response = client.models.generate_content(
                model="gemini-2.0-flash-exp",
                contents=[prompt],
            )

            # Extract image from response
            for part in response.parts:
                if part.inline_data is not None:
                    # Return the image bytes directly
                    return part.inline_data.data

            print("[Gemini Client] No image in response")
            return None

        except Exception as e:
            print(f"[Gemini Client] Error generating image: {e}")
            traceback.print_exc()
            return None


# Global client instance
_gemini_client = None

def get_gemini_client() -> GeminiClient:
    """Get or create global Gemini client instance"""
    global _gemini_client
    if _gemini_client is None:
        _gemini_client = GeminiClient()
    return _gemini_client


def initialize_gemini(api_key: str) -> GeminiClient:
    """
    Initialize Gemini client with specific API key

    Args:
        api_key: Google AI Studio API key

    Returns:
        Initialized GeminiClient instance
    """
    global _gemini_client
    _gemini_client = GeminiClient(api_key=api_key)
    return _gemini_client


if __name__ == "__main__":
    # Test the Gemini client
    print("Testing Gemini Client")
    print("=" * 60)

    # Test with environment variable
    client = GeminiClient()

    if client.is_available():
        # Test simple generation
        test_prompt = "Explain the Pythagorean theorem in one sentence."
        response = client.generate_simple(test_prompt)

        if response:
            print(f"✓ Test passed!")
            print(f"Prompt: {test_prompt}")
            print(f"Response: {response}")
        else:
            print("✗ Test failed: No response generated")
    else:
        print("✗ Gemini not available - set GEMINI_API_KEY environment variable")