diff --git a/frontend/app.py b/frontend/app.py index 696511c..e1159ff 100644 --- 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 index 414fabd..b5afcf4 100644 --- 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 index 1d17efb..9d96a14 100644 --- 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}