nl2sql-engine 0.1.0__py3-none-any.whl

nl2sql/aggregation/aggregator.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+from typing import Dict, List, Tuple
+
+import polars as pl
+from nl2sql.execution.contracts import ArtifactRef
+from nl2sql.pipeline.nodes.global_planner.schemas import ExecutionDAG, LogicalNode, LogicalEdge
+
+from .engines.polars_duckdb import PolarsDuckdbEngine
+
+
+class AggregationService:
+    def __init__(self, engine: PolarsDuckdbEngine):
+        self.engine = engine
+
+    def execute(
+        self,
+        dag: ExecutionDAG,
+        artifact_refs: Dict[str, ArtifactRef],
+    ) -> Dict[str, List[Dict]]:
+        if not dag:
+            raise ValueError("No ExecutionDAG found for aggregation.")
+
+        if not dag.nodes:
+            return {}
+
+        node_index: Dict[str, LogicalNode] = {n.node_id: n for n in dag.nodes}
+        edges: List[LogicalEdge] = dag.edges
+        incoming_edges: Dict[str, List[LogicalEdge]] = {}
+        outgoing_edges: Dict[str, List[LogicalEdge]] = {}
+        for edge in edges:
+            incoming_edges.setdefault(edge.to_id, []).append(edge)
+            outgoing_edges.setdefault(edge.from_id, []).append(edge)
+
+        computed: Dict[str, pl.DataFrame] = {}
+
+        for layer in dag.layers:
+            for node_id in layer:
+                node = node_index.get(node_id)
+                if not node:
+                    continue
+                if node.kind == "scan":
+                    artifact = artifact_refs.get(node_id)
+                    if not artifact:
+                        raise ValueError(
+                            f"Missing artifact for scan node {node_id}."
+                        )
+                    computed[node_id] = self.engine.load_scan(artifact)
+                    continue
+
+                if node.kind == "combine":
+                    upstream = incoming_edges.get(node_id, [])
+                    inputs = self._ordered_inputs(upstream, computed)
+                    computed[node_id] = self.engine.combine(
+                        node.attributes.get("operation"),
+                        inputs,
+                        node.attributes.get("join_keys", []),
+                    )
+                    continue
+
+                if node.kind.startswith("post_"):
+                    input_ids = node.inputs or []
+                    if len(input_ids) != 1 or input_ids[0] not in computed:
+                        raise ValueError(
+                            f"Post-combine node '{node_id}' expects a single input."
+                        )
+                    computed[node_id] = self.engine.post_op(
+                        node.attributes.get("operation"),
+                        computed[input_ids[0]],
+                        node.attributes,
+                    )
+                    continue
+
+                raise ValueError(
+                    f"Unsupported logical node kind '{node.kind}'."
+                )
+
+        terminal_nodes = sorted([n.node_id for n in dag.nodes if n.node_id not in outgoing_edges])
+        if not terminal_nodes:
+            return {}
+
+        return {
+            node_id: self.engine.to_rows(computed[node_id])
+            for node_id in terminal_nodes
+            if node_id in computed
+        }
+
+    def _ordered_inputs(
+        self,
+        edges: List[LogicalEdge],
+        computed: Dict[str, pl.DataFrame],
+    ) -> List[Tuple[str, pl.DataFrame]]:
+        def role_rank(role: str) -> int:
+            order = {"left": 0, "base": 0, "primary": 0, "right": 1, "compare": 1, "secondary": 1}
+            return order.get(role or "", 2)
+
+        ordered = sorted(edges, key=lambda e: (role_rank(e.role), e.from_id))
+        return [(edge.role or "", computed[edge.from_id]) for edge in ordered if edge.from_id in computed]

nl2sql/aggregation/engines/__init__.py

@@ -0,0 +1,3 @@
+from .polars_duckdb import PolarsDuckdbEngine
+
+__all__ = ["PolarsDuckdbEngine"]

nl2sql/aggregation/engines/polars_duckdb.py

