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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pgvector-template
-Version: 0.1.3
+Version: 0.2.1
 Summary: Template library for flexible PGVector RAG implementations
 Author-email: DL <v49t9zpqd@mozmail.com>
 License: MIT
@@ -20,6 +20,8 @@ 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"
+Provides-Extra: dev
+Requires-Dist: black>=23.0.0; extra == "dev"
 
 # PGVector-Template
 

{pgvector_template-0.1.3 → pgvector_template-0.2.1}/pgvector_template/core/__init__.py

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

{pgvector_template-0.1.3 → pgvector_template-0.2.1}/pgvector_template/core/manager.py

@@ -7,7 +7,11 @@ from uuid import UUID, uuid4
 from pydantic import BaseModel, Field
 from sqlalchemy.orm import Session
 
-from pgvector_template.core.document import BaseDocument, BaseDocumentMetadata, BaseDocumentOptionalProps
+from pgvector_template.core.document import (
+    BaseDocument,
+    BaseDocumentMetadata,
+    BaseDocumentOptionalProps,
+)
 from pgvector_template.core.embedder import BaseEmbeddingProvider
 
 
@@ -17,13 +21,12 @@ logger = getLogger(__name__)
 class BaseCorpusManagerConfig(BaseModel):
     """Base configuration for `Corpus` & `Document` management operations"""
 
-    document_cls: Type[BaseDocument] = Field(..., description="Document class, must be subclass of BaseDocument")
-    embedding_provider: BaseEmbeddingProvider | None = Field(
-        default=None, description="Embedding provider for insert operations"
-    )
-    document_metadata_cls: Type[BaseDocumentMetadata] | None = Field(
-        default=None, description="Document metadata class, must be subclass of BaseDocumentMetadata"
-    )
+    document_cls: Type[BaseDocument] = Field(...)
+    """Document class **type** (not an instance). Must be subclass of `BaseDocument`."""
+    embedding_provider: BaseEmbeddingProvider | None = Field(default=None)
+    """Instance of `BaseEmbeddingProvider` child class. Acts as embedding provider for insert operations."""
+    document_metadata_cls: Type[BaseDocumentMetadata] = Field(default=BaseDocumentMetadata)
+    """Document metadata class **type** (not an instance). Must be subclass of BaseDocumentMetadata."""
 
     model_config = {"arbitrary_types_allowed": True}
 
@@ -55,6 +58,18 @@ class BaseCorpusManager(ABC):
     def config(self) -> BaseCorpusManagerConfig:
         return self._cfg
 
+    @property
+    def document_metadata_class(self) -> Type[BaseDocumentMetadata]:
+        """Returns the document metadata class, raising an error if it's not set."""
+        return self.config.document_metadata_cls
+
+    @property
+    def embedding_provider(self) -> BaseEmbeddingProvider:
+        """Returns the embedding provider, raising an error if it's not set."""
+        if self.config.embedding_provider is None:
+            raise ValueError("embedding_provider must be provided in config for this operation")
+        return self.config.embedding_provider
+
     def __init__(
         self,
         session: Session,
@@ -62,12 +77,19 @@ class BaseCorpusManager(ABC):
     ) -> None:
         self.session = session
         self._cfg = config
+        if not self.config.embedding_provider:
+            logger.warning(
+                "EmbeddingProvider not provided in config. Insertion will be unavailable."
+            )
 
-    def get_full_corpus(self, corpus_id: str) -> Corpus | None:
+    def get_full_corpus(self, corpus_id: str, **kwargs) -> Corpus | None:
         """Reconstruct full corpus from its individual documents/chunks"""
         chunks = (
             self.session.query(self.config.document_cls)
-            .filter(self.config.document_cls.corpus_id == corpus_id, self.config.document_cls.is_deleted == False)
+            .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()
         )
@@ -89,6 +111,7 @@ class BaseCorpusManager(ABC):
         corpus_metadata: dict[str, Any],
         optional_props: BaseDocumentOptionalProps | None = None,
         corpus_id: UUID | str | None = None,
+        **kwargs,
     ) -> int:
         """
         Insert a new `Corpus`, which will be split into 1-or-more `Document`s, depending on its length.
@@ -102,13 +125,13 @@ class BaseCorpusManager(ABC):
         Returns:
             int: The number of **documents** inserted for the provided corpus
         """
