pgvector-template 0.1.0__tar.gz

pgvector_template-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 David Ge Liu
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

pgvector_template-0.1.0/PKG-INFO

@@ -0,0 +1,44 @@
+Metadata-Version: 2.1
+Name: pgvector-template
+Version: 0.1.0
+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: psycopg[binary]>=3.1.0
+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: pytest>=7.0.0; extra == "test"
+Requires-Dist: pytest-cov>=4.0.0; extra == "test"
+Requires-Dist: python-dotenv>=1.0.0; extra == "test"
+
+# PGVector-Template
+
+Template library for flexible PGVector RAG implementations
+
+
+## Testing
+
+Install dependencies (preferably in a virtualenv) before running tests:
+```bash
+pip install -e .[test]
+```
+
+### Unit tests
+```bash
+python -m unittest
+```
+
+### Integration tests
+```bash
+python -m unittest discover -s integ-tests
+```

pgvector_template-0.1.0/README.md

@@ -0,0 +1,21 @@
+# PGVector-Template
+
+Template library for flexible PGVector RAG implementations
+
+
+## Testing
+
+Install dependencies (preferably in a virtualenv) before running tests:
+```bash
+pip install -e .[test]
+```
+
+### Unit tests
+```bash
+python -m unittest
+```
+
+### Integration tests
+```bash
+python -m unittest discover -s integ-tests
+```

pgvector_template-0.1.0/pgvector_template/__init__.py

@@ -0,0 +1,6 @@
+# Import key classes for easier access
+from pgvector_template.db.connection import DatabaseManager
+from pgvector_template.db.document_db import DocumentDatabaseManager
+
+# Re-export at the top level
+__all__ = ["DatabaseManager", "DocumentDatabaseManager"]

pgvector_template-0.1.0/pgvector_template/core/__init__.py

@@ -0,0 +1,23 @@
+from pgvector_template.core.document import BaseDocument, BaseDocumentMetadata, Corpus, BaseDocumentOptionalProps
+from pgvector_template.core.embedder import BaseEmbeddingProvider
+from pgvector_template.core.manager import BaseCorpusManager, BaseCorpusManagerConfig
+from pgvector_template.core.retriever import RetrievalResult, SearchQuery
+from pgvector_template.core.search import BaseSearchClient
+
+
+__all__ = [
+    ### document
+    "BaseDocumentOptionalProps",
+    "BaseDocument",
+    "BaseDocumentMetadata",
+    "Corpus",
+    ### embedder
+    "BaseEmbeddingProvider",
+    ### manager
+    "BaseCorpusManager",
+    "BaseCorpusManagerConfig",
+    ### search
+    "RetrievalResult",
+    "SearchQuery",
+    "BaseSearchClient",
+]

pgvector_template-0.1.0/pgvector_template/core/document.py