@@ -0,0 +1,125 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List, Tuple
+
+import duckdb
+import polars as pl
+
+from nl2sql.execution.contracts import ArtifactRef
+from nl2sql.execution.artifacts import build_artifact_store
+
+
+class PolarsDuckdbEngine:
+    def __init__(self):
+        self.artifact_store = build_artifact_store()
+
+    def load_scan(self, artifact: ArtifactRef) -> pl.DataFrame:
+        return self.artifact_store.read_parquet(artifact)
+
+    def combine(
+        self,
+        operation: str,
+        inputs: List[Tuple[str, pl.DataFrame]],
+        join_keys: List[Dict[str, Any]],
+    ) -> pl.DataFrame:
+        frames = [frame for _, frame in inputs]
+        if not frames:
+            return pl.DataFrame()
+        if operation == "standalone":
+            return frames[0]
+        if operation == "union":
+            return pl.concat(frames, how="vertical")
+        if operation == "join":
+            if len(frames) < 2:
+                return frames[0]
+            left = frames[0]
+            right = frames[1]
+            left_on = [k.get("left") for k in join_keys]
+            right_on = [k.get("right") for k in join_keys]
+            return left.join(right, left_on=left_on, right_on=right_on, how="inner", suffix="_right")
+        if operation == "compare":
+            if len(frames) < 2:
+                return frames[0]
+            left = frames[0]
+            right = frames[1]
+            left_on = [k.get("left") for k in join_keys]
+            right_on = [k.get("right") for k in join_keys]
+            joined = left.join(right, left_on=left_on, right_on=right_on, how="inner", suffix="_right")
+            diff_cols = []
+            for col in left.columns:
+                if col in left_on:
+                    continue
+                right_col = f"{col}_right"
+                if right_col in joined.columns:
+                    diff_cols.append((col, right_col))
+            if not diff_cols:
+                return joined
+            diff_exprs = [(pl.col(l) != pl.col(r)) for l, r in diff_cols]
+            return joined.filter(pl.any_horizontal(diff_exprs))
+        raise ValueError(f"Unsupported combine operation '{operation}'.")
+
+    def post_op(self, operation: str, frame: pl.DataFrame, attributes: Dict[str, Any]) -> pl.DataFrame:
+        if operation == "filter":
+            rows = frame
+            for flt in attributes.get("filters", []):
+                attr = flt.get("attribute")
+                op = flt.get("operator")
+                val = flt.get("value")
+                col = pl.col(attr)
+                if op == "=":
+                    rows = rows.filter(col == val)
+                elif op == "!=":
+                    rows = rows.filter(col != val)
+                elif op == ">":
+                    rows = rows.filter(col > val)
+                elif op == ">=":
+                    rows = rows.filter(col >= val)
+                elif op == "<":
+                    rows = rows.filter(col < val)
+                elif op == "<=":
+                    rows = rows.filter(col <= val)
+                elif op == "between" and isinstance(val, list) and len(val) == 2:
+                    rows = rows.filter((col >= val[0]) & (col <= val[1]))
+                elif op == "in" and isinstance(val, list):
+                    rows = rows.filter(col.is_in(val))
+                elif op == "contains":
+                    rows = rows.filter(col.cast(pl.Utf8).str.contains(str(val)))
+            return rows
+        if operation == "aggregate":
+            group_by = [g.get("attribute") for g in attributes.get("group_by", []) if g.get("attribute")]
+            metrics = attributes.get("metrics", [])
+            agg_exprs = []
+            for metric in metrics:
+                name = metric.get("name")
+                agg = metric.get("aggregation")
+                col = pl.col(name)
+                if agg == "count":
+                    agg_exprs.append(col.count().alias(name))
+                elif agg == "sum":
+                    agg_exprs.append(col.sum().alias(name))
+                elif agg == "avg":
+                    agg_exprs.append(col.mean().alias(name))
+                elif agg == "min":
+                    agg_exprs.append(col.min().alias(name))
+                elif agg == "max":
+                    agg_exprs.append(col.max().alias(name))
+            if group_by:
+                return frame.groupby(group_by).agg(agg_exprs)
+            return frame.select(agg_exprs)
+        if operation == "project":
+            columns = [c.get("name") for c in attributes.get("expected_schema", []) if c.get("name")]
+            return frame.select(columns) if columns else frame
+        if operation == "sort":
+            order_by = attributes.get("order_by", [])
+            sort_cols = [o.get("attribute") for o in order_by if o.get("attribute")]
+            descending = [o.get("direction") == "desc" for o in order_by if o.get("attribute")]
+            if sort_cols:
+                return frame.sort(sort_cols, descending=descending)
+            return frame
+        if operation == "limit":
+            limit = attributes.get("limit")
+            return frame.head(limit) if limit is not None else frame
+        raise ValueError(f"Unsupported post-combine operation '{operation}'.")
+
+    def to_rows(self, frame: pl.DataFrame) -> List[Dict[str, Any]]:
+        return frame.to_dicts()

nl2sql/api/auth_api.py

