pgvector-haystack 3.0.1__py3-none-any.whl → 3.2.0__py3-none-any.whl

haystack_integrations/components/retrievers/pgvector/embedding_retriever.py

@@ -151,7 +151,8 @@ class PgvectorEmbeddingRetriever:
         :param top_k: Maximum number of Documents to return.
         :param vector_function: The similarity function to use when searching for similar embeddings.
 
-        :returns: List of Documents similar to `query_embedding`.
+        :returns: A dictionary with the following keys:
+            - `documents`: List of `Document`s that are similar to `query_embedding`.
         """
         filters = apply_filter_policy(self.filter_policy, self.filters, filters)
         top_k = top_k or self.top_k
@@ -164,3 +165,36 @@ class PgvectorEmbeddingRetriever:
             vector_function=vector_function,
         )
         return {"documents": docs}
+
+    @component.output_types(documents=List[Document])
+    async def run_async(
+        self,
+        query_embedding: List[float],
+        filters: Optional[Dict[str, Any]] = None,
+        top_k: Optional[int] = None,
+        vector_function: Optional[Literal["cosine_similarity", "inner_product", "l2_distance"]] = None,
+    ):
+        """
+        Asynchronously retrieve documents from the `PgvectorDocumentStore`, based on their embeddings.
+
+        :param query_embedding: Embedding of the query.
+        :param filters: Filters applied to the retrieved Documents. The way runtime filters are applied depends on
+                        the `filter_policy` chosen at retriever initialization. See init method docstring for more
+                        details.
+        :param top_k: Maximum number of Documents to return.
+        :param vector_function: The similarity function to use when searching for similar embeddings.
+
+        :returns: A dictionary with the following keys:
+            - `documents`: List of `Document`s that are similar to `query_embedding`.
+        """
+        filters = apply_filter_policy(self.filter_policy, self.filters, filters)
+        top_k = top_k or self.top_k
+        vector_function = vector_function or self.vector_function
+
+        docs = await self.document_store._embedding_retrieval_async(
+            query_embedding=query_embedding,
+            filters=filters,
+            top_k=top_k,
+            vector_function=vector_function,
+        )
+        return {"documents": docs}

haystack_integrations/components/retrievers/pgvector/keyword_retriever.py

@@ -136,3 +136,33 @@ class PgvectorKeywordRetriever:
             top_k=top_k,
         )
         return {"documents": docs}
+
+    @component.output_types(documents=List[Document])
+    async def run_async(
+        self,
+        query: str,
+        filters: Optional[Dict[str, Any]] = None,
+        top_k: Optional[int] = None,
+    ):
+        """
+        Asynchronously retrieve documents from the `PgvectorDocumentStore`, based on keywords.
+
+        :param query: String to search in `Document`s' content.
+        :param filters: Filters applied to the retrieved Documents. The way runtime filters are applied depends on
+                        the `filter_policy` chosen at retriever initialization. See init method docstring for more
+                        details.
+        :param top_k: Maximum number of Documents to return.
+
+        :returns: A dictionary with the following keys:
+            - `documents`: List of `Document`s that match the query.
+        """
+        filters = apply_filter_policy(self.filter_policy, self.filters, filters)
+
+        top_k = top_k or self.top_k
+
+        docs = await self.document_store._keyword_retrieval_async(
+            query=query,
+            filters=filters,
+            top_k=top_k,
+        )
+        return {"documents": docs}

haystack_integrations/document_stores/pgvector/converters.py

@@ -0,0 +1,69 @@
+from typing import Any, Dict, List
+
+from haystack import logging
+from haystack.dataclasses import ByteStream, Document
+from psycopg.types.json import Jsonb
+
+logger = logging.getLogger(__name__)
+
+
+def _from_haystack_to_pg_documents(documents: List[Document]) -> List[Dict[str, Any]]:
+    """
+    Internal method to convert a list of Haystack Documents to a list of dictionaries that can be used to insert
+    documents into the PgvectorDocumentStore.
+    """
+
+    db_documents = []
+    for document in documents:
+        db_document = {k: v for k, v in document.to_dict(flatten=False).items() if k not in ["score", "blob"]}
+
+        blob = document.blob
+        db_document["blob_data"] = blob.data if blob else None
+        db_document["blob_meta"] = Jsonb(blob.meta) if blob and blob.meta else None
+        db_document["blob_mime_type"] = blob.mime_type if blob and blob.mime_type else None
+        db_document["meta"] = Jsonb(db_document["meta"])
+
+        if "sparse_embedding" in db_document:
+            sparse_embedding = db_document.pop("sparse_embedding", None)
+            if sparse_embedding:
+                logger.warning(
+                    "Document {doc_id} has the `sparse_embedding` field set,"
+                    "but storing sparse embeddings in Pgvector is not currently supported."
+                    "The `sparse_embedding` field will be ignored.",
+                    doc_id=db_document["id"],
+                )
+
+        db_documents.append(db_document)
+
+    return db_documents
+
+
+def _from_pg_to_haystack_documents(documents: List[Dict[str, Any]]) -> List[Document]:
+    """
+    Internal method to convert a list of dictionaries from pgvector to a list of Haystack Documents.
+    """
+
+    haystack_documents = []
+    for document in documents:
+        haystack_dict = dict(document)
+        blob_data = haystack_dict.pop("blob_data")
+        blob_meta = haystack_dict.pop("blob_meta")
+        blob_mime_type = haystack_dict.pop("blob_mime_type")
+
+        # convert the embedding to a list of floats
+        # for strange reasons, halfvec and vector have different methods to convert the embedding to a list
+        if document.get("embedding") is not None:
+            if hasattr(document["embedding"], "tolist"):  # vector
+                haystack_dict["embedding"] = document["embedding"].tolist()
+            else:  # halfvec
+                haystack_dict["embedding"] = document["embedding"].to_list()
+
+        haystack_document = Document.from_dict(haystack_dict)
+
+        if blob_data:
+            blob = ByteStream(data=blob_data, meta=blob_meta, mime_type=blob_mime_type)
+            haystack_document.blob = blob
+
+        haystack_documents.append(haystack_document)
+
+    return haystack_documents