pgvector-haystack 1.0.0__tar.gz → 1.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 (23) hide show
  1. pgvector_haystack-1.2.0/CHANGELOG.md +96 -0
  2. {pgvector_haystack-1.0.0 → pgvector_haystack-1.2.0}/PKG-INFO +2 -3
  3. {pgvector_haystack-1.0.0 → pgvector_haystack-1.2.0}/pyproject.toml +11 -7
  4. {pgvector_haystack-1.0.0 → pgvector_haystack-1.2.0}/src/haystack_integrations/document_stores/pgvector/document_store.py +104 -28
  5. {pgvector_haystack-1.0.0 → pgvector_haystack-1.2.0}/tests/test_document_store.py +26 -0
  6. {pgvector_haystack-1.0.0 → pgvector_haystack-1.2.0}/tests/test_retrievers.py +8 -0
  7. pgvector_haystack-1.0.0/CHANGELOG.md +0 -51
  8. {pgvector_haystack-1.0.0 → pgvector_haystack-1.2.0}/.gitignore +0 -0
  9. {pgvector_haystack-1.0.0 → pgvector_haystack-1.2.0}/LICENSE.txt +0 -0
  10. {pgvector_haystack-1.0.0 → pgvector_haystack-1.2.0}/README.md +0 -0
  11. {pgvector_haystack-1.0.0 → pgvector_haystack-1.2.0}/examples/embedding_retrieval.py +0 -0
  12. {pgvector_haystack-1.0.0 → pgvector_haystack-1.2.0}/examples/hybrid_retrieval.py +0 -0
  13. {pgvector_haystack-1.0.0 → pgvector_haystack-1.2.0}/pydoc/config.yml +0 -0
  14. {pgvector_haystack-1.0.0 → pgvector_haystack-1.2.0}/src/haystack_integrations/components/retrievers/pgvector/__init__.py +0 -0
  15. {pgvector_haystack-1.0.0 → pgvector_haystack-1.2.0}/src/haystack_integrations/components/retrievers/pgvector/embedding_retriever.py +0 -0
  16. {pgvector_haystack-1.0.0 → pgvector_haystack-1.2.0}/src/haystack_integrations/components/retrievers/pgvector/keyword_retriever.py +0 -0
  17. {pgvector_haystack-1.0.0 → pgvector_haystack-1.2.0}/src/haystack_integrations/document_stores/pgvector/__init__.py +0 -0
  18. {pgvector_haystack-1.0.0 → pgvector_haystack-1.2.0}/src/haystack_integrations/document_stores/pgvector/filters.py +0 -0
  19. {pgvector_haystack-1.0.0 → pgvector_haystack-1.2.0}/tests/__init__.py +0 -0
  20. {pgvector_haystack-1.0.0 → pgvector_haystack-1.2.0}/tests/conftest.py +0 -0
  21. {pgvector_haystack-1.0.0 → pgvector_haystack-1.2.0}/tests/test_embedding_retrieval.py +0 -0
  22. {pgvector_haystack-1.0.0 → pgvector_haystack-1.2.0}/tests/test_filters.py +0 -0
  23. {pgvector_haystack-1.0.0 → pgvector_haystack-1.2.0}/tests/test_keyword_retrieval.py +0 -0
