app.py

import streamlit as st
from langchain.prompts import ChatPromptTemplate, FewShotChatMessagePromptTemplate
from huggingface_hub import InferenceClient

# Hugging Face API setup
client = InferenceClient(
    model="deepseek-ai/DeepSeek-V3-0324",
    api_key=st.secrets["HUGGINGFACEHUB_API_TOKEN"]
)

# Few-shot examples
examples = [
    {
        "input": "Disease: Flu\nSymptoms: fever, cough, fatigue",
        "output": "Precautions: rest, drink fluids, take antipyretics if needed"
    },
    {
        "input": "Disease: Dengue\nSymptoms: high fever, severe headache, muscle pain",
        "output": "Precautions: avoid mosquito bites, drink fluids, monitor platelet count"
    },
    {
        "input": "Disease: Asthma\nSymptoms: shortness of breath, wheezing",
        "output": "Precautions: avoid allergens, use inhalers as prescribed"
    }
]

# Templates
example_prompt = ChatPromptTemplate.from_messages([
    ("human", "{input}"),
    ("ai", "{output}")
])

few_shot_prompt = FewShotChatMessagePromptTemplate(
    example_prompt=example_prompt,
    examples=examples
)

final_prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a helpful medical assistant. Based on the disease and symptoms, suggest possible precautions."),
    few_shot_prompt,
    ("human", "{question}")
])

# Streamlit UI
st.title("🩺 AI Medical Assistant (Few-Shot Chat)")

user_input = st.text_area("Enter disease and symptoms (e.g., 'Disease: Malaria\nSymptoms: fever, chills, vomiting'):")

if st.button("Get Precautions") and user_input:
    messages = final_prompt.format_messages(question=user_input)
    chat_messages = [{"role": m.type, "content": m.content} for m in messages]

    with st.spinner("Thinking..."):
        response = client.chat.completions.create(
            model="deepseek-ai/DeepSeek-V3-0324",
            messages=chat_messages,
            max_tokens=256,
        )
        st.success(response.choices[0].message["content"])