update code

requirements.txt

@@ -1,3 +1,2 @@
-altair
-pandas
-streamlit
+streamlit
+requests

src/streamlit_app.py

@@ -1,40 +1,220 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 import streamlit as st
+import requests
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+# Page configuration
+st.set_page_config(
+    page_title="Wren AI Cloud API Demo",
+    page_icon="📊",
+    layout="wide"
+)
+
+def generate_sql(
+    api_key: str,
+    project_id: str,
+    query: str,
+    thread_id: str = "",
+) -> dict:
+    """Generate SQL from natural language query."""
+    base_url = "https://cloud.getwren.ai/api/v1"
+    endpoint = f"{base_url}/generate_sql"
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json"
+    }
+    payload = {
+        "projectId": project_id,
+        "question": query,
+    }
+    if thread_id:
+        payload["threadId"] = thread_id
+    
+    try:
+        response = requests.post(endpoint, json=payload, headers=headers)
+        response.raise_for_status()
+        return response.json()
+    except requests.exceptions.RequestException as e:
+        st.toast(f"Error generating SQL: {e}", icon="🚨")
+        return {}
+
+def generate_chart(
+    api_key: str,
+    project_id: str,
+    question: str,
+    sql: str,
+    sample_size: int = 1000,
+    thread_id: str = ""
+) -> dict:
+    """Generate a chart from query results."""
+    base_url = "https://cloud.getwren.ai/api/v1"
+    endpoint = f"{base_url}/generate_vega_chart"
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json"
+    }
+    payload = {
+        "projectId": project_id,
+        "question": question,
+        "sql": sql,
+        "sampleSize": sample_size
+    }
+    if thread_id:
+        payload["threadId"] = thread_id
+    
+    try:
+        response = requests.post(endpoint, json=payload, headers=headers)
+        response.raise_for_status()
+        return response.json()
+    except requests.exceptions.RequestException as e:
+        st.toast(f"Error generating chart: {e}", icon="🚨")
+        return {}
+
+def main():
+    st.title("📊 Wren AI Cloud API Demo")
+    st.markdown("Ask questions about your data and get both SQL queries and beautiful charts!")
+
+    # Sidebar for API configuration
+    with st.sidebar:
+        st.header("🔧 Configuration")
+        api_key = st.text_input(
+            "API Key", 
+            type="password", 
+            placeholder="sk-your-api-key-here",
+            help="Enter your Wren AI Cloud API key"
+        )
+        project_id = st.text_input(
+            "Project ID", 
+            placeholder="1234",
+            help="Enter your Wren AI Cloud project ID"
+        )
+        
+        # Sample size configuration
+        sample_size = st.slider(
+            "Chart Sample Size", 
+            min_value=100, 
+            max_value=10000, 
+            value=1000,
+            step=100,
+            help="Number of data points to include in charts"
+        )
+
+    # Main chat interface
+    if not api_key or not project_id:
+        st.warning("⚠️ Please enter your API Key and Project ID in the sidebar to get started.")
+        st.info("""
+        **How to get started:**
+        1. Enter your Wren AI API Key in the sidebar
+        2. Enter your Project ID
+        3. Ask questions about your data in natural language
+        4. Get SQL queries and interactive charts automatically!
+        """)
+        return
+
+    # Initialize chat history
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+    if "thread_id" not in st.session_state:
+        st.session_state.thread_id = ""
+
+    # Display chat history
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            if message["role"] == "user":
+                st.write(message["content"])
+            else:
+                st.write(message["content"])
+                if "sql" in message:
+                    with st.expander("📝 Generated SQL Query", expanded=False):
+                        st.code(message["sql"], language="sql")
+                if "vega_spec" in message:
+                    try:
+                        with st.expander("📊 Chart Specification", expanded=False):
+                            st.json(message["vega_spec"])
+                        st.vega_lite_chart(message["vega_spec"])
+                    except Exception as e:
+                        st.toast(f"Error rendering chart: {e}", icon="🚨")
+
+    # Chat input
+    if prompt := st.chat_input("Ask a question about your data..."):
+        # Add user message to chat history
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        
+        # Display user message
+        with st.chat_message("user"):
+            st.write(prompt)
+
+        # Generate response
+        with st.chat_message("assistant"):
+            with st.spinner("Generating SQL query..."):
+                sql_response = generate_sql(api_key, project_id, prompt, st.session_state.thread_id)
+            
+            if sql_response:
+                sql_query = sql_response.get("sql", "")
+                st.session_state.thread_id = sql_response.get("threadId", "")
+                
+                if sql_query:
+                    st.toast("SQL query generated successfully!", icon="🎉")
+                    
+                    # Store the response
+                    assistant_message = {
+                        "role": "assistant",
+                        "content": f"I've generated a SQL query for your question: '{prompt}'",
+                        "sql": sql_query
+                    }
+                    st.session_state.messages.append(assistant_message)
+                    st.write(assistant_message["content"])
+
+                    # Display SQL query
+                    with st.expander("📝 Generated SQL Query", expanded=False):
+                        st.code(sql_query, language="sql")
+                    
+                    # Generate chart
+                    with st.spinner("Generating chart..."):
+                        chart_response = generate_chart(api_key, project_id, prompt, sql_query, sample_size, st.session_state.thread_id)
+                    
+                    if chart_response:
+                        vega_spec = chart_response.get("vegaSpec", {})
+                        if vega_spec:
+                            st.toast("Chart generated successfully!", icon="🎉")
+                        
+                            assistant_message = {
+                                "role": "assistant",
+                                "content": f"I've generated a Chart for your question: '{prompt}'",
+                                "vega_spec": vega_spec
+                            }
+                            st.session_state.messages.append(assistant_message)
+                            st.write(assistant_message["content"])
+
+                            # Display chart
+                            try:
+                                # Show chart specification in expander
+                                with st.expander("📊 Chart Specification", expanded=False):
+                                    st.json(vega_spec)
+                                st.vega_lite_chart(vega_spec)
+                            except Exception as e:
+                                st.toast(f"Error rendering chart: {e}", icon="🚨")
+                        else:
+                            st.toast("Failed to generate chart. Please check your query and try again.", icon="🚨")
+                    else:
+                        st.toast("Failed to generate chart. Please check your query and try again.", icon="🚨")
+                else:
+                    st.toast("No SQL query was generated. Please try rephrasing your question.", icon="🚨")
+                    assistant_message = {
+                        "role": "assistant",
+                        "content": "I couldn't generate a SQL query for your question. Please try rephrasing it or make sure it's related to your data."
+                    }
+                    st.session_state.messages.append(assistant_message)
+            else:
+                assistant_message = {
+                    "role": "assistant", 
+                    "content": "Sorry, I couldn't process your request. Please check your API credentials and try again."
+                }
+                st.session_state.messages.append(assistant_message)
+
+    # Clear chat button
+    if st.sidebar.button("🗑️ Clear Chat History"):
+        st.session_state.messages = []
+        st.session_state.thread_id = ""
+        st.rerun()
+
+if __name__ == "__main__":
+    main()