@@ -0,0 +1,96 @@
1
+ # Changelog
2
+
3
+ ## [integrations/pgvector-v1.1.0] - 2024-11-21
4
+
5
+ ### 🚀 Features
6
+
7
+ - Add filter_policy to pgvector integration (#820)
8
+ - Add schema support to pgvector document store. (#1095)
9
+ - Pgvector - recreate the connection if it is no longer valid (#1202)
10
+
11
+ ### 🐛 Bug Fixes
12
+
13
+ - `PgVector` - Fallback to default filter policy when deserializing retrievers without the init parameter (#900)
14
+
15
+ ### 📚 Documentation
16
+
17
+ - Explain different connection string formats in the docstring (#1132)
18
+
19
+ ### 🧪 Testing
20
+
21
+ - Do not retry tests in `hatch run test` command (#954)
22
+
23
+ ### ⚙️ CI
24
+
25
+ - Retry tests to reduce flakyness (#836)
26
+ - Adopt uv as installer (#1142)
27
+
28
+ ### 🧹 Chores
29
+
30
+ - Update ruff invocation to include check parameter (#853)
31
+ - PgVector - remove legacy filter support (#1068)
32
+ - Update changelog after removing legacy filters (#1083)
33
+ - Update ruff linting scripts and settings (#1105)
34
+
35
+ ### 🌀 Miscellaneous
36
+
37
+ - Ci: install `pytest-rerunfailures` where needed; add retry config to `test-cov` script (#845)
38
+ - Chore: Minor retriever pydoc fix (#884)
39
+ - Chore: Update pgvector test for the new `apply_filter_policy` usage (#970)
40
+ - Chore: pgvector ruff update, don't ruff tests (#984)
41
+
42
+ ## [integrations/pgvector-v0.4.0] - 2024-06-20
43
+
44
+ ### 🚀 Features
45
+
46
+ - Defer the database connection to when it's needed (#773)
47
+ - Add customizable index names for pgvector (#818)
48
+
49
+ ### 🌀 Miscellaneous
50
+
51
+ - Docs: add missing api references (#728)
52
+ - [deepset-ai/haystack-core-integrations#727] (#738)
53
+
54
+ ## [integrations/pgvector-v0.2.0] - 2024-05-08
55
+
56
+ ### 🚀 Features
57
+
58
+ - `MongoDBAtlasEmbeddingRetriever` (#427)
59
+ - Implement keyword retrieval for pgvector integration (#644)
60
+
61
+ ### 🐛 Bug Fixes
62
+
63
+ - Fix order of API docs (#447)
64
+
65
+ ### 📚 Documentation
66
+
67
+ - Update category slug (#442)
68
+ - Disable-class-def (#556)
69
+
70
+ ### 🌀 Miscellaneous
71
+
72
+ - Pgvector - review docstrings and API reference (#502)
73
+ - Refactor tests (#574)
74
+ - Remove references to Python 3.7 (#601)
75
+ - Make Document Stores initially skip `SparseEmbedding` (#606)
76
+ - Chore: add license classifiers (#680)
77
+ - Type hints in pgvector document store updated for 3.8 compability (#704)
78
+ - Chore: change the pydoc renderer class (#718)
79
+
80
+ ## [integrations/pgvector-v0.1.0] - 2024-02-14
81
+
82
+ ### 🐛 Bug Fixes
83
+
84
+ - Pgvector: fix linting (#328)
85
+
86
+ ### 🌀 Miscellaneous
87
+
88
+ - Pgvector Document Store - minimal implementation (#239)
89
+ - Pgvector - filters (#257)
90
+ - Pgvector - embedding retrieval (#298)
91
+ - Pgvector - Embedding Retriever (#320)
92
+ - Pgvector: generate API docs (#325)
93
+ - Pgvector: add an example (#334)
94
+ - Adopt `Secret` to pgvector (#402)
95
+
96
+ <!-- generated by git-cliff -->
@@ -1,13 +1,12 @@
1
1
  Metadata-Version: 2.3
2
2
  Name: pgvector-haystack
3
- Version: 1.0.0
3
+ Version: 1.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
7
7
  Project-URL: Issues, https://github.com/deepset-ai/haystack-core-integrations/issues
8
8
  Author-email: deepset GmbH <info@deepset.ai>
9
- License-Expression: Apache-2.0
10
- License-File: LICENSE.txt
9
+ License: Apache-2.0
11
10
  Classifier: Development Status :: 4 - Beta
12
11
  Classifier: License :: OSI Approved :: Apache Software License
13
12
  Classifier: Programming Language :: Python
@@ -42,6 +42,7 @@ root = "../.."
42
42
  git_describe_command = 'git describe --tags --match="integrations/pgvector-v[0-9]*"'
43
43
 
44
44
  [tool.hatch.envs.default]
45
+ installer = "uv"
45
46
  dependencies = [
46
47
  "coverage[toml]>=6.5",
47
48
  "pytest",
@@ -62,12 +63,13 @@ docs = ["pydoc-markdown pydoc/config.yml"]
62
63
  python = ["3.8", "3.9", "3.10", "3.11", "3.12"]
63
64
 
64
65
  [tool.hatch.envs.lint]
66
+ installer = "uv"
65
67
  detached = true
66
- dependencies = ["black>=23.1.0", "mypy>=1.0.0", "ruff>=0.0.243"]
68
+ dependencies = ["pip", "black>=23.1.0", "mypy>=1.0.0", "ruff>=0.0.243"]
67
69
  [tool.hatch.envs.lint.scripts]
68
70
  typing = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}"
69
- style = ["ruff check {args:. --exclude tests/, examples/}", "black --check --diff {args:.}"]
70
- fmt = ["black {args:.}", "ruff --fix {args:. --exclude tests/, examples/}", "style"]
71
+ style = ["ruff check {args:.}", "black --check --diff {args:.}"]
72
+ fmt = ["black {args:.}", "ruff check --fix {args:.}", "style"]
71
73
  all = ["style", "typing"]
72
74
 
73
75
  [tool.black]
@@ -78,6 +80,8 @@ skip-string-normalization = true
78
80
  [tool.ruff]
79
81
  target-version = "py38"
80
82
  line-length = 120
83
+
84
+ [tool.ruff.lint]
81
85
  select = [
82
86
  "A",
83
87
  "ARG",
@@ -128,13 +132,13 @@ unfixable = [
128
132
  "F401",
129
133
  ]
130
134
 
131
- [tool.ruff.isort]
132
- known-first-party = ["src"]
135
+ [tool.ruff.lint.isort]
136
+ known-first-party = ["haystack_integrations"]
133
137
 
134
- [tool.ruff.flake8-tidy-imports]
138
+ [tool.ruff.lint.flake8-tidy-imports]
135
139
  ban-relative-imports = "parents"
136
140
 
137
- [tool.ruff.per-file-ignores]
141
+ [tool.ruff.lint.per-file-ignores]
138
142
  # Tests can use magic values, assertions, and relative imports
139
143
  "tests/**/*" = ["PLR2004", "S101", "TID252"]
140
144
  # examples can contain "print" commands
@@ -24,7 +24,7 @@ from .filters import _convert_filters_to_where_clause_and_params
24
24
  logger = logging.getLogger(__name__)
25
25
 
26
26
  CREATE_TABLE_STATEMENT = """
27
- CREATE TABLE IF NOT EXISTS {table_name} (
27
+ CREATE TABLE IF NOT EXISTS {schema_name}.{table_name} (
28
28
  id VARCHAR(128) PRIMARY KEY,
29
29
  embedding VECTOR({embedding_dimension}),
30
30
  content TEXT,
@@ -36,7 +36,7 @@ meta JSONB)
36
36
  """
37
37
 
38
38
  INSERT_STATEMENT = """
39
- INSERT INTO {table_name}
39
+ INSERT INTO {schema_name}.{table_name}
40
40
  (id, embedding, content, dataframe, blob_data, blob_meta, blob_mime_type, meta)
41
41
  VALUES (%(id)s, %(embedding)s, %(content)s, %(dataframe)s, %(blob_data)s, %(blob_meta)s, %(blob_mime_type)s, %(meta)s)
42
42
  """
@@ -54,7 +54,7 @@ meta = EXCLUDED.meta
54
54
 
55
55
  KEYWORD_QUERY = """
56
56
  SELECT {table_name}.*, ts_rank_cd(to_tsvector({language}, content), query) AS score
57
- FROM {table_name}, plainto_tsquery({language}, %s) query
57
+ FROM {schema_name}.{table_name}, plainto_tsquery({language}, %s) query
58
58
  WHERE to_tsvector({language}, content) @@ query
59
59
  """
60
60
 
@@ -78,6 +78,8 @@ class PgvectorDocumentStore:
78
78
  self,
79
79
  *,
80
80
  connection_string: Secret = Secret.from_env_var("PG_CONN_STR"),
81
+ create_extension: bool = True,
82
+ schema_name: str = "public",
81
83
  table_name: str = "haystack_documents",
82
84
  language: str = "english",
83
85
  embedding_dimension: int = 768,
@@ -96,7 +98,16 @@ class PgvectorDocumentStore:
96
98
  A specific table to store Haystack documents will be created if it doesn't exist yet.
97
99
 
98
100
  :param connection_string: The connection string to use to connect to the PostgreSQL database, defined as an
99
- environment variable, e.g.: `PG_CONN_STR="postgresql://USER:PASSWORD@HOST:PORT/DB_NAME"`
101
+ environment variable. It can be provided in either URI format
102
+ e.g.: `PG_CONN_STR="postgresql://USER:PASSWORD@HOST:PORT/DB_NAME"`, or keyword/value format
103
+ e.g.: `PG_CONN_STR="host=HOST port=PORT dbname=DBNAME user=USER password=PASSWORD"`
104
+ See [PostgreSQL Documentation](https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNSTRING)
105
+ for more details.
106
+ :param create_extension: Whether to create the pgvector extension if it doesn't exist.
107
+ Set this to `True` (default) to automatically create the extension if it is missing.
108
+ Creating the extension may require superuser privileges.
109
+ If set to `False`, ensure the extension is already installed; otherwise, an error will be raised.
110
+ :param schema_name: The name of the schema the table is created in. The schema must already exist.
100
111
  :param table_name: The name of the table to use to store Haystack documents.
101
112
  :param language: The language to be used to parse query and document content in keyword retrieval.
102
113
  To see the list of available languages, you can run the following SQL query in your PostgreSQL database:
@@ -132,7 +143,9 @@ class PgvectorDocumentStore:
132
143
  """
133
144
 
134
145
  self.connection_string = connection_string
146
+ self.create_extension = create_extension
135
147
  self.table_name = table_name
148
+ self.schema_name = schema_name
136
149
  self.embedding_dimension = embedding_dimension
137
150
  if vector_function not in VALID_VECTOR_FUNCTIONS:
138
151
  msg = f"vector_function must be one of {VALID_VECTOR_FUNCTIONS}, but got {vector_function}"
@@ -149,49 +162,86 @@ class PgvectorDocumentStore:
149
162
  self._connection = None
150
163
  self._cursor = None
151
164
  self._dict_cursor = None
165
+ self._table_initialized = False
152
166
 
153
167
  @property
154
168
  def cursor(self):
155
- if self._cursor is None:
169
+ if self._cursor is None or not self._connection_is_valid(self._connection):
156
170
  self._create_connection()
157
171
 
158
172
  return self._cursor
159
173
 
160
174
  @property
161
175
  def dict_cursor(self):
162
- if self._dict_cursor is None:
176
+ if self._dict_cursor is None or not self._connection_is_valid(self._connection):
163
177
  self._create_connection()
164
178
 
165
179
  return self._dict_cursor
166
180
 
167
181
  @property
168
182
  def connection(self):
169
- if self._connection is None:
183
+ if self._connection is None or not self._connection_is_valid(self._connection):
170
184
  self._create_connection()
171
185
 
172
186
  return self._connection
173
187
 
174
188
  def _create_connection(self):
189
+ """
190
+ Internal method to create a connection to the PostgreSQL database.
191
+ """
192
+
193
+ # close the connection if it already exists
194
+ if self._connection:
195
+ try:
196
+ self._connection.close()
197
+ except Error as e:
198
+ logger.debug("Failed to close connection: %s", str(e))
199
+
175
200
  conn_str = self.connection_string.resolve_value() or ""
176
201
  connection = connect(conn_str)
177
202
  connection.autocommit = True
178
- connection.execute("CREATE EXTENSION IF NOT EXISTS vector")
203
+ if self.create_extension:
204
+ connection.execute("CREATE EXTENSION IF NOT EXISTS vector")
179
205
  register_vector(connection) # Note: this must be called before creating the cursors.
180
206
 
181
207
  self._connection = connection
182
208
  self._cursor = self._connection.cursor()
183
209
  self._dict_cursor = self._connection.cursor(row_factory=dict_row)
184
210
 
185
- # Init schema
211
+ if not self._table_initialized:
212
+ self._initialize_table()
213
+
214
+ return self._connection
215
+
216
+ def _initialize_table(self):
217
+ """
218
+ Internal method to initialize the table.
219
+ """
186
220
  if self.recreate_table:
187
221
  self.delete_table()
222
+
188
223
  self._create_table_if_not_exists()
189
224
  self._create_keyword_index_if_not_exists()
190
225
 
191
226
  if self.search_strategy == "hnsw":
192
227
  self._handle_hnsw()
193
228
 
194
- return self._connection
229
+ self._table_initialized = True
230
+
231
+ @staticmethod
232
+ def _connection_is_valid(connection):
233
+ """
234
+ Internal method to check if the connection is still valid.
235
+ """
236
+
237
+ # implementation inspired to psycopg pool
238
+ # https://github.com/psycopg/psycopg/blob/d38cf7798b0c602ff43dac9f20bbab96237a9c38/psycopg_pool/psycopg_pool/pool.py#L528
239
+
240
+ try:
241
+ connection.execute("")
242
+ except Error:
243
+ return False
244
+ return True
195
245
 
196
246
  def to_dict(self) -> Dict[str, Any]:
197
247
  """
@@ -203,6 +253,8 @@ class PgvectorDocumentStore:
203
253
  return default_to_dict(
204
254
  self,
205
255
  connection_string=self.connection_string.to_dict(),
256
+ create_extension=self.create_extension,
257
+ schema_name=self.schema_name,
206
258
  table_name=self.table_name,
207
259
  embedding_dimension=self.embedding_dimension,
208
260
  vector_function=self.vector_function,
@@ -262,7 +314,9 @@ class PgvectorDocumentStore:
262
314
  """
263
315
 
264
316
  create_sql = SQL(CREATE_TABLE_STATEMENT).format(
265
- table_name=Identifier(self.table_name), embedding_dimension=SQLLiteral(self.embedding_dimension)
317
+ schema_name=Identifier(self.schema_name),
318
+ table_name=Identifier(self.table_name),
319
+ embedding_dimension=SQLLiteral(self.embedding_dimension),
266
320
  )
267
321
 
268
322
  self._execute_sql(create_sql, error_msg="Could not create table in PgvectorDocumentStore")
@@ -270,12 +324,18 @@ class PgvectorDocumentStore:
270
324
  def delete_table(self):
271
325
  """
272
326
  Deletes the table used to store Haystack documents.
273
- The name of the table (`table_name`) is defined when initializing the `PgvectorDocumentStore`.
327
+ The name of the schema (`schema_name`) and the name of the table (`table_name`)
328
+ are defined when initializing the `PgvectorDocumentStore`.
274
329
  """
330
+ delete_sql = SQL("DROP TABLE IF EXISTS {schema_name}.{table_name}").format(
331
+ schema_name=Identifier(self.schema_name),
332
+ table_name=Identifier(self.table_name),
333
+ )
275
334
 
276
- delete_sql = SQL("DROP TABLE IF EXISTS {table_name}").format(table_name=Identifier(self.table_name))
277
-
278
- self._execute_sql(delete_sql, error_msg=f"Could not delete table {self.table_name} in PgvectorDocumentStore")
335
+ self._execute_sql(
336
+ delete_sql,
337
+ error_msg=f"Could not delete table {self.schema_name}.{self.table_name} in PgvectorDocumentStore",
338
+ )
279
339
 
280
340
  def _create_keyword_index_if_not_exists(self):
281
341
  """
@@ -283,15 +343,16 @@ class PgvectorDocumentStore:
283
343
  """
284
344
  index_exists = bool(
285
345
  self._execute_sql(
286
- "SELECT 1 FROM pg_indexes WHERE tablename = %s AND indexname = %s",
287
- (self.table_name, self.keyword_index_name),
346
+ "SELECT 1 FROM pg_indexes WHERE schemaname = %s AND tablename = %s AND indexname = %s",
347
+ (self.schema_name, self.table_name, self.keyword_index_name),
288
348
  "Could not check if keyword index exists",
289
349
  ).fetchone()
290
350
  )
291
351
 
292
352
  sql_create_index = SQL(
293
- "CREATE INDEX {index_name} ON {table_name} USING GIN (to_tsvector({language}, content))"
353
+ "CREATE INDEX {index_name} ON {schema_name}.{table_name} USING GIN (to_tsvector({language}, content))"
294
354
  ).format(
355
+ schema_name=Identifier(self.schema_name),
295
356
  index_name=Identifier(self.keyword_index_name),
296
357
  table_name=Identifier(self.table_name),
297
358
  language=SQLLiteral(self.language),
@@ -314,8 +375,8 @@ class PgvectorDocumentStore:
314
375
 
315
376
  index_exists = bool(
316
377
  self._execute_sql(
317
- "SELECT 1 FROM pg_indexes WHERE tablename = %s AND indexname = %s",
318
- (self.table_name, self.hnsw_index_name),
378
+ "SELECT 1 FROM pg_indexes WHERE schemaname = %s AND tablename = %s AND indexname = %s",
379
+ (self.schema_name, self.table_name, self.hnsw_index_name),
319
380
  "Could not check if HNSW index exists",
320
381
  ).fetchone()
321
382
  )
@@ -345,8 +406,13 @@ class PgvectorDocumentStore:
345
406
  if key in HNSW_INDEX_CREATION_VALID_KWARGS
346
407
  }
347
408
 
348
- sql_create_index = SQL("CREATE INDEX {index_name} ON {table_name} USING hnsw (embedding {ops}) ").format(
349
- index_name=Identifier(self.hnsw_index_name), table_name=Identifier(self.table_name), ops=SQL(pg_ops)
409
+ sql_create_index = SQL(
410
+ "CREATE INDEX {index_name} ON {schema_name}.{table_name} USING hnsw (embedding {ops}) "
411
+ ).format(
412
+ schema_name=Identifier(self.schema_name),
413
+ index_name=Identifier(self.hnsw_index_name),
414
+ table_name=Identifier(self.table_name),
415
+ ops=SQL(pg_ops),
350
416
  )
351
417
 
352
418
  if actual_hnsw_index_creation_kwargs:
@@ -365,7 +431,9 @@ class PgvectorDocumentStore:
365
431
  Returns how many documents are present in the document store.
366
432
  """
367
433
 
368
- sql_count = SQL("SELECT COUNT(*) FROM {table_name}").format(table_name=Identifier(self.table_name))
434
+ sql_count = SQL("SELECT COUNT(*) FROM {schema_name}.{table_name}").format(
435
+ schema_name=Identifier(self.schema_name), table_name=Identifier(self.table_name)
436
+ )
369
437
 
370
438
  count = self._execute_sql(sql_count, error_msg="Could not count documents in PgvectorDocumentStore").fetchone()[
371
439
  0
@@ -391,7 +459,9 @@ class PgvectorDocumentStore:
391
459
  msg = "Invalid filter syntax. See https://docs.haystack.deepset.ai/docs/metadata-filtering for details."
392
460
  raise ValueError(msg)
393
461
 
394
- sql_filter = SQL("SELECT * FROM {table_name}").format(table_name=Identifier(self.table_name))
462
+ sql_filter = SQL("SELECT * FROM {schema_name}.{table_name}").format(
463
+ schema_name=Identifier(self.schema_name), table_name=Identifier(self.table_name)
464
+ )
395
465
 
396
466
  params = ()
397
467
  if filters:
@@ -430,7 +500,9 @@ class PgvectorDocumentStore:
430
500
 
431
501
  db_documents = self._from_haystack_to_pg_documents(documents)
432
502
 
433
- sql_insert = SQL(INSERT_STATEMENT).format(table_name=Identifier(self.table_name))
503
+ sql_insert = SQL(INSERT_STATEMENT).format(
504
+ schema_name=Identifier(self.schema_name), table_name=Identifier(self.table_name)
505
+ )
434
506
 
435
507
  if policy == DuplicatePolicy.OVERWRITE:
436
508
  sql_insert += SQL(UPDATE_STATEMENT)
@@ -539,8 +611,10 @@ class PgvectorDocumentStore:
539
611
 
540
612
  document_ids_str = ", ".join(f"'{document_id}'" for document_id in document_ids)
541
613
 
542
- delete_sql = SQL("DELETE FROM {table_name} WHERE id IN ({document_ids_str})").format(
543
- table_name=Identifier(self.table_name), document_ids_str=SQL(document_ids_str)
614
+ delete_sql = SQL("DELETE FROM {schema_name}.{table_name} WHERE id IN ({document_ids_str})").format(
615
+ schema_name=Identifier(self.schema_name),
616
+ table_name=Identifier(self.table_name),
617
+ document_ids_str=SQL(document_ids_str),
544
618
  )
545
619
 
546
620
  self._execute_sql(delete_sql, error_msg="Could not delete documents from PgvectorDocumentStore")
@@ -566,6 +640,7 @@ class PgvectorDocumentStore:
566
640
  raise ValueError(msg)
567
641
 
568
642
  sql_select = SQL(KEYWORD_QUERY).format(
643
+ schema_name=Identifier(self.schema_name),
569
644
  table_name=Identifier(self.table_name),
570
645
  language=SQLLiteral(self.language),
571
646
  query=SQLLiteral(query),
@@ -639,7 +714,8 @@ class PgvectorDocumentStore:
639
714
  elif vector_function == "l2_distance":
640
715
  score_definition = f"embedding <-> {query_embedding_for_postgres} AS score"
641
716
 
642
- sql_select = SQL("SELECT *, {score} FROM {table_name}").format(
717
+ sql_select = SQL("SELECT *, {score} FROM {schema_name}.{table_name}").format(
718
+ schema_name=Identifier(self.schema_name),
643
719
  table_name=Identifier(self.table_name),
644
720
  score=SQL(score_definition),
645
721
  )
@@ -41,12 +41,33 @@ class TestDocumentStore(CountDocumentsTest, WriteDocumentsTest, DeleteDocumentsT
41
41
  retrieved_docs = document_store.filter_documents()
42
42
  assert retrieved_docs == docs
43
43
 
44
+ def test_connection_check_and_recreation(self, document_store: PgvectorDocumentStore):
45
+ original_connection = document_store.connection
46
+
47
+ with patch.object(PgvectorDocumentStore, "_connection_is_valid", return_value=False):
48
+ new_connection = document_store.connection
49
+
50
+ # verify that a new connection is created
51
+ assert new_connection is not original_connection
52
+ assert document_store._connection == new_connection
53
+ assert original_connection.closed
54
+
55
+ assert document_store._cursor is not None
56
+ assert document_store._dict_cursor is not None
57
+
58
+ # test with new connection
59
+ with patch.object(PgvectorDocumentStore, "_connection_is_valid", return_value=True):
60
+ same_connection = document_store.connection
61
+ assert same_connection is document_store._connection
62
+
44
63
 
45
64
  @pytest.mark.usefixtures("patches_for_unit_tests")
46
65
  def test_init(monkeypatch):
47
66
  monkeypatch.setenv("PG_CONN_STR", "some_connection_string")
48
67
 
49
68
  document_store = PgvectorDocumentStore(
69
+ create_extension=True,
70
+ schema_name="my_schema",
50
71
  table_name="my_table",
51
72
  embedding_dimension=512,
52
73
  vector_function="l2_distance",
@@ -59,6 +80,8 @@ def test_init(monkeypatch):
59
80
  keyword_index_name="my_keyword_index",
60
81
  )
61
82
 
83
+ assert document_store.create_extension
84
+ assert document_store.schema_name == "my_schema"
62
85
  assert document_store.table_name == "my_table"
63
86
  assert document_store.embedding_dimension == 512
64
87
  assert document_store.vector_function == "l2_distance"
@@ -76,6 +99,7 @@ def test_to_dict(monkeypatch):
76
99
  monkeypatch.setenv("PG_CONN_STR", "some_connection_string")
77
100
 
78
101
  document_store = PgvectorDocumentStore(
102
+ create_extension=False,
79
103
  table_name="my_table",
80
104
  embedding_dimension=512,
81
105
  vector_function="l2_distance",
@@ -92,7 +116,9 @@ def test_to_dict(monkeypatch):
92
116
  "type": "haystack_integrations.document_stores.pgvector.document_store.PgvectorDocumentStore",
93
117
  "init_parameters": {
94
118
  "connection_string": {"env_vars": ["PG_CONN_STR"], "strict": True, "type": "env_var"},
119
+ "create_extension": False,
95
120
  "table_name": "my_table",
121
+ "schema_name": "public",
96
122
  "embedding_dimension": 512,
97
123
  "vector_function": "l2_distance",
98
124
  "recreate_table": True,
@@ -50,6 +50,8 @@ class TestEmbeddingRetriever:
50
50
  "type": "haystack_integrations.document_stores.pgvector.document_store.PgvectorDocumentStore",
51
51
  "init_parameters": {
52
52
  "connection_string": {"env_vars": ["PG_CONN_STR"], "strict": True, "type": "env_var"},
53
+ "create_extension": True,
54
+ "schema_name": "public",
53
55
  "table_name": "haystack",
54
56
  "embedding_dimension": 768,
55
57
  "vector_function": "cosine_similarity",
@@ -81,6 +83,7 @@ class TestEmbeddingRetriever:
81
83
  "type": "haystack_integrations.document_stores.pgvector.document_store.PgvectorDocumentStore",
82
84
  "init_parameters": {
83
85
  "connection_string": {"env_vars": ["PG_CONN_STR"], "strict": True, "type": "env_var"},
86
+ "create_extension": False,
84
87
  "table_name": "haystack_test_to_dict",
85
88
  "embedding_dimension": 768,
86
89
  "vector_function": "cosine_similarity",
@@ -105,6 +108,7 @@ class TestEmbeddingRetriever:
105
108
 
106
109
  assert isinstance(document_store, PgvectorDocumentStore)
107
110
  assert isinstance(document_store.connection_string, EnvVarSecret)
111
+ assert not document_store.create_extension
108
112
  assert document_store.table_name == "haystack_test_to_dict"
109
113
  assert document_store.embedding_dimension == 768
110
114
  assert document_store.vector_function == "cosine_similarity"
@@ -175,6 +179,8 @@ class TestKeywordRetriever:
175
179
  "type": "haystack_integrations.document_stores.pgvector.document_store.PgvectorDocumentStore",
176
180
  "init_parameters": {
177
181
  "connection_string": {"env_vars": ["PG_CONN_STR"], "strict": True, "type": "env_var"},
182
+ "create_extension": True,
183
+ "schema_name": "public",
178
184
  "table_name": "haystack",
179
185
  "embedding_dimension": 768,
180
186
  "vector_function": "cosine_similarity",
@@ -205,6 +211,7 @@ class TestKeywordRetriever:
205
211
  "type": "haystack_integrations.document_stores.pgvector.document_store.PgvectorDocumentStore",
206
212
  "init_parameters": {
207
213
  "connection_string": {"env_vars": ["PG_CONN_STR"], "strict": True, "type": "env_var"},
214
+ "create_extension": False,
208
215
  "table_name": "haystack_test_to_dict",
209
216
  "embedding_dimension": 768,
210
217
  "vector_function": "cosine_similarity",
@@ -228,6 +235,7 @@ class TestKeywordRetriever:
228
235
 
229
236
  assert isinstance(document_store, PgvectorDocumentStore)
230
237
  assert isinstance(document_store.connection_string, EnvVarSecret)
238
+ assert not document_store.create_extension
231
239
  assert document_store.table_name == "haystack_test_to_dict"
232
240
  assert document_store.embedding_dimension == 768
233
241
  assert document_store.vector_function == "cosine_similarity"
@@ -1,51 +0,0 @@
1
- # Changelog
2
-
3
- ## [unreleased]
4
-
5
- ### 🚀 Features
6
-
7
- - Add filter_policy to pgvector integration (#820)
8
-
9
- ### 🐛 Bug Fixes
10
-
11
- - `PgVector` - Fallback to default filter policy when deserializing retrievers without the init parameter (#900)
12
-
13
- ### ⚙️ Miscellaneous Tasks
14
-
15
- - Retry tests to reduce flakyness (#836)
16
- - Update ruff invocation to include check parameter (#853)
17
-
18
- ## [integrations/pgvector-v0.4.0] - 2024-06-20
19
-
20
- ### 🚀 Features
21
-
22
- - Defer the database connection to when it's needed (#773)
23
- - Add customizable index names for pgvector (#818)
24
-
25
- ## [integrations/pgvector-v0.2.0] - 2024-05-08
26
-
27
- ### 🚀 Features
28
-
29
- - `MongoDBAtlasEmbeddingRetriever` (#427)
30
- - Implement keyword retrieval for pgvector integration (#644)
31
-
32
- ### 🐛 Bug Fixes
33
-
34
- - Fix order of API docs (#447)
35
-
36
- This PR will also push the docs to Readme
37
-
38
- ### 📚 Documentation
39
-
40
- - Update category slug (#442)
41
- - Disable-class-def (#556)
42
-
43
- ## [integrations/pgvector-v0.1.0] - 2024-02-14
44
-
45
- ### 🐛 Bug Fixes
46
-
47
- - Fix linting (#328)
48
-
49
-
50
-
51
- <!-- generated by git-cliff -->