app.py

import gradio as gr
import os
import re
from groq import Groq
from tavily import TavilyClient
# Import your custom logic from language.py
from language import detect_language, translate_text

# 1. Clients
groq_client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
tavily_client = TavilyClient(api_key=os.environ.get("TAVILY_API_KEY"))

# 2. Optimized Medical Prompt
SYSTEM_PROMPT = """
You are a Medical Information Specialist. 
- Use the context to explain health concepts simply.
- STICK TO NON-DRUG ADVICE (Rest, hydration, hygiene).
- NO MEDICATIONS, NO DOSAGES, NO BRAND NAMES.
- Always include a section on 'When to see a doctor'.
- If the user just says hi/hello, respond with a friendly medical greeting and ask how you can help with their health queries.
"""

def medical_chat(message, history):
    try:
        # A. Language Detection
        user_lang = detect_language(message)
        
        # B. Handle Small Talk / Greetings
        # If message is very short, don't waste search credits
        small_talk_patterns = r"^(hi|hello|hey|namaste|vanakkam|hi there|hello there)[!?.]?$"
        if re.match(small_talk_patterns, message.strip().lower()):
            greeting = "Hello! I am your Indic-Medical Assistant. How can I help you with your health or medical questions today?"
            return translate_text(greeting, user_lang) if user_lang != "en" else greeting

        # C. Translation to English for the best search results
        eng_query = translate_text(message, "en") if user_lang != "en" else message

        # D. Search Trusted Official Sources
        search_result = tavily_client.search(
            query=f"medical guidance for {eng_query}",
            include_domains=["who.int", "mohfw.gov.in", "mayoclinic.org", "nhs.uk"],
            max_results=3
        )
        
        # Extract context and build a source list
        context_parts = []
        sources = []
        for i, r in enumerate(search_result['results']):
            context_parts.append(f"Source [{i+1}]: {r['content']}")
            sources.append(f"- [{r.get('title', 'Official Source')}]({r['url']})")
        
        context = "\n".join(context_parts)
        source_markdown = "\n\n**Sources / स्रोत:**\n" + "\n".join(sources)

        # E. Groq Reasoning (English)
        response = groq_client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": f"CONTEXT: {context}\n\nUSER QUESTION: {eng_query}"}
            ]
        )
        eng_response = response.choices[0].message.content

        # F. Final Translation back to User's Language
        final_output = translate_text(eng_response, user_lang) if user_lang != "en" else eng_response

        # G. Safety Check (Regex for dosages)
        if re.search(r'\d+\s?(mg|ml|tablet|pill)', final_output.lower()):
            final_output = "I cannot provide specific dosages. Please consult a doctor. / मैं खुराक की जानकारी नहीं दे सकता। कृपया डॉक्टर से सलाह लें।"

        # Combine output with sources
        return f"{final_output}\n\n---\n{source_markdown}"

    except Exception as e:
        return f"⚠️ Error: {str(e)}"

# 3. Launch UI (Gradio 6.0 compatible)
with gr.Blocks() as demo:
    gr.Markdown("# ⚡ Official Indic-Medical AI")
    gr.Markdown("Conversational medical guidance in Indian languages with official citations.")
    
    gr.ChatInterface(
        fn=medical_chat,
        examples=["बुखार के लक्षण क्या हैं?", "Signs of heart attack", "Explain diabetes in Tamil"],
    )

if __name__ == "__main__":
    demo.launch(ssr_mode=False, theme="soft")