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

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

@@ -1,6 +1,6 @@
 # Changelog
 
-## [unreleased]
+## [integrations/pgvector-v1.0.0] - 2024-09-12
 
 ### 🚀 Features
 
@@ -10,10 +10,15 @@
 
 - `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
 

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

@@ -1,13 +1,12 @@
 Metadata-Version: 2.3
 Name: pgvector-haystack
-Version: 1.0.0
+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-1.0.0 → 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",
@@ -62,12 +63,13 @@ docs = ["pydoc-markdown pydoc/config.yml"]
 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:. --exclude tests/, examples/}", "black --check --diff {args:.}"]
-fmt = ["black {args:.}", "ruff --fix {args:. --exclude tests/, examples/}", "style"]
+style = ["ruff check {args:.}", "black --check --diff {args:.}"]
+fmt = ["black {args:.}", "ruff check --fix {args:.}", "style"]
 all = ["style", "typing"]
 
 [tool.black]
@@ -78,6 +80,8 @@ skip-string-normalization = true
 [tool.ruff]
 target-version = "py38"
 line-length = 120
+
+[tool.ruff.lint]
 select = [
   "A",
   "ARG",
@@ -128,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-1.0.0 → pgvector_haystack-1.1.0}/src/haystack_integrations/document_stores/pgvector/document_store.py

@@ -24,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,
@@ -36,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)
 """
@@ -54,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
 """
 
@@ -78,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,
@@ -96,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:
@@ -133,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}"
@@ -149,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
@@ -182,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]:
         """
@@ -203,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,
@@ -262,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")
@@ -270,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):
         """
@@ -283,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),
@@ -314,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()
         )
@@ -345,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:
@@ -365,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
@@ -391,7 +451,9 @@ class PgvectorDocumentStore:
                 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-1.0.0 → pgvector_haystack-1.1.0}/tests/test_document_store.py

@@ -41,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",
@@ -59,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"
@@ -93,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-1.0.0 → pgvector_haystack-1.1.0}/tests/test_retrievers.py

@@ -50,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",
@@ -175,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",