@@ -0,0 +1,171 @@
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any, ClassVar, Type, TypeVar, Annotated
+from uuid import uuid4, UUID
+
+from pydantic import BaseModel, Field, field_validator, model_validator
+from sqlalchemy import (
+    Column,
+    String,
+    Text,
+    DateTime,
+    Boolean,
+    Integer,
+    Float,
+    Index,
+    text,
+)
+from sqlalchemy.orm import declarative_base
+from sqlalchemy.dialects.postgresql import UUID, JSONB
+from pgvector.sqlalchemy import Vector
+
+Base = declarative_base()
+
+
+class BaseDocumentOptionalProps(BaseModel):
+    """Optional properties for document creation"""
+
+    title: str | None = None
+    """Optional title or summary for the document"""
+    collection: str | None = Field(None, max_length=64)
+    """Collection name for grouping documents of the same type"""
+    original_url: str | None = Field(None, max_length=2048)
+    """Optional source URL for the document"""
+    language: str | None = Field("en", pattern=r"^[a-z]{2}(-[A-Z]{2})?$")
+    """Language of the content (ISO 639-1 code), e.g., 'en', 'es', 'zh'"""
+    score: float | None = Field(None, ge=0.0, le=1.0)
+    """Optional score assigned during ingestion (e.g., relevance, confidence)"""
+    tags: list[str] | None = None
+    """List of tags or keywords for filtering, categorization, or faceted search"""
+
+    @field_validator("tags")
+    @classmethod
+    def validate_tags(cls, v):
+        if v is not None:
+            # Ensure all tags are strings and not empty
+            if not all(isinstance(tag, str) and tag.strip() for tag in v):
+                raise ValueError("All tags must be non-empty strings")
+            # Remove duplicates while preserving order
+            return list(dict.fromkeys(v))
+        return v
+
+
+T = TypeVar("T", bound="BaseDocument")
+
+
+class BaseDocument(Base):
+    """
+    Template table for Documents, that works for all collection types.
+    Each row represents a single retrievable document (could be chunk or full doc).
+
+    Glossary:
+    - `corpus` - a full text document, consisting of 1-or-more documents.
+      - `corpus_id` is associated with these entries
+    - `document` - a chunk (or entirety) of an corpus. `id` is associated with these chunks
+    """
+
+    __abstract__ = True
+
+    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid4)
+    """Primary key of the Document table. Represents unique ID of a Document"""
+
+    # Hierarchy: original_id groups chunks from same source
+    collection = Column(String(64), nullable=True)
+    """Collection name. Used for filtering and grouping documents of the same type."""
+    corpus_id = Column(UUID(as_uuid=True), index=True)
+    """An `corpus` is the original, full text that chunks are a part (or all) of"""
+    chunk_index = Column(Integer, default=0)
+    """Index of this chunk within an `corpus`. Starts from 0."""
+
+    # Content
+    content = Column(Text, nullable=False)
+    """String content of the chunk"""
+    title = Column(String(500))
+    """Optional chunk title/summary"""
+    document_metadata = Column(JSONB, nullable=False, default=dict)
+    """Flexible metadata as JSON"""
+    origin_url = Column(String(2048), nullable=True)
+    """Optional source URL"""
+    language = Column(String(10), default="en")
+    """Language of the content (ISO 639-1 code), e.g., 'en', 'es', 'zh'."""
+    score = Column(Float, nullable=True)
+    """Optional score assigned during ingestion (e.g., relevance, confidence)."""
+    tags = Column(JSONB, nullable=True, default=list)
+    """List of tags or keywords for filtering, categorization, or faceted search."""
+
+    # Vector embedding
+    embedding = Column(Vector(1024))
+    """Embedding vector. 1024 dimensions by default. Adjust as-needed."""
+
+    # Audit fields
+    created_at = Column(DateTime, default=datetime.utcnow)
+    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
+    is_deleted = Column(Boolean, default=False)
+    """Entries can be logically marked for deletion before they are permanently deleted."""
+
+    @classmethod
+    def from_props(
+        cls: Type[T],
+        corpus_id: UUID,
+        chunk_index: int,
+        content: str,
+        embedding: list[float],
+        metadata: dict[str, Any] = {},
+        optional_props: BaseDocumentOptionalProps | None = None,
+    ) -> T:
+        """
+        Create a BaseDocument instance from mandatory and optional properties.
+
+        Args:
+            corpus_id: UUID of the corpus this document belongs to
+            chunk_index: Index of this chunk within the corpus
+            content: Text content of the document
+            embedding: Vector embedding of the content
+            optional_props: Optional properties for the document
+
+        Returns:
+            A new BaseDocument instance of the calling class type
+        """
+        if optional_props is None:
+            optional_props = BaseDocumentOptionalProps()
+
+        return cls(
+            corpus_id=corpus_id,
+            chunk_index=chunk_index,
+            content=content,
+            embedding=embedding,
+            title=optional_props.title,
+            document_metadata=metadata,
+            collection=optional_props.collection,
+            origin_url=optional_props.original_url,
+            language=optional_props.language,
+            score=optional_props.score,
+            tags=optional_props.tags,
+        )
+
+    # Index("ix_corpus_chunk", "corpus_id", "chunk_index")
+    # Index("ix_content_trgm", text("content gin_trgm_ops"), postgresql_using="gin")  # For fuzzy text search
+    # Index("ix_metadata_gin", "metadata", postgresql_using="gin")
+
+
+class BaseDocumentMetadata(BaseModel):
+    """Base metadata structure"""
+
+    document_type: str = Field(..., description="Description for type of document, e.g. markdown, pdf, etc")
+    schema_version: str = Field("1.0", description="Schema version for the metadata")
+
+    def to_dict(self) -> dict[str, Any]:
+        return self.model_dump()
+
+
+@dataclass
+class Corpus:
+    """
+    Logical grouping of one or more documents (chunks) belonging to the same original source.
+
+    Typically all documents in a corpus share the same `corpus_id` and are ordered by `chunk_index`.
+    """
+
+    corpus_id: UUID
+    documents: list[BaseDocument]
+    metadata: dict[str, Any]  # e.g. source, tags, etc.

