|
import streamlit as st |
|
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, ToolMessage |
|
from datetime import datetime |
|
|
|
from config.settings import settings |
|
from agent import get_agent_executor |
|
from models import ChatMessage, ChatSession, User |
|
from models.db import get_session_context |
|
from services.logger import app_logger |
|
from services.metrics import log_consultation_start |
|
|
|
st.set_page_config(page_title=f"Consult - {settings.APP_TITLE}", layout="wide") |
|
|
|
if not st.session_state.get("authenticated_user"): |
|
st.warning("Please log in to access the consultation page.") |
|
st.switch_page("app.py") |
|
|
|
|
|
try: |
|
agent_executor = get_agent_executor() |
|
except ValueError as e: |
|
st.error(f"Could not initialize AI Agent: {e}") |
|
st.stop() |
|
|
|
|
|
|
|
def load_chat_history(session_id: int) -> list: |
|
"""Loads chat history from DB for the current session""" |
|
messages = [] |
|
with get_session_context() as db: |
|
db_messages = db.query(ChatMessage).filter(ChatMessage.session_id == session_id).order_by(ChatMessage.timestamp).all() |
|
for msg in db_messages: |
|
if msg.role == "user": |
|
messages.append(HumanMessage(content=msg.content)) |
|
elif msg.role == "assistant": |
|
messages.append(AIMessage(content=msg.content)) |
|
|
|
|
|
|
|
return messages |
|
|
|
def save_chat_message(session_id: int, role: str, content: str, tool_call_id: Optional[str]=None, tool_name: Optional[str]=None): |
|
"""Saves a chat message to the database.""" |
|
with get_session_context() as db: |
|
chat_message = ChatMessage( |
|
session_id=session_id, |
|
role=role, |
|
content=content, |
|
timestamp=datetime.utcnow(), |
|
tool_call_id=tool_call_id, |
|
tool_name=tool_name |
|
) |
|
db.add(chat_message) |
|
db.commit() |
|
|
|
|
|
st.title("AI Consultation Room") |
|
st.markdown("Interact with the Quantum Health Navigator AI.") |
|
|
|
current_user: User = st.session_state.authenticated_user |
|
chat_session_id = st.session_state.get("current_chat_session_id") |
|
|
|
if not chat_session_id: |
|
st.error("No active chat session. Please re-login or contact support.") |
|
st.stop() |
|
|
|
|
|
|
|
if "agent_chat_history" not in st.session_state: |
|
st.session_state.agent_chat_history = load_chat_history(chat_session_id) |
|
if not st.session_state.agent_chat_history: |
|
log_consultation_start() |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
with get_session_context() as db: |
|
ui_messages = db.query(ChatMessage).filter(ChatMessage.session_id == chat_session_id).order_by(ChatMessage.timestamp).all() |
|
for msg in ui_messages: |
|
with st.chat_message(msg.role): |
|
st.markdown(msg.content) |
|
|
|
|
|
if prompt := st.chat_input("Ask the AI... (e.g., 'What is hypertension?' or 'Optimize treatment for patient X with diabetes')"): |
|
|
|
with st.chat_message("user"): |
|
st.markdown(prompt) |
|
save_chat_message(chat_session_id, "user", prompt) |
|
|
|
|
|
st.session_state.agent_chat_history.append(HumanMessage(content=prompt)) |
|
|
|
|
|
with st.spinner("AI is thinking..."): |
|
try: |
|
response = agent_executor.invoke({ |
|
"input": prompt, |
|
"chat_history": st.session_state.agent_chat_history |
|
}) |
|
ai_response_content = response['output'] |
|
|
|
|
|
with st.chat_message("assistant"): |
|
st.markdown(ai_response_content) |
|
save_chat_message(chat_session_id, "assistant", ai_response_content) |
|
|
|
|
|
st.session_state.agent_chat_history.append(AIMessage(content=ai_response_content)) |
|
|
|
|
|
|
|
|
|
|
|
|
|
except Exception as e: |
|
app_logger.error(f"Error during agent invocation: {e}") |
|
st.error(f"An error occurred: {e}") |
|
|
|
error_message = f"Sorry, I encountered an error: {str(e)[:200]}" |
|
with st.chat_message("assistant"): |
|
st.markdown(error_message) |
|
save_chat_message(chat_session_id, "assistant", error_message) |
|
st.session_state.agent_chat_history.append(AIMessage(content=error_message)) |
|
|
|
|
|
|