Update app error handling

src/streamlit_app.py

@@ -4,11 +4,10 @@
 
 import streamlit as st
 from dotenv import load_dotenv
-from httpx import ConnectError, ConnectTimeout
 from pydantic import ValidationError
 from streamlit.runtime.scriptrunner import get_script_run_ctx
 
-from client import AgentClient
+from client import AgentClient, AgentClientError
 from schema import ChatHistory, ChatMessage
 from schema.task_data import TaskData, TaskDataStatus
 
@@ -59,9 +58,11 @@ async def main() -> None:
             port = os.getenv("PORT", 80)
             agent_url = f"http://{host}:{port}"
         try:
-            st.session_state.agent_client = AgentClient(base_url=agent_url)
-        except (ConnectError, ConnectTimeout) as e:
+            with st.spinner("Connecting to agent service..."):
+                st.session_state.agent_client = AgentClient(base_url=agent_url)
+        except AgentClientError as e:
             st.error(f"Error connecting to agent service: {e}")
+            st.markdown("The service might be booting up. Try again in a few seconds.")
             st.stop()
     agent_client: AgentClient = st.session_state.agent_client
 
@@ -140,25 +141,29 @@ async def amessage_iter() -> AsyncGenerator[ChatMessage, None]:
     if user_input := st.chat_input():
         messages.append(ChatMessage(type="human", content=user_input))
         st.chat_message("human").write(user_input)
-        if use_streaming:
-            stream = agent_client.astream(
-                message=user_input,
-                model=model,
-                thread_id=st.session_state.thread_id,
-            )
-            await draw_messages(stream, is_new=True)
-        else:
-            response = await agent_client.ainvoke(
-                message=user_input,
-                model=model,
-                thread_id=st.session_state.thread_id,
-            )
-            messages.append(response)
-            st.chat_message("ai").write(response.content)
-        st.rerun()  # Clear stale containers
+        try:
+            if use_streaming:
+                stream = agent_client.astream(
+                    message=user_input,
+                    model=model,
+                    thread_id=st.session_state.thread_id,
+                )
+                await draw_messages(stream, is_new=True)
+            else:
+                response = await agent_client.ainvoke(
+                    message=user_input,
+                    model=model,
+                    thread_id=st.session_state.thread_id,
+                )
+                messages.append(response)
+                st.chat_message("ai").write(response.content)
+            st.rerun()  # Clear stale containers
+        except AgentClientError as e:
+            st.error(f"Error generating response: {e}")
+            st.stop()
 
     # If messages have been generated, show feedback widget
-    if len(messages) > 0:
+    if len(messages) > 0 and st.session_state.last_message:
         with st.session_state.last_message:
             await handle_feedback()
 
@@ -321,12 +326,16 @@ async def handle_feedback() -> None:
         normalized_score = (feedback + 1) / 5.0
 
         agent_client: AgentClient = st.session_state.agent_client
-        await agent_client.acreate_feedback(
-            run_id=latest_run_id,
-            key="human-feedback-stars",
-            score=normalized_score,
-            kwargs={"comment": "In-line human feedback"},
-        )
+        try:
+            await agent_client.acreate_feedback(
+                run_id=latest_run_id,
+                key="human-feedback-stars",
+                score=normalized_score,
+                kwargs={"comment": "In-line human feedback"},
+            )
+        except AgentClientError as e:
+            st.error(f"Error recording feedback: {e}")
+            st.stop()
         st.session_state.last_feedback = (latest_run_id, feedback)
         st.toast("Feedback recorded", icon=":material/reviews:")
 

tests/app/test_streamlit_app.py

@@ -4,6 +4,7 @@
 import pytest
 from streamlit.testing.v1 import AppTest
 
+from client import AgentClientError
 from schema import ChatHistory, ChatMessage
 from schema.models import OpenAIModelName
 
@@ -139,3 +140,23 @@ async def amessage_iter() -> AsyncGenerator[ChatMessage, None]:
     assert tool_status.markdown[2].value == "42"
     assert response.markdown[-1].value == "The answer is 42"
     assert not at.exception
+
+
+@pytest.mark.asyncio
+async def test_app_init_error(mock_agent_client):
+    """Test the app with an error in the agent initialization"""
+    at = AppTest.from_file("../../src/streamlit_app.py").run()
+
+    # Setup mock streaming response
+    PROMPT = "What is 6 * 7?"
+    mock_agent_client.astream.side_effect = AgentClientError("Error connecting to agent")
+
+    at.toggle[0].set_value(True)  # Use Streaming = True
+    at.chat_input[0].set_value(PROMPT).run()
+    print(at)
+
+    assert at.chat_message[0].avatar == "assistant"
+    assert at.chat_message[1].avatar == "user"
+    assert at.chat_message[1].markdown[0].value == PROMPT
+    assert at.error[0].value == "Error generating response: Error connecting to agent"
+    assert not at.exception