pgvector-haystack 0.1.0__tar.gz → 0.2.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (24) hide show
  1. {pgvector_haystack-0.1.0 → pgvector_haystack-0.2.0}/.gitignore +1 -0
  2. {pgvector_haystack-0.1.0 → pgvector_haystack-0.2.0}/PKG-INFO +3 -2
  3. pgvector_haystack-0.1.0/examples/example.py → pgvector_haystack-0.2.0/examples/embedding_retrieval.py +2 -2
  4. {pgvector_haystack-0.1.0 → pgvector_haystack-0.2.0}/pydoc/config.yml +4 -4
  5. {pgvector_haystack-0.1.0 → pgvector_haystack-0.2.0}/pyproject.toml +10 -4
  6. {pgvector_haystack-0.1.0 → pgvector_haystack-0.2.0}/src/haystack_integrations/components/retrievers/pgvector/__init__.py +2 -1
  7. {pgvector_haystack-0.1.0 → pgvector_haystack-0.2.0}/src/haystack_integrations/components/retrievers/pgvector/embedding_retriever.py +65 -16
  8. pgvector_haystack-0.2.0/src/haystack_integrations/components/retrievers/pgvector/keyword_retriever.py +123 -0
  9. {pgvector_haystack-0.1.0 → pgvector_haystack-0.2.0}/src/haystack_integrations/document_stores/pgvector/document_store.py +141 -27
  10. {pgvector_haystack-0.1.0 → pgvector_haystack-0.2.0}/src/haystack_integrations/document_stores/pgvector/filters.py +15 -13
  11. pgvector_haystack-0.2.0/tests/conftest.py +64 -0
  12. pgvector_haystack-0.2.0/tests/test_document_store.py +229 -0
  13. {pgvector_haystack-0.1.0 → pgvector_haystack-0.2.0}/tests/test_embedding_retrieval.py +1 -0
  14. pgvector_haystack-0.2.0/tests/test_filters.py +191 -0
  15. pgvector_haystack-0.2.0/tests/test_keyword_retrieval.py +50 -0
  16. pgvector_haystack-0.2.0/tests/test_retrievers.py +213 -0
  17. pgvector_haystack-0.1.0/tests/conftest.py +0 -26
  18. pgvector_haystack-0.1.0/tests/test_document_store.py +0 -219
  19. pgvector_haystack-0.1.0/tests/test_filters.py +0 -179
  20. pgvector_haystack-0.1.0/tests/test_retriever.py +0 -113
  21. {pgvector_haystack-0.1.0 → pgvector_haystack-0.2.0}/LICENSE.txt +0 -0
  22. {pgvector_haystack-0.1.0 → pgvector_haystack-0.2.0}/README.md +0 -0
  23. {pgvector_haystack-0.1.0 → pgvector_haystack-0.2.0}/src/haystack_integrations/document_stores/pgvector/__init__.py +0 -0
  24. {pgvector_haystack-0.1.0 → pgvector_haystack-0.2.0}/tests/__init__.py +0 -0
@@ -134,3 +134,4 @@ dmypy.json
134
134
 
135
135
  # Docs generation artifacts
136
136
  _readme_*.md
137
+ .idea
@@ -1,6 +1,6 @@
1
- Metadata-Version: 2.1
1
+ Metadata-Version: 2.3
2
2
  Name: pgvector-haystack
3
- Version: 0.1.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
@@ -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
- os.environ["PG_CONN_STR"] = "postgresql://postgres:postgres@localhost:5432/postgres"
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(
@@ -4,7 +4,6 @@ loaders:
4
4
  modules: [
5
5
  "haystack_integrations.components.retrievers.pgvector.embedding_retriever",
6
6
  "haystack_integrations.document_stores.pgvector.document_store",
7
- "haystack_integrations.document_stores.pgvector.filters",
8
7
  ]
9
8
  ignore_when_discovered: ["__init__"]
10
9
  processors:
@@ -16,14 +15,15 @@ processors:
16
15
  - type: smart
17
16
  - type: crossref
18
17
  renderer:
19
- type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer
18
+ type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer
20
19
  excerpt: Pgvector integration for Haystack
21
- category_slug: haystack-integrations
20
+ category_slug: integrations-api
22
21
  title: Pgvector
23
22
  slug: integrations-pgvector
24
- order: 140
23
+ order: 190
25
24
  markdown:
26
25
  descriptive_class_title: false
26
+ classdef_code_block: false
27
27
  descriptive_module_title: true
28
28
  add_method_class_prefix: true
29
29
  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 = ["py37"]
99
+ target-version = ["py38"]
99
100
  line-length = 120
100
101
  skip-string-normalization = true
101
102
 
102
103
  [tool.ruff]
103
- target-version = "py37"
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 = true
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, based on their dense embeddings.
15
+ Retrieves documents from the `PgvectorDocumentStore`, based on their dense embeddings.
16
16
 
17
- Needs to be connected to the PgvectorDocumentStore.
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
- Create the PgvectorEmbeddingRetriever component.
30
-
31
- :param document_store: An instance of PgvectorDocumentStore.
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, based on their embeddings.
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
- :type vector_function: Literal["cosine_similarity", "inner_product", "l2_distance"]
91
- :return: List of Documents similar to `query_embedding`.
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}