@@ -0,0 +1,60 @@
+"""
+Auth API for NL2SQL
+
+Provides functionality for authentication and role-based access control.
+"""
+
+from __future__ import annotations
+
+from typing import List, Optional
+
+from nl2sql.context import NL2SQLContext
+from nl2sql.auth.models import UserContext
+from nl2sql.auth.rbac import RBAC
+
+
+class AuthAPI:
+    """
+    API for authentication and role-based access control.
+    """
+    
+    def __init__(self, ctx: NL2SQLContext):
+        self._ctx = ctx
+        self._rbac: RBAC = ctx.rbac
+    
+    def check_permissions(
+        self,
+        user_context: UserContext,
+        datasource_id: str,
+        table: str
+    ) -> bool:
+        """
+        Check if a user has permission to access a specific resource.
+        
+        Args:
+            user_context: User context with roles
+            datasource_id: ID of the datasource
+            table: Name of the table
+            
+        Returns:
+            True if user has permission, False otherwise
+        """
+        return self._rbac.is_allowed(user_context, datasource_id, table)
+    
+    def get_allowed_resources(
+        self,
+        user_context: UserContext
+    ) -> dict:
+        """
+        Get resources a user has access to.
+        
+        Args:
+            user_context: User context with roles
+            
+        Returns:
+            Dictionary with allowed datasources and tables
+        """
+        return {
+            "datasources": self._rbac.get_allowed_datasources(user_context),
+            "tables": self._rbac.get_allowed_tables(user_context)
+        }

nl2sql/api/benchmark_api.py

@@ -0,0 +1,114 @@
+"""
+Benchmark API for NL2SQL.
+
+Provides public entry points for dataset benchmarking without exposing internal runners.
+"""
+
+from __future__ import annotations
+
+import pathlib
+from dataclasses import dataclass
+from typing import Dict, Optional
+
+import yaml
+
+from nl2sql.common.settings import settings
+from nl2sql.configs import ConfigManager
+from nl2sql.configs.llm import LLMFileConfig, AgentConfig
+from nl2sql.context import NL2SQLContext
+from nl2sql.datasources import DatasourceRegistry
+from nl2sql.evaluation.benchmark_runner import BenchmarkRunner, BenchmarkResult
+from nl2sql.evaluation.types import BenchmarkConfig
+from nl2sql.indexing.vector_store import VectorStore
+from nl2sql.llm import LLMRegistry
+from nl2sql.secrets import SecretManager
+
+
+@dataclass
+class BenchmarkMatrixResult:
+    """Aggregate results for a matrix benchmark run."""
+
+    results_by_config: Dict[str, BenchmarkResult]
+
+
+class BenchmarkAPI:
+    """
+    API for running dataset benchmarks using core evaluation tooling.
+    """
+
+    def __init__(self, ctx: Optional[NL2SQLContext] = None):
+        self._ctx = ctx
+
+    def run_matrix(
+        self,
+        config: BenchmarkConfig,
+        *,
+        progress_callback=None,
+    ) -> BenchmarkMatrixResult:
+        """
+        Run a benchmark suite against one or more LLM configs.
+
+        Args:
+            config: Benchmark configuration (dataset, datasource config, etc.).
+            progress_callback: Optional progress iterator wrapper.
+        """
+        cm = self._ctx.config_manager if self._ctx else ConfigManager()
+
+        secret_manager = SecretManager()
+        secrets_path = config.secrets_path or settings.secrets_config_path
+        if secrets_path and pathlib.Path(secrets_path).exists():
+            secret_configs = cm.load_secrets(pathlib.Path(secrets_path))
+            if secret_configs:
+                secret_manager.configure(secret_configs)
+
+        config_path = pathlib.Path(config.config_path) if config.config_path else pathlib.Path(settings.datasource_config_path)
+        ds_configs = cm.load_datasources(config_path)
+        ds_registry = DatasourceRegistry(secret_manager)
+        ds_registry.register_datasources(ds_configs)
+
+        vector_store = VectorStore(
+            collection_name=settings.vector_store_collection_name,
+            persist_directory=config.vector_store_path or settings.vector_store_path,
+        )
+
+        llm_configs = self._load_llm_configs(config, cm)
+        results: Dict[str, BenchmarkResult] = {}
+
+        for name, llm_cfg in llm_configs.items():
+            llm_registry = LLMRegistry(secret_manager)
+            agents = llm_cfg.agents or {}
+            agents["default"] = llm_cfg.default
+            llm_registry.register_llms(agents)
+
+            runner = BenchmarkRunner(config, ds_registry, vector_store, llm_registry)
+            results[name] = runner.run_dataset(
+                config_name=name,
+                progress_callback=progress_callback,
+            )
+
+        return BenchmarkMatrixResult(results_by_config=results)
+
+    def _load_llm_configs(self, config: BenchmarkConfig, cm: ConfigManager) -> Dict[str, LLMFileConfig]:
+        llm_configs: Dict[str, LLMFileConfig] = {}
+
+        if config.bench_config_path and pathlib.Path(config.bench_config_path).exists():
+            bench_data = yaml.safe_load(pathlib.Path(config.bench_config_path).read_text()) or {}
+            for name, cfg_data in bench_data.items():
+                if isinstance(cfg_data, dict):
+                    llm_configs[name] = LLMFileConfig.model_validate(cfg_data)
+
+        if not llm_configs:
+            if config.llm_config_path and pathlib.Path(config.llm_config_path).exists():
+                llm_configs["default"] = cm.load_llm(pathlib.Path(config.llm_config_path))
+            else:
+                llm_configs["default"] = LLMFileConfig(
+                    default=AgentConfig(provider="openai", model="gpt-4o")
+                )
+
+        if config.stub_llm:
+            for llm_cfg in llm_configs.values():
+                llm_cfg.default.provider = "stub"
+                for agent_cfg in llm_cfg.agents.values():
+                    agent_cfg.provider = "stub"
+
+        return llm_configs

