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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (26) hide show
  1. {pgvector_template-0.1.3 → pgvector_template-0.2.0}/PKG-INFO +3 -1
  2. {pgvector_template-0.1.3 → pgvector_template-0.2.0}/pgvector_template/core/__init__.py +3 -3
  3. {pgvector_template-0.1.3 → pgvector_template-0.2.0}/pgvector_template/core/manager.py +56 -31
  4. pgvector_template-0.2.0/pgvector_template/core/search.py +185 -0
  5. pgvector_template-0.2.0/pgvector_template/service/__init__.py +7 -0
  6. pgvector_template-0.2.0/pgvector_template/service/document_service.py +95 -0
  7. {pgvector_template-0.1.3 → pgvector_template-0.2.0}/pgvector_template.egg-info/PKG-INFO +3 -1
  8. {pgvector_template-0.1.3 → pgvector_template-0.2.0}/pgvector_template.egg-info/SOURCES.txt +2 -2
  9. {pgvector_template-0.1.3 → pgvector_template-0.2.0}/pgvector_template.egg-info/requires.txt +3 -0
  10. {pgvector_template-0.1.3 → pgvector_template-0.2.0}/pyproject.toml +5 -2
  11. pgvector_template-0.1.3/pgvector_template/core/retriever.py +0 -45
  12. pgvector_template-0.1.3/pgvector_template/core/search.py +0 -71
  13. pgvector_template-0.1.3/pgvector_template/utils/__init__.py +0 -0
  14. {pgvector_template-0.1.3 → pgvector_template-0.2.0}/LICENSE +0 -0
  15. {pgvector_template-0.1.3 → pgvector_template-0.2.0}/README.md +0 -0
  16. {pgvector_template-0.1.3 → pgvector_template-0.2.0}/pgvector_template/__init__.py +0 -0
  17. {pgvector_template-0.1.3 → pgvector_template-0.2.0}/pgvector_template/core/document.py +0 -0
  18. {pgvector_template-0.1.3 → pgvector_template-0.2.0}/pgvector_template/core/embedder.py +0 -0
  19. {pgvector_template-0.1.3 → pgvector_template-0.2.0}/pgvector_template/db/__init__.py +0 -0
  20. {pgvector_template-0.1.3 → pgvector_template-0.2.0}/pgvector_template/db/connection.py +0 -0
  21. {pgvector_template-0.1.3 → pgvector_template-0.2.0}/pgvector_template/db/document_db.py +0 -0
  22. {pgvector_template-0.1.3 → pgvector_template-0.2.0}/pgvector_template/types.py +0 -0
  23. {pgvector_template-0.1.3/pgvector_template/search → pgvector_template-0.2.0/pgvector_template/utils}/__init__.py +0 -0
  24. {pgvector_template-0.1.3 → pgvector_template-0.2.0}/pgvector_template.egg-info/dependency_links.txt +0 -0
  25. {pgvector_template-0.1.3 → pgvector_template-0.2.0}/pgvector_template.egg-info/top_level.txt +0 -0
  26. {pgvector_template-0.1.3 → pgvector_template-0.2.0}/setup.cfg +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: pgvector-template
3
- Version: 0.1.3
3
+ Version: 0.2.0
4
4
  Summary: Template library for flexible PGVector RAG implementations
5
5
  Author-email: DL <v49t9zpqd@mozmail.com>
6
6
  License: MIT
@@ -20,6 +20,8 @@ Requires-Dist: psycopg[binary]>=3.1.0; extra == "test"
20
20
  Requires-Dist: pytest>=7.0.0; extra == "test"
21
21
  Requires-Dist: pytest-cov>=4.0.0; extra == "test"
22
22
  Requires-Dist: python-dotenv>=1.0.0; extra == "test"
23
+ Provides-Extra: dev
24
+ Requires-Dist: black>=23.0.0; extra == "dev"
23
25
 
24
26
  # PGVector-Template
25
27
 
@@ -1,8 +1,7 @@
1
1
  from pgvector_template.core.document import BaseDocument, BaseDocumentMetadata, BaseDocumentOptionalProps
2
2
  from pgvector_template.core.embedder import BaseEmbeddingProvider
3
3
  from pgvector_template.core.manager import BaseCorpusManager, BaseCorpusManagerConfig, Corpus
4
- from pgvector_template.core.retriever import RetrievalResult, SearchQuery
5
- from pgvector_template.core.search import BaseSearchClient
4
+ from pgvector_template.core.search import SearchQuery, RetrievalResult, BaseSearchClientConfig, BaseSearchClient
6
5
 
