nl2sql-engine 0.1.0__py3-none-any.whl

nl2sql/evaluation/benchmark_runner.py

@@ -0,0 +1,320 @@
+import concurrent.futures
+from dataclasses import dataclass
+from typing import List, Dict, Any
+
+import yaml
+from nl2sql_adapter_sdk.contracts import ResultFrame
+
+from nl2sql.datasources import DatasourceRegistry
+from nl2sql.llm import LLMRegistry
+from nl2sql.indexing.vector_store import VectorStore
+from nl2sql.pipeline.runtime import run_with_graph
+from nl2sql.evaluation.evaluator import ModelEvaluator
+from nl2sql.evaluation.types import BenchmarkConfig
+
+
+@dataclass
+class BenchmarkResult:
+    """Standardized result object from a benchmark run."""
+
+    results: List[Dict[str, Any]]
+    metrics: Dict[str, Any]
+    iterations: int = 1
+
+
+class BenchmarkRunner:
+    """
+    Orchestrates the execution of the Benchmark suite.
+    Decoupled from CLI presentation logic.
+    """
+
+    def __init__(
+        self,
+        config: BenchmarkConfig,
+        datasource_registry: DatasourceRegistry,
+        vector_store: VectorStore,
+        llm_registry: LLMRegistry,
+    ):
+        self.config = config
+        self.ds_registry = datasource_registry
+        self.vector_store = vector_store
+        self.llm_registry = llm_registry
+
+    def run_dataset(self, config_name: str = "default", progress_callback=None) -> BenchmarkResult:
+        """
+        Runs the dataset evaluation.
+        """
+        dataset = self._load_dataset()
+
+        results = []
+        workers = 5
+        iterations = self.config.iterations if self.config.iterations else 1
+
+        with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as executor:
+            futures = []
+            for _ in range(iterations):
+                for item in dataset:
+                    futures.append(executor.submit(self._evaluate_case, item))
+
+            total_tasks = len(dataset) * iterations
+
+            # Use callback for progress bar if provided
+            iterator = concurrent.futures.as_completed(futures)
+            if progress_callback:
+                iterator = progress_callback(
+                    iterator,
+                    total=total_tasks,
+                    description=f"Evaluating ({workers} parallel, {iterations} runs)...",
+                )
+
+            for future in iterator:
+                results.append(future.result())
+
+        # Sort results by ID
+        results.sort(key=lambda x: x["id"])
+
+        # Calculate Metrics
+        metrics = ModelEvaluator.calculate_aggregate_metrics(results, len(results))
+
+        return BenchmarkResult(results=results, metrics=metrics, iterations=iterations)
+
+    def _load_dataset(self) -> List[Dict]:
+        """Loads and filters the dataset."""
+        dataset_path = self.config.dataset_path
+        if not dataset_path.exists():
+            raise FileNotFoundError(f"Dataset file not found: {dataset_path}")
+
+        dataset = yaml.safe_load(dataset_path.read_text())
+        if not isinstance(dataset, list):
+            raise ValueError("Dataset must be a list of test cases.")
+
+        if self.config.include_ids:
+            dataset = [item for item in dataset if item.get("id") in self.config.include_ids]
+            if not dataset:
+                raise ValueError(f"No test cases found matching IDs: {self.config.include_ids}")
+
+        return dataset
+
+    def _evaluate_case(self, item: dict) -> dict:
+        """Evaluates a single test case."""
+        q_id = item.get("id", "unknown")
+        question = item.get("question")
+        expected_sql = item.get("expected_sql")
+        expected_ds = item.get("datasource")
+        expected_layer = item.get("expected_routing_layer")
+
+        try:
+            state = run_with_graph(
+                registry=self.ds_registry,
+                llm_registry=self.llm_registry,
+                user_query=question,
+                datasource_id=None,
+                execute=not self.config.routing_only,
+                vector_store=self.vector_store,
+                vector_store_path=self.config.vector_store_path,
+            )
+        except Exception as e:
+            return {
+                "id": q_id,
+                "question": question,
+                "status": "ERROR",
+                "error": str(e),
+                "routing_match": False,
+                "sql_match": False,
+            }
+
+        actual_ds = state.get("datasource_id") or set()
+        expected_set = set(expected_ds) if expected_ds else set()
+        routing_match = actual_ds == expected_set
+
+        # --- Metrics Extraction ---
+        all_routing_info = state.get("routing_info", {})
+        primary_id = sorted(list(actual_ds))[0] if actual_ds else None
+        routing_info = all_routing_info.get(primary_id) if primary_id else None
+
+        def get_val(obj, key, default=None):
+            if isinstance(obj, dict):
+                return obj.get(key, default)
+            return getattr(obj, key, default)
+
+        if routing_info:
+            routing_layer = get_val(routing_info, "layer", "unknown")
+            routing_reasoning = get_val(routing_info, "reasoning", "")
+            routing_tokens = get_val(routing_info, "tokens", 0)
+            routing_latency = get_val(routing_info, "latency", 0)
+            l1_score = get_val(routing_info, "l1_score", 0.0)
+            candidates = get_val(routing_info, "candidates", [])
+            if candidates and not isinstance(candidates[0], dict):
+                candidates = [{"id": c.id, "score": c.score} for c in candidates]
+        else:
+            routing_layer = "unknown"
+            routing_reasoning = "No routing info"
+            routing_tokens = 0
+            routing_latency = 0
+            l1_score = 0.0
+            candidates = []
+
+        layer_match = routing_layer == expected_layer
+
+        if self.config.routing_only:
+            return {
+                "id": q_id,
+                "question": question,
+                "status": "PASS" if routing_match else "ROUTE_FAIL",
+                "routing_match": routing_match,
+                "sql_match": None,
+                "actual_ds": actual_ds,
+                "expected_ds": expected_ds,
+                "routing_layer": routing_layer,
+                "routing_reasoning": routing_reasoning,
+                "routing_tokens": routing_tokens,
+                "routing_latency": routing_latency,
+                "l1_score": l1_score,
+                "candidates": candidates,
+                "expected_layer": expected_layer,
+                "layer_match": layer_match,
+            }
+
+        generated_sql = None
+        execution_res = None
+        subgraph_outputs = state.get("subgraph_outputs") or {}
+        if subgraph_outputs:
+            first_output = next(iter(subgraph_outputs.values()))
+            generated_sql = (
+                first_output.get("sql_draft")
+                if isinstance(first_output, dict)
+                else getattr(first_output, "sql_draft", None)
+            )
+            execution_res = (
+                first_output.get("artifact")
+                if isinstance(first_output, dict)
+                else getattr(first_output, "artifact", None)
+            )
+
+        if not generated_sql:
+            generated_sql_data = state.get("sql_draft")
+            if isinstance(generated_sql_data, str):
+                generated_sql = generated_sql_data
+            else:
+                generated_sql = (
+                    generated_sql_data.get("sql")
+                    if isinstance(generated_sql_data, dict)
+                    else getattr(generated_sql_data, "sql", None)
+                )
+
+        if execution_res is None:
+            execution_res = state.get("execution")
+        generated_rows = execution_res.get("rows") if isinstance(execution_res, dict) else getattr(execution_res, "rows", [])
+        exec_error = execution_res.get("error") if isinstance(execution_res, dict) else getattr(execution_res, "error", None)
+
+        if exec_error:
+            return {
+                "id": q_id,
+                "question": question,
+                "status": "EXEC_FAIL",
+                "error": exec_error,
+                "routing_match": routing_match,
+                "sql_match": False,
+                "gen_sql": generated_sql,
+            }
+
+        if not generated_sql:
+            return {
+                "id": q_id,
+                "question": question,
+                "status": "NO_SQL",
+                "routing_match": routing_match,
+                "sql_match": False,
+            }
+
+        if not expected_sql:
+            return {
+                "id": q_id,
+                "question": question,
+                "status": "NO_GT",
+                "routing_match": routing_match,
+                "sql_match": None,
+                "semantic_sql_match": None,
+                "gen_sql": generated_sql,
+            }
+
+        if not expected_ds:
+            return {
+                "id": q_id,
+                "status": "BAD_CONFIG",
+                "error": "Dataset missing expected datasource",
+                "routing_match": routing_match,
+                "sql_match": False,
+            }
+
+        try:
+            expected_ds_id = expected_ds[0] if isinstance(expected_ds, list) else expected_ds
+            adapter = self.ds_registry.get_adapter(expected_ds_id)
+            if not hasattr(adapter, "execute_sql"):
+                raise ValueError(f"Datasource '{expected_ds_id}' does not support SQL execution")
+            result: ResultFrame = adapter.execute_sql(expected_sql)
+            expected_rows = result.to_row_dicts()
+        except Exception as e:
+            return {
+                "id": q_id,
+                "question": question,
+                "status": "GT_FAIL",
+                "error": str(e),
+                "routing_match": routing_match,
+                "sql_match": False,
+                "gen_sql": generated_sql,
+            }
+
+        try:
+            data_match = ModelEvaluator.compare_results(generated_rows, expected_rows, order_matters=False)
+
+            try:
+                semantic_sql_match = ModelEvaluator.compare_sql_semantic(generated_sql, expected_sql)
+            except ValueError as ve:
+                err_msg = str(ve)
+                if "Ground Truth" in err_msg:
+                    return {
+                        "id": q_id,
+                        "question": question,
+                        "status": "INVALID_GT",
+                        "error": err_msg,
+                        "routing_match": routing_match,
+                        "sql_match": None if data_match else False,
+                        "semantic_sql_match": None,
+                        "gen_sql": generated_sql,
+                    }
+                return {
+                    "id": q_id,
+                    "question": question,
+                    "status": "INVALID_SQL",
+                    "error": err_msg,
+                    "routing_match": routing_match,
+                    "sql_match": data_match,
+                    "semantic_sql_match": False,
+                    "gen_sql": generated_sql,
+                }
+
+            return {
+                "id": q_id,
+                "question": question,
+                "status": "PASS" if data_match else "DATA_MISMATCH",
+                "routing_match": routing_match,
+                "sql_match": data_match,
+                "semantic_sql_match": semantic_sql_match,
+                "gen_sql": generated_sql,
+                "exp_sql": expected_sql,
+                "gen_rows": len(generated_rows),
+                "exp_rows": len(expected_rows),
+                "expected_layer": expected_layer,
+                "layer_match": layer_match,
+            }
+
+        except Exception as e:
+            return {
+                "id": q_id,
+                "status": "COMPARE_FAIL",
+                "error": str(e),
+                "routing_match": routing_match,
+                "sql_match": False,
+                "gen_sql": generated_sql,
+            }