pgvector_template-0.1.0/pgvector_template/core/embedder.py

@@ -0,0 +1,20 @@
+from abc import ABC, abstractmethod
+
+
+class BaseEmbeddingProvider(ABC):
+    """Abstract base for embedding generation"""
+
+    @abstractmethod
+    def embed_text(self, text: str) -> list[float]:
+        """Generate embedding vector for text"""
+        raise NotImplementedError("Subclasses must implement embed_text method")
+
+    @abstractmethod
+    def embed_batch(self, texts: list[str]) -> list[list[float]]:
+        """Generate embeddings for multiple texts efficiently"""
+        raise NotImplementedError("Subclasses must implement embed_batch method")
+
+    @abstractmethod
+    def get_dimensions(self) -> int:
+        """Return embedding vector dimensions count"""
+        raise NotImplementedError("Subclasses must implement get_dimensions method")

pgvector_template-0.1.0/pgvector_template/core/manager.py

@@ -0,0 +1,167 @@
+from abc import ABC, abstractmethod
+from logging import getLogger
+from typing import Any, Optional, Type
+from uuid import UUID, uuid4
+
+from pydantic import BaseModel
+from sqlalchemy.orm import Session
+
+from pgvector_template.core.document import BaseDocument, BaseDocumentMetadata, BaseDocumentOptionalProps
+from pgvector_template.core.embedder import BaseEmbeddingProvider
+
+
+logger = getLogger(__name__)
+
+
+class BaseCorpusManagerConfig(BaseModel):
+    """Base configuration for `Corpus` & `Document` management operations"""
+
+    schema_name: str
+    document_cls: Type[BaseDocument]
+    embedding_provider: BaseEmbeddingProvider
+    document_metadata: Type[BaseDocumentMetadata]
+
+    model_config = {"arbitrary_types_allowed": True}
+
+
+class BaseCorpusManager(ABC):
+    """
+    Template class for `Corpus` & `Document` management operations.
+    Each instance should be able to handle multiple collections, of the same Corpus/Document type.
+    For example, if the document class is Jira tickets, multiple teams should be able to share the same
+    `CorpusManager` implementation, with a slightly different config.
+    """
+
+    @property
+    def config(self) -> BaseCorpusManagerConfig:
+        return self._cfg
+
+    def __init__(
+        self,
+        session: Session,
+        config: BaseCorpusManagerConfig,
+    ) -> None:
+        self.session = session
+        self._cfg = config
+        self.schema_name = config.schema_name
+
+    def get_full_corpus(self, corpus_id: str, chunk_delimiter: str = "\n") -> Optional[dict[str, Any]]:
+        """Reconstruct full corpus from its individual documents/chunks"""
+        chunks = (
+            self.session.query(BaseDocument)
+            .filter(BaseDocument.corpus_id == corpus_id, BaseDocument.is_deleted == False)
+            .order_by(BaseDocument.chunk_index)
+            .all()
+        )
+
+        if not chunks:
+            return None
+
+        # Full document is chunk_index = 0, or reconstruct from chunks
+        full_doc = next((c for c in chunks if c.chunk_index == 0), None)
+        if full_doc:
+            return {
+                "id": full_doc.original_id,
+                "content": full_doc.content,
+                "metadata": full_doc.document_metadata,
+                "chunks": [{"id": c.id, "index": c.chunk_index, "title": c.title} for c in chunks if c.chunk_index > 0],
+            }
+
+        # Reconstruct from chunks
+        reconstructed_content = chunk_delimiter.join([c.content for c in chunks])
+        return {
+            "id": corpus_id,
+            "content": reconstructed_content,
+            "metadata": chunks[0].document_metadata,  # Use first chunk's metadata
+            "chunks": [{"id": c.id, "index": c.chunk_index, "title": c.title} for c in chunks],
+        }
+
+    def insert_corpus(
+        self,
+        content: str,
+        corpus_metadata: dict[str, Any],
+        optional_props: BaseDocumentOptionalProps,
+    ) -> int:
+        """
+        Insert a new `Corpus`, which will be split into 1-or-more `Document`s, depending on its length.
+        Each `Document` chunk shall have its own embedding vector, but reference the parent corpus_id.
+
+        Args:
+            content: The text content to be inserted as a corpus
+            corpus_metadata: Dictionary of metadata associated with the corpus
+            optional_props: Optional properties for the documents (title, collection, etc.)
+
+        Returns:
+            int: The number of **documents** inserted for the provided corpus
+        """
+        corpus_id = uuid4()
+        document_contents = self._split_corpus(content)
+        document_embeddings = self.config.embedding_provider.embed_batch(document_contents)
+        return self.insert_documents(corpus_id, document_contents, document_embeddings, corpus_metadata, optional_props)
+
+    def insert_documents(
+        self,
+        corpus_id: UUID,
+        document_contents: list[str],
+        document_embeddings: list[list[float]],
+        corpus_metadata: dict[str, Any],
+        optional_props: BaseDocumentOptionalProps | None = None,
+    ) -> int:
+        """
+        Insert a list of documents (usually from a chunked + embedded corpus).
+
+        Args:
+            corpus_id: UUID of the corpus these documents belong to
+            document_contents: List of text content for each document
+            document_embeddings: List of embedding vectors corresponding to each document
+            corpus_metadata: Dictionary of metadata to associate with all documents
+            optional_props: Optional properties for the documents (title, collection, etc.)
+
+        Returns:
+            int: The number of documents inserted (0 if input lists are empty)
+
+        Raises:
+            ValueError: If the length of document_contents doesn't match document_embeddings
+        """
+        if len(document_contents) != len(document_embeddings):
+            raise ValueError("Number of embeddings does not match number of documents")
+        if len(document_contents) == 0:
+            return 0
+        documents_to_insert = []
+        for i in range(len(document_contents)):
+            chunk_md = self._extract_chunk_metadata(document_contents[i])
+            base_metadata = self.config.document_metadata(**(corpus_metadata | chunk_md))
+            documents_to_insert.append(
+                self.config.document_cls.from_props(
+                    corpus_id=corpus_id,
+                    chunk_index=i,
+                    content=document_contents[i],
+                    embedding=document_embeddings[i],
+                    metadata=base_metadata.model_dump(),
+                    optional_props=optional_props,
+                )
+            )
+        self.session.add_all(documents_to_insert)
+        self.session.commit()
+        return len(documents_to_insert)
+
+    def _split_corpus(self, content: str, **kwargs) -> list[str]:
+        """
+        **It is highly recommended to override this method.**
+        Split a corpus into chunks.
+        """
+        if self.__class__ is not BaseCorpusManager:
+            logger.warning("Using default _split_corpus. Override this method to improve performance.")
+        split_content = [content[i : i + 1000] for i in range(0, len(content), 1000)]
+        return [c for c in split_content if len(c.strip()) > 0]
+
+    def _extract_chunk_metadata(self, content: str) -> dict[str, Any]:
+        """
+        **It is highly recommended to override this method.**
+        Extract metadata from a chunk of content, to be appended to corpus metadata
+        """
+        if self.__class__ is not BaseCorpusManager:
+            logger.warning("Using default _extract_chunk_metadata. It is highly recommended to override this method.")
+        return {
+            "chunk_length": len(content),
+        }

