Update app.py

app.py

@@ -1,6 +1,5 @@
 import os
 import json
-import time
 import requests
 import pandas as pd
 import streamlit as st
@@ -18,7 +17,7 @@ CHAT_URL = get_env_variable("CHAT_URL")
 # API interaction
 def fetch_results(query: str) -> list[dict]:
     """
-    Send a query to the chat API and return the parsed JSON 'results' list.
+    Send a query to the chat API and return the parsed JSON response.
     """
     payload = {"query": query}
     response = requests.post(
@@ -26,27 +25,24 @@ def fetch_results(query: str) -> list[dict]:
         headers={"Content-Type": "application/json"},
         data=json.dumps(payload),
     )
-    # response.raise_for_status()
-    response = response.json()
-    return response
+    return response.json()
 
 # UI components
-def display_messages(history: list[dict]) -> None:
-    """Render chat messages from history."""
-    for msg in history:
-        with st.chat_message(msg["role"]):
-            st.markdown(msg["content"])
-
-
-def display_results_table(results: list[dict]) -> None:
+def display_results_table(results: dict) -> None:
     """Normalize and display API results as a Streamlit table."""
     if "results" not in results:
-        st.info(results['detail'])
+        st.info(results.get('detail', 'No results found.'))
         return
     df = pd.json_normalize(results['results'])
-    st.subheader("Results")
     st.dataframe(df)
 
+def display_messages(history: list[dict]) -> None:
+    """Render chat messages from history including tables."""
+    for msg in history:
+        with st.chat_message(msg["role"]):
+            st.markdown(msg["content"])
+            if "data" in msg:
+                display_results_table(msg["data"])
 
 def main() -> None:
     st.title("BookDB 📚 🤖")
@@ -55,30 +51,26 @@ def main() -> None:
     if "messages" not in st.session_state:
         st.session_state.messages = []
 
-    display_messages(st.session_state.messages)
-
-    # User input
+    # Get user input
     query = st.chat_input("Ask about books...")
-    if not query:
-        return
-
-    # Add and render user message
-    user_msg = query.strip()
-    st.session_state.messages.append({"role": "user", "content": user_msg})
-    with st.chat_message("user"):
-        st.markdown(user_msg)
-
-    # Fetch and display results
-    with st.chat_message("assistant"):
-        results = fetch_results(user_msg)
-        display_results_table(results)
-
-    # Store assistant placeholder in history
-    st.session_state.messages.append({
-        "role": "assistant",
-        "content": "Displayed results as a table."
-    })
-
+    
+    if query:
+        # Add user message to history
+        st.session_state.messages.append({
+            "role": "user",
+            "content": query.strip()
+        })
+        
+        # Fetch results and add assistant response
+        results = fetch_results(query)
+        st.session_state.messages.append({
+            "role": "assistant",
+            "content": "Here are the results:",
+            "data": results
+        })
+
+    # Display all messages
+    display_messages(st.session_state.messages)
 
 if __name__ == "__main__":
-    main()
+    main()