pgvector-template 0.3.5__tar.gz → 0.5__tar.gz

{pgvector_template-0.3.5 → pgvector_template-0.5}/PKG-INFO

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

{pgvector_template-0.3.5 → pgvector_template-0.5}/pgvector_template/core/document.py

@@ -1,8 +1,8 @@
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import Any, Type, TypeVar
 from uuid import uuid4, UUID as UuidLiteral
 
-from pydantic import BaseModel, Field, field_validator
+from pydantic import BaseModel, Field
 from sqlalchemy import (
     Column,
     String,
@@ -10,7 +10,6 @@ from sqlalchemy import (
     DateTime,
     Boolean,
     Integer,
-    Float,
     Index,
     UniqueConstraint,
 )
@@ -32,21 +31,6 @@ class BaseDocumentOptionalProps(BaseModel):
     """Optional source URL for the document"""
     language: str | None = Field(default="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(default=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")
@@ -88,18 +72,29 @@ class BaseDocument(Base):
     """Optional source URL"""
     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)."""
-    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."""
+    embedding_config = Column(JSONB, nullable=True)
+    """
+    Configuration used to generate the embedding vector. Stored as a JSON object so it
+    can capture provider-specific parameters beyond just the model ID — e.g. input_type,
+    dimensions, truncation strategy, or any future provider-specific options.
+
+    Recommended shape: {"model": "cohere.embed-v4:0", "input_type": "search_document", ...}
+
+    Intentionally nullable: rows embedded before this column was added will have NULL here.
+    Use this field for auditing and diagnosing mixed-model corpora, not for runtime dispatch.
+    """
 
     # Audit fields
-    created_at = Column(DateTime, default=datetime.utcnow)
-    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
+    created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
+    updated_at = Column(
+        DateTime(timezone=True),
+        default=lambda: datetime.now(timezone.utc),
+        onupdate=lambda: datetime.now(timezone.utc),
+    )
     is_deleted = Column(Boolean, default=False, index=True)
     """Entries can be logically marked for deletion before they are permanently deleted."""
 