7
6
 
8
7
  __all__ = [
@@ -17,7 +16,8 @@ __all__ = [
17
16
  "BaseCorpusManager",
18
17
  "BaseCorpusManagerConfig",
19
18
  ### search
20
- "RetrievalResult",
21
19
  "SearchQuery",
20
+ "RetrievalResult",
21
+ "BaseSearchClientConfig",
22
22
  "BaseSearchClient",
23
23
  ]
@@ -7,7 +7,11 @@ from uuid import UUID, uuid4
7
7
  from pydantic import BaseModel, Field
8
8
  from sqlalchemy.orm import Session
9
9
 
10
- from pgvector_template.core.document import BaseDocument, BaseDocumentMetadata, BaseDocumentOptionalProps
10
+ from pgvector_template.core.document import (
11
+ BaseDocument,
12
+ BaseDocumentMetadata,
13
+ BaseDocumentOptionalProps,
14
+ )
11
15
  from pgvector_template.core.embedder import BaseEmbeddingProvider
12
16
 
13
17
 
@@ -17,13 +21,12 @@ logger = getLogger(__name__)
17
21
  class BaseCorpusManagerConfig(BaseModel):
18
22
  """Base configuration for `Corpus` & `Document` management operations"""
19
23
 
20
- document_cls: Type[BaseDocument] = Field(..., description="Document class, must be subclass of BaseDocument")
21
- embedding_provider: BaseEmbeddingProvider | None = Field(
22
- default=None, description="Embedding provider for insert operations"
23
- )
24
- document_metadata_cls: Type[BaseDocumentMetadata] | None = Field(
25
- default=None, description="Document metadata class, must be subclass of BaseDocumentMetadata"
26
- )
24
+ document_cls: Type[BaseDocument] = Field(...)
25
+ """Document class **type** (not an instance). Must be subclass of `BaseDocument`."""
26
+ embedding_provider: BaseEmbeddingProvider | None = Field(default=None)
27
+ """Instance of `BaseEmbeddingProvider` child class. Acts as embedding provider for insert operations."""
28
+ document_metadata_cls: Type[BaseDocumentMetadata] = Field(default=BaseDocumentMetadata)
29
+ """Document metadata class **type** (not an instance). Must be subclass of BaseDocumentMetadata."""
27
30
 
28
31
  model_config = {"arbitrary_types_allowed": True}
29
32
 
@@ -55,6 +58,18 @@ class BaseCorpusManager(ABC):
55
58
  def config(self) -> BaseCorpusManagerConfig:
56
59
  return self._cfg
57
60
 
61
+ @property
62
+ def document_metadata_class(self) -> Type[BaseDocumentMetadata]:
63
+ """Returns the document metadata class, raising an error if it's not set."""
64
+ return self.config.document_metadata_cls
65
+
66
+ @property
67
+ def embedding_provider(self) -> BaseEmbeddingProvider:
68
+ """Returns the embedding provider, raising an error if it's not set."""
69
+ if self.config.embedding_provider is None:
70
+ raise ValueError("embedding_provider must be provided in config for this operation")
71
+ return self.config.embedding_provider
72
+
58
73
  def __init__(
59
74
  self,
60
75
  session: Session,
@@ -62,12 +77,19 @@ class BaseCorpusManager(ABC):
62
77
  ) -> None:
63
78
  self.session = session
64
79
  self._cfg = config
80
+ if not self.config.embedding_provider:
81
+ logger.warning(
82
+ "EmbeddingProvider not provided in config. Insertion will be unavailable."
83
+ )
65
84
 
66
- def get_full_corpus(self, corpus_id: str) -> Corpus | None:
85
+ def get_full_corpus(self, corpus_id: str, **kwargs) -> Corpus | None:
67
86
  """Reconstruct full corpus from its individual documents/chunks"""
68
87
  chunks = (
69
88
  self.session.query(self.config.document_cls)
70
- .filter(self.config.document_cls.corpus_id == corpus_id, self.config.document_cls.is_deleted == False)
89
+ .filter(
90
+ self.config.document_cls.corpus_id == corpus_id,
91
+ self.config.document_cls.is_deleted == False,
92
+ )
71
93
  .order_by(self.config.document_cls.chunk_index)
72
94
  .all()
73
95
  )
@@ -89,6 +111,7 @@ class BaseCorpusManager(ABC):
89
111
  corpus_metadata: dict[str, Any],
90
112
  optional_props: BaseDocumentOptionalProps | None = None,
91
113
  corpus_id: UUID | str | None = None,
114
+ **kwargs,
92
115
  ) -> int:
