yi
721 lines
32 KiB
Python
721 lines
32 KiB
Python
import asyncio
|
|
import sys
|
|
import os
|
|
import json
|
|
import time
|
|
import threading
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import List, Optional, Dict, Any, Callable, Awaitable
|
|
from pydantic import BaseModel, Field
|
|
import duckdb
|
|
import pandas as pd
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Add project root to sys.path to allow importing nanobot
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[3]
|
|
NANOBOT_ROOT = PROJECT_ROOT / "agent-core"
|
|
if str(NANOBOT_ROOT) not in sys.path:
|
|
sys.path.append(str(NANOBOT_ROOT))
|
|
|
|
from app.core.llm_provider import build_llm_provider
|
|
from app.connectors.postgres import postgres_connector
|
|
from app.connectors.clickhouse import clickhouse_connector
|
|
from app.connectors.factory import get_connector
|
|
from app.schemas.chart import ChartGenerationResponse
|
|
from app.agent.chart import generate_chart
|
|
from app.database import SessionLocal
|
|
from app.models.datasource import DataSource
|
|
from app.core.files import resolve_upload_file_path
|
|
from app.services.mdl import MDLService
|
|
from app.services.llm_cache import get_active_llm_config
|
|
from app.trace import trace_service
|
|
|
|
SCHEMA_CACHE_TTL_SECONDS = 300
|
|
CONNECTION_CACHE_TTL_SECONDS = 30
|
|
UPLOAD_CACHE_TTL_SECONDS = 900
|
|
MAX_UPLOAD_CACHE_ITEMS = 8
|
|
NL2SQL_MAX_TOKENS = 900
|
|
NL2SQL_TEMPERATURE = 0.1
|
|
NL2SQL_REASONING_EFFORT = "low"
|
|
NL2SQL_LLM_TIMEOUT_SECONDS = int(os.getenv("NL2SQL_LLM_TIMEOUT_SECONDS", "90"))
|
|
NL2SQL_LLM_REQUEST_TIMEOUT_SECONDS = int(os.getenv("NL2SQL_LLM_REQUEST_TIMEOUT_SECONDS", "45"))
|
|
NL2SQL_LLM_RETRY_COUNT = int(os.getenv("NL2SQL_LLM_RETRY_COUNT", "0"))
|
|
NL2SQL_SQL_EXEC_TIMEOUT_SECONDS = 60
|
|
NL2SQL_CHART_TIMEOUT_SECONDS = int(os.getenv("NL2SQL_CHART_TIMEOUT_SECONDS", "45"))
|
|
|
|
_schema_cache: Dict[str, Dict[str, Any]] = {}
|
|
_connection_cache: Dict[str, Dict[str, Any]] = {}
|
|
_upload_cache: Dict[str, Dict[str, Any]] = {}
|
|
_cache_lock = threading.Lock()
|
|
|
|
class NL2SQLRequest(BaseModel):
|
|
query: str = Field(..., description="User's natural language query")
|
|
source: str = Field(..., description="Data source to query (postgres, clickhouse, upload, ds:{id})")
|
|
file_url: Optional[str] = Field(None, description="Uploaded file URL when source is upload")
|
|
session_id: Optional[str] = Field(None, description="Conversation session identifier")
|
|
generate_chart: bool = Field(False, description="Whether to generate chart specification")
|
|
|
|
class NL2SQLResponse(BaseModel):
|
|
sql: str
|
|
result: List[Dict[str, Any]]
|
|
error: Optional[str] = None
|
|
chart: Optional[ChartGenerationResponse] = None
|
|
|
|
# WrenAI-inspired SQL Rules
|
|
DEFAULT_TEXT_TO_SQL_RULES = """
|
|
### SQL RULES ###
|
|
- ONLY USE SELECT statements, NO DELETE, UPDATE OR INSERT etc. statements that might change the data in the database.
|
|
- ONLY USE the tables and columns mentioned in the database schema.
|
|
- ONLY USE "*" if the user query asks for all the columns of a table.
|
|
- ONLY CHOOSE columns belong to the tables mentioned in the database schema.
|
|
- DON'T INCLUDE comments in the generated SQL query.
|
|
- YOU MUST USE "JOIN" if you choose columns from multiple tables!
|
|
- PREFER USING CTEs over subqueries.
|
|
- When generating SQL query, always:
|
|
- Put double quotes around column and table names.
|
|
- Put single quotes around string literals.
|
|
- Never quote numeric literals.
|
|
For example: SELECT "customers"."customer_name" FROM "customers" WHERE "customers"."city" = 'Taipei' and "customers"."year" = 1992;
|
|
- YOU MUST USE "lower(<table_name>.<column_name>) like lower(<value>)" function or "lower(<table_name>.<column_name>) = lower(<value>)" function for case-insensitive comparison!
|
|
- Use "lower(<table_name>.<column_name>) LIKE lower(<value>)" when:
|
|
- The user requests a pattern or partial match.
|
|
- The value is not specific enough to be a single, exact value.
|
|
- Wildcards (%) are needed to capture the pattern.
|
|
- Use "lower(<table_name>.<column_name>) = lower(<value>)" when:
|
|
- The user requests an exact, specific value.
|
|
- There is no ambiguity or pattern in the value.
|
|
- If the column is date/time related field, and it is a INT/BIGINT/DOUBLE/FLOAT type, please use the appropriate function mentioned in the SQL FUNCTIONS section to cast the column to "TIMESTAMP" type first before using it in the query
|
|
- ALWAYS CAST the date/time related field to "TIMESTAMP WITH TIME ZONE" type when using them in the query
|
|
- If the user asks for a specific date, please give the date range in SQL query
|
|
- Aggregate functions are not allowed in the WHERE clause. Instead, they belong in the HAVING clause, which is used to filter after aggregation.
|
|
- You can only add "ORDER BY" and "LIMIT" to the final "UNION" result.
|
|
- For the ranking problem, you must use the ranking function, `DENSE_RANK()` to rank the results and then use `WHERE` clause to filter the results.
|
|
- For the ranking problem, you must add the ranking column to the final SELECT clause.
|
|
"""
|
|
|
|
TABLE_SELECTOR_SYSTEM_PROMPT = """
|
|
You are a helpful assistant that identifies relevant database tables for a given natural language query.
|
|
|
|
Given the list of available tables and the user's question, return a list of table names that are likely to contain the information needed to answer the question.
|
|
|
|
### FINAL ANSWER FORMAT ###
|
|
The final answer must be a JSON array of strings:
|
|
[
|
|
"table_name1",
|
|
"table_name2"
|
|
]
|
|
"""
|
|
|
|
SQL_GENERATION_SYSTEM_PROMPT = """
|
|
You are a helpful assistant that converts natural language queries into ANSI SQL queries.
|
|
|
|
Given user's question and database schema, generate accurate ANSI SQL directly and concisely.
|
|
|
|
### GENERAL RULES ###
|
|
|
|
1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input.
|
|
2. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions.
|
|
|
|
""" + DEFAULT_TEXT_TO_SQL_RULES + """
|
|
|
|
### FEW-SHOT EXAMPLES ###
|
|
Example 1:
|
|
User's Question: 谁是去年前五个销售额最高的客户？
|
|
Database Schema: {"customers": [{"name": "customer_id", "type": "INT"}, {"name": "customer_name", "type": "TEXT"}], "orders": [{"name": "order_id", "type": "INT"}, {"name": "customer_id", "type": "INT"}, {"name": "amount", "type": "DECIMAL"}, {"name": "order_date", "type": "DATE"}]}
|
|
Final Answer:
|
|
{
|
|
"reasoning": "I need to join customers and orders, filter for last year (2025 if current is 2026), group by customer, sum the amount, and limit to top 5.",
|
|
"sql": "SELECT \\"customers\\".\\"customer_name\\", SUM(\\"orders\\".\\"amount\\") AS \\"total_sales\\" FROM \\"customers\\" JOIN \\"orders\\" ON \\"customers\\".\\"customer_id\\" = \\"orders\\".\\"customer_id\\" WHERE \\"orders\\".\\"order_date\" BETWEEN '2025-01-01' AND '2025-12-31' GROUP BY \\"customers\\".\\"customer_name\\" ORDER BY \\"total_sales\\" DESC LIMIT 5;"
|
|
}
|
|
|
|
### FINAL ANSWER FORMAT ###
|
|
The final answer must be a ANSI SQL query in JSON format:
|
|
|
|
{
|
|
"reasoning": <STEP_BY_STEP_REASONING_PLAN>,
|
|
"sql": <SQL_QUERY_STRING>
|
|
}
|
|
"""
|
|
|
|
def _resolve_upload_file_path(file_url: Optional[str]) -> Path:
|
|
try:
|
|
return resolve_upload_file_path(file_url)
|
|
except ValueError as e:
|
|
raise ValueError(f"Invalid uploaded file URL: {e}")
|
|
|
|
def _load_upload_dataframe_from_path(file_path: Path) -> pd.DataFrame:
|
|
suffix = file_path.suffix.lower()
|
|
if suffix == ".csv":
|
|
return pd.read_csv(file_path)
|
|
if suffix in [".xls", ".xlsx"]:
|
|
return pd.read_excel(file_path)
|
|
if suffix == ".parquet":
|
|
return pd.read_parquet(file_path)
|
|
raise ValueError(f"Unsupported uploaded file type: {suffix}")
|
|
|
|
def _build_upload_schema(df: pd.DataFrame) -> Dict[str, List[Dict[str, str]]]:
|
|
conn = duckdb.connect(":memory:")
|
|
conn.register("uploaded_file", df)
|
|
columns = conn.execute("DESCRIBE uploaded_file").fetchall()
|
|
schema = {"uploaded_file": [{"name": col[0], "type": col[1]} for col in columns]}
|
|
conn.close()
|
|
return schema
|
|
|
|
def _get_upload_payload(file_url: Optional[str]) -> Dict[str, Any]:
|
|
file_path = _resolve_upload_file_path(file_url)
|
|
stat = file_path.stat()
|
|
cache_key = f"{file_path}:{int(stat.st_mtime)}:{stat.st_size}"
|
|
now = time.time()
|
|
with _cache_lock:
|
|
cached = _upload_cache.get(cache_key)
|
|
if cached and now < cached["expires_at"]:
|
|
return {"df": cached["df"], "schema": cached["schema"]}
|
|
df = _load_upload_dataframe_from_path(file_path)
|
|
schema = _build_upload_schema(df)
|
|
with _cache_lock:
|
|
if len(_upload_cache) >= MAX_UPLOAD_CACHE_ITEMS:
|
|
oldest_key = min(_upload_cache.keys(), key=lambda key: _upload_cache[key]["expires_at"])
|
|
_upload_cache.pop(oldest_key, None)
|
|
_upload_cache[cache_key] = {
|
|
"df": df,
|
|
"schema": schema,
|
|
"expires_at": now + UPLOAD_CACHE_TTL_SECONDS,
|
|
}
|
|
return {"df": df, "schema": schema}
|
|
|
|
def _execute_upload_sql(sql_query: str, df: pd.DataFrame) -> List[Dict[str, Any]]:
|
|
conn = duckdb.connect(":memory:")
|
|
conn.register("uploaded_file", df)
|
|
result_df = conn.execute(sql_query).df()
|
|
conn.close()
|
|
return result_df.to_dict(orient="records")
|
|
|
|
def _to_number(value: Any) -> Optional[float]:
|
|
if isinstance(value, bool):
|
|
return None
|
|
if isinstance(value, (int, float)):
|
|
return float(value)
|
|
if isinstance(value, str):
|
|
text = value.strip().replace(",", "")
|
|
if not text:
|
|
return None
|
|
try:
|
|
return float(text)
|
|
except ValueError:
|
|
return None
|
|
return None
|
|
|
|
# _build_fallback_chart removed as per user request to not hardcode fallbacks
|
|
|
|
def _build_schema_cache_key(source: str, connector: Any) -> str:
|
|
# If source is ds:ID, that's already a good key
|
|
if source.startswith("ds:"):
|
|
return source
|
|
|
|
if source == "postgres":
|
|
return f"postgres:{getattr(connector, 'db_url', '')}"
|
|
if source == "clickhouse":
|
|
return (
|
|
f"clickhouse:{getattr(connector, 'host', '')}:{getattr(connector, 'port', '')}:"
|
|
f"{getattr(connector, 'user', '')}:{getattr(connector, 'database', '')}"
|
|
)
|
|
return source
|
|
|
|
def _get_cached_schema(source: str, connector: Any) -> Optional[Dict[str, List[Dict[str, str]]]]:
|
|
key = _build_schema_cache_key(source, connector)
|
|
now = time.time()
|
|
with _cache_lock:
|
|
cached = _schema_cache.get(key)
|
|
if cached and now < cached["expires_at"]:
|
|
return cached["schema"]
|
|
return None
|
|
|
|
def _set_cached_schema(source: str, connector: Any, schema: Dict[str, List[Dict[str, str]]]) -> None:
|
|
key = _build_schema_cache_key(source, connector)
|
|
with _cache_lock:
|
|
_schema_cache[key] = {"schema": schema, "expires_at": time.time() + SCHEMA_CACHE_TTL_SECONDS}
|
|
|
|
async def _check_connection_with_cache(source: str, connector: Any) -> bool:
|
|
cache_key = _build_schema_cache_key(source, connector)
|
|
now = time.time()
|
|
with _cache_lock:
|
|
cached = _connection_cache.get(cache_key)
|
|
if cached and now < cached["expires_at"]:
|
|
return bool(cached["ok"])
|
|
|
|
# Run synchronous test_connection in a separate thread to avoid blocking event loop
|
|
try:
|
|
ok = await asyncio.wait_for(
|
|
asyncio.to_thread(connector.test_connection),
|
|
timeout=15.0
|
|
)
|
|
except asyncio.TimeoutError:
|
|
print("Connection test failed or timed out: Timeout after 15 seconds")
|
|
ok = False
|
|
except Exception as e:
|
|
print(f"Connection test failed or timed out: {e}")
|
|
ok = False
|
|
|
|
with _cache_lock:
|
|
_connection_cache[cache_key] = {"ok": ok, "expires_at": now + CONNECTION_CACHE_TTL_SECONDS}
|
|
return ok
|
|
|
|
async def _select_relevant_tables(
|
|
query: str,
|
|
schema: Dict[str, Any],
|
|
provider: Any,
|
|
on_progress: Callable[[str], Awaitable[None]] | None = None
|
|
) -> List[str]:
|
|
"""Use LLM to select relevant tables from the schema to reduce context size."""
|
|
table_names = list(schema.keys())
|
|
if len(table_names) <= 5:
|
|
return table_names
|
|
|
|
if on_progress:
|
|
await on_progress("正在进行语义表搜索")
|
|
|
|
user_prompt = f"User's Question: {query}\nAvailable Tables: {', '.join(table_names)}"
|
|
messages = [
|
|
{"role": "system", "content": TABLE_SELECTOR_SYSTEM_PROMPT},
|
|
{"role": "user", "content": user_prompt}
|
|
]
|
|
|
|
try:
|
|
response = await asyncio.wait_for(
|
|
provider.chat(
|
|
messages=messages,
|
|
max_tokens=200,
|
|
temperature=0.0,
|
|
),
|
|
timeout=30.0
|
|
)
|
|
content = (response.content or "").strip()
|
|
if "```json" in content:
|
|
content = content.split("```json")[1].split("```")[0]
|
|
elif "```" in content:
|
|
content = content.split("```")[1].split("```")[0]
|
|
|
|
selected = json.loads(content.strip())
|
|
if isinstance(selected, list):
|
|
# Filter valid table names
|
|
valid_selected = [t for t in selected if t in table_names]
|
|
if valid_selected:
|
|
return valid_selected
|
|
except Exception as e:
|
|
logger.warning(f"Table selection failed: {e}")
|
|
|
|
return table_names
|
|
|
|
async def _fetch_sample_data(
|
|
connector: Any,
|
|
table_names: List[str],
|
|
on_progress: Callable[[str], Awaitable[None]] | None = None
|
|
) -> Dict[str, List[Dict[str, Any]]]:
|
|
"""Fetch sample rows for selected tables to help LLM understand data."""
|
|
samples = {}
|
|
if not connector or not hasattr(connector, "execute_query"):
|
|
return samples
|
|
|
|
if on_progress:
|
|
await on_progress(f"正在抓取 {len(table_names)} 張表的樣本數據")
|
|
|
|
for table in table_names:
|
|
try:
|
|
# We use a very small limit
|
|
query = f"SELECT * FROM \"{table}\" LIMIT 3"
|
|
results = await asyncio.wait_for(
|
|
asyncio.to_thread(connector.execute_query, query),
|
|
timeout=10.0
|
|
)
|
|
|
|
rows = []
|
|
if isinstance(results, list):
|
|
if results and isinstance(results[0], dict):
|
|
rows = results
|
|
elif results and isinstance(results[0], (list, tuple)):
|
|
rows = [list(row) for row in results]
|
|
else:
|
|
rows = results
|
|
elif isinstance(results, tuple) and len(results) == 2:
|
|
rows_raw, cols = results
|
|
col_names = [c[0] for c in cols]
|
|
rows = [dict(zip(col_names, row)) for row in rows_raw]
|
|
|
|
if rows:
|
|
samples[table] = rows
|
|
except Exception as e:
|
|
logger.warning(f"Failed to fetch sample for {table}: {e}")
|
|
|
|
return samples
|
|
|
|
async def process_nl2sql(
|
|
request: NL2SQLRequest,
|
|
on_progress: Callable[[str], Awaitable[None]] | None = None,
|
|
) -> NL2SQLResponse:
|
|
async def emit_progress(content: str) -> None:
|
|
if on_progress and content:
|
|
await on_progress(content)
|
|
|
|
total_started = time.perf_counter()
|
|
trace_base_attributes = {
|
|
"component": "nl2sql",
|
|
"source": request.source,
|
|
"session_id": request.session_id,
|
|
"generate_chart": request.generate_chart,
|
|
}
|
|
# 1. Get the connector and schema
|
|
connector = None
|
|
schema = {}
|
|
upload_df: Optional[pd.DataFrame] = None
|
|
|
|
if request.source == "postgres":
|
|
connector = postgres_connector
|
|
elif request.source == "clickhouse":
|
|
connector = clickhouse_connector
|
|
elif request.source == "upload":
|
|
try:
|
|
upload_started = time.perf_counter()
|
|
upload_payload = await asyncio.to_thread(_get_upload_payload, request.file_url)
|
|
upload_df = upload_payload["df"]
|
|
schema = upload_payload["schema"]
|
|
await emit_progress(f"上传文件加载完成 ({time.perf_counter() - upload_started:.2f}s)")
|
|
except Exception as e:
|
|
return NL2SQLResponse(sql="", result=[], error=f"Failed to load uploaded file: {e}")
|
|
elif request.source.startswith("ds:"):
|
|
try:
|
|
ds_started = time.perf_counter()
|
|
ds_id = int(request.source.split(":")[1])
|
|
|
|
def _get_ds_connector():
|
|
db = SessionLocal()
|
|
try:
|
|
ds = db.query(DataSource).filter(DataSource.id == ds_id).first()
|
|
if not ds:
|
|
return None
|
|
return get_connector(ds)
|
|
finally:
|
|
db.close()
|
|
|
|
connector = await asyncio.to_thread(_get_ds_connector)
|
|
if not connector:
|
|
return NL2SQLResponse(sql="", result=[], error=f"Data source not found: {request.source}")
|
|
|
|
await emit_progress(f"数据源配置读取完成 ({time.perf_counter() - ds_started:.2f}s)")
|
|
except ValueError:
|
|
return NL2SQLResponse(sql="", result=[], error=f"Invalid data source ID: {request.source}")
|
|
except Exception as e:
|
|
return NL2SQLResponse(sql="", result=[], error=f"Failed to load data source: {e}")
|
|
else:
|
|
return NL2SQLResponse(sql="", result=[], error=f"Unsupported data source: {request.source}")
|
|
|
|
if connector:
|
|
await emit_progress("正在检测数据源连通性")
|
|
cached_schema = _get_cached_schema(request.source, connector)
|
|
if cached_schema is not None:
|
|
schema = cached_schema
|
|
await emit_progress(f"命中 Schema 缓存，已加载 {len(schema)} 张表")
|
|
else:
|
|
conn_started = time.perf_counter()
|
|
if not await _check_connection_with_cache(request.source, connector):
|
|
return NL2SQLResponse(sql="", result=[], error=f"Failed to connect to {request.source}")
|
|
await emit_progress(f"连接检测完成 ({time.perf_counter() - conn_started:.2f}s)")
|
|
schema_started = time.perf_counter()
|
|
try:
|
|
schema = await asyncio.wait_for(
|
|
asyncio.to_thread(connector.get_schema),
|
|
timeout=120.0
|
|
)
|
|
except asyncio.TimeoutError:
|
|
return NL2SQLResponse(sql="", result=[], error="Failed to fetch schema: Timeout after 120 seconds. Data source might be too large or network is slow.")
|
|
except Exception as e:
|
|
return NL2SQLResponse(sql="", result=[], error=f"Failed to fetch schema: {e}")
|
|
|
|
_set_cached_schema(request.source, connector, schema)
|
|
await emit_progress(f"Schema 拉取完成，共 {len(schema)} 张表 ({time.perf_counter() - schema_started:.2f}s)")
|
|
|
|
# 2. Get the active LLM config
|
|
active_config = get_active_llm_config()
|
|
|
|
if not active_config:
|
|
return NL2SQLResponse(sql="", result=[], error="No active LLM configuration found")
|
|
|
|
# 3. Initialize Provider
|
|
try:
|
|
provider = build_llm_provider(
|
|
model=active_config.get("model"),
|
|
provider=active_config.get("provider"),
|
|
api_key=active_config.get("api_key"),
|
|
api_base=active_config.get("api_base"),
|
|
extra_headers=active_config.get("extra_headers") or {},
|
|
)
|
|
except Exception as e:
|
|
return NL2SQLResponse(sql="", result=[], error=f"Failed to initialize LLM provider: {e}")
|
|
|
|
# 4. Table Selection and Sample Data (Optimization)
|
|
relevant_tables = await _select_relevant_tables(request.query, schema, provider, emit_progress)
|
|
pruned_schema = {t: schema[t] for t in relevant_tables if t in schema}
|
|
|
|
samples = {}
|
|
if request.source != "upload": # For upload, df is already in memory and small
|
|
samples = await _fetch_sample_data(connector, relevant_tables, emit_progress)
|
|
|
|
schema_str = json.dumps(pruned_schema, ensure_ascii=False, separators=(",", ":"))
|
|
samples_str = json.dumps(samples, ensure_ascii=False, separators=(",", ":")) if samples else ""
|
|
|
|
# Try to load MDL context
|
|
mdl_context = ""
|
|
if request.source.startswith("ds:"):
|
|
try:
|
|
ds_id = int(request.source.split(":")[1])
|
|
mdl = await asyncio.to_thread(MDLService.get_mdl, ds_id)
|
|
if mdl:
|
|
mdl_lines = ["\n### SEMANTIC MODEL (WrenMDL) ###"]
|
|
|
|
mdl_lines.append("MODELS:")
|
|
for model in mdl.models:
|
|
# Only include relevant models
|
|
if model.name not in relevant_tables and (model.tableReference and model.tableReference.table not in relevant_tables):
|
|
continue
|
|
|
|
table_ref = model.tableReference.table if model.tableReference else model.name
|
|
desc = f" - Description: {model.properties.get('description', '')}" if model.properties.get('description') else ""
|
|
mdl_lines.append(f"- Model: {model.name} (Table: {table_ref}){desc}")
|
|
|
|
if model.columns:
|
|
mdl_lines.append(" Columns:")
|
|
for col in model.columns:
|
|
col_desc = f" ({col.properties.get('description')})" if col.properties.get('description') else ""
|
|
expr = f" [Calculated: {col.expression}]" if col.isCalculated else ""
|
|
mdl_lines.append(f" - {col.name} ({col.type}){col_desc}{expr}")
|
|
|
|
if mdl.relationships:
|
|
mdl_lines.append("\nRELATIONSHIPS:")
|
|
for rel in mdl.relationships:
|
|
# Only include relevant relationships
|
|
rel_models = rel.models if isinstance(rel.models, list) else []
|
|
if any(m in relevant_tables for m in rel_models):
|
|
mdl_lines.append(f"- {rel.name}: {rel.joinType} between {rel.models} ON {rel.condition}")
|
|
|
|
mdl_context = "\n".join(mdl_lines)
|
|
except Exception as e:
|
|
print(f"Failed to load MDL: {e}")
|
|
|
|
# 5. Construct Prompt
|
|
user_prompt = f"""
|
|
### DATABASE SCHEMA ###
|
|
{schema_str}
|
|
{mdl_context}
|
|
|
|
### SAMPLE DATA ###
|
|
{samples_str}
|
|
|
|
### INPUTS ###
|
|
User's Question: {request.query}
|
|
Language: Chinese (Simplified)
|
|
"""
|
|
|
|
messages = [
|
|
{"role": "system", "content": SQL_GENERATION_SYSTEM_PROMPT},
|
|
{"role": "user", "content": user_prompt}
|
|
]
|
|
|
|
# 5. Call LLM & 6. Execute SQL (with Self-Correction Loop)
|
|
MAX_SQL_EXEC_RETRIES = int(os.getenv("NL2SQL_MAX_EXEC_RETRIES", "2"))
|
|
sql_query = ""
|
|
formatted_results = []
|
|
chart_response = None
|
|
timeout_stage = "llm_generation"
|
|
|
|
for exec_attempt in range(MAX_SQL_EXEC_RETRIES + 1):
|
|
try:
|
|
llm_started = time.perf_counter()
|
|
if exec_attempt == 0:
|
|
await emit_progress("正在生成 SQL")
|
|
else:
|
|
await emit_progress(f"正在尝试修复 SQL ({exec_attempt}/{MAX_SQL_EXEC_RETRIES})")
|
|
|
|
response = None
|
|
last_error = ""
|
|
|
|
for attempt in range(NL2SQL_LLM_RETRY_COUNT + 1):
|
|
try:
|
|
with trace_service.start_span(
|
|
"nl2sql.llm_generation",
|
|
attributes={
|
|
**trace_base_attributes,
|
|
"exec_attempt": exec_attempt,
|
|
"retry_attempt": attempt,
|
|
"model": active_config.get("model"),
|
|
},
|
|
) as llm_span:
|
|
response = await asyncio.wait_for(
|
|
provider.chat(
|
|
messages=messages,
|
|
max_tokens=NL2SQL_MAX_TOKENS,
|
|
temperature=NL2SQL_TEMPERATURE,
|
|
reasoning_effort=NL2SQL_REASONING_EFFORT,
|
|
),
|
|
timeout=NL2SQL_LLM_TIMEOUT_SECONDS,
|
|
)
|
|
llm_span.update(output={"finish_reason": getattr(response, "finish_reason", None)})
|
|
except asyncio.TimeoutError:
|
|
last_error = f"LLM generation timeout after {NL2SQL_LLM_TIMEOUT_SECONDS}s"
|
|
if attempt < NL2SQL_LLM_RETRY_COUNT:
|
|
await emit_progress(f"SQL 生成超时，正在重试 ({attempt + 1}/{NL2SQL_LLM_RETRY_COUNT})")
|
|
continue
|
|
return NL2SQLResponse(sql=sql_query, result=[], error=last_error)
|
|
except Exception as e:
|
|
last_error = f"LLM generation failed: {e}"
|
|
if attempt < NL2SQL_LLM_RETRY_COUNT:
|
|
await emit_progress(f"SQL 生成失败，正在重试 ({attempt + 1}/{NL2SQL_LLM_RETRY_COUNT})")
|
|
continue
|
|
return NL2SQLResponse(sql=sql_query, result=[], error=last_error)
|
|
|
|
if response.finish_reason == "error":
|
|
last_error = response.content or "LLM Error"
|
|
if attempt < NL2SQL_LLM_RETRY_COUNT:
|
|
await emit_progress(f"模型返回错误，正在重试 ({attempt + 1}/{NL2SQL_LLM_RETRY_COUNT})")
|
|
continue
|
|
return NL2SQLResponse(sql=sql_query, result=[], error=last_error)
|
|
break
|
|
|
|
if response is None:
|
|
return NL2SQLResponse(sql=sql_query, result=[], error=last_error or "LLM generation failed")
|
|
|
|
content = (response.content or "").strip()
|
|
if not content:
|
|
return NL2SQLResponse(sql=sql_query, result=[], error="LLM returned empty response")
|
|
|
|
# Clean up code blocks
|
|
if "```json" in content:
|
|
content = content.split("```json")[1].split("```")[0]
|
|
elif "```" in content:
|
|
content = content.split("```")[1].split("```")[0]
|
|
|
|
content = content.strip()
|
|
|
|
try:
|
|
result_json = json.loads(content)
|
|
sql_query = result_json.get("sql", "").strip()
|
|
except json.JSONDecodeError:
|
|
# Fallback if LLM doesn't return valid JSON despite instructions
|
|
sql_query = content
|
|
|
|
logger.info(f"Generated SQL for query '{request.query}':\n{sql_query}")
|
|
|
|
# 格式化单行 SQL 用于在前端进度中展示
|
|
formatted_sql = sql_query.replace('\n', ' ')
|
|
if len(formatted_sql) > 150:
|
|
formatted_sql = formatted_sql[:147] + "..."
|
|
await emit_progress(f"SQL 生成完成: {formatted_sql}")
|
|
|
|
except Exception as e:
|
|
return NL2SQLResponse(sql=sql_query, result=[], error=f"LLM generation failed: {e}")
|
|
|
|
# 6. Execute SQL
|
|
try:
|
|
timeout_stage = "sql_execution"
|
|
sql_exec_started = time.perf_counter()
|
|
await emit_progress("正在执行 SQL 查询")
|
|
with trace_service.start_span(
|
|
"nl2sql.sql_execution",
|
|
attributes={
|
|
**trace_base_attributes,
|
|
"exec_attempt": exec_attempt,
|
|
},
|
|
input_payload={"sql": sql_query},
|
|
) as sql_span:
|
|
if request.source == "upload":
|
|
if upload_df is None:
|
|
upload_payload = await asyncio.to_thread(_get_upload_payload, request.file_url)
|
|
upload_df = upload_payload["df"]
|
|
formatted_results = await asyncio.wait_for(
|
|
asyncio.to_thread(_execute_upload_sql, sql_query, upload_df),
|
|
timeout=NL2SQL_SQL_EXEC_TIMEOUT_SECONDS,
|
|
)
|
|
else:
|
|
results = await asyncio.wait_for(
|
|
asyncio.to_thread(connector.execute_query, sql_query),
|
|
timeout=NL2SQL_SQL_EXEC_TIMEOUT_SECONDS,
|
|
)
|
|
formatted_results = []
|
|
if isinstance(results, list):
|
|
if results and isinstance(results[0], dict):
|
|
formatted_results = results
|
|
elif results and isinstance(results[0], (list, tuple)):
|
|
formatted_results = [list(row) for row in results]
|
|
else:
|
|
formatted_results = results
|
|
elif isinstance(results, tuple) and len(results) == 2:
|
|
rows, cols = results
|
|
col_names = [c[0] for c in cols]
|
|
formatted_results = [dict(zip(col_names, row)) for row in rows]
|
|
else:
|
|
formatted_results = []
|
|
sql_span.set_attributes({"result_rows": len(formatted_results)})
|
|
|
|
await emit_progress(f"SQL 执行完成，返回 {len(formatted_results)} 行 ({time.perf_counter() - sql_exec_started:.2f}s)")
|
|
break # Execution succeeded, break the retry loop
|
|
|
|
except asyncio.TimeoutError:
|
|
return NL2SQLResponse(sql=sql_query, result=[], error=f"SQL execution timeout after {NL2SQL_SQL_EXEC_TIMEOUT_SECONDS}s")
|
|
except Exception as e:
|
|
if exec_attempt < MAX_SQL_EXEC_RETRIES:
|
|
await emit_progress(f"SQL 执行失败，准备自动修复 ({exec_attempt + 1}/{MAX_SQL_EXEC_RETRIES})")
|
|
messages.append({"role": "assistant", "content": f"```json\n{{\"sql\": \"{sql_query}\"}}\n```"})
|
|
messages.append({
|
|
"role": "user",
|
|
"content": f"The generated SQL failed to execute. Database error:\n{str(e)}\n\nPlease fix the SQL query to resolve this error and provide the corrected version following the exact same JSON format."
|
|
})
|
|
continue
|
|
else:
|
|
return NL2SQLResponse(sql=sql_query, result=[], error=f"SQL execution failed after {MAX_SQL_EXEC_RETRIES} retries: {e}")
|
|
|
|
# 7. Generate Chart
|
|
if request.generate_chart and formatted_results:
|
|
try:
|
|
chart_started = time.perf_counter()
|
|
await emit_progress("正在生成可视化方案")
|
|
timeout_stage = "chart_generation"
|
|
with trace_service.start_span(
|
|
"nl2sql.chart_generation",
|
|
attributes=trace_base_attributes,
|
|
input_payload={"query": request.query, "rows": len(formatted_results)},
|
|
) as chart_span:
|
|
chart_response = await asyncio.wait_for(
|
|
generate_chart(formatted_results, request.query),
|
|
timeout=NL2SQL_CHART_TIMEOUT_SECONDS,
|
|
)
|
|
chart_span.set_attributes(
|
|
{
|
|
"chart.can_visualize": bool(getattr(chart_response, "can_visualize", False)),
|
|
"chart.type": getattr(chart_response, "chart_type", ""),
|
|
}
|
|
)
|
|
await emit_progress(f"可视化方案生成完成 ({time.perf_counter() - chart_started:.2f}s)")
|
|
except asyncio.TimeoutError:
|
|
fallback_chart = ChartGenerationResponse(
|
|
reasoning=f"Chart generation timeout after {NL2SQL_CHART_TIMEOUT_SECONDS}s",
|
|
chart_type="",
|
|
can_visualize=False,
|
|
chart_spec=None,
|
|
)
|
|
return NL2SQLResponse(sql=sql_query, result=formatted_results, chart=fallback_chart)
|
|
except Exception as e:
|
|
pass # Ignore chart generation errors, return data only
|
|
|
|
with trace_service.start_span(
|
|
"nl2sql.completed",
|
|
attributes={
|
|
**trace_base_attributes,
|
|
"total_seconds": round(time.perf_counter() - total_started, 4),
|
|
"result_rows": len(formatted_results),
|
|
"has_chart": bool(chart_response),
|
|
},
|
|
):
|
|
pass
|
|
await emit_progress(f"NL2SQL 总耗时 {time.perf_counter() - total_started:.2f}s")
|
|
return NL2SQLResponse(sql=sql_query, result=formatted_results, chart=chart_response)
|