pgvector-haystack 0.2.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

haystack_integrations/document_stores/pgvector/document_store.py

@@ -69,10 +69,6 @@ VECTOR_FUNCTION_TO_POSTGRESQL_OPS = {
 
 HNSW_INDEX_CREATION_VALID_KWARGS = ["m", "ef_construction"]
 
-HNSW_INDEX_NAME = "haystack_hnsw_index"
-
-KEYWORD_INDEX_NAME = "haystack_keyword_index"
-
 
 class PgvectorDocumentStore:
     """
@@ -91,7 +87,9 @@ class PgvectorDocumentStore:
         search_strategy: Literal["exact_nearest_neighbor", "hnsw"] = "exact_nearest_neighbor",
         hnsw_recreate_index_if_exists: bool = False,
         hnsw_index_creation_kwargs: Optional[Dict[str, int]] = None,
+        hnsw_index_name: str = "haystack_hnsw_index",
         hnsw_ef_search: Optional[int] = None,
+        keyword_index_name: str = "haystack_keyword_index",
     ):
         """
         Creates a new PgvectorDocumentStore instance.
@@ -127,9 +125,11 @@ class PgvectorDocumentStore:
         :param hnsw_index_creation_kwargs: Additional keyword arguments to pass to the HNSW index creation.
             Only used if search_strategy is set to `"hnsw"`. You can find the list of valid arguments in the
             [pgvector documentation](https://github.com/pgvector/pgvector?tab=readme-ov-file#hnsw)
+        :param hnsw_index_name: Index name for the HNSW index.
         :param hnsw_ef_search: The `ef_search` parameter to use at query time. Only used if search_strategy is set to
             `"hnsw"`. You can find more information about this parameter in the
             [pgvector documentation](https://github.com/pgvector/pgvector?tab=readme-ov-file#hnsw).
+        :param keyword_index_name: Index name for the Keyword index.
         """
 
         self.connection_string = connection_string
@@ -143,28 +143,57 @@ class PgvectorDocumentStore:
         self.search_strategy = search_strategy
         self.hnsw_recreate_index_if_exists = hnsw_recreate_index_if_exists
         self.hnsw_index_creation_kwargs = hnsw_index_creation_kwargs or {}
+        self.hnsw_index_name = hnsw_index_name
         self.hnsw_ef_search = hnsw_ef_search
+        self.keyword_index_name = keyword_index_name
         self.language = language
+        self._connection = None
+        self._cursor = None
+        self._dict_cursor = None
 
-        connection = connect(self.connection_string.resolve_value())
-        connection.autocommit = True
-        self._connection = connection
+    @property
+    def cursor(self):
+        if self._cursor is None:
+            self._create_connection()
+
+        return self._cursor
+
+    @property
+    def dict_cursor(self):
+        if self._dict_cursor is None:
+            self._create_connection()
 
-        # we create a generic cursor and another one that returns dictionaries
-        self._cursor = connection.cursor()
-        self._dict_cursor = connection.cursor(row_factory=dict_row)
+        return self._dict_cursor
 
+    @property
+    def connection(self):
+        if self._connection is None:
+            self._create_connection()
+
+        return self._connection
+
+    def _create_connection(self):
+        conn_str = self.connection_string.resolve_value() or ""
+        connection = connect(conn_str)
+        connection.autocommit = True
         connection.execute("CREATE EXTENSION IF NOT EXISTS vector")
-        register_vector(connection)
+        register_vector(connection)  # Note: this must be called before creating the cursors.
 
-        if recreate_table:
+        self._connection = connection
+        self._cursor = self._connection.cursor()
+        self._dict_cursor = self._connection.cursor(row_factory=dict_row)
+
+        # Init schema
+        if self.recreate_table:
             self.delete_table()
         self._create_table_if_not_exists()
         self._create_keyword_index_if_not_exists()
 
-        if search_strategy == "hnsw":
+        if self.search_strategy == "hnsw":
             self._handle_hnsw()
 
+        return self._connection
+
     def to_dict(self) -> Dict[str, Any]:
         """
         Serializes the component to a dictionary.
@@ -182,7 +211,9 @@ class PgvectorDocumentStore:
             search_strategy=self.search_strategy,
             hnsw_recreate_index_if_exists=self.hnsw_recreate_index_if_exists,
             hnsw_index_creation_kwargs=self.hnsw_index_creation_kwargs,
+            hnsw_index_name=self.hnsw_index_name,
             hnsw_ef_search=self.hnsw_ef_search,
+            keyword_index_name=self.keyword_index_name,
             language=self.language,
         )
 
@@ -208,11 +239,11 @@ class PgvectorDocumentStore:
         :param sql_query: The SQL query to execute.
         :param params: The parameters to pass to the SQL query.
         :param error_msg: The error message to use if an exception is raised.
-        :param cursor: The cursor to use to execute the SQL query. Defaults to self._cursor.
+        :param cursor: The cursor to use to execute the SQL query. Defaults to self.cursor.
         """
 
         params = params or ()
-        cursor = cursor or self._cursor
+        cursor = cursor or self.cursor
 
         sql_query_str = sql_query.as_string(cursor) if not isinstance(sql_query, str) else sql_query
         logger.debug("SQL query: %s\nParameters: %s", sql_query_str, params)
@@ -220,7 +251,7 @@ class PgvectorDocumentStore:
         try:
             result = cursor.execute(sql_query, params)
         except Error as e:
-            self._connection.rollback()
+            self.connection.rollback()
             detailed_error_msg = f"{error_msg}.\nYou can find the SQL query and the parameters in the debug logs."
             raise DocumentStoreError(detailed_error_msg) from e
 
@@ -254,7 +285,7 @@ class PgvectorDocumentStore:
         index_exists = bool(
             self._execute_sql(
                 "SELECT 1 FROM pg_indexes WHERE tablename = %s AND indexname = %s",
-                (self.table_name, KEYWORD_INDEX_NAME),
+                (self.table_name, self.keyword_index_name),
                 "Could not check if keyword index exists",
             ).fetchone()
         )
@@ -262,7 +293,7 @@ class PgvectorDocumentStore:
         sql_create_index = SQL(
             "CREATE INDEX {index_name} ON {table_name} USING GIN (to_tsvector({language}, content))"
         ).format(
-            index_name=Identifier(KEYWORD_INDEX_NAME),
+            index_name=Identifier(self.keyword_index_name),
             table_name=Identifier(self.table_name),
             language=SQLLiteral(self.language),
         )
@@ -285,7 +316,7 @@ class PgvectorDocumentStore:
         index_exists = bool(
             self._execute_sql(
                 "SELECT 1 FROM pg_indexes WHERE tablename = %s AND indexname = %s",
-                (self.table_name, HNSW_INDEX_NAME),
+                (self.table_name, self.hnsw_index_name),
                 "Could not check if HNSW index exists",
             ).fetchone()
         )
@@ -298,7 +329,7 @@ class PgvectorDocumentStore:
             )
             return
 
-        sql_drop_index = SQL("DROP INDEX IF EXISTS {index_name}").format(index_name=Identifier(HNSW_INDEX_NAME))
+        sql_drop_index = SQL("DROP INDEX IF EXISTS {index_name}").format(index_name=Identifier(self.hnsw_index_name))
         self._execute_sql(sql_drop_index, error_msg="Could not drop HNSW index")
 
         self._create_hnsw_index()
@@ -316,7 +347,7 @@ class PgvectorDocumentStore:
         }
 
         sql_create_index = SQL("CREATE INDEX {index_name} ON {table_name} USING hnsw (embedding {ops}) ").format(
-            index_name=Identifier(HNSW_INDEX_NAME), table_name=Identifier(self.table_name), ops=SQL(pg_ops)
+            index_name=Identifier(self.hnsw_index_name), table_name=Identifier(self.table_name), ops=SQL(pg_ops)
         )
 
         if actual_hnsw_index_creation_kwargs:
@@ -371,7 +402,7 @@ class PgvectorDocumentStore:
             sql_filter,
             params,
             error_msg="Could not filter documents from PgvectorDocumentStore.",
-            cursor=self._dict_cursor,
+            cursor=self.dict_cursor,
         )
 
         records = result.fetchall()
@@ -408,16 +439,16 @@ class PgvectorDocumentStore:
 
         sql_insert += SQL(" RETURNING id")
 
-        sql_query_str = sql_insert.as_string(self._cursor) if not isinstance(sql_insert, str) else sql_insert
+        sql_query_str = sql_insert.as_string(self.cursor) if not isinstance(sql_insert, str) else sql_insert
         logger.debug("SQL query: %s\nParameters: %s", sql_query_str, db_documents)
 
         try:
-            self._cursor.executemany(sql_insert, db_documents, returning=True)
+            self.cursor.executemany(sql_insert, db_documents, returning=True)
         except IntegrityError as ie:
-            self._connection.rollback()
+            self.connection.rollback()
             raise DuplicateDocumentError from ie
         except Error as e:
-            self._connection.rollback()
+            self.connection.rollback()
             error_msg = (
                 "Could not write documents to PgvectorDocumentStore. \n"
                 "You can find the SQL query and the parameters in the debug logs."
@@ -428,9 +459,9 @@ class PgvectorDocumentStore:
         # https://www.psycopg.org/psycopg3/docs/api/cursors.html#psycopg.Cursor.executemany
         written_docs = 0
         while True:
-            if self._cursor.fetchone():
+            if self.cursor.fetchone():
                 written_docs += 1
-            if not self._cursor.nextset():
+            if not self.cursor.nextset():
                 break
 
         return written_docs
@@ -483,8 +514,8 @@ class PgvectorDocumentStore:
 
             # postgresql returns the embedding as a string
             # so we need to convert it to a list of floats
-            if document.get("embedding"):
-                haystack_dict["embedding"] = [float(el) for el in document["embedding"].strip("[]").split(",")]
+            if document.get("embedding") is not None:
+                haystack_dict["embedding"] = document["embedding"].tolist()
 
             haystack_document = Document.from_dict(haystack_dict)
 
@@ -555,7 +586,7 @@ class PgvectorDocumentStore:
             sql_query,
             (query, *where_params),
             error_msg="Could not retrieve documents from PgvectorDocumentStore.",
-            cursor=self._dict_cursor,
+            cursor=self.dict_cursor,
         )
 
         records = result.fetchall()
@@ -633,7 +664,7 @@ class PgvectorDocumentStore:
             sql_query,
             params,
             error_msg="Could not retrieve documents from PgvectorDocumentStore.",
-            cursor=self._dict_cursor,
+            cursor=self.dict_cursor,
         )
 
         records = result.fetchall()

{pgvector_haystack-0.2.0.dist-info → pgvector_haystack-0.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: pgvector-haystack
-Version: 0.2.0
+Version: 0.4.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

{pgvector_haystack-0.2.0.dist-info → pgvector_haystack-0.4.0.dist-info}/RECORD

@@ -2,9 +2,9 @@ haystack_integrations/components/retrievers/pgvector/__init__.py,sha256=n4q4pUAp
 haystack_integrations/components/retrievers/pgvector/embedding_retriever.py,sha256=6LuQmF7JmPOCU_I10ksPcbDC6QMu-2SOq0--BQxBxpg,6523
 haystack_integrations/components/retrievers/pgvector/keyword_retriever.py,sha256=fAJ3M2l_Pt8zKqHES1l1L4gM16BOHeZWq8xc18Bhq7A,4624
 haystack_integrations/document_stores/pgvector/__init__.py,sha256=Ge-7ggarfr3x8UqNdFHUuFGiu52u49sOVe3l16WEwZU,197
-haystack_integrations/document_stores/pgvector/document_store.py,sha256=KY0FZGTjz4ejxKNqQ9UHDJUArsghwf7Iag45YfXgm2g,27011
+haystack_integrations/document_stores/pgvector/document_store.py,sha256=bvN4tYQeNNZ0s07Ur8VtJNDL_XtFQqHq0IYAUbkmTho,28018
 haystack_integrations/document_stores/pgvector/filters.py,sha256=751_Z1QnXJCG0GWC2jcfeZ25NNYfJewQEIUCeJS70PA,8212
-pgvector_haystack-0.2.0.dist-info/METADATA,sha256=fD4i8wnlapdeHvJMpj11zlmGS_MxN_d6GDFvc3W_i1A,2219
-pgvector_haystack-0.2.0.dist-info/WHEEL,sha256=zEMcRr9Kr03x1ozGwg5v9NQBKn3kndp6LSoSlVg-jhU,87
-pgvector_haystack-0.2.0.dist-info/licenses/LICENSE.txt,sha256=B05uMshqTA74s-0ltyHKI6yoPfJ3zYgQbvcXfDVGFf8,10280
-pgvector_haystack-0.2.0.dist-info/RECORD,,
+pgvector_haystack-0.4.0.dist-info/METADATA,sha256=dNAGH9BMjyn5qlnaAdQqBpMUElmt173XoS9PEoiuCVY,2219
+pgvector_haystack-0.4.0.dist-info/WHEEL,sha256=zEMcRr9Kr03x1ozGwg5v9NQBKn3kndp6LSoSlVg-jhU,87
+pgvector_haystack-0.4.0.dist-info/licenses/LICENSE.txt,sha256=B05uMshqTA74s-0ltyHKI6yoPfJ3zYgQbvcXfDVGFf8,10280
+pgvector_haystack-0.4.0.dist-info/RECORD,,