nl2sql-engine 0.1.0__py3-none-any.whl

nl2sql/indexing/vector_store.py

@@ -0,0 +1,422 @@
+from __future__ import annotations
+
+from typing import List, Optional, Dict, Any
+
+from langchain_chroma import Chroma
+from langchain_core.documents import Document
+from langchain_core.embeddings import Embeddings
+
+from nl2sql.indexing.embeddings import (
+    PROVIDER_BY_DIMENSION,
+    EmbeddingService,
+    describe_embeddings,
+)
+from nl2sql.common.exceptions import NL2SQLError
+from nl2sql.common.logger import get_logger
+from .models import BaseChunk
+
+logger = get_logger(__name__)
+
+
+class EmbeddingDimensionMismatchError(NL2SQLError):
+    """Raised when a persisted collection was built with a different embedder."""
+
+
+def check_embedding_dimension_compatibility(
+    persisted_dimension: Optional[int],
+    embeddings: Embeddings,
+    collection_name: str,
+) -> None:
+    """
+    Verifies that a persisted collection can be queried with the current embedder.
+
+    Vectors from different embedding providers have different dimensionality
+    (OpenAI ``text-embedding-3-small`` is 1536, the local model is 384), so an
+    index built with one provider cannot be read with the other.
+
+    Args:
+        persisted_dimension: Dimensionality of the persisted vectors, or None
+            when the collection is empty or the dimension cannot be determined.
+        embeddings: Embedding implementation configured for this process.
+        collection_name: Name of the Chroma collection.
+
+    Raises:
+        EmbeddingDimensionMismatchError: If the dimensions disagree.
+    """
+    provider, model, expected_dimension = describe_embeddings(embeddings)
+
+    if persisted_dimension is None or expected_dimension is None:
+        return
+    if persisted_dimension == expected_dimension:
+        return
+
+    built_with = PROVIDER_BY_DIMENSION.get(persisted_dimension)
+    origin = (
+        f"the '{built_with}' embedding provider"
+        if built_with
+        else "a different embedding provider"
+    )
+
+    raise EmbeddingDimensionMismatchError(
+        f"Vector store collection '{collection_name}' was indexed with {origin} "
+        f"({persisted_dimension}-dimensional vectors), but the configured provider "
+        f"is '{provider}' ({model}, {expected_dimension}-dimensional vectors). "
+        "A vector index cannot be shared across embedding providers: re-index with "
+        "'nl2sql index' after changing EMBEDDING_PROVIDER, or point VECTOR_STORE at "
+        "a separate directory per provider."
+    )
+
+
+class VectorStore:
+    """
+    Vector store for NL2SQL orchestration.
+
+    This store indexes schema chunks and provides staged retrieval
+    for datasource routing, schema grounding, and planning context.
+    """
+
+    def __init__(
+        self,
+        collection_name: str,
+        persist_directory: str,
+        embeddings: Optional[Embeddings] = None,
+    ):
+        """
+        Initializes the vector store.
+
+        Args:
+            collection_name: Name of the Chroma collection.
+            persist_directory: Directory used for persistence.
+            embeddings: Embedding implementation to use.
+        """
+        self.collection_name = collection_name
+        self.embeddings = embeddings or EmbeddingService.get_embeddings()
+        self.persist_directory = persist_directory
+        self._initialize_vector_store()
+
+    def _initialize_vector_store(self) -> None:
+        """
+        Initializes the underlying Chroma vector store.
+        """
+        self.vectorstore = Chroma(
+            collection_name=self.collection_name,
+            embedding_function=self.embeddings,
+            persist_directory=self.persist_directory,
+        )
+        self._dimension_checked = False
+
+    def _verify_embedding_dimensions(self) -> None:
+        """
+        Checks the persisted vectors against the configured embedder once per
+        store instance.
+
+        The check runs on the read path only: indexing clears the collection
+        before writing, so a provider switch is fixed by re-running
+        ``nl2sql index`` rather than blocked by it.
+        """
+        if self._dimension_checked:
+            return
+        check_embedding_dimension_compatibility(
+            self._persisted_dimension(),
+            self.embeddings,
+            self.collection_name,
+        )
+        self._dimension_checked = True
+
+    def _persisted_dimension(self) -> Optional[int]:
+        """
+        Reads the dimensionality of the vectors already persisted in the collection.
+
+        Returns:
+            Vector dimensionality, or None when the collection is empty or the
+            dimension cannot be determined.
+        """
+        try:
+            stored = self.vectorstore._collection.peek(limit=1).get("embeddings")
+            if stored is None or len(stored) == 0:
+                return None
+            return len(stored[0])
+        except Exception as exc:
+            logger.debug(f"Could not determine persisted embedding dimension: {exc}")
+            return None
+
+    def initialize_if_not_exists(self) -> None:
+        """
+        Initializes the vector store if it does not exist.
+        """
+        try:
+            _ = self.vectorstore._collection.count()
+        except Exception:
+            logger.info("Vector store not found, initializing new store.")
+            self._initialize_vector_store()
+
+        self._verify_embedding_dimensions()
+
+    def is_empty(self) -> bool:
+        """
+        Checks whether the vector store is empty.
+
+        Returns:
+            True if the store contains no documents.
+        """
+        try:
+            return self.vectorstore._collection.count() == 0
+        except Exception as exc:
+            logger.error(f"Failed to check vector store state: {exc}")
+            return True
+
+    def clear(self) -> None:
+        """
+        Deletes the entire vector collection.
+        """
+        try:
+            self.vectorstore.delete_collection()
+            self._initialize_vector_store()
+        except Exception as exc:
+            logger.error(f"Failed to clear vector store: {exc}")
+
+    def delete_documents(self, filter: Dict[str, Any]) -> None:
+        """
+        Deletes documents matching a metadata filter.
+
+        Args:
+            filter: Metadata filter used for deletion.
+        """
+        try:
+            where = (
+                {"$and": [{k: v} for k, v in filter.items()]}
+                if len(filter) > 1
+                else filter
+            )
+            self.vectorstore._collection.delete(where=where)
+        except Exception as exc:
+            logger.error(f"Failed to delete documents: {exc}")
+
+    def refresh_schema_chunks(
+        self,
+        datasource_id: str,
+        schema_version: str,
+        chunks: List[BaseChunk],
+        evicted_versions: List[str],
+    ) -> Dict[str, int]:
+        """
+        Indexes schema chunks for a datasource and evicts old versions.
+
+        Args:
+            datasource_id: Datasource identifier.
+            schema_version: Active schema version.
+            chunks: Schema chunks to index.
+            evicted_versions: Schema versions to remove.
+
+        Returns:
+            Indexing statistics by chunk type.
+        """
+        self._delete_evicted_versions(datasource_id, evicted_versions)
+
+        self.delete_documents(
+            {
+                "datasource_id": datasource_id,
+                "schema_version": schema_version,
+            }
+        )
+
+        documents = self._prepare_chunk_documents(chunks)
+
+        if documents:
+            self.vectorstore.add_documents(documents)
+
+        stats: Dict[str, Any] = {}
+        stats["datasource_id"] = datasource_id
+        stats["schema_version"] = schema_version
+        for chunk in chunks:
+            stats[chunk.type] = stats.get(chunk.type, 0) + 1
+
+        return stats
+
+    def _delete_evicted_versions(
+        self,
+        datasource_id: str,
+        evicted_versions: List[str],
+    ) -> None:
+        """
+        Deletes all documents belonging to evicted schema versions.
+
+        Args:
+            datasource_id: Datasource identifier.
+            evicted_versions: Schema versions to remove.
+        """
+        for version in evicted_versions:
+            self.delete_documents(
+                {
+                    "datasource_id": datasource_id,
+                    "schema_version": version,
+                }
+            )
+
+    def _prepare_chunk_documents(
+        self,
+        chunks: List[BaseChunk],
+    ) -> List[Document]:
+        """
+        Converts schema chunks into vector documents.
+
+        Args:
+            chunks: Schema chunks to convert.
+
+        Returns:
+            List of vector documents.
+        """
+        return [
+            Document(
+                page_content=chunk.get_page_content(),
+                metadata=chunk.get_metadata(),
+            )
+            for chunk in chunks
+        ]
+
+    def retrieve_datasource_candidates(
+        self,
+        query: str,
+        k: int = 3,
+    ) -> List[Document]:
+        """
+        Retrieves candidate datasources for a user query.
+
+        Args:
+            query: User query.
+            k: Number of datasource candidates to retrieve.
+
+        Returns:
+            Retrieved datasource documents.
+        """
+        self.initialize_if_not_exists()
+        from nl2sql.common.resilience import VECTOR_BREAKER
+
+        @VECTOR_BREAKER
+        def _execute():
+            return self.vectorstore.max_marginal_relevance_search(
+                query,
+                k=k,
+                fetch_k=k * 4,
+                lambda_mult=0.7,
+                filter={"type": "schema.datasource"},
+            )
+
+        return _execute()
+
+    def retrieve_schema_context(
+        self,
+        query: str,
+        datasource_id: str,
+        k: int = 8,
+    ) -> List[Document]:
+        """
+        Retrieves schema-level context for a datasource.
+
+        Args:
+            query: User query.
+            datasource_id: Selected datasource identifier.
+            k: Number of schema documents to retrieve.
+
+        Returns:
+            Retrieved schema documents.
+        """
+        self.initialize_if_not_exists()
+        from nl2sql.common.resilience import VECTOR_BREAKER
+
+
+        @VECTOR_BREAKER
+        def _execute():
+            return self.vectorstore.max_marginal_relevance_search(
+                query,
+                k=k,
+                fetch_k=k * 4,
+                lambda_mult=0.7,
+                filter={
+                    "$and": [
+                        {"datasource_id": datasource_id},
+                        {"type": {"$in": ["schema.table", "schema.metric"]}},
+                    ]
+                },
+            )
+
+        return _execute()
+
+    def retrieve_column_candidates(
+        self,
+        query: str,
+        datasource_id: str,
+        k: int = 8,
+    ) -> List[Document]:
+        """
+        Retrieves candidate column documents for a datasource.
+
+        Args:
+            query: User query.
+            datasource_id: Selected datasource identifier.
+            k: Number of column documents to retrieve.
+
+        Returns:
+            Retrieved column documents.
+        """
+        from nl2sql.common.resilience import VECTOR_BREAKER
+
+        self.initialize_if_not_exists()
+
+        @VECTOR_BREAKER
+        def _execute():
+            return self.vectorstore.max_marginal_relevance_search(
+                query,
+                k=k,
+                fetch_k=k * 4,
+                lambda_mult=0.7,
+                filter={
+                    "$and": [
+                        {"datasource_id": datasource_id},
+                        {"type": "schema.column"},
+                    ]
+                },
+            )
+
+        return _execute()
+
+    def retrieve_planning_context(
+        self,
+        query: str,
+        datasource_id: str,
+        tables: List[str],
+        k: int = 12,
+    ) -> List[Document]:
+        """
+        Retrieves planning-level context for selected tables.
+
+        Args:
+            query: User query.
+            datasource_id: Selected datasource identifier.
+            tables: Fully qualified table names.
+            k: Number of planning documents to retrieve.
+
+        Returns:
+            Retrieved planning documents.
+        """
+        from nl2sql.common.resilience import VECTOR_BREAKER
+
+        self.initialize_if_not_exists()
+    
+        @VECTOR_BREAKER
+        def _execute():
+            return self.vectorstore.max_marginal_relevance_search(
+                query,
+                k=k,
+                fetch_k=k * 4,
+                lambda_mult=0.7,
+                filter={
+                    "$and": [
+                        {"datasource_id": datasource_id},
+                        {"type": {"$in": ["schema.column", "schema.relationship"]}},
+                        {"table": {"$in": tables}},
+                    ]
+                },
+            )
+
+        return _execute()