pgvector_template-0.1.0/pgvector_template/core/retriever.py

@@ -0,0 +1,45 @@
+from dataclasses import dataclass, asdict
+from datetime import datetime
+from enum import Enum
+from typing import Any, Optional
+
+
+# class RetrievalStrategy(Enum):
+#     """Supported retrieval strategies"""
+#     SIMILARITY = "similarity"      # Pure vector similarity
+#     KEYWORD = "keyword"           # Metadata-based keyword search
+#     HYBRID = "hybrid"             # Combined similarity + keyword
+#     TEMPORAL = "temporal"         # Time-based search (if applicable)
+
+
+@dataclass
+class RetrievalResult:
+    """Standardized result structure for all retrieval operations"""
+
+    document_id: str
+    original_id: str
+    chunk_index: int
+    content: str
+    title: Optional[str]
+    metadata: dict[str, Any]
+    similarity_score: Optional[float] = None
+    keyword_matches: Optional[list[str]] = None
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+@dataclass
+class SearchQuery:
+    """Standardized search query structure"""
+
+    text: Optional[str] = None
+    keywords: Optional[list[str]] = None
+    metadata_filters: Optional[dict[str, Any]] = None
+    date_range: Optional[tuple[datetime, datetime]] = None
+    limit: int = 10
+    # strategy: RetrievalStrategy = RetrievalStrategy.HYBRID
+
+    def validate(self) -> bool:
+        """Ensure query has at least one search criterion"""
+        return any([self.text, self.keywords, self.metadata_filters, self.date_range])

