add new page and refine sql query ui

poetry.lock

@@ -974,6 +974,21 @@ files = [
     {file = "smmap-5.0.2.tar.gz", hash = "sha256:26ea65a03958fa0c8a1c7e8c7a58fdc77221b8910f6be2131affade476898ad5"},
 ]
 
+[[package]]
+name = "sqlparse"
+version = "0.5.3"
+description = "A non-validating SQL parser."
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "sqlparse-0.5.3-py3-none-any.whl", hash = "sha256:cf2196ed3418f3ba5de6af7e82c694a9fbdbfecccdfc72e281548517081f16ca"},
+    {file = "sqlparse-0.5.3.tar.gz", hash = "sha256:09f67787f56a0b16ecdbde1bfc7f5d9c3371ca683cfeaa8e6ff60b4807ec9272"},
+]
+
+[package.extras]
+dev = ["build", "hatch"]
+doc = ["sphinx"]
+
 [[package]]
 name = "streamlit"
 version = "1.46.1"
@@ -1139,4 +1154,4 @@ watchmedo = ["PyYAML (>=3.10)"]
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.12,<3.13"
-content-hash = "d3018cb9ea02785fe38a30f0f0e1196ad02058106619e7ba3a2a11c4a78753a2"
+content-hash = "67b448796799eb25e83725fc27c23ff400273d48c767151c6462d8e1545052fe"

pyproject.toml

@@ -12,6 +12,7 @@ python = ">=3.12,<3.13"
 streamlit = "^1.46.1"
 requests = "^2.32.4"
 watchdog = "^6.0.0"
+sqlparse = "^0.5.3"
 
 [tool.poetry.group.dev.dependencies]
 python-dotenv = "^1.1.1"

src/pages/1_Chart_Generation.py

@@ -1,6 +1,7 @@
 import streamlit as st
 
 from apis import generate_sql, generate_chart
+from utils import format_sql
 
 
 def main():
@@ -46,7 +47,7 @@ def main():
                 st.write(message["content"])
                 if "sql" in message:
                     with st.expander("📝 Generated SQL Query", expanded=False):
-                        st.code(message["sql"], language="sql")
+                        st.code(format_sql(message["sql"]), language="sql")
                 if "vega_spec" in message:
                     try:
                         with st.expander("📊 Chart Specification", expanded=False):
@@ -87,7 +88,7 @@ def main():
 
                     # Display SQL query
                     with st.expander("📝 Generated SQL Query", expanded=False):
-                        st.code(sql_query, language="sql")
+                        st.code(format_sql(sql_query), language="sql")
                     
                     # Generate chart
                     with st.spinner("Generating chart..."):

src/pages/2_Query_And_Answer.py

