app-rag.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 ---

# LEARNING MODE 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.
"""

# VIVA MODE PROMPTS (SPLIT)

# 1. TEXTBOOK - QUIZ GENERATOR
PROMPT_VIVA_TEXT_ASK = """
You are Anato-Mitra, an Examiner.
- Action: Use 'get_random_quiz_image' immediately.
- Output: Show the image and ask: "Identify the structure in the image and its function."
- Constraint: Do NOT grade previous answers. Just ask the question.
"""

# 2. TEXTBOOK - GRADER
PROMPT_VIVA_TEXT_GRADE = """
You are Anato-Mitra, a strict Professor.
- Task: Evaluate the user's answer based on the previous image shown.
- Format:
  1. Start with "🎉 CORRECT" or "❌ INCORRECT".
  2. If INCORRECT, immediately provide the **TRUE CORRECT ANSWER**.
  3. Provide a brief 1-sentence explanation from the text.
- Constraint: Do NOT generate a new image. Stop after the explanation.
"""

# 3. INTERNET - QUIZ GENERATOR
PROMPT_VIVA_NET_ASK = """
You are Anato-Mitra, an Examiner.
- Action: 
  1. Use 'get_random_anatomy_topic' to pick a subject.
  2. Use 'search_anatomy_diagrams' to find an image.
- Output: Show the image and ask: "Identify this structure."
- Constraint: Do NOT grade previous answers. Just ask the question.
"""

# 4. INTERNET - GRADER
PROMPT_VIVA_NET_GRADE = """
You are Anato-Mitra, a Medical Researcher.
- Task: Verify the user's answer.
- Format:
  1. Start with "🎉 CORRECT" or "❌ INCORRECT".
  2. If INCORRECT, explicitly state: "The correct structure is [NAME]".
  3. Provide a brief clinical fact.
- Constraint: Do NOT generate a new image. Stop after the explanation.
"""

# --- 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 ---
def fix_gradio_images(text):
    """
    Forces Gradio to load images via the API route.
    """
    if not text: return ""
    return text.replace("(extracted_images/", "(/gradio_api/file=extracted_images/")


# --- AGENT LOGIC (UPDATED) ---
async def run_agent(user_message, history, source, custom_prompt=None, force_no_tools=False):
    
    # Determine the Prompt
    if custom_prompt:
        system_prompt = custom_prompt
    else:
        system_prompt = PROMPT_LEARN_TEXT if source == "textbook" else PROMPT_LEARN_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: {source})...")
            
            # LOGIC FIX: If force_no_tools is True, we tell OpenAI "none"
            # This physically prevents the model from calling tools/images
            current_tool_choice = "none" if force_no_tools else "auto"

            response = client.chat.completions.create(
                model="meta-llama/Meta-Llama-3.1-70B-Instruct",
                messages=messages,
                tools=openai_tools,
                tool_choice=current_tool_choice 
            )

            final_response = response.choices[0].message.content or ""
            
            # Only process tools if we allowed them
            if not force_no_tools and 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 path fix
            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 ---
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 ===
    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)
                
                topic_dropdown = gr.Dropdown(
                    choices=[
                        "Draw and label the Brachial Plexus",
                        "Show the boundaries and contents of the Axilla",
                        "Illustrate the Cubital Fossa with its contents",
                        "Diagram the Femoral Triangle and its contents",
                        "Draw the anastomoses around the Scapula",
                        "Show the interior of the Right Atrium of the Heart",
                        "Illustrate the bronchopulmonary segments of both lungs",
                        "Draw the stomach bed with related structures",
                        "Show the porta hepatis and its contents",
                        "Illustrate the Circle of Willis",
                        "Draw the cavernous sinus and its contents",
                        "Show the orbit with extraocular muscles",
                        "Illustrate the nasal cavity and paranasal sinuses",
                        "Draw the inguinal canal and its boundaries",
                        "Show the popliteal fossa with neurovascular structures",
                        "Illustrate the carpal tunnel and its contents",
                        "Draw the plantar arches of the foot",
                        "Show the anatomical snuffbox boundaries",
                        "Illustrate the gluteal region muscles",
                        "Draw the mediastinum divisions"
                    ],
                    label="📚 Select a Topic from Textbook (All topics have diagrams)",
                    value=None,
                    interactive=True
                )
                
                msg_learn = gr.Textbox(placeholder="Or type your own question...")
                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 / Next Question ➡️", variant="primary")
            btn_home_viva = gr.Button("🏠 Main Menu")
        aud_viva = gr.Audio(visible=False, autoplay=True)

    # --- NAVIGATION LOGIC ---
    def go_home():
        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 ---
    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: LEARNING ---
    async def run_learn(msg, hist, src):
        # We pass None for prompt to use the default Learning logic inside run_agent
        txt, url = await run_agent(msg, hist, src, custom_prompt=None)
        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()

    def select_topic(topic): return topic if topic else ""
    
    topic_dropdown.change(select_topic, inputs=[topic_dropdown], outputs=[msg_learn])
    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)

    # --- CHAT INTERACTIONS: VIVA (SPLIT LOGIC) ---
    
    # 1. ASK (Start / Next) -> ALLOW TOOLS (force_no_tools=False)
    async def run_viva_next_question(hist, src):
        prompt = PROMPT_VIVA_TEXT_ASK if src == "textbook" else PROMPT_VIVA_NET_ASK
        
        # We send a hidden instruction to the agent to generate the question
        # force_no_tools=False because we NEED the image here
        txt, url = await run_agent("Generate next question", hist, src, custom_prompt=prompt, force_no_tools=False)
        
        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})"
        
        # We only append the Assistant's output (Question + Image)
        hist.append({"role": "assistant", "content": final})
        return hist

    # 2. GRADE (Submit Answer) -> BLOCK TOOLS (force_no_tools=True)
    async def run_viva_grade_answer(msg, hist, src):
        prompt = PROMPT_VIVA_TEXT_GRADE if src == "textbook" else PROMPT_VIVA_NET_GRADE
        
        # force_no_tools=True prevents the AI from fetching the next image while grading
        txt, _ = await run_agent(msg, hist, src, custom_prompt=prompt, force_no_tools=True)
        
        hist.append({"role": "user", "content": msg})
        hist.append({"role": "assistant", "content": txt})
        return "", hist
    # Viva Wiring
    msg_viva.submit(run_viva_grade_answer, [msg_viva, chat_viva, state_source], [msg_viva, chat_viva])
    btn_next_viva.click(run_viva_next_question, [chat_viva, state_source], [chat_viva])
    
    btn_speak_viva.click(generate_audio, chat_viva, aud_viva)
    mic_viva.stop_recording(transcribe_audio, mic_viva, msg_viva).then(run_viva_grade_answer, [msg_viva, chat_viva, state_source], [msg_viva, chat_viva])

if __name__ == "__main__":
    cwd = os.path.abspath(".")
    # Ensure specific folders exist or allow generic temp access
    allowed = [cwd, "/tmp"]
    if os.path.exists("extracted_images"): allowed.append(os.path.abspath("extracted_images"))
    
    demo.launch(share=True, allowed_paths=allowed)