pgvector-haystack 0.5.1__tar.gz → 1.1.0__tar.gz

{pgvector_haystack-0.5.1 → pgvector_haystack-1.1.0}/CHANGELOG.md

@@ -1,5 +1,25 @@
 # Changelog
 
+## [integrations/pgvector-v1.0.0] - 2024-09-12
+
+### 🚀 Features
+
+- Add filter_policy to pgvector integration (#820)
+
+### 🐛 Bug Fixes
+
+- `PgVector` - Fallback to default filter policy when deserializing retrievers without the init parameter (#900)
+
+### 🧪 Testing
+
+- Do not retry tests in `hatch run test` command (#954)
+
+### ⚙️ Miscellaneous Tasks
+
+- Retry tests to reduce flakyness (#836)
+- Update ruff invocation to include check parameter (#853)
+- PgVector - remove legacy filter support (#1068)
+
 ## [integrations/pgvector-v0.4.0] - 2024-06-20
 
 ### 🚀 Features

{pgvector_haystack-0.5.1 → pgvector_haystack-1.1.0}/PKG-INFO

@@ -1,13 +1,12 @@
 Metadata-Version: 2.3
 Name: pgvector-haystack
-Version: 0.5.1
+Version: 1.1.0
 Summary: An integration of pgvector (vector search extension for Postgres) with Haystack
 Project-URL: Source, https://github.com/deepset-ai/haystack-core-integrations
 Project-URL: Documentation, https://github.com/deepset-ai/haystack-core-integrations/blob/main/integrations/pgvector/README.md
 Project-URL: Issues, https://github.com/deepset-ai/haystack-core-integrations/issues
 Author-email: deepset GmbH <info@deepset.ai>
-License-Expression: Apache-2.0
-License-File: LICENSE.txt
+License: Apache-2.0
 Classifier: Development Status :: 4 - Beta
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python

{pgvector_haystack-0.5.1 → pgvector_haystack-1.1.0}/examples/embedding_retrieval.py

@@ -17,6 +17,7 @@ from haystack.components.converters import MarkdownToDocument
 from haystack.components.embedders import SentenceTransformersDocumentEmbedder, SentenceTransformersTextEmbedder
 from haystack.components.preprocessors import DocumentSplitter
 from haystack.components.writers import DocumentWriter
+
 from haystack_integrations.components.retrievers.pgvector import PgvectorEmbeddingRetriever
 from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
 

{pgvector_haystack-0.5.1 → pgvector_haystack-1.1.0}/examples/hybrid_retrieval.py

@@ -18,6 +18,7 @@ from haystack.components.embedders import SentenceTransformersDocumentEmbedder,
 from haystack.components.joiners import DocumentJoiner
 from haystack.components.preprocessors import DocumentSplitter
 from haystack.components.writers import DocumentWriter
+
 from haystack_integrations.components.retrievers.pgvector import PgvectorEmbeddingRetriever, PgvectorKeywordRetriever
 from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
 

{pgvector_haystack-0.5.1 → pgvector_haystack-1.1.0}/pyproject.toml

@@ -42,6 +42,7 @@ root = "../.."
 git_describe_command = 'git describe --tags --match="integrations/pgvector-v[0-9]*"'
 
 [tool.hatch.envs.default]
+installer = "uv"
 dependencies = [
   "coverage[toml]>=6.5",
   "pytest",
@@ -50,22 +51,25 @@ dependencies = [
   "haystack-pydoc-tools",
 ]
 [tool.hatch.envs.default.scripts]
-test = "pytest --reruns 3 --reruns-delay 30 -x {args:tests}"
-test-cov = "coverage run -m pytest --reruns 3 --reruns-delay 30 -x {args:tests}"
+test = "pytest {args:tests}"
+test-cov = "coverage run -m pytest {args:tests}"
+test-cov-retry = "test-cov --reruns 3 --reruns-delay 30 -x"
 cov-report = ["- coverage combine", "coverage report"]
 cov = ["test-cov", "cov-report"]
+cov-retry = ["test-cov-retry", "cov-report"]
 docs = ["pydoc-markdown pydoc/config.yml"]
 
 [[tool.hatch.envs.all.matrix]]
 python = ["3.8", "3.9", "3.10", "3.11", "3.12"]
 
 [tool.hatch.envs.lint]
+installer = "uv"
 detached = true
-dependencies = ["black>=23.1.0", "mypy>=1.0.0", "ruff>=0.0.243"]
+dependencies = ["pip", "black>=23.1.0", "mypy>=1.0.0", "ruff>=0.0.243"]
 [tool.hatch.envs.lint.scripts]
 typing = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}"
 style = ["ruff check {args:.}", "black --check --diff {args:.}"]
-fmt = ["black {args:.}", "ruff --fix {args:.}", "style"]
+fmt = ["black {args:.}", "ruff check --fix {args:.}", "style"]
 all = ["style", "typing"]
 
 [tool.black]
@@ -76,6 +80,8 @@ skip-string-normalization = true
 [tool.ruff]
 target-version = "py38"
 line-length = 120
+
+[tool.ruff.lint]
 select = [
   "A",
   "ARG",
@@ -126,13 +132,13 @@ unfixable = [
   "F401",
 ]
 
-[tool.ruff.isort]
-known-first-party = ["src"]
+[tool.ruff.lint.isort]
+known-first-party = ["haystack_integrations"]
 
-[tool.ruff.flake8-tidy-imports]
+[tool.ruff.lint.flake8-tidy-imports]
 ban-relative-imports = "parents"
 
-[tool.ruff.per-file-ignores]
+[tool.ruff.lint.per-file-ignores]
 # Tests can use magic values, assertions, and relative imports
 "tests/**/*" = ["PLR2004", "S101", "TID252"]
 # examples can contain "print" commands

{pgvector_haystack-0.5.1 → pgvector_haystack-1.1.0}/src/haystack_integrations/components/retrievers/pgvector/embedding_retriever.py

@@ -7,6 +7,7 @@ from haystack import component, default_from_dict, default_to_dict
 from haystack.dataclasses import Document
 from haystack.document_stores.types import FilterPolicy
 from haystack.document_stores.types.filter_policy import apply_filter_policy
+
 from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
 from haystack_integrations.document_stores.pgvector.document_store import VALID_VECTOR_FUNCTIONS
 

{pgvector_haystack-0.5.1 → pgvector_haystack-1.1.0}/src/haystack_integrations/components/retrievers/pgvector/keyword_retriever.py

@@ -7,6 +7,7 @@ from haystack import component, default_from_dict, default_to_dict
 from haystack.dataclasses import Document
 from haystack.document_stores.types import FilterPolicy
 from haystack.document_stores.types.filter_policy import apply_filter_policy
+
 from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
 
 

{pgvector_haystack-0.5.1 → pgvector_haystack-1.1.0}/src/haystack_integrations/document_stores/pgvector/document_store.py

@@ -9,7 +9,6 @@ from haystack.dataclasses.document import ByteStream, Document
 from haystack.document_stores.errors import DocumentStoreError, DuplicateDocumentError
 from haystack.document_stores.types import DuplicatePolicy
 from haystack.utils.auth import Secret, deserialize_secrets_inplace
-from haystack.utils.filters import convert
 from psycopg import Error, IntegrityError, connect
 from psycopg.abc import Query
 from psycopg.cursor import Cursor
@@ -25,7 +24,7 @@ from .filters import _convert_filters_to_where_clause_and_params
 logger = logging.getLogger(__name__)
 
 CREATE_TABLE_STATEMENT = """
-CREATE TABLE IF NOT EXISTS {table_name} (
+CREATE TABLE IF NOT EXISTS {schema_name}.{table_name} (
 id VARCHAR(128) PRIMARY KEY,
 embedding VECTOR({embedding_dimension}),
 content TEXT,
@@ -37,7 +36,7 @@ meta JSONB)
 """
 
 INSERT_STATEMENT = """
-INSERT INTO {table_name}
+INSERT INTO {schema_name}.{table_name}
 (id, embedding, content, dataframe, blob_data, blob_meta, blob_mime_type, meta)
 VALUES (%(id)s, %(embedding)s, %(content)s, %(dataframe)s, %(blob_data)s, %(blob_meta)s, %(blob_mime_type)s, %(meta)s)
 """
@@ -55,7 +54,7 @@ meta = EXCLUDED.meta
 
 KEYWORD_QUERY = """
 SELECT {table_name}.*, ts_rank_cd(to_tsvector({language}, content), query) AS score
-FROM {table_name}, plainto_tsquery({language}, %s) query
+FROM {schema_name}.{table_name}, plainto_tsquery({language}, %s) query
 WHERE to_tsvector({language}, content) @@ query
 """
 
@@ -79,6 +78,7 @@ class PgvectorDocumentStore:
         self,
         *,
         connection_string: Secret = Secret.from_env_var("PG_CONN_STR"),
+        schema_name: str = "public",
         table_name: str = "haystack_documents",
         language: str = "english",
         embedding_dimension: int = 768,
@@ -97,7 +97,12 @@ class PgvectorDocumentStore:
         A specific table to store Haystack documents will be created if it doesn't exist yet.
 
         :param connection_string: The connection string to use to connect to the PostgreSQL database, defined as an
-            environment variable, e.g.: `PG_CONN_STR="postgresql://USER:PASSWORD@HOST:PORT/DB_NAME"`
+            environment variable. It can be provided in either URI format
+            e.g.: `PG_CONN_STR="postgresql://USER:PASSWORD@HOST:PORT/DB_NAME"`, or keyword/value format
+            e.g.: `PG_CONN_STR="host=HOST port=PORT dbname=DBNAME user=USER password=PASSWORD"`
+            See [PostgreSQL Documentation](https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNSTRING)
+            for more details.
+        :param schema_name: The name of the schema the table is created in. The schema must already exist.
         :param table_name: The name of the table to use to store Haystack documents.
         :param language: The language to be used to parse query and document content in keyword retrieval.
             To see the list of available languages, you can run the following SQL query in your PostgreSQL database:
@@ -134,6 +139,7 @@ class PgvectorDocumentStore:
 
         self.connection_string = connection_string
         self.table_name = table_name
+        self.schema_name = schema_name
         self.embedding_dimension = embedding_dimension
         if vector_function not in VALID_VECTOR_FUNCTIONS:
             msg = f"vector_function must be one of {VALID_VECTOR_FUNCTIONS}, but got {vector_function}"
@@ -150,29 +156,41 @@ class PgvectorDocumentStore:
         self._connection = None
         self._cursor = None
         self._dict_cursor = None
+        self._table_initialized = False
 
     @property
     def cursor(self):
-        if self._cursor is None:
+        if self._cursor is None or not self._connection_is_valid(self._connection):
             self._create_connection()
 
         return self._cursor
 
     @property
     def dict_cursor(self):
-        if self._dict_cursor is None:
+        if self._dict_cursor is None or not self._connection_is_valid(self._connection):
             self._create_connection()
 
         return self._dict_cursor
 
     @property
     def connection(self):
-        if self._connection is None:
+        if self._connection is None or not self._connection_is_valid(self._connection):
             self._create_connection()
 
         return self._connection
 
     def _create_connection(self):
+        """
+        Internal method to create a connection to the PostgreSQL database.
+        """
+
+        # close the connection if it already exists
+        if self._connection:
+            try:
+                self._connection.close()
+            except Error as e:
+                logger.debug("Failed to close connection: %s", str(e))
+
         conn_str = self.connection_string.resolve_value() or ""
         connection = connect(conn_str)
         connection.autocommit = True
@@ -183,16 +201,40 @@ class PgvectorDocumentStore:
         self._cursor = self._connection.cursor()
         self._dict_cursor = self._connection.cursor(row_factory=dict_row)
 
-        # Init schema
+        if not self._table_initialized:
+            self._initialize_table()
+
+        return self._connection
+
+    def _initialize_table(self):
+        """
+        Internal method to initialize the table.
+        """
         if self.recreate_table:
             self.delete_table()
+
         self._create_table_if_not_exists()
         self._create_keyword_index_if_not_exists()
 
         if self.search_strategy == "hnsw":
             self._handle_hnsw()
 
-        return self._connection
+        self._table_initialized = True
+
+    @staticmethod
+    def _connection_is_valid(connection):
+        """
+        Internal method to check if the connection is still valid.
+        """
+
+        # implementation inspired to psycopg pool
+        # https://github.com/psycopg/psycopg/blob/d38cf7798b0c602ff43dac9f20bbab96237a9c38/psycopg_pool/psycopg_pool/pool.py#L528
+
+        try:
+            connection.execute("")
+        except Error:
+            return False
+        return True
 
     def to_dict(self) -> Dict[str, Any]:
         """
@@ -204,6 +246,7 @@ class PgvectorDocumentStore:
         return default_to_dict(
             self,
             connection_string=self.connection_string.to_dict(),
+            schema_name=self.schema_name,
             table_name=self.table_name,
             embedding_dimension=self.embedding_dimension,
             vector_function=self.vector_function,
@@ -263,7 +306,9 @@ class PgvectorDocumentStore:
         """
 
         create_sql = SQL(CREATE_TABLE_STATEMENT).format(
-            table_name=Identifier(self.table_name), embedding_dimension=SQLLiteral(self.embedding_dimension)
+            schema_name=Identifier(self.schema_name),
+            table_name=Identifier(self.table_name),
+            embedding_dimension=SQLLiteral(self.embedding_dimension),
         )
 
         self._execute_sql(create_sql, error_msg="Could not create table in PgvectorDocumentStore")
@@ -271,12 +316,18 @@ class PgvectorDocumentStore:
     def delete_table(self):
         """
         Deletes the table used to store Haystack documents.
-        The name of the table (`table_name`) is defined when initializing the `PgvectorDocumentStore`.
+        The name of the schema (`schema_name`) and the name of the table (`table_name`)
+        are defined when initializing the `PgvectorDocumentStore`.
         """
+        delete_sql = SQL("DROP TABLE IF EXISTS {schema_name}.{table_name}").format(
+            schema_name=Identifier(self.schema_name),
+            table_name=Identifier(self.table_name),
+        )
 
-        delete_sql = SQL("DROP TABLE IF EXISTS {table_name}").format(table_name=Identifier(self.table_name))
-
-        self._execute_sql(delete_sql, error_msg=f"Could not delete table {self.table_name} in PgvectorDocumentStore")
+        self._execute_sql(
+            delete_sql,
+            error_msg=f"Could not delete table {self.schema_name}.{self.table_name} in PgvectorDocumentStore",
+        )
 
     def _create_keyword_index_if_not_exists(self):
         """
@@ -284,15 +335,16 @@ class PgvectorDocumentStore:
         """
         index_exists = bool(
             self._execute_sql(
-                "SELECT 1 FROM pg_indexes WHERE tablename = %s AND indexname = %s",
-                (self.table_name, self.keyword_index_name),
+                "SELECT 1 FROM pg_indexes WHERE schemaname = %s AND tablename = %s AND indexname = %s",
+                (self.schema_name, self.table_name, self.keyword_index_name),
                 "Could not check if keyword index exists",
             ).fetchone()
         )
 
         sql_create_index = SQL(
-            "CREATE INDEX {index_name} ON {table_name} USING GIN (to_tsvector({language}, content))"
+            "CREATE INDEX {index_name} ON {schema_name}.{table_name} USING GIN (to_tsvector({language}, content))"
         ).format(
+            schema_name=Identifier(self.schema_name),
             index_name=Identifier(self.keyword_index_name),
             table_name=Identifier(self.table_name),
             language=SQLLiteral(self.language),
@@ -315,8 +367,8 @@ class PgvectorDocumentStore:
 
         index_exists = bool(
             self._execute_sql(
-                "SELECT 1 FROM pg_indexes WHERE tablename = %s AND indexname = %s",
-                (self.table_name, self.hnsw_index_name),
+                "SELECT 1 FROM pg_indexes WHERE schemaname = %s AND tablename = %s AND indexname = %s",
+                (self.schema_name, self.table_name, self.hnsw_index_name),
                 "Could not check if HNSW index exists",
             ).fetchone()
         )
@@ -346,8 +398,13 @@ class PgvectorDocumentStore:
             if key in HNSW_INDEX_CREATION_VALID_KWARGS
         }
 
-        sql_create_index = SQL("CREATE INDEX {index_name} ON {table_name} USING hnsw (embedding {ops}) ").format(
-            index_name=Identifier(self.hnsw_index_name), table_name=Identifier(self.table_name), ops=SQL(pg_ops)
+        sql_create_index = SQL(
+            "CREATE INDEX {index_name} ON {schema_name}.{table_name} USING hnsw (embedding {ops}) "
+        ).format(
+            schema_name=Identifier(self.schema_name),
+            index_name=Identifier(self.hnsw_index_name),
+            table_name=Identifier(self.table_name),
+            ops=SQL(pg_ops),
         )
 
         if actual_hnsw_index_creation_kwargs:
@@ -366,7 +423,9 @@ class PgvectorDocumentStore:
         Returns how many documents are present in the document store.
         """
 
-        sql_count = SQL("SELECT COUNT(*) FROM {table_name}").format(table_name=Identifier(self.table_name))
+        sql_count = SQL("SELECT COUNT(*) FROM {schema_name}.{table_name}").format(
+            schema_name=Identifier(self.schema_name), table_name=Identifier(self.table_name)
+        )
 
         count = self._execute_sql(sql_count, error_msg="Could not count documents in PgvectorDocumentStore").fetchone()[
             0
@@ -389,9 +448,12 @@ class PgvectorDocumentStore:
                 msg = "Filters must be a dictionary"
                 raise TypeError(msg)
             if "operator" not in filters and "conditions" not in filters:
-                filters = convert(filters)
+                msg = "Invalid filter syntax. See https://docs.haystack.deepset.ai/docs/metadata-filtering for details."
+                raise ValueError(msg)
 
-        sql_filter = SQL("SELECT * FROM {table_name}").format(table_name=Identifier(self.table_name))
+        sql_filter = SQL("SELECT * FROM {schema_name}.{table_name}").format(
+            schema_name=Identifier(self.schema_name), table_name=Identifier(self.table_name)
+        )
 
         params = ()
         if filters:
@@ -430,7 +492,9 @@ class PgvectorDocumentStore:
 
         db_documents = self._from_haystack_to_pg_documents(documents)
 
-        sql_insert = SQL(INSERT_STATEMENT).format(table_name=Identifier(self.table_name))
+        sql_insert = SQL(INSERT_STATEMENT).format(
+            schema_name=Identifier(self.schema_name), table_name=Identifier(self.table_name)
+        )
 
         if policy == DuplicatePolicy.OVERWRITE:
             sql_insert += SQL(UPDATE_STATEMENT)
@@ -539,8 +603,10 @@ class PgvectorDocumentStore:
 
         document_ids_str = ", ".join(f"'{document_id}'" for document_id in document_ids)
 
-        delete_sql = SQL("DELETE FROM {table_name} WHERE id IN ({document_ids_str})").format(
-            table_name=Identifier(self.table_name), document_ids_str=SQL(document_ids_str)
+        delete_sql = SQL("DELETE FROM {schema_name}.{table_name} WHERE id IN ({document_ids_str})").format(
+            schema_name=Identifier(self.schema_name),
+            table_name=Identifier(self.table_name),
+            document_ids_str=SQL(document_ids_str),
         )
 
         self._execute_sql(delete_sql, error_msg="Could not delete documents from PgvectorDocumentStore")
@@ -566,6 +632,7 @@ class PgvectorDocumentStore:
             raise ValueError(msg)
 
         sql_select = SQL(KEYWORD_QUERY).format(
+            schema_name=Identifier(self.schema_name),
             table_name=Identifier(self.table_name),
             language=SQLLiteral(self.language),
             query=SQLLiteral(query),
@@ -639,7 +706,8 @@ class PgvectorDocumentStore:
         elif vector_function == "l2_distance":
             score_definition = f"embedding <-> {query_embedding_for_postgres} AS score"
 
-        sql_select = SQL("SELECT *, {score} FROM {table_name}").format(
+        sql_select = SQL("SELECT *, {score} FROM {schema_name}.{table_name}").format(
+            schema_name=Identifier(self.schema_name),
             table_name=Identifier(self.table_name),
             score=SQL(score_definition),
         )

{pgvector_haystack-0.5.1 → pgvector_haystack-1.1.0}/tests/conftest.py

@@ -2,6 +2,7 @@ import os
 from unittest.mock import patch
 
 import pytest
+
 from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
 
 

{pgvector_haystack-0.5.1 → pgvector_haystack-1.1.0}/tests/test_document_store.py

@@ -11,9 +11,10 @@ from haystack.document_stores.errors import DuplicateDocumentError
 from haystack.document_stores.types import DuplicatePolicy
 from haystack.testing.document_store import CountDocumentsTest, DeleteDocumentsTest, WriteDocumentsTest
 from haystack.utils import Secret
-from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
 from pandas import DataFrame
 
+from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
+
 
 @pytest.mark.integration
 class TestDocumentStore(CountDocumentsTest, WriteDocumentsTest, DeleteDocumentsTest):
@@ -40,12 +41,32 @@ class TestDocumentStore(CountDocumentsTest, WriteDocumentsTest, DeleteDocumentsT
         retrieved_docs = document_store.filter_documents()
         assert retrieved_docs == docs
 
+    def test_connection_check_and_recreation(self, document_store: PgvectorDocumentStore):
+        original_connection = document_store.connection
+
+        with patch.object(PgvectorDocumentStore, "_connection_is_valid", return_value=False):
+            new_connection = document_store.connection
+
+        # verify that a new connection is created
+        assert new_connection is not original_connection
+        assert document_store._connection == new_connection
+        assert original_connection.closed
+
+        assert document_store._cursor is not None
+        assert document_store._dict_cursor is not None
+
+        # test with new connection
+        with patch.object(PgvectorDocumentStore, "_connection_is_valid", return_value=True):
+            same_connection = document_store.connection
+            assert same_connection is document_store._connection
+
 
 @pytest.mark.usefixtures("patches_for_unit_tests")
 def test_init(monkeypatch):
     monkeypatch.setenv("PG_CONN_STR", "some_connection_string")
 
     document_store = PgvectorDocumentStore(
+        schema_name="my_schema",
         table_name="my_table",
         embedding_dimension=512,
         vector_function="l2_distance",
@@ -58,6 +79,7 @@ def test_init(monkeypatch):
         keyword_index_name="my_keyword_index",
     )
 
+    assert document_store.schema_name == "my_schema"
     assert document_store.table_name == "my_table"
     assert document_store.embedding_dimension == 512
     assert document_store.vector_function == "l2_distance"
@@ -92,6 +114,7 @@ def test_to_dict(monkeypatch):
         "init_parameters": {
             "connection_string": {"env_vars": ["PG_CONN_STR"], "strict": True, "type": "env_var"},
             "table_name": "my_table",
+            "schema_name": "public",
             "embedding_dimension": 512,
             "vector_function": "l2_distance",
             "recreate_table": True,

{pgvector_haystack-0.5.1 → pgvector_haystack-1.1.0}/tests/test_embedding_retrieval.py

@@ -6,9 +6,10 @@ from typing import List
 
 import pytest
 from haystack.dataclasses.document import Document
-from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
 from numpy.random import rand
 
+from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
+
 
 @pytest.mark.integration
 class TestEmbeddingRetrieval:

{pgvector_haystack-0.5.1 → pgvector_haystack-1.1.0}/tests/test_filters.py

@@ -3,6 +3,10 @@ from typing import List
 import pytest
 from haystack.dataclasses.document import Document
 from haystack.testing.document_store import FilterDocumentsTest
+from pandas import DataFrame
+from psycopg.sql import SQL
+from psycopg.types.json import Jsonb
+
 from haystack_integrations.document_stores.pgvector.filters import (
     FilterError,
     _convert_filters_to_where_clause_and_params,
@@ -10,9 +14,6 @@ from haystack_integrations.document_stores.pgvector.filters import (
     _parse_logical_condition,
     _treat_meta_field,
 )
-from pandas import DataFrame
-from psycopg.sql import SQL
-from psycopg.types.json import Jsonb
 
 
 @pytest.mark.integration

{pgvector_haystack-0.5.1 → pgvector_haystack-1.1.0}/tests/test_keyword_retrieval.py

@@ -1,5 +1,6 @@
 import pytest
 from haystack.dataclasses.document import Document
+
 from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
 
 

{pgvector_haystack-0.5.1 → pgvector_haystack-1.1.0}/tests/test_retrievers.py

@@ -7,6 +7,7 @@ import pytest
 from haystack.dataclasses import Document
 from haystack.document_stores.types import FilterPolicy
 from haystack.utils.auth import EnvVarSecret
+
 from haystack_integrations.components.retrievers.pgvector import PgvectorEmbeddingRetriever, PgvectorKeywordRetriever
 from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
 
@@ -49,6 +50,7 @@ class TestEmbeddingRetriever:
                     "type": "haystack_integrations.document_stores.pgvector.document_store.PgvectorDocumentStore",
                     "init_parameters": {
                         "connection_string": {"env_vars": ["PG_CONN_STR"], "strict": True, "type": "env_var"},
+                        "schema_name": "public",
                         "table_name": "haystack",
                         "embedding_dimension": 768,
                         "vector_function": "cosine_similarity",
@@ -174,6 +176,7 @@ class TestKeywordRetriever:
                     "type": "haystack_integrations.document_stores.pgvector.document_store.PgvectorDocumentStore",
                     "init_parameters": {
                         "connection_string": {"env_vars": ["PG_CONN_STR"], "strict": True, "type": "env_var"},
+                        "schema_name": "public",
                         "table_name": "haystack",
                         "embedding_dimension": 768,
                         "vector_function": "cosine_similarity",
@@ -311,7 +314,7 @@ class TestKeywordRetriever:
         )
         res = retriever.run(query="test query", filters={"field2": "value2"})
 
-        mock_store._keyword_retrieval.assert_called_once_with(
+        mock_store.mock_calls[0].assert_called_once_with(
             query="test query", filters={"field": "value", "field2": "value2"}, top_k=10
         )