add new page and refactor

src/apis.py

@@ -3,6 +3,9 @@ from typing import Optional
 import requests
 
 
+API_BASE_URL = "https://cloud.getwren.ai/api/v1"
+
+
 def ask(
     api_key: str,
     project_id: str,
@@ -12,8 +15,7 @@ def ask(
     sample_size: int = 1000,
 ) -> tuple[dict, str]:
     """Ask a question and get an answer with SQL and explanation."""
-    base_url = "https://cloud.getwren.ai/api/v1"
-    endpoint = f"{base_url}/ask"
+    endpoint = f"{API_BASE_URL}/ask"
     headers = {
         "Authorization": f"Bearer {api_key}",
         "Content-Type": "application/json"
@@ -44,8 +46,7 @@ def generate_sql(
     return_sql_dialect: bool = False,
 ) -> tuple[dict, str]:
     """Generate SQL from natural language query."""
-    base_url = "https://cloud.getwren.ai/api/v1"
-    endpoint = f"{base_url}/generate_sql"
+    endpoint = f"{API_BASE_URL}/generate_sql"
     headers = {
         "Authorization": f"Bearer {api_key}",
         "Content-Type": "application/json"
@@ -75,8 +76,7 @@ def run_sql(
     limit: int = 1000,
 ) -> tuple[dict, str]:
     """Execute SQL query and return results."""
-    base_url = "https://cloud.getwren.ai/api/v1"
-    endpoint = f"{base_url}/run_sql"
+    endpoint = f"{API_BASE_URL}/run_sql"
     headers = {
         "Authorization": f"Bearer {api_key}",
         "Content-Type": "application/json"
@@ -107,8 +107,7 @@ def generate_summary(
     sample_size: int = 1000,
 ) -> tuple[dict, str]:
     """Generate a summary of query results."""
-    base_url = "https://cloud.getwren.ai/api/v1"
-    endpoint = f"{base_url}/generate_summary"
+    endpoint = f"{API_BASE_URL}/generate_summary"
     headers = {
         "Authorization": f"Bearer {api_key}",
         "Content-Type": "application/json"
@@ -136,8 +135,7 @@ def stream_explanation(
     query_id: str,
 ) -> tuple[dict, str]:
     """Stream explanation for a query."""
-    base_url = "https://cloud.getwren.ai/api/v1"
-    endpoint = f"{base_url}/stream_explanation/{query_id}"
+    endpoint = f"{API_BASE_URL}/stream_explanation/{query_id}"
     headers = {
         "Authorization": f"Bearer {api_key}",
         "Accept": "text/event-stream",
@@ -160,8 +158,7 @@ def generate_chart(
     sample_size: int = 1000,
 ) -> tuple[dict, str]:
     """Generate a chart from query results."""
-    base_url = "https://cloud.getwren.ai/api/v1"
-    endpoint = f"{base_url}/generate_vega_chart"
+    endpoint = f"{API_BASE_URL}/generate_vega_chart"
     headers = {
         "Authorization": f"Bearer {api_key}",
         "Content-Type": "application/json"
@@ -192,8 +189,7 @@ def stream_ask(
     sample_size: int = 1000,
 ) -> tuple[dict, str]:
     """Stream ask endpoint for real-time responses."""
-    base_url = "https://cloud.getwren.ai/api/v1"
-    endpoint = f"{base_url}/stream/ask"
+    endpoint = f"{API_BASE_URL}/stream/ask"
     headers = {
         "Authorization": f"Bearer {api_key}",
         "Content-Type": "application/json",
@@ -211,7 +207,7 @@ def stream_ask(
     try:
         response = requests.post(endpoint, json=payload, headers=headers, stream=True)
         response.raise_for_status()
-        return response.json(), ""
+        return response, ""
     except requests.exceptions.RequestException as e:
         return {}, str(e)
 
@@ -225,8 +221,7 @@ def stream_generate_sql(
     return_sql_dialect: bool = False,
 ) -> tuple[dict, str]:
     """Stream SQL generation endpoint for real-time responses."""
-    base_url = "https://cloud.getwren.ai/api/v1"
-    endpoint = f"{base_url}/stream/generate_sql"
+    endpoint = f"{API_BASE_URL}/stream/generate_sql"
     headers = {
         "Authorization": f"Bearer {api_key}",
         "Content-Type": "application/json",
@@ -254,8 +249,7 @@ def get_models(
     project_id: str,
 ) -> tuple[dict, str]:
     """Get latest deployed models for a project."""
-    base_url = "https://cloud.getwren.ai/api/v1"
-    endpoint = f"{base_url}/projects/{project_id}/models"
+    endpoint = f"{API_BASE_URL}/projects/{project_id}/models"
     headers = {
         "Authorization": f"Bearer {api_key}",
     }
@@ -273,8 +267,7 @@ def get_instructions(
     project_id: str,
 ) -> tuple[dict, str]:
     """Get all instructions for a project."""
-    base_url = "https://cloud.getwren.ai/api/v1"
-    endpoint = f"{base_url}/projects/{project_id}/knowledge/instructions"
+    endpoint = f"{API_BASE_URL}/projects/{project_id}/knowledge/instructions"
     headers = {
         "Authorization": f"Bearer {api_key}",
     }
@@ -295,8 +288,7 @@ def create_instruction(
     questions: Optional[list[str]] = None,
 ) -> tuple[dict, str]:
     """Create a new instruction for a project."""
-    base_url = "https://cloud.getwren.ai/api/v1"
-    endpoint = f"{base_url}/projects/{project_id}/knowledge/instructions"
+    endpoint = f"{API_BASE_URL}/projects/{project_id}/knowledge/instructions"
     headers = {
         "Authorization": f"Bearer {api_key}",
         "Content-Type": "application/json"
@@ -324,8 +316,7 @@ def update_instruction(
     questions: Optional[list[str]] = None,
 ) -> tuple[dict, str]:
     """Update an existing instruction."""
-    base_url = "https://cloud.getwren.ai/api/v1"
-    endpoint = f"{base_url}/projects/{project_id}/knowledge/instructions/{instruction_id}"
+    endpoint = f"{API_BASE_URL}/projects/{project_id}/knowledge/instructions/{instruction_id}"
     headers = {
         "Authorization": f"Bearer {api_key}",
         "Content-Type": "application/json"
@@ -350,8 +341,7 @@ def delete_instruction(
     instruction_id: str,
 ) -> str:
     """Delete an instruction."""
-    base_url = "https://cloud.getwren.ai/api/v1"
-    endpoint = f"{base_url}/projects/{project_id}/knowledge/instructions/{instruction_id}"
+    endpoint = f"{API_BASE_URL}/projects/{project_id}/knowledge/instructions/{instruction_id}"
     headers = {
         "Authorization": f"Bearer {api_key}",
     }
@@ -369,8 +359,7 @@ def get_sql_pairs(
     project_id: str,
 ) -> tuple[dict, str]:
     """Get all SQL pairs for a project."""
-    base_url = "https://cloud.getwren.ai/api/v1"
-    endpoint = f"{base_url}/projects/{project_id}/knowledge/sql_pairs"
+    endpoint = f"{API_BASE_URL}/projects/{project_id}/knowledge/sql_pairs"
     headers = {
         "Authorization": f"Bearer {api_key}",
     }
@@ -390,8 +379,7 @@ def create_sql_pair(
     sql: str,
 ) -> tuple[dict, str]:
     """Create a new SQL pair for a project."""
-    base_url = "https://cloud.getwren.ai/api/v1"
-    endpoint = f"{base_url}/projects/{project_id}/knowledge/sql_pairs"
+    endpoint = f"{API_BASE_URL}/projects/{project_id}/knowledge/sql_pairs"
     headers = {
         "Authorization": f"Bearer {api_key}",
         "Content-Type": "application/json"
@@ -417,8 +405,7 @@ def update_sql_pair(
     sql: str,
 ) -> tuple[dict, str]:
     """Update an existing SQL pair."""
-    base_url = "https://cloud.getwren.ai/api/v1"
-    endpoint = f"{base_url}/projects/{project_id}/knowledge/sql_pairs/{sql_pair_id}"
+    endpoint = f"{API_BASE_URL}/projects/{project_id}/knowledge/sql_pairs/{sql_pair_id}"
     headers = {
         "Authorization": f"Bearer {api_key}",
         "Content-Type": "application/json"
@@ -442,8 +429,7 @@ def delete_sql_pair(
     sql_pair_id: str,
 ) -> str:
     """Delete an SQL pair."""
-    base_url = "https://cloud.getwren.ai/api/v1"
-    endpoint = f"{base_url}/projects/{project_id}/knowledge/sql_pairs/{sql_pair_id}"
+    endpoint = f"{API_BASE_URL}/projects/{project_id}/knowledge/sql_pairs/{sql_pair_id}"
     headers = {
         "Authorization": f"Bearer {api_key}",
     }

src/pages/1_Query_And_Answer.py

@@ -23,7 +23,11 @@ def main():
     # Sidebar for API configuration
     with st.sidebar:
         st.header("🔧 Configuration")
-        # Sample size configuration
+        language = st.text_input(
+            "Language",
+            "English",
+            help="Language of the response. Default is English."
+        )
         sample_size = st.slider(
             "Sample Size", 
             min_value=100, 
@@ -86,7 +90,7 @@ def main():
         # 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)
+                ask_response, error = ask(api_key, project_id, prompt, st.session_state.qa_thread_id, sample_size=sample_size, language=language)
             
             if ask_response:
                 answer = ask_response.get("summary", "")

src/pages/2_Chart_Generation.py

@@ -22,7 +22,6 @@ def main():
     # Sidebar for API configuration
     with st.sidebar:
         st.header("🔧 Configuration")
-        # Sample size configuration
         sample_size = st.slider(
             "Chart Sample Size", 
             min_value=100, 

src/pages/5_Ask_Streaming.py

@@ -0,0 +1,236 @@
+import streamlit as st
+import pandas as pd
+import requests
+import json
+import time
+from typing import Iterator, Dict, Any
+
+from apis import stream_ask, run_sql
+from utils import format_sql
+
+
+def stream_ask_api(api_key: str, project_id: str, question: str, thread_id: str = "", 
+                   language: str = "English", sample_size: int = 1000) -> Iterator[Dict[str, Any]]:
+    """Stream ask endpoint with proper SSE handling."""
+    try:
+        response, error = stream_ask(api_key, project_id, question, thread_id, language, sample_size)
+        if error:
+            raise Exception(error)
+
+        for line in response.iter_lines():
+            if line:
+                line_str = line.decode('utf-8')
+                if line_str.startswith('data: '):
+                    data_str = line_str[6:]  # Remove 'data: ' prefix
+                    if data_str.strip() == '[DONE]':
+                        break
+                    try:
+                        data = json.loads(data_str)
+                        yield data
+                    except json.JSONDecodeError:
+                        continue
+    except requests.exceptions.RequestException as e:
+        yield {"error": str(e)}
+
+
+def main():
+    st.title("💬 Wren AI Cloud API Demo - Ask Streaming")
+
+    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: [Stream Ask](https://wrenai.readme.io/reference/post_stream-ask-1)')
+
+    # Sidebar for API configuration
+    with st.sidebar:
+        st.header("🔧 Configuration")
+        language = st.text_input(
+            "Language",
+            "English",
+            help="Language of the response. Default is English."
+        )
+        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 "streaming_messages" not in st.session_state:
+        st.session_state.streaming_messages = []
+    if "streaming_thread_id" not in st.session_state:
+        st.session_state.streaming_thread_id = ""
+
+    # Display chat history
+    for message in st.session_state.streaming_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.streaming_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}")
+
+    # Chat input
+    if prompt := st.chat_input("Ask a question about your data..."):
+        # Add user message to chat history
+        st.session_state.streaming_messages.append({"role": "user", "content": prompt})
+        
+        # Display user message
+        with st.chat_message("user"):
+            st.write(prompt)
+
+        # Generate response using stream ask API
+        with st.chat_message("assistant"):
+            # Create containers for streaming updates
+            status_container = st.empty()
+            progress_container = st.empty()
+            content_container = st.empty()
+            
+            final_answer = ""
+            final_sql = ""
+            current_state = ""
+            
+            # Stream the response
+            try:
+                for event in stream_ask_api(api_key, project_id, prompt, st.session_state.streaming_thread_id, sample_size=sample_size, language=language):
+                    if "error" in event:
+                        st.error(f"Error: {event['error']}")
+                        break
+                    
+                    # Handle different event types
+                    if event.get("type") == "state":
+                        current_state = event.get("data", {}).get("state", "")
+                        rephrased_question = event.get("data", {}).get("rephrasedQuestion", "")
+                        intent_reasoning = event.get("data", {}).get("intentReasoning", "")
+                        retrieved_tables = event.get("data", {}).get("retrievedTables", "")
+                        sql_generation_reasoning = event.get("data", {}).get("sqlGenerationReasoning", "")
+                        current_state_message = f"💬 {current_state}\n\n"
+                        if rephrased_question:
+                            current_state_message += f"\nRerephrased Question: \n{rephrased_question}\n"
+                        if intent_reasoning:
+                            current_state_message += f"\nIntent Reasoning: \n{intent_reasoning}\n"
+                        if retrieved_tables:
+                            current_state_message += f"\nRetrieved Tables: \n{retrieved_tables}\n"
+                        if sql_generation_reasoning:
+                            current_state_message += f"\nSQL Generation Reasoning: \n{sql_generation_reasoning}\n"
+
+                        status_container.info(current_state_message)
+                    elif event.get("type") == "content_block_delta":
+                        delta = event.get("delta", {})
+                        if delta.get("type") == "text_delta":
+                            final_answer += delta.get("text", "")
+                            content_container.markdown(final_answer)
+                    
+                    elif event.get("type") == "message_stop":
+                        # Extract final data from the event
+                        if "summary" in event:
+                            final_answer = event["summary"]
+                        if "sql" in event:
+                            final_sql = event["sql"]
+                        if "threadId" in event:
+                            st.session_state.streaming_thread_id = event["threadId"]
+                        break
+                    
+                    # Small delay to make streaming visible
+                    time.sleep(0.1)
+                
+                # Clear status and show final result
+                status_container.empty()
+                progress_container.empty()
+                
+                if final_answer:
+                    st.toast("Answer generated successfully!", icon="🎉")
+                    
+                    # Create unique message ID
+                    message_id = len(st.session_state.streaming_messages)
+                    
+                    # Store the response
+                    assistant_message = {
+                        "role": "assistant",
+                        "content": final_answer,
+                        "message_id": message_id
+                    }
+                    
+                    if final_sql:
+                        assistant_message["sql"] = final_sql
+                    
+                    st.session_state.streaming_messages.append(assistant_message)
+                    content_container.write(final_answer)
+
+                    # Display SQL query if available
+                    if final_sql:
+                        with st.expander("=Generated SQL Query", expanded=False):
+                            st.code(format_sql(final_sql), 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, final_sql, st.session_state.streaming_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.streaming_messages.append(assistant_message)
+                    content_container.write(assistant_message["content"])
+            
+            except Exception as e:
+                st.error(f"Error during streaming: {str(e)}")
+                assistant_message = {
+                    "role": "assistant", 
+                    "content": "Sorry, I couldn't process your request. Please check your API credentials and try again."
+                }
+                st.session_state.streaming_messages.append(assistant_message)
+                content_container.write(assistant_message["content"])
+
+    # Clear chat button
+    if st.sidebar.button("🧹 Clear Chat History"):
+        st.session_state.streaming_messages = []
+        st.session_state.streaming_thread_id = ""
+        st.rerun()
+
+
+if __name__ == "__main__":
+    main()