pgvector-template 0.4

Template library for flexible PGVector RAG implementations

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 DatabaseManager, DocumentDatabaseManager
from pgvector_template.core.document import BaseDocument
from sqlalchemy import create_engine
from uuid import uuid4

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

# 2. Set up database connection
engine = create_engine("postgresql://user:pass@localhost/mydb")
db_manager = DatabaseManager(engine)
db_manager.create_tables([MyDocument])

# 3. Create document manager
doc_manager = DocumentDatabaseManager(db_manager)

# 4. Insert a document
with db_manager.get_session() as session:
    doc = MyDocument.from_props(
        corpus_id=uuid4(),
        chunk_index=0,
        content="Your document content here",
        embedding=[0.1, 0.2, 0.3, ...],  # Your embedding vector
    )
    doc_manager.insert_document(session, doc)

# 5. Search similar documents
with db_manager.get_session() as session:
    results = doc_manager.search_similar(
        session, MyDocument, query_embedding=[0.1, 0.2, 0.3, ...], 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
• 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 sqlalchemy import create_engine
from pgvector_template import DatabaseManager

# Option 1: Direct connection string
engine = create_engine("postgresql://user:password@localhost:5432/mydb")
db_manager = DatabaseManager(engine)

# Option 2: From environment variable
import os
engine = create_engine(os.getenv("DATABASE_URL"))
db_manager = DatabaseManager(engine)

Production Configuration

from sqlalchemy import create_engine
from sqlalchemy.pool import QueuePool

engine = create_engine(
    "postgresql://user:password@localhost:5432/mydb",
    poolclass=QueuePool,
    pool_size=10,
    max_overflow=20,
    pool_pre_ping=True,
)

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 Documents

from pgvector_template.core.document import BaseDocumentOptionalProps
from uuid import uuid4

# Create document with metadata
metadata = MyDocumentMetadata(source_type="pdf", author="John Doe")
optional_props = BaseDocumentOptionalProps(
    title="Chapter 1: Introduction",
    collection="textbooks",
    language="en",
    tags=["education", "intro"]
)

doc = MyDocument.from_props(
    corpus_id=uuid4(),
    chunk_index=0,
    content="This is the document content...",
    embedding=your_embedding_vector,  # list[float] with 1024 dimensions
    metadata=metadata.to_dict(),
    optional_props=optional_props
)

with db_manager.get_session() as session:
    doc_manager.insert_document(session, doc)

3. Search Documents

# Basic similarity search
with db_manager.get_session() as session:
    results = doc_manager.search_similar(
        session=session,
        document_cls=MyDocument,
        query_embedding=query_vector,
        limit=10
    )

# Search with filters
from pgvector_template.models.search import MetadataFilter

filters = [
    MetadataFilter(key="source_type", value="pdf"),
    MetadataFilter(key="author", value="John Doe")
]

results = doc_manager.search_similar(
    session=session,
    document_cls=MyDocument,
    query_embedding=query_vector,
    limit=10,
    metadata_filters=filters,
    collection="textbooks"
)

4. Manage Collections

# Get all documents in a collection
with db_manager.get_session() as session:
    docs = doc_manager.get_documents_by_collection(
        session, MyDocument, "textbooks"
    )

# Get all chunks from a corpus
    corpus_docs = doc_manager.get_documents_by_corpus_id(
        session, MyDocument, corpus_id
    )

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
• DatabaseManager: Database connection and session management
• DocumentDatabaseManager: High-level document CRUD operations

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]
black .  # Format code

License

MIT License - see LICENSE file for details.

Links

• GitHub Repository
• PGVector Documentation