pgvector-haystack 0.0.1__tar.gz → 0.1.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 (18) hide show
  1. {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/PKG-INFO +13 -2
  2. {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/README.md +12 -1
  3. pgvector_haystack-0.1.0/examples/example.py +60 -0
  4. pgvector_haystack-0.1.0/pydoc/config.yml +30 -0
  5. {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/pyproject.toml +13 -6
  6. {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/src/haystack_integrations/components/retrievers/pgvector/embedding_retriever.py +2 -3
  7. {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/src/haystack_integrations/document_stores/pgvector/document_store.py +18 -10
  8. {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/tests/conftest.py +4 -2
  9. {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/tests/test_document_store.py +1 -4
  10. {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/tests/test_retriever.py +4 -3
  11. {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/.gitignore +0 -0
  12. {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/LICENSE.txt +0 -0
  13. {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/src/haystack_integrations/components/retrievers/pgvector/__init__.py +0 -0
  14. {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/src/haystack_integrations/document_stores/pgvector/__init__.py +0 -0
  15. {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/src/haystack_integrations/document_stores/pgvector/filters.py +0 -0
  16. {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/tests/__init__.py +0 -0
  17. {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/tests/test_embedding_retrieval.py +0 -0
  18. {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/tests/test_filters.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: pgvector-haystack
3
- Version: 0.0.1
3
+ Version: 0.1.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
@@ -45,12 +45,23 @@ pip install pgvector-haystack
45
45
 
46
46
  ## Testing
47
47
 
48
- TODO
48
+ Ensure that you have a PostgreSQL running with the `pgvector` extension. For a quick setup using Docker, run:
49
+ ```
50
+ docker run -d -p 5432:5432 -e POSTGRES_USER=postgres -e POSTGRES_PASSWORD=postgres -e POSTGRES_DB=postgres ankane/pgvector
51
+ ```
52
+
53
+ then run the tests:
49
54
 
50
55
  ```console
51
56
  hatch run test
52
57
  ```
53
58
 
59
+ To run the coverage report:
60
+
61
+ ```console
62
+ hatch run cov
63
+ ```
64
+
54
65
  ## License
55
66
 
56
67
  `pgvector-haystack` is distributed under the terms of the [Apache-2.0](https://spdx.org/licenses/Apache-2.0.html) license.
@@ -20,12 +20,23 @@ pip install pgvector-haystack
20
20
 
21
21
  ## Testing
22
22
 
23
- TODO
23
+ Ensure that you have a PostgreSQL running with the `pgvector` extension. For a quick setup using Docker, run:
24
+ ```
25
+ docker run -d -p 5432:5432 -e POSTGRES_USER=postgres -e POSTGRES_PASSWORD=postgres -e POSTGRES_DB=postgres ankane/pgvector
26
+ ```
27
+
28
+ then run the tests:
24
29
 
25
30
  ```console
26
31
  hatch run test
27
32
  ```
28
33
 
34
+ To run the coverage report:
35
+
36
+ ```console
37
+ hatch run cov
38
+ ```
39
+
29
40
  ## License
30
41
 
31
42
  `pgvector-haystack` is distributed under the terms of the [Apache-2.0](https://spdx.org/licenses/Apache-2.0.html) license.
@@ -0,0 +1,60 @@
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
+ import os
15
+
16
+ from haystack import Pipeline
17
+ from haystack.components.converters import MarkdownToDocument
18
+ from haystack.components.embedders import SentenceTransformersDocumentEmbedder, SentenceTransformersTextEmbedder
19
+ from haystack.components.preprocessors import DocumentSplitter
20
+ from haystack.components.writers import DocumentWriter
21
+ from haystack_integrations.components.retrievers.pgvector import PgvectorEmbeddingRetriever
22
+ from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
23
+
24
+ os.environ["PG_CONN_STR"] = "postgresql://postgres:postgres@localhost:5432/postgres"
25
+
26
+ # Initialize PgvectorDocumentStore
27
+ document_store = PgvectorDocumentStore(
28
+ table_name="haystack_test",
29
+ embedding_dimension=768,
30
+ vector_function="cosine_similarity",
31
+ recreate_table=True,
32
+ search_strategy="hnsw",
33
+ )
34
+
35
+ # Create the indexing Pipeline and index some documents
36
+ file_paths = glob.glob("neural-search-pills/pills/*.md")
37
+
38
+
39
+ indexing = Pipeline()
40
+ indexing.add_component("converter", MarkdownToDocument())
41
+ indexing.add_component("splitter", DocumentSplitter(split_by="sentence", split_length=2))
42
+ indexing.add_component("embedder", SentenceTransformersDocumentEmbedder())
43
+ indexing.add_component("writer", DocumentWriter(document_store))
44
+ indexing.connect("converter", "splitter")
45
+ indexing.connect("splitter", "embedder")
46
+ indexing.connect("embedder", "writer")
47
+
48
+ indexing.run({"converter": {"sources": file_paths}})
49
+
50
+ # Create the querying Pipeline and try a query
51
+ querying = Pipeline()
52
+ querying.add_component("embedder", SentenceTransformersTextEmbedder())
53
+ querying.add_component("retriever", PgvectorEmbeddingRetriever(document_store=document_store, top_k=3))
54
+ querying.connect("embedder", "retriever")
55
+
56
+ results = querying.run({"embedder": {"text": "What is a cross-encoder?"}})
57
+
58
+ for doc in results["retriever"]["documents"]:
59
+ print(doc)
60
+ print("-" * 10)
@@ -0,0 +1,30 @@
1
+ loaders:
2
+ - type: haystack_pydoc_tools.loaders.CustomPythonLoader
3
+ search_path: [../src]
4
+ modules: [
5
+ "haystack_integrations.components.retrievers.pgvector.embedding_retriever",
6
+ "haystack_integrations.document_stores.pgvector.document_store",
7
+ "haystack_integrations.document_stores.pgvector.filters",
8
+ ]
9
+ ignore_when_discovered: ["__init__"]
10
+ processors:
11
+ - type: filter
12
+ expression:
13
+ documented_only: true
14
+ do_not_filter_modules: false
15
+ skip_empty_modules: true
16
+ - type: smart
17
+ - type: crossref
18
+ renderer:
19
+ type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer
20
+ excerpt: Pgvector integration for Haystack
21
+ category_slug: haystack-integrations
22
+ title: Pgvector
23
+ slug: integrations-pgvector
24
+ order: 140
25
+ markdown:
26
+ descriptive_class_title: false
27
+ descriptive_module_title: true
28
+ add_method_class_prefix: true
29
+ add_member_class_prefix: false
30
+ filename: _readme_pgvector.md
@@ -51,6 +51,7 @@ dependencies = [
51
51
  "coverage[toml]>=6.5",
52
52
  "pytest",
53
53
  "ipython",
54
+ "haystack-pydoc-tools",
54
55
  ]
55
56
  [tool.hatch.envs.default.scripts]
56
57
  test = "pytest {args:tests}"
@@ -63,6 +64,9 @@ cov = [
63
64
  "test-cov",
64
65
  "cov-report",
65
66
  ]
67
+ docs = [
68
+ "pydoc-markdown pydoc/config.yml"
69
+ ]
66
70
 
67
71
  [[tool.hatch.envs.all.matrix]]
68
72
  python = ["3.8", "3.9", "3.10", "3.11", "3.12"]
@@ -134,6 +138,8 @@ ignore = [
134
138
  "S105", "S106", "S107",
135
139
  # Ignore complexity
136
140
  "C901", "PLR0911", "PLR0912", "PLR0913", "PLR0915",
141
+ # ignore function-call-in-default-argument
142
+ "B008",
137
143
  ]
138
144
  unfixable = [
139
145
  # Don't touch unused imports
@@ -149,24 +155,25 @@ ban-relative-imports = "parents"
149
155
  [tool.ruff.per-file-ignores]
150
156
  # Tests can use magic values, assertions, and relative imports
151
157
  "tests/**/*" = ["PLR2004", "S101", "TID252"]
158
+ # examples can contain "print" commands
159
+ "examples/**/*" = ["T201"]
160
+
152
161
 
153
162
  [tool.coverage.run]
154
- source_pkgs = ["src", "tests"]
163
+ source = ["haystack_integrations"]
155
164
  branch = true
156
165
  parallel = true
157
166
 
158
-
159
- [tool.coverage.paths]
160
- weaviate_haystack = ["src/haystack_integrations", "*/pgvector-haystack/src"]
161
- tests = ["tests", "*/pgvector-haystack/tests"]
162
-
163
167
  [tool.coverage.report]
168
+ omit = ["*/tests/*", "*/__init__.py"]
169
+ show_missing=true
164
170
  exclude_lines = [
165
171
  "no cov",
166
172
  "if __name__ == .__main__.:",
167
173
  "if TYPE_CHECKING:",
168
174
  ]
169
175
 
176
+
170
177
  [[tool.mypy.overrides]]
171
178
  module = [
172
179
  "haystack.*",
@@ -68,9 +68,8 @@ class PgvectorEmbeddingRetriever:
68
68
 
69
69
  @classmethod
70
70
  def from_dict(cls, data: Dict[str, Any]) -> "PgvectorEmbeddingRetriever":
71
- data["init_parameters"]["document_store"] = default_from_dict(
72
- PgvectorDocumentStore, data["init_parameters"]["document_store"]
73
- )
71
+ doc_store_params = data["init_parameters"]["document_store"]
72
+ data["init_parameters"]["document_store"] = PgvectorDocumentStore.from_dict(doc_store_params)
74
73
  return default_from_dict(cls, data)
75
74
 
76
75
  @component.output_types(documents=List[Document])
@@ -4,10 +4,11 @@
4
4
  import logging
5
5
  from typing import Any, Dict, List, Literal, Optional
6
6
 
7
- from haystack import default_to_dict
7
+ from haystack import default_from_dict, default_to_dict
8
8
  from haystack.dataclasses.document import ByteStream, Document
9
9
  from haystack.document_stores.errors import DocumentStoreError, DuplicateDocumentError
10
10
  from haystack.document_stores.types import DuplicatePolicy
11
+ from haystack.utils.auth import Secret, deserialize_secrets_inplace
11
12
  from haystack.utils.filters import convert
12
13
  from psycopg import Error, IntegrityError, connect
13
14
  from psycopg.abc import Query
@@ -69,7 +70,7 @@ class PgvectorDocumentStore:
69
70
  def __init__(
70
71
  self,
71
72
  *,
72
- connection_string: str,
73
+ connection_string: Secret = Secret.from_env_var("PG_CONN_STR"),
73
74
  table_name: str = "haystack_documents",
74
75
  embedding_dimension: int = 768,
75
76
  vector_function: Literal["cosine_similarity", "inner_product", "l2_distance"] = "cosine_similarity",
@@ -84,8 +85,8 @@ class PgvectorDocumentStore:
84
85
  It is meant to be connected to a PostgreSQL database with the pgvector extension installed.
85
86
  A specific table to store Haystack documents will be created if it doesn't exist yet.
86
87
 
87
- :param connection_string: The connection string to use to connect to the PostgreSQL database.
88
- e.g. "postgresql://USER:PASSWORD@HOST:PORT/DB_NAME"
88
+ :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"
89
90
  :param table_name: The name of the table to use to store Haystack documents. Defaults to "haystack_documents".
90
91
  :param embedding_dimension: The dimension of the embedding. Defaults to 768.
91
92
  :param vector_function: The similarity function to use when searching for similar embeddings.
@@ -130,7 +131,7 @@ class PgvectorDocumentStore:
130
131
  self.hnsw_index_creation_kwargs = hnsw_index_creation_kwargs or {}
131
132
  self.hnsw_ef_search = hnsw_ef_search
132
133
 
133
- connection = connect(connection_string)
134
+ connection = connect(self.connection_string.resolve_value())
134
135
  connection.autocommit = True
135
136
  self._connection = connection
136
137
 
@@ -151,7 +152,7 @@ class PgvectorDocumentStore:
151
152
  def to_dict(self) -> Dict[str, Any]:
152
153
  return default_to_dict(
153
154
  self,
154
- connection_string=self.connection_string,
155
+ connection_string=self.connection_string.to_dict(),
155
156
  table_name=self.table_name,
156
157
  embedding_dimension=self.embedding_dimension,
157
158
  vector_function=self.vector_function,
@@ -162,6 +163,11 @@ class PgvectorDocumentStore:
162
163
  hnsw_ef_search=self.hnsw_ef_search,
163
164
  )
164
165
 
166
+ @classmethod
167
+ def from_dict(cls, data: Dict[str, Any]) -> "PgvectorDocumentStore":
168
+ deserialize_secrets_inplace(data["init_parameters"], ["connection_string"])
169
+ return default_from_dict(cls, data)
170
+
165
171
  def _execute_sql(
166
172
  self, sql_query: Query, params: Optional[tuple] = None, error_msg: str = "", cursor: Optional[Cursor] = None
167
173
  ):
@@ -221,7 +227,7 @@ class PgvectorDocumentStore:
221
227
  )
222
228
  self._execute_sql(sql_set_hnsw_ef_search, error_msg="Could not set hnsw.ef_search")
223
229
 
224
- index_esists = bool(
230
+ index_exists = bool(
225
231
  self._execute_sql(
226
232
  "SELECT 1 FROM pg_indexes WHERE tablename = %s AND indexname = %s",
227
233
  (self.table_name, HNSW_INDEX_NAME),
@@ -229,7 +235,7 @@ class PgvectorDocumentStore:
229
235
  ).fetchone()
230
236
  )
231
237
 
232
- if index_esists and not self.hnsw_recreate_index_if_exists:
238
+ if index_exists and not self.hnsw_recreate_index_if_exists:
233
239
  logger.warning(
234
240
  "HNSW index already exists and won't be recreated. "
235
241
  "If you want to recreate it, pass 'hnsw_recreate_index_if_exists=True' to the "
@@ -373,7 +379,8 @@ class PgvectorDocumentStore:
373
379
 
374
380
  return written_docs
375
381
 
376
- def _from_haystack_to_pg_documents(self, documents: List[Document]) -> List[Dict[str, Any]]:
382
+ @staticmethod
383
+ def _from_haystack_to_pg_documents(documents: List[Document]) -> List[Dict[str, Any]]:
377
384
  """
378
385
  Internal method to convert a list of Haystack Documents to a list of dictionaries that can be used to insert
379
386
  documents into the PgvectorDocumentStore.
@@ -395,7 +402,8 @@ class PgvectorDocumentStore:
395
402
 
396
403
  return db_documents
397
404
 
398
- def _from_pg_to_haystack_documents(self, documents: List[Dict[str, Any]]) -> List[Document]:
405
+ @staticmethod
406
+ def _from_pg_to_haystack_documents(documents: List[Dict[str, Any]]) -> List[Document]:
399
407
  """
400
408
  Internal method to convert a list of dictionaries from pgvector to a list of Haystack Documents.
401
409
  """
@@ -1,10 +1,12 @@
1
+ import os
2
+
1
3
  import pytest
2
4
  from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
3
5
 
4
6
 
5
7
  @pytest.fixture
6
8
  def document_store(request):
7
- connection_string = "postgresql://postgres:postgres@localhost:5432/postgres"
9
+ os.environ["PG_CONN_STR"] = "postgresql://postgres:postgres@localhost:5432/postgres"
8
10
  table_name = f"haystack_{request.node.name}"
9
11
  embedding_dimension = 768
10
12
  vector_function = "cosine_similarity"
@@ -12,13 +14,13 @@ def document_store(request):
12
14
  search_strategy = "exact_nearest_neighbor"
13
15
 
14
16
  store = PgvectorDocumentStore(
15
- connection_string=connection_string,
16
17
  table_name=table_name,
17
18
  embedding_dimension=embedding_dimension,
18
19
  vector_function=vector_function,
19
20
  recreate_table=recreate_table,
20
21
  search_strategy=search_strategy,
21
22
  )
23
+
22
24
  yield store
23
25
 
24
26
  store.delete_table()
@@ -41,7 +41,6 @@ class TestDocumentStore(CountDocumentsTest, WriteDocumentsTest, DeleteDocumentsT
41
41
 
42
42
  def test_init(self):
43
43
  document_store = PgvectorDocumentStore(
44
- connection_string="postgresql://postgres:postgres@localhost:5432/postgres",
45
44
  table_name="my_table",
46
45
  embedding_dimension=512,
47
46
  vector_function="l2_distance",
@@ -52,7 +51,6 @@ class TestDocumentStore(CountDocumentsTest, WriteDocumentsTest, DeleteDocumentsT
52
51
  hnsw_ef_search=50,
53
52
  )
54
53
 
55
- assert document_store.connection_string == "postgresql://postgres:postgres@localhost:5432/postgres"
56
54
  assert document_store.table_name == "my_table"
57
55
  assert document_store.embedding_dimension == 512
58
56
  assert document_store.vector_function == "l2_distance"
@@ -64,7 +62,6 @@ class TestDocumentStore(CountDocumentsTest, WriteDocumentsTest, DeleteDocumentsT
64
62
 
65
63
  def test_to_dict(self):
66
64
  document_store = PgvectorDocumentStore(
67
- connection_string="postgresql://postgres:postgres@localhost:5432/postgres",
68
65
  table_name="my_table",
69
66
  embedding_dimension=512,
70
67
  vector_function="l2_distance",
@@ -78,7 +75,7 @@ class TestDocumentStore(CountDocumentsTest, WriteDocumentsTest, DeleteDocumentsT
78
75
  assert document_store.to_dict() == {
79
76
  "type": "haystack_integrations.document_stores.pgvector.document_store.PgvectorDocumentStore",
80
77
  "init_parameters": {
81
- "connection_string": "postgresql://postgres:postgres@localhost:5432/postgres",
78
+ "connection_string": {"env_vars": ["PG_CONN_STR"], "strict": True, "type": "env_var"},
82
79
  "table_name": "my_table",
83
80
  "embedding_dimension": 512,
84
81
  "vector_function": "l2_distance",
@@ -4,6 +4,7 @@
4
4
  from unittest.mock import Mock
5
5
 
6
6
  from haystack.dataclasses import Document
7
+ from haystack.utils.auth import EnvVarSecret
7
8
  from haystack_integrations.components.retrievers.pgvector import PgvectorEmbeddingRetriever
8
9
  from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
9
10
 
@@ -37,7 +38,7 @@ class TestRetriever:
37
38
  "document_store": {
38
39
  "type": "haystack_integrations.document_stores.pgvector.document_store.PgvectorDocumentStore",
39
40
  "init_parameters": {
40
- "connection_string": "postgresql://postgres:postgres@localhost:5432/postgres",
41
+ "connection_string": {"env_vars": ["PG_CONN_STR"], "strict": True, "type": "env_var"},
41
42
  "table_name": "haystack_test_to_dict",
42
43
  "embedding_dimension": 768,
43
44
  "vector_function": "cosine_similarity",
@@ -62,7 +63,7 @@ class TestRetriever:
62
63
  "document_store": {
63
64
  "type": "haystack_integrations.document_stores.pgvector.document_store.PgvectorDocumentStore",
64
65
  "init_parameters": {
65
- "connection_string": "postgresql://postgres:postgres@localhost:5432/postgres",
66
+ "connection_string": {"env_vars": ["PG_CONN_STR"], "strict": True, "type": "env_var"},
66
67
  "table_name": "haystack_test_to_dict",
67
68
  "embedding_dimension": 768,
68
69
  "vector_function": "cosine_similarity",
@@ -83,7 +84,7 @@ class TestRetriever:
83
84
  document_store = retriever.document_store
84
85
 
85
86
  assert isinstance(document_store, PgvectorDocumentStore)
86
- assert document_store.connection_string == "postgresql://postgres:postgres@localhost:5432/postgres"
87
+ assert isinstance(document_store.connection_string, EnvVarSecret)
87
88
  assert document_store.table_name == "haystack_test_to_dict"
88
89
  assert document_store.embedding_dimension == 768
89
90
  assert document_store.vector_function == "cosine_similarity"