ai_handler.py

import json
import google.generativeai as genai
from typing import Dict, List, Optional
from config import GEMINI_API_KEY, GEMINI_MODEL
from ai_prompts import PromptTemplates


class AIHandler:
    def __init__(self):
        if not GEMINI_API_KEY:
            raise ValueError("GEMINI_API_KEY not found in the .env file")
        
        genai.configure(api_key=GEMINI_API_KEY)
        self.model = genai.GenerativeModel(GEMINI_MODEL)
        self.prompts = PromptTemplates()
    
    def _call_api(self, prompt: str, response_format: Optional[Dict] = None) -> str:
        try:
            if response_format:
                prompt = prompt + "\n\nIMPORTANT: Respond with valid JSON only, no markdown code blocks or additional text."
            
            response = self.model.generate_content(prompt)
            text = response.text
            
            if response_format:
                text = text.strip()
                if text.startswith("```json"):
                    text = text[7:]
                elif text.startswith("```"):
                    text = text[3:]
                if text.endswith("```"):
                    text = text[:-3]
                text = text.strip()
            
            return text
        except Exception as e:
            raise Exception(f"AI API call failed: {str(e)}")
    
    def generate_summary(self, text: str, summary_length: str, 
                        output_format: str, tone: str) -> str:
        prompt = self.prompts.create_summary_prompt(
            text, summary_length, output_format, tone
        )
        return self._call_api(prompt)
    
    def generate_summary_for_chunks(self, chunks: List[str], summary_length: str,
                                   output_format: str, tone: str) -> str:
        if len(chunks) == 1:
            return self.generate_summary(chunks[0], summary_length, output_format, tone)
        
        chunk_summaries = []
        for chunk in chunks:
            summary = self.generate_summary(chunk, summary_length, output_format, tone)
            chunk_summaries.append(summary)
        
        merge_prompt = self.prompts.create_merge_summary_prompt(
            chunk_summaries, output_format, tone
        )
        return self._call_api(merge_prompt)
    
    def extract_key_terms(self, text: str) -> Dict:
        prompt = self.prompts.create_key_terms_prompt(text)
        response = self._call_api(prompt, response_format={"type": "json_object"})
        
        try:
            return json.loads(response)
        except json.JSONDecodeError:
            return {
                "key_terms": [],
                "exam_topics": []
            }
    
    def generate_quiz(self, text: str, difficulty: str = "medium") -> Dict:
        prompt = self.prompts.create_quiz_prompt(text, difficulty)
        response = self._call_api(prompt, response_format={"type": "json_object"})
        
        try:
            data = json.loads(response)
            if not data.get("multiple_choice") and not data.get("short_answer"):
                raise ValueError("Quiz response is empty")
            return data
        except (json.JSONDecodeError, ValueError) as e:
            raise Exception(f"Quiz parsing failed: {str(e)}. Response: {response[:200]}")