app.py

import streamlit as st
import sqlite3
import pandas as pd
from huggingface_hub import InferenceClient
from dotenv import load_dotenv
import os

# Load environment variables
load_dotenv()

HF_TOKEN = os.getenv("HF_TOKEN")

# Load InferenceClient
@st.cache_resource
def get_client():
    return InferenceClient(token=HF_TOKEN)

# Connect to SQLite database
conn = sqlite3.connect("health.db")

# Fetch schema
schema_query = """
SELECT sql
FROM sqlite_master
WHERE type='table'
AND name='health_indicators';
"""

schema = conn.execute(schema_query).fetchone()[0]

# Streamlit UI
st.set_page_config(
    page_title="MediQuery",
    page_icon="🩺",
    layout="centered"
)

st.title("🩺 MediQuery")

st.write(
    "Ask questions about CDC health data in plain English."
)

question = st.text_input(
    "Ask a question:",
    placeholder=""
)

# Example questions sidebar
with st.sidebar:
    st.header("Example Questions")
    examples = [
        "Which state has the highest obesity rate?",
        "What is the diabetes prevalence by state?",
        "Which state has the most cancer cases?",
        "Show alcohol-related indicators by state",
        "What are the top 5 states for asthma rates?",
        "Compare mental health indicators across states",
    ]
    for ex in examples:
        if st.button(ex, use_container_width=True):
            st.session_state["example_question"] = ex

# Use example question if clicked
if "example_question" in st.session_state:
    question = st.session_state.pop("example_question")

# Function to generate SQL
def generate_sql(prompt):

    client = get_client()

    response = client.chat.completions.create(
        model="Qwen/Qwen2.5-72B-Instruct",
        messages=[{"role": "user", "content": prompt}],
        max_tokens=200
    )

    return response.choices[0].message.content.strip()


# Function to generate natural language answer from results
def generate_answer(question, df):

    client = get_client()

    data_str = df.to_string(index=False)

    answer_prompt = f"""
You are a helpful health data assistant.

The user asked: "{question}"

Here is the data retrieved from the CDC health database:
{data_str}

Write a clear, concise answer in 1-2 sentences using the data above.
Do not mention SQL or databases. Just answer the question naturally.
"""

    response = client.chat.completions.create(
        model="Qwen/Qwen2.5-72B-Instruct",
        messages=[{"role": "user", "content": answer_prompt}],
        max_tokens=150
    )

    return response.choices[0].message.content.strip()

# Generate Answer
if st.button("Generate Answer"):

    if not question.strip():
        st.warning("Please enter a question.")
        st.stop()

    prompt = f"""
You are an expert SQLite SQL generator.

Convert the user's question into valid SQLite SQL.

IMPORTANT RULES:
1. ONLY return the raw SQL query
2. No markdown, no code fences, no explanation
3. Use only SQLite syntax
4. Table name is health_indicators
5. Use LIKE for partial text matches on topic/question columns
6. datavaluetype values are: 'Age-adjusted Prevalence', 'Crude Prevalence', 'Number'
7. Common topic values: 'Alcohol', 'Arthritis', 'Asthma', 'Cancer', 'Cardiovascular Disease', 'Chronic Kidney Disease', 'Chronic Obstructive Pulmonary Disease', 'Diabetes', 'Disability', 'Immunization', 'Mental Health', 'Nutrition, Physical Activity, and Weight Status', 'Oral Health', 'Overarching Conditions', 'Reproductive Health', 'Tobacco'
8. For obesity questions, use: topic = 'Nutrition, Physical Activity, and Weight Status' AND question LIKE '%Obesity%'

Database schema:
{schema}

User Question:
{question}
"""

    try:

        with st.spinner("Analyzing CDC health data..."):

            # Generate SQL
            sql_query = generate_sql(prompt)

            # Execute query
            df = pd.read_sql_query(sql_query, conn)

            # Remove empty rows
            df = df.dropna(how="all")

            # Fill null values
            df = df.fillna("N/A")

        if df.empty:
            st.warning("No results found for your question.")
        else:
            # Generate plain English answer
            with st.spinner("Generating answer..."):
                answer = generate_answer(question, df)

            st.subheader("Answer")
            st.success(answer)

            # Auto-detect numeric column for chart
            numeric_cols = df.select_dtypes(include="number").columns.tolist()
            label_cols = df.select_dtypes(exclude="number").columns.tolist()

            if numeric_cols and label_cols:
                st.subheader("Chart")
                chart_df = df.set_index(label_cols[0])[numeric_cols[0]]
                st.bar_chart(chart_df)

            # Show raw data expander
            with st.expander("View raw data"):
                st.dataframe(df)

            # Show generated SQL expander
            with st.expander("View generated SQL"):
                st.code(sql_query, language="sql")

    except Exception as e:
        st.error(f"Error: {e}")