nl2sql-engine 0.1.0__py3-none-any.whl

nl2sql/execution/contracts.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Dict, List, Optional, Literal, Any
+
+from pydantic import BaseModel, Field, ConfigDict
+
+from nl2sql.auth.models import UserContext
+from nl2sql.common.errors import PipelineError
+
+
+ArtifactBackend = Literal["local", "s3", "adls"]
+ArtifactFormat = Literal["parquet"]
+
+
+class ArtifactRef(BaseModel):
+    uri: str
+    backend: ArtifactBackend
+    format: ArtifactFormat
+    row_count: int
+    columns: List[str]
+    bytes: int
+    content_hash: str
+    created_at: datetime
+    schema_version: Optional[str] = None
+    path_template: str
+
+    model_config = ConfigDict(extra="ignore")
+
+
+class ExecutorRequest(BaseModel):
+    node_id: str
+    trace_id: str
+    subgraph_name: str
+    datasource_id: Optional[str] = None
+    schema_version: Optional[str] = None
+    sql: Optional[str] = None
+    user_context: Optional[UserContext] = None
+    tenant_id: str
+
+    model_config = ConfigDict(extra="ignore")
+
+
+class ExecutorResponse(BaseModel):
+    executor_name: str
+    subgraph_name: str
+    node_id: str
+    trace_id: str
+    datasource_id: Optional[str] = None
+    schema_version: Optional[str] = None
+    artifact: Optional[ArtifactRef] = None
+    metrics: Dict[str, float] = Field(default_factory=dict)
+    errors: List[PipelineError] = Field(default_factory=list)
+    reasoning: List[Dict[str, Any]] = Field(default_factory=list)
+    tenant_id: str
+
+    model_config = ConfigDict(extra="ignore")

nl2sql/execution/execution_store.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from threading import Lock
+from typing import Dict, Optional
+
+from .contracts import ArtifactRef
+
+
+class ExecutionStore:
+    """In-memory store for artifact references keyed by node id."""
+
+    def __init__(self) -> None:
+        self._artifacts: Dict[str, ArtifactRef] = {}
+        self._lock = Lock()
+
+    def put(self, node_id: str, artifact: ArtifactRef) -> None:
+        with self._lock:
+            self._artifacts[node_id] = artifact
+
+    def get(self, node_id: str) -> Optional[ArtifactRef]:
+        return self._artifacts.get(node_id)
+
+    def snapshot(self) -> Dict[str, ArtifactRef]:
+        with self._lock:
+            return dict(self._artifacts)

nl2sql/execution/executor/__init__.py

@@ -0,0 +1,3 @@
+from .sql_executor import SqlExecutorService
+
+__all__ = ["SqlExecutorService"]

