Update utils/llm_generator.py

changed settings for Gemini model selection

utils/llm_generator.py

@@ -1,297 +1,297 @@
-"""
-Real LLM-based generator using Groq or Google Gemini API.
-This ACTUALLY generates responses (unlike SimpleGenerator which just extracts text).
-"""
-
-import os
-from typing import List, Dict, Optional
-import streamlit as st
-
-try:
-    from groq import Groq
-    GROQ_AVAILABLE = True
-except ImportError:
-    GROQ_AVAILABLE = False
-
-try:
-    import google.generativeai as genai
-    GEMINI_AVAILABLE = True
-except ImportError:
-    GEMINI_AVAILABLE = False
-
-
-class LLMGenerator:
-    """
-    Actual LLM-based response generation using Groq (Llama-3-70B) or Gemini.
-    This is what NotebookLM uses - real AI generation, not text extraction.
-    """
-    
-    def __init__(self, provider: str = "groq", api_key: Optional[str] = None):
-        """
-        Initialize LLM generator.
-        
-        Args:
-            provider: "groq" or "gemini"
-            api_key: API key (if None, reads from environment or asks user)
-        """
-        self.provider = provider
-        self.client = None
-        self.ready = False
-        
-        # Get API key
-        if api_key:
-            self.api_key = api_key
-        elif provider == "groq":
-            self.api_key = os.getenv("GROQ_API_KEY", "")
-        elif provider == "gemini":
-            self.api_key = os.getenv("GEMINI_API_KEY", "")
-        else:
-            self.api_key = ""
-        
-        # Initialize client
-        self._initialize_client()
-    
-    def _initialize_client(self):
-        """Initialize the LLM client."""
-        if not self.api_key:
-            return
-        
-        try:
-            if self.provider == "groq" and GROQ_AVAILABLE:
-                # Initialize Groq client with explicit parameters
-                # Avoid potential proxies kwarg issue by not passing extra config
-                import os
-                os.environ["GROQ_API_KEY"] = self.api_key
-                self.client = Groq()  # Will read from environment
-                self.ready = True
-            elif self.provider == "gemini" and GEMINI_AVAILABLE:
-                genai.configure(api_key=self.api_key)
-                self.client = genai.GenerativeModel('gemini-1.5-flash')
-                self.ready = True
-        except Exception as e:
-            print(f"Failed to initialize {self.provider}: {e}")
-            self.ready = False
-    
-    def set_api_key(self, api_key: str):
-        """Update API key and reinitialize."""
-        self.api_key = api_key
-        self._initialize_client()
-    
-    def generate_response(
-        self,
-        prompt: str,
-        context: str = "",
-        use_case: str = "explanation",
-        metadatas: List[Dict] = None,
-        temperature: float = 0.7,
-        max_tokens: int = 1500,
-        **kwargs
-    ) -> str:
-        """
-        Generate response using actual LLM (NotebookLM-style).
-        
-        Args:
-            prompt: User's question
-            context: Retrieved context from documents
-            use_case: Response type (explanation, summary, qa, notes)
-            metadatas: Metadata for citations
-            temperature: LLM temperature (0.0-1.0)
-            max_tokens: Maximum response length
-            
-        Returns:
-            Generated response with inline citations
-        """
-        if not self.ready:
-            return (
-                "⚠️ **LLM not configured.** Please add your API key in the sidebar.\n\n"
-                "Get a free key:\n"
-                "- **Groq** (recommended, very fast): https://console.groq.com/keys\n"
-                "- **Gemini** (Google): https://makersuite.google.com/app/apikey"
-            )
-        
-        if not context:
-            return (
-                "I don't have enough information from your uploaded documents to answer this question. "
-                "Please upload relevant study materials first."
-            )
-        
-        # Build NotebookLM-style system prompt with strict source grounding
-        system_prompt = self._build_system_prompt(use_case)
-        
-        # Build user message with context
-        user_message = self._build_user_message(prompt, context, metadatas)
-        
-        try:
-            # Generate with LLM
-            if self.provider == "groq":
-                response = self._generate_groq(system_prompt, user_message, temperature, max_tokens)
-            elif self.provider == "gemini":
-                response = self._generate_gemini(system_prompt, user_message, temperature, max_tokens)
-            else:
-                return "Error: Unknown provider"
-            
-            return response
-        
-        except Exception as e:
-            return f"Error generating response: {str(e)}\n\nPlease check your API key and try again."
-    
-    def _build_system_prompt(self, use_case: str) -> str:
-        """Build specialized system prompt based on use case."""
-        base_prompt = (
-            "You are an expert academic assistant for students, acting like a highly intelligent study buddy. "
-            "⚠️ CRITICAL RULE: You MUST ONLY use information from the provided context below. "
-            "DO NOT use your training knowledge. DO NOT infer beyond what's explicitly stated. "
-            "If the context doesn't contain adequate information to answer the question, you MUST respond: "
-            "'I cannot find sufficient information about this in the uploaded documents. Please upload materials covering this topic or rephrase your question.'\n\n"
-            "⚠️ GROUNDING REQUIREMENT: Every statement must be traceable to the provided context. "
-            "If you cannot find it in the context below, DO NOT answer from general knowledge.\n\n"
-            "✨ FORMATTING RULES (NotebookLM Style):\n"
-            "- Use clean, hierarchical Markdown (### Headers, **Bold** terms).\n"
-            "- Break down long paragraphs into easily readable bullet points.\n"
-            "- Be direct and concise. Avoid conversational fluff like 'Certainly!' or 'Here is the answer'.\n"
-            "- If applicable to the prompt, always try to extract a **Real-World Example** from the text to aid understanding.\n\n"
-        )
-        
-        if use_case == "explanation":
-            base_prompt += (
-                "**Your task:** Explain the concept in a clear, step-by-step manner suitable for students.\n"
-                "1. Start with a concise, one-sentence definition.\n"
-                "2. Break down the core mechanics or components using bullet points.\n"
-                "3. Provide an example (only if found in the text).\n"
-                "4. Add a 'Key Takeaway' at the end.\n"
-            )
-        elif use_case == "summary":
-            base_prompt += (
-                "**Your task:** Create a highly structured summary.\n"
-                "- Start with a brief high-level overview (2 sentences max).\n"
-                "- Use '### Key Themes' and list the main points as bulleted items.\n"
-                "- Keep each point concise but factually dense.\n"
-            )
-        elif use_case == "qa":
-            base_prompt += (
-                "**Your task:** Answer the question directly and comprehensively.\n"
-                "- Provide the direct answer immediately in the first sentence.\n"
-                "- Use numbered lists or bullet points to provide supporting details from the context.\n"
-                "- Use **bold** for key facts, numbers, and formulas.\n"
-            )
-        elif use_case == "notes":
-            base_prompt += (
-                "**Your task:** Create comprehensive, structured study notes.\n"
-                "- Use clear section headers (###).\n"
-                "- Organize information hierarchically (using nested bullet points).\n"
-                "- Explicitly highlight **Definitions**, **Formulas**, and **Important Dates/Names**.\n"
-            )
-        
-        base_prompt += (
-            "\n**Citation Rules:**\n"
-            "- You MUST cite your source at the end of every major claim or paragraph using numbered brackets like **[1]**, **[2]** based on the Source number provided in the context.\n"
-            "- If a claim comes from multiple sources, use **[1, 2]**.\n"
-            "- Do NOT use the document filename in the citation, ONLY the number.\n"
-            "- Do NOT make up information - stick strictly to the provided context.\n"
-        )
-        
-        return base_prompt
-    
-    def _build_user_message(self, prompt: str, context: str, metadatas: List[Dict] = None) -> str:
-        """Build user message with context and question."""
-        # Extract source names from metadata
-        sources = []
-        if metadatas:
-            for meta in metadatas:
-                filename = meta.get('filename', 'Unknown')
-                clean_name = filename.replace('.pdf', '').replace('.docx', '').replace('.txt', '')
-                if clean_name not in sources:
-                    sources.append(clean_name)
-        
-        message = "**Available Sources (USE ONLY THESE):**\n"
-        for source in sources[:5]:  # Show up to 5 sources
-            message += f"- {source}\n"
-        
-        message += f"\n**===== START OF CONTEXT (ANSWER ONLY FROM THIS) =====**\n\n{context}\n\n"
-        message += f"**===== END OF CONTEXT =====**\n\n"
-        message += f"**Student's Question:** {prompt}\n\n"
-        message += "**Instructions:** Answer ONLY using the context between the markers above. If the context doesn't contain the answer, say you don't have that information. Cite sources in brackets."
-        
-        return message
-    
-    def _generate_groq(self, system_prompt: str, user_message: str, temperature: float, max_tokens: int) -> str:
-        """Generate using Groq API (Llama-3.3-70B)."""
-        completion = self.client.chat.completions.create(
-            model="llama-3.3-70b-versatile",  # Latest 70B model (Dec 2024)
-            messages=[
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": user_message}
-            ],
-            temperature=temperature,
-            max_tokens=max_tokens,
-            top_p=0.95,
-            stream=False
-        )
-        
-        return completion.choices[0].message.content
-    
-    def _generate_gemini(self, system_prompt: str, user_message: str, temperature: float, max_tokens: int) -> str:
-        """Generate using Google Gemini API."""
-        full_prompt = f"{system_prompt}\n\n{user_message}"
-        
-        response = self.client.generate_content(
-            full_prompt,
-            generation_config=genai.GenerationConfig(
-                temperature=temperature,
-                max_output_tokens=max_tokens,
-                top_p=0.95
-            )
-        )
-        
-        return response.text
-    
-    def is_ready(self) -> bool:
-        """Check if LLM is ready to generate."""
-        return self.ready
-    
-    def get_provider(self) -> str:
-        """Get current provider name."""
-        if self.provider == "groq":
-            return "Groq (Llama-3.3-70B)"
-        elif self.provider == "gemini":
-            return "Google Gemini 1.5 Flash"
-        return "Unknown"
-    
-    def generate(self, prompt: str, temperature: float = 0.3, max_tokens: int = 1500) -> str:
-        """
-        Simple wrapper for backend compatibility.
-        Generates response from a complete prompt that already includes context.
-        
-        Args:
-            prompt: Complete prompt with context already embedded
-            temperature: LLM temperature (0.0-1.0)
-            max_tokens: Maximum response length
-            
-        Returns:
-            Generated response
-        """
-        if not self.ready:
-            return (
-                "⚠️ **LLM not configured.** Please add your API key.\n\n"
-                "Get a free key:\n"
-                "- **Groq** (recommended, very fast): https://console.groq.com/keys\n"
-                "- **Gemini** (Google): https://makersuite.google.com/app/apikey"
-            )
-        
-        try:
-            if self.provider == "groq":
-                return self._generate_groq(
-                    system_prompt="You are a helpful AI assistant.",
-                    user_message=prompt,
-                    temperature=temperature,
-                    max_tokens=max_tokens
-                )
-            elif self.provider == "gemini":
-                return self._generate_gemini(
-                    system_prompt="You are a helpful AI assistant.",
-                    user_message=prompt,
-                    temperature=temperature,
-                    max_tokens=max_tokens
-                )
-        except Exception as e:
-            return f"Error generating response: {str(e)}"
+"""
+Real LLM-based generator using Groq or Google Gemini API.
+This ACTUALLY generates responses (unlike SimpleGenerator which just extracts text).
+"""
+
+import os
+from typing import List, Dict, Optional
+import streamlit as st
+
+try:
+    from groq import Groq
+    GROQ_AVAILABLE = True
+except ImportError:
+    GROQ_AVAILABLE = False
+
+try:
+    import google.generativeai as genai
+    GEMINI_AVAILABLE = True
+except ImportError:
+    GEMINI_AVAILABLE = False
+
+
+class LLMGenerator:
+    """
+    Actual LLM-based response generation using Groq (Llama-3-70B) or Gemini.
+    This is what NotebookLM uses - real AI generation, not text extraction.
+    """
+    
+    def __init__(self, provider: str = "groq", api_key: Optional[str] = None):
+        """
+        Initialize LLM generator.
+        
+        Args:
+            provider: "groq" or "gemini"
+            api_key: API key (if None, reads from environment or asks user)
+        """
+        self.provider = provider
+        self.client = None
+        self.ready = False
+        
+        # Get API key
+        if api_key:
+            self.api_key = api_key
+        elif provider == "groq":
+            self.api_key = os.getenv("GROQ_API_KEY", "")
+        elif provider == "gemini":
+            self.api_key = os.getenv("GEMINI_API_KEY", "")
+        else:
+            self.api_key = ""
+        
+        # Initialize client
+        self._initialize_client()
+    
+    def _initialize_client(self):
+        """Initialize the LLM client."""
+        if not self.api_key:
+            return
+        
+        try:
+            if self.provider == "groq" and GROQ_AVAILABLE:
+                # Initialize Groq client with explicit parameters
+                # Avoid potential proxies kwarg issue by not passing extra config
+                import os
+                os.environ["GROQ_API_KEY"] = self.api_key
+                self.client = Groq()  # Will read from environment
+                self.ready = True
+            elif self.provider == "gemini" and GEMINI_AVAILABLE:
+                genai.configure(api_key=self.api_key)
+                self.client = genai.GenerativeModel('gemini-2.5-flash')
+                self.ready = True
+        except Exception as e:
+            print(f"Failed to initialize {self.provider}: {e}")
+            self.ready = False
+    
+    def set_api_key(self, api_key: str):
+        """Update API key and reinitialize."""
+        self.api_key = api_key
+        self._initialize_client()
+    
+    def generate_response(
+        self,
+        prompt: str,
+        context: str = "",
+        use_case: str = "explanation",
+        metadatas: List[Dict] = None,
+        temperature: float = 0.7,
+        max_tokens: int = 1500,
+        **kwargs
+    ) -> str:
+        """
+        Generate response using actual LLM (NotebookLM-style).
+        
+        Args:
+            prompt: User's question
+            context: Retrieved context from documents
+            use_case: Response type (explanation, summary, qa, notes)
+            metadatas: Metadata for citations
+            temperature: LLM temperature (0.0-1.0)
+            max_tokens: Maximum response length
+            
+        Returns:
+            Generated response with inline citations
+        """
+        if not self.ready:
+            return (
+                "⚠️ **LLM not configured.** Please add your API key in the sidebar.\n\n"
+                "Get a free key:\n"
+                "- **Groq** (recommended, very fast): https://console.groq.com/keys\n"
+                "- **Gemini** (Google): https://makersuite.google.com/app/apikey"
+            )
+        
+        if not context:
+            return (
+                "I don't have enough information from your uploaded documents to answer this question. "
+                "Please upload relevant study materials first."
+            )
+        
+        # Build NotebookLM-style system prompt with strict source grounding
+        system_prompt = self._build_system_prompt(use_case)
+        
+        # Build user message with context
+        user_message = self._build_user_message(prompt, context, metadatas)
+        
+        try:
+            # Generate with LLM
+            if self.provider == "groq":
+                response = self._generate_groq(system_prompt, user_message, temperature, max_tokens)
+            elif self.provider == "gemini":
+                response = self._generate_gemini(system_prompt, user_message, temperature, max_tokens)
+            else:
+                return "Error: Unknown provider"
+            
+            return response
+        
+        except Exception as e:
+            return f"Error generating response: {str(e)}\n\nPlease check your API key and try again."
+    
+    def _build_system_prompt(self, use_case: str) -> str:
+        """Build specialized system prompt based on use case."""
+        base_prompt = (
+            "You are an expert academic assistant for students, acting like a highly intelligent study buddy. "
+            "⚠️ CRITICAL RULE: You MUST ONLY use information from the provided context below. "
+            "DO NOT use your training knowledge. DO NOT infer beyond what's explicitly stated. "
+            "If the context doesn't contain adequate information to answer the question, you MUST respond: "
+            "'I cannot find sufficient information about this in the uploaded documents. Please upload materials covering this topic or rephrase your question.'\n\n"
+            "⚠️ GROUNDING REQUIREMENT: Every statement must be traceable to the provided context. "
+            "If you cannot find it in the context below, DO NOT answer from general knowledge.\n\n"
+            "✨ FORMATTING RULES (NotebookLM Style):\n"
+            "- Use clean, hierarchical Markdown (### Headers, **Bold** terms).\n"
+            "- Break down long paragraphs into easily readable bullet points.\n"
+            "- Be direct and concise. Avoid conversational fluff like 'Certainly!' or 'Here is the answer'.\n"
+            "- If applicable to the prompt, always try to extract a **Real-World Example** from the text to aid understanding.\n\n"
+        )
+        
+        if use_case == "explanation":
+            base_prompt += (
+                "**Your task:** Explain the concept in a clear, step-by-step manner suitable for students.\n"
+                "1. Start with a concise, one-sentence definition.\n"
+                "2. Break down the core mechanics or components using bullet points.\n"
+                "3. Provide an example (only if found in the text).\n"
+                "4. Add a 'Key Takeaway' at the end.\n"
+            )
+        elif use_case == "summary":
+            base_prompt += (
+                "**Your task:** Create a highly structured summary.\n"
+                "- Start with a brief high-level overview (2 sentences max).\n"
+                "- Use '### Key Themes' and list the main points as bulleted items.\n"
+                "- Keep each point concise but factually dense.\n"
+            )
+        elif use_case == "qa":
+            base_prompt += (
+                "**Your task:** Answer the question directly and comprehensively.\n"
+                "- Provide the direct answer immediately in the first sentence.\n"
+                "- Use numbered lists or bullet points to provide supporting details from the context.\n"
+                "- Use **bold** for key facts, numbers, and formulas.\n"
+            )
+        elif use_case == "notes":
+            base_prompt += (
+                "**Your task:** Create comprehensive, structured study notes.\n"
+                "- Use clear section headers (###).\n"
+                "- Organize information hierarchically (using nested bullet points).\n"
+                "- Explicitly highlight **Definitions**, **Formulas**, and **Important Dates/Names**.\n"
+            )
+        
+        base_prompt += (
+            "\n**Citation Rules:**\n"
+            "- You MUST cite your source at the end of every major claim or paragraph using numbered brackets like **[1]**, **[2]** based on the Source number provided in the context.\n"
+            "- If a claim comes from multiple sources, use **[1, 2]**.\n"
+            "- Do NOT use the document filename in the citation, ONLY the number.\n"
+            "- Do NOT make up information - stick strictly to the provided context.\n"
+        )
+        
+        return base_prompt
+    
+    def _build_user_message(self, prompt: str, context: str, metadatas: List[Dict] = None) -> str:
+        """Build user message with context and question."""
+        # Extract source names from metadata
+        sources = []
+        if metadatas:
+            for meta in metadatas:
+                filename = meta.get('filename', 'Unknown')
+                clean_name = filename.replace('.pdf', '').replace('.docx', '').replace('.txt', '')
+                if clean_name not in sources:
+                    sources.append(clean_name)
+        
+        message = "**Available Sources (USE ONLY THESE):**\n"
+        for source in sources[:5]:  # Show up to 5 sources
+            message += f"- {source}\n"
+        
+        message += f"\n**===== START OF CONTEXT (ANSWER ONLY FROM THIS) =====**\n\n{context}\n\n"
+        message += f"**===== END OF CONTEXT =====**\n\n"
+        message += f"**Student's Question:** {prompt}\n\n"
+        message += "**Instructions:** Answer ONLY using the context between the markers above. If the context doesn't contain the answer, say you don't have that information. Cite sources in brackets."
+        
+        return message
+    
+    def _generate_groq(self, system_prompt: str, user_message: str, temperature: float, max_tokens: int) -> str:
+        """Generate using Groq API (Llama-3.3-70B)."""
+        completion = self.client.chat.completions.create(
+            model="llama-3.3-70b-versatile",  # Latest 70B model (Dec 2024)
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": user_message}
+            ],
+            temperature=temperature,
+            max_tokens=max_tokens,
+            top_p=0.95,
+            stream=False
+        )
+        
+        return completion.choices[0].message.content
+    
+    def _generate_gemini(self, system_prompt: str, user_message: str, temperature: float, max_tokens: int) -> str:
+        """Generate using Google Gemini API."""
+        full_prompt = f"{system_prompt}\n\n{user_message}"
+        
+        response = self.client.generate_content(
+            full_prompt,
+            generation_config=genai.GenerationConfig(
+                temperature=temperature,
+                max_output_tokens=max_tokens,
+                top_p=0.95
+            )
+        )
+        
+        return response.text
+    
+    def is_ready(self) -> bool:
+        """Check if LLM is ready to generate."""
+        return self.ready
+    
+    def get_provider(self) -> str:
+        """Get current provider name."""
+        if self.provider == "groq":
+            return "Groq (Llama-3.3-70B)"
+        elif self.provider == "gemini":
+            return "Google Gemini 1.5 Flash"
+        return "Unknown"
+    
+    def generate(self, prompt: str, temperature: float = 0.3, max_tokens: int = 1500) -> str:
+        """
+        Simple wrapper for backend compatibility.
+        Generates response from a complete prompt that already includes context.
+        
+        Args:
+            prompt: Complete prompt with context already embedded
+            temperature: LLM temperature (0.0-1.0)
+            max_tokens: Maximum response length
+            
+        Returns:
+            Generated response
+        """
+        if not self.ready:
+            return (
+                "⚠️ **LLM not configured.** Please add your API key.\n\n"
+                "Get a free key:\n"
+                "- **Groq** (recommended, very fast): https://console.groq.com/keys\n"
+                "- **Gemini** (Google): https://makersuite.google.com/app/apikey"
+            )
+        
+        try:
+            if self.provider == "groq":
+                return self._generate_groq(
+                    system_prompt="You are a helpful AI assistant.",
+                    user_message=prompt,
+                    temperature=temperature,
+                    max_tokens=max_tokens
+                )
+            elif self.provider == "gemini":
+                return self._generate_gemini(
+                    system_prompt="You are a helpful AI assistant.",
+                    user_message=prompt,
+                    temperature=temperature,
+                    max_tokens=max_tokens
+                )
+        except Exception as e:
+            return f"Error generating response: {str(e)}"