nl2sql-engine 0.1.0__py3-none-any.whl

nl2sql/configs/datasources.py

@@ -0,0 +1,10 @@
+
+from typing import Optional, Dict, Any, Union, List
+from pydantic import BaseModel, Field
+from nl2sql.datasources import DatasourceConfig, ConnectionConfig
+
+
+class DatasourceFileConfig(BaseModel):
+    """File-level schema for datasources.yaml."""
+    version: int = Field(1, description="Schema version")
+    datasources: List[DatasourceConfig]

nl2sql/configs/llm.py

@@ -0,0 +1,36 @@
+
+from typing import Optional, Dict
+from pydantic import BaseModel, Field, SecretStr, field_serializer
+
+class AgentConfig(BaseModel):
+    """Configuration for a specific agent's LLM.
+
+    The single definition. ``nl2sql.llm.models`` re-exports this class, so
+    ``ConfigManager``/``LLMGenerator`` and ``LLMRegistry`` share one model.
+    Do not add a second copy: the previous duplicate silently swallowed a fix
+    applied to only one of the two.
+    """
+    provider: str
+    model: str
+    temperature: float = 0.0
+    api_key: Optional[SecretStr] = None
+    base_url: Optional[str] = Field(
+        None,
+        description=(
+            "Override the provider endpoint. Defaults to the provider preset: "
+            "the OpenRouter gateway for 'openrouter', the local Ollama daemon "
+            "for 'ollama', the client default for 'openai'. Set it to reach any "
+            "other OpenAI-compatible endpoint."
+        ),
+    )
+    name: str = Field("default", description="Name of the agent")
+
+    @field_serializer("api_key", when_used="json")
+    def _serialize_api_key(self, value):
+        return value.get_secret_value() if value else None
+
+class LLMFileConfig(BaseModel):
+    """Global LLM configuration (File Envelope)."""
+    version: int = Field(1, description="Schema version")
+    default: AgentConfig
+    agents: Optional[Dict[str, AgentConfig]] = Field(default_factory=dict)

nl2sql/configs/manager.py