pgvector_template-0.1.0/pgvector_template/core/search.py

@@ -0,0 +1,71 @@
+from abc import ABC, abstractmethod
+from logging import getLogger
+from typing import Any, Generic, Optional
+
+from sqlalchemy import text
+
+from pgvector_template.core import BaseEmbeddingProvider, RetrievalResult, SearchQuery, BaseDocument
+from pgvector_template.types import DocumentType
+from sqlalchemy.orm import Session
+
+
+class BaseSearchClient(ABC, Generic[DocumentType]):
+    """Abstract base for document retrieval"""
+
+    def __init__(self, session: Session, embedding_provider: BaseEmbeddingProvider):
+        self.session = session
+        self.embedding_provider = embedding_provider
+        self.logger = getLogger(self.__class__.__name__)
+
+    @abstractmethod
+    def search(self, query: SearchQuery) -> list[RetrievalResult]:
+        """Main search interface"""
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def search_by_metadata(self, filters: dict[str, Any], limit: int = 10) -> list[BaseDocument]:
+        """Generic JSON-based metadata search"""
+        query = self.session.query(BaseDocument).filter(BaseDocument.is_deleted == False)
+
+        # Apply JSON-based filters
+        for key, value in filters.items():
+            if isinstance(value, list):
+                # Array contains search
+                query = query.filter(text(f"metadata->>'{key}' = ANY(:value)")).params(value=value)
+            elif isinstance(value, dict):
+                # Nested JSON search
+                for nested_key, nested_value in value.items():
+                    query = query.filter(text(f"metadata->'{key}'->>'{nested_key}' = :value")).params(
+                        value=nested_value
+                    )
+            else:
+                # Simple equality
+                query = query.filter(text(f"metadata->>'{key}' = :value")).params(value=str(value))
+
+        return query.limit(limit).all()
+
+    @abstractmethod
+    def get_full_document(self, original_id: str) -> Optional[dict[str, Any]]:
+        """Reconstruct full document from chunks"""
+        raise NotImplementedError("Subclasses must implement this method")
+
+    # Template methods with default implementations
+    def similarity_search(self, text: str, limit: int = 10) -> list[RetrievalResult]:
+        """Vector similarity search"""
+        if not text.strip():
+            return []
+
+        embedding = self.embedding_provider.embed_text(text)
+        return self._vector_search(embedding, limit)
+
+    def keyword_search(self, keywords: list[str], limit: int = 10) -> list[RetrievalResult]:
+        """Metadata-based keyword search"""
+        return self._metadata_search({"keywords": keywords}, limit)
+
+    def _vector_search(self, embedding: list[float], limit: int) -> list[RetrievalResult]:
+        """Internal vector similarity search"""
+        # This will be implemented by the template using the DocumentType
+        raise NotImplementedError("Subclasses must implement vector search")
+
+    def _metadata_search(self, filters: dict[str, Any], limit: int) -> list[RetrievalResult]:
+        """Internal metadata search"""
+        raise NotImplementedError("Subclasses must implement metadata search")

pgvector_template-0.1.0/pgvector_template/db/__init__.py

@@ -0,0 +1,4 @@
+from pgvector_template.db.connection import DatabaseManager
+from pgvector_template.db.document_db import DocumentDatabaseManager, TempDocumentDatabaseManager
+
+__all__ = ["DatabaseManager", "DocumentDatabaseManager", "TempDocumentDatabaseManager"]

pgvector_template-0.1.0/pgvector_template/db/connection.py