nl2sql/evaluation/evaluator.py

@@ -0,0 +1,134 @@
+import pandas as pd
+import sqlglot
+from sqlglot import exp
+from typing import List, Dict, Any, Optional
+
+class ModelEvaluator:
+    """Evaluates the correctness of AI-generated SQL and its execution results."""
+
+    @staticmethod
+    def compare_sql_semantic(generated_sql: str, expected_sql: str) -> bool:
+        """Compares two SQL queries semantically by normalizing them to ASTs.
+        
+        Strict Mode: Raises generic exceptions (ValueError) if parsing fails.
+
+        Args:
+            generated_sql (str): The SQL generated by the model.
+            expected_sql (str): The ground truth SQL.
+
+        Returns:
+            bool: True if semantically equivalent, False otherwise.
+
+        Raises:
+            ValueError: If either SQL query is invalid or unparseable.
+        """
+        if not generated_sql or not expected_sql:
+            return False
+            
+        if generated_sql.strip() == expected_sql.strip():
+            return True
+
+        try:
+            gen_ast = sqlglot.parse_one(generated_sql)
+        except Exception as e:
+            raise ValueError(f"Generated SQL is invalid/unparseable: {e}")
+
+        try:
+            exp_ast = sqlglot.parse_one(expected_sql)
+        except Exception as e:
+            raise ValueError(f"Ground Truth SQL is invalid/unparseable: {e}")
+            
+        return gen_ast.sql() == exp_ast.sql()
+
+    @staticmethod
+    def compare_results(
+        generated_rows: List[Dict[str, Any]], 
+        expected_rows: List[Dict[str, Any]], 
+        order_matters: bool = False
+    ) -> bool:
+        """Compares two result sets (lists of dicts).
+
+        Args:
+            generated_rows (List[Dict[str, Any]]): Rows returned by the AI query.
+            expected_rows (List[Dict[str, Any]]): Rows returned by the ground truth query.
+            order_matters (bool): If True, strictly enforces row order.
+
+        Returns:
+            bool: True if the result sets match, False otherwise.
+        """
+        if len(generated_rows) != len(expected_rows):
+            return False
+            
+        if not generated_rows and not expected_rows:
+            return True
+            
+        try:
+            df_gen = pd.DataFrame(generated_rows)
+            df_exp = pd.DataFrame(expected_rows)
+            
+            df_gen.columns = df_gen.columns.str.lower()
+            df_exp.columns = df_exp.columns.str.lower()
+            
+            df_gen = df_gen.reindex(sorted(df_gen.columns), axis=1)
+            df_exp = df_exp.reindex(sorted(df_exp.columns), axis=1)
+            
+            if list(df_gen.columns) != list(df_exp.columns):
+                return False
+            
+            if not order_matters:
+                df_gen = df_gen.sort_values(by=list(df_gen.columns)).reset_index(drop=True)
+                df_exp = df_exp.sort_values(by=list(df_exp.columns)).reset_index(drop=True)
+            
+
+            pd.testing.assert_frame_equal(df_gen, df_exp, check_dtype=False, check_like=True, atol=1e-5)
+            return True
+            
+        except AssertionError:
+            return False
+        except Exception:
+            # Fallback for complex types or other issues
+            return False
+
+    @staticmethod
+    def calculate_aggregate_metrics(results: List[Dict[str, Any]], total_samples: int) -> Dict[str, Any]:
+        """Calculates high-level metrics from a list of result dictionaries.
+
+        Args:
+            results (List[Dict[str, Any]]): List of evaluation result dictionaries.
+            total_samples (int): Total number of samples evaluated.
+
+        Returns:
+            Dict[str, Any]: Dictionary containing aggregate metrics:
+                - routing_accuracy
+                - execution_accuracy 
+                - semantic_sql_accuracy
+                - valid_sql_rate
+                - layer_distribution (count and percentage)
+        """
+        if not total_samples:
+            return {}
+
+        correct_routing = sum(1 for r in results if r.get("routing_match"))
+        correct_sql = sum(1 for r in results if r.get("sql_match"))
+        correct_semantic = sum(1 for r in results if r.get("semantic_sql_match"))
+        valid_sql = sum(1 for r in results if r.get("status") not in ["EXEC_FAIL", "NO_SQL", "ERROR", "ROUTE_FAIL", "BAD_CONFIG", "GT_FAIL", "NO_GT", "INVALID_GT", "INVALID_SQL"])
+        
+        layer_counts = {"layer_1": 0, "layer_2": 0, "layer_3": 0, "fallback": 0}
+        for r in results:
+            layer = r.get("routing_layer", "unknown")
+            if layer in layer_counts:
+                layer_counts[layer] += 1
+        
+
+        total_with_gt = sum(1 for r in results if r.get("status") not in ["NO_GT", "INVALID_GT"])
+        
+        metrics = {
+            "routing_accuracy": (correct_routing / total_samples) * 100,
+            "execution_accuracy": (correct_sql / total_with_gt) * 100 if total_with_gt > 0 else 0.0,
+            "semantic_sql_accuracy": (correct_semantic / total_with_gt) * 100 if total_with_gt > 0 else 0.0,
+            "valid_sql_rate": (valid_sql / total_samples) * 100,
+            "layer_distribution": layer_counts,
+            "layer_percentages": {k: (v / total_samples) * 100 for k, v in layer_counts.items()}
+        }
+        
+        return metrics

