Zum Inhalt

store

module wittgenstein_pgvector_qsql.store

QSQLVectorStore — LangChain PGVector wrapper for the Q→SQL 'sql' collection.

Replaces the Vanna 0.7.9 PG_VectorStore usage for training and retrieval of question→SQL pairs. Writes to the langchain_pg_embedding / langchain_pg_collection tables (collection name: 'sql'), exactly where Vanna wrote them.

PGVector creates its schema tables on first use — no migration needed.

Classes

  • QSQLVectorStore — Semantic Q→SQL pair store backed by LangChain PGVector ('sql' collection).

Functions

  • ensure_vector_index — Create an HNSW cosine-distance index on langchain_pg_embedding.embedding.

wittgenstein_pgvector_qsql.store.ensure_vector_index

ensure_vector_index(connection_or_engine: str | Engine) → None

Create an HNSW cosine-distance index on langchain_pg_embedding.embedding.

langchain_postgres.PGVector auto-creates its schema on first use, but EmbeddingStore.__table_args__ (in the installed langchain_postgres package) only defines a GIN index on the JSONB cmetadata column — nothing on embedding itself. Every get_similar call therefore runs an exact, brute-force cosine-distance scan (DistanceStrategy.COSINE, PGVector's default) across the whole table. This helper closes that gap without touching langchain_postgres's own schema-creation code.

SHARED TABLE, NOT JUST THIS LIBRARY'S DATA: langchain_pg_embedding is one physical table shared by every PGVector collection in the database — rows are scoped by collection_id, not by table. Running this once benefits (and adds index-maintenance overhead to) every other feature that stores vectors via PGVector in the same database, not only this library's "sql" collection. That's an intentional trade-off worth knowing about, not a side effect to hide.

Opt-in by design: this is NOT called automatically by QSQLVectorStore. HNSW build time grows with existing row count, so a caller should invoke this once during deliberate setup/migration for an environment rather than pay a surprise index-build cost on first use in a fresh one.

Requires pgvector >= 0.5.0 (the release that introduced the HNSW index type). Safe to call repeatedly — uses IF NOT EXISTS.

Parameters

  • connection_or_engine : str | Engine — either the same psycopg3 SQLAlchemy URL passed to QSQLVectorStore (e.g. 'postgresql+psycopg://user:pass@host/db') or an existing SQLAlchemy Engine pointed at the same database.

wittgenstein_pgvector_qsql.store.QSQLVectorStore

class QSQLVectorStore(connection_string: str, embedder: Embeddings, collection_name: str | None = None, pre_delete_collection: bool = False)

Semantic Q→SQL pair store backed by LangChain PGVector ('sql' collection).

Questions are embedded and stored; SQL is kept in metadata. Retrieval returns the most similar (question, sql) pairs for a given query.

Parameters

  • connection_string : str — psycopg3 URL, e.g. 'postgresql+psycopg://user:pass@host:port/dbname'

  • embedder : Embeddings — any LangChain Embeddings instance (OllamaEmbedder, OpenAIEmbedder, or any compatible class).

  • collection_name : str | None — PGVector collection to read/write. Defaults to COLLECTION_NAME ('sql') — the production collection. langchain_pg_embedding is one physical table shared by every collection (scoped by collection_id), so different models' embeddings must live in different collections: a query embedded with model A run against rows embedded with model B would error at query time (pgvector's <=> requires matching vector widths). Pass a per-model name (e.g. "sql__local__multilingual-e5-base") when comparing embedder choices — see training/reindex_embeddings.py.

  • pre_delete_collection : bool — wipe collection_name (only that collection, not the shared table) before use — makes a full reindex idempotent across re-runs. Never use this against COLLECTION_NAME ('sql'), the production collection.

Methods

  • add_pair — Index a question→SQL pair into the vector store.

  • get_similar — Return the N most similar (question, sql) pairs for a query.

wittgenstein_pgvector_qsql.store.QSQLVectorStore.add_pair

method QSQLVectorStore.add_pair(question: str, sql: str) → None

Index a question→SQL pair into the vector store.

The question text is embedded; the SQL is stored in document metadata so it can be returned alongside the question on retrieval.

wittgenstein_pgvector_qsql.store.QSQLVectorStore.get_similar

method QSQLVectorStore.get_similar(question: str, n: int = 3) → list[dict]

Return the N most similar (question, sql) pairs for a query.

Parameters

  • question : str — the user question to match against.

  • n : int — maximum number of results to return.

Returns

  • list[dict] — list of dicts with keys 'question' and 'sql'.