refactor

.gitignore

@@ -1 +1,2 @@
-.DS_Store
+.DS_Store
+__pycache__

Dockerfile

@@ -31,4 +31,4 @@ EXPOSE 8501
 
 HEALTHCHECK CMD curl --fail http://localhost:8501/_stcore/health
 
-ENTRYPOINT ["streamlit", "run", "src/streamlit_app.py", "--server.port=8501", "--server.address=0.0.0.0"]
+ENTRYPOINT ["streamlit", "run", "src/app.py", "--server.port=8501", "--server.address=0.0.0.0"]

Makefile

@@ -1,5 +1,5 @@
 run:
-	poetry run streamlit run src/streamlit_app.py
+	poetry run streamlit run src/app.py
 
 deps:
 	poetry export --without-hashes --format=requirements.txt > requirements.txt

README.md

@@ -14,7 +14,7 @@ license: mit
 
 # Welcome to Streamlit!
 
-Edit `/src/streamlit_app.py` to customize this app to your heart's desire. :heart:
+Edit `/src/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).

src/apis.py

@@ -0,0 +1,61 @@
+import requests
+
+
+def generate_sql(
+    api_key: str,
+    project_id: str,
+    query: str,
+    thread_id: str = "",
+) -> tuple[dict, str]:
+    """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:
+        return {}, e
+
+
+def generate_chart(
+    api_key: str,
+    project_id: str,
+    question: str,
+    sql: str,
+    sample_size: int = 1000,
+    thread_id: str = ""
+) -> tuple[dict, str]:
+    """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:
+        return {}, e

src/{streamlit_app.py → app.py}

@@ -1,5 +1,7 @@
 import streamlit as st
-import requests
+
+from apis import generate_sql, generate_chart
+
 
 # Page configuration
 st.set_page_config(
@@ -8,65 +10,6 @@ st.set_page_config(
     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")
@@ -102,7 +45,7 @@ def main():
         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
+        1. Enter your Wren AI Cloud 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!
@@ -145,7 +88,7 @@ def main():
         # 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)
+                sql_response, error = generate_sql(api_key, project_id, prompt, st.session_state.thread_id)
             
             if sql_response:
                 sql_query = sql_response.get("sql", "")
@@ -169,7 +112,7 @@ def main():
                     
                     # 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)
+                        chart_response, error = 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", {})
@@ -195,7 +138,7 @@ def main():
                         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="🚨")
+                        st.toast(f"Failed to generate chart. Please check your query and try again.: {error}", icon="🚨")
                 else:
                     st.toast("No SQL query was generated. Please try rephrasing your question.", icon="🚨")
                     assistant_message = {
@@ -204,6 +147,7 @@ def main():
                     }
                     st.session_state.messages.append(assistant_message)
             else:
+                st.toast(f"Error generating SQL: {error}", icon="🚨")
                 assistant_message = {
                     "role": "assistant", 
                     "content": "Sorry, I couldn't process your request. Please check your API credentials and try again."
@@ -216,5 +160,6 @@ def main():
         st.session_state.thread_id = ""
         st.rerun()
 
+
 if __name__ == "__main__":
     main()