File size: 7,024 Bytes
cd6f412
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8d7a112
 
cd6f412
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
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()