nl2sql/api/datasource_api.py

@@ -0,0 +1,132 @@
+"""
+Datasource API for NL2SQL
+
+Provides functionality for managing datasources programmatically or via config.
+"""
+
+from __future__ import annotations
+
+import pathlib
+from typing import Union, Dict, Any, List
+
+from nl2sql.context import NL2SQLContext
+from nl2sql.datasources.registry import DatasourceRegistry
+from nl2sql.datasources.models import DatasourceConfig
+
+
+class DatasourceAPI:
+    """
+    API for managing datasources programmatically or via config files.
+    """
+    
+    def __init__(self, ctx: NL2SQLContext):
+        self._ctx = ctx
+        self._registry: DatasourceRegistry = ctx.ds_registry
+    
+    def add_datasource(
+        self,
+        config: Union[DatasourceConfig, Dict[str, Any]]
+    ) -> None:
+        """
+        Programmatically add a datasource to the engine.
+        
+        Args:
+            config: Datasource configuration as either a DatasourceConfig object
+                   or a dictionary with the configuration
+        """
+        if isinstance(config, dict):
+            config = DatasourceConfig(**config)
+        
+        self._registry.register_datasource(config)
+    
+    def add_datasource_from_config(
+        self,
+        config_path: Union[str, pathlib.Path]
+    ) -> None:
+        """
+        Add datasources from a configuration file.
+        
+        Args:
+            config_path: Path to the datasource configuration file
+        """
+        from nl2sql.configs import ConfigManager
+        cm = ConfigManager()
+        config_path = pathlib.Path(config_path)
+        ds_configs = cm.load_datasources(config_path)
+        self._registry.register_datasources(ds_configs)
+    
+    def list_datasources(self) -> List[str]:
+        """
+        List all registered datasource IDs.
+        
+        Returns:
+            List of datasource IDs
+        """
+        return self._registry.list_ids()
+    
+    def get_adapter(self, datasource_id: str):
+        """
+        Get the adapter for a specific datasource.
+        
+        Args:
+            datasource_id: ID of the datasource
+            
+        Returns:
+            Datasource adapter object
+        """
+        return self._registry.get_adapter(datasource_id)
+    
+    def get_capabilities(self, datasource_id: str) -> List[str]:
+        """
+        Get the capabilities of a specific datasource.
+        
+        Args:
+            datasource_id: ID of the datasource
+            
+        Returns:
+            List of capability strings
+        """
+        return list(self._registry.get_capabilities(datasource_id))
+    
+
+    def validate_connection(
+        self,
+        ds_id: str,
+    ) -> bool:
+        """
+        Validate a datasource connection configuration without registering it.
+        
+        Args:
+            ds_id: ID of the datasource
+            
+        Returns:
+            True if the connection is valid, raises an exception otherwise
+        """
+        
+        adapter = self._registry.get_adapter(ds_id)
+        return adapter.test_connection()
+    
+    def get_datasource_details(
+        self,
+        datasource_id: str
+    ) -> Dict[str, Any]:
+        """
+        Get detailed information about a specific datasource.
+        
+        Args:
+            datasource_id: ID of the datasource
+            
+        Returns:
+            Dictionary with datasource details
+        """
+        adapter = self._registry.get_adapter(datasource_id)
+        details = {
+            "datasource_id": adapter.datasource_id,
+            "datasource_engine_type": adapter.datasource_engine_type,
+            "connection_args": adapter.connection_args,
+            "statement_timeout_ms": adapter.statement_timeout_ms,
+            "row_limit": adapter.row_limit,
+            "max_bytes": adapter.max_bytes,
+            "capabilities": list(self._registry.get_capabilities(datasource_id)),
+        }
+        return details