@@ -112,10 +107,11 @@ class BaseDocument(Base):
 
             base_table_args = (
                 Index(
-                    f"{table_name}_metadata_gin_idx", "document_metadata", postgresql_using="gin"
+                    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",
@@ -140,7 +136,8 @@ class BaseDocument(Base):
         chunk_index: int,
         content: str,
         embedding: list[float],
-        metadata: dict[str, Any] = {},
+        embedding_config: dict[str, Any] | None = None,
+        metadata: dict[str, Any] | None = None,
         optional_props: BaseDocumentOptionalProps | None = None,
     ) -> T:
         """
@@ -151,6 +148,10 @@ class BaseDocument(Base):
             chunk_index: Index of this chunk within the corpus
             content: Text content of the document
             embedding: Vector embedding of the content
+            embedding_config: Provider config used to generate the embedding. Should include
+                at minimum {"model": "<model_id>"} plus any provider-specific params
+                (e.g. {"model": "cohere.embed-v4:0", "input_type": "search_document"}).
+                Nullable — rows pre-dating this column will have NULL.
             optional_props: Optional properties for the document
 
         Returns:
@@ -165,13 +166,12 @@ class BaseDocument(Base):
             chunk_index=chunk_index,
             content=content,
             embedding=embedding,
+            embedding_config=embedding_config,
             title=optional_props.title,
-            document_metadata=metadata,
+            document_metadata=metadata or {},
             collection=optional_props.collection,
             origin_url=optional_props.original_url,
             language=optional_props.language,
-            score=optional_props.score,
-            tags=optional_props.tags,
         )
 
     @classmethod

{pgvector_template-0.3.5 → pgvector_template-0.5}/pgvector_template/core/embedder.py

@@ -1,9 +1,22 @@
 from abc import ABC, abstractmethod
+from typing import Any
 
 
 class BaseEmbeddingProvider(ABC):
     """Abstract base for embedding generation"""
 
+    def __init__(self, model_id: str, **kwargs):
+        self.model_id = model_id
+
+    def get_embedding_config(self) -> dict[str, Any]:
+        """
+        Return a JSON-serializable dict describing the full embedding configuration.
+        Stored in the `embedding_config` column on every document row.
+        Subclasses should override to include provider-specific params (e.g. input_type).
+        Default implementation returns just the model ID.
+        """
+        return {"model": self.model_id}
+
     @abstractmethod
     def embed_text(self, text: str) -> list[float]:
         """Generate embedding vector for text"""

{pgvector_template-0.3.5 → pgvector_template-0.5}/pgvector_template/core/manager.py

@@ -204,6 +204,7 @@ class BaseCorpusManager(ABC):
         optional_props: BaseDocumentOptionalProps | None,
     ) -> list[BaseDocument]:
         """Create document instances from the provided data"""
+        embedding_config = self.embedding_provider.get_embedding_config()
         documents = []
         for i, (content, embedding) in enumerate(zip(document_contents, document_embeddings)):
             chunk_md = self._extract_chunk_metadata(content)
@@ -214,6 +215,7 @@ class BaseCorpusManager(ABC):
                     chunk_index=i,
                     content=content,
                     embedding=embedding,
+                    embedding_config=embedding_config,
                     metadata=base_metadata.model_dump(),
                     optional_props=optional_props,
                 )

{pgvector_template-0.3.5 → pgvector_template-0.5}/pgvector_template/core/search.py

@@ -69,38 +69,43 @@ class BaseSearchClient:
             query: Search query containing text, metadata filters, and pagination.
 
         Returns:
-            List of retrieval results matching the search criteria.
+            List of retrieval results matching the search criteria. Each result's `score`
+            is a float in [0, 1] when `query.text` is set (1.0 = most similar), or `None`
+            for keyword/metadata-only queries.
         """
-        db_query = select(self.config.document_cls)
+        distance_expr: ColumnElement | None = None
 
         if query.text:
-            db_query = self._apply_semantic_search(db_query, query)
+            distance_expr = self._build_distance_expr(query)
+            db_query = select(self.config.document_cls, distance_expr).order_by(distance_expr)
+        else:
+            db_query = select(self.config.document_cls)
+
         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 query and return results
-        results = self.session.scalars(db_query).all()
-        return self._convert_to_retrieval_results(results)
+        if distance_expr is not None:
+            rows = self.session.execute(db_query).all()
+            return self._convert_to_retrieval_results(rows)
+        else:
+            results = self.session.scalars(db_query).all()
+            return self._convert_to_retrieval_results([(doc, None) for doc in results])
 
-    def _apply_semantic_search(self, db_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.
+    def _build_distance_expr(self, search_query: SearchQuery) -> ColumnElement:
+        """Build a labeled cosine-distance expression for the given query text.
 
-        Args:
-            query: The base SQLAlchemy query.
-            search_query: The search query containing the text to search for.
+        `embedding_provider` must be provided at instantiation, or a `ValueError` is raised.
+        In PGVector, `<=>` computes cosine distance (0 = identical, up to 2 = opposite).
 
         Returns:
-            Updated SQLAlchemy query with semantic search applied.
+            A SQLAlchemy labeled column expression suitable for SELECT and ORDER BY.
         """
-        if not search_query.text:
-            return db_query
+        assert search_query.text is not None
         query_embedding = self.embedding_provider.embed_text(search_query.text)
-        return db_query.order_by(
-            self.config.document_cls.embedding.cosine_distance(query_embedding)
+        return self.config.document_cls.embedding.cosine_distance(query_embedding).label(
+            "_score_distance"
         )
 
     def _apply_keyword_search(self, db_query: Select, search_query: SearchQuery) -> Select:
@@ -199,16 +204,18 @@ class BaseSearchClient:
             raise ValueError(f"Unsupported condition: {filter_obj.condition}")
 
     def _convert_to_retrieval_results(self, results: Sequence[Any]) -> list[RetrievalResult]:
-        """Convert database results to RetrievalResult objects.
+        """Convert (document, distance | None) pairs to RetrievalResult objects.
 
         Args:
-            results: Raw database results.
-            search_query: The original search query.
+            results: Sequence of (document, cosine_distance) tuples. Pass distance=None
+                     for keyword/metadata-only queries that have no semantic ranking.
 
         Returns:
-            List of RetrievalResult objects.
+            List of RetrievalResult objects with score = max(0.0, 1.0 - distance),
+            or score=None when distance is None.
         """
         retrieval_results = []
-        for result in results:
-            retrieval_results.append(RetrievalResult(document=result, score=1.0))
+        for document, distance in results:
+            score = max(0.0, 1.0 - distance / 2.0) if distance is not None else None
+            retrieval_results.append(RetrievalResult(document=document, score=score))
         return retrieval_results

{pgvector_template-0.3.5 → pgvector_template-0.5}/pgvector_template/models/search.py

@@ -1,5 +1,4 @@
 from dataclasses import dataclass, asdict
-from datetime import datetime
 from typing import Any, Literal
 
 from pydantic import BaseModel, ConfigDict, Field, model_validator
@@ -38,7 +37,13 @@ class SearchQuery(BaseModel):
     text: str | None = None
     """String to match against using in a semantic search, i.e. using vector distance."""
     keywords: list[str] = []
-    """List of keywords to **exact-match** in a keyword search."""
+    """
+    List of keywords for a case-insensitive **substring** search (SQL `ILIKE '%keyword%'`).
+    An entry matches if **any** keyword appears anywhere in its content (OR-combined).
+    Pass multiple variants, synonyms, or spellings to increase recall — e.g.
+    ["run", "running", "ran"] will match entries containing any of those forms.
+    Note: substring matching means "cat" also matches "category"; prefer distinctive stems.
+    """
     metadata_filters: list[MetadataFilter] = Field(
         default=[],
         json_schema_extra={"metadata_schema": BaseDocumentMetadata.model_json_schema()},
@@ -47,8 +52,6 @@ class SearchQuery(BaseModel):
     List of metadata conditions that must be matched.
     Refer to `metadata_schema` for the expected schema, as it exists in the database.
     """
-    date_range: tuple[datetime, datetime] | None = None
-    """Retrieve/limit results based on created_at & updated_at timestamps (i.e. database operations)"""
     limit: int = Field(
         ...,
         ge=1,
@@ -62,7 +65,7 @@ class SearchQuery(BaseModel):
 
     @model_validator(mode="after")
     def ensure_criterion(self):
-        if not any([self.text, self.keywords, self.metadata_filters, self.date_range]):
+        if not any([self.text, self.keywords, self.metadata_filters]):
             raise ValueError("At least one search criterion is required")
         return self
 
@@ -72,7 +75,12 @@ class RetrievalResult:
     """Standardized result structure for all retrieval operations"""
 
     document: BaseDocument
-    score: float
+    score: float | None
+    """
+    Semantic similarity score in [0, 1] when a `text` query was provided (1.0 = identical,
+    0.0 = maximally distant). `None` when the search had no semantic component (keywords /
+    metadata filters only) — results are still filtered correctly but have no meaningful ranking.
+    """
 
     def to_dict(self) -> dict[str, Any]:
         return asdict(self)

{pgvector_template-0.3.5 → pgvector_template-0.5}/pgvector_template.egg-info/PKG-INFO

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

{pgvector_template-0.3.5 → pgvector_template-0.5}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pgvector-template"
-version = "0.3.5"
+version = "0.5"
 description = "Template library for flexible PGVector RAG implementations"
 authors = [{ name="DL", email="v49t9zpqd@mozmail.com" }]
 license = { text = "MIT" }
@@ -52,4 +52,4 @@ line-length = 100
 target-version = "py311"
 
 [tool.ty.src]
-exclude = []
+exclude = ["integ-tests"]