93
116
  """
94
117
  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):
102
125
  Returns:
103
126
  int: The number of **documents** inserted for the provided corpus
104
127
  """
105
- self._check_insert_dependencies()
106
-
107
128
  if not corpus_id:
108
129
  corpus_id = uuid4()
109
130
  document_contents = self._split_corpus(content)
110
- document_embeddings = self.config.embedding_provider.embed_batch(document_contents)
111
- return self.insert_documents(corpus_id, document_contents, document_embeddings, corpus_metadata, optional_props)
131
+ document_embeddings = self.embedding_provider.embed_batch(document_contents)
132
+ return self.insert_documents(
133
+ corpus_id, document_contents, document_embeddings, corpus_metadata, optional_props
134
+ )
112
135
 
113
136
  def insert_documents(
114
137
  self,
@@ -117,6 +140,7 @@ class BaseCorpusManager(ABC):
117
140
  document_embeddings: list[list[float]],
118
141
  corpus_metadata: dict[str, Any],
119
142
  optional_props: BaseDocumentOptionalProps | None = None,
143
+ **kwargs,
120
144
  ) -> int:
121
145
  """
122
146
  Insert a list of documents (usually from a chunked + embedded corpus).
@@ -134,8 +158,6 @@ class BaseCorpusManager(ABC):
134
158
  Raises:
135
159
  ValueError: If the length of document_contents doesn't match document_embeddings
136
160
  """
137
- self._check_insert_dependencies()
138
-
139
161
  if len(document_contents) != len(document_embeddings):
140
162
  raise ValueError("Number of embeddings does not match number of documents")
141
163
  if len(document_contents) == 0:
@@ -143,7 +165,7 @@ class BaseCorpusManager(ABC):
143
165
  documents_to_insert = []
144
166
  for i in range(len(document_contents)):
145
167
  chunk_md = self._extract_chunk_metadata(document_contents[i])