@@ -0,0 +1,159 @@
+import streamlit as st
+import pandas as pd
+
+from apis import ask, run_sql
+from utils import format_sql
+
+
+def main():
+    st.title("💬 Wren AI Cloud API Demo - Query and Answer")
+
+    if "api_key" not in st.session_state or "project_id" not in st.session_state:
+        st.error("Please enter your API Key and Project ID in the sidebar of Home page to get started.")
+        return
+    if not st.session_state.api_key or not st.session_state.project_id:
+        st.error("Please enter your API Key and Project ID in the sidebar of Home page to get started.")
+        return
+    
+    api_key = st.session_state.api_key
+    project_id = st.session_state.project_id
+
+    st.markdown('Using APIs: [Ask](https://wrenai.readme.io/reference/post_ask-1), [Run SQL](https://wrenai.readme.io/reference/cloud_post_run-sql)')
+
+    # Sidebar for API configuration
+    with st.sidebar:
+        st.header("🔧 Configuration")
+        # Sample size configuration
+        sample_size = st.slider(
+            "Sample Size", 
+            min_value=100, 
+            max_value=10000, 
+            value=1000,
+            step=100,
+            help="Number of data points to include in results"
+        )
+        
+    # Initialize chat history
+    if "qa_messages" not in st.session_state:
+        st.session_state.qa_messages = []
+    if "qa_thread_id" not in st.session_state:
+        st.session_state.qa_thread_id = ""
+
+    # Display chat history
+    for message in st.session_state.qa_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(format_sql(message["sql"]), language="sql")
+                        
+                        # Add button to run SQL
+                        if st.button("🔄 Run SQL Query", key=f"run_sql_{message.get('message_id', 'unknown')}"):
+                            with st.spinner("Executing SQL query..."):
+                                sql_result, error = run_sql(api_key, project_id, message["sql"], st.session_state.qa_thread_id, sample_size)
+                            
+                            if sql_result:
+                                data = sql_result.get("records", [])                                
+                                if data:
+                                    # Convert to DataFrame for better display
+                                    df = pd.DataFrame(data)
+                                    st.success("SQL query executed successfully!")
+                                    st.dataframe(df, use_container_width=True)
+                                else:
+                                    st.info("Query executed but returned no data.")
+                            else:
+                                st.error(f"Error executing SQL: {error}")
+                
+                if "sql_results" in message:
+                    st.subheader("🔍 Query Results")
+                    if message["sql_results"]:
+                        st.dataframe(message["sql_results"], use_container_width=True)
+                    else:
+                        st.info("No results returned from the query.")
+
+    # Chat input
+    if prompt := st.chat_input("Ask a question about your data..."):
+        # Add user message to chat history
+        st.session_state.qa_messages.append({"role": "user", "content": prompt})
+        
+        # Display user message
+        with st.chat_message("user"):
+            st.write(prompt)
+
+        # Generate response using ask API
+        with st.chat_message("assistant"):
+            with st.spinner("Generating answer..."):
+                ask_response, error = ask(api_key, project_id, prompt, st.session_state.qa_thread_id, sample_size=sample_size)
+            
+            if ask_response:
+                answer = ask_response.get("summary", "")
+                sql_query = ask_response.get("sql", "")
+                st.session_state.qa_thread_id = ask_response.get("threadId", "")
+                
+                if answer:
+                    st.toast("Answer generated successfully!", icon="🎉")
+                    
+                    # Create unique message ID
+                    message_id = len(st.session_state.qa_messages)
+                    
+                    # Store the response
+                    assistant_message = {
+                        "role": "assistant",
+                        "content": answer,
+                        "message_id": message_id
+                    }
+                    
+                    if sql_query:
+                        assistant_message["sql"] = sql_query
+                    
+                    st.session_state.qa_messages.append(assistant_message)
+                    st.write(answer)
+
+                    # Display SQL query if available
+                    if sql_query:
+                        with st.expander("🔍 Generated SQL Query", expanded=False):
+                            st.code(format_sql(sql_query), language="sql")
+                            
+                            # Add button to run SQL
+                            if st.button("🔄 Run SQL Query", key=f"run_sql_{message_id}"):
+                                with st.spinner("Executing SQL query..."):
+                                    sql_result, error = run_sql(api_key, project_id, sql_query, st.session_state.qa_thread_id, sample_size)
+                                
+                                if sql_result:
+                                    data = sql_result.get("records", [])                                    
+                                    if data:
+                                        # Convert to DataFrame for better display
+                                        df = pd.DataFrame(data)
+                                        st.success("SQL query executed successfully!")
+                                        st.dataframe(df, use_container_width=True)
+                                    else:
+                                        st.info("Query executed but returned no data.")
+                                else:
+                                    st.error(f"Error executing SQL: {error}")
+                else:
+                    st.toast("No answer was generated. Please try rephrasing your question.", icon="🤔")
+                    assistant_message = {
+                        "role": "assistant",
+                        "content": "I couldn't generate an answer for your question. Please try rephrasing it or make sure it's related to your data."
+                    }
+                    st.session_state.qa_messages.append(assistant_message)
+            else:
+                st.toast(f"Error generating answer: {error}", icon="🤔")
+                assistant_message = {
+                    "role": "assistant", 
+                    "content": "Sorry, I couldn't process your request. Please check your API credentials and try again."
+                }
+                st.session_state.qa_messages.append(assistant_message)
+
+    # Clear chat button
+    if st.sidebar.button("🧹 Clear Chat History"):
+        st.session_state.qa_messages = []
+        st.session_state.qa_thread_id = ""
+        st.rerun()
+
+
+if __name__ == "__main__":
+    main()

src/utils.py

@@ -0,0 +1,9 @@
+import sqlparse
+
+
+def format_sql(sql: str) -> str:
+    return sqlparse.format(
+        sql,
+        reindent=True,
+        keyword_case="upper",
+    )