pgvector-haystack 0.2.0__tar.gz → 0.3.0__tar.gz

{pgvector_haystack-0.2.0 → pgvector_haystack-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: pgvector-haystack
-Version: 0.2.0
+Version: 0.3.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.3.0/examples/hybrid_retrieval.py

@@ -0,0 +1,69 @@
+# Before running this example, ensure you have PostgreSQL installed with the pgvector extension.
+# For a quick setup using Docker:
+# docker run -d -p 5432:5432 -e POSTGRES_USER=postgres -e POSTGRES_PASSWORD=postgres
+# -e POSTGRES_DB=postgres ankane/pgvector
+
+# Install required packages for this example, including pgvector-haystack and other libraries needed
+# for Markdown conversion and embeddings generation. Use the following command:
+# pip install pgvector-haystack markdown-it-py mdit_plain "sentence-transformers>=2.2.0"
+
+# Download some Markdown files to index.
+# git clone https://github.com/anakin87/neural-search-pills
+
+import glob
+
+from haystack import Pipeline
+from haystack.components.converters import MarkdownToDocument
+from haystack.components.embedders import SentenceTransformersDocumentEmbedder, SentenceTransformersTextEmbedder
+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
+
+# Set an environment variable `PG_CONN_STR` with the connection string to your PostgreSQL database.
+# e.g., "postgresql://USER:PASSWORD@HOST:PORT/DB_NAME"
+
+# Initialize PgvectorDocumentStore
+document_store = PgvectorDocumentStore(
+    table_name="haystack_test",
+    embedding_dimension=768,
+    vector_function="cosine_similarity",
+    recreate_table=True,
+    search_strategy="hnsw",
+)
+
+# Create the indexing Pipeline and index some documents
+file_paths = glob.glob("neural-search-pills/pills/*.md")
+
+
+indexing = Pipeline()
+indexing.add_component("converter", MarkdownToDocument())
+indexing.add_component("splitter", DocumentSplitter(split_by="sentence", split_length=2))
+indexing.add_component("document_embedder", SentenceTransformersDocumentEmbedder())
+indexing.add_component("writer", DocumentWriter(document_store))
+indexing.connect("converter", "splitter")
+indexing.connect("splitter", "document_embedder")
+indexing.connect("document_embedder", "writer")
+
+indexing.run({"converter": {"sources": file_paths}})
+
+# Create the querying Pipeline and try a query
+querying = Pipeline()
+querying.add_component("text_embedder", SentenceTransformersTextEmbedder())
+querying.add_component("retriever", PgvectorEmbeddingRetriever(document_store=document_store, top_k=3))
+querying.add_component("keyword_retriever", PgvectorKeywordRetriever(document_store=document_store, top_k=3))
+querying.add_component(
+    "joiner",
+    DocumentJoiner(join_mode="reciprocal_rank_fusion", top_k=3),
+)
+querying.connect("text_embedder", "retriever")
+querying.connect("keyword_retriever", "joiner")
+querying.connect("retriever", "joiner")
+
+query = "cross-encoder"
+results = querying.run({"text_embedder": {"text": query}, "keyword_retriever": {"query": query}})
+
+for doc in results["joiner"]["documents"]:
+    print(doc)
+    print("-" * 10)

{pgvector_haystack-0.2.0 → pgvector_haystack-0.3.0}/pydoc/config.yml

@@ -3,6 +3,7 @@ loaders:
     search_path: [../src]
     modules: [
       "haystack_integrations.components.retrievers.pgvector.embedding_retriever",
+      "haystack_integrations.components.retrievers.pgvector.keyword_retriever",
       "haystack_integrations.document_stores.pgvector.document_store",
     ]
     ignore_when_discovered: ["__init__"]

{pgvector_haystack-0.2.0 → pgvector_haystack-0.3.0}/src/haystack_integrations/document_stores/pgvector/document_store.py

@@ -145,26 +145,53 @@ class PgvectorDocumentStore:
         self.hnsw_index_creation_kwargs = hnsw_index_creation_kwargs or {}
         self.hnsw_ef_search = hnsw_ef_search
         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()
+
+        return self._dict_cursor
 
-        # we create a generic cursor and another one that returns dictionaries
-        self._cursor = connection.cursor()
-        self._dict_cursor = connection.cursor(row_factory=dict_row)
+    @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.
@@ -208,11 +235,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 +247,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
 
@@ -371,7 +398,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 +435,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 +455,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 +510,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 +582,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 +660,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 → pgvector_haystack-0.3.0}/tests/conftest.py

@@ -40,7 +40,6 @@ def patches_for_unit_tests():
     ) as mock_create_kw_index, patch(
         "haystack_integrations.document_stores.pgvector.document_store.PgvectorDocumentStore._handle_hnsw"
     ) as mock_hnsw:
-
         yield mock_connect, mock_register, mock_delete, mock_create, mock_create_kw_index, mock_hnsw
 
 

{pgvector_haystack-0.2.0 → pgvector_haystack-0.3.0}/tests/test_document_store.py

@@ -4,11 +4,13 @@
 
 from unittest.mock import patch
 
+import numpy as np
 import pytest
 from haystack.dataclasses.document import ByteStream, Document
 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
 
@@ -170,7 +172,7 @@ def test_from_pg_to_haystack_documents():
             "blob_meta": None,
             "blob_mime_type": None,
             "meta": {"meta_key": "meta_value"},
-            "embedding": "[0.1, 0.2, 0.3]",
+            "embedding": np.array([0.1, 0.2, 0.3]),
         },
         {
             "id": "2",
@@ -180,7 +182,7 @@ def test_from_pg_to_haystack_documents():
             "blob_meta": None,
             "blob_mime_type": None,
             "meta": {"meta_key": "meta_value"},
-            "embedding": "[0.4, 0.5, 0.6]",
+            "embedding": np.array([0.4, 0.5, 0.6]),
         },
         {
             "id": "3",
@@ -190,16 +192,11 @@ def test_from_pg_to_haystack_documents():
             "blob_meta": {"blob_meta_key": "blob_meta_value"},
             "blob_mime_type": "mime_type",
             "meta": {"meta_key": "meta_value"},
-            "embedding": "[0.7, 0.8, 0.9]",
+            "embedding": np.array([0.7, 0.8, 0.9]),
         },
     ]
 
-    with patch(
-        "haystack_integrations.document_stores.pgvector.document_store.PgvectorDocumentStore.__init__"
-    ) as mock_init:
-        mock_init.return_value = None
-        ds = PgvectorDocumentStore(connection_string="test")
-
+    ds = PgvectorDocumentStore(connection_string=Secret.from_token("test"))
     haystack_docs = ds._from_pg_to_haystack_documents(pg_docs)
 
     assert haystack_docs[0].id == "1"