app.py

# app.py
import re
import gradio as gr
import os
import time
import json
import sys
import requests 
import shutil   
import speech_recognition as sr
from openai import OpenAI
from dotenv import load_dotenv
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from elevenlabs.client import ElevenLabs

load_dotenv()

# --- CONFIGURATION ---
client = OpenAI(
    base_url="https://api.hyperbolic.xyz/v1", 
    api_key=os.getenv("HYPERBOLIC_API_KEY"), 
)

eleven = ElevenLabs(api_key=os.getenv("ELEVENLABS_API_KEY"))
VOICE_ID = "JBFqnCBsd6RMkjVDRZzb" 

# --- PROMPTS ---
PROMPT_LEARN_TEXT = """
You are Anato-Mitra, a strict Professor using ONLY the approved textbook.
- Source: Uploaded PDF.
- Tool: ALWAYS use 'consult_medical_textbook'.
- If answer is not in the book, say so. Do not hallucinate.
"""

PROMPT_LEARN_NET = """
You are Anato-Mitra, a Medical Researcher using the Web.
- Source: The Internet.
- Tool: ALWAYS use 'search_anatomy_diagrams'.
- Provide broad, clinical information.
"""

PROMPT_VIVA_TEXT = """
You are Anato-Mitra, an Examiner using the Course Textbook.
- Action: Use 'get_random_quiz_image' immediately.
- Ask: Identify the structure in the image and its function.
- Feedback: Start with "🎉 CORRECT!" or "❌ INCORRECT.".
"""

PROMPT_VIVA_NET = """
You are Anato-Mitra, an Examiner using the Web.
- Action: 
  1. Use 'get_random_anatomy_topic' to pick a subject.
  2. Use 'search_anatomy_diagrams' to find an image of it.
  3. Show the image and ask: "Identify this structure."
- Feedback: Start with "🎉 CORRECT!" or "❌ INCORRECT.".
"""

# --- HELPER: BYPASS 403 BLOCKS (Web Images) ---
def get_safe_image_path(url_or_path):
    if not url_or_path: return None
    
    # If it's a local path (starts with extracted_images or /gradio_api...)
    if not url_or_path.startswith("http"):
        # Clean it up to get the actual file system path
        clean = url_or_path.replace("/gradio_api/file=", "").replace("/file=", "")
        if os.path.exists(clean):
            return os.path.abspath(clean)
        return None

    # If it's a web URL, download it
    try:
        headers = {'User-Agent': 'Mozilla/5.0'}
        response = requests.get(url_or_path, headers=headers, stream=True, timeout=5)
        if response.status_code == 200:
            ext = ".svg" if ".svg" in url_or_path else ".png"
            temp_filename = f"temp_downloaded_img{ext}"
            with open(temp_filename, 'wb') as f:
                response.raw.decode_content = True
                shutil.copyfileobj(response.raw, f)
            return os.path.abspath(temp_filename)
    except: return None
    return None

# --- HELPER: FIX IMAGE PATHS (YOUR FIX) ---
def fix_gradio_images(text):
    """
    Forces Gradio to load images via the API route.
    """
    if not text: return ""
    # The specific fix you requested:
    return text.replace("(extracted_images/", "(/gradio_api/file=extracted_images/")

