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

{pgvector_haystack-0.0.1 → pgvector_haystack-0.2.0}/.gitignore

@@ -134,3 +134,4 @@ dmypy.json
 
 # Docs generation artifacts
 _readme_*.md
+.idea

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

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: pgvector-haystack
-Version: 0.0.1
+Version: 0.2.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
@@ -9,6 +9,7 @@ Author-email: deepset GmbH <info@deepset.ai>
 License-Expression: Apache-2.0
 License-File: LICENSE.txt
 Classifier: Development Status :: 4 - Beta
+Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python
 Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
@@ -45,12 +46,23 @@ pip install pgvector-haystack
 
 ## Testing
 
-TODO
+Ensure that you have a PostgreSQL running with the `pgvector` extension. For a quick setup using Docker, run:
+```
+docker run -d -p 5432:5432 -e POSTGRES_USER=postgres -e POSTGRES_PASSWORD=postgres -e POSTGRES_DB=postgres ankane/pgvector
+```
+
+then run the tests:
 
 ```console
 hatch run test
 ```
 
+To run the coverage report:
+
+```console
+hatch run cov
+```
+
 ## License
 
 `pgvector-haystack` is distributed under the terms of the [Apache-2.0](https://spdx.org/licenses/Apache-2.0.html) license.

{pgvector_haystack-0.0.1 → pgvector_haystack-0.2.0}/README.md

@@ -20,12 +20,23 @@ pip install pgvector-haystack
 
 ## Testing
 
-TODO
+Ensure that you have a PostgreSQL running with the `pgvector` extension. For a quick setup using Docker, run:
+```
+docker run -d -p 5432:5432 -e POSTGRES_USER=postgres -e POSTGRES_PASSWORD=postgres -e POSTGRES_DB=postgres ankane/pgvector
+```
+
+then run the tests:
 
 ```console
 hatch run test
 ```
 
+To run the coverage report:
+
+```console
+hatch run cov
+```
+
 ## License
 
 `pgvector-haystack` is distributed under the terms of the [Apache-2.0](https://spdx.org/licenses/Apache-2.0.html) license.

pgvector_haystack-0.2.0/examples/embedding_retrieval.py

@@ -0,0 +1,60 @@
+# 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.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
+
+# 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("embedder", SentenceTransformersDocumentEmbedder())
+indexing.add_component("writer", DocumentWriter(document_store))
+indexing.connect("converter", "splitter")
+indexing.connect("splitter", "embedder")
+indexing.connect("embedder", "writer")
+
+indexing.run({"converter": {"sources": file_paths}})
+
+# Create the querying Pipeline and try a query
+querying = Pipeline()
+querying.add_component("embedder", SentenceTransformersTextEmbedder())
+querying.add_component("retriever", PgvectorEmbeddingRetriever(document_store=document_store, top_k=3))
+querying.connect("embedder", "retriever")
+
+results = querying.run({"embedder": {"text": "What is a cross-encoder?"}})
+
+for doc in results["retriever"]["documents"]:
+    print(doc)
+    print("-" * 10)

pgvector_haystack-0.2.0/pydoc/config.yml

@@ -0,0 +1,30 @@
+loaders:
+  - type: haystack_pydoc_tools.loaders.CustomPythonLoader
+    search_path: [../src]
+    modules: [
+      "haystack_integrations.components.retrievers.pgvector.embedding_retriever",
+      "haystack_integrations.document_stores.pgvector.document_store",
+    ]
+    ignore_when_discovered: ["__init__"]
+processors:
+  - type: filter
+    expression:
+    documented_only: true
+    do_not_filter_modules: false
+    skip_empty_modules: true
+  - type: smart
+  - type: crossref
+renderer:
+  type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer
+  excerpt: Pgvector integration for Haystack
+  category_slug: integrations-api
+  title: Pgvector
+  slug: integrations-pgvector
+  order: 190
+  markdown:
+    descriptive_class_title: false
+    classdef_code_block: false
+    descriptive_module_title: true
+    add_method_class_prefix: true
+    add_member_class_prefix: false
+    filename: _readme_pgvector.md

{pgvector_haystack-0.0.1 → pgvector_haystack-0.2.0}/pyproject.toml

@@ -14,6 +14,7 @@ authors = [
   { name = "deepset GmbH", email = "info@deepset.ai" },
 ]
 classifiers = [
+  "License :: OSI Approved :: Apache Software License",
   "Development Status :: 4 - Beta",
   "Programming Language :: Python",
   "Programming Language :: Python :: 3.8",
@@ -51,6 +52,7 @@ dependencies = [
   "coverage[toml]>=6.5",
   "pytest",
   "ipython",
+  "haystack-pydoc-tools",
 ]
 [tool.hatch.envs.default.scripts]
 test = "pytest {args:tests}"
@@ -63,6 +65,9 @@ cov = [
   "test-cov",
   "cov-report",
 ]
+docs = [
+  "pydoc-markdown pydoc/config.yml"
+]
 
 [[tool.hatch.envs.all.matrix]]
 python = ["3.8", "3.9", "3.10", "3.11", "3.12"]
@@ -91,12 +96,12 @@ all = [
 ]
 
 [tool.black]
-target-version = ["py37"]
+target-version = ["py38"]
 line-length = 120
 skip-string-normalization = true
 
 [tool.ruff]
-target-version = "py37"
+target-version = "py38"
 line-length = 120
 select = [
   "A",
@@ -134,6 +139,8 @@ ignore = [
   "S105", "S106", "S107",
   # Ignore complexity
   "C901", "PLR0911", "PLR0912", "PLR0913", "PLR0915",
+  # ignore function-call-in-default-argument
+  "B008",
 ]
 unfixable = [
   # Don't touch unused imports
@@ -149,24 +156,30 @@ ban-relative-imports = "parents"
 [tool.ruff.per-file-ignores]
 # Tests can use magic values, assertions, and relative imports
 "tests/**/*" = ["PLR2004", "S101", "TID252"]
+# examples can contain "print" commands
+"examples/**/*" = ["T201"]
+
 
 [tool.coverage.run]
-source_pkgs = ["src", "tests"]
+source = ["haystack_integrations"]
 branch = true
-parallel = true
-
-
-[tool.coverage.paths]
-weaviate_haystack = ["src/haystack_integrations", "*/pgvector-haystack/src"]
-tests = ["tests", "*/pgvector-haystack/tests"]
+parallel = false
 
 [tool.coverage.report]
+omit = ["*/tests/*", "*/__init__.py"]
+show_missing=true
 exclude_lines = [
   "no cov",
   "if __name__ == .__main__.:",
   "if TYPE_CHECKING:",
 ]
 
+[tool.pytest.ini_options]
+markers = [
+  "integration: integration tests"
+]
+
+
 [[tool.mypy.overrides]]
 module = [
   "haystack.*",

{pgvector_haystack-0.0.1 → pgvector_haystack-0.2.0}/src/haystack_integrations/components/retrievers/pgvector/__init__.py

@@ -2,5 +2,6 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 from .embedding_retriever import PgvectorEmbeddingRetriever
+from .keyword_retriever import PgvectorKeywordRetriever
 
-__all__ = ["PgvectorEmbeddingRetriever"]
+__all__ = ["PgvectorEmbeddingRetriever", "PgvectorKeywordRetriever"]

{pgvector_haystack-0.0.1 → pgvector_haystack-0.2.0}/src/haystack_integrations/components/retrievers/pgvector/embedding_retriever.py

@@ -12,9 +12,47 @@ from haystack_integrations.document_stores.pgvector.document_store import VALID_
 @component
 class PgvectorEmbeddingRetriever:
     """
-    Retrieves documents from the PgvectorDocumentStore, based on their dense embeddings.
+    Retrieves documents from the `PgvectorDocumentStore`, based on their dense embeddings.
 
-    Needs to be connected to the PgvectorDocumentStore.
+    Example usage:
+    ```python
+    from haystack.document_stores import DuplicatePolicy
+    from haystack import Document, Pipeline
+    from haystack.components.embedders import SentenceTransformersTextEmbedder, SentenceTransformersDocumentEmbedder
+
+    from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
+    from haystack_integrations.components.retrievers.pgvector import PgvectorEmbeddingRetriever
+
+    # Set an environment variable `PG_CONN_STR` with the connection string to your PostgreSQL database.
+    # e.g., "postgresql://USER:PASSWORD@HOST:PORT/DB_NAME"
+
+    document_store = PgvectorDocumentStore(
+        embedding_dimension=768,
+        vector_function="cosine_similarity",
+        recreate_table=True,
+    )
+
+    documents = [Document(content="There are over 7,000 languages spoken around the world today."),
+                 Document(content="Elephants have been observed to behave in a way that indicates..."),
+                 Document(content="In certain places, you can witness the phenomenon of bioluminescent waves.")]
+
+    document_embedder = SentenceTransformersDocumentEmbedder()
+    document_embedder.warm_up()
+    documents_with_embeddings = document_embedder.run(documents)
+
+    document_store.write_documents(documents_with_embeddings.get("documents"), policy=DuplicatePolicy.OVERWRITE)
+
+    query_pipeline = Pipeline()
+    query_pipeline.add_component("text_embedder", SentenceTransformersTextEmbedder())
+    query_pipeline.add_component("retriever", PgvectorEmbeddingRetriever(document_store=document_store))
+    query_pipeline.connect("text_embedder.embedding", "retriever.query_embedding")
+
+    query = "How many languages are there?"
+
+    res = query_pipeline.run({"text_embedder": {"text": query}})
+
+    assert res['retriever']['documents'][0].content == "There are over 7,000 languages spoken around the world today."
+    ```
     """
 
     def __init__(
@@ -26,23 +64,20 @@ class PgvectorEmbeddingRetriever:
         vector_function: Optional[Literal["cosine_similarity", "inner_product", "l2_distance"]] = None,
     ):
         """
-        Create the PgvectorEmbeddingRetriever component.
-
-        :param document_store: An instance of PgvectorDocumentStore.
-        :param filters: Filters applied to the retrieved Documents. Defaults to None.
-        :param top_k: Maximum number of Documents to return, defaults to 10.
+        :param document_store: An instance of `PgvectorDocumentStore`.
+        :param filters: Filters applied to the retrieved Documents.
+        :param top_k: Maximum number of Documents to return.
         :param vector_function: The similarity function to use when searching for similar embeddings.
             Defaults to the one set in the `document_store` instance.
-            "cosine_similarity" and "inner_product" are similarity functions and
+            `"cosine_similarity"` and `"inner_product"` are similarity functions and
             higher scores indicate greater similarity between the documents.
-            "l2_distance" returns the straight-line distance between vectors,
+            `"l2_distance"` returns the straight-line distance between vectors,
             and the most similar documents are the ones with the smallest score.
-
-            Important: if the document store is using the "hnsw" search strategy, the vector function
+            **Important**: if the document store is using the `"hnsw"` search strategy, the vector function
             should match the one utilized during index creation to take advantage of the index.
-        :type vector_function: Literal["cosine_similarity", "inner_product", "l2_distance"]
 
-        :raises ValueError: If `document_store` is not an instance of PgvectorDocumentStore.
+        :raises ValueError: If `document_store` is not an instance of `PgvectorDocumentStore` or if `vector_function`
+            is not one of the valid options.
         """
         if not isinstance(document_store, PgvectorDocumentStore):
             msg = "document_store must be an instance of PgvectorDocumentStore"
@@ -58,6 +93,12 @@ class PgvectorEmbeddingRetriever:
         self.vector_function = vector_function or document_store.vector_function
 
     def to_dict(self) -> Dict[str, Any]:
+        """
+        Serializes the component to a dictionary.
+
+        :returns:
+            Dictionary with serialized data.
+        """
         return default_to_dict(
             self,
             filters=self.filters,
@@ -68,9 +109,16 @@ class PgvectorEmbeddingRetriever:
 
     @classmethod
     def from_dict(cls, data: Dict[str, Any]) -> "PgvectorEmbeddingRetriever":
-        data["init_parameters"]["document_store"] = default_from_dict(
-            PgvectorDocumentStore, data["init_parameters"]["document_store"]
-        )
+        """
+        Deserializes the component from a dictionary.
+
+        :param data:
+            Dictionary to deserialize from.
+        :returns:
+            Deserialized component.
+        """
+        doc_store_params = data["init_parameters"]["document_store"]
+        data["init_parameters"]["document_store"] = PgvectorDocumentStore.from_dict(doc_store_params)
         return default_from_dict(cls, data)
 
     @component.output_types(documents=List[Document])
@@ -82,14 +130,14 @@ class PgvectorEmbeddingRetriever:
         vector_function: Optional[Literal["cosine_similarity", "inner_product", "l2_distance"]] = None,
     ):
         """
-        Retrieve documents from the PgvectorDocumentStore, based on their embeddings.
+        Retrieve documents from the `PgvectorDocumentStore`, based on their embeddings.
 
         :param query_embedding: Embedding of the query.
         :param filters: Filters applied to the retrieved Documents.
         :param top_k: Maximum number of Documents to return.
         :param vector_function: The similarity function to use when searching for similar embeddings.
-        :type vector_function: Literal["cosine_similarity", "inner_product", "l2_distance"]
-        :return: List of Documents similar to `query_embedding`.
+
+        :returns: List of Documents similar to `query_embedding`.
         """
         filters = filters or self.filters
         top_k = top_k or self.top_k

pgvector_haystack-0.2.0/src/haystack_integrations/components/retrievers/pgvector/keyword_retriever.py

@@ -0,0 +1,123 @@
+# SPDX-FileCopyrightText: 2023-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+from typing import Any, Dict, List, Optional
+
+from haystack import component, default_from_dict, default_to_dict
+from haystack.dataclasses import Document
+from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
+
+
+@component
+class PgvectorKeywordRetriever:
+    """
+    Retrieve documents from the `PgvectorDocumentStore`, based on keywords.
+
+    To rank the documents, the `ts_rank_cd` function of PostgreSQL is used.
+    It considers how often the query terms appear in the document, how close together the terms are in the document,
+    and how important is the part of the document where they occur.
+    For more details, see
+    [Postgres documentation](https://www.postgresql.org/docs/current/textsearch-controls.html#TEXTSEARCH-RANKING).
+
+    Usage example:
+    ```python
+    from haystack.document_stores import DuplicatePolicy
+    from haystack import Document
+
+    from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
+    from haystack_integrations.components.retrievers.pgvector import PgvectorKeywordRetriever
+
+    # Set an environment variable `PG_CONN_STR` with the connection string to your PostgreSQL database.
+    # e.g., "postgresql://USER:PASSWORD@HOST:PORT/DB_NAME"
+
+    document_store = PgvectorDocumentStore(language="english", recreate_table=True)
+
+    documents = [Document(content="There are over 7,000 languages spoken around the world today."),
+        Document(content="Elephants have been observed to behave in a way that indicates..."),
+        Document(content="In certain places, you can witness the phenomenon of bioluminescent waves.")]
+
+    document_store.write_documents(documents_with_embeddings.get("documents"), policy=DuplicatePolicy.OVERWRITE)
+
+    retriever = PgvectorKeywordRetriever(document_store=document_store)
+
+    result = retriever.run(query="languages")
+
+    assert res['retriever']['documents'][0].content == "There are over 7,000 languages spoken around the world today."
+    """
+
+    def __init__(
+        self,
+        *,
+        document_store: PgvectorDocumentStore,
+        filters: Optional[Dict[str, Any]] = None,
+        top_k: int = 10,
+    ):
+        """
+        :param document_store: An instance of `PgvectorDocumentStore`.
+        :param filters: Filters applied to the retrieved Documents.
+        :param top_k: Maximum number of Documents to return.
+
+        :raises ValueError: If `document_store` is not an instance of `PgvectorDocumentStore`.
+        """
+        if not isinstance(document_store, PgvectorDocumentStore):
+            msg = "document_store must be an instance of PgvectorDocumentStore"
+            raise ValueError(msg)
+
+        self.document_store = document_store
+        self.filters = filters or {}
+        self.top_k = top_k
+
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Serializes the component to a dictionary.
+
+        :returns:
+            Dictionary with serialized data.
+        """
+        return default_to_dict(
+            self,
+            filters=self.filters,
+            top_k=self.top_k,
+            document_store=self.document_store.to_dict(),
+        )
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "PgvectorKeywordRetriever":
+        """
+        Deserializes the component from a dictionary.
+
+        :param data:
+            Dictionary to deserialize from.
+        :returns:
+            Deserialized component.
+        """
+        doc_store_params = data["init_parameters"]["document_store"]
+        data["init_parameters"]["document_store"] = PgvectorDocumentStore.from_dict(doc_store_params)
+        return default_from_dict(cls, data)
+
+    @component.output_types(documents=List[Document])
+    def run(
+        self,
+        query: str,
+        filters: Optional[Dict[str, Any]] = None,
+        top_k: Optional[int] = None,
+    ):
+        """
+        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.
+        :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 = filters or self.filters
+        top_k = top_k or self.top_k
+
+        docs = self.document_store._keyword_retrieval(
+            query=query,
+            filters=filters,
+            top_k=top_k,
+        )
+        return {"documents": docs}