pgvector-haystack 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl
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.
- haystack_integrations/components/retrievers/pgvector/__init__.py +2 -1
- haystack_integrations/components/retrievers/pgvector/embedding_retriever.py +65 -16
- haystack_integrations/components/retrievers/pgvector/keyword_retriever.py +123 -0
- haystack_integrations/document_stores/pgvector/document_store.py +141 -27
- haystack_integrations/document_stores/pgvector/filters.py +15 -13
- {pgvector_haystack-0.1.0.dist-info → pgvector_haystack-0.2.0.dist-info}/METADATA +3 -2
- pgvector_haystack-0.2.0.dist-info/RECORD +10 -0
- {pgvector_haystack-0.1.0.dist-info → pgvector_haystack-0.2.0.dist-info}/WHEEL +1 -1
- pgvector_haystack-0.1.0.dist-info/RECORD +0 -9
- {pgvector_haystack-0.1.0.dist-info → pgvector_haystack-0.2.0.dist-info}/licenses/LICENSE.txt +0 -0
|
@@ -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}
|
|
@@ -53,6 +53,12 @@ blob_mime_type = EXCLUDED.blob_mime_type,
|
|
|
53
53
|
meta = EXCLUDED.meta
|
|
54
54
|
"""
|
|
55
55
|
|
|
56
|
+
KEYWORD_QUERY = """
|
|
57
|
+
SELECT {table_name}.*, ts_rank_cd(to_tsvector({language}, content), query) AS score
|
|
58
|
+
FROM {table_name}, plainto_tsquery({language}, %s) query
|
|
59
|
+
WHERE to_tsvector({language}, content) @@ query
|
|
60
|
+
"""
|
|
61
|
+
|
|
56
62
|
VALID_VECTOR_FUNCTIONS = ["cosine_similarity", "inner_product", "l2_distance"]
|
|
57
63
|
|
|
58
64
|
VECTOR_FUNCTION_TO_POSTGRESQL_OPS = {
|
|
@@ -65,13 +71,20 @@ HNSW_INDEX_CREATION_VALID_KWARGS = ["m", "ef_construction"]
|
|
|
65
71
|
|
|
66
72
|
HNSW_INDEX_NAME = "haystack_hnsw_index"
|
|
67
73
|
|
|
74
|
+
KEYWORD_INDEX_NAME = "haystack_keyword_index"
|
|
75
|
+
|
|
68
76
|
|
|
69
77
|
class PgvectorDocumentStore:
|
|
78
|
+
"""
|
|
79
|
+
A Document Store using PostgreSQL with the [pgvector extension](https://github.com/pgvector/pgvector) installed.
|
|
80
|
+
"""
|
|
81
|
+
|
|
70
82
|
def __init__(
|
|
71
83
|
self,
|
|
72
84
|
*,
|
|
73
85
|
connection_string: Secret = Secret.from_env_var("PG_CONN_STR"),
|
|
74
86
|
table_name: str = "haystack_documents",
|
|
87
|
+
language: str = "english",
|
|
75
88
|
embedding_dimension: int = 768,
|
|
76
89
|
vector_function: Literal["cosine_similarity", "inner_product", "l2_distance"] = "cosine_similarity",
|
|
77
90
|
recreate_table: bool = False,
|
|
@@ -86,36 +99,37 @@ class PgvectorDocumentStore:
|
|
|
86
99
|
A specific table to store Haystack documents will be created if it doesn't exist yet.
|
|
87
100
|
|
|
88
101
|
:param connection_string: The connection string to use to connect to the PostgreSQL database, defined as an
|
|
89
|
-
environment variable, e.g.: PG_CONN_STR="postgresql://USER:PASSWORD@HOST:PORT/DB_NAME"
|
|
90
|
-
:param table_name: The name of the table to use to store Haystack documents.
|
|
91
|
-
:param
|
|
102
|
+
environment variable, e.g.: `PG_CONN_STR="postgresql://USER:PASSWORD@HOST:PORT/DB_NAME"`
|
|
103
|
+
:param table_name: The name of the table to use to store Haystack documents.
|
|
104
|
+
:param language: The language to be used to parse query and document content in keyword retrieval.
|
|
105
|
+
To see the list of available languages, you can run the following SQL query in your PostgreSQL database:
|
|
106
|
+
`SELECT cfgname FROM pg_ts_config;`.
|
|
107
|
+
More information can be found in this [StackOverflow answer](https://stackoverflow.com/a/39752553).
|
|
108
|
+
:param embedding_dimension: The dimension of the embedding.
|
|
92
109
|
:param vector_function: The similarity function to use when searching for similar embeddings.
|
|
93
|
-
|
|
110
|
+
`"cosine_similarity"` and `"inner_product"` are similarity functions and
|
|
94
111
|
higher scores indicate greater similarity between the documents.
|
|
95
|
-
"l2_distance" returns the straight-line distance between vectors,
|
|
112
|
+
`"l2_distance"` returns the straight-line distance between vectors,
|
|
96
113
|
and the most similar documents are the ones with the smallest score.
|
|
97
|
-
|
|
98
|
-
Important: when using the "hnsw" search strategy, an index will be created that depends on the
|
|
114
|
+
**Important**: when using the `"hnsw"` search strategy, an index will be created that depends on the
|
|
99
115
|
`vector_function` passed here. Make sure subsequent queries will keep using the same
|
|
100
116
|
vector similarity function in order to take advantage of the index.
|
|
101
|
-
:
|
|
102
|
-
:param recreate_table: Whether to recreate the table if it already exists. Defaults to False.
|
|
117
|
+
:param recreate_table: Whether to recreate the table if it already exists.
|
|
103
118
|
:param search_strategy: The search strategy to use when searching for similar embeddings.
|
|
104
|
-
|
|
119
|
+
`"exact_nearest_neighbor"` provides perfect recall but can be slow for large numbers of documents.
|
|
120
|
+
`"hnsw"` is an approximate nearest neighbor search strategy,
|
|
105
121
|
which trades off some accuracy for speed; it is recommended for large numbers of documents.
|
|
106
|
-
|
|
107
|
-
Important: when using the "hnsw" search strategy, an index will be created that depends on the
|
|
122
|
+
**Important**: when using the `"hnsw"` search strategy, an index will be created that depends on the
|
|
108
123
|
`vector_function` passed here. Make sure subsequent queries will keep using the same
|
|
109
124
|
vector similarity function in order to take advantage of the index.
|
|
110
|
-
:type search_strategy: Literal["exact_nearest_neighbor", "hnsw"]
|
|
111
125
|
:param hnsw_recreate_index_if_exists: Whether to recreate the HNSW index if it already exists.
|
|
112
|
-
|
|
126
|
+
Only used if search_strategy is set to `"hnsw"`.
|
|
113
127
|
:param hnsw_index_creation_kwargs: Additional keyword arguments to pass to the HNSW index creation.
|
|
114
|
-
Only used if search_strategy is set to "hnsw"
|
|
115
|
-
pgvector documentation
|
|
116
|
-
:param hnsw_ef_search: The ef_search parameter to use at query time. Only used if search_strategy is set to
|
|
117
|
-
"hnsw"
|
|
118
|
-
https://github.com/pgvector/pgvector?tab=readme-ov-file#hnsw
|
|
128
|
+
Only used if search_strategy is set to `"hnsw"`. You can find the list of valid arguments in the
|
|
129
|
+
[pgvector documentation](https://github.com/pgvector/pgvector?tab=readme-ov-file#hnsw)
|
|
130
|
+
:param hnsw_ef_search: The `ef_search` parameter to use at query time. Only used if search_strategy is set to
|
|
131
|
+
`"hnsw"`. You can find more information about this parameter in the
|
|
132
|
+
[pgvector documentation](https://github.com/pgvector/pgvector?tab=readme-ov-file#hnsw).
|
|
119
133
|
"""
|
|
120
134
|
|
|
121
135
|
self.connection_string = connection_string
|
|
@@ -130,6 +144,7 @@ class PgvectorDocumentStore:
|
|
|
130
144
|
self.hnsw_recreate_index_if_exists = hnsw_recreate_index_if_exists
|
|
131
145
|
self.hnsw_index_creation_kwargs = hnsw_index_creation_kwargs or {}
|
|
132
146
|
self.hnsw_ef_search = hnsw_ef_search
|
|
147
|
+
self.language = language
|
|
133
148
|
|
|
134
149
|
connection = connect(self.connection_string.resolve_value())
|
|
135
150
|
connection.autocommit = True
|
|
@@ -145,11 +160,18 @@ class PgvectorDocumentStore:
|
|
|
145
160
|
if recreate_table:
|
|
146
161
|
self.delete_table()
|
|
147
162
|
self._create_table_if_not_exists()
|
|
163
|
+
self._create_keyword_index_if_not_exists()
|
|
148
164
|
|
|
149
165
|
if search_strategy == "hnsw":
|
|
150
166
|
self._handle_hnsw()
|
|
151
167
|
|
|
152
168
|
def to_dict(self) -> Dict[str, Any]:
|
|
169
|
+
"""
|
|
170
|
+
Serializes the component to a dictionary.
|
|
171
|
+
|
|
172
|
+
:returns:
|
|
173
|
+
Dictionary with serialized data.
|
|
174
|
+
"""
|
|
153
175
|
return default_to_dict(
|
|
154
176
|
self,
|
|
155
177
|
connection_string=self.connection_string.to_dict(),
|
|
@@ -161,10 +183,19 @@ class PgvectorDocumentStore:
|
|
|
161
183
|
hnsw_recreate_index_if_exists=self.hnsw_recreate_index_if_exists,
|
|
162
184
|
hnsw_index_creation_kwargs=self.hnsw_index_creation_kwargs,
|
|
163
185
|
hnsw_ef_search=self.hnsw_ef_search,
|
|
186
|
+
language=self.language,
|
|
164
187
|
)
|
|
165
188
|
|
|
166
189
|
@classmethod
|
|
167
190
|
def from_dict(cls, data: Dict[str, Any]) -> "PgvectorDocumentStore":
|
|
191
|
+
"""
|
|
192
|
+
Deserializes the component from a dictionary.
|
|
193
|
+
|
|
194
|
+
:param data:
|
|
195
|
+
Dictionary to deserialize from.
|
|
196
|
+
:returns:
|
|
197
|
+
Deserialized component.
|
|
198
|
+
"""
|
|
168
199
|
deserialize_secrets_inplace(data["init_parameters"], ["connection_string"])
|
|
169
200
|
return default_from_dict(cls, data)
|
|
170
201
|
|
|
@@ -209,16 +240,40 @@ class PgvectorDocumentStore:
|
|
|
209
240
|
def delete_table(self):
|
|
210
241
|
"""
|
|
211
242
|
Deletes the table used to store Haystack documents.
|
|
243
|
+
The name of the table (`table_name`) is defined when initializing the `PgvectorDocumentStore`.
|
|
212
244
|
"""
|
|
213
245
|
|
|
214
246
|
delete_sql = SQL("DROP TABLE IF EXISTS {table_name}").format(table_name=Identifier(self.table_name))
|
|
215
247
|
|
|
216
248
|
self._execute_sql(delete_sql, error_msg=f"Could not delete table {self.table_name} in PgvectorDocumentStore")
|
|
217
249
|
|
|
250
|
+
def _create_keyword_index_if_not_exists(self):
|
|
251
|
+
"""
|
|
252
|
+
Internal method to create the keyword index if not exists.
|
|
253
|
+
"""
|
|
254
|
+
index_exists = bool(
|
|
255
|
+
self._execute_sql(
|
|
256
|
+
"SELECT 1 FROM pg_indexes WHERE tablename = %s AND indexname = %s",
|
|
257
|
+
(self.table_name, KEYWORD_INDEX_NAME),
|
|
258
|
+
"Could not check if keyword index exists",
|
|
259
|
+
).fetchone()
|
|
260
|
+
)
|
|
261
|
+
|
|
262
|
+
sql_create_index = SQL(
|
|
263
|
+
"CREATE INDEX {index_name} ON {table_name} USING GIN (to_tsvector({language}, content))"
|
|
264
|
+
).format(
|
|
265
|
+
index_name=Identifier(KEYWORD_INDEX_NAME),
|
|
266
|
+
table_name=Identifier(self.table_name),
|
|
267
|
+
language=SQLLiteral(self.language),
|
|
268
|
+
)
|
|
269
|
+
|
|
270
|
+
if not index_exists:
|
|
271
|
+
self._execute_sql(sql_create_index, error_msg="Could not create keyword index on table")
|
|
272
|
+
|
|
218
273
|
def _handle_hnsw(self):
|
|
219
274
|
"""
|
|
220
275
|
Internal method to handle the HNSW index creation.
|
|
221
|
-
It also sets the hnsw.ef_search parameter for queries if it is specified.
|
|
276
|
+
It also sets the `hnsw.ef_search` parameter for queries if it is specified.
|
|
222
277
|
"""
|
|
223
278
|
|
|
224
279
|
if self.hnsw_ef_search:
|
|
@@ -295,7 +350,8 @@ class PgvectorDocumentStore:
|
|
|
295
350
|
refer to the [documentation](https://docs.haystack.deepset.ai/v2.0/docs/metadata-filtering)
|
|
296
351
|
|
|
297
352
|
:param filters: The filters to apply to the document list.
|
|
298
|
-
:
|
|
353
|
+
:raises TypeError: If `filters` is not a dictionary.
|
|
354
|
+
:returns: A list of Documents that match the given filters.
|
|
299
355
|
"""
|
|
300
356
|
if filters:
|
|
301
357
|
if not isinstance(filters, dict):
|
|
@@ -324,13 +380,13 @@ class PgvectorDocumentStore:
|
|
|
324
380
|
|
|
325
381
|
def write_documents(self, documents: List[Document], policy: DuplicatePolicy = DuplicatePolicy.NONE) -> int:
|
|
326
382
|
"""
|
|
327
|
-
Writes documents
|
|
383
|
+
Writes documents to the document store.
|
|
328
384
|
|
|
329
385
|
:param documents: A list of Documents to write to the document store.
|
|
330
386
|
:param policy: The duplicate policy to use when writing documents.
|
|
331
387
|
:raises DuplicateDocumentError: If a document with the same id already exists in the document store
|
|
332
|
-
and the policy is set to DuplicatePolicy.FAIL (or not specified).
|
|
333
|
-
:
|
|
388
|
+
and the policy is set to `DuplicatePolicy.FAIL` (or not specified).
|
|
389
|
+
:returns: The number of documents written to the document store.
|
|
334
390
|
"""
|
|
335
391
|
|
|
336
392
|
if len(documents) > 0:
|
|
@@ -398,6 +454,16 @@ class PgvectorDocumentStore:
|
|
|
398
454
|
db_document["dataframe"] = Jsonb(db_document["dataframe"]) if db_document["dataframe"] else None
|
|
399
455
|
db_document["meta"] = Jsonb(db_document["meta"])
|
|
400
456
|
|
|
457
|
+
if "sparse_embedding" in db_document:
|
|
458
|
+
sparse_embedding = db_document.pop("sparse_embedding", None)
|
|
459
|
+
if sparse_embedding:
|
|
460
|
+
logger.warning(
|
|
461
|
+
"Document %s has the `sparse_embedding` field set,"
|
|
462
|
+
"but storing sparse embeddings in Pgvector is not currently supported."
|
|
463
|
+
"The `sparse_embedding` field will be ignored.",
|
|
464
|
+
db_document["id"],
|
|
465
|
+
)
|
|
466
|
+
|
|
401
467
|
db_documents.append(db_document)
|
|
402
468
|
|
|
403
469
|
return db_documents
|
|
@@ -432,7 +498,7 @@ class PgvectorDocumentStore:
|
|
|
432
498
|
|
|
433
499
|
def delete_documents(self, document_ids: List[str]) -> None:
|
|
434
500
|
"""
|
|
435
|
-
Deletes
|
|
501
|
+
Deletes documents that match the provided `document_ids` from the document store.
|
|
436
502
|
|
|
437
503
|
:param document_ids: the document ids to delete
|
|
438
504
|
"""
|
|
@@ -448,6 +514,54 @@ class PgvectorDocumentStore:
|
|
|
448
514
|
|
|
449
515
|
self._execute_sql(delete_sql, error_msg="Could not delete documents from PgvectorDocumentStore")
|
|
450
516
|
|
|
517
|
+
def _keyword_retrieval(
|
|
518
|
+
self,
|
|
519
|
+
query: str,
|
|
520
|
+
*,
|
|
521
|
+
filters: Optional[Dict[str, Any]] = None,
|
|
522
|
+
top_k: int = 10,
|
|
523
|
+
) -> List[Document]:
|
|
524
|
+
"""
|
|
525
|
+
Retrieves documents that are most similar to the query using a full-text search.
|
|
526
|
+
|
|
527
|
+
This method is not meant to be part of the public interface of
|
|
528
|
+
`PgvectorDocumentStore` and it should not be called directly.
|
|
529
|
+
`PgvectorKeywordRetriever` uses this method directly and is the public interface for it.
|
|
530
|
+
|
|
531
|
+
:returns: List of Documents that are most similar to `query`
|
|
532
|
+
"""
|
|
533
|
+
if not query:
|
|
534
|
+
msg = "query must be a non-empty string"
|
|
535
|
+
raise ValueError(msg)
|
|
536
|
+
|
|
537
|
+
sql_select = SQL(KEYWORD_QUERY).format(
|
|
538
|
+
table_name=Identifier(self.table_name),
|
|
539
|
+
language=SQLLiteral(self.language),
|
|
540
|
+
query=SQLLiteral(query),
|
|
541
|
+
)
|
|
542
|
+
|
|
543
|
+
where_params = ()
|
|
544
|
+
sql_where_clause = SQL("")
|
|
545
|
+
if filters:
|
|
546
|
+
sql_where_clause, where_params = _convert_filters_to_where_clause_and_params(
|
|
547
|
+
filters=filters, operator="AND"
|
|
548
|
+
)
|
|
549
|
+
|
|
550
|
+
sql_sort = SQL(" ORDER BY score DESC LIMIT {top_k}").format(top_k=SQLLiteral(top_k))
|
|
551
|
+
|
|
552
|
+
sql_query = sql_select + sql_where_clause + sql_sort
|
|
553
|
+
|
|
554
|
+
result = self._execute_sql(
|
|
555
|
+
sql_query,
|
|
556
|
+
(query, *where_params),
|
|
557
|
+
error_msg="Could not retrieve documents from PgvectorDocumentStore.",
|
|
558
|
+
cursor=self._dict_cursor,
|
|
559
|
+
)
|
|
560
|
+
|
|
561
|
+
records = result.fetchall()
|
|
562
|
+
docs = self._from_pg_to_haystack_documents(records)
|
|
563
|
+
return docs
|
|
564
|
+
|
|
451
565
|
def _embedding_retrieval(
|
|
452
566
|
self,
|
|
453
567
|
query_embedding: List[float],
|
|
@@ -462,8 +576,8 @@ class PgvectorDocumentStore:
|
|
|
462
576
|
This method is not meant to be part of the public interface of
|
|
463
577
|
`PgvectorDocumentStore` and it should not be called directly.
|
|
464
578
|
`PgvectorEmbeddingRetriever` uses this method directly and is the public interface for it.
|
|
465
|
-
|
|
466
|
-
:
|
|
579
|
+
|
|
580
|
+
:returns: List of Documents that are most similar to `query_embedding`
|
|
467
581
|
"""
|
|
468
582
|
|
|
469
583
|
if not query_embedding:
|
|
@@ -3,7 +3,7 @@
|
|
|
3
3
|
# SPDX-License-Identifier: Apache-2.0
|
|
4
4
|
from datetime import datetime
|
|
5
5
|
from itertools import chain
|
|
6
|
-
from typing import Any, Dict, List
|
|
6
|
+
from typing import Any, Dict, List, Literal, Tuple
|
|
7
7
|
|
|
8
8
|
from haystack.errors import FilterError
|
|
9
9
|
from pandas import DataFrame
|
|
@@ -22,7 +22,9 @@ PYTHON_TYPES_TO_PG_TYPES = {
|
|
|
22
22
|
NO_VALUE = "no_value"
|
|
23
23
|
|
|
24
24
|
|
|
25
|
-
def _convert_filters_to_where_clause_and_params(
|
|
25
|
+
def _convert_filters_to_where_clause_and_params(
|
|
26
|
+
filters: Dict[str, Any], operator: Literal["WHERE", "AND"] = "WHERE"
|
|
27
|
+
) -> Tuple[SQL, Tuple]:
|
|
26
28
|
"""
|
|
27
29
|
Convert Haystack filters to a WHERE clause and a tuple of params to query PostgreSQL.
|
|
28
30
|
"""
|
|
@@ -31,13 +33,13 @@ def _convert_filters_to_where_clause_and_params(filters: Dict[str, Any]) -> tupl
|
|
|
31
33
|
else:
|
|
32
34
|
query, values = _parse_logical_condition(filters)
|
|
33
35
|
|
|
34
|
-
where_clause = SQL("
|
|
36
|
+
where_clause = SQL(f" {operator} ") + SQL(query)
|
|
35
37
|
params = tuple(value for value in values if value != NO_VALUE)
|
|
36
38
|
|
|
37
39
|
return where_clause, params
|
|
38
40
|
|
|
39
41
|
|
|
40
|
-
def _parse_logical_condition(condition: Dict[str, Any]) ->
|
|
42
|
+
def _parse_logical_condition(condition: Dict[str, Any]) -> Tuple[str, List[Any]]:
|
|
41
43
|
if "operator" not in condition:
|
|
42
44
|
msg = f"'operator' key missing in {condition}"
|
|
43
45
|
raise FilterError(msg)
|
|
@@ -77,7 +79,7 @@ def _parse_logical_condition(condition: Dict[str, Any]) -> tuple[str, List[Any]]
|
|
|
77
79
|
return sql_query, values
|
|
78
80
|
|
|
79
81
|
|
|
80
|
-
def _parse_comparison_condition(condition: Dict[str, Any]) ->
|
|
82
|
+
def _parse_comparison_condition(condition: Dict[str, Any]) -> Tuple[str, List[Any]]:
|
|
81
83
|
field: str = condition["field"]
|
|
82
84
|
if "operator" not in condition:
|
|
83
85
|
msg = f"'operator' key missing in {condition}"
|
|
@@ -132,20 +134,20 @@ def _treat_meta_field(field: str, value: Any) -> str:
|
|
|
132
134
|
return field
|
|
133
135
|
|
|
134
136
|
|
|
135
|
-
def _equal(field: str, value: Any) ->
|
|
137
|
+
def _equal(field: str, value: Any) -> Tuple[str, Any]:
|
|
136
138
|
if value is None:
|
|
137
139
|
# NO_VALUE is a placeholder that will be removed in _convert_filters_to_where_clause_and_params
|
|
138
140
|
return f"{field} IS NULL", NO_VALUE
|
|
139
141
|
return f"{field} = %s", value
|
|
140
142
|
|
|
141
143
|
|
|
142
|
-
def _not_equal(field: str, value: Any) ->
|
|
144
|
+
def _not_equal(field: str, value: Any) -> Tuple[str, Any]:
|
|
143
145
|
# we use IS DISTINCT FROM to correctly handle NULL values
|
|
144
146
|
# (not handled by !=)
|
|
145
147
|
return f"{field} IS DISTINCT FROM %s", value
|
|
146
148
|
|
|
147
149
|
|
|
148
|
-
def _greater_than(field: str, value: Any) ->
|
|
150
|
+
def _greater_than(field: str, value: Any) -> Tuple[str, Any]:
|
|
149
151
|
if isinstance(value, str):
|
|
150
152
|
try:
|
|
151
153
|
datetime.fromisoformat(value)
|
|
@@ -162,7 +164,7 @@ def _greater_than(field: str, value: Any) -> tuple[str, Any]:
|
|
|
162
164
|
return f"{field} > %s", value
|
|
163
165
|
|
|
164
166
|
|
|
165
|
-
def _greater_than_equal(field: str, value: Any) ->
|
|
167
|
+
def _greater_than_equal(field: str, value: Any) -> Tuple[str, Any]:
|
|
166
168
|
if isinstance(value, str):
|
|
167
169
|
try:
|
|
168
170
|
datetime.fromisoformat(value)
|
|
@@ -179,7 +181,7 @@ def _greater_than_equal(field: str, value: Any) -> tuple[str, Any]:
|
|
|
179
181
|
return f"{field} >= %s", value
|
|
180
182
|
|
|
181
183
|
|
|
182
|
-
def _less_than(field: str, value: Any) ->
|
|
184
|
+
def _less_than(field: str, value: Any) -> Tuple[str, Any]:
|
|
183
185
|
if isinstance(value, str):
|
|
184
186
|
try:
|
|
185
187
|
datetime.fromisoformat(value)
|
|
@@ -196,7 +198,7 @@ def _less_than(field: str, value: Any) -> tuple[str, Any]:
|
|
|
196
198
|
return f"{field} < %s", value
|
|
197
199
|
|
|
198
200
|
|
|
199
|
-
def _less_than_equal(field: str, value: Any) ->
|
|
201
|
+
def _less_than_equal(field: str, value: Any) -> Tuple[str, Any]:
|
|
200
202
|
if isinstance(value, str):
|
|
201
203
|
try:
|
|
202
204
|
datetime.fromisoformat(value)
|
|
@@ -213,7 +215,7 @@ def _less_than_equal(field: str, value: Any) -> tuple[str, Any]:
|
|
|
213
215
|
return f"{field} <= %s", value
|
|
214
216
|
|
|
215
217
|
|
|
216
|
-
def _not_in(field: str, value: Any) ->
|
|
218
|
+
def _not_in(field: str, value: Any) -> Tuple[str, List]:
|
|
217
219
|
if not isinstance(value, list):
|
|
218
220
|
msg = f"{field}'s value must be a list when using 'not in' comparator in Pinecone"
|
|
219
221
|
raise FilterError(msg)
|
|
@@ -221,7 +223,7 @@ def _not_in(field: str, value: Any) -> tuple[str, List]:
|
|
|
221
223
|
return f"{field} IS NULL OR {field} != ALL(%s)", [value]
|
|
222
224
|
|
|
223
225
|
|
|
224
|
-
def _in(field: str, value: Any) ->
|
|
226
|
+
def _in(field: str, value: Any) -> Tuple[str, List]:
|
|
225
227
|
if not isinstance(value, list):
|
|
226
228
|
msg = f"{field}'s value must be a list when using 'in' comparator in Pinecone"
|
|
227
229
|
raise FilterError(msg)
|
|
@@ -1,6 +1,6 @@
|
|
|
1
|
-
Metadata-Version: 2.
|
|
1
|
+
Metadata-Version: 2.3
|
|
2
2
|
Name: pgvector-haystack
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.2.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
|
|
@@ -0,0 +1,10 @@
|
|
|
1
|
+
haystack_integrations/components/retrievers/pgvector/__init__.py,sha256=n4q4pUApJ9DZakJ_JL-egwOzlefdnI_ZxjU1Lu44Sqo,296
|
|
2
|
+
haystack_integrations/components/retrievers/pgvector/embedding_retriever.py,sha256=6LuQmF7JmPOCU_I10ksPcbDC6QMu-2SOq0--BQxBxpg,6523
|
|
3
|
+
haystack_integrations/components/retrievers/pgvector/keyword_retriever.py,sha256=fAJ3M2l_Pt8zKqHES1l1L4gM16BOHeZWq8xc18Bhq7A,4624
|
|
4
|
+
haystack_integrations/document_stores/pgvector/__init__.py,sha256=Ge-7ggarfr3x8UqNdFHUuFGiu52u49sOVe3l16WEwZU,197
|
|
5
|
+
haystack_integrations/document_stores/pgvector/document_store.py,sha256=KY0FZGTjz4ejxKNqQ9UHDJUArsghwf7Iag45YfXgm2g,27011
|
|
6
|
+
haystack_integrations/document_stores/pgvector/filters.py,sha256=751_Z1QnXJCG0GWC2jcfeZ25NNYfJewQEIUCeJS70PA,8212
|
|
7
|
+
pgvector_haystack-0.2.0.dist-info/METADATA,sha256=fD4i8wnlapdeHvJMpj11zlmGS_MxN_d6GDFvc3W_i1A,2219
|
|
8
|
+
pgvector_haystack-0.2.0.dist-info/WHEEL,sha256=zEMcRr9Kr03x1ozGwg5v9NQBKn3kndp6LSoSlVg-jhU,87
|
|
9
|
+
pgvector_haystack-0.2.0.dist-info/licenses/LICENSE.txt,sha256=B05uMshqTA74s-0ltyHKI6yoPfJ3zYgQbvcXfDVGFf8,10280
|
|
10
|
+
pgvector_haystack-0.2.0.dist-info/RECORD,,
|
|
@@ -1,9 +0,0 @@
|
|
|
1
|
-
haystack_integrations/components/retrievers/pgvector/__init__.py,sha256=Ecyh_BCqILZVpuYVw9m8d7nwv0PoPomZxjfBSIYDVkI,212
|
|
2
|
-
haystack_integrations/components/retrievers/pgvector/embedding_retriever.py,sha256=n_byd731W5T47PjHQbGM7DSPkKeTgKq5awBATEN4bW0,4536
|
|
3
|
-
haystack_integrations/document_stores/pgvector/__init__.py,sha256=Ge-7ggarfr3x8UqNdFHUuFGiu52u49sOVe3l16WEwZU,197
|
|
4
|
-
haystack_integrations/document_stores/pgvector/document_store.py,sha256=ydhBeypmCzwNiA5TjtTBMFwroJ5n10HwdI1B41XAIkk,22749
|
|
5
|
-
haystack_integrations/document_stores/pgvector/filters.py,sha256=_yj-5ZVccFc_uzDcwdVxa3UpiI9PzaNUErwEbp-lSo0,8139
|
|
6
|
-
pgvector_haystack-0.1.0.dist-info/METADATA,sha256=r4UtLicr9GOmQYWFOiDx6Pb1vEkIjiG9w3FHWVB0bpk,2156
|
|
7
|
-
pgvector_haystack-0.1.0.dist-info/WHEEL,sha256=TJPnKdtrSue7xZ_AVGkp9YXcvDrobsjBds1du3Nx6dc,87
|
|
8
|
-
pgvector_haystack-0.1.0.dist-info/licenses/LICENSE.txt,sha256=B05uMshqTA74s-0ltyHKI6yoPfJ3zYgQbvcXfDVGFf8,10280
|
|
9
|
-
pgvector_haystack-0.1.0.dist-info/RECORD,,
|
{pgvector_haystack-0.1.0.dist-info → pgvector_haystack-0.2.0.dist-info}/licenses/LICENSE.txt
RENAMED
|
File without changes
|