nl2sql-engine 0.1.0__py3-none-any.whl

nl2sql/indexing/enrichment_service.py

@@ -0,0 +1,316 @@
+from __future__ import annotations
+
+from typing import Dict, List, Optional, Tuple
+
+from pydantic import BaseModel, Field
+from langchain_core.prompts import ChatPromptTemplate
+
+from nl2sql.schema import SchemaSnapshot
+from nl2sql_adapter_sdk.schema import TableMetadata, ColumnMetadata
+from nl2sql.common.logger import get_logger
+
+logger = get_logger("indexing_enrichment")
+
+# Name of the agent in configs/llm.yaml that enrichment calls.
+ENRICHMENT_LLM_NAME = "indexing_enrichment"
+
+
+class DatasourceEnrichment(BaseModel):
+    description: Optional[str] = None
+    domains: List[str] = Field(default_factory=list)
+    sample_questions: List[str] = Field(default_factory=list)
+    citations: List[str] = Field(default_factory=list)
+
+
+class TableEnrichment(BaseModel):
+    description: Optional[str] = None
+    citations: List[str] = Field(default_factory=list)
+
+
+class ColumnEnrichment(BaseModel):
+    description: Optional[str] = None
+    synonyms: List[str] = Field(default_factory=list)
+    citations: List[str] = Field(default_factory=list)
+
+
+class SchemaEnrichment(BaseModel):
+    datasource: DatasourceEnrichment
+    tables: Dict[str, TableEnrichment] = Field(default_factory=dict)
+    columns: Dict[str, ColumnEnrichment] = Field(default_factory=dict)
+
+
+ENRICHMENT_PROMPT = """You are enriching database schema metadata for retrieval.
+Use ONLY the evidence provided. Do not introduce any tables, columns, or facts
+that are not present in evidence. If evidence is insufficient, return empty values.
+
+Return a JSON object that matches this schema:
+SchemaEnrichment = {{
+  "datasource": {{
+    "description": string|null,
+    "domains": [string],
+    "sample_questions": [string],
+    "citations": [string]
+  }},
+  "tables": {{
+    "<table_key>": {{
+      "description": string|null,
+      "citations": [string]
+    }}
+  }},
+  "columns": {{
+    "<table_key>.<column_name>": {{
+      "description": string|null,
+      "synonyms": [string],
+      "citations": [string]
+    }}
+  }}
+}}
+
+Evidence (JSON):
+{evidence_json}
+"""
+
+
+def build_evidence(
+    snapshot: SchemaSnapshot,
+    datasource_description: Optional[str],
+    existing_questions: List[str],
+) -> Dict[str, object]:
+    contract = snapshot.contract
+    metadata = snapshot.metadata
+
+    tables_payload = []
+    relationships_payload = []
+
+    for table_key, table_contract in contract.tables.items():
+        table_md = metadata.tables.get(table_key)
+        columns_payload = []
+        for col_name, col_contract in table_contract.columns.items():
+            col_md = table_md.columns.get(col_name) if table_md else None
+            stats = col_md.statistics.model_dump() if col_md and col_md.statistics else {}
+            columns_payload.append(
+                {
+                    "name": col_name,
+                    "type": col_contract.data_type,
+                    "description": col_md.description if col_md else None,
+                    "statistics": stats,
+                    "sample_values": stats.get("sample_values", []),
+                    "pii": bool(col_md.pii) if col_md else False,
+                }
+            )
+
+        tables_payload.append(
+            {
+                "table_key": table_key,
+                "description": table_md.description if table_md else None,
+                "row_count": table_md.row_count if table_md else None,
+                "columns": columns_payload,
+            }
+        )
+
+        for fk in table_contract.foreign_keys:
+            relationships_payload.append(
+                {
+                    "from_table": table_key,
+                    "to_table": fk.referred_table.full_name,
+                    "from_columns": fk.constrained_columns,
+                    "to_columns": fk.referred_columns,
+                }
+            )
+
+    return {
+        "datasource_id": contract.datasource_id,
+        "engine_type": contract.engine_type,
+        "datasource_description": datasource_description or metadata.description or "",
+        "domains": metadata.domains or [],
+        "existing_questions": existing_questions,
+        "tables": tables_payload,
+        "relationships": relationships_payload,
+    }
+
+
+def _normalize_text(value: Optional[str]) -> Optional[str]:
+    if value is None:
+        return None
+    cleaned = value.strip()
+    return cleaned or None
+
+
+def _dedupe_list(items: List[str]) -> List[str]:
+    seen = set()
+    result = []
+    for item in items:
+        cleaned = item.strip()
+        if not cleaned:
+            continue
+        key = cleaned.lower()
+        if key in seen:
+            continue
+        seen.add(key)
+        result.append(cleaned)
+    return result
+
+
+def _question_mentions_schema(question: str, table_names: List[str], column_names: List[str]) -> bool:
+    text = question.lower()
+    return any(name in text for name in table_names + column_names)
+
+
+def sanitize_enrichment(
+    snapshot: SchemaSnapshot,
+    enrichment: SchemaEnrichment,
+    max_questions: int = 100,
+) -> SchemaEnrichment:
+    contract = snapshot.contract
+    valid_tables = set(contract.tables.keys())
+    valid_columns = set()
+    table_names = []
+    column_names = []
+
+    for table_key, table_contract in contract.tables.items():
+        table_names.append(table_contract.table.table_name.lower())
+        table_names.append(table_key.lower())
+        for col_name in table_contract.columns.keys():
+            valid_columns.add(f"{table_key}.{col_name}")
+            column_names.append(col_name.lower())
+
+    sanitized_tables: Dict[str, TableEnrichment] = {}
+    for table_key, table_enrichment in enrichment.tables.items():
+        if table_key not in valid_tables:
+            continue
+        sanitized_tables[table_key] = TableEnrichment(
+            description=_normalize_text(table_enrichment.description),
+            citations=_dedupe_list(table_enrichment.citations),
+        )
+
+    sanitized_columns: Dict[str, ColumnEnrichment] = {}
+    for column_key, column_enrichment in enrichment.columns.items():
+        if column_key not in valid_columns:
+            continue
+        sanitized_columns[column_key] = ColumnEnrichment(
+            description=_normalize_text(column_enrichment.description),
+            synonyms=_dedupe_list(column_enrichment.synonyms),
+            citations=_dedupe_list(column_enrichment.citations),
+        )
+
+    questions = [
+        q.strip()
+        for q in enrichment.datasource.sample_questions
+        if q and q.strip()
+    ]
+    questions = [
+        q for q in questions if _question_mentions_schema(q, table_names, column_names)
+    ]
+    questions = questions[:max_questions]
+
+    return SchemaEnrichment(
+        datasource=DatasourceEnrichment(
+            description=_normalize_text(enrichment.datasource.description),
+            domains=_dedupe_list(enrichment.datasource.domains),
+            sample_questions=questions,
+            citations=_dedupe_list(enrichment.datasource.citations),
+        ),
+        tables=sanitized_tables,
+        columns=sanitized_columns,
+    )
+
+
+def apply_enrichment(
+    snapshot: SchemaSnapshot,
+    enrichment: SchemaEnrichment,
+) -> SchemaSnapshot:
+    metadata = snapshot.metadata.model_copy(deep=True)
+
+    if enrichment.datasource.description:
+        metadata.description = enrichment.datasource.description
+    if enrichment.datasource.domains:
+        metadata.domains = enrichment.datasource.domains
+
+    for table_key, table_contract in snapshot.contract.tables.items():
+        table_md = metadata.tables.get(table_key)
+        if not table_md:
+            table_md = TableMetadata(table=table_contract.table, columns={})
+            metadata.tables[table_key] = table_md
+
+        table_enrichment = enrichment.tables.get(table_key)
+        if table_enrichment and table_enrichment.description:
+            table_md.description = table_enrichment.description
+
+        for col_name in table_contract.columns.keys():
+            column_md = table_md.columns.get(col_name)
+            if not column_md:
+                column_md = ColumnMetadata()
+                table_md.columns[col_name] = column_md
+
+            column_key = f"{table_key}.{col_name}"
+            column_enrichment = enrichment.columns.get(column_key)
+            if not column_enrichment:
+                continue
+            if column_enrichment.description:
+                column_md.description = column_enrichment.description
+            if column_enrichment.synonyms:
+                column_md.synonyms = column_enrichment.synonyms
+
+    return snapshot.model_copy(update={"metadata": metadata})
+
+
+def enrich_schema_snapshot(
+    snapshot: SchemaSnapshot,
+    llm_registry,
+    datasource_description: Optional[str],
+    existing_questions: List[str],
+    max_questions: int = 100,
+) -> Tuple[SchemaSnapshot, List[str]]:
+    """Adds LLM-written descriptions to a schema snapshot, if an LLM is usable.
+
+    Enrichment is an embellishment, never a prerequisite: indexing has to work
+    without any LLM at all. Every step that can involve the model - resolving
+    the client, binding structured output, the call itself - therefore sits
+    inside one guard, and any failure returns the inputs untouched so the
+    caller indexes an unenriched schema instead of failing.
+
+    Args:
+        snapshot: Schema snapshot to enrich.
+        llm_registry: Registry the ``indexing_enrichment`` client is built from.
+        datasource_description: Operator-supplied description, if any.
+        existing_questions: Example questions already configured.
+        max_questions: Upper bound on the merged question list.
+
+    Returns:
+        The enriched snapshot and merged questions, or the originals unchanged
+        when enrichment could not run.
+    """
+    datasource_id = snapshot.contract.datasource_id
+
+    try:
+        llm = llm_registry.get_llm(ENRICHMENT_LLM_NAME)
+        evidence = build_evidence(snapshot, datasource_description, existing_questions)
+        prompt = ChatPromptTemplate.from_template(ENRICHMENT_PROMPT)
+        chain = prompt | llm.with_structured_output(
+            SchemaEnrichment, method="function_calling"
+        )
+        enrichment = chain.invoke({"evidence_json": evidence})
+    except ValueError as exc:
+        # No usable LLM is configured - typically no API key. Expected on the
+        # key-free path, so it is reported without a stack trace.
+        logger.info(
+            f"Schema enrichment skipped for '{datasource_id}': {exc} "
+            "Indexing continues without LLM-written descriptions."
+        )
+        return snapshot, existing_questions
+    except Exception as exc:
+        # Anything else is a genuine surprise and stays fully visible, but it
+        # still must not take indexing down with it.
+        logger.warning(
+            f"Schema enrichment failed for '{datasource_id}': {exc}. "
+            "Indexing continues without LLM-written descriptions.",
+            exc_info=True,
+        )
+        return snapshot, existing_questions
+
+    sanitized = sanitize_enrichment(snapshot, enrichment, max_questions=max_questions)
+    updated_snapshot = apply_enrichment(snapshot, sanitized)
+
+    merged_questions = _dedupe_list(existing_questions + sanitized.datasource.sample_questions)
+    merged_questions = merged_questions[:max_questions]
+    return updated_snapshot, merged_questions

