Metadata-Version: 2.4
Name: pgvector-template
Version: 0.6.1
Summary: Template library for flexible PGVector RAG implementations
Author-email: DL <v49t9zpqd@mozmail.com>
License: MIT
Project-URL: Homepage, https://github.com/DavidLiuGit/PGVector-Template
Classifier: Programming Language :: Python :: 3
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: pgvector>=0.2.0
Requires-Dist: pydantic<3.0,>=2.11
Requires-Dist: sqlalchemy>=2.0.0
Requires-Dist: typing-extensions>=4.0.0
Provides-Extra: test
Requires-Dist: psycopg[binary]>=3.1.0; extra == "test"
Requires-Dist: pytest>=7.0.0; extra == "test"
Requires-Dist: pytest-cov>=4.0.0; extra == "test"
Requires-Dist: python-dotenv>=1.0.0; extra == "test"
Provides-Extra: dev
Requires-Dist: ruff>=0.9.0; extra == "dev"
Requires-Dist: ty>=0.0.1a0; extra == "dev"
Provides-Extra: dist
Requires-Dist: build>=1.2.2; extra == "dist"
Requires-Dist: twine>=6.1.0; extra == "dist"
Dynamic: license-file

# 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

```python
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

```bash
pip install pgvector-template
```

### With Database Driver

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

```bash
# 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:**

```sql
CREATE EXTENSION IF NOT EXISTS vector;
```

3. **Set up your connection:**

```python
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

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

## Usage Examples

### 1. Define Your Document Model

```python
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

```python
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:

```python
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

```python
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

```python
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:
```bash
pip install -e .[test]
```

### Unit Tests
```bash
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`:

```bash
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

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

## License

MIT License - see [LICENSE](LICENSE) file for details.

## Links

- [GitHub Repository](https://github.com/DavidLiuGit/PGVector-Template)
- [PGVector Documentation](https://github.com/pgvector/pgvector)
