app.py

import os
import json
from typing import List, Dict, Any, Generator

# Gradio & LLM/Agent Imports
import gradio as gr
from gradio.components.chatbot import ChatMessage
import openai # Using the standard OpenAI client (will be configured for Nebius)

# LlamaIndex Imports (Essential for RAG and Agent Context)
# In a real app, you would import LlamaIndex components here, e.g.,
# from llama_index.core.agent import FunctionCallingAgentWorker
# from llama_index.core.query_engine import RetrieverQueryEngine

# --- CONFIGURATION PLACEHOLDERS ---
# You will need to configure the OpenAI client to point to the Nebius Token Factory endpoint.
# NEBIUS_API_KEY should be set as a Hugging Face Secret.
NEBIUS_API_KEY = os.getenv("NEBIUS_API_KEY")
NEBIUS_BASE_URL = os.getenv("NEBIUS_BASE_URL", "https://api.nebiustokenfactory.com/v1")
LLM_MODEL = "openai/gpt-oss-120b" # Using the specified Nebius model ID

# --- MOCK RAG SETUP (Needs replacement with real LlamaIndex code) ---
class MockRAG:
    """A placeholder for the real LlamaIndex RAG query engine."""
    def query(self, query: str) -> str:
        if "quantum physics" in query.lower() or "explain" in query.lower():
            return "Quantum physics studies matter and energy at the most fundamental level, where classical mechanics breaks down. Key concepts include wave-particle duality and the uncertainty principle. This explanation is grounded in the latest 2025 educational materials."
        else:
            return "Hello! I'm your AI Tutor. Let's learn about that! Ask me to explain a topic, show a video, or give you a quiz!"

# --- MCP TOOL DEFINITIONS ---

def video_search_tool(topic: str) -> str:
    """
    Searches YouTube for a highly rated educational video relevant to the user's topic.

    Args:
        topic (str): The subject or question provided by the user.

    Returns:
        str: An HTML snippet embedding the YouTube video in the chat.
    """
    # Placeholder Logic: A real tool would call the YouTube API or Google Search
    video_map = {
        "quantum physics": "https://www.youtube.com/embed/gIWy5p4MZVE", # Kurzgesagt
        "biology": "https://www.youtube.com/embed/A8gNf1p60pE",        # Amoeba Sisters
        "history": "https://www.youtube.com/embed/Yocja_N5s1I",        # Crash Course
    }
    
    url = next((v for k, v in video_map.items() if k in topic.lower()), video_map["quantum physics"])
    
    # Use HTML to embed the video directly in the Chatbot
    html_embed = f"""
    <div style="padding: 10px; border: 1px solid #e0e0e0; border-radius: 8px; background: #f9f9f9; max-width: 100%; border-radius: 8px;">
        <p style="font-weight: bold; margin-bottom: 5px; color: #1f2937;">🎥 Tutor Found a Video Resource!</p>
        <iframe width="100%" height="315" src="{url}" frameborder="0" allowfullscreen style="border-radius: 6px;"></iframe>
        <p style="font-style: italic; font-size: 0.9em; margin-top: 5px; color: #6b7280;">Click play above to watch the lesson.</p>
    </div>
    """
    return html_embed

def quiz_generator_tool(topic: str) -> str:
    """
    Generates a short, three-question multiple-choice quiz on a specific subject.

    Args:
        topic (str): The subject or knowledge area for the quiz.

    Returns:
        str: A Markdown/HTML formatted quiz for the user to answer.
    """
    # Placeholder Logic: A real tool would use the LLM to generate the quiz JSON
    quiz_markdown = """
    <div style="background-color: #ecfdf5; padding: 15px; border-radius: 8px; border-left: 5px solid #059669; color: #065f46;">
        <h4 style="margin-top: 0; color: #059669;">📝 Quiz Time: Understanding the Fundamentals!</h4>
        <ol style="padding-left: 20px;">
            <li style="margin-bottom: 10px;">**Which of these phenomena demonstrates wave-particle duality?**
                <ul><li>A. An apple falling from a tree</li><li>B. The double-slit experiment</li><li>C. Water boiling at 100°C</li></ul>
            </li>
            <li style="margin-bottom: 10px;">**What does the Heisenberg Uncertainty Principle state?**
                <ul><li>A. It's impossible to know a particle's position and velocity simultaneously.</li><li>B. Energy is conserved in a closed system.</li><li>C. Light travels fastest in a vacuum.</li></ul>
            </li>
            <li style="margin-bottom: 0;">**In quantum mechanics, what is an "orbital"?**
                <ul><li>A. The fixed path of an electron around the nucleus</li><li>B. A region of space where an electron is most likely to be found</li><li>C. A subatomic particle of light</li></ul>
            </li>
        </ol>
        <p style="font-style: italic; font-size: 0.85em; margin-top: 15px;">**Hint:** Your tutor will check your answers in the next message!</p>
    </div>
    """
    return quiz_markdown

# --- AGENT CORE (SIMULATED FOR TOOL-CALLING) ---