@@ -0,0 +1,176 @@
+
+import yaml
+import json
+import shutil
+import pathlib
+from typing import List, Dict, Optional, Union
+from pydantic import ValidationError
+
+from nl2sql.common.settings import settings
+from .datasources import DatasourceConfig, DatasourceFileConfig
+from .llm import LLMFileConfig
+from .policies import PolicyFileConfig
+from .secrets import SecretProviderConfig, SecretsFileConfig
+from .sample_questions import SampleQuestionsFileConfig
+
+class ConfigManager:
+    """
+    Centralized manager for reading and writing application configuration.
+    Enforces consistency and handles file I/O for Datasources, LLMs, Policies, and Secrets.
+
+    Secret references such as ``${env:VAR}`` are returned verbatim. Resolution
+    happens at point of use, in ``DatasourceRegistry`` and ``LLMRegistry``.
+    Do not resolve them here: a resolved config written back to disk would
+    persist real credentials in plaintext.
+    """
+
+    def __init__(self, project_root: Optional[pathlib.Path] = None):
+        """
+        Args:
+            project_root: Optional override for project root. 
+                          If None, uses settings paths or CWD resolution strategy.
+        """
+        self.project_root = project_root
+        
+        # Resolve root: Use override, or CWD
+        root = self.project_root or pathlib.Path.cwd()
+        
+        # Default paths from settings if not overridden
+        self._ds_path = root / settings.datasource_config_path
+        self._llm_path = root / settings.llm_config_path
+        self._policy_path = root / settings.policies_config_path
+        self._secrets_path = root / settings.secrets_config_path
+        self._sample_questions_path = root / settings.sample_questions_path
+
+    def ensure_config_dirs(self) -> None:
+        """Ensures that the configuration directories exist."""
+        for path in [self._ds_path, self._llm_path, self._policy_path, self._secrets_path]:
+            if not path.parent.exists():
+                path.parent.mkdir(parents=True, exist_ok=True)
+
+    def load_datasources(self, path: Optional[pathlib.Path] = None) -> List[DatasourceConfig]:
+        """
+        Loads datasource configurations from YAML.
+        Handles Legacy (Dict) vs V2 (List under 'datasources' key) formats.
+        """
+        target_path = path or self._ds_path
+        
+        if not target_path.exists():
+            raise FileNotFoundError(f"Datasource config not found: {target_path}")
+
+        try:
+            content = target_path.read_text(encoding="utf-8")
+            raw = yaml.safe_load(content) or {}
+        except ImportError as exc:
+            raise RuntimeError("PyYAML is required to load datasource configs") from exc
+        except Exception as e:
+            raise ValueError(f"Failed to parse YAML from {target_path}: {e}")
+
+        try:
+            # Structurally validate the file envelope
+            file_config = DatasourceFileConfig.model_validate(raw)
+            return file_config.datasources
+        except ValidationError as e:
+            raise ValueError(f"Datasource Configuration Invalid: {e}")
+
+    def load_llm(self, path: Optional[pathlib.Path] = None) -> LLMFileConfig:
+        """
+        Loads LLM configuration. 
+        Returns nl2sql.configs.LLMFileConfig object.
+        """
+        target_path = path or self._llm_path
+        
+        if not target_path.exists():
+             raise FileNotFoundError(f"LLM config not found: {target_path}")
+
+        try:
+            data = yaml.safe_load(target_path.read_text()) or {}
+            
+            # Use File Config for validation
+            return LLMFileConfig.model_validate(data)
+        except ValidationError as e:
+            raise ValueError(f"LLM Configuration Invalid: {e}")
+
+
+
+    def load_policies(self, path: Optional[pathlib.Path] = None) -> PolicyFileConfig:
+        """
+        Loads Policy configuration.
+        Returns nl2sql.configs.PolicyFileConfig object.
+        """
+        target_path = path or self._policy_path
+        
+        if not target_path.exists():
+            raise FileNotFoundError(f"Policy config not found: {target_path}")
+            
+        try:
+            with open(target_path, "r") as f:
+                raw_json = f.read()
+                
+            data = json.loads(raw_json)
+            # Use File Config for validation
+            return PolicyFileConfig.model_validate(data)
+        except ValidationError as ve:
+            raise ValueError(f"Policy Schema Validation Failed: {ve}")
+        except Exception as e:
+            raise ValueError(f"Failed to load policies: {e}")
+
+
+
+    def load_secrets(self, path: Optional[pathlib.Path] = None) -> List[SecretProviderConfig]:
+        """Loads Secret configurations."""
+        target_path = path or self._secrets_path
+        if not target_path.exists():
+            return []
+            
+        try:
+            content = target_path.read_text(encoding="utf-8")
+            raw = yaml.safe_load(content) or []
+            
+            file_config = SecretsFileConfig.model_validate(raw)
+            return file_config.providers
+        except ValidationError as e:
+            raise ValueError(f"Secret Configuration Invalid: {e}")
+        except Exception as e:
+            raise ValueError(f"Failed to load secrets: {e}")
+
+
+    def load_sample_questions(self, path: Optional[pathlib.Path] = None, ds_id: Optional[str] = None) -> Union[Dict[str, List[str]], List[str]]:
+        """Loads sample questions from a YAML file."""
+        target_path = path or self._sample_questions_path
+        
+        if not target_path.exists():
+            raise FileNotFoundError(f"Sample questions file not found: {target_path}")
+            
+        try:
+            raw = yaml.safe_load(target_path.read_text(encoding="utf-8")) or {}
+            
+            if ds_id:
+                return raw.get(ds_id, [])
+
+            return raw
+        except Exception as e:
+            raise ValueError(f"Failed to load sample questions: {e}")
+
+    def get_example_questions(self, datasource_id: str) -> List[str]:
+        """Returns example questions for a datasource, if configured."""
+        try:
+            questions = self.load_sample_questions(ds_id=datasource_id)
+            return questions or []
+        except FileNotFoundError:
+            return []
+
+    def get_datasource_description(self, datasource_id: str) -> Optional[str]:
+        """Returns the configured datasource description if available."""
+        try:
+            datasources = self.load_datasources()
+        except FileNotFoundError:
+            return None
+        for ds in datasources:
+            if ds.id == datasource_id:
+                return ds.description
+        return None
+
+
+
+

