pgvector-template 0.1.0__tar.gz → 0.1.2__tar.gz

{pgvector_template-0.1.0 → pgvector_template-0.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pgvector-template
-Version: 0.1.0
+Version: 0.1.2
 Summary: Template library for flexible PGVector RAG implementations
 Author-email: DL <v49t9zpqd@mozmail.com>
 License: MIT
@@ -11,12 +11,12 @@ 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: 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"

{pgvector_template-0.1.0 → pgvector_template-0.1.2}/pgvector_template/core/__init__.py

@@ -1,6 +1,6 @@
-from pgvector_template.core.document import BaseDocument, BaseDocumentMetadata, Corpus, BaseDocumentOptionalProps
+from pgvector_template.core.document import BaseDocument, BaseDocumentMetadata, BaseDocumentOptionalProps
 from pgvector_template.core.embedder import BaseEmbeddingProvider
-from pgvector_template.core.manager import BaseCorpusManager, BaseCorpusManagerConfig
+from pgvector_template.core.manager import BaseCorpusManager, BaseCorpusManagerConfig, Corpus
 from pgvector_template.core.retriever import RetrievalResult, SearchQuery
 from pgvector_template.core.search import BaseSearchClient
 

{pgvector_template-0.1.0 → pgvector_template-0.1.2}/pgvector_template/core/document.py

@@ -1,7 +1,6 @@
-from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Any, ClassVar, Type, TypeVar, Annotated
-from uuid import uuid4, UUID
+from typing import Any, Type, TypeVar
+from uuid import uuid4, UUID as UuidLiteral
 
 from pydantic import BaseModel, Field, field_validator, model_validator
 from sqlalchemy import (
@@ -13,7 +12,6 @@ from sqlalchemy import (
     Integer,
     Float,
     Index,
-    text,
 )
 from sqlalchemy.orm import declarative_base
 from sqlalchemy.dialects.postgresql import UUID, JSONB
@@ -106,7 +104,7 @@ class BaseDocument(Base):
     @classmethod
     def from_props(
         cls: Type[T],
-        corpus_id: UUID,
+        corpus_id: UuidLiteral | str,
         chunk_index: int,
         content: str,
         embedding: list[float],
@@ -117,7 +115,7 @@ class BaseDocument(Base):
         Create a BaseDocument instance from mandatory and optional properties.
 
         Args:
-            corpus_id: UUID of the corpus this document belongs to
+            corpus_id: UUID or string (max 64 chars) 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
@@ -127,8 +125,9 @@ class BaseDocument(Base):
             A new BaseDocument instance of the calling class type
         """
         if optional_props is None:
-            optional_props = BaseDocumentOptionalProps()
+            optional_props = BaseDocumentOptionalProps()  # type: ignore
 
+        # SQLAlchemy handles string-to-UUID conversion automatically with as_uuid=True
         return cls(
             corpus_id=corpus_id,
             chunk_index=chunk_index,
@@ -149,23 +148,14 @@ class BaseDocument(Base):
 
 
 class BaseDocumentMetadata(BaseModel):
-    """Base metadata structure"""
+    """
+    Base metadata structure.
+    It is generally expected that every `BaseDocument`'s metadata follows this exact schema,
+    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")
     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-0.1.2}/pgvector_template/core/manager.py

@@ -1,9 +1,10 @@
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
 from logging import getLogger
-from typing import Any, Optional, Type
+from typing import Any, Type
 from uuid import UUID, uuid4
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 from sqlalchemy.orm import Session
 
 from pgvector_template.core.document import BaseDocument, BaseDocumentMetadata, BaseDocumentOptionalProps
@@ -16,14 +17,32 @@ 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]
+    document_cls: Type[BaseDocument] = Field(..., description="Document class, must be subclass of BaseDocument")
+    embedding_provider: BaseEmbeddingProvider | None = Field(
+        None, description="Embedding provider for insert operations"
+    )
+    document_metadata_cls: Type[BaseDocumentMetadata] | None = Field(
+        None, description="Document metadata class, must be subclass of BaseDocumentMetadata"
+    )
 
     model_config = {"arbitrary_types_allowed": True}
 
 
+@dataclass
+class Corpus:
+    """
+    Serves as the return type for `BaseCorpusManager.get_full_corpus()`.
+
+    Logical grouping of one or more documents (chunks) belonging to the same original source (corpus).
+    Typically all documents in a corpus share the same `corpus_id` and are ordered by `chunk_index`.
+    """
+
+    corpus_id: UUID | str
+    content: str
+    metadata: dict[str, Any]  # e.g. source, tags, etc.
+    documents: list[BaseDocument]
+
+
 class BaseCorpusManager(ABC):
     """
     Template class for `Corpus` & `Document` management operations.
@@ -43,44 +62,33 @@ class BaseCorpusManager(ABC):
     ) -> 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]]:
+    def get_full_corpus(self, corpus_id: str) -> Corpus | None:
         """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)