nl2sql/evaluation/types.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+import pathlib
+from typing import Optional, List
+
+from pydantic import BaseModel
+
+
+class BenchmarkConfig(BaseModel):
+    """Configuration for running benchmarks."""
+
+    dataset_path: pathlib.Path
+    config_path: Optional[pathlib.Path] = None
+    bench_config_path: Optional[pathlib.Path] = None
+    llm_config_path: Optional[pathlib.Path] = None
+    vector_store_path: Optional[str] = None
+    secrets_path: Optional[pathlib.Path] = None
+    iterations: int = 3
+    routing_only: bool = False
+    include_ids: Optional[List[str]] = None
+    export_path: Optional[pathlib.Path] = None
+    stub_llm: bool = False

nl2sql/execution/__init__.py

@@ -0,0 +1,4 @@
+from .contracts import ArtifactRef, ExecutorResponse, ExecutorRequest
+from .execution_store import ExecutionStore
+
+__all__ = ["ArtifactRef", "ExecutorResponse", "ExecutorRequest", "ExecutionStore"]

nl2sql/execution/artifacts/__init__.py

@@ -0,0 +1,3 @@
+from .store import ArtifactStore, ArtifactStoreConfig, build_artifact_store
+
+__all__ = ["ArtifactStore", "ArtifactStoreConfig", "build_artifact_store"]