@@ -0,0 +1,66 @@
+from contextlib import contextmanager
+from logging import getLogger
+from typing import Generator, Type
+
+from sqlalchemy import create_engine, text
+from sqlalchemy.orm import sessionmaker, Session
+from sqlalchemy.ext.declarative import DeclarativeMeta
+
+
+class DatabaseManager:
+    """Manages database connections and schema setup"""
+
+    def __init__(self, database_url: str):
+        self.database_url = database_url
+        self.engine = None
+        self.SessionLocal = None
+        self.logger = getLogger(self.__class__.__name__)
+
+    def initialize(self):
+        """Initialize database connection and session factory"""
+        self.engine = create_engine(
+            self.database_url, pool_pre_ping=True, pool_recycle=300, echo=False  # Set to True for SQL debugging
+        )
+        self.SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=self.engine)
+        self._ensure_pgvector_extension()
+
+    def create_schema(self, schema_name: str) -> None:
+        """Create a new schema for a collection type"""
+        with self.engine.connect() as conn:
+            conn.execute(text(f"CREATE SCHEMA IF NOT EXISTS {schema_name}"))
+            conn.commit()
+
+        self.logger.info(f"Created schema: {schema_name}")
+
+    def create_tables(self, base_class: Type[DeclarativeMeta], schema_name: str) -> None:
+        """Create tables for a specific schema"""
+        # Ensure all tables in this base class use the specified schema
+        for table in base_class.metadata.tables.values():
+            table.schema = schema_name
+        base_class.metadata.create_all(self.engine, checkfirst=True)
+        self.logger.info(f"Created tables for schema: {schema_name}")
+
+    def _ensure_pgvector_extension(self):
+        """Ensure pgvector extension is available"""
+        with self.engine.connect() as conn:
+            conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
+            conn.commit()
+
+    @contextmanager
+    def get_session(self) -> Generator[Session, None, None]:
+        """
+        Get a database session with automatic cleanup. Usage:
+
+        ```python
+        with db_mgr.get_session() as session:
+            # do stuff with the Session here
+        ```
+        """
+        session = self.SessionLocal()
+        try:
+            yield session
+        except Exception:
+            session.rollback()
+            raise
+        finally:
+            session.close()

pgvector_template-0.1.0/pgvector_template/db/document_db.py

@@ -0,0 +1,85 @@
+from typing import List, Type
+
+from sqlalchemy import text
+
+from pgvector_template.core.document import BaseDocument
+from pgvector_template.db.connection import DatabaseManager
+
+
+class DocumentDatabaseManager(DatabaseManager):
+    """
+    Specialized database manager for document collections with streamlined setup
+
+    Usage example:
+    ```python
+    db = DocumentDatabaseManager(
+        database_url="postgresql://user:pass@localhost/db",
+        schema_suffix="fpga_operations",
+        document_classes=[TextDocument, ImageDocument, AudioDocument]
+    )
+    db.setup()
+    ```
+    """
+
+    SCHEMA_PREFIX = "knowledge_base_"
+
+    def __init__(self, database_url: str, schema_suffix: str, document_classes: List[Type[BaseDocument]]):
+        """
+        Initialize a document-oriented database manager
+
+        Args:
+            database_url: PostgreSQL connection string
+            schema_suffix: Suffix for schema name (will be prefixed with SCHEMA_PREFIX)
+            document_classes: List of concrete subclasses of BaseDocument
+        """
+        super().__init__(database_url)
+        self.schema_suffix = schema_suffix
+        self.schema_name = f"{self.SCHEMA_PREFIX}{schema_suffix}"
+        self.document_classes = document_classes
+
+    def setup(self) -> None:
+        """One-step setup: initialize connection, create schema and tables for all document classes"""
+        self.initialize()
+        self.create_schema(self.schema_name)
+
+        for doc_class in self.document_classes:
+            # Set schema for each document class
+            doc_class.__table__.schema = self.schema_name
+            self.create_tables(doc_class, self.schema_name)
+
+        self.logger.info(
+            f"Document database setup complete for {self.schema_name} with {len(self.document_classes)} tables"
+        )
+
+
+class TempDocumentDatabaseManager(DocumentDatabaseManager):
+    def setup(self) -> str:
+        """
+        Create a temporary schema with a unique name for testing
+        Format: `temp_knowledge_base_<schema_suffix>_<uuid_snippet>`
+        """
+        from uuid import uuid4
+
+        temp_schema_name = f"temp_{self.schema_name}_{uuid4().hex[:8]}"
+
+        self.initialize()
+        self.create_schema(temp_schema_name)
+
+        for doc_class in self.document_classes:
+            # Create a copy of the table with the temporary schema
+            doc_class.__table__.schema = temp_schema_name
+            self.create_tables(doc_class, temp_schema_name)
+
+        self.logger.info(f"Created temporary schema: {temp_schema_name}")
+        return temp_schema_name
+
+    def cleanup(self, schema_name: str) -> None:
+        """Drop a temporary schema and all its objects"""
+        if not schema_name.startswith("temp_"):
+            raise ValueError("Can only drop schemas with 'temp_' prefix for safety")
+
+        with self.engine.connect() as conn:
+            conn.execute(text(f"DROP SCHEMA IF EXISTS {schema_name} CASCADE"))
+            conn.commit()
+
+        self.logger.info(f"Dropped temporary schema: {schema_name}")

