| """Service orchestration for text-to-SQL query processing with Consensus Voting and Self-Correction."""
|
|
|
| import time
|
| import uuid
|
| from typing import Any, Dict, List
|
|
|
| from src.t2s.data.prompts import build_prompt_messages
|
| from src.t2s.gen.backends import get_backend
|
| from src.t2s.safety.executor import SafeSQLExecutor
|
| from src.t2s.safety.policy import SafetyPolicy
|
| from src.t2s.safety.validator import SQLValidator
|
| from src.t2s.service.self_correct import SelfCorrectionRunner
|
| from src.t2s.service.voting import ConsensusVoter
|
| from src.t2s.serving.dbs import DatabaseRegistry
|
|
|
|
|
| class Text2SQLService:
|
| def __init__(
|
| self,
|
| backend_name: str = "mock",
|
| model_path: str = "models/qlora_adapter",
|
| db_base_dir: str = "data/raw/spider/database",
|
| ) -> None:
|
| self.db_registry = DatabaseRegistry(db_base_dir=db_base_dir)
|
| self.policy = SafetyPolicy.load()
|
| self.validator = SQLValidator(self.policy)
|
| self.executor = SafeSQLExecutor(
|
| timeout_seconds=self.policy.timeout_seconds,
|
| max_rows=self.policy.max_row_limit,
|
| )
|
| self.self_corrector = SelfCorrectionRunner(
|
| validator=self.validator,
|
| executor=self.executor,
|
| max_retries=2,
|
| )
|
| self.consensus_voter = ConsensusVoter(
|
| validator=self.validator,
|
| executor=self.executor,
|
| per_candidate_timeout_seconds=self.policy.timeout_seconds,
|
| total_budget_seconds=15.0,
|
| max_concurrent=4,
|
| )
|
| try:
|
| self.backend = get_backend(backend_name, model_name=model_path, model_path=model_path)
|
| except Exception:
|
| self.backend = get_backend("mock")
|
|
|
| def process_query(
|
| self,
|
| db_id: str,
|
| question: str,
|
| include_sample_rows: bool = False,
|
| enable_consensus: bool = False,
|
| n_samples: int = 5,
|
| ) -> Dict[str, Any]:
|
| if enable_consensus and n_samples > 1:
|
| return self.process_query_consensus(
|
| db_id=db_id,
|
| question=question,
|
| include_sample_rows=include_sample_rows,
|
| n_samples=n_samples,
|
| )
|
|
|
| request_id = str(uuid.uuid4())
|
| start_time = time.perf_counter()
|
|
|
| timings: Dict[str, float] = {}
|
|
|
|
|
| t0 = time.perf_counter()
|
| schema_ddl = self.db_registry.get_schema_ddl(db_id, include_sample_rows=include_sample_rows)
|
| timings["schema_fetch_ms"] = round((time.perf_counter() - t0) * 1000.0, 2)
|
|
|
|
|
| prompt_msgs = build_prompt_messages(schema_ddl, question)
|
|
|
|
|
| t1 = time.perf_counter()
|
| gen_output = self.backend.generate([prompt_msgs])[0]
|
| timings["inference_ms"] = round((time.perf_counter() - t1) * 1000.0, 2)
|
|
|
| raw_sql = gen_output.get("text", "") if isinstance(gen_output, dict) else str(gen_output)
|
| clean_sql = raw_sql.strip()
|
| if ";" in clean_sql:
|
| clean_sql = clean_sql.split(";")[0] + ";"
|
|
|
| db_path = str(self.db_registry.get_db_path(db_id))
|
|
|
| def generate_repair(msgs: List[Dict[str, str]]) -> str:
|
| res = self.backend.generate([msgs])[0]
|
| return res.get("text", "") if isinstance(res, dict) else str(res)
|
|
|
| t_corr = time.perf_counter()
|
| corr_res = self.self_corrector.run(
|
| db_path=db_path,
|
| initial_sql=clean_sql,
|
| generate_fn=generate_repair,
|
| base_prompt_messages=prompt_msgs,
|
| )
|
| timings["self_correction_ms"] = round((time.perf_counter() - t_corr) * 1000.0, 2)
|
| timings["total_ms"] = round((time.perf_counter() - start_time) * 1000.0, 2)
|
|
|
| client_attempts = [
|
| {"attempt_index": att.attempt_index, "sql": att.sql, "status": att.status}
|
| for att in corr_res.attempts
|
| ]
|
|
|
| if not corr_res.allowed:
|
| return {
|
| "request_id": request_id,
|
| "db_id": db_id,
|
| "question": question,
|
| "sql": corr_res.final_sql,
|
| "allowed": False,
|
| "reasons": corr_res.verdict.reasons,
|
| "attempts": client_attempts,
|
| "columns": [],
|
| "rows": [],
|
| "row_count": 0,
|
| "truncated": False,
|
| "retried": corr_res.retried,
|
| "consensus_score": "1/1",
|
| "timings_ms": timings,
|
| "model_version": "Qwen2.5-Coder-3B-QLoRA",
|
| }
|
|
|
| exec_res = corr_res.exec_res
|
| columns = exec_res.columns if exec_res else []
|
| rows = [list(r) for r in exec_res.rows] if exec_res else []
|
| row_count = exec_res.row_count if exec_res else 0
|
| truncated = exec_res.truncated if exec_res else False
|
|
|
| return {
|
| "request_id": request_id,
|
| "db_id": db_id,
|
| "question": question,
|
| "sql": corr_res.final_sql,
|
| "allowed": True,
|
| "reasons": corr_res.verdict.reasons,
|
| "attempts": client_attempts,
|
| "columns": columns,
|
| "rows": rows,
|
| "row_count": row_count,
|
| "truncated": truncated,
|
| "retried": corr_res.retried,
|
| "consensus_score": "1/1",
|
| "timings_ms": timings,
|
| "model_version": "Qwen2.5-Coder-3B-QLoRA",
|
| }
|
|
|
| def process_query_consensus(
|
| self,
|
| db_id: str,
|
| question: str,
|
| include_sample_rows: bool = False,
|
| n_samples: int = 5,
|
| ) -> Dict[str, Any]:
|
| """Execution-Guided Self-Correction with Majority Voting (Consensus Decoding)."""
|
| request_id = str(uuid.uuid4())
|
| start_time = time.perf_counter()
|
| timings: Dict[str, float] = {}
|
|
|
| t0 = time.perf_counter()
|
| schema_ddl = self.db_registry.get_schema_ddl(db_id, include_sample_rows=include_sample_rows)
|
| timings["schema_fetch_ms"] = round((time.perf_counter() - t0) * 1000.0, 2)
|
|
|
| prompt_msgs = build_prompt_messages(schema_ddl, question)
|
| prompt_list = [prompt_msgs] * n_samples
|
|
|
| t1 = time.perf_counter()
|
| gen_outputs = self.backend.generate(prompt_list)
|
| timings["inference_ms"] = round((time.perf_counter() - t1) * 1000.0, 2)
|
|
|
| db_path = str(self.db_registry.get_db_path(db_id))
|
|
|
| candidate_sqls = []
|
| for out in gen_outputs:
|
| raw_sql = out.get("text", "") if isinstance(out, dict) else str(out)
|
| candidate_sqls.append(raw_sql)
|
|
|
| t_vote = time.perf_counter()
|
| vote_res = self.consensus_voter.select_best_candidate(
|
| db_path=db_path,
|
| candidate_sqls=candidate_sqls,
|
| )
|
| timings["voting_execution_ms"] = round((time.perf_counter() - t_vote) * 1000.0, 2)
|
| timings["total_ms"] = round((time.perf_counter() - start_time) * 1000.0, 2)
|
|
|
| client_attempts = []
|
| for c in vote_res.candidates:
|
| err = c.exec_res.error if c.exec_res and not c.exec_res.success else None
|
| status = (
|
| "blocked_by_policy"
|
| if not c.verdict.allowed
|
| else (
|
| "ok"
|
| if c.exec_res and c.exec_res.success
|
| else f"exec_error:{c.exec_res.error[:30] if c.exec_res else 'failed'}"
|
| )
|
| )
|
| from src.t2s.service.self_correct import sanitize_attempt_status
|
|
|
| status = sanitize_attempt_status(
|
| c.verdict.allowed, bool(c.exec_res and c.exec_res.success), err
|
| )
|
| client_attempts.append({"attempt_index": c.index, "sql": c.sql, "status": status})
|
|
|
| if not vote_res.winner_verdict.allowed:
|
| return {
|
| "request_id": request_id,
|
| "db_id": db_id,
|
| "question": question,
|
| "sql": vote_res.winner_sql,
|
| "allowed": False,
|
| "reasons": vote_res.winner_verdict.reasons,
|
| "attempts": client_attempts,
|
| "columns": [],
|
| "rows": [],
|
| "row_count": 0,
|
| "truncated": False,
|
| "retried": False,
|
| "consensus_score": vote_res.consensus_score,
|
| "timings_ms": timings,
|
| "model_version": "Qwen2.5-Coder-3B-QLoRA",
|
| }
|
|
|
| exec_res = vote_res.winner_exec_res
|
| columns = exec_res.columns if exec_res else []
|
| rows = [list(r) for r in exec_res.rows] if exec_res else []
|
| row_count = exec_res.row_count if exec_res else 0
|
| truncated = exec_res.truncated if exec_res else False
|
|
|
| return {
|
| "request_id": request_id,
|
| "db_id": db_id,
|
| "question": question,
|
| "sql": vote_res.winner_sql,
|
| "allowed": True,
|
| "reasons": vote_res.winner_verdict.reasons,
|
| "attempts": client_attempts,
|
| "columns": columns,
|
| "rows": rows,
|
| "row_count": row_count,
|
| "truncated": truncated,
|
| "retried": False,
|
| "consensus_score": vote_res.consensus_score,
|
| "timings_ms": timings,
|
| "model_version": "Qwen2.5-Coder-3B-QLoRA",
|
| }
|
|
|
