PGVector-Template

A flexible, production-ready template library for building Retrieval-Augmented Generation (RAG) applications using PostgreSQL with PGVector extensions.

Overview

PGVector-Template provides a robust foundation for implementing vector-based document storage and retrieval systems. It offers a clean abstraction layer over PostgreSQL's PGVector extension, making it easy to build scalable RAG applications with proper document management, metadata handling, and efficient vector search capabilities.

Quick Start

from pgvector_template import DocumentDatabaseManager
from pgvector_template.core.document import BaseDocument
from pgvector_template.models.search import SearchQuery
from pgvector_template.service.document_service import DocumentService, DocumentServiceConfig


# 1. Define your document model
class MyDocument(BaseDocument):
    __tablename__ = "my_documents"


# 2. Set up the database: connects, enables pgvector, creates the schema and tables
db = DocumentDatabaseManager(
    database_url="postgresql://user:pass@localhost/mydb",
    schema_suffix="my_kb",
    document_classes=[MyDocument],
)
db.setup()

# 3. Create a service. `my_embedding_provider` is your `BaseEmbeddingProvider` subclass
config = DocumentServiceConfig(document_cls=MyDocument, embedding_provider=my_embedding_provider)

with db.get_session() as session:
    service = DocumentService(session, config)

    # 4. Insert a corpus: it is split into chunks, embedded and stored
    service.corpus_manager.insert_corpus("Your document content here", {}, corpus_id="doc-1")

    # 5. Search
    results = service.search_client.search(SearchQuery(text="your query", limit=5))

Key Concepts

Corpus vs Document: Understanding the hierarchy is essential:

• Corpus: A complete source document (e.g., a full PDF, article, or book)
• Document: A chunk or segment of a corpus that fits within embedding limits
• Collection: A logical grouping of related corpora (e.g., "legal_docs", "user_manuals")

Example: A 50-page PDF (corpus) might be split into 200 documents (chunks), all sharing the same corpus_id but with different chunk_index values.

Key Features

• Flexible Document Model: Abstract base classes for customizable document schemas
• Vector Search: Optimized HNSW indexing for fast similarity search
• Metadata Management: JSON-based flexible metadata with GIN indexing
• Collection Support: Organize documents into logical collections
• Chunk Management: Handle long content by chunking into retrievable documents
• Smart Upserts: Re-inserting a corpus only re-embeds chunks whose content changed, and skips the write entirely if nothing did
• Database Abstraction: Clean SQLAlchemy-based database layer with schema creation API
• Type Safety: Full Pydantic validation and type hints
• Production Ready: Comprehensive testing and error handling

Architecture

The library is organized into several key components:

• Core: Document models, embedders, search functionality
• Database: Connection management and document database operations
• Service: High-level document service layer
• Types: Shared type definitions and schemas

Installation

Basic Installation

pip install pgvector-template

With Database Driver

For production use, you'll also need a PostgreSQL driver:

# For binary driver (recommended)
pip install pgvector-template psycopg[binary]

# Or for source driver
pip install pgvector-template psycopg

Prerequisites

• Python 3.11+
• PostgreSQL 12+ with PGVector extension
• For development: Additional test dependencies

Configuration

Database Setup

1. Install PostgreSQL with PGVector extension
2. Create your database and enable the vector extension:

CREATE EXTENSION IF NOT EXISTS vector;

3. Set up your connection:

from pgvector_template import DatabaseManager

db_manager = DatabaseManager("postgresql://user:password@localhost:5432/mydb")
db_manager.initialize()  # connects and enables the vector extension

create_tables also adds nullable model columns that are missing from existing tables (for example content_hash). Only the column name and type are applied: indexes, foreign keys, unique constraints and server defaults are not. This is a stopgap until proper migrations are added.

Environment Variables

# Development
DATABASE_URL=postgresql://user:password@localhost:5432/dev_db

Usage Examples

1. Define Your Document Model

from pgvector_template.core.document import BaseDocument, BaseDocumentMetadata
from pydantic import Field


class MyDocumentMetadata(BaseDocumentMetadata):
    source_type: str = Field(..., description="Type of source document")
    author: str = Field(default="unknown", description="Document author")


class MyDocument(BaseDocument):
    __tablename__ = "my_documents"

2. Insert a Corpus

from pgvector_template.core.document import BaseDocumentOptionalProps

optional_props = BaseDocumentOptionalProps(
    title="Chapter 1: Introduction",
    collection="textbooks",
    language="en",
)

with db.get_session() as session:
    service = DocumentService(session, config)
    result = service.corpus_manager.insert_corpus(
        "This is the document content...",
        {"source_type": "pdf", "author": "John Doe"},  # corpus metadata
        optional_props,
        corpus_id="intro-chapter",
    )

To insert chunks you have already split and embedded, use insert_documents().

Calling insert_corpus() again with the same corpus_id replaces the corpus. Unchanged chunks reuse their stored vector, and the write is skipped if nothing changed:

result = corpus_manager.insert_corpus(text, {"source": "pdf"}, corpus_id="handbook")
# UpsertResult(corpus_hash=..., skipped=False, embedded=12, reused=0, total=12)

result = corpus_manager.insert_corpus(edited_text, {"source": "pdf"}, corpus_id="handbook")
# UpsertResult(corpus_hash=..., skipped=False, embedded=1, reused=11, total=12)

Rows written before content_hash existed have no hash, so their first upsert re-embeds them.

3. Search Documents

from pgvector_template.models.search import MetadataFilter, SearchQuery

# Semantic search
results = service.search_client.search(SearchQuery(text="how do I reset my password", limit=10))

# Keyword search combined with a metadata filter
results = service.search_client.search(
    SearchQuery(
        keywords=["password", "reset"],
        metadata_filters=[MetadataFilter(field_name="source_type", condition="eq", value="pdf")],
        limit=10,
    )
)

for r in results:
    print(r.score, r.document.content)

4. Read a Corpus Back

corpus = service.corpus_manager.get_full_corpus("intro-chapter")
print(corpus.content, corpus.metadata, len(corpus.documents))

Concept Reference

Core Classes

• BaseDocument: Abstract SQLAlchemy model for documents with vector embeddings
• BaseDocumentOptionalProps: Pydantic model for optional document properties
• BaseDocumentMetadata: Base schema for structured document metadata
• BaseEmbeddingProvider: Interface to implement for your embedding model
• BaseCorpusManager: Splits, embeds and upserts corpora (insert_corpus, get_full_corpus)
• BaseSearchClient: Semantic, keyword and metadata search (search)
• DocumentService: Bundles a corpus manager and a search client over one session
• DatabaseManager / DocumentDatabaseManager: Connection, schema and table setup

Testing

Install dependencies (preferably in a virtualenv) before running tests:

pip install -e .[test]

Unit Tests

python -m unittest

Integration Tests

Integration tests require a PostgreSQL database with PGVector extension. Set up your test database and configure the connection in integ-tests/.env:

python -m unittest discover -s integ-tests

Contributing

1. Fork the repository
2. Create a feature branch
3. Make your changes with tests
4. Run the test suite
5. Submit a pull request

Development Setup

pip install -e .[dev,test]
ruff check . && ruff format .
ty check .

License

MIT License - see LICENSE file for details.

Links

• GitHub Repository
• PGVector Documentation