aetherdialect 0.1.0__py3-none-any.whl

text2sql/text2sql.py

@@ -0,0 +1,726 @@
+"""High-level API for text-to-SQL pipeline orchestration and mode management.
+
+Provides the ``Text2SQL`` class as the primary entry point for all pipeline operations: engine configuration, API key management, schema loading, template persistence, and execution mode switching between interactive, simulator, and question-generation (qsim) modes.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+from dotenv import load_dotenv
+
+from .config import (
+    DatabricksRuntimeConfig,
+    EngineConfig,
+    PostgresRuntimeConfig,
+    QSimConfig,
+)
+from .contracts_base import (
+    QSimRange,
+    QSimSummary,
+    RejectedTemplateInfo,
+    SchemaGraph,
+    SimulatorSummary,
+    TemplateInfo,
+)
+from .dialect import get_dialect
+from .main_execution import (
+    get_artifacts_dir,
+    get_qsim,
+    get_questions_only,
+    get_rejected_templates_list,
+    get_simulator_summary_from_dir,
+    get_templates_list,
+    interactive_run_once,
+    load_generated_questions,
+    qsim_run_once,
+    resolve_qsim_path,
+    simulator_run_once,
+)
+from .schema import load_or_create_schema_graph
+from .templates import (
+    load_template_store,
+    store_to_rejected_templates,
+    store_to_templates,
+)
+
+
+def _first_non_empty_env(names: list[str]) -> str | None:
+    """Return the first non-empty environment value for the given names."""
+    for name in names:
+        value = os.environ.get(name)
+        if value:
+            return value
+    return None
+
+
+def _resolve_postgres_runtime_config(
+    host: str | None,
+    port: int | None,
+    user: str | None,
+    password: str | None,
+    database: str | None,
+    schema: str | None,
+) -> tuple[str | None, int | None, str | None, str | None, str | None, str | None]:
+    """Resolve PostgreSQL runtime configuration from arguments and environment."""
+    resolved_host = host or _first_non_empty_env(
+        ["POSTGRESQL_HOST", "PGHOST", "DB_HOST"]
+    )
+    port_env = _first_non_empty_env(
+        ["POSTGRESQL_PORT", "PGPORT", "DB_PORT"]
+    )
+    resolved_port = port
+    if resolved_port is None and port_env is not None:
+        try:
+            resolved_port = int(port_env)
+        except ValueError:
+            resolved_port = None
+    resolved_user = user or _first_non_empty_env(
+        ["POSTGRESQL_USER", "PGUSER", "DB_USER"]
+    )
+    resolved_password = password or _first_non_empty_env(
+        ["POSTGRESQL_PASSWORD", "PGPASSWORD", "DB_PASSWORD"]
+    )
+    resolved_database = database or _first_non_empty_env(
+        ["POSTGRESQL_DB", "POSTGRESQL_DATABASE", "PGDATABASE", "DB_NAME", "DB_DATABASE"]
+    )
+    resolved_schema = schema or _first_non_empty_env(
+        ["POSTGRESQL_SCHEMA", "DB_SCHEMA"]
+    )
+    return (
+        resolved_host,
+        resolved_port,
+        resolved_user,
+        resolved_password,
+        resolved_database,
+        resolved_schema,
+    )
+
+
+def _resolve_databricks_runtime_config(
+    catalog: str | None,
+    schema: str | None,
+    server_hostname: str | None,
+    http_path: str | None,
+    access_token: str | None,
+) -> tuple[str | None, str | None, str | None, str | None, str | None]:
+    """Resolve Databricks runtime configuration from arguments and environment."""
+    resolved_catalog = catalog or _first_non_empty_env(
+        ["DATABRICKS_CATALOG", "DBR_CATALOG"]
+    )
+    resolved_schema = schema or _first_non_empty_env(
+        ["DATABRICKS_SCHEMA", "DBR_SCHEMA"]
+    )
+    resolved_server = server_hostname or _first_non_empty_env(
+        [
+            "DATABRICKS_SERVER_HOSTNAME",
+            "DATABRICKS_HOST",
+            "DATABRICKS_SERVER",
+            "DBR_SERVER_HOSTNAME",
+            "DBR_HOST",
+        ]
+    )
+    resolved_http_path = http_path or _first_non_empty_env(
+        [
+            "DATABRICKS_HTTP_PATH",
+            "DATABRICKS_WAREHOUSE_HTTP_PATH",
+            "DATABRICKS_SQL_HTTP_PATH",
+            "DBR_HTTP_PATH",
+        ]
+    )
+    resolved_token = access_token or _first_non_empty_env(
+        [
+            "DATABRICKS_ACCESS_TOKEN",
+            "DATABRICKS_PAT",
+            "ACCESS_TOKEN",
+            "PERSONAL_ACCESS_TOKEN",
+            "DBR_ACCESS_TOKEN",
+            "DBR_PAT",
+            "DATABRICKS_TOKEN",
+        ]
+    )
+    return (
+        resolved_catalog,
+        resolved_schema,
+        resolved_server,
+        resolved_http_path,
+        resolved_token,
+    )
+
+
+class Text2SQL:
+    """Primary entry point for the text-to-SQL pipeline.
+
+    Manages engine configuration, API key resolution, schema loading,
+    template persistence, and dispatches execution to the selected
+    operational mode (interactive, simulator, or qsim).
+    """
+
+    def __init__(
+        self,
+        engine: str = "postgresql",
+        mode: str = "interactive",
+        env_file: str | None = None,
+        openai_api_key: str | None = None,
+        host: str | None = None,
+        port: int | None = None,
+        user: str | None = None,
+        password: str | None = None,
+        database: str | None = None,
+        schema: str | None = None,
+        catalog: str | None = None,
+        sql_file: str | None = None,
+        server_hostname: str | None = None,
+        http_path: str | None = None,
+        access_token: str | None = None,
+    ) -> None:
+        """Initialise the deterministic Text-to-SQL system.
+
+        Args:
+
+            engine: ``"postgresql"`` or ``"databricks"``.
+
+            mode: ``"interactive"``, ``"simulator"``, or ``"qsim"``.
+
+            env_file: Path to ``.env`` containing ``OPENAI_API_KEY``.
+
+            openai_api_key: Direct API key (overrides *env_file* and environment variable).
+
+            host: PostgreSQL host.
+
+            port: PostgreSQL port.
+
+            user: PostgreSQL user.
+
+            password: PostgreSQL password.
+
+            database: PostgreSQL database name.
+
+            schema: PostgreSQL or Databricks schema name.
+
+            catalog: Databricks catalog name.
+
+            sql_file: Optional ``CREATE TABLE`` file for offline schema loading.
+
+            server_hostname: Databricks server hostname for ``databricks-sql-connector``.
+
+            http_path: Databricks HTTP path for ``databricks-sql-connector``.
+
+            access_token: Databricks personal access token for ``databricks-sql-connector``.
+
+        Raises:
+
+            ValueError: If *engine* or *mode* is invalid.
+        """
+        if engine not in ("postgresql", "databricks"):
+            raise ValueError(f"Invalid engine: {engine}. Must be 'postgresql' or 'databricks'.")
+        if mode not in ("interactive", "simulator", "qsim"):
+            raise ValueError(f"Invalid mode: {mode}. Must be 'interactive', 'simulator', or 'qsim'.")
+
+        self._mode = mode
+        self._api_key_set = False
+
+        if openai_api_key:
+            self.set_openai_api_key(openai_api_key)
+        elif env_file:
+            self.set_env(env_file)
+        elif os.environ.get("OPENAI_API_KEY"):
+            EngineConfig.API_TOKEN = os.environ.get("OPENAI_API_KEY")
+            self._api_key_set = True
+
+        EngineConfig.TYPE = engine
+
+        pg_host = host
+        pg_port = port
+        pg_user = user
+        pg_password = password
+        pg_database = database
+        pg_schema = schema
+        dbr_server_hostname = server_hostname
+        dbr_http_path = http_path
+        dbr_access_token = access_token
+        dbr_catalog = catalog
+        dbr_schema = schema
+
+        if engine == "postgresql":
+            (
+                pg_host,
+                pg_port,
+                pg_user,
+                pg_password,
+                pg_database,
+                pg_schema,
+            ) = _resolve_postgres_runtime_config(
+                host,
+                port,
+                user,
+                password,
+                database,
+                schema,
+            )
+        elif engine == "databricks":
+            (
+                dbr_catalog,
+                dbr_schema,
+                dbr_server_hostname,
+                dbr_http_path,
+                dbr_access_token,
+            ) = _resolve_databricks_runtime_config(
+                catalog,
+                schema,
+                server_hostname,
+                http_path,
+                access_token,
+            )
+
+        self.artifacts_dir = get_artifacts_dir(
+            engine,
+            pg_host if engine == "postgresql" else host,
+            pg_database if engine == "postgresql" else database,
+            pg_schema if engine == "postgresql" else dbr_schema,
+            dbr_catalog,
+        )
+        if not os.path.exists(self.artifacts_dir):
+            os.makedirs(self.artifacts_dir, exist_ok=True)
+
+        EngineConfig.SCHEMA_JSON_PATH = os.path.join(self.artifacts_dir, "schema_graph.json")
+        EngineConfig.TEMPLATE_JSON_PATH = os.path.join(self.artifacts_dir, "intent_templates.json")
+        QSimConfig.QUESTIONS_OUTPUT_PATH = os.path.join(self.artifacts_dir, QSimConfig.QUESTIONS_OUTPUT_PATH)
+        QSimConfig.SKELETONS_JSON_PATH = os.path.join(self.artifacts_dir, "qsim_skeletons.json")
+
+        if engine == "postgresql":
+            EngineConfig.RUNTIME = PostgresRuntimeConfig
+            if pg_host is not None:
+                PostgresRuntimeConfig.HOST = pg_host
+            if pg_port is not None:
+                PostgresRuntimeConfig.PORT = pg_port
+            if pg_user is not None:
+                PostgresRuntimeConfig.USER = pg_user
+            if pg_password is not None:
+                PostgresRuntimeConfig.PASSWORD = pg_password
+            if pg_database is not None:
+                PostgresRuntimeConfig.DATABASE = pg_database
+            if pg_schema is not None:
+                PostgresRuntimeConfig.SCHEMA = pg_schema
+            if sql_file is not None:
+                PostgresRuntimeConfig.SQL_FILE_PATH = sql_file
+        elif engine == "databricks":
+            EngineConfig.RUNTIME = DatabricksRuntimeConfig
+            if dbr_catalog is not None:
+                DatabricksRuntimeConfig.CATALOG = dbr_catalog
+            if dbr_schema is not None:
+                DatabricksRuntimeConfig.SCHEMA = dbr_schema
+            if sql_file is not None:
+                DatabricksRuntimeConfig.SQL_FILE_PATH = sql_file
+            if dbr_server_hostname is not None:
+                DatabricksRuntimeConfig.SERVER_HOSTNAME = dbr_server_hostname
+            if dbr_http_path is not None:
+                DatabricksRuntimeConfig.HTTP_PATH = dbr_http_path
+            if dbr_access_token is not None:
+                DatabricksRuntimeConfig.ACCESS_TOKEN = dbr_access_token
+            DatabricksRuntimeConfig.validate()
+
+        self.dialect = get_dialect(EngineConfig.TYPE, EngineConfig.RUNTIME)
+
+        if EngineConfig.TYPE == "postgresql":
+            from sqlalchemy import create_engine
+
+            self.engine = create_engine(EngineConfig.RUNTIME.db_url(), future=True)
+        elif EngineConfig.TYPE == "databricks":
+            self.engine = None
+
+        try:
+            self._schema = load_or_create_schema_graph(self.engine)
+            self._store = load_template_store(self._schema.schema_hash)
+            self._templates = store_to_templates(self._store)
+            self._rejected = store_to_rejected_templates(self._store)
+            self._schema_terms = set(self._schema.tables.keys())
+            for tinfo in self._schema.tables.values():
+                self._schema_terms.update(tinfo.columns)
+                for col in tinfo.columns:
+                    self._schema_terms.add(col.lower())
+            self._schema_stats = self._schema.schema_stats or {}
+        except Exception:
+            raise
+
+    @property
+    def schema(self) -> SchemaGraph:
+        """The loaded schema graph for the connected database."""
+        return self._schema
+
+    @property
+    def store(self) -> dict[str, Any]:
+        """The in-memory template store keyed by intent fingerprint."""
+        return self._store
+
+    @store.setter
+    def store(self, value: dict[str, Any]) -> None:
+        """Replace the template store."""
+        self._store = value
+
+    @property
+    def templates(self) -> dict[str, Any]:
+        """Dict of accepted templates derived from the template
+        store."""
+        return self._templates
+
+    @templates.setter
+    def templates(self, value: dict[str, Any]) -> None:
+        """Replace the accepted templates dict."""
+        self._templates = value
+
+    @property
+    def rejected(self) -> dict[str, Any]:
+        """Dict of rejected templates derived from the template
+        store."""
+        return self._rejected
+
+    @rejected.setter
+    def rejected(self, value: dict[str, Any]) -> None:
+        """Replace the rejected templates dict."""
+        self._rejected = value
+
+    @property
+    def schema_terms(self) -> set[str]:
+        """Set of all table and column name tokens from the loaded
+        schema."""
+        return self._schema_terms
+
+    def _compute_num_intents_range(self) -> tuple[int, int]:
+        """Compute an adaptive ``num_intents`` range based on schema complexity.
+
+        Returns:
+
+            A 2-tuple ``(min_intents, max_intents)`` derived from the table count.
+        """
+        table_count = self._schema_stats["table_count"]
+        min_intents = max(5, table_count)
+        max_intents = min(200, table_count * 10)
+        return (min_intents, max_intents)
+
+    def _compute_num_questions_range(self) -> tuple[int, int]:
+        """Compute an adaptive ``num_questions`` range based on schema variance potential.
+
+        Returns:
+
+            A 2-tuple ``(min_questions, max_questions)`` derived from the total number of filterable columns in the schema.
+        """
+        min_intents, _ = self._compute_num_intents_range()
+        total_filterable = self._schema_stats["total_filterable"]
+        min_questions = max(min_intents, 10)
+        max_questions = min(2000, total_filterable * 20)
+        return (min_questions, max_questions)
+
+    def get_qsim_range(self) -> QSimRange:
+        """Return valid ``num_intents`` and ``num_questions`` ranges for this schema."""
+        return QSimRange(
+            num_intents_range=self._compute_num_intents_range(),
+            num_questions_range=self._compute_num_questions_range(),
+        )
+
+    def set_mode(self, mode: str) -> None:
+        """Switch the operational mode.
+
+        Args:
+
+            mode: ``"interactive"``, ``"simulator"``, or ``"qsim"``.
+
+        Raises:
+
+            ValueError: If *mode* is invalid.
+        """
+        if mode not in ("interactive", "simulator", "qsim"):
+            raise ValueError(f"Invalid mode: {mode}. Must be 'interactive', 'simulator', or 'qsim'.")
+        self._mode = mode
+
+    def set_env(self, env_file: str) -> None:
+        """Load ``OPENAI_API_KEY`` from a ``.env`` file.
+
+        Args:
+
+            env_file: Path to the ``.env`` file.
+
+        Raises:
+
+            FileNotFoundError: If *env_file* does not exist.
+
+            ValueError: If ``OPENAI_API_KEY`` is missing from the file.
+        """
+        if not os.path.exists(env_file):
+            raise FileNotFoundError(f"Environment file not found: {env_file}")
+        load_dotenv(env_file)
+        api_key = os.environ.get("OPENAI_API_KEY")
+        if not api_key:
+            raise ValueError("OPENAI_API_KEY not found in environment file")
+        EngineConfig.API_TOKEN = api_key
+        self._api_key_set = True
+
+    def set_openai_api_key(self, key: str) -> None:
+        """Set the OpenAI API key directly.
+
+        Args:
+
+            key: Non-empty API key string.
+
+        Raises:
+
+            ValueError: If *key* is empty.
+        """
+        if not key or not key.strip():
+            raise ValueError("API key cannot be empty")
+        EngineConfig.API_TOKEN = key.strip()
+        self._api_key_set = True
+
+    def _ensure_api_key(self) -> None:
+        """Raise ``RuntimeError`` if no API key has been configured.
+
+        Raises:
+
+            RuntimeError: If neither ``set_openai_api_key``, ``set_env``, nor the ``openai_api_key`` or ``env_file`` constructor parameters were used to provide a key.
+        """
+        if not self._api_key_set and not EngineConfig.API_TOKEN:
+            raise RuntimeError(
+                "OpenAI API key not set. Use set_openai_api_key(), set_env(), or pass openai_api_key/env_file to __init__."
+            )
+
+    def _validate_num_intents(self, value: int) -> None:
+        """Validate *value* is within the schema-adaptive ``num_intents`` range.
+
+        Args:
+
+            value: The requested number of intents.
+
+        Raises:
+
+            ValueError: If *value* is outside ``[min_intents, max_intents]``.
+        """
+        min_intents, max_intents = self._compute_num_intents_range()
+        if not (min_intents <= value <= max_intents):
+            raise ValueError(
+                f"num_intents must be {min_intents}-{max_intents} for this schema ({self._schema_stats['table_count']} tables)"
+            )
+
+    def _validate_num_questions(self, value: int) -> None:
+        """Validate *value* is within the schema-adaptive ``num_questions`` range.
+
+        Args:
+
+            value: The requested number of questions.
+
+        Raises:
+
+            ValueError: If *value* is outside ``[min_questions, max_questions]``.
+        """
+        min_questions, max_questions = self._compute_num_questions_range()
+        if not (min_questions <= value <= max_questions):
+            raise ValueError(
+                f"num_questions must be {min_questions}-{max_questions} for this schema ({self._schema_stats['total_filterable']} total filterable columns)"
+            )
+
+    def run_interactive(self, debug: bool | None = None, verbose: bool | None = None) -> None:
+        """Run the interactive question-answering loop with template reuse.
+
+        Args:
+
+            debug: Override instance debug flag for this run.
+
+            verbose: Override instance verbose flag for this run.
+
+        Raises:
+
+            RuntimeError: If mode is not ``"interactive"`` or no API key is set.
+        """
+        if self._mode != "interactive":
+            raise RuntimeError(f"Current mode is '{self._mode}'. Call set_mode('interactive') first.")
+        self._ensure_api_key()
+
+        debug = self._debug if debug is None else debug
+        verbose = self._verbose if verbose is None else verbose
+
+        EngineConfig.DEBUG = debug
+        EngineConfig._propagate_debug()
+
+        prev_verbose_state = EngineConfig.VERBOSE
+        EngineConfig.VERBOSE = verbose
+        try:
+            interactive_run_once(
+                self._schema,
+                self._store,
+                self._templates,
+                self._rejected,
+                self._schema_terms,
+            )
+        finally:
+            EngineConfig.VERBOSE = prev_verbose_state
+
+    def run_simulator(
+        self,
+        seed_filepath: str,
+        interactive_gold: bool = True,
+        debug: bool | None = None,
+        verbose: bool | None = None,
+    ) -> SimulatorSummary:
+        """Generate synthetic templates from seed questions.
+
+        Args:
+
+            seed_filepath: Path to seed questions text file.
+
+            interactive_gold: Interactively confirm each gold intent when ``True``.
+
+            debug: Override instance debug flag for this run.
+
+            verbose: Override instance verbose flag for this run.
+
+        Returns:
+
+            ``SimulatorSummary`` describing the results.
+
+        Raises:
+
+            RuntimeError: If mode is not ``"simulator"`` or no API key is set.
+        """
+        if self._mode != "simulator":
+            raise RuntimeError(f"Current mode is '{self._mode}'. Call set_mode('simulator') first.")
+        self._ensure_api_key()
+
+        debug = self._debug if debug is None else debug
+        verbose = self._verbose if verbose is None else verbose
+
+        EngineConfig.DEBUG = debug
+        EngineConfig._propagate_debug()
+
+        prev_verbose_state = EngineConfig.VERBOSE
+        EngineConfig.VERBOSE = verbose
+        try:
+            return simulator_run_once(
+                schema=self.schema,
+                dialect=self.dialect,
+                seed_filepath=seed_filepath,
+                output_dir=self.artifacts_dir,
+                store=self.store,
+                templates=self.templates,
+                interactive_gold=interactive_gold,
+            )
+        finally:
+            EngineConfig.VERBOSE = prev_verbose_state
+
+    def run_qsim(
+        self,
+        num_intents: int = 20,
+        num_questions: int = 100,
+        seed: int | None = None,
+        debug: bool | None = None,
+        verbose: bool | None = None,
+    ) -> QSimSummary:
+        """Generate synthetic NL questions from schema-derived intent skeletons.
+
+        Args:
+
+            num_intents: Number of distinct intent types to generate.
+
+            num_questions: Total NL question variants to produce.
+
+            seed: Random seed for reproducible generation.
+
+            debug: Override instance debug flag for this run.
+
+            verbose: Override instance verbose flag for this run.
+
+        Returns:
+
+            ``QSimSummary`` describing the generation results.
+
+        Raises:
+
+            RuntimeError: If mode is not ``"qsim"`` or no API key is set.
+
+            ValueError: If *num_intents* or *num_questions* are out of range.
+        """
+        if self._mode != "qsim":
+            raise RuntimeError(f"Current mode is '{self._mode}'. Call set_mode('qsim') first.")
+        self._ensure_api_key()
+
+        self._validate_num_intents(num_intents)
+        self._validate_num_questions(num_questions)
+
+        debug = self._debug if debug is None else debug
+        verbose = self._verbose if verbose is None else verbose
+
+        EngineConfig.DEBUG = debug
+        EngineConfig._propagate_debug()
+
+        prev_verbose_state = EngineConfig.VERBOSE
+        EngineConfig.VERBOSE = verbose
+
+        try:
+            return qsim_run_once(
+                num_intents=num_intents,
+                num_questions=num_questions,
+                seed=seed,
+                artifacts_dir=self.artifacts_dir,
+                schema=self._schema,
+            )
+        finally:
+            EngineConfig.VERBOSE = prev_verbose_state
+
+    def get_qsim(self) -> list[QSimSummary]:
+        """Return all QSim run summaries oldest-first."""
+        return get_qsim(self.artifacts_dir)
+
+    def get_questions_only(self, timestamp_or_result: str | QSimSummary, output_path: str | None = None) -> None:
+        """Print and save NL questions from a QSim run.
+
+        Args:
+
+            timestamp_or_result: Timestamp string or ``QSimSummary`` identifying the run.
+
+            output_path: Output file path; defaults to ``qsim_<timestamp>_questions.txt``.
+        """
+        if isinstance(timestamp_or_result, QSimSummary):
+            timestamp = timestamp_or_result.timestamp
+        else:
+            timestamp = timestamp_or_result
+
+        path = resolve_qsim_path(timestamp_or_result, self.artifacts_dir)
+        results = load_generated_questions(path)
+
+        if output_path is None:
+            output_path = f"qsim_{timestamp}_questions.txt"
+
+        get_questions_only(results, output_path)
+
+    def get_stats(self) -> dict[str, Any]:
+        """Return template store statistics (counts, schema hash, table count)."""
+        return {
+            "templates_count": len(self.templates),
+            "rejected_count": len(self.rejected),
+            "schema_hash": self.schema.schema_hash,
+            "table_count": len(self.schema.tables),
+        }
+
+    def get_schema_created_at(self) -> str:
+        """Return the ISO-format creation timestamp of the loaded schema."""
+        return self.schema.created_at
+
+    def get_templates(self) -> list[TemplateInfo]:
+        """Return all accepted templates as ``TemplateInfo`` objects."""
+        return get_templates_list(self.templates)
+
+    def get_rejected_templates(self) -> list[RejectedTemplateInfo]:
+        """Return all rejected templates as ``RejectedTemplateInfo`` objects."""
+        return get_rejected_templates_list(self.rejected)
+
+    def get_schema_tables(self) -> list[str]:
+        """Return the list of table names in the loaded schema."""
+        return list(self.schema.tables.keys())
+
+    def get_simulator_summary(self, version: int) -> SimulatorSummary:
+        """Return the ``SimulatorSummary`` for a specific run version.
+
+        Args:
+
+            version: Simulator run version number.
+        """
+        return get_simulator_summary_from_dir(self.artifacts_dir, version)