pgvector_template-0.1.0/pgvector_template/types.py

@@ -0,0 +1,6 @@
+from typing import TypeVar
+
+from pgvector_template.core.document import BaseDocument, BaseDocumentMetadata
+
+T = TypeVar("T", bound="BaseDocumentMetadata")
+DocumentType = TypeVar("DocumentType", bound="BaseDocument")

pgvector_template-0.1.0/pgvector_template.egg-info/PKG-INFO

@@ -0,0 +1,44 @@
+Metadata-Version: 2.1
+Name: pgvector-template
+Version: 0.1.0
+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: psycopg[binary]>=3.1.0
+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: pytest>=7.0.0; extra == "test"
+Requires-Dist: pytest-cov>=4.0.0; extra == "test"
+Requires-Dist: python-dotenv>=1.0.0; extra == "test"
+
+# PGVector-Template
+
+Template library for flexible PGVector RAG implementations
+
+
+## Testing
+
+Install dependencies (preferably in a virtualenv) before running tests:
+```bash
+pip install -e .[test]
+```
+
+### Unit tests
+```bash
+python -m unittest
+```
+
+### Integration tests
+```bash
+python -m unittest discover -s integ-tests
+```

pgvector_template-0.1.0/pgvector_template.egg-info/SOURCES.txt

@@ -0,0 +1,21 @@
+LICENSE
+README.md
+pyproject.toml
+pgvector_template/__init__.py
+pgvector_template/types.py
+pgvector_template.egg-info/PKG-INFO
+pgvector_template.egg-info/SOURCES.txt
+pgvector_template.egg-info/dependency_links.txt
+pgvector_template.egg-info/requires.txt
+pgvector_template.egg-info/top_level.txt
+pgvector_template/core/__init__.py
+pgvector_template/core/document.py
+pgvector_template/core/embedder.py
+pgvector_template/core/manager.py
+pgvector_template/core/retriever.py
+pgvector_template/core/search.py
+pgvector_template/db/__init__.py
+pgvector_template/db/connection.py
+pgvector_template/db/document_db.py
+pgvector_template/search/__init__.py
+pgvector_template/utils/__init__.py

pgvector_template-0.1.0/pgvector_template.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

pgvector_template-0.1.0/pgvector_template.egg-info/requires.txt

@@ -0,0 +1,10 @@
+psycopg[binary]>=3.1.0
+pgvector>=0.2.0
+pydantic<3.0,>=2.11
+sqlalchemy>=2.0.0
+typing-extensions>=4.0.0
+
+[test]
+pytest>=7.0.0
+pytest-cov>=4.0.0
+python-dotenv>=1.0.0

pgvector_template-0.1.0/pgvector_template.egg-info/top_level.txt

@@ -0,0 +1 @@
+pgvector_template

pgvector_template-0.1.0/pyproject.toml

@@ -0,0 +1,44 @@
+[build-system]
+requires = ["setuptools", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "pgvector-template"
+version = "0.1.0"
+description = "Template library for flexible PGVector RAG implementations"
+authors = [{ name="DL", email="v49t9zpqd@mozmail.com" }]
+license = { text = "MIT" }
+dependencies = [
+    "psycopg[binary]>=3.1.0",
+    "pgvector>=0.2.0",
+    "pydantic>=2.11,<3.0",
+    "sqlalchemy>=2.0.0",
+    "typing-extensions>=4.0.0",
+]
+readme = "README.md"
+requires-python = ">=3.11"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+
+[project.optional-dependencies]
+test = [
+    "pytest>=7.0.0",
+    "pytest-cov>=4.0.0",
+    "python-dotenv>=1.0.0",
+]
+
+[project.urls]
+Homepage = "https://github.com/DavidLiuGit/PGVector-Template"
+
+
+[tool.setuptools.packages.find]
+where = ["."]  # Look for packages in the root directory
+include = ["pgvector_template*"]  # Include all packages starting with pgvector_template
+
+
+[tool.black]
+line-length = 120
+target-version = ["py311", "py312", "py313"]

pgvector_template-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+