nl2sql/llm/__init__.py

@@ -0,0 +1,8 @@
+
+from .registry import LLMRegistry
+from .models import AgentConfig
+
+__all__ = [
+    "LLMRegistry",
+    "AgentConfig",
+]

nl2sql/llm/models.py

@@ -0,0 +1,10 @@
+"""Re-export of the canonical ``AgentConfig``.
+
+The model lives in :mod:`nl2sql.configs.llm`, which owns the on-disk file
+schemas. This module keeps the ``nl2sql.llm.models`` import path working for
+``LLMRegistry`` and existing callers.
+"""
+
+from nl2sql.configs.llm import AgentConfig
+
+__all__ = ["AgentConfig"]

nl2sql/llm/registry.py

@@ -0,0 +1,214 @@
+import os
+from threading import RLock
+from typing import Any, Dict, NamedTuple, Optional
+
+from langchain_openai import ChatOpenAI
+
+from nl2sql.secrets import SecretManager
+from .models import AgentConfig
+
+
+class ProviderPreset(NamedTuple):
+    """Endpoint and credential defaults for one OpenAI-compatible provider.
+
+    Attributes:
+        base_url: Endpoint the provider is reached on, or None to let the
+            OpenAI client resolve its own default.
+        api_key_env: Environment variable named in errors and used as the
+            last-resort source of the key.
+        api_key_placeholder: Stand-in key for providers that authenticate
+            nothing. Non-None means the provider needs no real credential.
+    """
+
+    base_url: Optional[str]
+    api_key_env: Optional[str]
+    api_key_placeholder: Optional[str] = None
+
+
+# OpenAI, OpenRouter and Ollama all speak the OpenAI wire protocol, so a single
+# ChatOpenAI client serves all three; only the endpoint differs. A
+# config-supplied ``base_url`` overrides the preset, which is what lets the same
+# path serve vLLM, LiteLLM or any other OpenAI-compatible endpoint.
+#
+# Ollama needs no credential, but ChatOpenAI refuses to construct without an
+# ``api_key`` (``openai.OpenAIError: Missing credentials``), so its preset
+# supplies a placeholder the local daemon ignores.
+PROVIDER_PRESETS: Dict[str, ProviderPreset] = {
+    "openai": ProviderPreset(
+        base_url=None,
+        api_key_env="OPENAI_API_KEY",
+    ),
+    "openrouter": ProviderPreset(
+        base_url="https://openrouter.ai/api/v1",
+        api_key_env="OPENROUTER_API_KEY",
+    ),
+    "ollama": ProviderPreset(
+        base_url="http://localhost:11434/v1",
+        api_key_env=None,
+        api_key_placeholder="ollama",
+    ),
+}
+
+
+class LLMRegistry:
+
+    def __init__(self, secret_manager: SecretManager):
+        self.secret_manager = secret_manager
+        self.llms = {}
+        self._configs: Dict[str, AgentConfig] = {}
+        self._lock = RLock()
+
+    def register_llms(self, config: Dict[str, AgentConfig]):
+        for agent in config.values():
+            self.register_llm(agent)
+
+    def register_llm(self, agent: AgentConfig):
+        """Validates an agent's configuration and records it for later use.
+
+        The client itself is built on first ``get_llm`` so that constructing a
+        context does not require credentials for every configured agent. What is
+        genuinely misconfiguration - an unknown provider, a missing model - still
+        fails here rather than mid-query.
+
+        Args:
+            agent: Configuration for one agent.
+
+        Raises:
+            ValueError: If the provider is unknown or the model is empty.
+        """
+        if agent.provider not in PROVIDER_PRESETS:
+            raise ValueError(
+                f"Unsupported LLM provider: {agent.provider}. "
+                f"Valid providers are: {', '.join(sorted(PROVIDER_PRESETS))}."
+            )
+        if not agent.model or not agent.model.strip():
+            raise ValueError(
+                f"LLM agent '{agent.name}' has no model configured. "
+                "Set 'model' for it in configs/llm.yaml."
+            )
+
+        with self._lock:
+            self._configs[agent.name] = agent
+            # A re-registration replaces any client built from the old config.
+            self.llms.pop(agent.name, None)
+
+    def get_llm(self, name: str) -> ChatOpenAI:
+        """Returns the client for an agent, building it on first use.
+
+        Args:
+            name: Agent name; falls back to the 'default' agent.
+
+        Returns:
+            ChatOpenAI: The cached client for that agent.
+
+        Raises:
+            ValueError: If neither the named agent nor a 'default' agent is
+                registered, or if the provider needs an API key that cannot be
+                resolved.
+        """
+        with self._lock:
+            if name in self.llms:
+                return self.llms[name]
+
+            config = self._configs.get(name) or self._configs.get("default")
+            if config is None:
+                raise ValueError(
+                    f"No LLM named '{name}' is configured and no 'default' LLM has "
+                    "been registered. Add it to configs/llm.yaml (under 'agents', or "
+                    "as the 'default' agent)."
+                )
+
+            if config.name in self.llms:
+                return self.llms[config.name]
+
+            client = self._build_client(config)
+            self.llms[config.name] = client
+            return client
+
+    def _build_client(self, agent: AgentConfig) -> ChatOpenAI:
+        """Builds the ChatOpenAI client for one agent from its provider preset.
+
+        Args:
+            agent: Validated configuration for the agent.
+
+        Returns:
+            ChatOpenAI: A client pointed at the configured endpoint.
+        """
+        preset = PROVIDER_PRESETS[agent.provider]
+        api_key = self._resolve_api_key(agent, preset)
+
+        base_url = agent.base_url or preset.base_url
+        kwargs = {"base_url": base_url} if base_url else {}
+
+        return ChatOpenAI(
+            model=agent.model,
+            api_key=api_key,
+            temperature=agent.temperature,
+            tags=[agent.name],
+            seed=42,
+            **kwargs,
+        )
+
+    def _resolve_api_key(self, agent: AgentConfig, preset: ProviderPreset):
+        """Resolves the API key for an agent, or explains what is missing.
+
+        Args:
+            agent: Configuration for the agent.
+            preset: Preset for the agent's provider.
+
+        Returns:
+            The resolved key, the provider's placeholder for key-free providers,
+            or the value of the provider's environment variable.
+
+        Raises:
+            ValueError: If the provider requires a key and none can be found.
+        """
+        try:
+            resolved = self.secret_manager.resolve_object(agent.api_key)
+        except ValueError:
+            # An unresolvable "${env:...}" reference is the same situation as no
+            # key at all, and is reported as such below.
+            resolved = None
+
+        if resolved is not None and resolved.get_secret_value():
+            return resolved
+
+        if preset.api_key_placeholder:
+            return preset.api_key_placeholder
+
+        env_var = self._api_key_env_var(agent, preset)
+        from_env = os.environ.get(env_var) if env_var else None
+        if from_env:
+            return from_env
+
+        raise ValueError(
+            f"LLM agent '{agent.name}' uses provider '{agent.provider}', which "
+            f"requires an API key, but none could be resolved. Set {env_var} in "
+            "the environment, or give the agent an 'api_key' in configs/llm.yaml."
+        )
+
+    @staticmethod
+    def _api_key_env_var(agent: AgentConfig, preset: ProviderPreset) -> Optional[str]:
+        """Returns the environment variable the agent's key should come from.
+
+        A config that already says ``${env:SOME_VAR}`` names its own variable;
+        anything else falls back to the provider's conventional one.
+        """
+        raw = agent.api_key.get_secret_value() if agent.api_key else ""
+        if raw and raw.startswith("${env:") and raw.endswith("}"):
+            return raw[len("${env:") : -1]
+        return preset.api_key_env
+
+    def get_llm_config(self, name: str) -> Dict[str, Any]:
+        with self._lock:
+            if name not in self._configs:
+                name = "default"
+            config = self._configs[name]
+            return config.model_dump(exclude={"api_key"})
+
+    def list_llms(self) -> Dict[str, Dict[str, Any]]:
+        with self._lock:
+            return {
+                name: config.model_dump(exclude={"api_key"})
+                for name, config in self._configs.items()
+            }

nl2sql/pipeline/__init__.py

@@ -0,0 +1 @@
+"""Core pipeline module containing the graph definition and state management."""