pgvector-haystack 0.3.0__tar.gz → 0.5.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.
- {pgvector_haystack-0.3.0 → pgvector_haystack-0.5.0}/.gitignore +9 -0
- pgvector_haystack-0.5.0/CHANGELOG.md +36 -0
- {pgvector_haystack-0.3.0 → pgvector_haystack-0.5.0}/PKG-INFO +1 -1
- {pgvector_haystack-0.3.0 → pgvector_haystack-0.5.0}/pyproject.toml +24 -51
- {pgvector_haystack-0.3.0 → pgvector_haystack-0.5.0}/src/haystack_integrations/components/retrievers/pgvector/embedding_retriever.py +14 -4
- {pgvector_haystack-0.3.0 → pgvector_haystack-0.5.0}/src/haystack_integrations/components/retrievers/pgvector/keyword_retriever.py +15 -4
- {pgvector_haystack-0.3.0 → pgvector_haystack-0.5.0}/src/haystack_integrations/document_stores/pgvector/document_store.py +13 -9
- {pgvector_haystack-0.3.0 → pgvector_haystack-0.5.0}/tests/test_document_store.py +8 -0
- {pgvector_haystack-0.3.0 → pgvector_haystack-0.5.0}/tests/test_retrievers.py +57 -0
- {pgvector_haystack-0.3.0 → pgvector_haystack-0.5.0}/LICENSE.txt +0 -0
- {pgvector_haystack-0.3.0 → pgvector_haystack-0.5.0}/README.md +0 -0
- {pgvector_haystack-0.3.0 → pgvector_haystack-0.5.0}/examples/embedding_retrieval.py +0 -0
- {pgvector_haystack-0.3.0 → pgvector_haystack-0.5.0}/examples/hybrid_retrieval.py +0 -0
- {pgvector_haystack-0.3.0 → pgvector_haystack-0.5.0}/pydoc/config.yml +0 -0
- {pgvector_haystack-0.3.0 → pgvector_haystack-0.5.0}/src/haystack_integrations/components/retrievers/pgvector/__init__.py +0 -0
- {pgvector_haystack-0.3.0 → pgvector_haystack-0.5.0}/src/haystack_integrations/document_stores/pgvector/__init__.py +0 -0
- {pgvector_haystack-0.3.0 → pgvector_haystack-0.5.0}/src/haystack_integrations/document_stores/pgvector/filters.py +0 -0
- {pgvector_haystack-0.3.0 → pgvector_haystack-0.5.0}/tests/__init__.py +0 -0
- {pgvector_haystack-0.3.0 → pgvector_haystack-0.5.0}/tests/conftest.py +0 -0
- {pgvector_haystack-0.3.0 → pgvector_haystack-0.5.0}/tests/test_embedding_retrieval.py +0 -0
- {pgvector_haystack-0.3.0 → pgvector_haystack-0.5.0}/tests/test_filters.py +0 -0
- {pgvector_haystack-0.3.0 → pgvector_haystack-0.5.0}/tests/test_keyword_retrieval.py +0 -0
|
@@ -0,0 +1,36 @@
|
|
|
1
|
+
# Changelog
|
|
2
|
+
|
|
3
|
+
## [integrations/pgvector-v0.4.0] - 2024-06-20
|
|
4
|
+
|
|
5
|
+
### 🚀 Features
|
|
6
|
+
|
|
7
|
+
- Defer the database connection to when it's needed (#773)
|
|
8
|
+
- Add customizable index names for pgvector (#818)
|
|
9
|
+
|
|
10
|
+
## [integrations/pgvector-v0.2.0] - 2024-05-08
|
|
11
|
+
|
|
12
|
+
### 🚀 Features
|
|
13
|
+
|
|
14
|
+
- `MongoDBAtlasEmbeddingRetriever` (#427)
|
|
15
|
+
- Implement keyword retrieval for pgvector integration (#644)
|
|
16
|
+
|
|
17
|
+
### 🐛 Bug Fixes
|
|
18
|
+
|
|
19
|
+
- Fix order of API docs (#447)
|
|
20
|
+
|
|
21
|
+
This PR will also push the docs to Readme
|
|
22
|
+
|
|
23
|
+
### 📚 Documentation
|
|
24
|
+
|
|
25
|
+
- Update category slug (#442)
|
|
26
|
+
- Disable-class-def (#556)
|
|
27
|
+
|
|
28
|
+
## [integrations/pgvector-v0.1.0] - 2024-02-14
|
|
29
|
+
|
|
30
|
+
### 🐛 Bug Fixes
|
|
31
|
+
|
|
32
|
+
- Fix linting (#328)
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
<!-- generated by git-cliff -->
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.3
|
|
2
2
|
Name: pgvector-haystack
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.5.0
|
|
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,7 @@ 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
|
+
data["init_parameters"]["filter_policy"] = FilterPolicy.from_str(data["init_parameters"]["filter_policy"])
|
|
122
130
|
return default_from_dict(cls, data)
|
|
123
131
|
|
|
124
132
|
@component.output_types(documents=List[Document])
|
|
@@ -133,13 +141,15 @@ class PgvectorEmbeddingRetriever:
|
|
|
133
141
|
Retrieve documents from the `PgvectorDocumentStore`, based on their embeddings.
|
|
134
142
|
|
|
135
143
|
:param query_embedding: Embedding of the query.
|
|
136
|
-
:param filters: Filters applied to the retrieved Documents.
|
|
144
|
+
:param filters: Filters applied to the retrieved Documents. The way runtime filters are applied depends on
|
|
145
|
+
the `filter_policy` chosen at retriever initialization. See init method docstring for more
|
|
146
|
+
details.
|
|
137
147
|
:param top_k: Maximum number of Documents to return.
|
|
138
148
|
:param vector_function: The similarity function to use when searching for similar embeddings.
|
|
139
149
|
|
|
140
150
|
:returns: List of Documents similar to `query_embedding`.
|
|
141
151
|
"""
|
|
142
|
-
filters =
|
|
152
|
+
filters = apply_filter_policy(self.filter_policy, self.filters, filters)
|
|
143
153
|
top_k = top_k or self.top_k
|
|
144
154
|
vector_function = vector_function or self.vector_function
|
|
145
155
|
|
|
@@ -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,7 @@ 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
|
+
data["init_parameters"]["filter_policy"] = FilterPolicy.from_str(data["init_parameters"]["filter_policy"])
|
|
96
104
|
return default_from_dict(cls, data)
|
|
97
105
|
|
|
98
106
|
@component.output_types(documents=List[Document])
|
|
@@ -106,13 +114,16 @@ class PgvectorKeywordRetriever:
|
|
|
106
114
|
Retrieve documents from the `PgvectorDocumentStore`, based on keywords.
|
|
107
115
|
|
|
108
116
|
:param query: String to search in `Document`s' content.
|
|
109
|
-
:param filters: Filters applied to the retrieved Documents.
|
|
117
|
+
:param filters: Filters applied to the retrieved Documents. The way runtime filters are applied depends on
|
|
118
|
+
the `filter_policy` chosen at retriever initialization. See init method docstring for more
|
|
119
|
+
details.
|
|
110
120
|
:param top_k: Maximum number of Documents to return.
|
|
111
121
|
|
|
112
122
|
:returns: A dictionary with the following keys:
|
|
113
123
|
- `documents`: List of `Document`s that match the query.
|
|
114
124
|
"""
|
|
115
|
-
filters =
|
|
125
|
+
filters = apply_filter_policy(self.filter_policy, self.filters, filters)
|
|
126
|
+
|
|
116
127
|
top_k = top_k or self.top_k
|
|
117
128
|
|
|
118
129
|
docs = self.document_store._keyword_retrieval(
|
|
@@ -69,10 +69,6 @@ VECTOR_FUNCTION_TO_POSTGRESQL_OPS = {
|
|
|
69
69
|
|
|
70
70
|
HNSW_INDEX_CREATION_VALID_KWARGS = ["m", "ef_construction"]
|
|
71
71
|
|
|
72
|
-
HNSW_INDEX_NAME = "haystack_hnsw_index"
|
|
73
|
-
|
|
74
|
-
KEYWORD_INDEX_NAME = "haystack_keyword_index"
|
|
75
|
-
|
|
76
72
|
|
|
77
73
|
class PgvectorDocumentStore:
|
|
78
74
|
"""
|
|
@@ -91,7 +87,9 @@ class PgvectorDocumentStore:
|
|
|
91
87
|
search_strategy: Literal["exact_nearest_neighbor", "hnsw"] = "exact_nearest_neighbor",
|
|
92
88
|
hnsw_recreate_index_if_exists: bool = False,
|
|
93
89
|
hnsw_index_creation_kwargs: Optional[Dict[str, int]] = None,
|
|
90
|
+
hnsw_index_name: str = "haystack_hnsw_index",
|
|
94
91
|
hnsw_ef_search: Optional[int] = None,
|
|
92
|
+
keyword_index_name: str = "haystack_keyword_index",
|
|
95
93
|
):
|
|
96
94
|
"""
|
|
97
95
|
Creates a new PgvectorDocumentStore instance.
|
|
@@ -127,9 +125,11 @@ class PgvectorDocumentStore:
|
|
|
127
125
|
:param hnsw_index_creation_kwargs: Additional keyword arguments to pass to the HNSW index creation.
|
|
128
126
|
Only used if search_strategy is set to `"hnsw"`. You can find the list of valid arguments in the
|
|
129
127
|
[pgvector documentation](https://github.com/pgvector/pgvector?tab=readme-ov-file#hnsw)
|
|
128
|
+
:param hnsw_index_name: Index name for the HNSW index.
|
|
130
129
|
:param hnsw_ef_search: The `ef_search` parameter to use at query time. Only used if search_strategy is set to
|
|
131
130
|
`"hnsw"`. You can find more information about this parameter in the
|
|
132
131
|
[pgvector documentation](https://github.com/pgvector/pgvector?tab=readme-ov-file#hnsw).
|
|
132
|
+
:param keyword_index_name: Index name for the Keyword index.
|
|
133
133
|
"""
|
|
134
134
|
|
|
135
135
|
self.connection_string = connection_string
|
|
@@ -143,7 +143,9 @@ class PgvectorDocumentStore:
|
|
|
143
143
|
self.search_strategy = search_strategy
|
|
144
144
|
self.hnsw_recreate_index_if_exists = hnsw_recreate_index_if_exists
|
|
145
145
|
self.hnsw_index_creation_kwargs = hnsw_index_creation_kwargs or {}
|
|
146
|
+
self.hnsw_index_name = hnsw_index_name
|
|
146
147
|
self.hnsw_ef_search = hnsw_ef_search
|
|
148
|
+
self.keyword_index_name = keyword_index_name
|
|
147
149
|
self.language = language
|
|
148
150
|
self._connection = None
|
|
149
151
|
self._cursor = None
|
|
@@ -209,7 +211,9 @@ class PgvectorDocumentStore:
|
|
|
209
211
|
search_strategy=self.search_strategy,
|
|
210
212
|
hnsw_recreate_index_if_exists=self.hnsw_recreate_index_if_exists,
|
|
211
213
|
hnsw_index_creation_kwargs=self.hnsw_index_creation_kwargs,
|
|
214
|
+
hnsw_index_name=self.hnsw_index_name,
|
|
212
215
|
hnsw_ef_search=self.hnsw_ef_search,
|
|
216
|
+
keyword_index_name=self.keyword_index_name,
|
|
213
217
|
language=self.language,
|
|
214
218
|
)
|
|
215
219
|
|
|
@@ -281,7 +285,7 @@ class PgvectorDocumentStore:
|
|
|
281
285
|
index_exists = bool(
|
|
282
286
|
self._execute_sql(
|
|
283
287
|
"SELECT 1 FROM pg_indexes WHERE tablename = %s AND indexname = %s",
|
|
284
|
-
(self.table_name,
|
|
288
|
+
(self.table_name, self.keyword_index_name),
|
|
285
289
|
"Could not check if keyword index exists",
|
|
286
290
|
).fetchone()
|
|
287
291
|
)
|
|
@@ -289,7 +293,7 @@ class PgvectorDocumentStore:
|
|
|
289
293
|
sql_create_index = SQL(
|
|
290
294
|
"CREATE INDEX {index_name} ON {table_name} USING GIN (to_tsvector({language}, content))"
|
|
291
295
|
).format(
|
|
292
|
-
index_name=Identifier(
|
|
296
|
+
index_name=Identifier(self.keyword_index_name),
|
|
293
297
|
table_name=Identifier(self.table_name),
|
|
294
298
|
language=SQLLiteral(self.language),
|
|
295
299
|
)
|
|
@@ -312,7 +316,7 @@ class PgvectorDocumentStore:
|
|
|
312
316
|
index_exists = bool(
|
|
313
317
|
self._execute_sql(
|
|
314
318
|
"SELECT 1 FROM pg_indexes WHERE tablename = %s AND indexname = %s",
|
|
315
|
-
(self.table_name,
|
|
319
|
+
(self.table_name, self.hnsw_index_name),
|
|
316
320
|
"Could not check if HNSW index exists",
|
|
317
321
|
).fetchone()
|
|
318
322
|
)
|
|
@@ -325,7 +329,7 @@ class PgvectorDocumentStore:
|
|
|
325
329
|
)
|
|
326
330
|
return
|
|
327
331
|
|
|
328
|
-
sql_drop_index = SQL("DROP INDEX IF EXISTS {index_name}").format(index_name=Identifier(
|
|
332
|
+
sql_drop_index = SQL("DROP INDEX IF EXISTS {index_name}").format(index_name=Identifier(self.hnsw_index_name))
|
|
329
333
|
self._execute_sql(sql_drop_index, error_msg="Could not drop HNSW index")
|
|
330
334
|
|
|
331
335
|
self._create_hnsw_index()
|
|
@@ -343,7 +347,7 @@ class PgvectorDocumentStore:
|
|
|
343
347
|
}
|
|
344
348
|
|
|
345
349
|
sql_create_index = SQL("CREATE INDEX {index_name} ON {table_name} USING hnsw (embedding {ops}) ").format(
|
|
346
|
-
index_name=Identifier(
|
|
350
|
+
index_name=Identifier(self.hnsw_index_name), table_name=Identifier(self.table_name), ops=SQL(pg_ops)
|
|
347
351
|
)
|
|
348
352
|
|
|
349
353
|
if actual_hnsw_index_creation_kwargs:
|
|
@@ -53,7 +53,9 @@ def test_init(monkeypatch):
|
|
|
53
53
|
search_strategy="hnsw",
|
|
54
54
|
hnsw_recreate_index_if_exists=True,
|
|
55
55
|
hnsw_index_creation_kwargs={"m": 32, "ef_construction": 128},
|
|
56
|
+
hnsw_index_name="my_hnsw_index",
|
|
56
57
|
hnsw_ef_search=50,
|
|
58
|
+
keyword_index_name="my_keyword_index",
|
|
57
59
|
)
|
|
58
60
|
|
|
59
61
|
assert document_store.table_name == "my_table"
|
|
@@ -63,7 +65,9 @@ def test_init(monkeypatch):
|
|
|
63
65
|
assert document_store.search_strategy == "hnsw"
|
|
64
66
|
assert document_store.hnsw_recreate_index_if_exists
|
|
65
67
|
assert document_store.hnsw_index_creation_kwargs == {"m": 32, "ef_construction": 128}
|
|
68
|
+
assert document_store.hnsw_index_name == "my_hnsw_index"
|
|
66
69
|
assert document_store.hnsw_ef_search == 50
|
|
70
|
+
assert document_store.keyword_index_name == "my_keyword_index"
|
|
67
71
|
|
|
68
72
|
|
|
69
73
|
@pytest.mark.usefixtures("patches_for_unit_tests")
|
|
@@ -78,7 +82,9 @@ def test_to_dict(monkeypatch):
|
|
|
78
82
|
search_strategy="hnsw",
|
|
79
83
|
hnsw_recreate_index_if_exists=True,
|
|
80
84
|
hnsw_index_creation_kwargs={"m": 32, "ef_construction": 128},
|
|
85
|
+
hnsw_index_name="my_hnsw_index",
|
|
81
86
|
hnsw_ef_search=50,
|
|
87
|
+
keyword_index_name="my_keyword_index",
|
|
82
88
|
)
|
|
83
89
|
|
|
84
90
|
assert document_store.to_dict() == {
|
|
@@ -93,7 +99,9 @@ def test_to_dict(monkeypatch):
|
|
|
93
99
|
"hnsw_recreate_index_if_exists": True,
|
|
94
100
|
"language": "english",
|
|
95
101
|
"hnsw_index_creation_kwargs": {"m": 32, "ef_construction": 128},
|
|
102
|
+
"hnsw_index_name": "my_hnsw_index",
|
|
96
103
|
"hnsw_ef_search": 50,
|
|
104
|
+
"keyword_index_name": "my_keyword_index",
|
|
97
105
|
},
|
|
98
106
|
}
|
|
99
107
|
|
|
@@ -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):
|
|
@@ -48,12 +57,15 @@ class TestEmbeddingRetriever:
|
|
|
48
57
|
"hnsw_recreate_index_if_exists": False,
|
|
49
58
|
"language": "english",
|
|
50
59
|
"hnsw_index_creation_kwargs": {},
|
|
60
|
+
"hnsw_index_name": "haystack_hnsw_index",
|
|
51
61
|
"hnsw_ef_search": None,
|
|
62
|
+
"keyword_index_name": "haystack_keyword_index",
|
|
52
63
|
},
|
|
53
64
|
},
|
|
54
65
|
"filters": {"field": "value"},
|
|
55
66
|
"top_k": 5,
|
|
56
67
|
"vector_function": "l2_distance",
|
|
68
|
+
"filter_policy": "replace",
|
|
57
69
|
},
|
|
58
70
|
}
|
|
59
71
|
|
|
@@ -75,12 +87,15 @@ class TestEmbeddingRetriever:
|
|
|
75
87
|
"search_strategy": "exact_nearest_neighbor",
|
|
76
88
|
"hnsw_recreate_index_if_exists": False,
|
|
77
89
|
"hnsw_index_creation_kwargs": {},
|
|
90
|
+
"hnsw_index_name": "haystack_hnsw_index",
|
|
78
91
|
"hnsw_ef_search": None,
|
|
92
|
+
"keyword_index_name": "haystack_keyword_index",
|
|
79
93
|
},
|
|
80
94
|
},
|
|
81
95
|
"filters": {"field": "value"},
|
|
82
96
|
"top_k": 5,
|
|
83
97
|
"vector_function": "l2_distance",
|
|
98
|
+
"filter_policy": "replace",
|
|
84
99
|
},
|
|
85
100
|
}
|
|
86
101
|
|
|
@@ -96,10 +111,13 @@ class TestEmbeddingRetriever:
|
|
|
96
111
|
assert document_store.search_strategy == "exact_nearest_neighbor"
|
|
97
112
|
assert not document_store.hnsw_recreate_index_if_exists
|
|
98
113
|
assert document_store.hnsw_index_creation_kwargs == {}
|
|
114
|
+
assert document_store.hnsw_index_name == "haystack_hnsw_index"
|
|
99
115
|
assert document_store.hnsw_ef_search is None
|
|
116
|
+
assert document_store.keyword_index_name == "haystack_keyword_index"
|
|
100
117
|
|
|
101
118
|
assert retriever.filters == {"field": "value"}
|
|
102
119
|
assert retriever.top_k == 5
|
|
120
|
+
assert retriever.filter_policy == FilterPolicy.REPLACE
|
|
103
121
|
assert retriever.vector_function == "l2_distance"
|
|
104
122
|
|
|
105
123
|
def test_run(self):
|
|
@@ -124,12 +142,27 @@ class TestKeywordRetriever:
|
|
|
124
142
|
assert retriever.filters == {}
|
|
125
143
|
assert retriever.top_k == 10
|
|
126
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
|
+
|
|
127
151
|
def test_init(self, mock_store):
|
|
128
152
|
retriever = PgvectorKeywordRetriever(document_store=mock_store, filters={"field": "value"}, top_k=5)
|
|
129
153
|
assert retriever.document_store == mock_store
|
|
130
154
|
assert retriever.filters == {"field": "value"}
|
|
131
155
|
assert retriever.top_k == 5
|
|
132
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
|
+
|
|
133
166
|
def test_to_dict(self, mock_store):
|
|
134
167
|
retriever = PgvectorKeywordRetriever(document_store=mock_store, filters={"field": "value"}, top_k=5)
|
|
135
168
|
res = retriever.to_dict()
|
|
@@ -149,11 +182,14 @@ class TestKeywordRetriever:
|
|
|
149
182
|
"hnsw_recreate_index_if_exists": False,
|
|
150
183
|
"language": "english",
|
|
151
184
|
"hnsw_index_creation_kwargs": {},
|
|
185
|
+
"hnsw_index_name": "haystack_hnsw_index",
|
|
152
186
|
"hnsw_ef_search": None,
|
|
187
|
+
"keyword_index_name": "haystack_keyword_index",
|
|
153
188
|
},
|
|
154
189
|
},
|
|
155
190
|
"filters": {"field": "value"},
|
|
156
191
|
"top_k": 5,
|
|
192
|
+
"filter_policy": "replace",
|
|
157
193
|
},
|
|
158
194
|
}
|
|
159
195
|
|
|
@@ -175,11 +211,14 @@ class TestKeywordRetriever:
|
|
|
175
211
|
"search_strategy": "exact_nearest_neighbor",
|
|
176
212
|
"hnsw_recreate_index_if_exists": False,
|
|
177
213
|
"hnsw_index_creation_kwargs": {},
|
|
214
|
+
"hnsw_index_name": "haystack_hnsw_index",
|
|
178
215
|
"hnsw_ef_search": None,
|
|
216
|
+
"keyword_index_name": "haystack_keyword_index",
|
|
179
217
|
},
|
|
180
218
|
},
|
|
181
219
|
"filters": {"field": "value"},
|
|
182
220
|
"top_k": 5,
|
|
221
|
+
"filter_policy": "replace",
|
|
183
222
|
},
|
|
184
223
|
}
|
|
185
224
|
|
|
@@ -195,7 +234,9 @@ class TestKeywordRetriever:
|
|
|
195
234
|
assert document_store.search_strategy == "exact_nearest_neighbor"
|
|
196
235
|
assert not document_store.hnsw_recreate_index_if_exists
|
|
197
236
|
assert document_store.hnsw_index_creation_kwargs == {}
|
|
237
|
+
assert document_store.hnsw_index_name == "haystack_hnsw_index"
|
|
198
238
|
assert document_store.hnsw_ef_search is None
|
|
239
|
+
assert document_store.keyword_index_name == "haystack_keyword_index"
|
|
199
240
|
|
|
200
241
|
assert retriever.filters == {"field": "value"}
|
|
201
242
|
assert retriever.top_k == 5
|
|
@@ -211,3 +252,19 @@ class TestKeywordRetriever:
|
|
|
211
252
|
mock_store._keyword_retrieval.assert_called_once_with(query="test query", filters={}, top_k=10)
|
|
212
253
|
|
|
213
254
|
assert res == {"documents": [doc]}
|
|
255
|
+
|
|
256
|
+
def test_run_with_filters(self):
|
|
257
|
+
mock_store = Mock(spec=PgvectorDocumentStore)
|
|
258
|
+
doc = Document(content="Test doc", embedding=[0.1, 0.2])
|
|
259
|
+
mock_store._keyword_retrieval.return_value = [doc]
|
|
260
|
+
|
|
261
|
+
retriever = PgvectorKeywordRetriever(
|
|
262
|
+
document_store=mock_store, filter_policy=FilterPolicy.MERGE, filters={"field": "value"}
|
|
263
|
+
)
|
|
264
|
+
res = retriever.run(query="test query", filters={"field2": "value2"})
|
|
265
|
+
|
|
266
|
+
mock_store._keyword_retrieval.assert_called_once_with(
|
|
267
|
+
query="test query", filters={"field": "value", "field2": "value2"}, top_k=10
|
|
268
|
+
)
|
|
269
|
+
|
|
270
|
+
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
|