146
- base_metadata = self.config.document_metadata_cls(**(corpus_metadata | chunk_md))
168
+ base_metadata = self.document_metadata_class(**(corpus_metadata | chunk_md))
147
169
  documents_to_insert.append(
148
170
  self.config.document_cls.from_props(
149
171
  corpus_id=corpus_id,
@@ -164,11 +186,15 @@ class BaseCorpusManager(ABC):
164
186
  Split a corpus' string content into smaller chunks.
165
187
  """
166
188
  if self.__class__ is not BaseCorpusManager:
167
- logger.warning("Using default _split_corpus. Override this method to improve performance.")
189
+ logger.warning(
190
+ "Using default _split_corpus. Override this method to improve performance."
191
+ )
168
192
  split_content = [content[i : i + 1000] for i in range(0, len(content), 1000)]
169
193
  return [c for c in split_content if len(c.strip()) > 0]
170
194
 
171
- def _join_documents(self, documents: list[BaseDocument]) -> tuple[str, dict[str, Any]]:
195
+ def _join_documents(
196
+ self, documents: list[BaseDocument], **kwargs
197
+ ) -> tuple[str, dict[str, Any]]:
172
198
  """
173
199
  **It is highly recommended to override this method.**
174
200
  **This method should effectively reverse the `_split_corpus` method.**
@@ -177,43 +203,42 @@ class BaseCorpusManager(ABC):
177
203
  `BaseDocumentMetadata`/`document_metadata_cls` are relevant to the corpus.
178
204
  """
179
205
  if self.__class__ is not BaseCorpusManager:
180
- logger.warning("Using default _join_documents. Override this method to improve functionality.")
206
+ logger.warning(
207
+ "Using default _join_documents. Override this method to improve functionality."
208
+ )
181
209
  documents.sort(key=lambda d: d.chunk_index) # type: ignore
182
210
  # since _split_corpus performs a simple split on every 1000 chars, we can simply call `join`
183
211
  corpus_content = "".join(d.content for d in documents) # type: ignore
184
212
  corpus_metadata = self._infer_corpus_metadata(documents)
185
213
  return corpus_content, corpus_metadata
186
214
 
187
- def _extract_chunk_metadata(self, content: str) -> dict[str, Any]:
215
+ def _extract_chunk_metadata(self, content: str, **kwargs) -> dict[str, Any]:
188
216
  """
189
217
  **It is highly recommended to override this method.**
190
218
  Extract metadata from a chunk of content, to be appended to corpus metadata.
191
219
  Note: returning a key-value pair here does NOT guarantee its inclusion when it's added to the database.
192
220
  """
193
221
  if self.__class__ is not BaseCorpusManager:
194
- logger.warning("Using default _extract_chunk_metadata. It is highly recommended to override this method.")
222
+ logger.warning(
223
+ "Using default _extract_chunk_metadata. It is highly recommended to override this method."
224
+ )
195
225
  # this is simply an example. Since the document metadata that gets saved
196
226
  return {
197
227
  "chunk_length": len(content),
198
228
  }
199
229
 
200
- def _infer_corpus_metadata(self, documents: list[BaseDocument]) -> dict[str, Any]:
230
+ def _infer_corpus_metadata(self, documents: list[BaseDocument], **kwargs) -> dict[str, Any]:
201
231
  """
202
232
  **It is highly recommended to override this method.**
203
233
  **This method should be a best-effort reversal of `extract_chunk_metadata()`**
204
234
  Infer metadata for the corpus from the constituent `BaseDocument`s.
205
235
  """
206
236
  if self.__class__ is not BaseCorpusManager:
207
- logger.warning("Using default _infer_corpus_metadata. It is highly recommended to override this method.")
237
+ logger.warning(
238
+ "Using default _infer_corpus_metadata. It is highly recommended to override this method."
239
+ )
208
240
  # merge all document.document_metadata together, and return
209
241
  merged = {}
210
242
  for d in documents:
211
243
  merged.update(d.document_metadata) # type: ignore
212
244
  return merged
213
-
214
- def _check_insert_dependencies(self) -> None:
215
- """Check that required dependencies for insert operations are available"""
216
- if not self.config.embedding_provider:
217
- raise ValueError("embedding_provider must be provided in config for insert operations")
218
- if not self.config.document_metadata_cls:
219
- raise ValueError("document_metadata_cls must be provided in config for insert operations")
@@ -0,0 +1,185 @@
1
+ from dataclasses import dataclass, asdict
2
+ from datetime import datetime
3
+ from logging import getLogger
4
+ from typing import Any, Type, Sequence
5
+
6
+ from pydantic import BaseModel, Field, model_validator
7
+ from sqlalchemy import text, select, or_, and_
8
+ from sqlalchemy.sql import Select
9
+ from pgvector.sqlalchemy import Vector
10
+
11
+ from pgvector_template.core import (
12
+ BaseEmbeddingProvider,
13
+ BaseDocument,
14
+ BaseDocumentMetadata,
15
+ )
16
+ from sqlalchemy.orm import Session
17
+
18
+
19
+ logger = getLogger(__name__)
20
+
21
+
22
+ class SearchQuery(BaseModel):
23
+ """Standardized search query structure. At least 1 search criterion is required."""
24
+
25
+ text: str | None = None
26
+ """String to approximate-search (using vector distance) in a semantic search."""
27
+ keywords: list[str] | None = None
28
+ """List of keywords to exact-match in a keyword search."""
29
+ metadata_filters: dict[str, Any] | None = None
30
+ """Strict metadata filters that must be matched."""
31
+ date_range: tuple[datetime, datetime] | None = None
32
+ """Retrieve/limit results based on created_at & updated_at timestamps"""
33
+ limit: int = Field(
34
+ ...,
35
+ ge=1,
36
+ )
37
+ """Maximum number of results to return."""
38
+
39
+ @model_validator(mode="after")
40
+ def ensure_criterion(self):
41
+ if not any([self.text, self.keywords, self.metadata_filters, self.date_range]):
42
+ raise ValueError("At least one search criterion is required")
43
+ return self
44
+
45
+
46
+ @dataclass
47
+ class RetrievalResult:
48
+ """Standardized result structure for all retrieval operations"""
49
+
50
+ document: BaseDocument
51
+ score: float
52
+
53
+ def to_dict(self) -> dict[str, Any]:
54
+ return asdict(self)
55
+
56
+
57
+ class BaseSearchClientConfig(BaseModel):
58
+ """Config obj for `BaseSearchClient`."""
59
+
60
+ document_cls: Type[BaseDocument] = Field(default=BaseDocument)
61
+ """Document class **type** (not an instance). Must be subclass of `BaseDocument`."""
62
+ embedding_provider: BaseEmbeddingProvider | None = Field(default=None)
63
+ """Instance of `BaseEmbeddingProvider` child class. Acts as embedding provider for semantic search."""
64
+ document_metadata_cls: Type[BaseDocumentMetadata] = Field(default=BaseDocumentMetadata)
65
+ """Document metadata class type. Used for metadata search operations."""
66
+
67
+ model_config = {"arbitrary_types_allowed": True}
68
+
69
+
70
+ class BaseSearchClient:
71
+ """Minimum-viable implementation of document retrieval for PGVector"""
72
+
73
+ @property
74
+ def config(self) -> BaseSearchClientConfig:
75
+ return self._cfg
76
+
77
+ @property
78
+ def document_metadata_class(self) -> Type[BaseDocumentMetadata]:
79
+ """Returns the document metadata class, raising an error if it's not set."""
80
+ return self.config.document_metadata_cls
81
+
82
+ @property
83
+ def embedding_provider(self) -> BaseEmbeddingProvider:
84
+ """Returns the embedding provider, raising an error if it's not set."""
85
+ if self.config.embedding_provider is None:
86
+ raise ValueError("embedding_provider must be provided in config for this operation")
87
+ return self.config.embedding_provider
88
+
89
+ def __init__(
90
+ self,
91
+ session: Session,
92
+ config: BaseSearchClientConfig,
93
+ ):
94
+ self.session = session
95
+ self._cfg = config
96
+ if not self.config.embedding_provider:
97
+ logger.warning(
98
+ "EmbeddingProvider not provided in config. Vector (semantic) search will be unavailable."
99
+ )
100
+
101
+ def search(self, query: SearchQuery) -> list[RetrievalResult]:
102
+ """Search for documents based on the provided query.
103
+
104
+ Args:
105
+ query: Search query containing text, metadata filters, and pagination.
106
+
107
+ Returns:
108
+ List of retrieval results matching the search criteria.
109
+ """
110
+ db_query = select(self.config.document_cls)
111
+
112
+ # if query.text:
113
+ # db_query = self._apply_semantic_search(db_query, query)
114
+ db_query = self._apply_keyword_search(db_query, query)
115
+ # if query.metadata_filters:
116
+ # db_query = self._apply_metadata_filters(db_query, query)
117
+ db_query = db_query.limit(query.limit)
118
+
119
+ # Execute the query
120
+ results = self.session.scalars(db_query).all()
121
+
122
+ # Convert to RetrievalResult objects
123
+ return self._convert_to_retrieval_results(results)
124
+
125
+ def _apply_semantic_search(self, query: Select, search_query: SearchQuery) -> Select:
126
+ """Apply semantic (vector) search criteria to the query.
127
+ `embedding_provider` must be provided at instantiation, or an `ValueError` will be raised.
128
+
129
+ Args:
130
+ query: The base SQLAlchemy query.
131
+ search_query: The search query containing the text to search for.
132
+
133
+ Returns:
134
+ Updated SQLAlchemy query with semantic search applied.
135
+ """
136
+ if not self.config.embedding_provider:
137
+ raise ValueError(
138
+ "EmbeddingProvider not provided in config. Vector (semantic) search is unavailable."
139
+ )
140
+ raise NotImplementedError
141
+
142
+ def _apply_keyword_search(self, db_query: Select, search_query: SearchQuery) -> Select:
143
+ """Apply keyword (full-text) search criteria to the query.
144
+ Search against `BaseDocument.content`.
145
+ Args:
146
+ db_query: The base SQLAlchemy query.
147
+ search_query: The search query containing the text to search for.
148
+ Returns:
149
+ Updated SQLAlchemy query with keyword search applied.
150
+ """
151
+ if not search_query.keywords:
152
+ return db_query
153
+
154
+ conditions = []
155
+ for keyword in search_query.keywords:
156
+ conditions.append(self.config.document_cls.content.ilike(f"%{keyword}%"))
157
+ return db_query.where(or_(*conditions))
158
+
159
+ def _apply_metadata_filters(self, query: Select, search_query: SearchQuery) -> Select:
160
+ """Apply metadata filters to the query.
161
+
162
+ Args:
163
+ query: The base SQLAlchemy query.
164
+ search_query: The search query containing metadata filters.
165
+
166
+ Returns:
167
+ Updated SQLAlchemy query with metadata filters applied.
168
+ """
169
+ raise NotImplementedError
170
+
171
+ def _convert_to_retrieval_results(self, results: Sequence[Any]) -> list[RetrievalResult]:
172
+ """Convert database results to RetrievalResult objects.
173
+
174
+ Args:
175
+ results: Raw database results.
176
+ search_query: The original search query.
177
+
178
+ Returns:
179
+ List of RetrievalResult objects.
180
+ """
181
+ retrieval_results = []
182
+ for result in results:
183
+ doc = result[0] if isinstance(result, tuple) else result
184
+ retrieval_results.append(RetrievalResult(document=result, score=1.0))
185
+ return retrieval_results
@@ -0,0 +1,7 @@
1
+ """
2
+ Service layer for document operations.
3
+ """
4
+
5
+ from pgvector_template.service.document_service import DocumentService, DocumentServiceConfig
6
+
7
+ __all__ = ["DocumentService", "DocumentServiceConfig"]
@@ -0,0 +1,95 @@
1
+ """
2
+ Document service layer combining corpus management and search capabilities.
3
+ """
4
+
5
+ from logging import getLogger
6
+ from typing import Any, Generic, Type, TypeVar
7
+ from uuid import UUID
8
+
9
+ from pydantic import BaseModel, Field
10
+ from sqlalchemy.orm import Session
11
+
12
+ from pgvector_template.core.document import BaseDocument, BaseDocumentMetadata
13
+ from pgvector_template.core.embedder import BaseEmbeddingProvider
14
+ from pgvector_template.core.manager import BaseCorpusManager, BaseCorpusManagerConfig
15
+ from pgvector_template.core.search import BaseSearchClient, BaseSearchClientConfig
16
+
17
+ # Type variable for document type
18
+ T = TypeVar("T", bound=BaseDocument)
19
+
20
+ logger = getLogger(__name__)
21
+
22
+
23
+ class DocumentServiceConfig(BaseModel):
24
+ """Configuration for DocumentService"""
25
+
26
+ # Required fields
27
+ document_cls: Type[BaseDocument] = Field(...)
28
+ """Document class **type** (not an instance). Must be subclass of `BaseDocument`."""
29
+ corpus_manager_cls: Type[BaseCorpusManager] = Field(default=BaseCorpusManager)
30
+ """CorpusManager class **type** (not an instance). Must be child of `BaseCorpusManager`."""
31
+ search_client_cls: Type[BaseSearchClient] = Field(default=BaseSearchClient)
32
+ """SearchClient class **type** (not an instance). Must be child of `BaseSearchClient`."""
33
+
34
+ # Optional fields with defaults
35
+ embedding_provider: BaseEmbeddingProvider | None = Field(default=None)
36
+ """Embedding provider for insert & vector-search operations."""
37
+ document_metadata_cls: Type[BaseDocumentMetadata] = Field(default=BaseDocumentMetadata)
38
+ """Document metadata schema. Must be child of `BaseDocumentMetadata`."""
39
+ corpus_manager_cfg: BaseCorpusManagerConfig = Field(
40
+ default=BaseCorpusManagerConfig(document_cls=BaseDocument)
41
+ )
42
+ """Instance of `BaseCorpusManagerConfig` or a child. Used to instantiate a CorpusManager."""
43
+ search_client_cfg: BaseSearchClientConfig = Field(
44
+ default=BaseSearchClientConfig(document_cls=BaseDocument)
45
+ )
46
+ """Instance of `BaseSearchClientConfig` or a child. Used to instantiate a SearchClient."""
47
+
48
+ def model_post_init(self, _):
49
+ # coerce document_cls onto corpus_manager_cfg & search_client_cfg.document_cls,
50
+ # iff either config is an instance of their respective base config classes (and not a subclass)
51
+ if type(self.corpus_manager_cfg) is BaseCorpusManagerConfig:
52
+ self.corpus_manager_cfg.document_cls = self.document_cls
53
+ if type(self.search_client_cfg) is BaseSearchClientConfig:
54
+ self.search_client_cfg.document_cls = self.document_cls
55
+
56
+ model_config = {"arbitrary_types_allowed": True}
57
+
58
+
59
+ class DocumentService(Generic[T]):
60
+ """Service layer for document operations combining management and search capabilities"""
61
+
62
+ @property
63
+ def config(self) -> DocumentServiceConfig:
64
+ return self._cfg
65
+
66
+ @property
67
+ def corpus_manager(self) -> BaseCorpusManager:
68
+ """CorpusManager instance"""
69
+ return self._corpus_manager
70
+
71
+ @property
72
+ def search_client(self) -> BaseSearchClient:
73
+ """`SearchClient` instance"""
74
+ return self._search_client
75
+
76
+ def __init__(self, session: Session, config: DocumentServiceConfig):
77
+ self.session = session
78
+ self._cfg = config
79
+ self._setup()
80
+
81
+ def _setup(self):
82
+ """Initialize CorpusManager and SearchClient"""
83
+ self._corpus_manager = self._create_corpus_manager()
84
+ self._search_client = self._setup_search()
85
+
86
+ def _create_corpus_manager(self) -> BaseCorpusManager:
87
+ """Initialize CorpusManager. Override this to provide custom instantiation logic."""
88
+ return self.config.corpus_manager_cls(self.session, self.config.corpus_manager_cfg)
89
+
90
+ def _setup_search(self) -> BaseSearchClient:
91
+ """Initialize search client - to be implemented. Override this to provide custom instantiation logic."""
92
+ return self.config.search_client_cls(
93
+ self.session,
94
+ self.config.search_client_cfg,
95
+ )
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: pgvector-template
3
- Version: 0.1.3
3
+ Version: 0.2.0
4
4
  Summary: Template library for flexible PGVector RAG implementations
5
5
  Author-email: DL <v49t9zpqd@mozmail.com>
6
6
  License: MIT
@@ -20,6 +20,8 @@ Requires-Dist: psycopg[binary]>=3.1.0; extra == "test"
20
20
  Requires-Dist: pytest>=7.0.0; extra == "test"
21
21
  Requires-Dist: pytest-cov>=4.0.0; extra == "test"
22
22
  Requires-Dist: python-dotenv>=1.0.0; extra == "test"
23
+ Provides-Extra: dev
24
+ Requires-Dist: black>=23.0.0; extra == "dev"
23
25
 
24
26
  # PGVector-Template
25
27
 
@@ -12,10 +12,10 @@ pgvector_template/core/__init__.py
12
12
  pgvector_template/core/document.py
13
13
  pgvector_template/core/embedder.py
14
14
  pgvector_template/core/manager.py
15
- pgvector_template/core/retriever.py
16
15
  pgvector_template/core/search.py
17
16
  pgvector_template/db/__init__.py
18
17
  pgvector_template/db/connection.py
19
18
  pgvector_template/db/document_db.py
20
- pgvector_template/search/__init__.py
19
+ pgvector_template/service/__init__.py
20
+ pgvector_template/service/document_service.py
21
21
  pgvector_template/utils/__init__.py
@@ -3,6 +3,9 @@ pydantic<3.0,>=2.11
3
3
  sqlalchemy>=2.0.0
4
4
  typing-extensions>=4.0.0
5
5
 
6
+ [dev]
7
+ black>=23.0.0
8
+
6
9
  [test]
7
10
  psycopg[binary]>=3.1.0
8
11
  pytest>=7.0.0
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "pgvector-template"
7
- version = "0.1.3"
7
+ version = "0.2.0"
8
8
  description = "Template library for flexible PGVector RAG implementations"
9
9
  authors = [{ name="DL", email="v49t9zpqd@mozmail.com" }]
10
10
  license = { text = "MIT" }
@@ -29,6 +29,9 @@ test = [
29
29
  "pytest-cov>=4.0.0",
30
30
  "python-dotenv>=1.0.0",
31
31
  ]
32
+ dev = [
33
+ "black>=23.0.0",
34
+ ]
32
35
 
33
36
  [project.urls]
34
37
  Homepage = "https://github.com/DavidLiuGit/PGVector-Template"
@@ -40,5 +43,5 @@ include = ["pgvector_template*"] # Include all packages starting with pgvector_
40
43
 
41
44
 
42
45
  [tool.black]
43
- line-length = 120
46
+ line-length = 100
44
47
  target-version = ["py311", "py312", "py313"]
@@ -1,45 +0,0 @@
1
- from dataclasses import dataclass, asdict
2
- from datetime import datetime
3
- from enum import Enum
4
- from typing import Any, Optional
5
-
6
-
7
- # class RetrievalStrategy(Enum):
8
- # """Supported retrieval strategies"""
9
- # SIMILARITY = "similarity" # Pure vector similarity
10
- # KEYWORD = "keyword" # Metadata-based keyword search
11
- # HYBRID = "hybrid" # Combined similarity + keyword
12
- # TEMPORAL = "temporal" # Time-based search (if applicable)
13
-
14
-
15
- @dataclass
16
- class RetrievalResult:
17
- """Standardized result structure for all retrieval operations"""
18
-
19
- document_id: str
20
- original_id: str
21
- chunk_index: int
22
- content: str
23
- title: Optional[str]
24
- metadata: dict[str, Any]
25
- similarity_score: Optional[float] = None
26
- keyword_matches: Optional[list[str]] = None
27
-
28
- def to_dict(self) -> dict[str, Any]:
29
- return asdict(self)
30
-
31
-
32
- @dataclass
33
- class SearchQuery:
34
- """Standardized search query structure"""
35
-
36
- text: Optional[str] = None
37
- keywords: Optional[list[str]] = None
38
- metadata_filters: Optional[dict[str, Any]] = None
39
- date_range: Optional[tuple[datetime, datetime]] = None
40
- limit: int = 10
41
- # strategy: RetrievalStrategy = RetrievalStrategy.HYBRID
42
-
43
- def validate(self) -> bool:
44
- """Ensure query has at least one search criterion"""
45
- return any([self.text, self.keywords, self.metadata_filters, self.date_range])
@@ -1,71 +0,0 @@
1
- from abc import ABC, abstractmethod
2
- from logging import getLogger
3
- from typing import Any, Generic, Optional
4
-
5
- from sqlalchemy import text
6
-
7
- from pgvector_template.core import BaseEmbeddingProvider, RetrievalResult, SearchQuery, BaseDocument
8
- from pgvector_template.types import DocumentType
9
- from sqlalchemy.orm import Session
10
-
11
-
12
- class BaseSearchClient(ABC, Generic[DocumentType]):
13
- """Abstract base for document retrieval"""
14
-
15
- def __init__(self, session: Session, embedding_provider: BaseEmbeddingProvider):
16
- self.session = session
17
- self.embedding_provider = embedding_provider
18
- self.logger = getLogger(self.__class__.__name__)
19
-
20
- @abstractmethod
21
- def search(self, query: SearchQuery) -> list[RetrievalResult]:
22
- """Main search interface"""
23
- raise NotImplementedError("Subclasses must implement this method")
24
-
25
- def search_by_metadata(self, filters: dict[str, Any], limit: int = 10) -> list[BaseDocument]:
26
- """Generic JSON-based metadata search"""
27
- query = self.session.query(BaseDocument).filter(BaseDocument.is_deleted == False)
28
-
29
- # Apply JSON-based filters
30
- for key, value in filters.items():
31
- if isinstance(value, list):
32
- # Array contains search
33
- query = query.filter(text(f"metadata->>'{key}' = ANY(:value)")).params(value=value)
34
- elif isinstance(value, dict):
35
- # Nested JSON search
36
- for nested_key, nested_value in value.items():
37
- query = query.filter(text(f"metadata->'{key}'->>'{nested_key}' = :value")).params(
38
- value=nested_value
39
- )
40
- else:
41
- # Simple equality
42
- query = query.filter(text(f"metadata->>'{key}' = :value")).params(value=str(value))
43
-
44
- return query.limit(limit).all()
45
-
46
- @abstractmethod
47
- def get_full_document(self, original_id: str) -> Optional[dict[str, Any]]:
48
- """Reconstruct full document from chunks"""
49
- raise NotImplementedError("Subclasses must implement this method")
50
-
51
- # Template methods with default implementations
52
- def similarity_search(self, text: str, limit: int = 10) -> list[RetrievalResult]:
53
- """Vector similarity search"""
54
- if not text.strip():
55
- return []
56
-
57
- embedding = self.embedding_provider.embed_text(text)
58
- return self._vector_search(embedding, limit)
59
-
60
- def keyword_search(self, keywords: list[str], limit: int = 10) -> list[RetrievalResult]:
61
- """Metadata-based keyword search"""
62
- return self._metadata_search({"keywords": keywords}, limit)
63
-
64
- def _vector_search(self, embedding: list[float], limit: int) -> list[RetrievalResult]:
65
- """Internal vector similarity search"""
66
- # This will be implemented by the template using the DocumentType
67
- raise NotImplementedError("Subclasses must implement vector search")
68
-
69
- def _metadata_search(self, filters: dict[str, Any], limit: int) -> list[RetrievalResult]:
70
- """Internal metadata search"""
71
- raise NotImplementedError("Subclasses must implement metadata search")