pgvector-template 0.2.0__tar.gz → 0.2.2__tar.gz

{pgvector_template-0.2.0 → pgvector_template-0.2.2}/PKG-INFO

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

{pgvector_template-0.2.0 → pgvector_template-0.2.2}/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"""
@@ -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.0 → pgvector_template-0.2.2}/pgvector_template/core/search.py

@@ -109,22 +109,21 @@ class BaseSearchClient:
         """
         db_query = select(self.config.document_cls)
 
-        # if query.text:
-        #     db_query = self._apply_semantic_search(db_query, query)
+        if query.text:
+            db_query = self._apply_semantic_search(db_query, query)
         db_query = self._apply_keyword_search(db_query, query)
         # if query.metadata_filters:
         #     db_query = self._apply_metadata_filters(db_query, query)
         db_query = db_query.limit(query.limit)
 
-        # Execute the query
+        # execute query and return results
         results = self.session.scalars(db_query).all()
-
-        # Convert to RetrievalResult objects
         return self._convert_to_retrieval_results(results)
 
     def _apply_semantic_search(self, query: Select, search_query: SearchQuery) -> Select:
         """Apply semantic (vector) search criteria to the query.
         `embedding_provider` must be provided at instantiation, or an `ValueError` will be raised.
+        In PGVector, `<=>` operator is used to compare cosine distance. Lower = more similar.
 
         Args:
             query: The base SQLAlchemy query.
@@ -133,11 +132,10 @@ class BaseSearchClient:
         Returns:
             Updated SQLAlchemy query with semantic search applied.
         """
-        if not self.config.embedding_provider:
-            raise ValueError(
-                "EmbeddingProvider not provided in config. Vector (semantic) search is unavailable."
-            )
-        raise NotImplementedError
+        if not search_query.text:
+            return query
+        query_embedding = self.embedding_provider.embed_text(search_query.text)
+        return query.order_by(self.config.document_cls.embedding.cosine_distance(query_embedding))
 
     def _apply_keyword_search(self, db_query: Select, search_query: SearchQuery) -> Select:
         """Apply keyword (full-text) search criteria to the query.

{pgvector_template-0.2.0 → pgvector_template-0.2.2}/pgvector_template/service/document_service.py

@@ -53,6 +53,12 @@ class DocumentServiceConfig(BaseModel):
         if type(self.search_client_cfg) is BaseSearchClientConfig:
             self.search_client_cfg.document_cls = self.document_cls
 
+        # assign embedding_provider to CorpusManager & SearchClient configs
+        if not self.corpus_manager_cfg.embedding_provider:
+            self.corpus_manager_cfg.embedding_provider = self.embedding_provider
+        if not self.search_client_cfg.embedding_provider:
+            self.search_client_cfg.embedding_provider = self.embedding_provider
+
     model_config = {"arbitrary_types_allowed": True}
 
 

{pgvector_template-0.2.0 → pgvector_template-0.2.2}/pgvector_template.egg-info/PKG-INFO

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

{pgvector_template-0.2.0 → pgvector_template-0.2.2}/pyproject.toml

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