# --- AGENT LOGIC ---
async def run_agent(user_message, history, source, mode):
    if mode == "learning":
        system_prompt = PROMPT_LEARN_TEXT if source == "textbook" else PROMPT_LEARN_NET
    else:
        system_prompt = PROMPT_VIVA_TEXT if source == "textbook" else PROMPT_VIVA_NET
    
    server_params = StdioServerParameters(
        command=sys.executable, 
        args=["super_server.py"], 
    )

    async with stdio_client(server_params) as (read, write):
        async with ClientSession(read, write) as session:
            await session.initialize()
            tools_list = await session.list_tools()
            
            openai_tools = [{
                "type": "function",
                "function": {
                    "name": t.name,
                    "description": t.description,
                    "parameters": t.inputSchema
                }
            } for t in tools_list.tools]

            messages = [{"role": "system", "content": system_prompt}]
            for msg in history:
                messages.append({"role": "user" if msg['role'] == "user" else "assistant", "content": msg['content']})
            messages.append({"role": "user", "content": user_message})

            print(f"🧠 Thinking ({source}/{mode})...")
            response = client.chat.completions.create(
                model="meta-llama/Meta-Llama-3.1-70B-Instruct",
                messages=messages,
                tools=openai_tools,
                tool_choice="auto"
            )

            final_response = response.choices[0].message.content or ""
            
            if response.choices[0].message.tool_calls:
                tool_call = response.choices[0].message.tool_calls[0]
                fn_name = tool_call.function.name
                fn_args = tool_call.function.arguments
                
                args_dict = json.loads(fn_args)
                result = await session.call_tool(fn_name, arguments=args_dict)
                tool_output = result.content[0].text
                
                messages.append({
                    "role": "user", 
                    "content": f"SYSTEM DATA from '{fn_name}': {tool_output}\n\nProceed."
                })
                
                final_response_obj = client.chat.completions.create(
                    model="meta-llama/Meta-Llama-3.1-70B-Instruct",
                    messages=messages
                )
                final_response_w_tool = final_response_obj.choices[0].message.content
                
                if "![" in tool_output and "![" not in final_response_w_tool:
                    final_response_w_tool += f"\n\n{tool_output}"
                
                final_response = final_response_w_tool

            # Apply your path fix here
            final_response = fix_gradio_images(final_response)

            extracted_image_url = None
            img_match = re.search(r"!\[.*?\]\((.*?)\)", final_response)
            if img_match:
                extracted_image_url = img_match.group(1)

            return final_response, extracted_image_url

# --- AUDIO & STT ---
def generate_audio(chat_history):
    if not chat_history: return None
    last_msg = chat_history[-1]['content']
    clean = re.sub(r'!\[.*?\]\(.*?\)', '', last_msg)
    clean = re.sub(r'https?://\S+', '', clean).replace("🎉", "").replace("❌", "").replace("/gradio_api/file=", "").strip()
    try:
        gen = eleven.text_to_speech.convert(text=clean, voice_id=VOICE_ID, model_id="eleven_multilingual_v2")
        path = f"voice_{int(time.time())}.mp3"
        with open(path, "wb") as f:
            for chunk in gen: f.write(chunk)
        return path
    except: return None

def transcribe_audio(audio_path):
    if not audio_path: return ""
    rec = sr.Recognizer()
    try:
        with sr.AudioFile(audio_path) as s: return rec.recognize_google(rec.record(s))
    except: return ""

