Update app.py

app.py

@@ -2,8 +2,6 @@ import streamlit as st
 import asyncio
 import xai_sdk
 import os
-import base64
-import io
 from datetime import datetime
 
 # Set up the XAI client
@@ -35,6 +33,12 @@ for message in st.session_state.messages:
 
 prompt = st.chat_input("What would you like to ask?")
 
+async def generate_response(prompt, max_len):
+    full_response = ""
+    async for token in client.sampler.sample(prompt="", inputs=(prompt,), max_len=max_len):
+        full_response += token.token_str
+        yield full_response
+
 if prompt:
     st.session_state.messages.append({"role": "user", "content": prompt})
     with st.chat_message("user"):
@@ -45,14 +49,10 @@ if prompt:
         full_response = ""
 
         # Use asyncio to run the asynchronous XAI SDK
-        async def generate_response():
-            nonlocal full_response
-            async for token in client.sampler.sample(prompt="", inputs=(prompt,), max_len=max_tokens):
-                full_response += token.token_str
-                message_placeholder.markdown(full_response + "▌")
-            message_placeholder.markdown(full_response)
-
-        asyncio.run(generate_response())
+        for response in asyncio.run(generate_response(prompt, max_tokens)):
+            message_placeholder.markdown(response + "▌")
+            full_response = response
+        message_placeholder.markdown(full_response)
 
     st.session_state.messages.append({"role": "assistant", "content": full_response})
 
@@ -60,14 +60,12 @@ if prompt:
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
     filename = f"output_{timestamp}_{prompt[:20].replace(' ', '_')}.txt"
     
-    output = io.StringIO()
-    output.write(f"Prompt: {prompt}\n\n")
-    output.write(f"Response: {full_response}\n")
+    output_content = f"Prompt: {prompt}\n\nResponse: {full_response}\n"
     
     # Create download button
     st.download_button(
         label="📥 Download Output",
-        data=output.getvalue(),
+        data=output_content,
         file_name=filename,
         mime="text/plain",
     )