src/apis.py

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