nl2sql/indexing/models.py

@@ -0,0 +1,209 @@
+from __future__ import annotations
+from typing import List, Dict, Optional, Literal, Any
+from pydantic import BaseModel, Field
+from nl2sql_adapter_sdk.schema import TableRef, ColumnRef
+
+
+class BaseChunk(BaseModel):
+    """
+    Base class for all schema chunks.
+    Chunk IDs MUST be deterministic and stable across runs.
+    """
+    id: str
+    type: str
+
+    def get_page_content(self) -> str:
+        raise NotImplementedError("Each chunk must implement get_page_content()")
+
+    def get_metadata(self) -> Dict[str, Any]:
+        return {
+            "id": self.id,
+            "type": self.type,
+        }
+
+
+class DatasourceChunk(BaseChunk):
+    type: Literal["schema.datasource"] = Field(
+        default="schema.datasource", frozen=True
+    )
+
+    datasource_id: str
+    description: str
+    domains: Optional[List[str]] = None
+    schema_version: str
+    examples: Optional[List[str]] = None
+
+    def get_page_content(self) -> str:
+        domains = ", ".join(self.domains) if self.domains else "N/A"
+        examples = ", ".join(self.examples) if self.examples else "N/A"
+        return (
+            f"Datasource: {self.datasource_id}\n"
+            f"{self.description}\n"
+            f"Domains: {domains}\n"
+            f"Examples: {examples}"
+        )
+
+    def get_metadata(self) -> Dict[str, Any]:
+        return {
+            **super().get_metadata(),
+            "datasource_id": self.datasource_id,
+            "schema_version": self.schema_version,
+        }
+
+
+class TableChunk(BaseChunk):
+    type: Literal["schema.table"] = Field(
+        default="schema.table", frozen=True
+    )
+
+    datasource_id: str
+    table: TableRef
+    description: Optional[str] = None
+    primary_key: List[str] = Field(default_factory=list)
+    columns: List[str] = Field(default_factory=list)
+    foreign_keys: List[str] = Field(
+        default_factory=list,
+        description="Human-readable FK summaries (retrieval only)"
+    )
+    row_count: Optional[int] = None
+    schema_version: str
+
+    def get_page_content(self) -> str:
+        pk = ", ".join(self.primary_key) if self.primary_key else "None"
+        column_list = ", ".join(self.columns) if self.columns else "None"
+        return (
+            f"Table: {self.table.full_name}\n"
+            f"{self.description or ''}\n"
+            f"Primary Key: {pk}\n"
+            f"Columns: {column_list}"
+        )
+
+    def get_metadata(self) -> Dict[str, Any]:
+        return {
+            **super().get_metadata(),
+            "datasource_id": self.datasource_id,
+            "table": self.table.full_name,
+            "row_count": self.row_count,
+            "schema_version": self.schema_version,
+            "description": self.description,
+            "primary_key": ','.join(self.primary_key),
+            "foreign_keys": ','.join(self.foreign_keys),
+        }
+
+
+class ColumnChunk(BaseChunk):
+    type: Literal["schema.column"] = Field(
+        default="schema.column", frozen=True
+    )
+
+    datasource_id: str
+    column: ColumnRef
+    dtype: str
+    description: Optional[str] = None
+    column_stats: Dict[str, Any] = Field(default_factory=dict)
+    synonyms: Optional[List[str]] = None
+    pii: bool = False
+    schema_version: str
+
+    def get_page_content(self) -> str:
+        stats = f"Stats: {self.column_stats}" if self.column_stats else ""
+        synonyms = (
+            f"Synonyms: {', '.join(self.synonyms)}"
+            if self.synonyms
+            else ""
+        )
+        return (
+            f"Table: {self.column.table.full_name}\n"
+            f"Column: {self.column.column_name}\n"
+            f"Type: {self.dtype}\n"
+            f"{self.description or ''}\n"
+            f"{stats}\n"
+            f"{synonyms}"
+        )
+
+    def get_metadata(self) -> Dict[str, Any]:
+        return {
+            **super().get_metadata(),
+            "datasource_id": self.datasource_id,
+            "column": self.column.column_name,
+            "table": self.column.table.full_name,
+            "dtype": self.dtype,
+            "pii": self.pii,
+            "schema_version": self.schema_version,
+            "description": self.description,
+        }
+
+class RelationshipChunk(BaseChunk):
+    type: Literal["schema.relationship"] = Field(
+        default="schema.relationship", frozen=True
+    )
+
+    datasource_id: str
+    from_table: TableRef
+    to_table: TableRef
+
+    from_columns: Optional[List[str]] = None
+    to_columns: Optional[List[str]] = None
+
+    cardinality: Literal[
+        "one-to-one",
+        "one-to-many",
+        "many-to-one",
+        "many-to-many",
+        "unknown"
+    ] = "unknown"
+
+    business_meaning: Optional[str] = None
+    schema_version: str
+
+    def get_page_content(self) -> str:
+        return (
+            f"Relationship between {self.from_table.full_name} "
+            f"and {self.to_table.full_name}.\n"
+            f"{self.business_meaning or ''}"
+        )
+
+    def get_metadata(self) -> Dict[str, Any]:
+        return {
+            **super().get_metadata(),
+            "datasource_id": self.datasource_id,
+            "from_table": self.from_table.full_name,
+            "to_table": self.to_table.full_name,
+            "table": self.from_table.full_name,
+            "from_columns": self.from_columns or [],
+            "to_columns": self.to_columns or [],
+            "cardinality": self.cardinality,
+            "schema_version": self.schema_version,
+        }
+
+
+class MetricChunk(BaseChunk):
+    type: Literal["schema.metric"] = Field(
+        default="schema.metric", frozen=True
+    )
+
+    datasource_id: str
+    name: str
+    definition: Optional[str] = None
+    grain: Optional[str] = None
+    business_meaning: Optional[str] = None
+    owner: Optional[str] = None
+    version: str = "v1"
+    schema_version: str
+
+    def get_page_content(self) -> str:
+        return (
+            f"Metric: {self.name}\n"
+            f"{self.business_meaning or ''}\n"
+            f"Definition: {self.definition or ''}\n"
+            f"Grain: {self.grain or 'N/A'}"
+        )
+
+    def get_metadata(self) -> Dict[str, Any]:
+        return {
+            **super().get_metadata(),
+            "datasource_id": self.datasource_id,
+            "name": self.name,
+            "version": self.version,
+            "schema_version": self.schema_version,
+        }