nl2sql/configs/policies.py

@@ -0,0 +1,14 @@
+
+from typing import List, Dict, Optional
+from pydantic import BaseModel, Field, field_validator
+from nl2sql.auth import RolePolicy
+
+
+
+
+class PolicyFileConfig(BaseModel):
+    """File-level schema for policies.json."""
+    version: int = Field(1, description="Schema version")
+    roles: Dict[str, RolePolicy]
+
+

nl2sql/configs/sample_questions.py

@@ -0,0 +1,11 @@
+from typing import Dict, List
+
+from pydantic import BaseModel, Field, ConfigDict
+
+
+class SampleQuestionsFileConfig(BaseModel):
+    """File-level schema for sample_questions.yaml."""
+
+    datasources: Dict[str, List[str]] = Field(default_factory=dict)
+
+    model_config = ConfigDict(extra="ignore")

nl2sql/configs/secrets.py

@@ -0,0 +1,11 @@
+
+from typing import List
+
+from pydantic import BaseModel, Field
+from nl2sql.secrets.models import SecretProviderConfig
+
+class SecretsFileConfig(BaseModel):
+    """File-level schema for secrets.yaml."""
+    version: int = Field(1, description="Schema version")
+    providers: List[SecretProviderConfig]
+

nl2sql/context.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+import pathlib
+from typing import Optional
+
+from nl2sql.configs import ConfigManager
+from nl2sql.datasources import DatasourceRegistry
+from nl2sql.llm import LLMRegistry
+from nl2sql.indexing.vector_store import VectorStore
+from nl2sql.secrets import SecretManager
+from nl2sql.common.settings import settings
+from nl2sql.auth import RBAC
+
+from nl2sql.schema import build_schema_store
+from nl2sql.execution import ExecutionStore
+from nl2sql.execution.artifacts import build_artifact_store
+
+from nl2sql.common.logger import get_logger
+logger = get_logger(__name__)
+
+
+class NL2SQLContext:
+    """
+    Centralized application context that manages the initialization lifecycle.
+    
+    Ensures that secrets are loaded BEFORE datasources, and handles
+    registry instantiation in a consistent order.
+    """
+
+    def __init__(
+        self,
+        ds_config_path: Optional[pathlib.Path] = None,
+        secrets_config_path: Optional[pathlib.Path] = None,
+        llm_config_path: Optional[pathlib.Path] = None,
+        vector_store_path: Optional[pathlib.Path] = None,
+        policies_config_path: Optional[pathlib.Path] = None,
+    ) :
+        """
+        Factory method to create a context from configuration paths.
+        Resolves defaults from global settings if paths are not provided.
+        """
+        ds_config_path = ds_config_path or pathlib.Path(settings.datasource_config_path)
+        secrets_config_path = secrets_config_path or pathlib.Path(settings.secrets_config_path)
+        llm_config_path = llm_config_path or pathlib.Path(settings.llm_config_path)
+        policies_config_path = policies_config_path or pathlib.Path(settings.policies_config_path)
+
+
+        logger.info(f"Loading datasource configuration from {ds_config_path}")
+        logger.info(f"Loading secrets configuration from {secrets_config_path}")
+        logger.info(f"Loading LLM configuration from {llm_config_path}")
+        logger.info(f"Loading policies configuration from {policies_config_path}")
+        logger.info(f"Loading vector store configuration from {vector_store_path}")
+
+        # Validate vector store configuration first so misconfiguration fails fast,
+        # before secrets, datasource and LLM registries are constructed.
+        vector_store_collection_name = settings.vector_store_collection_name
+        if not vector_store_collection_name:
+            raise ValueError(
+                "VECTOR_STORE_COLLECTION name is not configured. "
+                "Set vector_store_collection_name in settings."
+            )
+
+        vector_store_path = vector_store_path or (
+            pathlib.Path(settings.vector_store_path) if settings.vector_store_path else None
+        )
+        if not vector_store_path:
+            raise ValueError(
+                "VECTOR_STORE path is not configured. "
+                "Pass vector_store_path or set vector_store_path in settings."
+            )
+
+        cm = ConfigManager()
+        self.tenant_id = settings.tenant_id
+        self.config_manager = cm
+
+        secret_configs = cm.load_secrets(secrets_config_path)
+        secret_manager = SecretManager()
+        if secret_configs:
+            secret_manager.configure(secret_configs)
+
+        ds_configs = cm.load_datasources(ds_config_path)
+        self.ds_registry = DatasourceRegistry(secret_manager)
+        self.ds_registry.register_datasources(ds_configs)
+
+        llm_cfg = cm.load_llm(llm_config_path)
+        self.llm_registry = LLMRegistry(secret_manager)
+
+        agents = llm_cfg.agents or {}
+        agents["default"] = llm_cfg.default
+        self.llm_registry.register_llms(agents)
+
+        self.policies_cfg = cm.load_policies(policies_config_path)
+        self.rbac = RBAC(self.policies_cfg.roles)
+
+        self.vector_store = VectorStore(
+            collection_name=vector_store_collection_name,
+            persist_directory=vector_store_path,
+        )
+        self.schema_store = build_schema_store(
+            settings.schema_store_backend,
+            settings.schema_store_max_versions,
+            path=pathlib.Path(settings.schema_store_path),
+        )
+        self.execution_store = ExecutionStore()
+        self.artifact_store = build_artifact_store()
+
+       