# --- UI ---
# --- UI ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    # State Variables
    state_source = gr.State("") 
    state_mode = gr.State("")   
    
    gr.Markdown("# 🩻 Anato-Mitra: AI Medical Companion")
    
    # === DEFINING THE COLUMNS ===
    # Note: We set visible=True/False here, but we will also FORCE it on load
    with gr.Column(visible=True) as view_source:
        gr.Markdown("### Step 1: Select Data Source")
        with gr.Row():
            btn_src_text = gr.Button("📖 Textbook (Strict)", variant="secondary", size="lg")
            btn_src_net = gr.Button("🌐 Internet (Broad)", variant="secondary", size="lg")

    with gr.Column(visible=False) as view_mode:
        lbl_source_display = gr.Markdown("### Source: ???")
        gr.Markdown("### Step 2: Select Activity")
        with gr.Row():
            btn_mode_learn = gr.Button("🎓 Learning Mode", variant="secondary", size="lg")
            btn_mode_viva = gr.Button("🎲 VIVA Mode", variant="primary", size="lg")
        btn_back_source = gr.Button("⬅️ Change Source")

    with gr.Column(visible=False) as view_learn:
        lbl_learn_header = gr.Markdown("### Learning Mode")
        with gr.Row():
            with gr.Column(scale=1):
                chat_learn = gr.Chatbot(type="messages", height=500)
                msg_learn = gr.Textbox(placeholder="Ask a question...")
                gr.Examples(
                    examples=["Describe the Heart", "Brachial Plexus", "Axilla Boundaries"],
                    inputs=msg_learn, label="Try:"
                )
                with gr.Row():
                    btn_speak_learn = gr.Button("🔊 Speak")
                    btn_home_learn = gr.Button("🏠 Main Menu")
                aud_learn = gr.Audio(visible=False, autoplay=True)
            with gr.Column(scale=1):
                img_learn = gr.Image(label="Diagram", interactive=False, height=500)

    with gr.Column(visible=False) as view_viva:
        lbl_viva_header = gr.Markdown("### VIVA Mode")
        chat_viva = gr.Chatbot(type="messages", height=500)
        with gr.Row():
            msg_viva = gr.Textbox(placeholder="Answer here...", scale=4)
            mic_viva = gr.Audio(sources=["microphone"], type="filepath", scale=1)
        with gr.Row():
            btn_speak_viva = gr.Button("🔊 Speak")
            btn_next_viva = gr.Button("🚀 Start Quiz", variant="primary")
            btn_home_viva = gr.Button("🏠 Main Menu")
        aud_viva = gr.Audio(visible=False, autoplay=True)

    # --- NAVIGATION LOGIC ---
    def go_home():
        # Force reset: Source=True, others=False
        return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)

    def select_textbook(): return gr.update(visible=False), gr.update(visible=True), "textbook", "### Source: 📖 Textbook"
    def select_internet(): return gr.update(visible=False), gr.update(visible=True), "internet", "### Source: 🌐 Internet"
    
    def select_learn(src): 
        header = f"### 🎓 Learning ({src.capitalize()})"
        return gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), "learning", header

    def select_viva(src): 
        header = f"### 🎲 VIVA ({src.capitalize()})"
        return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True), "viva", header

    def go_back_source(): return gr.update(visible=True), gr.update(visible=False)

    # --- EVENT WIRING ---
    
    # CRITICAL FIX: Ensure correct visibility on initial load
    demo.load(go_home, outputs=[view_source, view_mode, view_learn, view_viva])

    btn_src_text.click(select_textbook, outputs=[view_source, view_mode, state_source, lbl_source_display])
    btn_src_net.click(select_internet, outputs=[view_source, view_mode, state_source, lbl_source_display])
    
    btn_mode_learn.click(select_learn, inputs=[state_source], outputs=[view_mode, view_learn, view_viva, state_mode, lbl_learn_header])
    btn_mode_viva.click(select_viva, inputs=[state_source], outputs=[view_mode, view_learn, view_viva, state_mode, lbl_viva_header])

    btn_back_source.click(go_back_source, outputs=[view_source, view_mode])
    
    btn_home_learn.click(go_home, outputs=[view_source, view_mode, view_learn, view_viva])
    btn_home_viva.click(go_home, outputs=[view_source, view_mode, view_learn, view_viva])

    # --- CHAT INTERACTIONS (Same as before) ---
    async def run_learn(msg, hist, src):
        txt, url = await run_agent(msg, hist, src, "learning")
        hist.append({"role": "user", "content": msg})
        hist.append({"role": "assistant", "content": txt})
        safe_path = get_safe_image_path(url)
        return "", hist, safe_path if safe_path else gr.update()

    msg_learn.submit(run_learn, [msg_learn, chat_learn, state_source], [msg_learn, chat_learn, img_learn])
    btn_speak_learn.click(generate_audio, chat_learn, aud_learn)

    async def run_viva_start(hist, src):
        txt, url = await run_agent("Start Quiz", hist, src, "viva")
        final = txt
        if url and "![" not in txt:
             final += f"\n\n![Quiz Image](/file={url})" if not url.startswith("/file=") else f"\n\n![Quiz Image]({url})"
        hist.append({"role": "user", "content": "Next Question"})
        hist.append({"role": "assistant", "content": final})
        return hist, gr.update(value="🎲 Next Question")

    async def run_viva_answer(msg, hist, src):
        txt, _ = await run_agent(msg, hist, src, "viva")
        hist.append({"role": "user", "content": msg})
        hist.append({"role": "assistant", "content": txt})
        return "", hist

    msg_viva.submit(run_viva_answer, [msg_viva, chat_viva, state_source], [msg_viva, chat_viva])
    btn_next_viva.click(run_viva_start, [chat_viva, state_source], [chat_viva, btn_next_viva])
    btn_speak_viva.click(generate_audio, chat_viva, aud_viva)
    mic_viva.stop_recording(transcribe_audio, mic_viva, msg_viva).then(run_viva_answer, [msg_viva, chat_viva, state_source], [msg_viva, chat_viva])

if __name__ == "__main__":
    cwd = os.path.abspath(".")
    img_dir = os.path.abspath("extracted_images")
    book_dir = os.path.abspath("textbooks")
    demo.launch(share=False, allowed_paths=[cwd, img_dir, book_dir, "/tmp"])


if __name__ == "__main__":
    cwd = os.path.abspath(".")
    img_dir = os.path.abspath("extracted_images")
    book_dir = os.path.abspath("textbooks")
    #demo.launch(share=True, allowed_paths=[cwd, img_dir, book_dir, "/tmp"])
    demo.launch(share=True, allowed_paths=[cwd, "/tmp"])