nl2sql/indexing/orchestrator.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+from typing import Dict, TYPE_CHECKING
+
+from nl2sql.datasources.protocols import DatasourceAdapterProtocol
+from nl2sql.common.logger import get_logger
+from nl2sql.indexing.chunk_builder import SchemaChunkBuilder
+from nl2sql.indexing.enrichment_service import enrich_schema_snapshot
+
+if TYPE_CHECKING:
+    from nl2sql.context import NL2SQLContext
+
+logger = get_logger("indexing_orchestrator")
+
+
+class IndexingOrchestrator:
+    """
+    Orchestrates schema indexing for datasources.
+
+    This class coordinates schema snapshot retrieval, schema version
+    registration, chunk construction, and vector store refresh.
+    """
+
+    def __init__(self, ctx: NL2SQLContext):
+        """
+        Initializes the indexing orchestrator.
+
+        Args:
+            ctx: Initialized NL2SQLContext.
+        """
+        self.vector_store = ctx.vector_store
+        self.schema_store = ctx.schema_store
+        self.config_manager = ctx.config_manager
+        self.llm_registry = ctx.llm_registry
+
+    def clear_store(self) -> None:
+        """
+        Clears the vector store.
+        """
+        self.vector_store.clear()
+
+    def index_datasource(
+        self,
+        adapter: DatasourceAdapterProtocol,
+    ) -> Dict[str, int]:
+        """
+        Indexes schema chunks for a datasource.
+
+        Args:
+            adapter: SQLAlchemy adapter for the datasource.
+
+        Returns:
+            Indexing statistics by chunk type.
+        """
+        schema_snapshot = adapter.fetch_schema_snapshot()
+
+        questions = self.config_manager.get_example_questions(adapter.datasource_id)
+        datasource_description = self.config_manager.get_datasource_description(
+            adapter.datasource_id
+        )
+
+        # Best-effort: enrich_schema_snapshot owns resolving its own LLM and
+        # degrades to the unenriched snapshot when none is usable, so indexing
+        # never depends on an API key being present.
+        schema_snapshot, questions = enrich_schema_snapshot(
+            snapshot=schema_snapshot,
+            llm_registry=self.llm_registry,
+            datasource_description=datasource_description,
+            existing_questions=questions,
+        )
+
+        schema_version, evicted_versions = self.schema_store.register_snapshot(
+            schema_snapshot
+        )
+
+        chunk_builder = SchemaChunkBuilder(
+            ds_id=adapter.datasource_id,
+            schema_snapshot=schema_snapshot,
+            schema_version=schema_version,
+            questions=questions,
+        )
+
+        chunks = chunk_builder.build()
+
+        return self.vector_store.refresh_schema_chunks(
+            datasource_id=adapter.datasource_id,
+            schema_version=schema_version,
+            chunks=chunks,
+            evicted_versions=evicted_versions,
+        )