nl2sql/datasources/__init__.py

@@ -0,0 +1,21 @@
+"""Datasource management, configuration, and adapter discovery."""
+from nl2sql.datasources.registry import DatasourceRegistry
+from nl2sql.datasources.discovery import discover_adapters
+from nl2sql.datasources.models import (
+    DatasourceConfig,
+    ConnectionConfig,
+)
+from nl2sql.datasources.protocols import DatasourceAdapterProtocol
+from nl2sql_adapter_sdk.contracts import AdapterRequest, ResultFrame
+from nl2sql_adapter_sdk.capabilities import DatasourceCapability
+
+__all__ = [
+    "DatasourceRegistry", 
+    "discover_adapters", 
+    "DatasourceConfig",
+    "ConnectionConfig",
+    "DatasourceAdapterProtocol",
+    "AdapterRequest",
+    "DatasourceCapability",
+    "ResultFrame",
+]

nl2sql/datasources/discovery.py

@@ -0,0 +1,28 @@
+from importlib.metadata import entry_points
+from typing import Dict, Type
+
+from nl2sql.datasources.protocols import DatasourceAdapterProtocol
+from nl2sql.common.logger import get_logger
+
+logger = get_logger(__name__)
+
+def discover_adapters() -> Dict[str, Type[DatasourceAdapterProtocol]]:
+    """Discovers installed adapters via 'nl2sql.adapters' entry points.
+    
+    Returns:
+        Dict[str, Type[DatasourceAdapterProtocol]]: Dict mapping adapter name (e.g., 'postgres')
+            to the Adapter Class.
+    """
+    adapters = {}
+    try:
+        eps = entry_points(group="nl2sql.adapters")
+    except TypeError:
+        eps = entry_points().get("nl2sql.adapters", [])
+    for ep in eps:
+        try:
+            AdapterCls = ep.load()
+            adapters[ep.name] = AdapterCls
+        except Exception as e:
+            logger.error(f"Failed to load adapter {ep.name}: {e}")
+            
+    return adapters

