use write_stream when consuming the response generator

nwaughachukwuma · Oct 26, 2024 · 911ed1e · 911ed1e
1 parent 73a1f5a
commit 911ed1e
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 36 deletions.
diff --git a/frontend/app.py b/frontend/app.py
@@ -47,33 +47,33 @@
     with st.chat_message("user"):
         st.write(prompt)
 
-    ai_message = chat_request(prompt, content_type)
-
-    if ai_message:
-        st.session_state.messages.append({"role": "assistant", "content": ai_message})
-
-        with st.chat_message("assistant"):
-            st.write(ai_message)
-
-            # Show generate button if enough context
-            if len(st.session_state.messages) >= 2:
-                if st.button("Generate Audiocast"):
-                    with st.spinner("Generating your audiocast..."):
-                        # Generate audiocast
-                        audiocast_response = httpx.post(
-                            f"{BACKEND_URL}/api/generate-audiocast",
-                            json={
-                                "query": prompt,
-                                "type": content_type,
-                                "chat_history": st.session_state.messages,
-                            },
-                        )
-
-                        if audiocast_response.status_code == 200:
-                            st.session_state.current_audiocast = (
-                                audiocast_response.json()
-                            )
-                            st.rerun()
+    with st.chat_message("assistant"):
+        response_generator = chat_request(prompt, content_type)
+
+        ai_message = st.write_stream(response_generator)
+
+        if ai_message:
+            st.session_state.messages.append(
+                {"role": "assistant", "content": ai_message}
+            )
+
+        # Show generate button if enough context
+        if len(st.session_state.messages) >= 2:
+            if st.button("Generate Audiocast"):
+                with st.spinner("Generating your audiocast..."):
+                    # Generate audiocast
+                    audiocast_response = httpx.post(
+                        f"{BACKEND_URL}/api/generate-audiocast",
+                        json={
+                            "query": prompt,
+                            "type": content_type,
+                            "chat_history": st.session_state.messages,
+                        },
+                    )
+
+                    if audiocast_response.status_code == 200:
+                        st.session_state.current_audiocast = audiocast_response.json()
+                        st.rerun()
 
 # Display current audiocast if available
 if st.session_state.current_audiocast:

diff --git a/frontend/chat_utils.py b/frontend/chat_utils.py
@@ -1,4 +1,4 @@
-from typing import Dict, List, Literal
+from typing import Any, Callable, Dict, List, Literal, Optional
 
 import httpx
 import streamlit as st
@@ -16,7 +16,11 @@
 }
 
 
-def chat_request(prompt: str, content_type: ContentType):
+def chat_request(
+    prompt: str,
+    content_type: ContentType,
+    on_finish: Optional[Callable[[str], Any]] = None,
+):
     """
     Send a chat request to the backend server and return the AI response.
     """
@@ -30,9 +34,4 @@ def chat_request(prompt: str, content_type: ContentType):
     )
 
     response.raise_for_status()
-
-    ai_message = ""
-    for line in response.iter_lines():
-        ai_message += line
-
-    return ai_message
+    return response.iter_lines()
diff --git a/frontend/example_utils.py b/frontend/example_utils.py
@@ -35,7 +35,9 @@ def display_example_cards():
                 # Add selected example to messages and trigger rerun to enter chat mode
                 st.session_state.messages.append({"role": "user", "content": example})
 
-                ai_message = chat_request(example, content_type)
+                response_generator = chat_request(example, content_type)
+
+                ai_message = st.write_stream(response_generator)
                 if ai_message:
                     st.session_state.messages.append(
                         {"role": "assistant", "content": ai_message}