nl2sql/execution/executor/sql_executor.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+from typing import Optional, Dict, Any
+
+from nl2sql.common.errors import PipelineError, ErrorSeverity, ErrorCode
+from nl2sql.common.logger import get_logger
+from nl2sql.execution.contracts import ExecutorRequest, ExecutorResponse
+from nl2sql.execution.artifacts import build_artifact_store
+from nl2sql_adapter_sdk.capabilities import DatasourceCapability
+from nl2sql.datasources import DatasourceRegistry
+from nl2sql_adapter_sdk.contracts import AdapterRequest
+
+logger = get_logger("sql_executor")
+
+
+class SqlExecutorService:
+    def __init__(self, ds_registry: DatasourceRegistry):
+        self.ds_registry = ds_registry
+        self.artifact_store = build_artifact_store()
+
+    def validate_request(self, request: ExecutorRequest) -> list[PipelineError]:
+        errors = []
+        if not request.sql:
+            errors.append(
+                PipelineError(
+                    node="sql_executor",
+                    message="No SQL to execute.",
+                    severity=ErrorSeverity.ERROR,
+                    error_code=ErrorCode.MISSING_SQL,
+                )
+            )
+
+        if not request.datasource_id:
+            errors.append(
+                PipelineError(
+                    node="sql_executor",
+                    message="No datasource_id provided.",
+                    severity=ErrorSeverity.ERROR,
+                    error_code=ErrorCode.MISSING_DATASOURCE_ID
+                )
+            )
+
+        ds_id = request.datasource_id
+        caps = self.ds_registry.get_capabilities(ds_id)
+        if DatasourceCapability.SUPPORTS_SQL.value not in caps:
+            errors.append(
+                PipelineError(
+                    node="sql_executor",
+                    message=f"Datasource '{ds_id}' does not support SQL execution.",
+                    severity=ErrorSeverity.ERROR,
+                    error_code=ErrorCode.INVALID_STATE,
+                )
+            )
+        return errors
+
+    def execute(self, request: ExecutorRequest) -> ExecutorResponse:
+        errors = self.validate_request(request)
+        if errors:
+            return ExecutorResponse(
+                executor_name="sql_executor",
+                subgraph_name=request.subgraph_name,
+                node_id=request.node_id,
+                trace_id=request.trace_id,
+                datasource_id=request.datasource_id,
+                schema_version=request.schema_version,
+                errors=errors,
+                tenant_id=request.tenant_id,
+            )
+        
+        ds_id = request.datasource_id
+        adapter = self.ds_registry.get_adapter(ds_id)
+
+        adapter_request = AdapterRequest(
+            plan_type="sql",
+            payload={"sql": request.sql},
+            limits={}
+        )
+
+        result_frame = adapter.execute(adapter_request)
+
+        if not result_frame.success:
+            error_msg = result_frame.error.safe_message if result_frame.error else "SQL execution failed."
+            error = PipelineError(
+                node="sql_executor",
+                message=error_msg,
+                severity=ErrorSeverity.ERROR,
+                error_code=ErrorCode.EXECUTION_FAILED,
+            )
+            return ExecutorResponse(
+                executor_name="sql_executor",
+                subgraph_name=request.subgraph_name,
+                node_id=request.node_id,
+                trace_id=request.trace_id,
+                datasource_id=request.datasource_id,
+                schema_version=request.schema_version,
+                errors=[error],
+                tenant_id=request.tenant_id,
+            )
+    
+        artifact_ref = self.artifact_store.create_artifact_ref(
+            result_frame, {"schema_version": request.schema_version, "request_id": request.trace_id, "tenant_id": request.tenant_id} )
+        
+        return ExecutorResponse(
+            executor_name="sql_executor",
+            subgraph_name=request.subgraph_name,
+            node_id=request.node_id,
+            trace_id=request.trace_id,
+            datasource_id=request.datasource_id,
+            schema_version=request.schema_version,
+            artifact=artifact_ref,
+            metrics={
+                "row_count": result_frame.row_count,
+                "bytes_returned": result_frame.bytes or 0,
+            },
+            tenant_id=request.tenant_id,
+        )

nl2sql/indexing/__init__.py

@@ -0,0 +1,7 @@
+from .orchestrator import IndexingOrchestrator
+from .vector_store import VectorStore
+
+__all__ = [
+    "IndexingOrchestrator",
+    "VectorStore"
+]

nl2sql/indexing/chunk_builder.py

