pgvector-haystack 0.4.0__tar.gz → 0.5.1__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.
- {pgvector_haystack-0.4.0 → pgvector_haystack-0.5.1}/.gitignore +9 -0
- {pgvector_haystack-0.4.0 → pgvector_haystack-0.5.1}/CHANGELOG.md +2 -1
- {pgvector_haystack-0.4.0 → pgvector_haystack-0.5.1}/PKG-INFO +1 -1
- {pgvector_haystack-0.4.0 → pgvector_haystack-0.5.1}/pyproject.toml +24 -51
- {pgvector_haystack-0.4.0 → pgvector_haystack-0.5.1}/src/haystack_integrations/components/retrievers/pgvector/embedding_retriever.py +17 -4
- {pgvector_haystack-0.4.0 → pgvector_haystack-0.5.1}/src/haystack_integrations/components/retrievers/pgvector/keyword_retriever.py +18 -4
- {pgvector_haystack-0.4.0 → pgvector_haystack-0.5.1}/tests/test_retrievers.py +93 -0
- {pgvector_haystack-0.4.0 → pgvector_haystack-0.5.1}/LICENSE.txt +0 -0
- {pgvector_haystack-0.4.0 → pgvector_haystack-0.5.1}/README.md +0 -0
- {pgvector_haystack-0.4.0 → pgvector_haystack-0.5.1}/examples/embedding_retrieval.py +0 -0
- {pgvector_haystack-0.4.0 → pgvector_haystack-0.5.1}/examples/hybrid_retrieval.py +0 -0
- {pgvector_haystack-0.4.0 → pgvector_haystack-0.5.1}/pydoc/config.yml +0 -0
- {pgvector_haystack-0.4.0 → pgvector_haystack-0.5.1}/src/haystack_integrations/components/retrievers/pgvector/__init__.py +0 -0
- {pgvector_haystack-0.4.0 → pgvector_haystack-0.5.1}/src/haystack_integrations/document_stores/pgvector/__init__.py +0 -0
- {pgvector_haystack-0.4.0 → pgvector_haystack-0.5.1}/src/haystack_integrations/document_stores/pgvector/document_store.py +0 -0
- {pgvector_haystack-0.4.0 → pgvector_haystack-0.5.1}/src/haystack_integrations/document_stores/pgvector/filters.py +0 -0
- {pgvector_haystack-0.4.0 → pgvector_haystack-0.5.1}/tests/__init__.py +0 -0
- {pgvector_haystack-0.4.0 → pgvector_haystack-0.5.1}/tests/conftest.py +0 -0
- {pgvector_haystack-0.4.0 → pgvector_haystack-0.5.1}/tests/test_document_store.py +0 -0
- {pgvector_haystack-0.4.0 → pgvector_haystack-0.5.1}/tests/test_embedding_retrieval.py +0 -0
- {pgvector_haystack-0.4.0 → pgvector_haystack-0.5.1}/tests/test_filters.py +0 -0
- {pgvector_haystack-0.4.0 → pgvector_haystack-0.5.1}/tests/test_keyword_retrieval.py +0 -0
|
@@ -1,10 +1,11 @@
|
|
|
1
1
|
# Changelog
|
|
2
2
|
|
|
3
|
-
## [
|
|
3
|
+
## [integrations/pgvector-v0.4.0] - 2024-06-20
|
|
4
4
|
|
|
5
5
|
### 🚀 Features
|
|
6
6
|
|
|
7
7
|
- Defer the database connection to when it's needed (#773)
|
|
8
|
+
- Add customizable index names for pgvector (#818)
|
|
8
9
|
|
|
9
10
|
## [integrations/pgvector-v0.2.0] - 2024-05-08
|
|
10
11
|
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.3
|
|
2
2
|
Name: pgvector-haystack
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.5.1
|
|
4
4
|
Summary: An integration of pgvector (vector search extension for Postgres) with Haystack
|
|
5
5
|
Project-URL: Source, https://github.com/deepset-ai/haystack-core-integrations
|
|
6
6
|
Project-URL: Documentation, https://github.com/deepset-ai/haystack-core-integrations/blob/main/integrations/pgvector/README.md
|
|
@@ -10,9 +10,7 @@ readme = "README.md"
|
|
|
10
10
|
requires-python = ">=3.8"
|
|
11
11
|
license = "Apache-2.0"
|
|
12
12
|
keywords = []
|
|
13
|
-
authors = [
|
|
14
|
-
{ name = "deepset GmbH", email = "info@deepset.ai" },
|
|
15
|
-
]
|
|
13
|
+
authors = [{ name = "deepset GmbH", email = "info@deepset.ai" }]
|
|
16
14
|
classifiers = [
|
|
17
15
|
"License :: OSI Approved :: Apache Software License",
|
|
18
16
|
"Development Status :: 4 - Beta",
|
|
@@ -25,11 +23,7 @@ classifiers = [
|
|
|
25
23
|
"Programming Language :: Python :: Implementation :: CPython",
|
|
26
24
|
"Programming Language :: Python :: Implementation :: PyPy",
|
|
27
25
|
]
|
|
28
|
-
dependencies = [
|
|
29
|
-
"haystack-ai",
|
|
30
|
-
"pgvector",
|
|
31
|
-
"psycopg[binary]"
|
|
32
|
-
]
|
|
26
|
+
dependencies = ["haystack-ai", "pgvector", "psycopg[binary]"]
|
|
33
27
|
|
|
34
28
|
[project.urls]
|
|
35
29
|
Source = "https://github.com/deepset-ai/haystack-core-integrations"
|
|
@@ -51,49 +45,28 @@ git_describe_command = 'git describe --tags --match="integrations/pgvector-v[0-9
|
|
|
51
45
|
dependencies = [
|
|
52
46
|
"coverage[toml]>=6.5",
|
|
53
47
|
"pytest",
|
|
48
|
+
"pytest-rerunfailures",
|
|
54
49
|
"ipython",
|
|
55
50
|
"haystack-pydoc-tools",
|
|
56
51
|
]
|
|
57
52
|
[tool.hatch.envs.default.scripts]
|
|
58
|
-
test = "pytest {args:tests}"
|
|
59
|
-
test-cov = "coverage run -m pytest {args:tests}"
|
|
60
|
-
cov-report = [
|
|
61
|
-
|
|
62
|
-
|
|
63
|
-
]
|
|
64
|
-
cov = [
|
|
65
|
-
"test-cov",
|
|
66
|
-
"cov-report",
|
|
67
|
-
]
|
|
68
|
-
docs = [
|
|
69
|
-
"pydoc-markdown pydoc/config.yml"
|
|
70
|
-
]
|
|
53
|
+
test = "pytest --reruns 3 --reruns-delay 30 -x {args:tests}"
|
|
54
|
+
test-cov = "coverage run -m pytest --reruns 3 --reruns-delay 30 -x {args:tests}"
|
|
55
|
+
cov-report = ["- coverage combine", "coverage report"]
|
|
56
|
+
cov = ["test-cov", "cov-report"]
|
|
57
|
+
docs = ["pydoc-markdown pydoc/config.yml"]
|
|
71
58
|
|
|
72
59
|
[[tool.hatch.envs.all.matrix]]
|
|
73
60
|
python = ["3.8", "3.9", "3.10", "3.11", "3.12"]
|
|
74
61
|
|
|
75
62
|
[tool.hatch.envs.lint]
|
|
76
63
|
detached = true
|
|
77
|
-
dependencies = [
|
|
78
|
-
"black>=23.1.0",
|
|
79
|
-
"mypy>=1.0.0",
|
|
80
|
-
"ruff>=0.0.243",
|
|
81
|
-
]
|
|
64
|
+
dependencies = ["black>=23.1.0", "mypy>=1.0.0", "ruff>=0.0.243"]
|
|
82
65
|
[tool.hatch.envs.lint.scripts]
|
|
83
66
|
typing = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}"
|
|
84
|
-
style = [
|
|
85
|
-
|
|
86
|
-
|
|
87
|
-
]
|
|
88
|
-
fmt = [
|
|
89
|
-
"black {args:.}",
|
|
90
|
-
"ruff --fix {args:.}",
|
|
91
|
-
"style",
|
|
92
|
-
]
|
|
93
|
-
all = [
|
|
94
|
-
"style",
|
|
95
|
-
"typing",
|
|
96
|
-
]
|
|
67
|
+
style = ["ruff check {args:.}", "black --check --diff {args:.}"]
|
|
68
|
+
fmt = ["black {args:.}", "ruff --fix {args:.}", "style"]
|
|
69
|
+
all = ["style", "typing"]
|
|
97
70
|
|
|
98
71
|
[tool.black]
|
|
99
72
|
target-version = ["py38"]
|
|
@@ -136,9 +109,15 @@ ignore = [
|
|
|
136
109
|
# Allow boolean positional values in function calls, like `dict.get(... True)`
|
|
137
110
|
"FBT003",
|
|
138
111
|
# Ignore checks for possible passwords
|
|
139
|
-
"S105",
|
|
112
|
+
"S105",
|
|
113
|
+
"S106",
|
|
114
|
+
"S107",
|
|
140
115
|
# Ignore complexity
|
|
141
|
-
"C901",
|
|
116
|
+
"C901",
|
|
117
|
+
"PLR0911",
|
|
118
|
+
"PLR0912",
|
|
119
|
+
"PLR0913",
|
|
120
|
+
"PLR0915",
|
|
142
121
|
# ignore function-call-in-default-argument
|
|
143
122
|
"B008",
|
|
144
123
|
]
|
|
@@ -167,17 +146,11 @@ parallel = false
|
|
|
167
146
|
|
|
168
147
|
[tool.coverage.report]
|
|
169
148
|
omit = ["*/tests/*", "*/__init__.py"]
|
|
170
|
-
show_missing=true
|
|
171
|
-
exclude_lines = [
|
|
172
|
-
"no cov",
|
|
173
|
-
"if __name__ == .__main__.:",
|
|
174
|
-
"if TYPE_CHECKING:",
|
|
175
|
-
]
|
|
149
|
+
show_missing = true
|
|
150
|
+
exclude_lines = ["no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:"]
|
|
176
151
|
|
|
177
152
|
[tool.pytest.ini_options]
|
|
178
|
-
markers = [
|
|
179
|
-
"integration: integration tests"
|
|
180
|
-
]
|
|
153
|
+
markers = ["integration: integration tests"]
|
|
181
154
|
|
|
182
155
|
|
|
183
156
|
[[tool.mypy.overrides]]
|
|
@@ -186,6 +159,6 @@ module = [
|
|
|
186
159
|
"haystack_integrations.*",
|
|
187
160
|
"pgvector.*",
|
|
188
161
|
"psycopg.*",
|
|
189
|
-
"pytest.*"
|
|
162
|
+
"pytest.*",
|
|
190
163
|
]
|
|
191
164
|
ignore_missing_imports = true
|
|
@@ -1,10 +1,12 @@
|
|
|
1
1
|
# SPDX-FileCopyrightText: 2023-present deepset GmbH <info@deepset.ai>
|
|
2
2
|
#
|
|
3
3
|
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
-
from typing import Any, Dict, List, Literal, Optional
|
|
4
|
+
from typing import Any, Dict, List, Literal, Optional, Union
|
|
5
5
|
|
|
6
6
|
from haystack import component, default_from_dict, default_to_dict
|
|
7
7
|
from haystack.dataclasses import Document
|
|
8
|
+
from haystack.document_stores.types import FilterPolicy
|
|
9
|
+
from haystack.document_stores.types.filter_policy import apply_filter_policy
|
|
8
10
|
from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
|
|
9
11
|
from haystack_integrations.document_stores.pgvector.document_store import VALID_VECTOR_FUNCTIONS
|
|
10
12
|
|
|
@@ -62,6 +64,7 @@ class PgvectorEmbeddingRetriever:
|
|
|
62
64
|
filters: Optional[Dict[str, Any]] = None,
|
|
63
65
|
top_k: int = 10,
|
|
64
66
|
vector_function: Optional[Literal["cosine_similarity", "inner_product", "l2_distance"]] = None,
|
|
67
|
+
filter_policy: Union[str, FilterPolicy] = FilterPolicy.REPLACE,
|
|
65
68
|
):
|
|
66
69
|
"""
|
|
67
70
|
:param document_store: An instance of `PgvectorDocumentStore`.
|
|
@@ -75,7 +78,7 @@ class PgvectorEmbeddingRetriever:
|
|
|
75
78
|
and the most similar documents are the ones with the smallest score.
|
|
76
79
|
**Important**: if the document store is using the `"hnsw"` search strategy, the vector function
|
|
77
80
|
should match the one utilized during index creation to take advantage of the index.
|
|
78
|
-
|
|
81
|
+
:param filter_policy: Policy to determine how filters are applied.
|
|
79
82
|
:raises ValueError: If `document_store` is not an instance of `PgvectorDocumentStore` or if `vector_function`
|
|
80
83
|
is not one of the valid options.
|
|
81
84
|
"""
|
|
@@ -91,6 +94,9 @@ class PgvectorEmbeddingRetriever:
|
|
|
91
94
|
self.filters = filters or {}
|
|
92
95
|
self.top_k = top_k
|
|
93
96
|
self.vector_function = vector_function or document_store.vector_function
|
|
97
|
+
self.filter_policy = (
|
|
98
|
+
filter_policy if isinstance(filter_policy, FilterPolicy) else FilterPolicy.from_str(filter_policy)
|
|
99
|
+
)
|
|
94
100
|
|
|
95
101
|
def to_dict(self) -> Dict[str, Any]:
|
|
96
102
|
"""
|
|
@@ -104,6 +110,7 @@ class PgvectorEmbeddingRetriever:
|
|
|
104
110
|
filters=self.filters,
|
|
105
111
|
top_k=self.top_k,
|
|
106
112
|
vector_function=self.vector_function,
|
|
113
|
+
filter_policy=self.filter_policy.value,
|
|
107
114
|
document_store=self.document_store.to_dict(),
|
|
108
115
|
)
|
|
109
116
|
|
|
@@ -119,6 +126,10 @@ class PgvectorEmbeddingRetriever:
|
|
|
119
126
|
"""
|
|
120
127
|
doc_store_params = data["init_parameters"]["document_store"]
|
|
121
128
|
data["init_parameters"]["document_store"] = PgvectorDocumentStore.from_dict(doc_store_params)
|
|
129
|
+
# Pipelines serialized with old versions of the component might not
|
|
130
|
+
# have the filter_policy field.
|
|
131
|
+
if filter_policy := data["init_parameters"].get("filter_policy"):
|
|
132
|
+
data["init_parameters"]["filter_policy"] = FilterPolicy.from_str(filter_policy)
|
|
122
133
|
return default_from_dict(cls, data)
|
|
123
134
|
|
|
124
135
|
@component.output_types(documents=List[Document])
|
|
@@ -133,13 +144,15 @@ class PgvectorEmbeddingRetriever:
|
|
|
133
144
|
Retrieve documents from the `PgvectorDocumentStore`, based on their embeddings.
|
|
134
145
|
|
|
135
146
|
:param query_embedding: Embedding of the query.
|
|
136
|
-
:param filters: Filters applied to the retrieved Documents.
|
|
147
|
+
:param filters: Filters applied to the retrieved Documents. The way runtime filters are applied depends on
|
|
148
|
+
the `filter_policy` chosen at retriever initialization. See init method docstring for more
|
|
149
|
+
details.
|
|
137
150
|
:param top_k: Maximum number of Documents to return.
|
|
138
151
|
:param vector_function: The similarity function to use when searching for similar embeddings.
|
|
139
152
|
|
|
140
153
|
:returns: List of Documents similar to `query_embedding`.
|
|
141
154
|
"""
|
|
142
|
-
filters =
|
|
155
|
+
filters = apply_filter_policy(self.filter_policy, self.filters, filters)
|
|
143
156
|
top_k = top_k or self.top_k
|
|
144
157
|
vector_function = vector_function or self.vector_function
|
|
145
158
|
|
|
@@ -1,10 +1,12 @@
|
|
|
1
1
|
# SPDX-FileCopyrightText: 2023-present deepset GmbH <info@deepset.ai>
|
|
2
2
|
#
|
|
3
3
|
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
-
from typing import Any, Dict, List, Optional
|
|
4
|
+
from typing import Any, Dict, List, Optional, Union
|
|
5
5
|
|
|
6
6
|
from haystack import component, default_from_dict, default_to_dict
|
|
7
7
|
from haystack.dataclasses import Document
|
|
8
|
+
from haystack.document_stores.types import FilterPolicy
|
|
9
|
+
from haystack.document_stores.types.filter_policy import apply_filter_policy
|
|
8
10
|
from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
|
|
9
11
|
|
|
10
12
|
|
|
@@ -51,12 +53,13 @@ class PgvectorKeywordRetriever:
|
|
|
51
53
|
document_store: PgvectorDocumentStore,
|
|
52
54
|
filters: Optional[Dict[str, Any]] = None,
|
|
53
55
|
top_k: int = 10,
|
|
56
|
+
filter_policy: Union[str, FilterPolicy] = FilterPolicy.REPLACE,
|
|
54
57
|
):
|
|
55
58
|
"""
|
|
56
59
|
:param document_store: An instance of `PgvectorDocumentStore`.
|
|
57
60
|
:param filters: Filters applied to the retrieved Documents.
|
|
58
61
|
:param top_k: Maximum number of Documents to return.
|
|
59
|
-
|
|
62
|
+
:param filter_policy: Policy to determine how filters are applied.
|
|
60
63
|
:raises ValueError: If `document_store` is not an instance of `PgvectorDocumentStore`.
|
|
61
64
|
"""
|
|
62
65
|
if not isinstance(document_store, PgvectorDocumentStore):
|
|
@@ -66,6 +69,9 @@ class PgvectorKeywordRetriever:
|
|
|
66
69
|
self.document_store = document_store
|
|
67
70
|
self.filters = filters or {}
|
|
68
71
|
self.top_k = top_k
|
|
72
|
+
self.filter_policy = (
|
|
73
|
+
filter_policy if isinstance(filter_policy, FilterPolicy) else FilterPolicy.from_str(filter_policy)
|
|
74
|
+
)
|
|
69
75
|
|
|
70
76
|
def to_dict(self) -> Dict[str, Any]:
|
|
71
77
|
"""
|
|
@@ -78,6 +84,7 @@ class PgvectorKeywordRetriever:
|
|
|
78
84
|
self,
|
|
79
85
|
filters=self.filters,
|
|
80
86
|
top_k=self.top_k,
|
|
87
|
+
filter_policy=self.filter_policy.value,
|
|
81
88
|
document_store=self.document_store.to_dict(),
|
|
82
89
|
)
|
|
83
90
|
|
|
@@ -93,6 +100,10 @@ class PgvectorKeywordRetriever:
|
|
|
93
100
|
"""
|
|
94
101
|
doc_store_params = data["init_parameters"]["document_store"]
|
|
95
102
|
data["init_parameters"]["document_store"] = PgvectorDocumentStore.from_dict(doc_store_params)
|
|
103
|
+
# Pipelines serialized with old versions of the component might not
|
|
104
|
+
# have the filter_policy field.
|
|
105
|
+
if filter_policy := data["init_parameters"].get("filter_policy"):
|
|
106
|
+
data["init_parameters"]["filter_policy"] = FilterPolicy.from_str(filter_policy)
|
|
96
107
|
return default_from_dict(cls, data)
|
|
97
108
|
|
|
98
109
|
@component.output_types(documents=List[Document])
|
|
@@ -106,13 +117,16 @@ class PgvectorKeywordRetriever:
|
|
|
106
117
|
Retrieve documents from the `PgvectorDocumentStore`, based on keywords.
|
|
107
118
|
|
|
108
119
|
:param query: String to search in `Document`s' content.
|
|
109
|
-
:param filters: Filters applied to the retrieved Documents.
|
|
120
|
+
:param filters: Filters applied to the retrieved Documents. The way runtime filters are applied depends on
|
|
121
|
+
the `filter_policy` chosen at retriever initialization. See init method docstring for more
|
|
122
|
+
details.
|
|
110
123
|
:param top_k: Maximum number of Documents to return.
|
|
111
124
|
|
|
112
125
|
:returns: A dictionary with the following keys:
|
|
113
126
|
- `documents`: List of `Document`s that match the query.
|
|
114
127
|
"""
|
|
115
|
-
filters =
|
|
128
|
+
filters = apply_filter_policy(self.filter_policy, self.filters, filters)
|
|
129
|
+
|
|
116
130
|
top_k = top_k or self.top_k
|
|
117
131
|
|
|
118
132
|
docs = self.document_store._keyword_retrieval(
|
|
@@ -5,6 +5,7 @@ from unittest.mock import Mock
|
|
|
5
5
|
|
|
6
6
|
import pytest
|
|
7
7
|
from haystack.dataclasses import Document
|
|
8
|
+
from haystack.document_stores.types import FilterPolicy
|
|
8
9
|
from haystack.utils.auth import EnvVarSecret
|
|
9
10
|
from haystack_integrations.components.retrievers.pgvector import PgvectorEmbeddingRetriever, PgvectorKeywordRetriever
|
|
10
11
|
from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
|
|
@@ -16,8 +17,15 @@ class TestEmbeddingRetriever:
|
|
|
16
17
|
assert retriever.document_store == mock_store
|
|
17
18
|
assert retriever.filters == {}
|
|
18
19
|
assert retriever.top_k == 10
|
|
20
|
+
assert retriever.filter_policy == FilterPolicy.REPLACE
|
|
19
21
|
assert retriever.vector_function == mock_store.vector_function
|
|
20
22
|
|
|
23
|
+
retriever = PgvectorEmbeddingRetriever(document_store=mock_store, filter_policy="merge")
|
|
24
|
+
assert retriever.filter_policy == FilterPolicy.MERGE
|
|
25
|
+
|
|
26
|
+
with pytest.raises(ValueError):
|
|
27
|
+
PgvectorEmbeddingRetriever(document_store=mock_store, filter_policy="invalid")
|
|
28
|
+
|
|
21
29
|
def test_init(self, mock_store):
|
|
22
30
|
retriever = PgvectorEmbeddingRetriever(
|
|
23
31
|
document_store=mock_store, filters={"field": "value"}, top_k=5, vector_function="l2_distance"
|
|
@@ -25,6 +33,7 @@ class TestEmbeddingRetriever:
|
|
|
25
33
|
assert retriever.document_store == mock_store
|
|
26
34
|
assert retriever.filters == {"field": "value"}
|
|
27
35
|
assert retriever.top_k == 5
|
|
36
|
+
assert retriever.filter_policy == FilterPolicy.REPLACE
|
|
28
37
|
assert retriever.vector_function == "l2_distance"
|
|
29
38
|
|
|
30
39
|
def test_to_dict(self, mock_store):
|
|
@@ -56,6 +65,7 @@ class TestEmbeddingRetriever:
|
|
|
56
65
|
"filters": {"field": "value"},
|
|
57
66
|
"top_k": 5,
|
|
58
67
|
"vector_function": "l2_distance",
|
|
68
|
+
"filter_policy": "replace",
|
|
59
69
|
},
|
|
60
70
|
}
|
|
61
71
|
|
|
@@ -85,6 +95,7 @@ class TestEmbeddingRetriever:
|
|
|
85
95
|
"filters": {"field": "value"},
|
|
86
96
|
"top_k": 5,
|
|
87
97
|
"vector_function": "l2_distance",
|
|
98
|
+
"filter_policy": "replace",
|
|
88
99
|
},
|
|
89
100
|
}
|
|
90
101
|
|
|
@@ -106,6 +117,7 @@ class TestEmbeddingRetriever:
|
|
|
106
117
|
|
|
107
118
|
assert retriever.filters == {"field": "value"}
|
|
108
119
|
assert retriever.top_k == 5
|
|
120
|
+
assert retriever.filter_policy == FilterPolicy.REPLACE
|
|
109
121
|
assert retriever.vector_function == "l2_distance"
|
|
110
122
|
|
|
111
123
|
def test_run(self):
|
|
@@ -130,12 +142,27 @@ class TestKeywordRetriever:
|
|
|
130
142
|
assert retriever.filters == {}
|
|
131
143
|
assert retriever.top_k == 10
|
|
132
144
|
|
|
145
|
+
retriever = PgvectorKeywordRetriever(document_store=mock_store, filter_policy="merge")
|
|
146
|
+
assert retriever.filter_policy == FilterPolicy.MERGE
|
|
147
|
+
|
|
148
|
+
with pytest.raises(ValueError):
|
|
149
|
+
PgvectorKeywordRetriever(document_store=mock_store, filter_policy="invalid")
|
|
150
|
+
|
|
133
151
|
def test_init(self, mock_store):
|
|
134
152
|
retriever = PgvectorKeywordRetriever(document_store=mock_store, filters={"field": "value"}, top_k=5)
|
|
135
153
|
assert retriever.document_store == mock_store
|
|
136
154
|
assert retriever.filters == {"field": "value"}
|
|
137
155
|
assert retriever.top_k == 5
|
|
138
156
|
|
|
157
|
+
def test_init_with_filter_policy(self, mock_store):
|
|
158
|
+
retriever = PgvectorKeywordRetriever(
|
|
159
|
+
document_store=mock_store, filters={"field": "value"}, top_k=5, filter_policy=FilterPolicy.MERGE
|
|
160
|
+
)
|
|
161
|
+
assert retriever.document_store == mock_store
|
|
162
|
+
assert retriever.filters == {"field": "value"}
|
|
163
|
+
assert retriever.top_k == 5
|
|
164
|
+
assert retriever.filter_policy == FilterPolicy.MERGE
|
|
165
|
+
|
|
139
166
|
def test_to_dict(self, mock_store):
|
|
140
167
|
retriever = PgvectorKeywordRetriever(document_store=mock_store, filters={"field": "value"}, top_k=5)
|
|
141
168
|
res = retriever.to_dict()
|
|
@@ -162,6 +189,7 @@ class TestKeywordRetriever:
|
|
|
162
189
|
},
|
|
163
190
|
"filters": {"field": "value"},
|
|
164
191
|
"top_k": 5,
|
|
192
|
+
"filter_policy": "replace",
|
|
165
193
|
},
|
|
166
194
|
}
|
|
167
195
|
|
|
@@ -190,6 +218,7 @@ class TestKeywordRetriever:
|
|
|
190
218
|
},
|
|
191
219
|
"filters": {"field": "value"},
|
|
192
220
|
"top_k": 5,
|
|
221
|
+
"filter_policy": "replace",
|
|
193
222
|
},
|
|
194
223
|
}
|
|
195
224
|
|
|
@@ -212,6 +241,54 @@ class TestKeywordRetriever:
|
|
|
212
241
|
assert retriever.filters == {"field": "value"}
|
|
213
242
|
assert retriever.top_k == 5
|
|
214
243
|
|
|
244
|
+
@pytest.mark.usefixtures("patches_for_unit_tests")
|
|
245
|
+
def test_from_dict_without_filter_policy(self, monkeypatch):
|
|
246
|
+
monkeypatch.setenv("PG_CONN_STR", "some-connection-string")
|
|
247
|
+
t = "haystack_integrations.components.retrievers.pgvector.keyword_retriever.PgvectorKeywordRetriever"
|
|
248
|
+
data = {
|
|
249
|
+
"type": t,
|
|
250
|
+
"init_parameters": {
|
|
251
|
+
"document_store": {
|
|
252
|
+
"type": "haystack_integrations.document_stores.pgvector.document_store.PgvectorDocumentStore",
|
|
253
|
+
"init_parameters": {
|
|
254
|
+
"connection_string": {"env_vars": ["PG_CONN_STR"], "strict": True, "type": "env_var"},
|
|
255
|
+
"table_name": "haystack_test_to_dict",
|
|
256
|
+
"embedding_dimension": 768,
|
|
257
|
+
"vector_function": "cosine_similarity",
|
|
258
|
+
"recreate_table": True,
|
|
259
|
+
"search_strategy": "exact_nearest_neighbor",
|
|
260
|
+
"hnsw_recreate_index_if_exists": False,
|
|
261
|
+
"hnsw_index_creation_kwargs": {},
|
|
262
|
+
"hnsw_index_name": "haystack_hnsw_index",
|
|
263
|
+
"hnsw_ef_search": None,
|
|
264
|
+
"keyword_index_name": "haystack_keyword_index",
|
|
265
|
+
},
|
|
266
|
+
},
|
|
267
|
+
"filters": {"field": "value"},
|
|
268
|
+
"top_k": 5,
|
|
269
|
+
},
|
|
270
|
+
}
|
|
271
|
+
|
|
272
|
+
retriever = PgvectorKeywordRetriever.from_dict(data)
|
|
273
|
+
document_store = retriever.document_store
|
|
274
|
+
|
|
275
|
+
assert isinstance(document_store, PgvectorDocumentStore)
|
|
276
|
+
assert isinstance(document_store.connection_string, EnvVarSecret)
|
|
277
|
+
assert document_store.table_name == "haystack_test_to_dict"
|
|
278
|
+
assert document_store.embedding_dimension == 768
|
|
279
|
+
assert document_store.vector_function == "cosine_similarity"
|
|
280
|
+
assert document_store.recreate_table
|
|
281
|
+
assert document_store.search_strategy == "exact_nearest_neighbor"
|
|
282
|
+
assert not document_store.hnsw_recreate_index_if_exists
|
|
283
|
+
assert document_store.hnsw_index_creation_kwargs == {}
|
|
284
|
+
assert document_store.hnsw_index_name == "haystack_hnsw_index"
|
|
285
|
+
assert document_store.hnsw_ef_search is None
|
|
286
|
+
assert document_store.keyword_index_name == "haystack_keyword_index"
|
|
287
|
+
|
|
288
|
+
assert retriever.filters == {"field": "value"}
|
|
289
|
+
assert retriever.filter_policy == FilterPolicy.REPLACE # defaults to REPLACE
|
|
290
|
+
assert retriever.top_k == 5
|
|
291
|
+
|
|
215
292
|
def test_run(self):
|
|
216
293
|
mock_store = Mock(spec=PgvectorDocumentStore)
|
|
217
294
|
doc = Document(content="Test doc", embedding=[0.1, 0.2])
|
|
@@ -223,3 +300,19 @@ class TestKeywordRetriever:
|
|
|
223
300
|
mock_store._keyword_retrieval.assert_called_once_with(query="test query", filters={}, top_k=10)
|
|
224
301
|
|
|
225
302
|
assert res == {"documents": [doc]}
|
|
303
|
+
|
|
304
|
+
def test_run_with_filters(self):
|
|
305
|
+
mock_store = Mock(spec=PgvectorDocumentStore)
|
|
306
|
+
doc = Document(content="Test doc", embedding=[0.1, 0.2])
|
|
307
|
+
mock_store._keyword_retrieval.return_value = [doc]
|
|
308
|
+
|
|
309
|
+
retriever = PgvectorKeywordRetriever(
|
|
310
|
+
document_store=mock_store, filter_policy=FilterPolicy.MERGE, filters={"field": "value"}
|
|
311
|
+
)
|
|
312
|
+
res = retriever.run(query="test query", filters={"field2": "value2"})
|
|
313
|
+
|
|
314
|
+
mock_store._keyword_retrieval.assert_called_once_with(
|
|
315
|
+
query="test query", filters={"field": "value", "field2": "value2"}, top_k=10
|
|
316
|
+
)
|
|
317
|
+
|
|
318
|
+
assert res == {"documents": [doc]}
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|