nl2sql/execution/artifacts/parquet.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+import polars as pl
+
+from nl2sql_adapter_sdk.contracts import ResultFrame
+
+
+def result_frame_to_polars(frame: ResultFrame) -> pl.DataFrame:
+    rows = frame.to_row_dicts()
+    columns = frame.columns
+    if columns:
+        return pl.DataFrame(rows, schema=columns)
+    return pl.DataFrame(rows)
+
+
+def polars_to_result_frame(df: pl.DataFrame) -> ResultFrame:
+    rows = df.to_dicts()
+    columns = df.columns
+    return ResultFrame.from_row_dicts(rows, columns=columns, row_count=len(rows), success=True)
+
+
+def write_parquet(
+    df: pl.DataFrame,
+    target: Any,
+    storage_options: Optional[Dict[str, Any]] = None,
+) -> None:
+    if storage_options is None:
+        df.write_parquet(target)
+    else:
+        df.write_parquet(target, storage_options=storage_options)
+
+
+def read_parquet(
+    source: Any,
+    storage_options: Optional[Dict[str, Any]] = None,
+) -> pl.DataFrame:
+    if storage_options is None:
+        return pl.read_parquet(source)
+    return pl.read_parquet(source, storage_options=storage_options)

nl2sql/execution/artifacts/store.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+import hashlib
+import json
+import re
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import polars as pl
+
+from nl2sql.common.settings import settings
+from nl2sql.execution.contracts import ArtifactRef
+from nl2sql_adapter_sdk.contracts import ResultFrame
+
+from .parquet import polars_to_result_frame, read_parquet, result_frame_to_polars, write_parquet
+
+_PLACEHOLDER = re.compile(r"<([a-zA-Z0-9_]+)>")
+
+
+@dataclass(frozen=True)
+class ArtifactStoreConfig:
+    backend: str
+    base_uri: str
+    path_template: str
+    s3_bucket: Optional[str] = None
+    s3_prefix: Optional[str] = None
+    adls_account: Optional[str] = None
+    adls_container: Optional[str] = None
+    adls_connection_string: Optional[str] = None
+
+
+class ArtifactStore:
+    """Writes and reads Parquet result artifacts for every supported backend.
+
+    Backends differ only in how the target URI is built; polars handles ``s3://``
+    and ``abfs://`` natively, so the read/write path itself is shared.
+    """
+
+    def __init__(self, config: ArtifactStoreConfig) -> None:
+        self.config = config
+
+    def create_artifact_ref(self, frame: ResultFrame, metadata: Dict[str, str]) -> ArtifactRef:
+        uri = self._build_uri(metadata)
+        df = result_frame_to_polars(frame)
+        write_parquet(df, uri, storage_options=self._storage_options())
+
+        payload = {
+            "columns": frame.columns,
+            "row_count": frame.row_count,
+            "path": uri,
+        }
+        content_hash = hashlib.sha256(
+            json.dumps(payload, sort_keys=True, separators=(",", ":"), ensure_ascii=True).encode("utf-8")
+        ).hexdigest()
+
+        return self._build_artifact_ref(
+            uri=uri,
+            frame=frame,
+            content_hash=content_hash,
+            bytes_written=self._bytes_written(uri, df),
+            schema_version=metadata.get("schema_version"),
+        )
+
+    def read_result_frame(self, artifact: ArtifactRef) -> ResultFrame:
+        return polars_to_result_frame(self.read_parquet(artifact))
+
+    def read_parquet(self, artifact: ArtifactRef) -> pl.DataFrame:
+        return read_parquet(artifact.uri, storage_options=self._storage_options())
+
+    def _render_path(self, metadata: Dict[str, str]) -> str:
+        def substitute(match: re.Match) -> str:
+            key = match.group(1)
+            value = metadata.get(key)
+            if value in (None, ""):
+                raise ValueError(
+                    f"Cannot render artifact path template '{self.config.path_template}': "
+                    f"no value for placeholder '<{key}>'. Provide '{key}' in the artifact metadata "
+                    f"or remove it from RESULT_ARTIFACT_PATH_TEMPLATE."
+                )
+            return str(value)
+
+        return _PLACEHOLDER.sub(substitute, self.config.path_template)
+
+    def _build_uri(self, metadata: Dict[str, str]) -> str:
+        backend = self.config.backend
+        relative_path = self._render_path(metadata)
+
+        if backend == "local":
+            target = Path(self.config.base_uri) / relative_path
+            target.parent.mkdir(parents=True, exist_ok=True)
+            return str(target.resolve())
+
+        if backend == "s3":
+            if not self.config.s3_bucket:
+                raise ValueError(
+                    "S3 artifact backend requires a bucket; set RESULT_ARTIFACT_S3_BUCKET."
+                )
+            prefix = (self.config.s3_prefix or "").strip("/")
+            key = f"{prefix}/{relative_path}" if prefix else relative_path
+            return f"s3://{self.config.s3_bucket}/{key}"
+
+        if backend == "adls":
+            if not self.config.adls_container:
+                raise ValueError(
+                    "ADLS artifact backend requires a container; set RESULT_ARTIFACT_ADLS_CONTAINER."
+                )
+            if not self.config.adls_account:
+                raise ValueError(
+                    "ADLS artifact backend requires a storage account; set RESULT_ARTIFACT_ADLS_ACCOUNT."
+                )
+            host = f"{self.config.adls_account}.dfs.core.windows.net"
+            return f"abfs://{self.config.adls_container}@{host}/{relative_path}"
+
+        raise ValueError(
+            f"Unsupported artifact backend '{backend}'. Expected one of: local, s3, adls."
+        )
+
+    def _storage_options(self) -> Optional[Dict[str, Any]]:
+        if self.config.backend == "adls" and self.config.adls_connection_string:
+            return {"connection_string": self.config.adls_connection_string}
+        return None
+
+    def _bytes_written(self, uri: str, df: pl.DataFrame) -> int:
+        if self.config.backend == "local":
+            path = Path(uri)
+            return path.stat().st_size if path.exists() else 0
+        return df.estimated_size()
+
+    def _build_artifact_ref(
+        self,
+        uri: str,
+        frame: ResultFrame,
+        content_hash: str,
+        bytes_written: int,
+        schema_version: Optional[str],
+    ) -> ArtifactRef:
+        return ArtifactRef(
+            uri=uri,
+            backend=self.config.backend,
+            format="parquet",
+            row_count=frame.row_count or len(frame.rows),
+            columns=frame.columns,
+            bytes=bytes_written,
+            content_hash=content_hash,
+            created_at=datetime.utcnow(),
+            schema_version=schema_version,
+            path_template=self.config.path_template,
+        )
+
+
+def build_artifact_store() -> ArtifactStore:
+    return ArtifactStore(
+        ArtifactStoreConfig(
+            backend=settings.result_artifact_backend,
+            base_uri=settings.result_artifact_base_uri,
+            path_template=settings.result_artifact_path_template,
+            s3_bucket=settings.result_artifact_s3_bucket,
+            s3_prefix=settings.result_artifact_s3_prefix,
+            adls_account=settings.result_artifact_adls_account,
+            adls_container=settings.result_artifact_adls_container,
+            adls_connection_string=settings.result_artifact_adls_connection_string,
+        )
+    )