Update app.py

app.py

@@ -1,4 +1,87 @@
 import streamlit as st
+import asyncio
+import xai_sdk
+import os
+import base64
+import io
+from datetime import datetime
 
-x = st.slider('Select a value')
-st.write(x, 'squared is', x * x)
+# Set up the XAI client
+api_key = st.sidebar.text_input("Enter your XAI API key", type="password")
+os.environ["XAI_API_KEY"] = api_key
+
+# Initialize the XAI client
+@st.cache_resource
+def get_xai_client():
+    return xai_sdk.Client()
+
+client = get_xai_client()
+
+# Streamlit app
+st.title("🤖 XAI Streamlit Client")
+
+# Sidebar
+st.sidebar.header("Settings")
+max_tokens = st.sidebar.slider("Max Tokens", 1, 100, 50)
+temperature = st.sidebar.slider("Temperature", 0.0, 1.0, 0.7)
+
+# Chat interface
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+prompt = st.chat_input("What would you like to ask?")
+
+if prompt:
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    with st.chat_message("user"):
+        st.markdown(prompt)
+
+    with st.chat_message("assistant"):
+        message_placeholder = st.empty()
+        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())
+
+    st.session_state.messages.append({"role": "assistant", "content": full_response})
+
+    # Generate output file
+    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")
+    
+    # Create download button
+    st.download_button(
+        label="📥 Download Output",
+        data=output.getvalue(),
+        file_name=filename,
+        mime="text/plain",
+    )
+
+# Add some helpful information
+st.sidebar.markdown("---")
+st.sidebar.subheader("📌 How to use")
+st.sidebar.markdown("""
+1. Enter your XAI API key in the sidebar
+2. Adjust the settings as needed
+3. Type your prompt in the chat input
+4. Download the output if desired
+""")
+
+# Footer
+st.sidebar.markdown("---")
+st.sidebar.markdown("Made with ❤️ using Streamlit and XAI SDK")