+            self.session.query(self.config.document_cls)
+            .filter(self.config.document_cls.corpus_id == corpus_id, self.config.document_cls.is_deleted == False)
+            .order_by(self.config.document_cls.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],
-        }
+        content, metadata = self._join_documents(chunks)
+        return Corpus(
+            corpus_id=corpus_id,
+            content=content,
+            metadata=metadata,
+            documents=chunks,
+        )
 
     def insert_corpus(
         self,
         content: str,
         corpus_metadata: dict[str, Any],
-        optional_props: BaseDocumentOptionalProps,
+        optional_props: BaseDocumentOptionalProps | None = None,
+        corpus_id: UUID | str | None = None,
     ) -> int:
         """
         Insert a new `Corpus`, which will be split into 1-or-more `Document`s, depending on its length.
@@ -94,14 +102,17 @@ class BaseCorpusManager(ABC):
         Returns:
             int: The number of **documents** inserted for the provided corpus
         """
-        corpus_id = uuid4()
+        self._check_insert_dependencies()
+
+        if not corpus_id:
+            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,
+        corpus_id: UUID | str,
         document_contents: list[str],
         document_embeddings: list[list[float]],
         corpus_metadata: dict[str, Any],
@@ -123,6 +134,8 @@ class BaseCorpusManager(ABC):
         Raises:
             ValueError: If the length of document_contents doesn't match document_embeddings
         """
+        self._check_insert_dependencies()
+
         if len(document_contents) != len(document_embeddings):
             raise ValueError("Number of embeddings does not match number of documents")
         if len(document_contents) == 0:
@@ -130,7 +143,7 @@ class BaseCorpusManager(ABC):
         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))
+            base_metadata = self.config.document_metadata_cls(**(corpus_metadata | chunk_md))
             documents_to_insert.append(
                 self.config.document_cls.from_props(
                     corpus_id=corpus_id,
@@ -148,20 +161,59 @@ class BaseCorpusManager(ABC):
     def _split_corpus(self, content: str, **kwargs) -> list[str]:
         """
         **It is highly recommended to override this method.**
-        Split a corpus into chunks.
+        Split a corpus' string content into smaller 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 _join_documents(self, documents: list[BaseDocument]) -> tuple[str, dict[str, Any]]:
+        """
+        **It is highly recommended to override this method.**
+        **This method should effectively reverse the `_split_corpus` method.**
+        Join a list of documents back into a single corpus string.
+        Return an instance of corpus metadata. This is a best effort, since not all properties in
+        `BaseDocumentMetadata`/`document_metadata_cls` are relevant to the corpus.
+        """
+        if self.__class__ is not BaseCorpusManager:
+            logger.warning("Using default _join_documents. Override this method to improve functionality.")
+        documents.sort(key=lambda d: d.chunk_index)  # type: ignore
+        # since _split_corpus performs a simple split on every 1000 chars, we can simply call `join`
+        corpus_content = "".join(d.content for d in documents)  # type: ignore
+        corpus_metadata = self._infer_corpus_metadata(documents)
+        return corpus_content, corpus_metadata
+
     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
+        Extract metadata from a chunk of content, to be appended to corpus metadata.
+        Note: returning a key-value pair here does NOT guarantee its inclusion when it's added to the database.
         """
         if self.__class__ is not BaseCorpusManager:
             logger.warning("Using default _extract_chunk_metadata. It is highly recommended to override this method.")
+        # this is simply an example. Since the document metadata that gets saved
         return {
             "chunk_length": len(content),
         }
+
+    def _infer_corpus_metadata(self, documents: list[BaseDocument]) -> dict[str, Any]:
+        """
+        **It is highly recommended to override this method.**
+        **This method should be a best-effort reversal of `extract_chunk_metadata()`**
+        Infer metadata for the corpus from the constituent `BaseDocument`s.
+        """
+        if self.__class__ is not BaseCorpusManager:
+            logger.warning("Using default _infer_corpus_metadata. It is highly recommended to override this method.")
+        # merge all document.document_metadata together, and return
+        merged = {}
+        for d in documents:
+            merged.update(d.document_metadata)  # type: ignore
+        return merged
+
+    def _check_insert_dependencies(self) -> None:
+        """Check that required dependencies for insert operations are available"""
+        if not self.config.embedding_provider:
+            raise ValueError("embedding_provider must be provided in config for insert operations")
+        if not self.config.document_metadata_cls:
+            raise ValueError("document_metadata_cls must be provided in config for insert operations")

{pgvector_template-0.1.0 → pgvector_template-0.1.2}/pgvector_template.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pgvector-template
-Version: 0.1.0
+Version: 0.1.2
 Summary: Template library for flexible PGVector RAG implementations
 Author-email: DL <v49t9zpqd@mozmail.com>
 License: MIT
@@ -11,12 +11,12 @@ 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: 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"

{pgvector_template-0.1.0 → pgvector_template-0.1.2}/pyproject.toml

@@ -4,12 +4,11 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pgvector-template"
-version = "0.1.0"
+version = "0.1.2"
 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",
@@ -25,6 +24,7 @@ classifiers = [
 
 [project.optional-dependencies]
 test = [
+    "psycopg[binary]>=3.1.0",
     "pytest>=7.0.0",
     "pytest-cov>=4.0.0",
     "python-dotenv>=1.0.0",

{pgvector_template-0.1.0 → pgvector_template-0.1.2}/pgvector_template.egg-info/requires.txt

@@ -1,10 +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]
+psycopg[binary]>=3.1.0
 pytest>=7.0.0
 pytest-cov>=4.0.0
 python-dotenv>=1.0.0