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.
- {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/PKG-INFO +13 -2
- {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/README.md +12 -1
- pgvector_haystack-0.1.0/examples/example.py +60 -0
- pgvector_haystack-0.1.0/pydoc/config.yml +30 -0
- {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/pyproject.toml +13 -6
- {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/src/haystack_integrations/components/retrievers/pgvector/embedding_retriever.py +2 -3
- {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/src/haystack_integrations/document_stores/pgvector/document_store.py +18 -10
- {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/tests/conftest.py +4 -2
- {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/tests/test_document_store.py +1 -4
- {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/tests/test_retriever.py +4 -3
- {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/.gitignore +0 -0
- {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/LICENSE.txt +0 -0
- {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/src/haystack_integrations/components/retrievers/pgvector/__init__.py +0 -0
- {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/src/haystack_integrations/document_stores/pgvector/__init__.py +0 -0
- {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/src/haystack_integrations/document_stores/pgvector/filters.py +0 -0
- {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/tests/__init__.py +0 -0
- {pgvector_haystack-0.0.1 → pgvector_haystack-0.1.0}/tests/test_embedding_retrieval.py +0 -0
- {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
|
|
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
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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"]
|
|
72
|
-
|
|
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:
|
|
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
|
|
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
|
-
|
|
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
|
|
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
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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": "
|
|
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": "
|
|
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": "
|
|
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
|
|
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"
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|