nl2sql/api/indexing_api.py

@@ -0,0 +1,59 @@
+"""
+Indexing API for NL2SQL
+
+Provides functionality for indexing schemas for datasources.
+"""
+
+from __future__ import annotations
+
+from typing import Dict, Any
+
+from nl2sql.context import NL2SQLContext
+from nl2sql.indexing.orchestrator import IndexingOrchestrator
+
+
+class IndexingAPI:
+    """
+    API for indexing schemas for datasources.
+    """
+    
+    def __init__(self, ctx: NL2SQLContext):
+        self._ctx = ctx
+        self._orchestrator = IndexingOrchestrator(ctx)
+    
+    def index_datasource(
+        self,
+        datasource_id: str
+    ) -> Dict[str, int]:
+        """
+        Index schema for a specific datasource.
+        
+        Args:
+            datasource_id: ID of the datasource to index
+            
+        Returns:
+            Dictionary with indexing statistics
+        """
+        adapter = self._ctx.ds_registry.get_adapter(datasource_id)
+        return self._orchestrator.index_datasource(adapter)
+    
+    def index_all_datasources(self) -> Dict[str, Dict[str, int]]:
+        """
+        Index schema for all registered datasources.
+        
+        Returns:
+            Dictionary mapping datasource IDs to indexing statistics
+        """
+        results = {}
+        for datasource_id in self._ctx.ds_registry.list_ids():
+            try:
+                results[datasource_id] = self.index_datasource(datasource_id)
+            except Exception as e:
+                results[datasource_id] = {"error": str(e)}
+        return results
+    
+    def clear_index(self) -> None:
+        """
+        Clear the vector store index.
+        """
+        self._orchestrator.clear_store()

nl2sql/api/llm_api.py

@@ -0,0 +1,82 @@
+"""
+LLM API for NL2SQL
+
+Provides functionality for configuring LLMs programmatically or via config.
+"""
+
+from __future__ import annotations
+
+import pathlib
+from typing import Union, Dict, Any
+
+from nl2sql.context import NL2SQLContext
+from nl2sql.llm.registry import LLMRegistry
+from nl2sql.llm.models import AgentConfig
+
+
+class LLM_API:
+    """
+    API for configuring LLMs programmatically or via config files.
+    """
+    
+    def __init__(self, ctx: NL2SQLContext):
+        self._ctx = ctx
+        self._registry = ctx.llm_registry
+    
+    def configure_llm(
+        self,
+        config: Union[AgentConfig, Dict[str, Any]]
+    ) -> None:
+        """
+        Programmatically configure an LLM.
+        
+        Args:
+            config: LLM configuration as either an AgentConfig object
+                   or a dictionary with the configuration
+        """
+        if isinstance(config, dict):
+            if 'name' not in config:
+                config['name'] = 'default'
+            config = AgentConfig(**config)
+        
+        self._registry.register_llm(config)
+    
+    def configure_llm_from_config(
+        self,
+        config_path: Union[str, pathlib.Path]
+    ) -> None:
+        """
+        Configure LLMs from a configuration file.
+        
+        Args:
+            config_path: Path to the LLM configuration file
+        """
+        from nl2sql.configs import ConfigManager
+        cm = ConfigManager()
+        config_path = pathlib.Path(config_path)
+        llm_cfg = cm.load_llm(config_path)
+        
+        agents = llm_cfg.agents or {}
+        agents["default"] = llm_cfg.default
+        self._registry.register_llms(agents)
+    
+    def get_llm(self, name: str):
+        """
+        Get a specific LLM by name.
+        
+        Args:
+            name: Name of the LLM to retrieve
+            
+        Returns:
+            LLM instance
+        """
+        return self._registry.get_llm_config(name)
+    
+    def list_llms(self) -> dict:
+        """
+        List all configured LLMs.
+        
+        Returns:
+            List of LLM names
+        """
+        return self._registry.list_llms()