@@ -0,0 +1,227 @@
+from typing import List
+
+from .models import (
+    BaseChunk,
+    DatasourceChunk,
+    TableChunk,
+    ColumnChunk,
+    RelationshipChunk,
+    MetricChunk,
+)
+from nl2sql.schema import SchemaSnapshot
+from nl2sql_adapter_sdk.schema import TableRef, ColumnRef
+
+
+class SchemaChunkBuilder:
+    """
+    Builds schema chunks from a SchemaSnapshot.
+
+    This class performs a pure transformation from SchemaSnapshot
+    to a list of schema chunks used for retrieval and grounding.
+
+    No inference, optimization, or database access is performed.
+    """
+
+    def __init__(
+        self,
+        ds_id: str,
+        schema_snapshot: SchemaSnapshot,
+        schema_version: str,
+        questions: List[str],
+    ):
+        """
+        Initializes the SchemaChunkBuilder.
+
+        Args:
+            ds_id: Datasource identifier.
+            schema_snapshot: Snapshot containing schema contract and metadata.
+            schema_version: Version identifier for the schema snapshot.
+            questions: Example user questions for grounding the datasource.
+        """
+        self.ds_id = ds_id
+        self.schema_snapshot = schema_snapshot
+        self.schema_version = schema_version
+        self.questions = questions
+
+    def build(self) -> List[BaseChunk]:
+        """
+        Builds all schema chunks.
+
+        Returns:
+            List of schema chunks derived from the schema snapshot.
+        """
+        chunks: List[BaseChunk] = []
+
+        chunks.extend(self._build_datasource_chunks())
+        chunks.extend(self._build_table_chunks())
+        chunks.extend(self._build_column_chunks())
+        chunks.extend(self._build_relationship_chunks())
+        chunks.extend(self._build_metric_chunks())
+
+        return chunks
+
+    def _build_datasource_chunks(self) -> List[DatasourceChunk]:
+        """
+        Builds datasource-level chunks.
+
+        Returns:
+            List containing a single DatasourceChunk.
+        """
+        md = self.schema_snapshot.metadata
+
+        return [
+            DatasourceChunk(
+                id=f"schema.datasource:{md.datasource_id}:{self.schema_version}",
+                datasource_id=md.datasource_id,
+                description=md.description or "",
+                domains=md.domains,
+                schema_version=self.schema_version,
+                examples=self.questions,
+            )
+        ]
+
+    def _build_table_chunks(self) -> List[TableChunk]:
+        """
+        Builds table-level chunks.
+
+        Returns:
+            List of TableChunk objects.
+        """
+        chunks: List[TableChunk] = []
+
+        contract = self.schema_snapshot.contract
+        metadata = self.schema_snapshot.metadata
+
+        for table_key, table_contract in contract.tables.items():
+            table_md = metadata.tables.get(table_key)
+            table_ref = TableRef(
+                schema_name=table_contract.table.schema_name,
+                table_name=table_contract.table.table_name,
+            )
+
+            primary_keys = [
+                c.name for c in table_contract.columns.values() if c.is_primary_key
+            ]
+
+            column_names = sorted(table_contract.columns.keys())
+
+            foreign_keys = [
+                f"{table_ref.full_name} -> {fk.referred_table.full_name}"
+                for fk in table_contract.foreign_keys
+            ]
+
+            chunks.append(
+                TableChunk(
+                    id=f"schema.table:{table_ref.full_name}:{self.schema_version}",
+                    datasource_id=self.ds_id,
+                    table=table_ref,
+                    description=table_md.description if table_md else None,
+                    primary_key=primary_keys,
+                        columns=column_names,
+                    foreign_keys=foreign_keys,
+                    row_count=table_md.row_count if table_md else None,
+                    schema_version=self.schema_version,
+                )
+            )
+
+        return chunks
+
+    def _build_column_chunks(self) -> List[ColumnChunk]:
+        """
+        Builds column-level chunks.
+
+        Returns:
+            List of ColumnChunk objects.
+        """
+        chunks: List[ColumnChunk] = []
+
+        contract = self.schema_snapshot.contract
+        metadata = self.schema_snapshot.metadata
+
+        for table_key, table_contract in contract.tables.items():
+            table_md = metadata.tables.get(table_key)
+            table_ref = TableRef(
+                schema_name=table_contract.table.schema_name,
+                table_name=table_contract.table.table_name,
+            )
+
+            for column_name, column_contract in table_contract.columns.items():
+                column_md = table_md.columns.get(column_name) if table_md else None
+
+                column_ref = ColumnRef(
+                    table=table_ref,
+                    column_name=column_name,
+                )
+
+                chunks.append(
+                    ColumnChunk(
+                        id=f"schema.column:{column_ref.table.full_name}:{column_ref.column_name}:{self.schema_version}",
+                        datasource_id=self.ds_id,
+                        column=column_ref,
+                        dtype=column_contract.data_type,
+                        description=column_md.description if column_md else None,
+                        column_stats=(
+                            column_md.statistics.model_dump()
+                            if column_md and column_md.statistics
+                            else {}
+                        ),
+                        synonyms=column_md.synonyms if column_md else None,
+                        pii=column_md.pii if column_md else False,
+                        schema_version=self.schema_version,
+                    )
+                )
+
+        return chunks
+
+    def _build_relationship_chunks(self) -> List[RelationshipChunk]:
+        """
+        Builds relationship-level chunks.
+
+        Returns:
+            List of RelationshipChunk objects with unknown cardinality.
+        """
+        chunks: List[RelationshipChunk] = []
+
+        contract = self.schema_snapshot.contract
+
+        for table_contract in contract.tables.values():
+            from_table = TableRef(
+                schema_name=table_contract.table.schema_name,
+                table_name=table_contract.table.table_name,
+            )
+
+            for fk in table_contract.foreign_keys:
+                to_table = TableRef(
+                    schema_name=fk.referred_table.schema_name,
+                    table_name=fk.referred_table.table_name,
+                )
+
+                chunks.append(
+                    RelationshipChunk(
+                        id=(
+                            f"schema.relationship:"
+                            f"{from_table.full_name}"
+                            f"->{to_table.full_name}:"
+                            f"{self.schema_version}"
+                        ),
+                        datasource_id=self.ds_id,
+                        from_table=from_table,
+                        to_table=to_table,
+                        from_columns=fk.constrained_columns,
+                        to_columns=fk.referred_columns,
+                        cardinality="unknown",
+                        business_meaning=None,
+                        schema_version=self.schema_version,
+                    )
+                )
+
+        return chunks
+
+    def _build_metric_chunks(self) -> List[MetricChunk]:
+        """
+        Builds metric-level chunks.
+
+        Returns:
+            Empty list, as metrics are not part of SchemaSnapshot.
+        """
+        return []

