import streamlit as st import requests import json import logging import uuid from pathlib import Path import sys # --- Path Setup --- project_root = str(Path(__file__).parent.parent) if project_root not in sys.path: sys.path.append(project_root) # --- Imports from our backend logic --- from insucompass.core.agent_orchestrator import app as unified_orchestrator_app from insucompass.services.zip_client import get_geo_data_from_zip # Configure logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) # --- Streamlit Page Configuration --- st.set_page_config( page_title="InsuCompass AI", page_icon="🧭", layout="wide", ) # --- Header --- LOGO_PATH = Path(__file__).parent / "assets" / "InsuCompass_Logo.png" INSU_AVATAR = Path(__file__).parent / "assets" / "robot-chat.jpg" # Create two columns col1, col2 = st.columns([0.4, 5]) # Adjust the width ratio as needed with col1: if LOGO_PATH.exists(): st.image(str(LOGO_PATH), width=70) with col2: st.title("InsuCompass") st.caption("Your AI guide to U.S. Health Insurance") # --- Session State Initialization --- if "thread_id" not in st.session_state: st.session_state.thread_id = str(uuid.uuid4()) st.session_state.user_profile = {} st.session_state.chat_history = [] st.session_state.is_profile_complete = False st.session_state.phase = "zip" # --- Helper Functions --- @st.cache_data def get_geodata(_zip_code: str): """Calls the zip_client function directly.""" return get_geo_data_from_zip(_zip_code) def run_orchestrator(user_message: str): """Calls the LangGraph orchestrator directly.""" thread_config = {"configurable": {"thread_id": st.session_state.thread_id}} inputs = { "user_profile": st.session_state.user_profile, "user_message": user_message, "is_profile_complete": st.session_state.is_profile_complete, "conversation_history": st.session_state.chat_history, } logger.info(f"Invoking orchestrator for thread_id: {st.session_state.thread_id}") try: final_state = unified_orchestrator_app.invoke(inputs, config=thread_config) st.session_state.user_profile = final_state["user_profile"] st.session_state.chat_history = final_state["conversation_history"] st.session_state.is_profile_complete = final_state["is_profile_complete"] except Exception as e: logger.error(f"Error invoking orchestrator: {e}", exc_info=True) st.error("An internal error occurred in the AI agent. Please try again.") # --- UI Rendering Functions --- def display_zip_form(): st.header("Let's Start with Your Location") zip_code_input = st.text_input("Enter your 5-digit ZIP code:", max_chars=5) if st.button("Confirm ZIP Code"): if zip_code_input and len(zip_code_input) == 5 and zip_code_input.isdigit(): with st.spinner("Verifying ZIP code..."): geo_data = get_geodata(zip_code_input) if geo_data: st.session_state.user_profile.update(geo_data.to_dict()) st.session_state.phase = "basic_profile" st.rerun() else: st.error("Invalid ZIP code or could not retrieve location data.") else: st.error("Please enter a valid 5-digit ZIP code.") def display_basic_profile_form(): st.header("Tell Us More About You") st.write(f"**Location:** {st.session_state.user_profile.get('city', '')}, {st.session_state.user_profile.get('state', '')}") with st.form("basic_profile_form"): st.subheader("Please provide some basic information:") # (CORRECTED) Fields now start empty and are validated. age = st.number_input("Your Age", min_value=1, max_value=120, value=None, placeholder="Enter your age") gender = st.selectbox("Your Gender", [None, "Male", "Female", "Non-binary", "Prefer not to say"], index=0, format_func=lambda x: "Select an option" if x is None else x) household_size = st.number_input("Household Size (including yourself)", min_value=1, value=None, placeholder="e.g., 1") income = st.number_input("Estimated Annual Household Income ($)", min_value=0, step=1000, value=None, placeholder="e.g., 55000") employment_status = st.selectbox("Your Employment Status", [None, "Employed with employer coverage", "Employed without coverage", "Unemployed", "Retired", "Student", "Self-employed"], index=0, format_func=lambda x: "Select an option" if x is None else x) citizenship = st.selectbox("Your Citizenship Status", [None, "US Citizen", "Lawful Permanent Resident", "Other legal resident", "Non-resident"], index=0, format_func=lambda x: "Select an option" if x is None else x) submitted = st.form_submit_button("Start My Personalized Session") if submitted: # Validation check if not all([age, gender, household_size, income is not None, employment_status, citizenship]): st.error("Please fill out all the fields to continue.") else: st.session_state.user_profile.update({ "age": age, "gender": gender, "household_size": household_size, "income": income, "employment_status": employment_status, "citizenship": citizenship, "medical_history": None, "medications": None, "special_cases": None }) st.session_state.phase = "chat" # Trigger the first profile question by calling the orchestrator with st.spinner("Starting your personalized profile conversation..."): run_orchestrator("START_PROFILE_BUILDING") st.rerun() def display_chat_interface(): st.header("Let's Chat!") # Display chat history from session state for message in st.session_state.chat_history: if ":" in message: role, content = message.split(":", 1) if role.lower() == "user": with st.chat_message("user", avatar="👤"): st.markdown(content.strip()) else: # Agent with st.chat_message("assistant", avatar=INSU_AVATAR): st.markdown(content.strip()) # User input if prompt := st.chat_input("Your message..."): with st.spinner("InsuCompass AI is thinking..."): run_orchestrator(prompt) st.rerun() # --- Main Application Flow Control --- if st.session_state.phase == "zip": display_zip_form() elif st.session_state.phase == "basic_profile": display_basic_profile_form() else: # 'chat' phase display_chat_interface() # --- Sidebar for Debugging --- with st.sidebar: st.header("Session State") st.json(st.session_state.to_dict(), expanded=False) if st.button("Reset Session"): st.session_state.clear() st.rerun()