nl2sql/datasources/models.py

@@ -0,0 +1,21 @@
+
+
+from pydantic import BaseModel, Field
+from typing import Optional, Dict, Any, List
+
+
+class ConnectionConfig(BaseModel):
+    """Database connection details."""
+    type: str 
+    
+    model_config = {"extra": "allow"}
+
+class DatasourceConfig(BaseModel):
+    """Configuration for a single datasource."""
+    id: str
+    description: Optional[str] = None
+    connection: ConnectionConfig
+    options: Dict[str, Any] = Field(default_factory=dict)
+
+
+

nl2sql/datasources/protocols.py

@@ -0,0 +1,3 @@
+from nl2sql_adapter_sdk.protocols import DatasourceAdapterProtocol
+
+__all__ = ["DatasourceAdapterProtocol"]

nl2sql/datasources/registry.py

@@ -0,0 +1,172 @@
+from typing import Dict, List, Any, Set
+from threading import RLock
+
+from nl2sql_adapter_sdk.capabilities import DatasourceCapability
+from nl2sql.datasources.discovery import discover_adapters
+from nl2sql.datasources.protocols import DatasourceAdapterProtocol
+from nl2sql.secrets import SecretManager
+from .models import DatasourceConfig, ConnectionConfig
+
+class DatasourceRegistry:
+    """Manages a collection of active DatasourceAdapters.
+    
+    Acts as the factory and cache for DataSourceAdapter instances.
+    """
+
+    def __init__(self, secret_manager: SecretManager):
+        """Initializes the registry by eagerly creating adapters for all configs."""
+        self._adapters: Dict[str, DatasourceAdapterProtocol] = {}
+        self._capabilities: Dict[str, Set[str]] = {}
+        self._available_adapters = discover_adapters()
+        self._secret_manager = secret_manager
+        self._lock = RLock()
+
+    def find_and_resolve_secret(self, key: str) -> str:
+        """Attempts to resolve a secret using the secret manager.
+
+        Args:
+            key (str): The name of the secret to resolve.
+
+        Returns:
+            str: The resolved secret value.
+
+        Raises:
+            ValueError: If the secret cannot be resolved.
+        """
+        return self._secret_manager.resolve(key)
+
+    def resolved_connection(self, unresolved_connection: ConnectionConfig) -> ConnectionConfig:
+        """Resolves all secrets in a connection dictionary.
+
+        Args:
+            unresolved_connection (ConnectionConfig): The connection dictionary with potential secrets.
+
+        Returns:
+            ConnectionConfig: The connection dictionary with resolved secrets.
+        """
+        from pydantic import SecretStr
+        
+        resolved_connection = unresolved_connection.model_dump()
+        for key, value in unresolved_connection.model_dump().items():
+            if isinstance(value, str) and value.startswith("${") and value.endswith("}"):
+                secret_val = self.find_and_resolve_secret(value)
+                resolved_connection[key] = SecretStr(secret_val)
+        return ConnectionConfig(**resolved_connection)
+
+    def register_datasources(self, configs: List[DatasourceConfig]):
+        for config in configs:
+            try:
+                self.register_datasource(config)
+            except Exception as e:
+                raise ValueError(f"Failed to initialize adapter for '{config.id}': {e}") from e
+
+    def _normalize_capabilities(self, caps: Any) -> Set[str]:
+        if not caps:
+            return set()
+        normalized = set()
+        for cap in caps:
+            if isinstance(cap, DatasourceCapability):
+                normalized.add(cap.value)
+            else:
+                normalized.add(str(cap))
+        return normalized
+
+    def register_datasource(self, config: DatasourceConfig) -> DatasourceAdapterProtocol:
+        """Registers a new datasource dynamically.
+
+        Args:
+            config: The datasource configuration dictionary or object.
+
+        Returns:
+            DatasourceAdapterProtocol: The created and registered adapter.
+
+        Raises:
+            ValueError: If configuration is invalid or adapter type is unknown.
+        """
+        ds_id = config.id
+        if not ds_id:
+            raise ValueError("Datasource ID is required. Please check your configuration.")
+
+        connection = config.connection
+        conn_type = connection.type.lower()
+        resolved_connection = self.resolved_connection(connection)
+        connection_args = resolved_connection.model_dump()
+
+        if conn_type in self._available_adapters:
+            adapter_cls = self._available_adapters[conn_type]
+
+            adapter = adapter_cls(
+                datasource_id=ds_id,
+                datasource_engine_type=conn_type,
+                connection_args=connection_args,
+                statement_timeout_ms=config.options.get("statement_timeout_ms"),
+                row_limit=config.options.get("row_limit"),
+                max_bytes=config.options.get("max_bytes"),
+            )
+            with self._lock:
+                self._adapters[ds_id] = adapter
+                if hasattr(adapter, "capabilities"):
+                    self._capabilities[ds_id] = self._normalize_capabilities(
+                        adapter.capabilities()
+                    )
+                else:
+                    self._capabilities[ds_id] = {DatasourceCapability.SUPPORTS_SQL.value}
+            return adapter
+        else:
+            raise ValueError(
+                f"No adapter found for engine type: '{conn_type}' in datasource '{ds_id}'"
+            )
+
+
+    def get_adapter(self, datasource_id: str) -> DatasourceAdapterProtocol:
+        """Retrieves the adapter for a datasource.
+
+        Args:
+            datasource_id: The ID of the datasource.
+
+        Returns:
+            DatasourceAdapterProtocol: The active adapter instance.
+
+        Raises:
+            ValueError: If the datasource ID is unknown.
+        """
+        with self._lock:
+            if datasource_id not in self._adapters:
+                raise ValueError(f"Unknown datasource ID: {datasource_id}")
+            return self._adapters[datasource_id]
+
+    def get_dialect(self, datasource_id: str) -> str:
+        """Returns a normalized dialect string from the adapter.
+
+        Args:
+            datasource_id: The ID of the datasource.
+
+        Returns:
+            str: The dialect string (e.g., 'postgres').
+        """
+        return self.get_adapter(datasource_id).get_dialect()
+
+    def get_capabilities(self, datasource_id: str) -> Set[str]:
+        """Returns the capability flags for a datasource."""
+        with self._lock:
+            if datasource_id not in self._capabilities:
+                raise ValueError(f"Unknown datasource ID: {datasource_id}")
+            return set(self._capabilities[datasource_id])
+
+    def list_adapters(self) -> List[DatasourceAdapterProtocol]:
+        """Returns a list of all registered adapters.
+
+        Returns:
+            List[DatasourceAdapterProtocol]: All active adapters.
+        """
+        with self._lock:
+            return list(self._adapters.values())
+
+    def list_ids(self) -> List[str]:
+        """Returns a list of all registered datasource IDs.
+
+        Returns:
+            List[str]: All registered IDs.
+        """
+        with self._lock:
+            return list(self._adapters.keys())

nl2sql/evaluation/__init__.py

@@ -0,0 +1,6 @@
+"""Evaluation tools for the NL2SQL pipeline."""
+from nl2sql.evaluation.evaluator import ModelEvaluator
+from nl2sql.evaluation.types import BenchmarkConfig
+from nl2sql.evaluation.benchmark_runner import BenchmarkRunner
+
+__all__ = ["ModelEvaluator", "BenchmarkConfig", "BenchmarkRunner"]