pgvector-haystack 0.1.0__tar.gz → 0.3.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.1.0 → pgvector_haystack-0.3.0}/.gitignore +1 -0
- {pgvector_haystack-0.1.0 → pgvector_haystack-0.3.0}/PKG-INFO +3 -2
- pgvector_haystack-0.1.0/examples/example.py → pgvector_haystack-0.3.0/examples/embedding_retrieval.py +2 -2
- pgvector_haystack-0.3.0/examples/hybrid_retrieval.py +69 -0
- {pgvector_haystack-0.1.0 → pgvector_haystack-0.3.0}/pydoc/config.yml +5 -4
- {pgvector_haystack-0.1.0 → pgvector_haystack-0.3.0}/pyproject.toml +10 -4
- {pgvector_haystack-0.1.0 → pgvector_haystack-0.3.0}/src/haystack_integrations/components/retrievers/pgvector/__init__.py +2 -1
- {pgvector_haystack-0.1.0 → pgvector_haystack-0.3.0}/src/haystack_integrations/components/retrievers/pgvector/embedding_retriever.py +65 -16
- pgvector_haystack-0.3.0/src/haystack_integrations/components/retrievers/pgvector/keyword_retriever.py +123 -0
- {pgvector_haystack-0.1.0 → pgvector_haystack-0.3.0}/src/haystack_integrations/document_stores/pgvector/document_store.py +190 -49
- {pgvector_haystack-0.1.0 → pgvector_haystack-0.3.0}/src/haystack_integrations/document_stores/pgvector/filters.py +15 -13
- pgvector_haystack-0.3.0/tests/conftest.py +63 -0
- pgvector_haystack-0.3.0/tests/test_document_store.py +226 -0
- {pgvector_haystack-0.1.0 → pgvector_haystack-0.3.0}/tests/test_embedding_retrieval.py +1 -0
- pgvector_haystack-0.3.0/tests/test_filters.py +191 -0
- pgvector_haystack-0.3.0/tests/test_keyword_retrieval.py +50 -0
- pgvector_haystack-0.3.0/tests/test_retrievers.py +213 -0
- pgvector_haystack-0.1.0/tests/conftest.py +0 -26
- pgvector_haystack-0.1.0/tests/test_document_store.py +0 -219
- pgvector_haystack-0.1.0/tests/test_filters.py +0 -179
- pgvector_haystack-0.1.0/tests/test_retriever.py +0 -113
- {pgvector_haystack-0.1.0 → pgvector_haystack-0.3.0}/LICENSE.txt +0 -0
- {pgvector_haystack-0.1.0 → pgvector_haystack-0.3.0}/README.md +0 -0
- {pgvector_haystack-0.1.0 → pgvector_haystack-0.3.0}/src/haystack_integrations/document_stores/pgvector/__init__.py +0 -0
- {pgvector_haystack-0.1.0 → pgvector_haystack-0.3.0}/tests/__init__.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
|
-
Metadata-Version: 2.
|
|
1
|
+
Metadata-Version: 2.3
|
|
2
2
|
Name: pgvector-haystack
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.3.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
|
|
@@ -9,6 +9,7 @@ Author-email: deepset GmbH <info@deepset.ai>
|
|
|
9
9
|
License-Expression: Apache-2.0
|
|
10
10
|
License-File: LICENSE.txt
|
|
11
11
|
Classifier: Development Status :: 4 - Beta
|
|
12
|
+
Classifier: License :: OSI Approved :: Apache Software License
|
|
12
13
|
Classifier: Programming Language :: Python
|
|
13
14
|
Classifier: Programming Language :: Python :: 3.8
|
|
14
15
|
Classifier: Programming Language :: Python :: 3.9
|
|
@@ -11,7 +11,6 @@
|
|
|
11
11
|
# git clone https://github.com/anakin87/neural-search-pills
|
|
12
12
|
|
|
13
13
|
import glob
|
|
14
|
-
import os
|
|
15
14
|
|
|
16
15
|
from haystack import Pipeline
|
|
17
16
|
from haystack.components.converters import MarkdownToDocument
|
|
@@ -21,7 +20,8 @@ from haystack.components.writers import DocumentWriter
|
|
|
21
20
|
from haystack_integrations.components.retrievers.pgvector import PgvectorEmbeddingRetriever
|
|
22
21
|
from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
|
|
23
22
|
|
|
24
|
-
|
|
23
|
+
# Set an environment variable `PG_CONN_STR` with the connection string to your PostgreSQL database.
|
|
24
|
+
# e.g., "postgresql://USER:PASSWORD@HOST:PORT/DB_NAME"
|
|
25
25
|
|
|
26
26
|
# Initialize PgvectorDocumentStore
|
|
27
27
|
document_store = PgvectorDocumentStore(
|
|
@@ -0,0 +1,69 @@
|
|
|
1
|
+
# Before running this example, ensure you have PostgreSQL installed with the pgvector extension.
|
|
2
|
+
# For a quick setup using Docker:
|
|
3
|
+
# docker run -d -p 5432:5432 -e POSTGRES_USER=postgres -e POSTGRES_PASSWORD=postgres
|
|
4
|
+
# -e POSTGRES_DB=postgres ankane/pgvector
|
|
5
|
+
|
|
6
|
+
# Install required packages for this example, including pgvector-haystack and other libraries needed
|
|
7
|
+
# for Markdown conversion and embeddings generation. Use the following command:
|
|
8
|
+
# pip install pgvector-haystack markdown-it-py mdit_plain "sentence-transformers>=2.2.0"
|
|
9
|
+
|
|
10
|
+
# Download some Markdown files to index.
|
|
11
|
+
# git clone https://github.com/anakin87/neural-search-pills
|
|
12
|
+
|
|
13
|
+
import glob
|
|
14
|
+
|
|
15
|
+
from haystack import Pipeline
|
|
16
|
+
from haystack.components.converters import MarkdownToDocument
|
|
17
|
+
from haystack.components.embedders import SentenceTransformersDocumentEmbedder, SentenceTransformersTextEmbedder
|
|
18
|
+
from haystack.components.joiners import DocumentJoiner
|
|
19
|
+
from haystack.components.preprocessors import DocumentSplitter
|
|
20
|
+
from haystack.components.writers import DocumentWriter
|
|
21
|
+
from haystack_integrations.components.retrievers.pgvector import PgvectorEmbeddingRetriever, PgvectorKeywordRetriever
|
|
22
|
+
from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
|
|
23
|
+
|
|
24
|
+
# Set an environment variable `PG_CONN_STR` with the connection string to your PostgreSQL database.
|
|
25
|
+
# e.g., "postgresql://USER:PASSWORD@HOST:PORT/DB_NAME"
|
|
26
|
+
|
|
27
|
+
# Initialize PgvectorDocumentStore
|
|
28
|
+
document_store = PgvectorDocumentStore(
|
|
29
|
+
table_name="haystack_test",
|
|
30
|
+
embedding_dimension=768,
|
|
31
|
+
vector_function="cosine_similarity",
|
|
32
|
+
recreate_table=True,
|
|
33
|
+
search_strategy="hnsw",
|
|
34
|
+
)
|
|
35
|
+
|
|
36
|
+
# Create the indexing Pipeline and index some documents
|
|
37
|
+
file_paths = glob.glob("neural-search-pills/pills/*.md")
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
indexing = Pipeline()
|
|
41
|
+
indexing.add_component("converter", MarkdownToDocument())
|
|
42
|
+
indexing.add_component("splitter", DocumentSplitter(split_by="sentence", split_length=2))
|
|
43
|
+
indexing.add_component("document_embedder", SentenceTransformersDocumentEmbedder())
|
|
44
|
+
indexing.add_component("writer", DocumentWriter(document_store))
|
|
45
|
+
indexing.connect("converter", "splitter")
|
|
46
|
+
indexing.connect("splitter", "document_embedder")
|
|
47
|
+
indexing.connect("document_embedder", "writer")
|
|
48
|
+
|
|
49
|
+
indexing.run({"converter": {"sources": file_paths}})
|
|
50
|
+
|
|
51
|
+
# Create the querying Pipeline and try a query
|
|
52
|
+
querying = Pipeline()
|
|
53
|
+
querying.add_component("text_embedder", SentenceTransformersTextEmbedder())
|
|
54
|
+
querying.add_component("retriever", PgvectorEmbeddingRetriever(document_store=document_store, top_k=3))
|
|
55
|
+
querying.add_component("keyword_retriever", PgvectorKeywordRetriever(document_store=document_store, top_k=3))
|
|
56
|
+
querying.add_component(
|
|
57
|
+
"joiner",
|
|
58
|
+
DocumentJoiner(join_mode="reciprocal_rank_fusion", top_k=3),
|
|
59
|
+
)
|
|
60
|
+
querying.connect("text_embedder", "retriever")
|
|
61
|
+
querying.connect("keyword_retriever", "joiner")
|
|
62
|
+
querying.connect("retriever", "joiner")
|
|
63
|
+
|
|
64
|
+
query = "cross-encoder"
|
|
65
|
+
results = querying.run({"text_embedder": {"text": query}, "keyword_retriever": {"query": query}})
|
|
66
|
+
|
|
67
|
+
for doc in results["joiner"]["documents"]:
|
|
68
|
+
print(doc)
|
|
69
|
+
print("-" * 10)
|
|
@@ -3,8 +3,8 @@ loaders:
|
|
|
3
3
|
search_path: [../src]
|
|
4
4
|
modules: [
|
|
5
5
|
"haystack_integrations.components.retrievers.pgvector.embedding_retriever",
|
|
6
|
+
"haystack_integrations.components.retrievers.pgvector.keyword_retriever",
|
|
6
7
|
"haystack_integrations.document_stores.pgvector.document_store",
|
|
7
|
-
"haystack_integrations.document_stores.pgvector.filters",
|
|
8
8
|
]
|
|
9
9
|
ignore_when_discovered: ["__init__"]
|
|
10
10
|
processors:
|
|
@@ -16,14 +16,15 @@ processors:
|
|
|
16
16
|
- type: smart
|
|
17
17
|
- type: crossref
|
|
18
18
|
renderer:
|
|
19
|
-
type: haystack_pydoc_tools.renderers.
|
|
19
|
+
type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer
|
|
20
20
|
excerpt: Pgvector integration for Haystack
|
|
21
|
-
category_slug:
|
|
21
|
+
category_slug: integrations-api
|
|
22
22
|
title: Pgvector
|
|
23
23
|
slug: integrations-pgvector
|
|
24
|
-
order:
|
|
24
|
+
order: 190
|
|
25
25
|
markdown:
|
|
26
26
|
descriptive_class_title: false
|
|
27
|
+
classdef_code_block: false
|
|
27
28
|
descriptive_module_title: true
|
|
28
29
|
add_method_class_prefix: true
|
|
29
30
|
add_member_class_prefix: false
|
|
@@ -14,6 +14,7 @@ authors = [
|
|
|
14
14
|
{ name = "deepset GmbH", email = "info@deepset.ai" },
|
|
15
15
|
]
|
|
16
16
|
classifiers = [
|
|
17
|
+
"License :: OSI Approved :: Apache Software License",
|
|
17
18
|
"Development Status :: 4 - Beta",
|
|
18
19
|
"Programming Language :: Python",
|
|
19
20
|
"Programming Language :: Python :: 3.8",
|
|
@@ -51,7 +52,7 @@ dependencies = [
|
|
|
51
52
|
"coverage[toml]>=6.5",
|
|
52
53
|
"pytest",
|
|
53
54
|
"ipython",
|
|
54
|
-
"haystack-pydoc-tools",
|
|
55
|
+
"haystack-pydoc-tools",
|
|
55
56
|
]
|
|
56
57
|
[tool.hatch.envs.default.scripts]
|
|
57
58
|
test = "pytest {args:tests}"
|
|
@@ -95,12 +96,12 @@ all = [
|
|
|
95
96
|
]
|
|
96
97
|
|
|
97
98
|
[tool.black]
|
|
98
|
-
target-version = ["
|
|
99
|
+
target-version = ["py38"]
|
|
99
100
|
line-length = 120
|
|
100
101
|
skip-string-normalization = true
|
|
101
102
|
|
|
102
103
|
[tool.ruff]
|
|
103
|
-
target-version = "
|
|
104
|
+
target-version = "py38"
|
|
104
105
|
line-length = 120
|
|
105
106
|
select = [
|
|
106
107
|
"A",
|
|
@@ -162,7 +163,7 @@ ban-relative-imports = "parents"
|
|
|
162
163
|
[tool.coverage.run]
|
|
163
164
|
source = ["haystack_integrations"]
|
|
164
165
|
branch = true
|
|
165
|
-
parallel =
|
|
166
|
+
parallel = false
|
|
166
167
|
|
|
167
168
|
[tool.coverage.report]
|
|
168
169
|
omit = ["*/tests/*", "*/__init__.py"]
|
|
@@ -173,6 +174,11 @@ exclude_lines = [
|
|
|
173
174
|
"if TYPE_CHECKING:",
|
|
174
175
|
]
|
|
175
176
|
|
|
177
|
+
[tool.pytest.ini_options]
|
|
178
|
+
markers = [
|
|
179
|
+
"integration: integration tests"
|
|
180
|
+
]
|
|
181
|
+
|
|
176
182
|
|
|
177
183
|
[[tool.mypy.overrides]]
|
|
178
184
|
module = [
|
|
@@ -2,5 +2,6 @@
|
|
|
2
2
|
#
|
|
3
3
|
# SPDX-License-Identifier: Apache-2.0
|
|
4
4
|
from .embedding_retriever import PgvectorEmbeddingRetriever
|
|
5
|
+
from .keyword_retriever import PgvectorKeywordRetriever
|
|
5
6
|
|
|
6
|
-
__all__ = ["PgvectorEmbeddingRetriever"]
|
|
7
|
+
__all__ = ["PgvectorEmbeddingRetriever", "PgvectorKeywordRetriever"]
|
|
@@ -12,9 +12,47 @@ from haystack_integrations.document_stores.pgvector.document_store import VALID_
|
|
|
12
12
|
@component
|
|
13
13
|
class PgvectorEmbeddingRetriever:
|
|
14
14
|
"""
|
|
15
|
-
Retrieves documents from the PgvectorDocumentStore
|
|
15
|
+
Retrieves documents from the `PgvectorDocumentStore`, based on their dense embeddings.
|
|
16
16
|
|
|
17
|
-
|
|
17
|
+
Example usage:
|
|
18
|
+
```python
|
|
19
|
+
from haystack.document_stores import DuplicatePolicy
|
|
20
|
+
from haystack import Document, Pipeline
|
|
21
|
+
from haystack.components.embedders import SentenceTransformersTextEmbedder, SentenceTransformersDocumentEmbedder
|
|
22
|
+
|
|
23
|
+
from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
|
|
24
|
+
from haystack_integrations.components.retrievers.pgvector import PgvectorEmbeddingRetriever
|
|
25
|
+
|
|
26
|
+
# Set an environment variable `PG_CONN_STR` with the connection string to your PostgreSQL database.
|
|
27
|
+
# e.g., "postgresql://USER:PASSWORD@HOST:PORT/DB_NAME"
|
|
28
|
+
|
|
29
|
+
document_store = PgvectorDocumentStore(
|
|
30
|
+
embedding_dimension=768,
|
|
31
|
+
vector_function="cosine_similarity",
|
|
32
|
+
recreate_table=True,
|
|
33
|
+
)
|
|
34
|
+
|
|
35
|
+
documents = [Document(content="There are over 7,000 languages spoken around the world today."),
|
|
36
|
+
Document(content="Elephants have been observed to behave in a way that indicates..."),
|
|
37
|
+
Document(content="In certain places, you can witness the phenomenon of bioluminescent waves.")]
|
|
38
|
+
|
|
39
|
+
document_embedder = SentenceTransformersDocumentEmbedder()
|
|
40
|
+
document_embedder.warm_up()
|
|
41
|
+
documents_with_embeddings = document_embedder.run(documents)
|
|
42
|
+
|
|
43
|
+
document_store.write_documents(documents_with_embeddings.get("documents"), policy=DuplicatePolicy.OVERWRITE)
|
|
44
|
+
|
|
45
|
+
query_pipeline = Pipeline()
|
|
46
|
+
query_pipeline.add_component("text_embedder", SentenceTransformersTextEmbedder())
|
|
47
|
+
query_pipeline.add_component("retriever", PgvectorEmbeddingRetriever(document_store=document_store))
|
|
48
|
+
query_pipeline.connect("text_embedder.embedding", "retriever.query_embedding")
|
|
49
|
+
|
|
50
|
+
query = "How many languages are there?"
|
|
51
|
+
|
|
52
|
+
res = query_pipeline.run({"text_embedder": {"text": query}})
|
|
53
|
+
|
|
54
|
+
assert res['retriever']['documents'][0].content == "There are over 7,000 languages spoken around the world today."
|
|
55
|
+
```
|
|
18
56
|
"""
|
|
19
57
|
|
|
20
58
|
def __init__(
|
|
@@ -26,23 +64,20 @@ class PgvectorEmbeddingRetriever:
|
|
|
26
64
|
vector_function: Optional[Literal["cosine_similarity", "inner_product", "l2_distance"]] = None,
|
|
27
65
|
):
|
|
28
66
|
"""
|
|
29
|
-
|
|
30
|
-
|
|
31
|
-
:param
|
|
32
|
-
:param filters: Filters applied to the retrieved Documents. Defaults to None.
|
|
33
|
-
:param top_k: Maximum number of Documents to return, defaults to 10.
|
|
67
|
+
:param document_store: An instance of `PgvectorDocumentStore`.
|
|
68
|
+
:param filters: Filters applied to the retrieved Documents.
|
|
69
|
+
:param top_k: Maximum number of Documents to return.
|
|
34
70
|
:param vector_function: The similarity function to use when searching for similar embeddings.
|
|
35
71
|
Defaults to the one set in the `document_store` instance.
|
|
36
|
-
"cosine_similarity" and "inner_product" are similarity functions and
|
|
72
|
+
`"cosine_similarity"` and `"inner_product"` are similarity functions and
|
|
37
73
|
higher scores indicate greater similarity between the documents.
|
|
38
|
-
"l2_distance" returns the straight-line distance between vectors,
|
|
74
|
+
`"l2_distance"` returns the straight-line distance between vectors,
|
|
39
75
|
and the most similar documents are the ones with the smallest score.
|
|
40
|
-
|
|
41
|
-
Important: if the document store is using the "hnsw" search strategy, the vector function
|
|
76
|
+
**Important**: if the document store is using the `"hnsw"` search strategy, the vector function
|
|
42
77
|
should match the one utilized during index creation to take advantage of the index.
|
|
43
|
-
:type vector_function: Literal["cosine_similarity", "inner_product", "l2_distance"]
|
|
44
78
|
|
|
45
|
-
:raises ValueError: If `document_store` is not an instance of PgvectorDocumentStore
|
|
79
|
+
:raises ValueError: If `document_store` is not an instance of `PgvectorDocumentStore` or if `vector_function`
|
|
80
|
+
is not one of the valid options.
|
|
46
81
|
"""
|
|
47
82
|
if not isinstance(document_store, PgvectorDocumentStore):
|
|
48
83
|
msg = "document_store must be an instance of PgvectorDocumentStore"
|
|
@@ -58,6 +93,12 @@ class PgvectorEmbeddingRetriever:
|
|
|
58
93
|
self.vector_function = vector_function or document_store.vector_function
|
|
59
94
|
|
|
60
95
|
def to_dict(self) -> Dict[str, Any]:
|
|
96
|
+
"""
|
|
97
|
+
Serializes the component to a dictionary.
|
|
98
|
+
|
|
99
|
+
:returns:
|
|
100
|
+
Dictionary with serialized data.
|
|
101
|
+
"""
|
|
61
102
|
return default_to_dict(
|
|
62
103
|
self,
|
|
63
104
|
filters=self.filters,
|
|
@@ -68,6 +109,14 @@ class PgvectorEmbeddingRetriever:
|
|
|
68
109
|
|
|
69
110
|
@classmethod
|
|
70
111
|
def from_dict(cls, data: Dict[str, Any]) -> "PgvectorEmbeddingRetriever":
|
|
112
|
+
"""
|
|
113
|
+
Deserializes the component from a dictionary.
|
|
114
|
+
|
|
115
|
+
:param data:
|
|
116
|
+
Dictionary to deserialize from.
|
|
117
|
+
:returns:
|
|
118
|
+
Deserialized component.
|
|
119
|
+
"""
|
|
71
120
|
doc_store_params = data["init_parameters"]["document_store"]
|
|
72
121
|
data["init_parameters"]["document_store"] = PgvectorDocumentStore.from_dict(doc_store_params)
|
|
73
122
|
return default_from_dict(cls, data)
|
|
@@ -81,14 +130,14 @@ class PgvectorEmbeddingRetriever:
|
|
|
81
130
|
vector_function: Optional[Literal["cosine_similarity", "inner_product", "l2_distance"]] = None,
|
|
82
131
|
):
|
|
83
132
|
"""
|
|
84
|
-
Retrieve documents from the PgvectorDocumentStore
|
|
133
|
+
Retrieve documents from the `PgvectorDocumentStore`, based on their embeddings.
|
|
85
134
|
|
|
86
135
|
:param query_embedding: Embedding of the query.
|
|
87
136
|
:param filters: Filters applied to the retrieved Documents.
|
|
88
137
|
:param top_k: Maximum number of Documents to return.
|
|
89
138
|
:param vector_function: The similarity function to use when searching for similar embeddings.
|
|
90
|
-
|
|
91
|
-
:
|
|
139
|
+
|
|
140
|
+
:returns: List of Documents similar to `query_embedding`.
|
|
92
141
|
"""
|
|
93
142
|
filters = filters or self.filters
|
|
94
143
|
top_k = top_k or self.top_k
|
|
@@ -0,0 +1,123 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: 2023-present deepset GmbH <info@deepset.ai>
|
|
2
|
+
#
|
|
3
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
+
from typing import Any, Dict, List, Optional
|
|
5
|
+
|
|
6
|
+
from haystack import component, default_from_dict, default_to_dict
|
|
7
|
+
from haystack.dataclasses import Document
|
|
8
|
+
from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
@component
|
|
12
|
+
class PgvectorKeywordRetriever:
|
|
13
|
+
"""
|
|
14
|
+
Retrieve documents from the `PgvectorDocumentStore`, based on keywords.
|
|
15
|
+
|
|
16
|
+
To rank the documents, the `ts_rank_cd` function of PostgreSQL is used.
|
|
17
|
+
It considers how often the query terms appear in the document, how close together the terms are in the document,
|
|
18
|
+
and how important is the part of the document where they occur.
|
|
19
|
+
For more details, see
|
|
20
|
+
[Postgres documentation](https://www.postgresql.org/docs/current/textsearch-controls.html#TEXTSEARCH-RANKING).
|
|
21
|
+
|
|
22
|
+
Usage example:
|
|
23
|
+
```python
|
|
24
|
+
from haystack.document_stores import DuplicatePolicy
|
|
25
|
+
from haystack import Document
|
|
26
|
+
|
|
27
|
+
from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
|
|
28
|
+
from haystack_integrations.components.retrievers.pgvector import PgvectorKeywordRetriever
|
|
29
|
+
|
|
30
|
+
# Set an environment variable `PG_CONN_STR` with the connection string to your PostgreSQL database.
|
|
31
|
+
# e.g., "postgresql://USER:PASSWORD@HOST:PORT/DB_NAME"
|
|
32
|
+
|
|
33
|
+
document_store = PgvectorDocumentStore(language="english", recreate_table=True)
|
|
34
|
+
|
|
35
|
+
documents = [Document(content="There are over 7,000 languages spoken around the world today."),
|
|
36
|
+
Document(content="Elephants have been observed to behave in a way that indicates..."),
|
|
37
|
+
Document(content="In certain places, you can witness the phenomenon of bioluminescent waves.")]
|
|
38
|
+
|
|
39
|
+
document_store.write_documents(documents_with_embeddings.get("documents"), policy=DuplicatePolicy.OVERWRITE)
|
|
40
|
+
|
|
41
|
+
retriever = PgvectorKeywordRetriever(document_store=document_store)
|
|
42
|
+
|
|
43
|
+
result = retriever.run(query="languages")
|
|
44
|
+
|
|
45
|
+
assert res['retriever']['documents'][0].content == "There are over 7,000 languages spoken around the world today."
|
|
46
|
+
"""
|
|
47
|
+
|
|
48
|
+
def __init__(
|
|
49
|
+
self,
|
|
50
|
+
*,
|
|
51
|
+
document_store: PgvectorDocumentStore,
|
|
52
|
+
filters: Optional[Dict[str, Any]] = None,
|
|
53
|
+
top_k: int = 10,
|
|
54
|
+
):
|
|
55
|
+
"""
|
|
56
|
+
:param document_store: An instance of `PgvectorDocumentStore`.
|
|
57
|
+
:param filters: Filters applied to the retrieved Documents.
|
|
58
|
+
:param top_k: Maximum number of Documents to return.
|
|
59
|
+
|
|
60
|
+
:raises ValueError: If `document_store` is not an instance of `PgvectorDocumentStore`.
|
|
61
|
+
"""
|
|
62
|
+
if not isinstance(document_store, PgvectorDocumentStore):
|
|
63
|
+
msg = "document_store must be an instance of PgvectorDocumentStore"
|
|
64
|
+
raise ValueError(msg)
|
|
65
|
+
|
|
66
|
+
self.document_store = document_store
|
|
67
|
+
self.filters = filters or {}
|
|
68
|
+
self.top_k = top_k
|
|
69
|
+
|
|
70
|
+
def to_dict(self) -> Dict[str, Any]:
|
|
71
|
+
"""
|
|
72
|
+
Serializes the component to a dictionary.
|
|
73
|
+
|
|
74
|
+
:returns:
|
|
75
|
+
Dictionary with serialized data.
|
|
76
|
+
"""
|
|
77
|
+
return default_to_dict(
|
|
78
|
+
self,
|
|
79
|
+
filters=self.filters,
|
|
80
|
+
top_k=self.top_k,
|
|
81
|
+
document_store=self.document_store.to_dict(),
|
|
82
|
+
)
|
|
83
|
+
|
|
84
|
+
@classmethod
|
|
85
|
+
def from_dict(cls, data: Dict[str, Any]) -> "PgvectorKeywordRetriever":
|
|
86
|
+
"""
|
|
87
|
+
Deserializes the component from a dictionary.
|
|
88
|
+
|
|
89
|
+
:param data:
|
|
90
|
+
Dictionary to deserialize from.
|
|
91
|
+
:returns:
|
|
92
|
+
Deserialized component.
|
|
93
|
+
"""
|
|
94
|
+
doc_store_params = data["init_parameters"]["document_store"]
|
|
95
|
+
data["init_parameters"]["document_store"] = PgvectorDocumentStore.from_dict(doc_store_params)
|
|
96
|
+
return default_from_dict(cls, data)
|
|
97
|
+
|
|
98
|
+
@component.output_types(documents=List[Document])
|
|
99
|
+
def run(
|
|
100
|
+
self,
|
|
101
|
+
query: str,
|
|
102
|
+
filters: Optional[Dict[str, Any]] = None,
|
|
103
|
+
top_k: Optional[int] = None,
|
|
104
|
+
):
|
|
105
|
+
"""
|
|
106
|
+
Retrieve documents from the `PgvectorDocumentStore`, based on keywords.
|
|
107
|
+
|
|
108
|
+
:param query: String to search in `Document`s' content.
|
|
109
|
+
:param filters: Filters applied to the retrieved Documents.
|
|
110
|
+
:param top_k: Maximum number of Documents to return.
|
|
111
|
+
|
|
112
|
+
:returns: A dictionary with the following keys:
|
|
113
|
+
- `documents`: List of `Document`s that match the query.
|
|
114
|
+
"""
|
|
115
|
+
filters = filters or self.filters
|
|
116
|
+
top_k = top_k or self.top_k
|
|
117
|
+
|
|
118
|
+
docs = self.document_store._keyword_retrieval(
|
|
119
|
+
query=query,
|
|
120
|
+
filters=filters,
|
|
121
|
+
top_k=top_k,
|
|
122
|
+
)
|
|
123
|
+
return {"documents": docs}
|