app.py

import os
import json
import streamlit as st
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
from groq import Groq
from reportlab.lib.pagesizes import A4
from reportlab.platypus import Paragraph, SimpleDocTemplate, Spacer
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.lib import colors

# Load milestones data
with open('milestones.json', 'r') as f:
    milestones = json.load(f)

# Age categories for dropdown selection
age_categories = {
    "Up to 2 months": 2,
    "Up to 4 months": 4,
    "Up to 6 months": 6,
    "Up to 9 months": 9,
    "Up to 1 year": 12,
    "Up to 15 months": 15,
    "Up to 18 months": 18,
    "Up to 2 years": 24,
    "Up to 30 months": 30,
    "Up to 3 years": 36,
    "Up to 4 years": 48,
    "Up to 5 years": 60
}

# Initialize FAISS and Sentence Transformer
model = SentenceTransformer('all-MiniLM-L6-v2')

def create_faiss_index(data):
    descriptions = []
    age_keys = []
    for age, categories in data.items():
        for entry in categories:
            descriptions.append(entry['description'])
            age_keys.append(int(age))  # Convert age to int for sorting
    
    embeddings = model.encode(descriptions, convert_to_numpy=True)
    index = faiss.IndexFlatL2(embeddings.shape[1])
    index.add(embeddings)
    return index, descriptions, age_keys

index, descriptions, age_keys = create_faiss_index(milestones)

# Function to retrieve the closest milestone
def retrieve_milestone(user_input):
    user_embedding = model.encode([user_input], convert_to_numpy=True)
    _, indices = index.search(user_embedding, 1)  # Get the closest match
    return descriptions[indices[0][0]] if indices[0][0] < len(descriptions) else "No relevant milestone found."

# Initialize Groq API
client = Groq(
    api_key=os.environ.get("GROQ_API_KEY"),
)

def generate_response(user_input, child_age):
    relevant_milestone = retrieve_milestone(user_input)
    prompt = (f"The child is {child_age} months old. Based on the given traits: {user_input}, "
              f"determine whether the child is meeting expected milestones. "
              f"Relevant milestone: {relevant_milestone}. "
              "If there are any concerns, suggest steps the parents can take. ")
    chat_completion = client.chat.completions.create(
        messages=[{"role": "user", "content": prompt}],
        model="llama-3.3-70b-versatile",
    )
    return chat_completion.choices[0].message.content

# Streamlit UI Styling
st.set_page_config(page_title="Tiny Triumphs Tracker", page_icon="👶", layout="wide")
st.markdown("""
    <style>
        .stApp { background-color: #1e1e2e; color: #ffffff; }
        .stTitle { text-align: center; color: #ffcc00; font-size: 36px; font-weight: bold; }
        .stButton > button { background-color: #ffcc00; color: #000; border-radius: 5px; font-weight: bold; }
        .stSelectbox, .stTextArea { background-color: #2e2e42; color: #ffffff; border-radius: 5px; }
        .placeholder { color: #aaaaaa; font-style: italic; }
    </style>
""", unsafe_allow_html=True)

st.markdown("<h1 class='stTitle'>👶 Tiny Triumphs Tracker</h1>", unsafe_allow_html=True)
st.markdown("Track your child's key growth milestones from birth to 5 years and detect early developmental concerns.", unsafe_allow_html=True)

# User selects child's age
selected_age = st.selectbox("📅 Select child's age:", list(age_categories.keys()))
child_age = age_categories[selected_age]

# User input for traits and skills
placeholder_text = "For example, your child might say simple words like 'mama' and 'dada' and smile when spoken to. They may grasp small objects with their fingers and show excitement during playtime."
user_input = st.text_area("✍️ Enter child's behavioral traits and skills:", placeholder=placeholder_text)

def generate_pdf_report(ai_response):
    pdf_file = "progress_report.pdf"
    doc = SimpleDocTemplate(pdf_file, pagesize=A4)
    styles = getSampleStyleSheet()
    
    elements = []
    
    # Title
    elements.append(Paragraph("Child Development Progress Report", styles['Title']))
    elements.append(Spacer(1, 12))
    
    # AI-generated insights
    elements.append(Paragraph("Development Insights:", styles['Heading2']))
    elements.append(Spacer(1, 10))
    
    # Format AI response into structured content
    response_parts = ai_response.split('\n')
    for part in response_parts:
        part = part.strip().lstrip('0123456789.- ')  # Remove numbers and unwanted prefixes
        if part:
            elements.append(Paragraph(f"• {part}", styles['Normal']))
            elements.append(Spacer(1, 5))
    
    # Disclaimer
    disclaimer = ("This report is AI-generated and is for informational purposes only. "
                  "It should not be considered a substitute for professional medical advice. "
                  "Always consult a qualified pediatrician for expert guidance on your child's development.")
    elements.append(Spacer(1, 12))
    elements.append(Paragraph(disclaimer, styles['Italic']))
    
    doc.build(elements)
    return pdf_file

if st.button("🔍 Analyze", help="Click to analyze the child's development milestones"):
    ai_response = generate_response(user_input, child_age)
    
    st.subheader("📊 Development Insights:")
    st.markdown(f"<div style='background-color:#44475a; color:#ffffff; padding: 15px; border-radius: 10px;'>{ai_response}</div>", unsafe_allow_html=True)
    
    # Download Report Button
    pdf_file = generate_pdf_report(ai_response)
    with open(pdf_file, "rb") as f:
        st.download_button(label="📥 Download Progress Report", data=f, file_name="progress_report.pdf", mime="application/pdf")

# Disclaimer
st.warning("⚠️ The results provided are generated by AI and should be interpreted with caution. Please consult a pediatrician for professional advice.")