nl2sql/indexing/embeddings.py

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+from typing import Dict, List, Optional, Tuple
+
+from langchain_core.embeddings import Embeddings
+from langchain_openai import OpenAIEmbeddings
+
+from nl2sql.common.logger import get_logger
+from nl2sql.common.settings import settings
+
+logger = get_logger(__name__)
+
+SUPPORTED_EMBEDDING_PROVIDERS: Tuple[str, ...] = ("openai", "local")
+
+LOCAL_EMBEDDING_MODEL = "all-MiniLM-L6-v2"
+LOCAL_EMBEDDING_DIMENSION = 384
+
+# Output dimensions of the embedding models this project knows about. Used only
+# to detect an index built with a different embedder; unknown models are skipped.
+KNOWN_EMBEDDING_DIMENSIONS: Dict[str, int] = {
+    "text-embedding-3-small": 1536,
+    "text-embedding-3-large": 3072,
+    "text-embedding-ada-002": 1536,
+    LOCAL_EMBEDDING_MODEL: LOCAL_EMBEDDING_DIMENSION,
+}
+
+PROVIDER_BY_DIMENSION: Dict[int, str] = {
+    LOCAL_EMBEDDING_DIMENSION: "local",
+    1536: "openai",
+    3072: "openai",
+}
+
+
+class LocalEmbeddings(Embeddings):
+    """
+    Key-free embeddings backed by the ONNX all-MiniLM-L6-v2 model bundled with
+    chromadb. Produces 384-dimensional vectors and requires no API key.
+    """
+
+    def __init__(self) -> None:
+        self._embedding_function = None
+
+    def _get_embedding_function(self):
+        """
+        Lazily builds chromadb's default embedding function.
+
+        Returns:
+            The chromadb embedding function instance.
+        """
+        if self._embedding_function is None:
+            # Imported here so nothing pays for chromadb's ONNX runtime unless
+            # local embeddings are actually selected.
+            from chromadb.utils.embedding_functions import DefaultEmbeddingFunction
+
+            logger.info(
+                f"Initializing local embeddings ({LOCAL_EMBEDDING_MODEL}). "
+                "The first run downloads the ONNX model (~79 MB) into the local "
+                "cache directory, which can take a few minutes."
+            )
+            self._embedding_function = DefaultEmbeddingFunction()
+        return self._embedding_function
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """
+        Embeds a batch of documents.
+
+        Args:
+            texts: Documents to embed.
+
+        Returns:
+            One plain float vector per document.
+        """
+        vectors = self._get_embedding_function()(list(texts))
+        return [[float(value) for value in vector] for vector in vectors]
+
+    def embed_query(self, text: str) -> List[float]:
+        """
+        Embeds a single query.
+
+        Args:
+            text: Query to embed.
+
+        Returns:
+            A flat plain float vector.
+        """
+        return self.embed_documents([text])[0]
+
+
+def describe_embeddings(embeddings: Embeddings) -> Tuple[str, str, Optional[int]]:
+    """
+    Describes an embedder for diagnostics.
+
+    Args:
+        embeddings: Embedding implementation in use.
+
+    Returns:
+        Tuple of provider name, model name and known output dimension. The
+        dimension is None when the model is not a known one.
+    """
+    if isinstance(embeddings, LocalEmbeddings):
+        return "local", LOCAL_EMBEDDING_MODEL, LOCAL_EMBEDDING_DIMENSION
+
+    model = getattr(embeddings, "model", None)
+    provider = "openai" if isinstance(embeddings, OpenAIEmbeddings) else type(embeddings).__name__
+    dimension = KNOWN_EMBEDDING_DIMENSIONS.get(model) if model else None
+    return provider, model or "unknown", dimension
+
+
+class EmbeddingService:
+    """
+    Centralized service for managing embedding models.
+    Ensures consistency across the application.
+    """
+
+    _instance: Optional[Embeddings] = None
+    _instance_provider: Optional[str] = None
+
+    @classmethod
+    def get_embeddings(cls) -> Embeddings:
+        """
+        Returns the embeddings instance for the configured provider.
+
+        The instance is cached per provider so a runtime settings reload
+        (``reload_settings``) cannot hand back an embedder for the old provider.
+        """
+        provider = cls._resolve_provider()
+        if cls._instance is None or cls._instance_provider != provider:
+            cls._instance = cls._build_embeddings(provider)
+            cls._instance_provider = provider
+        return cls._instance
+
+    @classmethod
+    def reset(cls) -> None:
+        """Drops the cached embedder so the next call rebuilds it.
+
+        The cache lives on the class and outlives any single caller, so tests
+        and long-lived processes that change the configured provider need a way
+        to discard it.
+        """
+        cls._instance = None
+        cls._instance_provider = None
+
+    @classmethod
+    def _resolve_provider(cls) -> str:
+        """Returns the normalized embedding provider from settings."""
+        return (settings.embedding_provider or "openai").strip().lower()
+
+    @classmethod
+    def _build_embeddings(cls, provider: str) -> Embeddings:
+        """
+        Builds an embeddings instance for a provider.
+
+        Args:
+            provider: Normalized provider name.
+
+        Returns:
+            The embeddings implementation.
+
+        Raises:
+            ValueError: If the provider is not recognized.
+        """
+        if provider == "openai":
+            return OpenAIEmbeddings(
+                model=settings.embedding_model,
+                api_key=settings.openai_api_key,
+            )
+        if provider == "local":
+            return LocalEmbeddings()
+
+        raise ValueError(
+            f"Unknown EMBEDDING_PROVIDER '{provider}'. "
+            f"Valid options are: {', '.join(SUPPORTED_EMBEDDING_PROVIDERS)}."
+        )
+
+    @classmethod
+    def get_model_name(cls) -> str:
+        """Returns the name of the configured embedding model."""
+        if cls._resolve_provider() == "local":
+            return LOCAL_EMBEDDING_MODEL
+        return settings.embedding_model