agents/viva_agent.py

# agents/viva_agent.py
"""
Viva Practice Agent - Handles mock interview sessions with state management.
"""
import re
import json

class VivaAgent:
    def __init__(self, gemini_model=None):
        """
        Initializes the agent with the Gemini model.

        Args:
            gemini_model: An instance of the Gemini model client.
        """
        self.model = gemini_model

    def _get_viva_question(self, topic: str, difficulty: str, asked_questions: list, file_context: str) -> str:
        """Calls the Gemini API to generate a single, unique viva question."""
        
        context_source = ""
        if file_context:
            context_source =f"---\nCONTEXT FROM KNOWLEDGE BASE:\n{file_context}\n---" if file_context else ""
        else:
            context_source = f"The question should be about the B.Pharmacy topic: **{topic}**."

        prompt = f"""
You are a "Viva Coach," a professional and encouraging AI examiner for B.Pharmacy students like Gurus from ancient India.
**CRITICAL INSTRUCTION FOR CITATIONS:** When you use information from the KNOWLEDGE BASE CONTEXT, you MUST cite the source at the end of the relevant sentence using the format `[Source: filename, Page: page_number]`.
**Your Task:**
Generate a SINGLE, insightful, open-ended viva question.

**Instructions:**
1.  **Context:** {context_source}
2.  **Difficulty:** The question's difficulty should be **{difficulty}**.
3.  **Avoid Repetition:** DO NOT ask any of these previously asked questions: {json.dumps(asked_questions)}.
4.  **Format:** Return ONLY the question itself. No conversational text, no "Here is your question:", just the question.

Example: "Can you explain the primary mechanism of action for beta-blockers?"
"""
        try:
            response = self.model.generate_content(prompt)
            # Basic cleanup of the response
            return response.text.strip().replace("\"", "")
        except Exception as e:
            print(f"Viva Agent AI Error: {e}")
            return "I'm having trouble thinking of a question right now. Please try again."

    def process_query(self, query: str, file_context: str = "", viva_state: dict = None):
        """
        Processes a query to manage a viva session.

        Args:
            query (str): The user's command (e.g., "start viva", "next question").
            file_context (str): Optional text from an uploaded file to base the viva on.
            viva_state (dict): The current state of the viva session from the user's session.
        
        Returns:
            dict: A dictionary containing the response and updated viva state.
        """
        if not self.model:
            return {
                'message': "🗣️ **Viva Coach**\n\nThe examiner is unavailable! The Gemini API key is missing, so I can't conduct a viva. Please configure the API key.",
                'agent_type': 'viva_practice',
                'status': 'error_no_api_key',
                'viva_state': None # No state change
            }

        query = query.lower().strip()
        viva_state = viva_state or {'active': False, 'topic': '', 'asked_questions': [], 'difficulty': 'easy'}

        # --- Command Handling ---

        # Command: END VIVA
        if viva_state['active'] and any(cmd in query for cmd in ["end viva", "stop viva", "exit viva"]):
            message = "Great session! You did well. Keep practicing! 👏\n\nThe viva session has now ended."
            return {'message': message, 'agent_used': 'viva_practice', 'status': 'session_ended', 'viva_state': {'active': False}}

        # Command: START VIVA
        if any(cmd in query for cmd in ["start viva", "begin viva"]):
            topic_match = re.search(r'viva on (.+)', query)
            topic = topic_match.group(1).title() if topic_match else "General Pharmacy"
            
            # Reset state for a new session
            new_state = {'active': True, 'topic': topic, 'asked_questions': [], 'difficulty': 'easy'}
            
            first_question = self._get_viva_question(topic, new_state['difficulty'], [], file_context)
            new_state['asked_questions'].append(first_question)
            
            message = f"Alright, let's begin your viva on **{topic}**! Here is your first question:\n\n**Q1:** {first_question}"
            return {'message': message, 'agent_used': 'viva_practice', 'status': 'session_started', 'viva_state': new_state}
            
        # Command: NEXT QUESTION (only if a session is active)
        if viva_state['active'] and any(cmd in query for cmd in ["next question", "next one", "ask another"]):
            question_number = len(viva_state['asked_questions']) + 1
            # Slightly increase difficulty over time
            if question_number > 5: viva_state['difficulty'] = 'hard'
            elif question_number > 2: viva_state['difficulty'] = 'medium'

            new_question = self._get_viva_question(viva_state['topic'], viva_state['difficulty'], viva_state['asked_questions'], file_context)
            viva_state['asked_questions'].append(new_question)
            
            message = f"Excellent. Here is your next question:\n\n**Q{question_number}:** {new_question}"
            return {'message': message, 'agent_used': 'viva_practice', 'status': 'next_question', 'viva_state': viva_state}

        # Default response if no command is recognized
        message = """🗣️ **Viva Coach Ready!**

You can start a session by typing:
* `start viva on [your topic]` (e.g., `start viva on Pharmacokinetics`)
* Or just `start viva` for general questions.

Once started, you can say:
* `next question`
* `end viva`
"""
        return {'message': message, 'agent_type': 'viva_practice', 'status': 'idle', 'viva_state': viva_state}