-        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)
+        document_embeddings = self.embedding_provider.embed_batch(document_contents)
+        return self.insert_documents(
+            corpus_id, document_contents, document_embeddings, corpus_metadata, optional_props
+        )
 
     def insert_documents(
         self,
@@ -117,6 +140,7 @@ class BaseCorpusManager(ABC):
         document_embeddings: list[list[float]],
         corpus_metadata: dict[str, Any],
         optional_props: BaseDocumentOptionalProps | None = None,
+        **kwargs,
     ) -> int:
         """
         Insert a list of documents (usually from a chunked + embedded corpus).
@@ -134,8 +158,6 @@ 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:
@@ -143,7 +165,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_cls(**(corpus_metadata | chunk_md))
+            base_metadata = self.document_metadata_class(**(corpus_metadata | chunk_md))
             documents_to_insert.append(
                 self.config.document_cls.from_props(
                     corpus_id=corpus_id,
@@ -164,11 +186,15 @@ class BaseCorpusManager(ABC):
         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.")
+            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]]:
+    def _join_documents(
+        self, documents: list[BaseDocument], **kwargs
+    ) -> tuple[str, dict[str, Any]]:
         """
         **It is highly recommended to override this method.**
         **This method should effectively reverse the `_split_corpus` method.**
@@ -177,43 +203,42 @@ class BaseCorpusManager(ABC):
         `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.")
+            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]:
+    def _extract_chunk_metadata(self, content: str, **kwargs) -> dict[str, Any]:
         """
         **It is highly recommended to override this method.**
         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.")
+            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]:
+    def _infer_corpus_metadata(self, documents: list[BaseDocument], **kwargs) -> 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.")
+            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.2.1/pgvector_template/core/search.py

@@ -0,0 +1,183 @@
+from dataclasses import dataclass, asdict
+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 sqlalchemy.sql import Select
+from pgvector.sqlalchemy import Vector
+
+from pgvector_template.core import (
+    BaseEmbeddingProvider,
+    BaseDocument,
+    BaseDocumentMetadata,
+)
+from sqlalchemy.orm import Session
+
+
+logger = getLogger(__name__)
+
+
+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."""
+    keywords: list[str] | None = None
+    """List of keywords to exact-match in a keyword search."""
+    metadata_filters: dict[str, Any] | None = None
+    """Strict metadata filters that must be matched."""
+    date_range: tuple[datetime, datetime] | None = None
+    """Retrieve/limit results based on created_at & updated_at timestamps"""
+    limit: int = Field(
+        ...,
+        ge=1,
+    )
+    """Maximum number of results to return."""
+
+    @model_validator(mode="after")
+    def ensure_criterion(self):
+        if not any([self.text, self.keywords, self.metadata_filters, self.date_range]):
+            raise ValueError("At least one search criterion is required")
+        return self
+
+
+@dataclass
+class RetrievalResult:
+    """Standardized result structure for all retrieval operations"""
+
+    document: BaseDocument
+    score: float
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+class BaseSearchClientConfig(BaseModel):
+    """Config obj for `BaseSearchClient`."""
+
+    document_cls: Type[BaseDocument] = Field(default=BaseDocument)
+    """Document class **type** (not an instance). Must be subclass of `BaseDocument`."""
+    embedding_provider: BaseEmbeddingProvider | None = Field(default=None)
+    """Instance of `BaseEmbeddingProvider` child class. Acts as embedding provider for semantic search."""
+    document_metadata_cls: Type[BaseDocumentMetadata] = Field(default=BaseDocumentMetadata)
+    """Document metadata class type. Used for metadata search operations."""
+
+    model_config = {"arbitrary_types_allowed": True}
+
+
+class BaseSearchClient:
+    """Minimum-viable implementation of document retrieval for PGVector"""
+
+    @property
+    def config(self) -> BaseSearchClientConfig:
+        return self._cfg
+
+    @property
+    def document_metadata_class(self) -> Type[BaseDocumentMetadata]:
+        """Returns the document metadata class, raising an error if it's not set."""
+        return self.config.document_metadata_cls
+
+    @property
+    def embedding_provider(self) -> BaseEmbeddingProvider:
+        """Returns the embedding provider, raising an error if it's not set."""
+        if self.config.embedding_provider is None:
+            raise ValueError("embedding_provider must be provided in config for this operation")
+        return self.config.embedding_provider
+
+    def __init__(
+        self,
+        session: Session,
+        config: BaseSearchClientConfig,
+    ):
+        self.session = session
+        self._cfg = config
+        if not self.config.embedding_provider:
+            logger.warning(
+                "EmbeddingProvider not provided in config. Vector (semantic) search will be unavailable."
+            )
+
+    def search(self, query: SearchQuery) -> list[RetrievalResult]:
+        """Search for documents based on the provided query.
+
+        Args:
+            query: Search query containing text, metadata filters, and pagination.
+
+        Returns:
+            List of retrieval results matching the search criteria.
+        """
+        db_query = select(self.config.document_cls)
+
+        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 query and return results
+        results = self.session.scalars(db_query).all()
+        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.
+            search_query: The search query containing the text to search for.
+
+        Returns:
+            Updated SQLAlchemy query with semantic search applied.
+        """
+        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.
+        Search against `BaseDocument.content`.
+        Args:
+            db_query: The base SQLAlchemy query.
+            search_query: The search query containing the text to search for.
+        Returns:
+            Updated SQLAlchemy query with keyword search applied.
+        """
+        if not search_query.keywords:
+            return db_query
+
+        conditions = []
+        for keyword in search_query.keywords:
+            conditions.append(self.config.document_cls.content.ilike(f"%{keyword}%"))
+        return db_query.where(or_(*conditions))
+
+    def _apply_metadata_filters(self, query: Select, search_query: SearchQuery) -> Select:
+        """Apply metadata filters to the query.
+
+        Args:
+            query: The base SQLAlchemy query.
+            search_query: The search query containing metadata filters.
+
+        Returns:
+            Updated SQLAlchemy query with metadata filters applied.
+        """
+        raise NotImplementedError
+
+    def _convert_to_retrieval_results(self, results: Sequence[Any]) -> list[RetrievalResult]:
+        """Convert database results to RetrievalResult objects.
+
+        Args:
+            results: Raw database results.
+            search_query: The original search query.
+
+        Returns:
+            List of RetrievalResult objects.
+        """
+        retrieval_results = []
+        for result in results:
+            doc = result[0] if isinstance(result, tuple) else result
+            retrieval_results.append(RetrievalResult(document=result, score=1.0))
+        return retrieval_results

pgvector_template-0.2.1/pgvector_template/service/__init__.py

@@ -0,0 +1,7 @@
+"""
+Service layer for document operations.
+"""
+
+from pgvector_template.service.document_service import DocumentService, DocumentServiceConfig
+
+__all__ = ["DocumentService", "DocumentServiceConfig"]

pgvector_template-0.2.1/pgvector_template/service/document_service.py

@@ -0,0 +1,101 @@
+"""
+Document service layer combining corpus management and search capabilities.
+"""
+
+from logging import getLogger
+from typing import Any, Generic, Type, TypeVar
+from uuid import UUID
+
+from pydantic import BaseModel, Field
+from sqlalchemy.orm import Session
+
+from pgvector_template.core.document import BaseDocument, BaseDocumentMetadata
+from pgvector_template.core.embedder import BaseEmbeddingProvider
+from pgvector_template.core.manager import BaseCorpusManager, BaseCorpusManagerConfig
+from pgvector_template.core.search import BaseSearchClient, BaseSearchClientConfig
+
+# Type variable for document type
+T = TypeVar("T", bound=BaseDocument)
+
+logger = getLogger(__name__)
+
+
+class DocumentServiceConfig(BaseModel):
+    """Configuration for DocumentService"""
+
+    # Required fields
+    document_cls: Type[BaseDocument] = Field(...)
+    """Document class **type** (not an instance). Must be subclass of `BaseDocument`."""
+    corpus_manager_cls: Type[BaseCorpusManager] = Field(default=BaseCorpusManager)
+    """CorpusManager class **type** (not an instance). Must be child of `BaseCorpusManager`."""
+    search_client_cls: Type[BaseSearchClient] = Field(default=BaseSearchClient)
+    """SearchClient class **type** (not an instance). Must be child of `BaseSearchClient`."""
+
+    # Optional fields with defaults
+    embedding_provider: BaseEmbeddingProvider | None = Field(default=None)
+    """Embedding provider for insert & vector-search operations."""
+    document_metadata_cls: Type[BaseDocumentMetadata] = Field(default=BaseDocumentMetadata)
+    """Document metadata schema. Must be child of `BaseDocumentMetadata`."""
+    corpus_manager_cfg: BaseCorpusManagerConfig = Field(
+        default=BaseCorpusManagerConfig(document_cls=BaseDocument)
+    )
+    """Instance of `BaseCorpusManagerConfig` or a child. Used to instantiate a CorpusManager."""
+    search_client_cfg: BaseSearchClientConfig = Field(
+        default=BaseSearchClientConfig(document_cls=BaseDocument)
+    )
+    """Instance of `BaseSearchClientConfig` or a child. Used to instantiate a SearchClient."""
+
+    def model_post_init(self, _):
+        # coerce document_cls onto corpus_manager_cfg & search_client_cfg.document_cls,
+        # 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
+        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}
+
+
+class DocumentService(Generic[T]):
+    """Service layer for document operations combining management and search capabilities"""
+
+    @property
+    def config(self) -> DocumentServiceConfig:
+        return self._cfg
+
+    @property
+    def corpus_manager(self) -> BaseCorpusManager:
+        """CorpusManager instance"""
+        return self._corpus_manager
+
+    @property
+    def search_client(self) -> BaseSearchClient:
+        """`SearchClient` instance"""
+        return self._search_client
+
+    def __init__(self, session: Session, config: DocumentServiceConfig):
+        self.session = session
+        self._cfg = config
+        self._setup()
+
+    def _setup(self):
+        """Initialize CorpusManager and SearchClient"""
+        self._corpus_manager = self._create_corpus_manager()
+        self._search_client = self._setup_search()
+
+    def _create_corpus_manager(self) -> BaseCorpusManager:
+        """Initialize CorpusManager. Override this to provide custom instantiation logic."""
+        return self.config.corpus_manager_cls(self.session, self.config.corpus_manager_cfg)
+
+    def _setup_search(self) -> BaseSearchClient:
+        """Initialize search client - to be implemented. Override this to provide custom instantiation logic."""
+        return self.config.search_client_cls(
+            self.session,
+            self.config.search_client_cfg,
+        )

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pgvector-template
-Version: 0.1.3
+Version: 0.2.1
 Summary: Template library for flexible PGVector RAG implementations
 Author-email: DL <v49t9zpqd@mozmail.com>
 License: MIT
@@ -20,6 +20,8 @@ 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"
+Provides-Extra: dev
+Requires-Dist: black>=23.0.0; extra == "dev"
 
 # PGVector-Template
 

{pgvector_template-0.1.3 → pgvector_template-0.2.1}/pgvector_template.egg-info/SOURCES.txt

@@ -12,10 +12,10 @@ 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/service/__init__.py
+pgvector_template/service/document_service.py
 pgvector_template/utils/__init__.py

{pgvector_template-0.1.3 → pgvector_template-0.2.1}/pgvector_template.egg-info/requires.txt

@@ -3,6 +3,9 @@ pydantic<3.0,>=2.11
 sqlalchemy>=2.0.0
 typing-extensions>=4.0.0
 
+[dev]
+black>=23.0.0
+
 [test]
 psycopg[binary]>=3.1.0
 pytest>=7.0.0

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pgvector-template"
-version = "0.1.3"
+version = "0.2.1"
 description = "Template library for flexible PGVector RAG implementations"
 authors = [{ name="DL", email="v49t9zpqd@mozmail.com" }]
 license = { text = "MIT" }
@@ -29,6 +29,9 @@ test = [
     "pytest-cov>=4.0.0",
     "python-dotenv>=1.0.0",
 ]
+dev = [
+    "black>=23.0.0",
+]
 
 [project.urls]
 Homepage = "https://github.com/DavidLiuGit/PGVector-Template"
@@ -40,5 +43,5 @@ include = ["pgvector_template*"]  # Include all packages starting with pgvector_
 
 
 [tool.black]
-line-length = 120
+line-length = 100
 target-version = ["py311", "py312", "py313"]

pgvector_template-0.1.3/pgvector_template/core/retriever.py

@@ -1,45 +0,0 @@
-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.3/pgvector_template/core/search.py

@@ -1,71 +0,0 @@
-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")