update

.dockerignore

@@ -0,0 +1 @@
+.env

.env.example

@@ -0,0 +1,2 @@
+WRENAI_CLOUD_API_KEY=
+WRENAI_CLOUD_PROJECT_ID=

.gitignore

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

Dockerfile

@@ -18,7 +18,6 @@ ENV PATH $HOME/.local/bin:$PATH
 WORKDIR $HOME
 RUN mkdir app
 WORKDIR $HOME/app
-COPY . $HOME/app
 
 RUN mkdir -p $HOME/app/.streamlit
 COPY .streamlit/config.toml $HOME/app/.streamlit/

Makefile

@@ -2,4 +2,4 @@ run:
 	poetry run streamlit run src/app.py
 
 deps:
-	poetry export --without-hashes --format=requirements.txt > requirements.txt
+	poetry export --without-hashes --without=dev --format=requirements.txt > requirements.txt

poetry.lock

@@ -764,6 +764,20 @@ files = [
 [package.dependencies]
 six = ">=1.5"
 
+[[package]]
+name = "python-dotenv"
+version = "1.1.1"
+description = "Read key-value pairs from a .env file and set them as environment variables"
+optional = false
+python-versions = ">=3.9"
+files = [
+    {file = "python_dotenv-1.1.1-py3-none-any.whl", hash = "sha256:31f23644fe2602f88ff55e1f5c79ba497e01224ee7737937930c448e4d0e24dc"},
+    {file = "python_dotenv-1.1.1.tar.gz", hash = "sha256:a8a6399716257f45be6a007360200409fce5cda2661e3dec71d23dc15f6189ab"},
+]
+
+[package.extras]
+cli = ["click (>=5.0)"]
+
 [[package]]
 name = "pytz"
 version = "2025.2"
@@ -1125,4 +1139,4 @@ watchmedo = ["PyYAML (>=3.10)"]
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.12,<3.13"
-content-hash = "0ea34f1300b1f1a7efcf97c0a55dd9d2b2d0b989ac9f8f94d808c57e1bac1af0"
+content-hash = "d3018cb9ea02785fe38a30f0f0e1196ad02058106619e7ba3a2a11c4a78753a2"

pyproject.toml

@@ -11,7 +11,10 @@ package-mode = false
 python = ">=3.12,<3.13"
 streamlit = "^1.46.1"
 requests = "^2.32.4"
+watchdog = "^6.0.0"
 
+[tool.poetry.group.dev.dependencies]
+python-dotenv = "^1.1.1"
 
 [build-system]
 requires = ["poetry-core"]

requirements.txt

@@ -35,4 +35,4 @@ tornado==6.5.1 ; python_version >= "3.12" and python_version < "3.13"
 typing-extensions==4.14.0 ; python_version >= "3.12" and python_version < "3.13"
 tzdata==2025.2 ; python_version >= "3.12" and python_version < "3.13"
 urllib3==2.5.0 ; python_version >= "3.12" and python_version < "3.13"
-watchdog==6.0.0 ; python_version >= "3.12" and python_version < "3.13" and platform_system != "Darwin"
+watchdog==6.0.0 ; python_version >= "3.12" and python_version < "3.13"

src/apis.py

@@ -1,11 +1,47 @@
+from typing import Optional
+
 import requests
 
 
+def ask(
+    api_key: str,
+    project_id: str,
+    question: str,
+    thread_id: str = "",
+    language: str = "English",
+    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"
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json"
+    }
+    payload = {
+        "projectId": project_id,
+        "question": question,
+        "sampleSize": sample_size,
+        "language": language,
+    }
+    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 {}, str(e)
+
+
 def generate_sql(
     api_key: str,
     project_id: str,
-    query: str,
+    question: str,
     thread_id: str = "",
+    language: str = "English",
+    return_sql_dialect: bool = False,
 ) -> tuple[dict, str]:
     """Generate SQL from natural language query."""
     base_url = "https://cloud.getwren.ai/api/v1"
@@ -16,7 +52,39 @@ def generate_sql(
     }
     payload = {
         "projectId": project_id,
-        "question": query,
+        "question": question,
+        "language": language,
+        "returnSqlDialect": return_sql_dialect,
+    }
+    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 {}, str(e)
+
+
+def run_sql(
+    api_key: str,
+    project_id: str,
+    sql: str,
+    thread_id: str = "",
+    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"
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json"
+    }
+    payload = {
+        "projectId": project_id,
+        "sql": sql,
+        "limit": limit,
     }
     if thread_id:
         payload["threadId"] = thread_id
@@ -26,7 +94,60 @@ def generate_sql(
         response.raise_for_status()
         return response.json(), ""
     except requests.exceptions.RequestException as e:
-        return {}, e
+        return {}, str(e)
+
+
+def generate_summary(
+    api_key: str,
+    project_id: str,
+    question: str,
+    sql: str,
+    thread_id: str = "",
+    language: str = "English",
+    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"
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json"
+    }
+    payload = {
+        "projectId": project_id,
+        "question": question,
+        "sql": sql,
+        "language": language,
+        "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 {}, str(e)
+
+
+def stream_explanation(
+    api_key: str
+) -> tuple[dict, str]:
+    """Stream explanation for a query."""
+    base_url = "https://cloud.getwren.ai/api/v1"
+    endpoint = f"{base_url}/stream_explanation"
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Accept": "text/event-stream",
+    }
+
+    try:
+        response = requests.get(endpoint, headers=headers, stream=True)
+        response.raise_for_status()
+        return response.json(), ""
+    except requests.exceptions.RequestException as e:
+        return {}, str(e)
 
 
 def generate_chart(
@@ -34,8 +155,8 @@ def generate_chart(
     project_id: str,
     question: str,
     sql: str,
+    thread_id: 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"
@@ -58,4 +179,277 @@ def generate_chart(
         response.raise_for_status()
         return response.json(), ""
     except requests.exceptions.RequestException as e:
-        return {}, e
+        return {}, str(e)
+
+
+def stream_ask(
+    api_key: str,
+    project_id: str,
+    question: str,
+    thread_id: str = "",
+    language: str = "English",
+    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"
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json",
+        "Accept": "text/event-stream",
+    }
+    payload = {
+        "projectId": project_id,
+        "question": question,
+        "language": language,
+        "sampleSize": sample_size,
+    }
+    if thread_id:
+        payload["threadId"] = thread_id
+    
+    try:
+        response = requests.post(endpoint, json=payload, headers=headers, stream=True)
+        response.raise_for_status()
+        return response.json(), ""
+    except requests.exceptions.RequestException as e:
+        return {}, str(e)
+
+
+def stream_generate_sql(
+    api_key: str,
+    project_id: str,
+    question: str,
+    thread_id: str = "",
+    language: str = "English",
+    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"
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json",
+        "Accept": "text/event-stream",
+    }
+    payload = {
+        "projectId": project_id,
+        "question": question,
+        "language": language,
+        "returnSqlDialect": return_sql_dialect,
+    }
+    if thread_id:
+        payload["threadId"] = thread_id
+    
+    try:
+        response = requests.post(endpoint, json=payload, headers=headers, stream=True)
+        response.raise_for_status()
+        return response.json(), ""
+    except requests.exceptions.RequestException as e:
+        return {}, str(e)
+
+
+def get_models(
+    api_key: str,
+    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"
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+    }
+    
+    try:
+        response = requests.get(endpoint, headers=headers)
+        response.raise_for_status()
+        return response.json(), ""
+    except requests.exceptions.RequestException as e:
+        return {}, str(e)
+
+
+def get_instructions(
+    api_key: str,
+    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"
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+    }
+    
+    try:
+        response = requests.get(endpoint, headers=headers)
+        response.raise_for_status()
+        return response.json(), ""
+    except requests.exceptions.RequestException as e:
+        return {}, str(e)
+
+
+def create_instruction(
+    api_key: str,
+    project_id: str,
+    instruction: str,
+    is_global: bool = True,
+    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"
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json"
+    }
+    payload = {
+        "instruction": instruction,
+        "isGlobal": is_global,
+        "questions": questions or [],
+    }
+    
+    try:
+        response = requests.post(endpoint, json=payload, headers=headers)
+        response.raise_for_status()
+        return response.json(), ""
+    except requests.exceptions.RequestException as e:
+        return {}, str(e)
+
+
+def update_instruction(
+    api_key: str,
+    project_id: str,
+    instruction_id: str,
+    instruction: str,
+    is_global: bool = True,
+    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}"
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json"
+    }
+    payload = {
+        "instruction": instruction,
+        "isGlobal": is_global,
+        "questions": questions or [],
+    }
+    
+    try:
+        response = requests.put(endpoint, json=payload, headers=headers)
+        response.raise_for_status()
+        return response.json(), ""
+    except requests.exceptions.RequestException as e:
+        return {}, str(e)
+
+
+def delete_instruction(
+    api_key: str,
+    project_id: str,
+    instruction_id: str,
+) -> tuple[dict, str]:
+    """Delete an instruction."""
+    base_url = "https://cloud.getwren.ai/api/v1"
+    endpoint = f"{base_url}/projects/{project_id}/knowledge/instructions/{instruction_id}"
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+    }
+    
+    try:
+        response = requests.delete(endpoint, headers=headers)
+        response.raise_for_status()
+        return response.json(), ""
+    except requests.exceptions.RequestException as e:
+        return {}, str(e)
+
+
+def get_sql_pairs(
+    api_key: str,
+    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"
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+    }
+    
+    try:
+        response = requests.get(endpoint, headers=headers)
+        response.raise_for_status()
+        return response.json(), ""
+    except requests.exceptions.RequestException as e:
+        return {}, str(e)
+
+
+def create_sql_pair(
+    api_key: str,
+    project_id: str,
+    question: str,
+    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"
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json"
+    }
+    payload = {
+        "question": question,
+        "sql": sql,
+    }
+    
+    try:
+        response = requests.post(endpoint, json=payload, headers=headers)
+        response.raise_for_status()
+        return response.json(), ""
+    except requests.exceptions.RequestException as e:
+        return {}, str(e)
+
+
+def update_sql_pair(
+    api_key: str,
+    project_id: str,
+    sql_pair_id: str,
+    question: str,
+    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}"
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json"
+    }
+    payload = {
+        "question": question,
+        "sql": sql,
+    }
+    
+    try:
+        response = requests.put(endpoint, json=payload, headers=headers)
+        response.raise_for_status()
+        return response.json(), ""
+    except requests.exceptions.RequestException as e:
+        return {}, str(e)
+
+
+def delete_sql_pair(
+    api_key: str,
+    project_id: str,
+    sql_pair_id: str,
+) -> tuple[dict, 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}"
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+    }
+    
+    try:
+        response = requests.delete(endpoint, headers=headers)
+        response.raise_for_status()
+        return response.json(), ""
+    except requests.exceptions.RequestException as e:
+        return {}, str(e)

src/app.py

@@ -112,7 +112,14 @@ def main():
                     
                     # Generate chart
                     with st.spinner("Generating chart..."):
-                        chart_response, error = 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,
+                            thread_id=st.session_state.thread_id,
+                            sample_size=sample_size,
+                        )
                     
                     if chart_response:
                         vega_spec = chart_response.get("vegaSpec", {})