pgvector-template 0.2.1__tar.gz → 0.2.3__tar.gz

{pgvector_template-0.2.1 → pgvector_template-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pgvector-template
-Version: 0.2.1
+Version: 0.2.3
 Summary: Template library for flexible PGVector RAG implementations
 Author-email: DL <v49t9zpqd@mozmail.com>
 License: MIT

{pgvector_template-0.2.1 → pgvector_template-0.2.3}/pgvector_template/core/document.py

@@ -12,6 +12,7 @@ from sqlalchemy import (
     Integer,
     Float,
     Index,
+    UniqueConstraint,
 )
 from sqlalchemy.orm import declarative_base
 from sqlalchemy.dialects.postgresql import UUID, JSONB
@@ -63,6 +64,7 @@ class BaseDocument(Base):
     """
 
     __abstract__ = True
+    __tablename__ = "base_document"
 
     id = Column(UUID(as_uuid=True), primary_key=True, default=uuid4)
     """Primary key of the Document table. Represents unique ID of a Document"""
@@ -84,7 +86,7 @@ class BaseDocument(Base):
     """Flexible metadata as JSON"""
     origin_url = Column(String(2048), nullable=True)
     """Optional source URL"""
-    language = Column(String(10), default="en")
+    language = Column(String(2), 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)."""
@@ -98,9 +100,39 @@ class BaseDocument(Base):
     # 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)
+    is_deleted = Column(Boolean, default=False, index=True)
     """Entries can be logically marked for deletion before they are permanently deleted."""
 
+    def __init_subclass__(cls, **kwargs):
+        """Simply declare `__table_args__` to add more args in child classes."""
+        super().__init_subclass__(**kwargs)
+
+        if hasattr(cls, "__tablename__"):
+            table_name = cls.__tablename__
+
+            base_table_args = (
+                Index(
+                    f"{table_name}_metadata_gin_idx", "document_metadata", postgresql_using="gin"
+                ),
+                cls.get_embedding_index(table_name),
+                Index(f"{table_name}_tags_gin_idx", "tags", postgresql_using="gin"),
+                Index(f"{table_name}_collection_corpus_idx", "collection", "corpus_id"),
+                UniqueConstraint(
+                    "collection",
+                    "corpus_id",
+                    "chunk_index",
+                    name=f"{table_name}_collection_corpus_chunk_unique",
+                ),
+            )
+
+            subclass_args = getattr(cls, "__table_args__", ())
+            if isinstance(subclass_args, dict):
+                cls.__table_args__ = base_table_args + (subclass_args,)
+            else:
+                if not isinstance(subclass_args, tuple):
+                    subclass_args = (subclass_args,) if subclass_args else ()
+                cls.__table_args__ = base_table_args + subclass_args
+
     @classmethod
     def from_props(
         cls: Type[T],
@@ -142,9 +174,15 @@ class BaseDocument(Base):
             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")
+    @classmethod
+    def get_embedding_index(cls, table_name: str) -> Index:
+        """Override this method to customize the embedding index."""
+        return Index(
+            f"{table_name}_embedding_hnsw_idx",
+            "embedding",
+            postgresql_using="hnsw",
+            postgresql_ops={"embedding": "vector_cosine_ops"},
+        )
 
 
 class BaseDocumentMetadata(BaseModel):
@@ -154,7 +192,9 @@ class BaseDocumentMetadata(BaseModel):
     without any extraneous properties, or any missing properties, to avoid ambiguity.
     """
 
-    document_type: str = Field(..., description="Description for type of document, e.g. markdown, pdf, etc")
+    document_type: str = Field(
+        ..., description="Description for type of document, e.g. markdown, pdf, etc"
+    )
     schema_version: str = Field(default="1.0", description="Schema version for the metadata")
 
     def to_dict(self) -> dict[str, Any]:

{pgvector_template-0.2.1 → pgvector_template-0.2.3}/pgvector_template/core/search.py

@@ -3,10 +3,9 @@ from datetime import datetime
 from logging import getLogger
 from typing import Any, Type, Sequence
 
-from pydantic import BaseModel, Field, model_validator
-from sqlalchemy import text, select, or_, and_
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+from sqlalchemy import select, or_
 from sqlalchemy.sql import Select
-from pgvector.sqlalchemy import Vector
 
 from pgvector_template.core import (
     BaseEmbeddingProvider,
@@ -23,7 +22,7 @@ class SearchQuery(BaseModel):
     """Standardized search query structure. At least 1 search criterion is required."""
 
     text: str | None = None
-    """String to approximate-search (using vector distance) in a semantic search."""
+    """String to match against using in a semantic search, i.e. using vector distance."""
     keywords: list[str] | None = None
     """List of keywords to exact-match in a keyword search."""
     metadata_filters: dict[str, Any] | None = None
@@ -36,6 +35,8 @@ class SearchQuery(BaseModel):
     )
     """Maximum number of results to return."""
 
+    model_config = ConfigDict(use_attribute_docstrings=True)
+
     @model_validator(mode="after")
     def ensure_criterion(self):
         if not any([self.text, self.keywords, self.metadata_filters, self.date_range]):

{pgvector_template-0.2.1 → pgvector_template-0.2.3}/pgvector_template/db/document_db.py

@@ -37,7 +37,7 @@ class DocumentDatabaseManager(DatabaseManager):
         self.schema_name = f"{self.SCHEMA_PREFIX}{schema_suffix}"
         self.document_classes = document_classes
 
-    def setup(self) -> None:
+    def setup(self) -> str:
         """One-step setup: initialize connection, create schema and tables for all document classes"""
         self.initialize()
         self.create_schema(self.schema_name)
@@ -50,6 +50,7 @@ class DocumentDatabaseManager(DatabaseManager):
         self.logger.info(
             f"Document database setup complete for {self.schema_name} with {len(self.document_classes)} tables"
         )
+        return self.schema_name
 
 
 class TempDocumentDatabaseManager(DocumentDatabaseManager):

{pgvector_template-0.2.1 → pgvector_template-0.2.3}/pgvector_template/service/document_service.py

@@ -50,8 +50,10 @@ class DocumentServiceConfig(BaseModel):
         # iff either config is an instance of their respective base config classes (and not a subclass)
         if type(self.corpus_manager_cfg) is BaseCorpusManagerConfig:
             self.corpus_manager_cfg.document_cls = self.document_cls
+            self.corpus_manager_cfg.document_metadata_cls = self.document_metadata_cls
         if type(self.search_client_cfg) is BaseSearchClientConfig:
             self.search_client_cfg.document_cls = self.document_cls
+            self.search_client_cfg.document_metadata_cls = self.document_metadata_cls
 
         # assign embedding_provider to CorpusManager & SearchClient configs
         if not self.corpus_manager_cfg.embedding_provider:

{pgvector_template-0.2.1 → pgvector_template-0.2.3}/pgvector_template.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pgvector-template
-Version: 0.2.1
+Version: 0.2.3
 Summary: Template library for flexible PGVector RAG implementations
 Author-email: DL <v49t9zpqd@mozmail.com>
 License: MIT

{pgvector_template-0.2.1 → pgvector_template-0.2.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pgvector-template"
-version = "0.2.1"
+version = "0.2.3"
 description = "Template library for flexible PGVector RAG implementations"
 authors = [{ name="DL", email="v49t9zpqd@mozmail.com" }]
 license = { text = "MIT" }