class NebiusAIAgent:
    """
    Simulates the Agentic workflow using the OpenAI Client (configured for Nebius) 
    and LlamaIndex tools.
    """
    def __init__(self, tools: List, rag_query_engine: Any):
        self.rag_query_engine = rag_query_engine
        self.tools = {t.__name__: t for t in tools}
        self.tool_descriptions = [
            {"type": "function", "function": {
                "name": "video_search_tool",
                "description": "Searches for an educational video to help the student learn a topic visually.",
                "parameters": {"type": "object", "properties": {"topic": {"type": "string", "description": "The educational topic to search for."}}, "required": ["topic"]}
            }},
            {"type": "function", "function": {
                "name": "quiz_generator_tool",
                "description": "Generates a structured quiz (3-5 questions) to test the student's knowledge on a specific topic.",
                "parameters": {"type": "object", "properties": {"topic": {"type": "string", "description": "The topic to base the quiz on."}}, "required": ["topic"]}
            }}
        ]
        
        # Initialize the OpenAI Client, pointing to Nebius Token Factory
        self.client = openai.OpenAI(
            api_key=NEBIUS_API_KEY,
            base_url=NEBIUS_BASE_URL,
        )

    def stream_chat(self, prompt: str, history: List) -> Generator[str, None, None]:
        
        # 1. Build Messages
        messages = [{"role": "system", "content": "You are a patient, world-class AI Tutor. Use your available tools to find videos or create quizzes when requested. Otherwise, use your RAG knowledge to explain topics. Be encouraging and use markdown formatting."}]
        # Convert history format
        for user_msg, assistant_msg in history:
            if user_msg: messages.append({"role": "user", "content": user_msg})
            if assistant_msg: messages.append({"role": "assistant", "content": assistant_msg})
        messages.append({"role": "user", "content": prompt})

        try:
            # 2. Call the LLM (Nebius via OpenAI client)
            # NOTE: In a real app, you would stream the response using 'stream=True' 
            # and handle the tool call logic from the response object.
            response = self.client.chat.completions.create(
                model=LLM_MODEL, 
                messages=messages,
                tools=self.tool_descriptions,
                tool_choice="auto" # Let the model decide whether to call a tool
            )
            
            # --- MOCK TOOL CALLING LOGIC (Simplified) ---
            # NOTE: For this mock, we manually check keywords to simulate tool calling
            # to ensure the demo works without a live agent model that can interpret tool descriptions.
            
            tool_call = None
            if "video" in prompt.lower() or "show me" in prompt.lower():
                tool_call = {"function": {"name": "video_search_tool", "arguments": json.dumps({"topic": prompt})}}
            elif "quiz" in prompt.lower() or "test me" in prompt.lower():
                tool_call = {"function": {"name": "quiz_generator_tool", "arguments": json.dumps({"topic": prompt})}}

            if tool_call:
                # Agent Thought: Tool Call
                yield ChatMessage(role="assistant", content="", metadata={"title": f"🧠 Agent Thought: Decided to call **{tool_call['function']['name']}**."})
                
                # Execute the tool (local mock execution)
                tool_func = self.tools[tool_call['function']['name']]
                args = json.loads(tool_call['function']['arguments'])
                tool_output = tool_func(args['topic'])
                
                # Display Tool Result
                yield ChatMessage(role="assistant", content=tool_output, metadata={"title": "Tool Result"})
                
            else:
                # Agent Thought: RAG/General LLM
                yield ChatMessage(role="assistant", content="", metadata={"title": "🧠 Agent Thought: Retrieving context from knowledge base (RAG)."})
                
                # Simulate RAG and response synthesis
                rag_result = self.rag_query_engine.query(prompt)
                
                # Stream the final response
                full_response = f"**Tutor's Explanation:**\n\n{rag_result}"
                for chunk in full_response.split():
                    yield chunk + " "

        except Exception as e:
            error_message = f"🚨 **Error during chat processing:** Ensure your Nebius API Key and Base URL are correctly configured. Error: {e}"
            for chunk in error_message.split():
                yield chunk + " "
            
# --- INITIALIZATION ---

# 1. Initialize RAG
rag_knowledge_base = MockRAG()

# 2. Define the list of tools available to the Agent
available_tools = [video_search_tool, quiz_generator_tool]

# 3. Initialize the Agent Executor
agent_executor = NebiusAIAgent(available_tools, rag_knowledge_base)

# The function that the Gradio ChatInterface will call
def tutor_chat_function(message: str, history: List[List[str]]) -> Generator[str, None, None]:
    # This block handles the streaming output from the agent
    # Stream the full response to the Gradio Chatbot
    response_stream = agent_executor.stream_chat(message, history)
    for chunk in response_stream:
        yield chunk

# --- GRADIO UI ---

with gr.Blocks(theme=gr.themes.Monochrome(), title="AI Tutor Agent (MCP in Action)") as demo:
    gr.Markdown("# 🎓 AI Tutor Agent (Nebius/OpenAI Client & Gradio MCP)")
    gr.Markdown(
        "**Ask me anything about science, history, or tech!** "
        "Try saying things like: *'Explain quantum physics'* or *'Test me on biology'* or *'Show me a video on the Cold War'*."
        "<br>_This agent uses RAG for articles and is tool-enabled via MCP. **Remember to set your NEBIUS_API_KEY secret.**_"
    )

    # Use gr.ChatInterface for the main interaction
    gr.ChatInterface(
        fn=tutor_chat_function,
        textbox=gr.Textbox(placeholder="What would you like to learn today?"),
        chatbot=gr.Chatbot(
            label="AI Tutor",
            show_copy_button=True, 
            height=600,
            # Added a class for mobile responsiveness
            elem_classes=["min-h-[70vh]", "mobile-full-width"]
        ),
        # Sets the initial user message and prompt for the user
        examples=[
            "Explain the water cycle in detail.", 
            "Give me a quiz on World War 2 history.",
            "Show me a video about photosynthesis."
        ]
    )

if __name__ == "__main__":
    demo.launch()