diff --git a/.github/workflows/CI_readme_sync.yml b/.github/workflows/CI_readme_sync.yml index 27f225e74..c6204a9be 100644 --- a/.github/workflows/CI_readme_sync.yml +++ b/.github/workflows/CI_readme_sync.yml @@ -4,6 +4,7 @@ on: push: tags: - "**-v[0-9].[0-9]+.[0-9]+" + workflow_dispatch: # Activate this workflow manually inputs: tag: @@ -16,8 +17,30 @@ env: TAG: ${{ inputs.tag || github.ref_name }} jobs: + get-versions: + runs-on: ubuntu-latest + outputs: + versions: ${{ steps.version_finder.outputs.versions }} + steps: + - name: Get Haystack Docs versions + id: version_finder + run: | + curl -s https://dash.readme.com/api/v1/version --header 'authorization: Basic ${{ secrets.README_API_KEY }}' > out + VERSIONS=$(jq '[ .[] | select(.version | startswith("2."))| .version ]' out) + { + echo 'versions<> "$GITHUB_OUTPUT" + sync: runs-on: ubuntu-latest + needs: get-versions + strategy: + fail-fast: false + max-parallel: 1 + matrix: + hs-docs-version: ${{ fromJSON(needs.get-versions.outputs.versions) }} steps: - name: Checkout this repo uses: actions/checkout@v4 @@ -39,7 +62,7 @@ jobs: import os project_path = os.environ["TAG"].rsplit("-", maxsplit=1)[0] with open(os.environ['GITHUB_OUTPUT'], 'a') as f: - print(f'project_path={project_path}', file=f) + print(f'project_path={project_path}', file=f) - name: Generate docs working-directory: ${{ steps.pathfinder.outputs.project_path }} @@ -48,13 +71,16 @@ jobs: # from Readme.io as we need them to associate the slug # in config files with their id. README_API_KEY: ${{ secrets.README_API_KEY }} + # The same category has a different id on different readme docs versions. + # This is the docs version on readme that we'll use to get the category id. + PYDOC_TOOLS_HAYSTACK_DOC_VERSION: ${{ matrix.hs-docs-version }} run: | hatch run docs mkdir tmp find . -name "_readme_*.md" -exec cp "{}" tmp \; ls tmp - - name: Sync API docs + - name: Sync API docs with Haystack docs version ${{ matrix.hs-docs-version }} uses: readmeio/rdme@v8 with: - rdme: docs ${{ steps.pathfinder.outputs.project_path }}/tmp --key=${{ secrets.README_API_KEY }} --version=2.0 + rdme: docs ${{ steps.pathfinder.outputs.project_path }}/tmp --key=${{ secrets.README_API_KEY }} --version=${{ matrix.hs-docs-version }} diff --git a/.github/workflows/nvidia.yml b/.github/workflows/nvidia.yml index 316e509be..34e6a3c0e 100644 --- a/.github/workflows/nvidia.yml +++ b/.github/workflows/nvidia.yml @@ -22,6 +22,7 @@ env: PYTHONUNBUFFERED: "1" FORCE_COLOR: "1" NVIDIA_API_KEY: ${{ secrets.NVIDIA_API_KEY }} + NVIDIA_CATALOG_API_KEY: ${{ secrets.NVIDIA_CATALOG_API_KEY }} jobs: run: @@ -73,7 +74,7 @@ jobs: uses: ./.github/actions/send_failure with: title: | - core-integrations failure: + core-integrations failure: ${{ (steps.tests.conclusion == 'nightly-haystack-main') && 'nightly-haystack-main' || 'tests' }} - ${{ github.workflow }} api-key: ${{ secrets.CORE_DATADOG_API_KEY }} diff --git a/README.md b/README.md index d91e615d6..10b4b251c 100644 --- a/README.md +++ b/README.md @@ -51,5 +51,5 @@ Please check out our [Contribution Guidelines](CONTRIBUTING.md) for all the deta | [qdrant-haystack](integrations/qdrant/) | Document Store | [![PyPI - Version](https://img.shields.io/pypi/v/qdrant-haystack.svg?color=orange)](https://pypi.org/project/qdrant-haystack) | [![Test / qdrant](https://github.com/deepset-ai/haystack-core-integrations/actions/workflows/qdrant.yml/badge.svg)](https://github.com/deepset-ai/haystack-core-integrations/actions/workflows/qdrant.yml) | | [ragas-haystack](integrations/ragas/) | Evaluator | [![PyPI - Version](https://img.shields.io/pypi/v/ragas-haystack.svg)](https://pypi.org/project/ragas-haystack) | [![Test / ragas](https://github.com/deepset-ai/haystack-core-integrations/actions/workflows/ragas.yml/badge.svg)](https://github.com/deepset-ai/haystack-core-integrations/actions/workflows/ragas.yml) | | [unstructured-fileconverter-haystack](integrations/unstructured/) | File converter | [![PyPI - Version](https://img.shields.io/pypi/v/unstructured-fileconverter-haystack.svg)](https://pypi.org/project/unstructured-fileconverter-haystack) | [![Test / unstructured](https://github.com/deepset-ai/haystack-core-integrations/actions/workflows/unstructured.yml/badge.svg)](https://github.com/deepset-ai/haystack-core-integrations/actions/workflows/unstructured.yml) | -| [uptrain-haystack](https://github.com/deepset-ai/haystack-core-integrations/tree/staging/integrations/uptrain) | Evaluator | [![PyPI - Version](https://img.shields.io/pypi/v/uptrain-haystack.svg)](https://pypi.org/project/uptrain-haystack) | Staged | +| [uptrain-haystack](https://github.com/deepset-ai/haystack-core-integrations/tree/staging/integrations/uptrain) | Evaluator | [![PyPI - Version](https://img.shields.io/pypi/v/uptrain-haystack.svg)](https://pypi.org/project/uptrain-haystack) | [Staged](https://docs.haystack.deepset.ai/docs/breaking-change-policy#discontinuing-an-integration) | | [weaviate-haystack](integrations/weaviate/) | Document Store | [![PyPI - Version](https://img.shields.io/pypi/v/weaviate-haystack.svg)](https://pypi.org/project/weaviate-haystack) | [![Test / weaviate](https://github.com/deepset-ai/haystack-core-integrations/actions/workflows/weaviate.yml/badge.svg)](https://github.com/deepset-ai/haystack-core-integrations/actions/workflows/weaviate.yml) | diff --git a/integrations/amazon_bedrock/pydoc/config.yml b/integrations/amazon_bedrock/pydoc/config.yml index c719c7cfd..6cb05d6f3 100644 --- a/integrations/amazon_bedrock/pydoc/config.yml +++ b/integrations/amazon_bedrock/pydoc/config.yml @@ -20,7 +20,7 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: Amazon Bedrock integration for Haystack category_slug: integrations-api title: Amazon Bedrock diff --git a/integrations/amazon_sagemaker/pydoc/config.yml b/integrations/amazon_sagemaker/pydoc/config.yml index 20d51b25e..950e949f7 100644 --- a/integrations/amazon_sagemaker/pydoc/config.yml +++ b/integrations/amazon_sagemaker/pydoc/config.yml @@ -14,7 +14,7 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: Amazon Sagemaker integration for Haystack category_slug: integrations-api title: Amazon Sagemaker diff --git a/integrations/anthropic/pydoc/config.yml b/integrations/anthropic/pydoc/config.yml index 553dfcaef..9c1e39daf 100644 --- a/integrations/anthropic/pydoc/config.yml +++ b/integrations/anthropic/pydoc/config.yml @@ -15,12 +15,12 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: Anthropic integration for Haystack category_slug: integrations-api title: Anthropic slug: integrations-anthropic - order: 22 + order: 23 markdown: descriptive_class_title: false descriptive_module_title: true diff --git a/integrations/astra/pydoc/config.yml b/integrations/astra/pydoc/config.yml index 61fec0523..ed35427e6 100644 --- a/integrations/astra/pydoc/config.yml +++ b/integrations/astra/pydoc/config.yml @@ -16,7 +16,7 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: Astra integration for Haystack category_slug: integrations-api title: Astra diff --git a/integrations/astra/pyproject.toml b/integrations/astra/pyproject.toml index 3a36df238..129570788 100644 --- a/integrations/astra/pyproject.toml +++ b/integrations/astra/pyproject.toml @@ -80,12 +80,12 @@ dependencies = [ [tool.hatch.envs.lint.scripts] typing = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}" style = [ - "ruff {args:.}", + "ruff check {args:.}", "black --check --diff {args:.}", ] fmt = [ "black {args:.}", - "ruff --fix {args:.}", + "ruff check --fix {args:.}", "style", ] all = [ @@ -104,7 +104,7 @@ skip-string-normalization = true [tool.ruff] target-version = "py38" line-length = 120 -select = [ +lint.select = [ "A", "ARG", "B", @@ -131,7 +131,7 @@ select = [ "W", "YTT", ] -ignore = [ +lint.ignore = [ # Allow non-abstract empty methods in abstract base classes "B027", # Allow boolean positional values in function calls, like `dict.get(... True)` @@ -141,19 +141,19 @@ ignore = [ # Ignore complexity "C901", "PLR0911", "PLR0912", "PLR0913", "PLR0915", ] -unfixable = [ +lint.unfixable = [ # Don't touch unused imports "F401", ] -exclude = ["example"] +lint.exclude = ["example"] -[tool.ruff.isort] +[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"] diff --git a/integrations/astra/tests/test_document_store.py b/integrations/astra/tests/test_document_store.py index 3650ffd61..c475e6ae2 100644 --- a/integrations/astra/tests/test_document_store.py +++ b/integrations/astra/tests/test_document_store.py @@ -14,7 +14,13 @@ from haystack_integrations.document_stores.astra import AstraDocumentStore -def test_namespace_init(): +@pytest.fixture +def mock_auth(monkeypatch): + monkeypatch.setenv("ASTRA_DB_API_ENDPOINT", "http://example.com") + monkeypatch.setenv("ASTRA_DB_APPLICATION_TOKEN", "test_token") + + +def test_namespace_init(mock_auth): # noqa with mock.patch("haystack_integrations.document_stores.astra.astra_client.AstraDB") as client: AstraDocumentStore() assert "namespace" in client.call_args.kwargs @@ -25,7 +31,7 @@ def test_namespace_init(): assert client.call_args.kwargs["namespace"] == "foo" -def test_to_dict(): +def test_to_dict(mock_auth): # noqa with mock.patch("haystack_integrations.document_stores.astra.astra_client.AstraDB"): ds = AstraDocumentStore() result = ds.to_dict() diff --git a/integrations/chroma/pydoc/config.yml b/integrations/chroma/pydoc/config.yml index c28902080..1e678b4cc 100644 --- a/integrations/chroma/pydoc/config.yml +++ b/integrations/chroma/pydoc/config.yml @@ -17,7 +17,7 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: Chroma integration for Haystack category_slug: integrations-api title: Chroma diff --git a/integrations/cohere/pydoc/config.yml b/integrations/cohere/pydoc/config.yml index 5d4e747f5..53c54b664 100644 --- a/integrations/cohere/pydoc/config.yml +++ b/integrations/cohere/pydoc/config.yml @@ -19,7 +19,7 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: Cohere integration for Haystack category_slug: integrations-api title: Cohere diff --git a/integrations/cohere/src/haystack_integrations/components/embedders/cohere/document_embedder.py b/integrations/cohere/src/haystack_integrations/components/embedders/cohere/document_embedder.py index f3e5d2d56..59a04cf3c 100644 --- a/integrations/cohere/src/haystack_integrations/components/embedders/cohere/document_embedder.py +++ b/integrations/cohere/src/haystack_integrations/components/embedders/cohere/document_embedder.py @@ -42,7 +42,6 @@ def __init__( api_base_url: str = "https://api.cohere.com", truncate: str = "END", use_async_client: bool = False, - max_retries: int = 3, timeout: int = 120, batch_size: int = 32, progress_bar: bool = True, @@ -67,7 +66,6 @@ def __init__( If "NONE" is selected, when the input exceeds the maximum input token length an error will be returned. :param use_async_client: flag to select the AsyncClient. It is recommended to use AsyncClient for applications with many concurrent calls. - :param max_retries: maximal number of retries for requests. :param timeout: request timeout in seconds. :param batch_size: number of Documents to encode at once. :param progress_bar: whether to show a progress bar or not. Can be helpful to disable in production deployments @@ -82,7 +80,6 @@ def __init__( self.api_base_url = api_base_url self.truncate = truncate self.use_async_client = use_async_client - self.max_retries = max_retries self.timeout = timeout self.batch_size = batch_size self.progress_bar = progress_bar @@ -104,7 +101,6 @@ def to_dict(self) -> Dict[str, Any]: api_base_url=self.api_base_url, truncate=self.truncate, use_async_client=self.use_async_client, - max_retries=self.max_retries, timeout=self.timeout, batch_size=self.batch_size, progress_bar=self.progress_bar, @@ -170,7 +166,6 @@ def run(self, documents: List[Document]): cohere_client = AsyncClient( api_key, base_url=self.api_base_url, - max_retries=self.max_retries, timeout=self.timeout, client_name="haystack", ) @@ -181,7 +176,6 @@ def run(self, documents: List[Document]): cohere_client = Client( api_key, base_url=self.api_base_url, - max_retries=self.max_retries, timeout=self.timeout, client_name="haystack", ) diff --git a/integrations/cohere/src/haystack_integrations/components/embedders/cohere/text_embedder.py b/integrations/cohere/src/haystack_integrations/components/embedders/cohere/text_embedder.py index 0f9ba7b28..80ede51bf 100644 --- a/integrations/cohere/src/haystack_integrations/components/embedders/cohere/text_embedder.py +++ b/integrations/cohere/src/haystack_integrations/components/embedders/cohere/text_embedder.py @@ -39,7 +39,6 @@ def __init__( api_base_url: str = "https://api.cohere.com", truncate: str = "END", use_async_client: bool = False, - max_retries: int = 3, timeout: int = 120, ): """ @@ -60,7 +59,6 @@ def __init__( If "NONE" is selected, when the input exceeds the maximum input token length an error will be returned. :param use_async_client: flag to select the AsyncClient. It is recommended to use AsyncClient for applications with many concurrent calls. - :param max_retries: maximum number of retries for requests. :param timeout: request timeout in seconds. """ @@ -70,7 +68,6 @@ def __init__( self.api_base_url = api_base_url self.truncate = truncate self.use_async_client = use_async_client - self.max_retries = max_retries self.timeout = timeout def to_dict(self) -> Dict[str, Any]: @@ -88,7 +85,6 @@ def to_dict(self) -> Dict[str, Any]: api_base_url=self.api_base_url, truncate=self.truncate, use_async_client=self.use_async_client, - max_retries=self.max_retries, timeout=self.timeout, ) @@ -132,7 +128,6 @@ def run(self, text: str): cohere_client = AsyncClient( api_key, base_url=self.api_base_url, - max_retries=self.max_retries, timeout=self.timeout, client_name="haystack", ) @@ -143,7 +138,6 @@ def run(self, text: str): cohere_client = Client( api_key, base_url=self.api_base_url, - max_retries=self.max_retries, timeout=self.timeout, client_name="haystack", ) diff --git a/integrations/cohere/src/haystack_integrations/components/embedders/cohere/utils.py b/integrations/cohere/src/haystack_integrations/components/embedders/cohere/utils.py index 631dadf4f..a5c20cb35 100644 --- a/integrations/cohere/src/haystack_integrations/components/embedders/cohere/utils.py +++ b/integrations/cohere/src/haystack_integrations/components/embedders/cohere/utils.py @@ -62,7 +62,7 @@ def get_response( desc="Calculating embeddings", ): batch = texts[i : i + batch_size] - response = cohere_client.embed(batch, model=model_name, input_type=input_type, truncate=truncate) + response = cohere_client.embed(texts=batch, model=model_name, input_type=input_type, truncate=truncate) for emb in response.embeddings: all_embeddings.append(emb) if response.meta is not None: diff --git a/integrations/cohere/src/haystack_integrations/components/generators/cohere/chat/chat_generator.py b/integrations/cohere/src/haystack_integrations/components/generators/cohere/chat/chat_generator.py index fbe17c9e3..d5cc640db 100644 --- a/integrations/cohere/src/haystack_integrations/components/generators/cohere/chat/chat_generator.py +++ b/integrations/cohere/src/haystack_integrations/components/generators/cohere/chat/chat_generator.py @@ -163,21 +163,32 @@ def run(self, messages: List[ChatMessage], generation_kwargs: Optional[Dict[str, chat_history=chat_history, **generation_kwargs, ) - for chunk in response: - if chunk.event_type == "text-generation": - stream_chunk = self._build_chunk(chunk) + + response_text = "" + finish_response = None + for event in response: + if event.event_type == "text-generation": + stream_chunk = self._build_chunk(event) self.streaming_callback(stream_chunk) - chat_message = ChatMessage.from_assistant(content=response.texts) - chat_message.meta.update( - { - "model": self.model, - "usage": response.token_count, - "index": 0, - "finish_reason": response.finish_reason, - "documents": response.documents, - "citations": response.citations, - } - ) + response_text += event.text + elif event.event_type == "stream-end": + finish_response = event.response + chat_message = ChatMessage.from_assistant(content=response_text) + + if finish_response and finish_response.meta: + if finish_response.meta.billed_units: + tokens_in = finish_response.meta.billed_units.input_tokens or -1 + tokens_out = finish_response.meta.billed_units.output_tokens or -1 + chat_message.meta["usage"] = tokens_in + tokens_out + chat_message.meta.update( + { + "model": self.model, + "index": 0, + "finish_reason": finish_response.finish_reason, + "documents": finish_response.documents, + "citations": finish_response.citations, + } + ) else: response = self.client.chat( message=messages[-1].content, @@ -195,7 +206,7 @@ def _build_chunk(self, chunk) -> StreamingChunk: :param choice: The choice returned by the OpenAI API. :returns: The StreamingChunk. """ - chat_message = StreamingChunk(content=chunk.text, meta={"index": chunk.index, "event_type": chunk.event_type}) + chat_message = StreamingChunk(content=chunk.text, meta={"event_type": chunk.event_type}) return chat_message def _build_message(self, cohere_response): @@ -206,10 +217,12 @@ def _build_message(self, cohere_response): """ content = cohere_response.text message = ChatMessage.from_assistant(content=content) + + total_tokens = cohere_response.meta.billed_units.input_tokens + cohere_response.meta.billed_units.output_tokens message.meta.update( { "model": self.model, - "usage": cohere_response.token_count, + "usage": total_tokens, "index": 0, "finish_reason": None, "documents": cohere_response.documents, diff --git a/integrations/cohere/src/haystack_integrations/components/generators/cohere/generator.py b/integrations/cohere/src/haystack_integrations/components/generators/cohere/generator.py index 2b6a2dd48..d87d3035a 100644 --- a/integrations/cohere/src/haystack_integrations/components/generators/cohere/generator.py +++ b/integrations/cohere/src/haystack_integrations/components/generators/cohere/generator.py @@ -2,14 +2,14 @@ # # SPDX-License-Identifier: Apache-2.0 import logging -from typing import Any, Callable, Dict, List, Optional, cast +from typing import Any, Callable, Dict, List, Optional from haystack import component, default_from_dict, default_to_dict from haystack.components.generators.utils import deserialize_callback_handler, serialize_callback_handler from haystack.dataclasses import StreamingChunk from haystack.utils import Secret, deserialize_secrets_inplace -from cohere import Client, Generation +from cohere import Client logger = logging.getLogger(__name__) @@ -75,6 +75,10 @@ def __init__( - 'logit_bias': Used to prevent the model from generating unwanted tokens or to incentivize it to include desired tokens. The format is {token_id: bias} where bias is a float between -10 and 10. """ + logger.warning( + "The 'generate' API is marked as Legacy and is no longer maintained by Cohere. " + "We recommend to use the CohereChatGenerator instead." + ) if not api_base_url: api_base_url = "https://api.cohere.com" @@ -132,20 +136,29 @@ def run(self, prompt: str): if self.streaming_callback: response = self.client.generate_stream(model=self.model, prompt=prompt, **self.model_parameters) metadata_dict: Dict[str, Any] = {} - for chunk in response: - stream_chunk = self._build_chunk(chunk) - self.streaming_callback(stream_chunk) - replies = response.texts - metadata_dict["finish_reason"] = response.finish_reason - metadata = [metadata_dict] - self._check_truncated_answers(metadata) - return {"replies": replies, "meta": metadata} + generated_text = "" + for event in response: + if event.event_type == "text-generation": + generated_text += event.text + stream_chunk = self._build_chunk(event) + self.streaming_callback(stream_chunk) + elif event.event_type == "stream-end": + metadata_dict["finish_reason"] = event.finish_reason + if event.finish_reason == "MAX_TOKENS": + logger.warning( + "Responses have been truncated before reaching a natural stopping point. " + "Increase the max_tokens parameter to allow for longer completions." + ) + return {"replies": [generated_text], "meta": [metadata_dict]} response = self.client.generate(model=self.model, prompt=prompt, **self.model_parameters) - metadata = [{"finish_reason": resp.finish_reason} for resp in cast(Generation, response)] - replies = [resp.text for resp in response] - self._check_truncated_answers(metadata) - return {"replies": replies, "meta": metadata} + metadata = {"finish_reason": response.finish_reason} + if response.finish_reason == "MAX_TOKENS": + logger.warning( + "Responses have been truncated before reaching a natural stopping point. " + "Increase the max_tokens parameter to allow for longer completions." + ) + return {"replies": [response.text], "meta": [metadata]} def _build_chunk(self, chunk) -> StreamingChunk: """ @@ -155,15 +168,3 @@ def _build_chunk(self, chunk) -> StreamingChunk: """ streaming_chunk = StreamingChunk(content=chunk.text, meta={"index": chunk.index}) return streaming_chunk - - def _check_truncated_answers(self, metadata: List[Dict[str, Any]]): - """ - Check the `finish_reason` returned with the Cohere response. - If the `finish_reason` is `MAX_TOKEN`, log a warning to the user. - :param metadata: The metadata returned by the Cohere API. - """ - if metadata[0]["finish_reason"] == "MAX_TOKENS": - logger.warning( - "Responses have been truncated before reaching a natural stopping point. " - "Increase the max_tokens parameter to allow for longer completions." - ) diff --git a/integrations/cohere/tests/test_cohere_chat_generator.py b/integrations/cohere/tests/test_cohere_chat_generator.py index b64f3e14f..b0dcf8e8a 100644 --- a/integrations/cohere/tests/test_cohere_chat_generator.py +++ b/integrations/cohere/tests/test_cohere_chat_generator.py @@ -1,8 +1,8 @@ import os -from unittest.mock import Mock, patch +from unittest.mock import Mock -import cohere import pytest +from cohere.core import ApiError from haystack.components.generators.utils import print_streaming_chunk from haystack.dataclasses import ChatMessage, ChatRole, StreamingChunk from haystack.utils import Secret @@ -11,51 +11,6 @@ pytestmark = pytest.mark.chat_generators -@pytest.fixture -def mock_chat_response(): - """ - Mock the Cohere API response and reuse it for tests - """ - with patch("cohere.Client.chat", autospec=True) as mock_chat_response: - # mimic the response from the Cohere API - - mock_response = Mock() - mock_response.text = "I'm fine, thanks." - mock_response.token_count = { - "prompt_tokens": 66, - "response_tokens": 78, - "total_tokens": 144, - "billed_tokens": 133, - } - mock_response.meta = { - "api_version": {"version": "1"}, - "billed_units": {"input_tokens": 55, "output_tokens": 78}, - } - mock_chat_response.return_value = mock_response - yield mock_chat_response - - -@pytest.fixture -def mock_chat_streaming_response(): - with patch("cohere.Client.chat_stream", autospec=True) as mock_chat_stream_response: - # mimic the response from the Cohere API - - mock_response = Mock() - mock_response.text = "I'm fine, thanks." - mock_response.token_count = { - "prompt_tokens": 66, - "response_tokens": 78, - "total_tokens": 144, - "billed_tokens": 133, - } - mock_response.meta = { - "api_version": {"version": "1"}, - "billed_units": {"input_tokens": 55, "output_tokens": 78}, - } - mock_chat_stream_response.return_value = mock_response - yield mock_chat_stream_response - - def streaming_chunk(text: str): """ Mock chunks of streaming responses from the Cohere API @@ -196,73 +151,11 @@ def test_from_dict_fail_wo_env_var(self, monkeypatch): with pytest.raises(ValueError): CohereChatGenerator.from_dict(data) - def test_run(self, chat_messages, mock_chat_response): # noqa: ARG002 - component = CohereChatGenerator(api_key=Secret.from_token("test-api-key")) - response = component.run(chat_messages) - - # check that the component returns the correct ChatMessage response - assert isinstance(response, dict) - assert "replies" in response - assert isinstance(response["replies"], list) - assert len(response["replies"]) == 1 - assert [isinstance(reply, ChatMessage) for reply in response["replies"]] - def test_message_to_dict(self, chat_messages): obj = CohereChatGenerator(api_key=Secret.from_token("test-api-key")) dictionary = [obj._message_to_dict(message) for message in chat_messages] assert dictionary == [{"user_name": "Chatbot", "text": "What's the capital of France"}] - def test_run_with_params(self, chat_messages, mock_chat_response): - component = CohereChatGenerator( - api_key=Secret.from_token("test-api-key"), generation_kwargs={"max_tokens": 10, "temperature": 0.5} - ) - response = component.run(chat_messages) - - # check that the component calls the Cohere API with the correct parameters - _, kwargs = mock_chat_response.call_args - assert kwargs["max_tokens"] == 10 - assert kwargs["temperature"] == 0.5 - - # check that the component returns the correct response - assert isinstance(response, dict) - assert "replies" in response - assert isinstance(response["replies"], list) - assert len(response["replies"]) == 1 - assert [isinstance(reply, ChatMessage) for reply in response["replies"]] - - def test_run_streaming(self, chat_messages, mock_chat_streaming_response): - streaming_call_count = 0 - - # Define the streaming callback function and assert that it is called with StreamingChunk objects - def streaming_callback_fn(chunk: StreamingChunk): - nonlocal streaming_call_count - streaming_call_count += 1 - assert isinstance(chunk, StreamingChunk) - - generator = CohereChatGenerator( - api_key=Secret.from_token("test-api-key"), streaming_callback=streaming_callback_fn - ) - - # Create a fake streamed response - # self needed here, don't remove - def mock_iter(self): # noqa: ARG001 - yield streaming_chunk("Hello") - yield streaming_chunk("How are you?") - - mock_response = Mock(**{"__iter__": mock_iter}) - mock_chat_streaming_response.return_value = mock_response - - response = generator.run(chat_messages) - - # Assert that the streaming callback was called twice - assert streaming_call_count == 2 - - # Assert that the response contains the generated replies - assert "replies" in response - assert isinstance(response["replies"], list) - assert len(response["replies"]) > 0 - assert [isinstance(reply, ChatMessage) for reply in response["replies"]] - @pytest.mark.skipif( not os.environ.get("COHERE_API_KEY", None) and not os.environ.get("CO_API_KEY", None), reason="Export an env var called COHERE_API_KEY/CO_API_KEY containing the Cohere API key to run this test.", @@ -283,7 +176,7 @@ def test_live_run(self): @pytest.mark.integration def test_live_run_wrong_model(self, chat_messages): component = CohereChatGenerator(model="something-obviously-wrong") - with pytest.raises(cohere.CohereAPIError): + with pytest.raises(ApiError): component.run(chat_messages) @pytest.mark.skipif( @@ -309,7 +202,7 @@ def __call__(self, chunk: StreamingChunk) -> None: assert len(results["replies"]) == 1 message: ChatMessage = results["replies"][0] - assert "Paris" in message.content[0] + assert "Paris" in message.content assert message.meta["finish_reason"] == "COMPLETE" @@ -353,7 +246,7 @@ def __call__(self, chunk: StreamingChunk) -> None: assert len(results["replies"]) == 1 message: ChatMessage = results["replies"][0] - assert "Paris" in message.content[0] + assert "Paris" in message.content assert message.meta["finish_reason"] == "COMPLETE" diff --git a/integrations/cohere/tests/test_cohere_generators.py b/integrations/cohere/tests/test_cohere_generators.py index fa97a03fb..5cff7cee0 100644 --- a/integrations/cohere/tests/test_cohere_generators.py +++ b/integrations/cohere/tests/test_cohere_generators.py @@ -4,6 +4,7 @@ import os import pytest +from cohere.core import ApiError from haystack.components.generators.utils import print_streaming_chunk from haystack.utils import Secret from haystack_integrations.components.generators.cohere import CohereGenerator @@ -119,15 +120,6 @@ def test_from_dict(self, monkeypatch): assert component.api_base_url == "test-base-url" assert component.model_parameters == {"max_tokens": 10, "some_test_param": "test-params"} - def test_check_truncated_answers(self, caplog): - component = CohereGenerator(api_key=Secret.from_token("test-api-key")) - meta = [{"finish_reason": "MAX_TOKENS"}] - component._check_truncated_answers(meta) - assert caplog.records[0].message == ( - "Responses have been truncated before reaching a natural stopping point. " - "Increase the max_tokens parameter to allow for longer completions." - ) - @pytest.mark.skipif( not os.environ.get("COHERE_API_KEY", None) and not os.environ.get("CO_API_KEY", None), reason="Export an env var called COHERE_API_KEY/CO_API_KEY containing the Cohere API key to run this test.", @@ -147,10 +139,8 @@ def test_cohere_generator_run(self): ) @pytest.mark.integration def test_cohere_generator_run_wrong_model(self): - import cohere - component = CohereGenerator(model="something-obviously-wrong") - with pytest.raises(cohere.CohereAPIError): + with pytest.raises(ApiError): component.run(prompt="What's the capital of France?") @pytest.mark.skipif( diff --git a/integrations/cohere/tests/test_document_embedder.py b/integrations/cohere/tests/test_document_embedder.py index 8891d4313..ffbf280e9 100644 --- a/integrations/cohere/tests/test_document_embedder.py +++ b/integrations/cohere/tests/test_document_embedder.py @@ -21,7 +21,6 @@ def test_init_default(self): assert embedder.api_base_url == COHERE_API_URL assert embedder.truncate == "END" assert embedder.use_async_client is False - assert embedder.max_retries == 3 assert embedder.timeout == 120 assert embedder.batch_size == 32 assert embedder.progress_bar is True @@ -36,7 +35,6 @@ def test_init_with_parameters(self): api_base_url="https://custom-api-base-url.com", truncate="START", use_async_client=True, - max_retries=5, timeout=60, batch_size=64, progress_bar=False, @@ -49,7 +47,6 @@ def test_init_with_parameters(self): assert embedder.api_base_url == "https://custom-api-base-url.com" assert embedder.truncate == "START" assert embedder.use_async_client is True - assert embedder.max_retries == 5 assert embedder.timeout == 60 assert embedder.batch_size == 64 assert embedder.progress_bar is False @@ -68,7 +65,6 @@ def test_to_dict(self): "api_base_url": COHERE_API_URL, "truncate": "END", "use_async_client": False, - "max_retries": 3, "timeout": 120, "batch_size": 32, "progress_bar": True, @@ -85,7 +81,6 @@ def test_to_dict_with_custom_init_parameters(self): api_base_url="https://custom-api-base-url.com", truncate="START", use_async_client=True, - max_retries=5, timeout=60, batch_size=64, progress_bar=False, @@ -102,7 +97,6 @@ def test_to_dict_with_custom_init_parameters(self): "api_base_url": "https://custom-api-base-url.com", "truncate": "START", "use_async_client": True, - "max_retries": 5, "timeout": 60, "batch_size": 64, "progress_bar": False, diff --git a/integrations/cohere/tests/test_text_embedder.py b/integrations/cohere/tests/test_text_embedder.py index 3a579a849..b4f3e234c 100644 --- a/integrations/cohere/tests/test_text_embedder.py +++ b/integrations/cohere/tests/test_text_embedder.py @@ -24,7 +24,6 @@ def test_init_default(self): assert embedder.api_base_url == COHERE_API_URL assert embedder.truncate == "END" assert embedder.use_async_client is False - assert embedder.max_retries == 3 assert embedder.timeout == 120 def test_init_with_parameters(self): @@ -38,7 +37,6 @@ def test_init_with_parameters(self): api_base_url="https://custom-api-base-url.com", truncate="START", use_async_client=True, - max_retries=5, timeout=60, ) assert embedder.api_key == Secret.from_token("test-api-key") @@ -47,7 +45,6 @@ def test_init_with_parameters(self): assert embedder.api_base_url == "https://custom-api-base-url.com" assert embedder.truncate == "START" assert embedder.use_async_client is True - assert embedder.max_retries == 5 assert embedder.timeout == 60 def test_to_dict(self): @@ -65,7 +62,6 @@ def test_to_dict(self): "api_base_url": COHERE_API_URL, "truncate": "END", "use_async_client": False, - "max_retries": 3, "timeout": 120, }, } @@ -81,7 +77,6 @@ def test_to_dict_with_custom_init_parameters(self): api_base_url="https://custom-api-base-url.com", truncate="START", use_async_client=True, - max_retries=5, timeout=60, ) component_dict = embedder_component.to_dict() @@ -94,7 +89,6 @@ def test_to_dict_with_custom_init_parameters(self): "api_base_url": "https://custom-api-base-url.com", "truncate": "START", "use_async_client": True, - "max_retries": 5, "timeout": 60, }, } diff --git a/integrations/deepeval/pydoc/config.yml b/integrations/deepeval/pydoc/config.yml index affa23acd..b3372f42c 100644 --- a/integrations/deepeval/pydoc/config.yml +++ b/integrations/deepeval/pydoc/config.yml @@ -18,7 +18,7 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: DeepEval integration for Haystack category_slug: integrations-api title: DeepEval diff --git a/integrations/elasticsearch/pydoc/config.yml b/integrations/elasticsearch/pydoc/config.yml index 04e20f992..39ffb2e5f 100644 --- a/integrations/elasticsearch/pydoc/config.yml +++ b/integrations/elasticsearch/pydoc/config.yml @@ -17,7 +17,7 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: Elasticsearch integration for Haystack category_slug: integrations-api title: Elasticsearch diff --git a/integrations/fastembed/pydoc/config.yml b/integrations/fastembed/pydoc/config.yml index c8bd11762..aad50e52c 100644 --- a/integrations/fastembed/pydoc/config.yml +++ b/integrations/fastembed/pydoc/config.yml @@ -18,7 +18,7 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: FastEmbed integration for Haystack category_slug: integrations-api title: FastEmbed diff --git a/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_document_embedder.py b/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_document_embedder.py index ec0b918d9..3b0ea83f3 100644 --- a/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_document_embedder.py +++ b/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_document_embedder.py @@ -29,14 +29,18 @@ class FastembedDocumentEmbedder: # Text taken from PubMed QA Dataset (https://huggingface.co/datasets/pubmed_qa) document_list = [ Document( - content="Oxidative stress generated within inflammatory joints can produce autoimmune phenomena and joint destruction. Radical species with oxidative activity, including reactive nitrogen species, represent mediators of inflammation and cartilage damage.", + content=("Oxidative stress generated within inflammatory joints can produce autoimmune phenomena and joint " + "destruction. Radical species with oxidative activity, including reactive nitrogen species, " + "represent mediators of inflammation and cartilage damage."), meta={ "pubid": "25,445,628", "long_answer": "yes", }, ), Document( - content="Plasma levels of pancreatic polypeptide (PP) rise upon food intake. Although other pancreatic islet hormones, such as insulin and glucagon, have been extensively investigated, PP secretion and actions are still poorly understood.", + content=("Plasma levels of pancreatic polypeptide (PP) rise upon food intake. Although other pancreatic " + "islet hormones, such as insulin and glucagon, have been extensively investigated, PP secretion " + "and actions are still poorly understood."), meta={ "pubid": "25,445,712", "long_answer": "yes", @@ -49,7 +53,7 @@ class FastembedDocumentEmbedder: print(f"Document Embedding: {result['documents'][0].embedding}") print(f"Embedding Dimension: {len(result['documents'][0].embedding)}") ``` - """ # noqa: E501 + """ def __init__( self, diff --git a/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_document_embedder.py b/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_document_embedder.py index ed5a3208b..caac3c8d0 100644 --- a/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_document_embedder.py +++ b/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_document_embedder.py @@ -12,30 +12,31 @@ class FastembedSparseDocumentEmbedder: Usage example: ```python - # To use this component, install the "fastembed-haystack" package. - # pip install fastembed-haystack - from haystack_integrations.components.embedders.fastembed import FastembedSparseDocumentEmbedder from haystack.dataclasses import Document - doc_embedder = FastembedSparseDocumentEmbedder( + sparse_doc_embedder = FastembedSparseDocumentEmbedder( model="prithvida/Splade_PP_en_v1", batch_size=32, ) - doc_embedder.warm_up() + sparse_doc_embedder.warm_up() # Text taken from PubMed QA Dataset (https://huggingface.co/datasets/pubmed_qa) document_list = [ Document( - content="Oxidative stress generated within inflammatory joints can produce autoimmune phenomena and joint destruction. Radical species with oxidative activity, including reactive nitrogen species, represent mediators of inflammation and cartilage damage.", + content=("Oxidative stress generated within inflammatory joints can produce autoimmune phenomena and joint " + "destruction. Radical species with oxidative activity, including reactive nitrogen species, " + "represent mediators of inflammation and cartilage damage."), meta={ "pubid": "25,445,628", "long_answer": "yes", }, ), Document( - content="Plasma levels of pancreatic polypeptide (PP) rise upon food intake. Although other pancreatic islet hormones, such as insulin and glucagon, have been extensively investigated, PP secretion and actions are still poorly understood.", + content=("Plasma levels of pancreatic polypeptide (PP) rise upon food intake. Although other pancreatic " + "islet hormones, such as insulin and glucagon, have been extensively investigated, PP secretion " + "and actions are still poorly understood."), meta={ "pubid": "25,445,712", "long_answer": "yes", @@ -43,12 +44,12 @@ class FastembedSparseDocumentEmbedder: ), ] - result = doc_embedder.run(document_list) + result = sparse_doc_embedder.run(document_list) print(f"Document Text: {result['documents'][0].content}") - print(f"Document Embedding: {result['documents'][0].sparse_embedding}") - print(f"Embedding Dimension: {len(result['documents'][0].sparse_embedding)}") + print(f"Document Sparse Embedding: {result['documents'][0].sparse_embedding}") + print(f"Sparse Embedding Dimension: {len(result['documents'][0].sparse_embedding)}") ``` - """ # noqa: E501 + """ def __init__( self, diff --git a/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_text_embedder.py b/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_text_embedder.py index b31677785..345df3f69 100644 --- a/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_text_embedder.py +++ b/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_text_embedder.py @@ -13,28 +13,25 @@ class FastembedSparseTextEmbedder: Usage example: ```python - # To use this component, install the "fastembed-haystack" package. - # pip install fastembed-haystack - from haystack_integrations.components.embedders.fastembed import FastembedSparseTextEmbedder - text = "It clearly says online this will work on a Mac OS system. The disk comes and it does not, only Windows. Do Not order this if you have a Mac!!" + text = ("It clearly says online this will work on a Mac OS system. " + "The disk comes and it does not, only Windows. Do Not order this if you have a Mac!!") - text_embedder = FastembedSparseTextEmbedder( + sparse_text_embedder = FastembedSparseTextEmbedder( model="prithvida/Splade_PP_en_v1" ) - text_embedder.warm_up() + sparse_text_embedder.warm_up() - embedding = text_embedder.run(text)["embedding"] + sparse_embedding = sparse_text_embedder.run(text)["sparse_embedding"] ``` - """ # noqa: E501 + """ def __init__( self, model: str = "prithvida/Splade_PP_en_v1", cache_dir: Optional[str] = None, threads: Optional[int] = None, - batch_size: int = 32, progress_bar: bool = True, parallel: Optional[int] = None, ): @@ -46,7 +43,6 @@ def __init__( Can be set using the `FASTEMBED_CACHE_PATH` env variable. Defaults to `fastembed_cache` in the system's temp directory. :param threads: The number of threads single onnxruntime session can use. Defaults to None. - :param batch_size: Number of strings to encode at once. :param progress_bar: If true, displays progress bar during embedding. :param parallel: If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets. @@ -57,7 +53,6 @@ def __init__( self.model_name = model self.cache_dir = cache_dir self.threads = threads - self.batch_size = batch_size self.progress_bar = progress_bar self.parallel = parallel @@ -73,7 +68,6 @@ def to_dict(self) -> Dict[str, Any]: model=self.model_name, cache_dir=self.cache_dir, threads=self.threads, - batch_size=self.batch_size, progress_bar=self.progress_bar, parallel=self.parallel, ) @@ -110,7 +104,6 @@ def run(self, text: str): embedding = self.embedding_backend.embed( [text], - batch_size=self.batch_size, show_progress_bar=self.progress_bar, parallel=self.parallel, )[0] diff --git a/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_text_embedder.py b/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_text_embedder.py index 9bc4475a5..246e1f866 100644 --- a/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_text_embedder.py +++ b/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_text_embedder.py @@ -12,12 +12,10 @@ class FastembedTextEmbedder: Usage example: ```python - # To use this component, install the "fastembed-haystack" package. - # pip install fastembed-haystack - from haystack_integrations.components.embedders.fastembed import FastembedTextEmbedder - text = "It clearly says online this will work on a Mac OS system. The disk comes and it does not, only Windows. Do Not order this if you have a Mac!!" + text = ("It clearly says online this will work on a Mac OS system. " + "The disk comes and it does not, only Windows. Do Not order this if you have a Mac!!") text_embedder = FastembedTextEmbedder( model="BAAI/bge-small-en-v1.5" @@ -26,7 +24,7 @@ class FastembedTextEmbedder: embedding = text_embedder.run(text)["embedding"] ``` - """ # noqa: E501 + """ def __init__( self, diff --git a/integrations/fastembed/tests/test_fastembed_sparse_text_embedder.py b/integrations/fastembed/tests/test_fastembed_sparse_text_embedder.py index 3751eea14..6f3f4e6cf 100644 --- a/integrations/fastembed/tests/test_fastembed_sparse_text_embedder.py +++ b/integrations/fastembed/tests/test_fastembed_sparse_text_embedder.py @@ -18,7 +18,6 @@ def test_init_default(self): assert embedder.model_name == "prithvida/Splade_PP_en_v1" assert embedder.cache_dir is None assert embedder.threads is None - assert embedder.batch_size == 32 assert embedder.progress_bar is True assert embedder.parallel is None @@ -30,14 +29,12 @@ def test_init_with_parameters(self): model="prithvida/Splade_PP_en_v1", cache_dir="fake_dir", threads=2, - batch_size=64, progress_bar=False, parallel=1, ) assert embedder.model_name == "prithvida/Splade_PP_en_v1" assert embedder.cache_dir == "fake_dir" assert embedder.threads == 2 - assert embedder.batch_size == 64 assert embedder.progress_bar is False assert embedder.parallel == 1 @@ -53,7 +50,6 @@ def test_to_dict(self): "model": "prithvida/Splade_PP_en_v1", "cache_dir": None, "threads": None, - "batch_size": 32, "progress_bar": True, "parallel": None, }, @@ -67,7 +63,6 @@ def test_to_dict_with_custom_init_parameters(self): model="prithvida/Splade_PP_en_v1", cache_dir="fake_dir", threads=2, - batch_size=64, progress_bar=False, parallel=1, ) @@ -78,7 +73,6 @@ def test_to_dict_with_custom_init_parameters(self): "model": "prithvida/Splade_PP_en_v1", "cache_dir": "fake_dir", "threads": 2, - "batch_size": 64, "progress_bar": False, "parallel": 1, }, @@ -94,7 +88,6 @@ def test_from_dict(self): "model": "prithvida/Splade_PP_en_v1", "cache_dir": None, "threads": None, - "batch_size": 32, "progress_bar": True, "parallel": None, }, @@ -103,7 +96,6 @@ def test_from_dict(self): assert embedder.model_name == "prithvida/Splade_PP_en_v1" assert embedder.cache_dir is None assert embedder.threads is None - assert embedder.batch_size == 32 assert embedder.progress_bar is True assert embedder.parallel is None @@ -117,7 +109,6 @@ def test_from_dict_with_custom_init_parameters(self): "model": "prithvida/Splade_PP_en_v1", "cache_dir": "fake_dir", "threads": 2, - "batch_size": 64, "progress_bar": False, "parallel": 1, }, @@ -126,7 +117,6 @@ def test_from_dict_with_custom_init_parameters(self): assert embedder.model_name == "prithvida/Splade_PP_en_v1" assert embedder.cache_dir == "fake_dir" assert embedder.threads == 2 - assert embedder.batch_size == 64 assert embedder.progress_bar is False assert embedder.parallel == 1 diff --git a/integrations/google_ai/pydoc/config.yml b/integrations/google_ai/pydoc/config.yml index 977a91fab..c2939a812 100644 --- a/integrations/google_ai/pydoc/config.yml +++ b/integrations/google_ai/pydoc/config.yml @@ -15,7 +15,7 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: Google AI integration for Haystack category_slug: integrations-api title: Google AI diff --git a/integrations/google_vertex/pydoc/config.yml b/integrations/google_vertex/pydoc/config.yml index bee97fdb8..6e23164b9 100644 --- a/integrations/google_vertex/pydoc/config.yml +++ b/integrations/google_vertex/pydoc/config.yml @@ -20,7 +20,7 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: Google Vertex integration for Haystack category_slug: integrations-api title: Google Vertex diff --git a/integrations/gradient/pydoc/config.yml b/integrations/gradient/pydoc/config.yml index a0ec5f72d..cb4b7a834 100644 --- a/integrations/gradient/pydoc/config.yml +++ b/integrations/gradient/pydoc/config.yml @@ -16,7 +16,7 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: Cohere integration for Haystack category_slug: integrations-api title: Gradient diff --git a/integrations/instructor_embedders/pydoc/config.yml b/integrations/instructor_embedders/pydoc/config.yml index dd5e38faa..a9ccc243c 100644 --- a/integrations/instructor_embedders/pydoc/config.yml +++ b/integrations/instructor_embedders/pydoc/config.yml @@ -16,7 +16,7 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: Instructor embedders integration for Haystack category_slug: integrations-api title: Instructor Embedders diff --git a/integrations/jina/pydoc/config.yml b/integrations/jina/pydoc/config.yml index 67788d26d..8c7a241f6 100644 --- a/integrations/jina/pydoc/config.yml +++ b/integrations/jina/pydoc/config.yml @@ -1,7 +1,7 @@ loaders: - type: haystack_pydoc_tools.loaders.CustomPythonLoader search_path: [../src] - modules: + modules: [ "haystack_integrations.components.embedders.jina.document_embedder", "haystack_integrations.components.embedders.jina.text_embedder", @@ -17,7 +17,7 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: Jina integration for Haystack category_slug: integrations-api title: Jina diff --git a/integrations/langfuse/pydoc/config.yml b/integrations/langfuse/pydoc/config.yml index dcfddfec8..c08bb35c3 100644 --- a/integrations/langfuse/pydoc/config.yml +++ b/integrations/langfuse/pydoc/config.yml @@ -15,12 +15,12 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: Langfuse integration for Haystack category_slug: integrations-api title: langfuse slug: integrations-langfuse - order: 135 + order: 136 markdown: descriptive_class_title: false classdef_code_block: false diff --git a/integrations/langfuse/src/haystack_integrations/tracing/langfuse/tracer.py b/integrations/langfuse/src/haystack_integrations/tracing/langfuse/tracer.py index 857081274..f4786b9bd 100644 --- a/integrations/langfuse/src/haystack_integrations/tracing/langfuse/tracer.py +++ b/integrations/langfuse/src/haystack_integrations/tracing/langfuse/tracer.py @@ -9,29 +9,50 @@ class LangfuseSpan(Span): + """ + Internal class representing a bridge between the Haystack span tracing API and Langfuse. + """ + def __init__(self, span: "Union[langfuse.client.StatefulSpanClient, langfuse.client.StatefulTraceClient]") -> None: + """ + Initialize a LangfuseSpan instance. + + :param span: The span instance managed by Langfuse. + """ self._span = span # locally cache tags self._data: Dict[str, Any] = {} def set_tag(self, key: str, value: Any) -> None: + """ + Set a generic tag for this span. + + :param key: The tag key. + :param value: The tag value. + """ coerced_value = tracing_utils.coerce_tag_value(value) self._span.update(metadata={key: coerced_value}) self._data[key] = value def set_content_tag(self, key: str, value: Any) -> None: + """ + Set a content-specific tag for this span. + + :param key: The content tag key. + :param value: The content tag value. + """ if not tracer.is_content_tracing_enabled: return if key.endswith(".input"): if "messages" in value: - messages = [self.to_openai_format(m) for m in value["messages"]] + messages = [m.to_openai_format() for m in value["messages"]] self._span.update(input=messages) else: self._span.update(input=value) elif key.endswith(".output"): if "replies" in value: if all(isinstance(r, ChatMessage) for r in value["replies"]): - replies = [self.to_openai_format(m) for m in value["replies"]] + replies = [m.to_openai_format() for m in value["replies"]] else: replies = value["replies"] self._span.update(output=replies) @@ -41,24 +62,33 @@ def set_content_tag(self, key: str, value: Any) -> None: self._data[key] = value def raw_span(self) -> Any: + """ + Return the underlying span instance. + + :return: The Langfuse span instance. + """ return self._span def get_correlation_data_for_logs(self) -> Dict[str, Any]: return {} - def to_openai_format(self, m: ChatMessage) -> Dict[str, Any]: - """ - Remove after haystack 2.0.1 has been released and use the `to_openai_format` method from the ChatMessage class - """ - msg = {"role": m.role.value, "content": m.content} - if m.name: - msg["name"] = m.name - - return msg - class LangfuseTracer(Tracer): + """ + Internal class representing a bridge between the Haystack tracer and Langfuse. + """ + def __init__(self, tracer: "langfuse.Langfuse", name: str = "Haystack", public: bool = False) -> None: + """ + Initialize a LangfuseTracer instance. + + :param tracer: The Langfuse tracer instance. + :param name: The name of the pipeline or component. This name will be used to identify the tracing run on the + Langfuse dashboard. + :param public: Whether the tracing data should be public or private. If set to `True`, the tracing data will + be publicly accessible to anyone with the tracing URL. If set to `False`, the tracing data will be private + and only accessible to the Langfuse account owner. + """ self._tracer = tracer self._context: list[LangfuseSpan] = [] self._name = name @@ -66,6 +96,12 @@ def __init__(self, tracer: "langfuse.Langfuse", name: str = "Haystack", public: @contextlib.contextmanager def trace(self, operation_name: str, tags: Optional[Dict[str, Any]] = None) -> Iterator[Span]: + """ + Start and manage a new trace span. + :param operation_name: The name of the operation. + :param tags: A dictionary of tags to attach to the span. + :return: A context manager yielding the span. + """ tags = tags or {} span_name = tags.get("haystack.component.name", operation_name) @@ -104,10 +140,19 @@ def trace(self, operation_name: str, tags: Optional[Dict[str, Any]] = None) -> I self._tracer.flush() def current_span(self) -> Span: + """ + Return the currently active span. + + :return: The currently active span. + """ if not self._context: # The root span has to be a trace self._context.append(LangfuseSpan(self._tracer.trace(name=self._name, public=self._public))) return self._context[-1] def get_trace_url(self) -> str: + """ + Return the URL to the tracing data. + :return: The URL to the tracing data. + """ return self._tracer.get_trace_url() diff --git a/integrations/langfuse/tests/test_langfuse_span.py b/integrations/langfuse/tests/test_langfuse_span.py new file mode 100644 index 000000000..a5a5f2c13 --- /dev/null +++ b/integrations/langfuse/tests/test_langfuse_span.py @@ -0,0 +1,65 @@ +import os + +os.environ["HAYSTACK_CONTENT_TRACING_ENABLED"] = "true" + +from unittest.mock import Mock +from haystack.dataclasses import ChatMessage +from haystack_integrations.tracing.langfuse.tracer import LangfuseSpan + + +class TestLangfuseSpan: + + # LangfuseSpan can be initialized with a span object + def test_initialized_with_span_object(self): + mock_span = Mock() + span = LangfuseSpan(mock_span) + assert span.raw_span() == mock_span + + # set_tag method can update metadata of the span object + def test_set_tag_updates_metadata(self): + mock_span = Mock() + span = LangfuseSpan(mock_span) + + span.set_tag("key", "value") + mock_span.update.assert_called_once_with(metadata={"key": "value"}) + assert span._data["key"] == "value" + + # set_content_tag method can update input and output of the span object + def test_set_content_tag_updates_input_and_output(self): + mock_span = Mock() + + span = LangfuseSpan(mock_span) + span.set_content_tag("input_key", "input_value") + assert span._data["input_key"] == "input_value" + + mock_span.reset_mock() + span.set_content_tag("output_key", "output_value") + assert span._data["output_key"] == "output_value" + + # set_content_tag method can update input and output of the span object with messages/replies + def test_set_content_tag_updates_input_and_output_with_messages(self): + mock_span = Mock() + + # test message input + span = LangfuseSpan(mock_span) + span.set_content_tag("key.input", {"messages": [ChatMessage.from_user("message")]}) + assert mock_span.update.call_count == 1 + # check we converted ChatMessage to OpenAI format + assert mock_span.update.call_args_list[0][1] == {"input": [{"role": "user", "content": "message"}]} + assert span._data["key.input"] == {"messages": [ChatMessage.from_user("message")]} + + # test replies ChatMessage list + mock_span.reset_mock() + span.set_content_tag("key.output", {"replies": [ChatMessage.from_system("reply")]}) + assert mock_span.update.call_count == 1 + # check we converted ChatMessage to OpenAI format + assert mock_span.update.call_args_list[0][1] == {"output": [{"role": "system", "content": "reply"}]} + assert span._data["key.output"] == {"replies": [ChatMessage.from_system("reply")]} + + # test replies string list + mock_span.reset_mock() + span.set_content_tag("key.output", {"replies": ["reply1", "reply2"]}) + assert mock_span.update.call_count == 1 + # check we handle properly string list replies + assert mock_span.update.call_args_list[0][1] == {"output": ["reply1", "reply2"]} + assert span._data["key.output"] == {"replies": ["reply1", "reply2"]} diff --git a/integrations/langfuse/tests/test_tracer.py b/integrations/langfuse/tests/test_tracer.py new file mode 100644 index 000000000..35a24c4b8 --- /dev/null +++ b/integrations/langfuse/tests/test_tracer.py @@ -0,0 +1,81 @@ +from unittest.mock import Mock, MagicMock, patch + +from haystack_integrations.tracing.langfuse.tracer import LangfuseTracer + + +class TestLangfuseTracer: + + # LangfuseTracer can be initialized with a Langfuse instance, a name and a boolean value for public. + def test_initialization(self): + langfuse_instance = Mock() + tracer = LangfuseTracer(tracer=langfuse_instance, name="Haystack", public=True) + assert tracer._tracer == langfuse_instance + assert tracer._context == [] + assert tracer._name == "Haystack" + assert tracer._public + + # check that the trace method is called on the tracer instance with the provided operation name and tags + # check that the span is added to the context and removed after the context manager exits + def test_create_new_span(self): + mock_raw_span = MagicMock() + mock_raw_span.operation_name = "operation_name" + mock_raw_span.metadata = {"tag1": "value1", "tag2": "value2"} + + with patch("haystack_integrations.tracing.langfuse.tracer.LangfuseSpan") as MockLangfuseSpan: + mock_span_instance = MockLangfuseSpan.return_value + mock_span_instance.raw_span.return_value = mock_raw_span + + mock_context_manager = MagicMock() + mock_context_manager.__enter__.return_value = mock_span_instance + + mock_tracer = MagicMock() + mock_tracer.trace.return_value = mock_context_manager + + tracer = LangfuseTracer(tracer=mock_tracer, name="Haystack", public=False) + + with tracer.trace("operation_name", tags={"tag1": "value1", "tag2": "value2"}) as span: + assert len(tracer._context) == 2, "The trace span should have been added to the the root context span" + assert span.raw_span().operation_name == "operation_name" + assert span.raw_span().metadata == {"tag1": "value1", "tag2": "value2"} + + assert len(tracer._context) == 1, "The trace span should have been popped, leaving root span in the context" + + # check that update method is called on the span instance with the provided key value pairs + def test_update_span_with_pipeline_input_output_data(self): + class MockTracer: + + def trace(self, name, **kwargs): + return MockSpan() + + def flush(self): + pass + + class MockSpan: + def __init__(self): + self._data = {} + self._span = self + self.operation_name = "operation_name" + + def raw_span(self): + return self + + def span(self, name=None): + # assert correct operation name passed to the span + assert name == "operation_name" + return self + + def update(self, **kwargs): + self._data.update(kwargs) + + def generation(self, name=None): + return self + + def end(self): + pass + + tracer = LangfuseTracer(tracer=MockTracer(), name="Haystack", public=False) + with tracer.trace(operation_name="operation_name", tags={"haystack.pipeline.input_data": "hello"}) as span: + assert span.raw_span()._data["metadata"] == {"haystack.pipeline.input_data": "hello"} + + with tracer.trace(operation_name="operation_name", tags={"haystack.pipeline.output_data": "bye"}) as span: + assert span.raw_span()._data["metadata"] == {"haystack.pipeline.output_data": "bye"} diff --git a/integrations/llama_cpp/pydoc/config.yml b/integrations/llama_cpp/pydoc/config.yml index 98068e672..a2b46c099 100644 --- a/integrations/llama_cpp/pydoc/config.yml +++ b/integrations/llama_cpp/pydoc/config.yml @@ -14,7 +14,7 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: Llama.cpp integration for Haystack category_slug: integrations-api title: Llama.cpp diff --git a/integrations/mistral/pydoc/config.yml b/integrations/mistral/pydoc/config.yml index 86ad5f1d0..c26843a54 100644 --- a/integrations/mistral/pydoc/config.yml +++ b/integrations/mistral/pydoc/config.yml @@ -16,7 +16,7 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: Mistral integration for Haystack category_slug: integrations-api title: Mistral diff --git a/integrations/mongodb_atlas/pydoc/config.yml b/integrations/mongodb_atlas/pydoc/config.yml index 85694c57f..a38b0a449 100644 --- a/integrations/mongodb_atlas/pydoc/config.yml +++ b/integrations/mongodb_atlas/pydoc/config.yml @@ -16,7 +16,7 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: MongoDB Atlas integration for Haystack category_slug: integrations-api title: MongoDB Atlas diff --git a/integrations/nvidia/pydoc/config.yml b/integrations/nvidia/pydoc/config.yml index 7e9811d25..80bb212c5 100644 --- a/integrations/nvidia/pydoc/config.yml +++ b/integrations/nvidia/pydoc/config.yml @@ -17,12 +17,12 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: Nvidia integration for Haystack category_slug: integrations-api title: Nvidia slug: integrations-nvidia - order: 50 + order: 165 markdown: descriptive_class_title: false classdef_code_block: false diff --git a/integrations/nvidia/pyproject.toml b/integrations/nvidia/pyproject.toml index 753f4f938..6d823407b 100644 --- a/integrations/nvidia/pyproject.toml +++ b/integrations/nvidia/pyproject.toml @@ -117,7 +117,6 @@ unfixable = [ # Don't touch unused imports "F401", ] -extend-exclude = ["tests", "example"] [tool.ruff.isort] known-first-party = ["src"] diff --git a/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/__init__.py b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/__init__.py index 588aca2e6..bc2d9372c 100644 --- a/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/__init__.py +++ b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/__init__.py @@ -1,7 +1,5 @@ from .document_embedder import NvidiaDocumentEmbedder from .text_embedder import NvidiaTextEmbedder +from .truncate import EmbeddingTruncateMode -__all__ = [ - "NvidiaDocumentEmbedder", - "NvidiaTextEmbedder", -] +__all__ = ["NvidiaDocumentEmbedder", "NvidiaTextEmbedder", "EmbeddingTruncateMode"] diff --git a/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/_nim_backend.py b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/_nim_backend.py index 27e0dbeac..26ba33c71 100644 --- a/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/_nim_backend.py +++ b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/_nim_backend.py @@ -1,6 +1,7 @@ from typing import Any, Dict, List, Optional, Tuple import requests +from haystack.utils import Secret from .backend import EmbedderBackend @@ -12,12 +13,17 @@ def __init__( self, model: str, api_url: str, + api_key: Optional[Secret] = None, model_kwargs: Optional[Dict[str, Any]] = None, ): headers = { "Content-Type": "application/json", "accept": "application/json", } + + if api_key: + headers["authorization"] = f"Bearer {api_key.resolve_value()}" + self.session = requests.Session() self.session.headers.update(headers) diff --git a/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/_nvcf_backend.py b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/_nvcf_backend.py index 7d4b07dca..65371de54 100644 --- a/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/_nvcf_backend.py +++ b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/_nvcf_backend.py @@ -1,3 +1,4 @@ +import warnings from dataclasses import asdict, dataclass from typing import Any, Dict, List, Literal, Optional, Tuple, Union @@ -17,6 +18,7 @@ def __init__( api_key: Secret, model_kwargs: Optional[Dict[str, Any]] = None, ): + warnings.warn("Nvidia NGC is deprecated, use Nvidia NIM instead.", DeprecationWarning, stacklevel=2) if not model.startswith("playground_"): model = f"playground_{model}" diff --git a/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/document_embedder.py b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/document_embedder.py index da181bd22..45680acce 100644 --- a/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/document_embedder.py +++ b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/document_embedder.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple, Union from haystack import Document, component, default_from_dict, default_to_dict from haystack.utils import Secret, deserialize_secrets_inplace @@ -7,6 +7,7 @@ from ._nim_backend import NimBackend from ._nvcf_backend import NvcfBackend from .backend import EmbedderBackend +from .truncate import EmbeddingTruncateMode @component @@ -41,6 +42,7 @@ def __init__( progress_bar: bool = True, meta_fields_to_embed: Optional[List[str]] = None, embedding_separator: str = "\n", + truncate: Optional[Union[EmbeddingTruncateMode, str]] = None, ): """ Create a NvidiaTextEmbedder component. @@ -64,6 +66,9 @@ def __init__( List of meta fields that should be embedded along with the Document text. :param embedding_separator: Separator used to concatenate the meta fields to the Document text. + :param truncate: + Specifies how inputs longer that the maximum token length should be truncated. + If None the behavior is model-dependent, see the official documentation for more information. """ self.api_key = api_key @@ -76,6 +81,10 @@ def __init__( self.meta_fields_to_embed = meta_fields_to_embed or [] self.embedding_separator = embedding_separator + if isinstance(truncate, str): + truncate = EmbeddingTruncateMode.from_str(truncate) + self.truncate = truncate + self.backend: Optional[EmbedderBackend] = None self._initialized = False @@ -93,7 +102,15 @@ def warm_up(self): self.backend = NvcfBackend(self.model, api_key=self.api_key, model_kwargs={"model": "passage"}) else: - self.backend = NimBackend(self.model, api_url=self.api_url, model_kwargs={"input_type": "passage"}) + model_kwargs = {"input_type": "passage"} + if self.truncate is not None: + model_kwargs["truncate"] = str(self.truncate) + self.backend = NimBackend( + self.model, + api_url=self.api_url, + api_key=self.api_key, + model_kwargs=model_kwargs, + ) self._initialized = True @@ -115,6 +132,7 @@ def to_dict(self) -> Dict[str, Any]: progress_bar=self.progress_bar, meta_fields_to_embed=self.meta_fields_to_embed, embedding_separator=self.embedding_separator, + truncate=str(self.truncate) if self.truncate is not None else None, ) @classmethod diff --git a/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/text_embedder.py b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/text_embedder.py index 6af5ba25f..b3ad4544e 100644 --- a/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/text_embedder.py +++ b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/text_embedder.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Union from haystack import component, default_from_dict, default_to_dict from haystack.utils import Secret, deserialize_secrets_inplace @@ -6,6 +6,7 @@ from ._nim_backend import NimBackend from ._nvcf_backend import NvcfBackend from .backend import EmbedderBackend +from .truncate import EmbeddingTruncateMode @component @@ -38,6 +39,7 @@ def __init__( api_url: Optional[str] = None, prefix: str = "", suffix: str = "", + truncate: Optional[Union[EmbeddingTruncateMode, str]] = None, ): """ Create a NvidiaTextEmbedder component. @@ -52,6 +54,9 @@ def __init__( A string to add to the beginning of each text. :param suffix: A string to add to the end of each text. + :param truncate: + Specifies how inputs longer that the maximum token length should be truncated. + If None the behavior is model-dependent, see the official documentation for more information. """ self.api_key = api_key @@ -60,6 +65,10 @@ def __init__( self.prefix = prefix self.suffix = suffix + if isinstance(truncate, str): + truncate = EmbeddingTruncateMode.from_str(truncate) + self.truncate = truncate + self.backend: Optional[EmbedderBackend] = None self._initialized = False @@ -77,7 +86,15 @@ def warm_up(self): self.backend = NvcfBackend(self.model, api_key=self.api_key, model_kwargs={"model": "query"}) else: - self.backend = NimBackend(self.model, api_url=self.api_url, model_kwargs={"input_type": "query"}) + model_kwargs = {"input_type": "query"} + if self.truncate is not None: + model_kwargs["truncate"] = str(self.truncate) + self.backend = NimBackend( + self.model, + api_url=self.api_url, + api_key=self.api_key, + model_kwargs=model_kwargs, + ) self._initialized = True @@ -95,6 +112,7 @@ def to_dict(self) -> Dict[str, Any]: api_url=self.api_url, prefix=self.prefix, suffix=self.suffix, + truncate=str(self.truncate) if self.truncate is not None else None, ) @classmethod diff --git a/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/truncate.py b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/truncate.py new file mode 100644 index 000000000..2c32eabb1 --- /dev/null +++ b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/truncate.py @@ -0,0 +1,32 @@ +from enum import Enum + + +class EmbeddingTruncateMode(Enum): + """ + Specifies how inputs to the NVIDIA embedding components are truncated. + If START, the input will be truncated from the start. + If END, the input will be truncated from the end. + """ + + START = "START" + END = "END" + + def __str__(self): + return self.value + + @classmethod + def from_str(cls, string: str) -> "EmbeddingTruncateMode": + """ + Create an truncate mode from a string. + + :param string: + String to convert. + :returns: + Truncate mode. + """ + enum_map = {e.value: e for e in EmbeddingTruncateMode} + opt_mode = enum_map.get(string) + if opt_mode is None: + msg = f"Unknown truncate mode '{string}'. Supported modes are: {list(enum_map.keys())}" + raise ValueError(msg) + return opt_mode diff --git a/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/_nim_backend.py b/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/_nim_backend.py index 499a60b78..879fe8a6b 100644 --- a/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/_nim_backend.py +++ b/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/_nim_backend.py @@ -1,6 +1,7 @@ from typing import Any, Dict, List, Optional, Tuple import requests +from haystack.utils import Secret from .backend import GeneratorBackend @@ -12,12 +13,17 @@ def __init__( self, model: str, api_url: str, + api_key: Optional[Secret] = None, model_kwargs: Optional[Dict[str, Any]] = None, ): headers = { "Content-Type": "application/json", "accept": "application/json", } + + if api_key: + headers["authorization"] = f"Bearer {api_key.resolve_value()}" + self.session = requests.Session() self.session.headers.update(headers) @@ -26,8 +32,9 @@ def __init__( self.model_kwargs = model_kwargs or {} def generate(self, prompt: str) -> Tuple[List[str], List[Dict[str, Any]]]: - # We're using the chat completion endpoint as the local containers don't support + # We're using the chat completion endpoint as the NIM API doesn't support # the /completions endpoint. So both the non-chat and chat generator will use this. + # This is the same for local containers and the cloud API. url = f"{self.api_url}/chat/completions" res = self.session.post( @@ -57,13 +64,17 @@ def generate(self, prompt: str) -> Tuple[List[str], List[Dict[str, Any]]]: replies.append(message["content"]) choice_meta = { "role": message["role"], - "finish_reason": choice["finish_reason"], "usage": { "prompt_tokens": completions["usage"]["prompt_tokens"], - "completion_tokens": completions["usage"]["completion_tokens"], "total_tokens": completions["usage"]["total_tokens"], }, } + # These fields could be null, the others will always be present + if "finish_reason" in choice: + choice_meta["finish_reason"] = choice["finish_reason"] + if "completion_tokens" in completions["usage"]: + choice_meta["usage"]["completion_tokens"] = completions["usage"]["completion_tokens"] + meta.append(choice_meta) return replies, meta diff --git a/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/_nvcf_backend.py b/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/_nvcf_backend.py index c0686c132..95d024fb8 100644 --- a/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/_nvcf_backend.py +++ b/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/_nvcf_backend.py @@ -1,3 +1,4 @@ +import warnings from dataclasses import asdict, dataclass from typing import Any, Dict, List, Optional, Tuple @@ -14,6 +15,7 @@ def __init__( api_key: Secret, model_kwargs: Optional[Dict[str, Any]] = None, ): + warnings.warn("Nvidia NGC is deprecated, use Nvidia NIM instead.", DeprecationWarning, stacklevel=2) if not model.startswith("playground_"): model = f"playground_{model}" diff --git a/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/generator.py b/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/generator.py index b6db399e6..d20478d93 100644 --- a/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/generator.py +++ b/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/generator.py @@ -85,6 +85,7 @@ def warm_up(self): self._backend = NimBackend( self._model, api_url=self._api_url, + api_key=self._api_key, model_kwargs=self._model_arguments, ) diff --git a/integrations/nvidia/tests/test_document_embedder.py b/integrations/nvidia/tests/test_document_embedder.py index 7ac89d5e2..28858dcf7 100644 --- a/integrations/nvidia/tests/test_document_embedder.py +++ b/integrations/nvidia/tests/test_document_embedder.py @@ -4,7 +4,7 @@ import pytest from haystack import Document from haystack.utils import Secret -from haystack_integrations.components.embedders.nvidia import NvidiaDocumentEmbedder +from haystack_integrations.components.embedders.nvidia import EmbeddingTruncateMode, NvidiaDocumentEmbedder class TestNvidiaDocumentEmbedder: @@ -64,6 +64,7 @@ def test_to_dict(self, monkeypatch): "progress_bar": True, "meta_fields_to_embed": [], "embedding_separator": "\n", + "truncate": None, }, } @@ -78,6 +79,7 @@ def test_to_dict_with_custom_init_parameters(self, monkeypatch): progress_bar=False, meta_fields_to_embed=["test_field"], embedding_separator=" | ", + truncate=EmbeddingTruncateMode.END, ) data = component.to_dict() assert data == { @@ -92,9 +94,38 @@ def test_to_dict_with_custom_init_parameters(self, monkeypatch): "progress_bar": False, "meta_fields_to_embed": ["test_field"], "embedding_separator": " | ", + "truncate": "END", }, } + def from_dict(self, monkeypatch): + monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key") + data = { + "type": "haystack_integrations.components.embedders.nvidia.document_embedder.NvidiaDocumentEmbedder", + "init_parameters": { + "api_key": {"env_vars": ["NVIDIA_API_KEY"], "strict": True, "type": "env_var"}, + "api_url": "https://example.com", + "model": "playground_nvolveqa_40k", + "prefix": "prefix", + "suffix": "suffix", + "batch_size": 10, + "progress_bar": False, + "meta_fields_to_embed": ["test_field"], + "embedding_separator": " | ", + "truncate": "START", + }, + } + component = NvidiaDocumentEmbedder.from_dict(data) + assert component.model == "nvolveqa_40k" + assert component.api_url is None + assert component.prefix == "prefix" + assert component.suffix == "suffix" + assert component.batch_size == 32 + assert component.progress_bar + assert component.meta_fields_to_embed == [] + assert component.embedding_separator == "\n" + assert component.truncate == EmbeddingTruncateMode.START + def test_prepare_texts_to_embed_w_metadata(self): documents = [ Document(content=f"document number {i}:\ncontent", meta={"meta_field": f"meta_value {i}"}) for i in range(5) @@ -355,3 +386,30 @@ def test_run_integration_with_nim_backend(self): for doc in docs_with_embeddings: assert isinstance(doc.embedding, list) assert isinstance(doc.embedding[0], float) + + @pytest.mark.skipif( + not os.environ.get("NVIDIA_CATALOG_API_KEY", None), + reason="Export an env var called NVIDIA_CATALOG_API_KEY containing the Nvidia API key to run this test.", + ) + @pytest.mark.integration + def test_run_integration_with_api_catalog(self): + embedder = NvidiaDocumentEmbedder( + model="NV-Embed-QA", + api_url="https://ai.api.nvidia.com/v1/retrieval/nvidia", + api_key=Secret.from_env_var("NVIDIA_CATALOG_API_KEY"), + ) + embedder.warm_up() + + docs = [ + Document(content="I love cheese", meta={"topic": "Cuisine"}), + Document(content="A transformer is a deep learning architecture", meta={"topic": "ML"}), + ] + + result = embedder.run(docs) + docs_with_embeddings = result["documents"] + + assert isinstance(docs_with_embeddings, list) + assert len(docs_with_embeddings) == len(docs) + for doc in docs_with_embeddings: + assert isinstance(doc.embedding, list) + assert isinstance(doc.embedding[0], float) diff --git a/integrations/nvidia/tests/test_generator.py b/integrations/nvidia/tests/test_generator.py index 9a157a9d1..102ef3508 100644 --- a/integrations/nvidia/tests/test_generator.py +++ b/integrations/nvidia/tests/test_generator.py @@ -202,3 +202,23 @@ def test_run_integration_with_nim_backend(self): assert result["replies"] assert result["meta"] + + @pytest.mark.skipif( + not os.environ.get("NVIDIA_CATALOG_API_KEY", None), + reason="Export an env var called NVIDIA_CATALOG_API_KEY containing the Nvidia API key to run this test.", + ) + @pytest.mark.integration + def test_run_integration_with_api_catalog(self): + generator = NvidiaGenerator( + model="meta/llama3-70b-instruct", + api_url="https://integrate.api.nvidia.com/v1", + api_key=Secret.from_env_var("NVIDIA_CATALOG_API_KEY"), + model_arguments={ + "temperature": 0.2, + }, + ) + generator.warm_up() + result = generator.run(prompt="What is the answer?") + + assert result["replies"] + assert result["meta"] diff --git a/integrations/nvidia/tests/test_text_embedder.py b/integrations/nvidia/tests/test_text_embedder.py index 39ee02206..b85ac39bf 100644 --- a/integrations/nvidia/tests/test_text_embedder.py +++ b/integrations/nvidia/tests/test_text_embedder.py @@ -3,7 +3,7 @@ import pytest from haystack.utils import Secret -from haystack_integrations.components.embedders.nvidia import NvidiaTextEmbedder +from haystack_integrations.components.embedders.nvidia import EmbeddingTruncateMode, NvidiaTextEmbedder class TestNvidiaTextEmbedder: @@ -46,6 +46,7 @@ def test_to_dict(self, monkeypatch): "model": "nvolveqa_40k", "prefix": "", "suffix": "", + "truncate": None, }, } @@ -55,6 +56,7 @@ def test_to_dict_with_custom_init_parameters(self, monkeypatch): model="nvolveqa_40k", prefix="prefix", suffix="suffix", + truncate=EmbeddingTruncateMode.START, ) data = component.to_dict() assert data == { @@ -65,9 +67,30 @@ def test_to_dict_with_custom_init_parameters(self, monkeypatch): "model": "nvolveqa_40k", "prefix": "prefix", "suffix": "suffix", + "truncate": "START", }, } + def from_dict(self, monkeypatch): + monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key") + data = { + "type": "haystack_integrations.components.embedders.nvidia.text_embedder.NvidiaTextEmbedder", + "init_parameters": { + "api_key": {"env_vars": ["NVIDIA_API_KEY"], "strict": True, "type": "env_var"}, + "api_url": None, + "model": "nvolveqa_40k", + "prefix": "prefix", + "suffix": "suffix", + "truncate": "START", + }, + } + component = NvidiaTextEmbedder.from_dict(data) + assert component.model == "nvolveqa_40k" + assert component.api_url is None + assert component.prefix == "prefix" + assert component.suffix == "suffix" + assert component.truncate == "START" + @patch("haystack_integrations.components.embedders.nvidia._nvcf_backend.NvidiaCloudFunctionsClient") def test_run(self, mock_client_class): embedder = NvidiaTextEmbedder( @@ -150,3 +173,23 @@ def test_run_integration_with_nim_backend(self): assert all(isinstance(x, float) for x in embedding) assert "usage" in meta + + @pytest.mark.skipif( + not os.environ.get("NVIDIA_CATALOG_API_KEY", None), + reason="Export an env var called NVIDIA_CATALOG_API_KEY containing the Nvidia API key to run this test.", + ) + @pytest.mark.integration + def test_run_integration_with_api_catalog(self): + embedder = NvidiaTextEmbedder( + model="NV-Embed-QA", + api_url="https://ai.api.nvidia.com/v1/retrieval/nvidia", + api_key=Secret.from_env_var("NVIDIA_CATALOG_API_KEY"), + ) + embedder.warm_up() + + result = embedder.run("A transformer is a deep learning architecture") + embedding = result["embedding"] + meta = result["meta"] + + assert all(isinstance(x, float) for x in embedding) + assert "usage" in meta diff --git a/integrations/ollama/pydoc/config.yml b/integrations/ollama/pydoc/config.yml index 4207ea997..e8f2ca6e5 100644 --- a/integrations/ollama/pydoc/config.yml +++ b/integrations/ollama/pydoc/config.yml @@ -17,7 +17,7 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: Ollama integration for Haystack category_slug: integrations-api title: Ollama diff --git a/integrations/opensearch/pydoc/config.yml b/integrations/opensearch/pydoc/config.yml index dfcb23b5f..7b2e20d83 100644 --- a/integrations/opensearch/pydoc/config.yml +++ b/integrations/opensearch/pydoc/config.yml @@ -17,7 +17,7 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: OpenSearch integration for Haystack category_slug: integrations-api title: OpenSearch diff --git a/integrations/optimum/pydoc/config.yml b/integrations/optimum/pydoc/config.yml index 8597b07ad..62edc9502 100644 --- a/integrations/optimum/pydoc/config.yml +++ b/integrations/optimum/pydoc/config.yml @@ -19,7 +19,7 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: Optimum integration for Haystack category_slug: integrations-api title: Optimum diff --git a/integrations/pgvector/examples/example.py b/integrations/pgvector/examples/embedding_retrieval.py similarity index 100% rename from integrations/pgvector/examples/example.py rename to integrations/pgvector/examples/embedding_retrieval.py diff --git a/integrations/pgvector/pydoc/config.yml b/integrations/pgvector/pydoc/config.yml index 1be4a1662..11be3aa4a 100644 --- a/integrations/pgvector/pydoc/config.yml +++ b/integrations/pgvector/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__"] @@ -15,7 +16,7 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: Pgvector integration for Haystack category_slug: integrations-api title: Pgvector diff --git a/integrations/pgvector/pyproject.toml b/integrations/pgvector/pyproject.toml index 39e2183cb..b440cf28e 100644 --- a/integrations/pgvector/pyproject.toml +++ b/integrations/pgvector/pyproject.toml @@ -174,6 +174,11 @@ exclude_lines = [ "if TYPE_CHECKING:", ] +[tool.pytest.ini_options] +markers = [ + "integration: integration tests" +] + [[tool.mypy.overrides]] module = [ diff --git a/integrations/pgvector/src/haystack_integrations/components/retrievers/pgvector/__init__.py b/integrations/pgvector/src/haystack_integrations/components/retrievers/pgvector/__init__.py index ec0cf0dc4..ea9fa8fe7 100644 --- a/integrations/pgvector/src/haystack_integrations/components/retrievers/pgvector/__init__.py +++ b/integrations/pgvector/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"] diff --git a/integrations/pgvector/src/haystack_integrations/components/retrievers/pgvector/embedding_retriever.py b/integrations/pgvector/src/haystack_integrations/components/retrievers/pgvector/embedding_retriever.py index 6085545cb..be894dcf7 100644 --- a/integrations/pgvector/src/haystack_integrations/components/retrievers/pgvector/embedding_retriever.py +++ b/integrations/pgvector/src/haystack_integrations/components/retrievers/pgvector/embedding_retriever.py @@ -64,7 +64,7 @@ def __init__( vector_function: Optional[Literal["cosine_similarity", "inner_product", "l2_distance"]] = None, ): """ - :param document_store: An instance of `PgvectorDocumentStore}. + :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. diff --git a/integrations/pgvector/src/haystack_integrations/components/retrievers/pgvector/keyword_retriever.py b/integrations/pgvector/src/haystack_integrations/components/retrievers/pgvector/keyword_retriever.py new file mode 100644 index 000000000..c09ac9bb5 --- /dev/null +++ b/integrations/pgvector/src/haystack_integrations/components/retrievers/pgvector/keyword_retriever.py @@ -0,0 +1,123 @@ +# SPDX-FileCopyrightText: 2023-present deepset GmbH +# +# 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} diff --git a/integrations/pgvector/src/haystack_integrations/document_stores/pgvector/document_store.py b/integrations/pgvector/src/haystack_integrations/document_stores/pgvector/document_store.py index da08a5f19..bb663f936 100644 --- a/integrations/pgvector/src/haystack_integrations/document_stores/pgvector/document_store.py +++ b/integrations/pgvector/src/haystack_integrations/document_stores/pgvector/document_store.py @@ -53,6 +53,12 @@ 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 +WHERE to_tsvector({language}, content) @@ query +""" + VALID_VECTOR_FUNCTIONS = ["cosine_similarity", "inner_product", "l2_distance"] VECTOR_FUNCTION_TO_POSTGRESQL_OPS = { @@ -65,6 +71,8 @@ HNSW_INDEX_NAME = "haystack_hnsw_index" +KEYWORD_INDEX_NAME = "haystack_keyword_index" + class PgvectorDocumentStore: """ @@ -76,6 +84,7 @@ def __init__( *, connection_string: Secret = Secret.from_env_var("PG_CONN_STR"), table_name: str = "haystack_documents", + language: str = "english", embedding_dimension: int = 768, vector_function: Literal["cosine_similarity", "inner_product", "l2_distance"] = "cosine_similarity", recreate_table: bool = False, @@ -92,6 +101,10 @@ def __init__( :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"` :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: + `SELECT cfgname FROM pg_ts_config;`. + More information can be found in this [StackOverflow answer](https://stackoverflow.com/a/39752553). :param embedding_dimension: The dimension of the embedding. :param vector_function: The similarity function to use when searching for similar embeddings. `"cosine_similarity"` and `"inner_product"` are similarity functions and @@ -116,7 +129,7 @@ def __init__( [pgvector documentation](https://github.com/pgvector/pgvector?tab=readme-ov-file#hnsw) :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) + [pgvector documentation](https://github.com/pgvector/pgvector?tab=readme-ov-file#hnsw). """ self.connection_string = connection_string @@ -131,6 +144,7 @@ def __init__( self.hnsw_recreate_index_if_exists = hnsw_recreate_index_if_exists self.hnsw_index_creation_kwargs = hnsw_index_creation_kwargs or {} self.hnsw_ef_search = hnsw_ef_search + self.language = language connection = connect(self.connection_string.resolve_value()) connection.autocommit = True @@ -146,6 +160,7 @@ def __init__( if recreate_table: self.delete_table() self._create_table_if_not_exists() + self._create_keyword_index_if_not_exists() if search_strategy == "hnsw": self._handle_hnsw() @@ -168,6 +183,7 @@ def to_dict(self) -> Dict[str, Any]: hnsw_recreate_index_if_exists=self.hnsw_recreate_index_if_exists, hnsw_index_creation_kwargs=self.hnsw_index_creation_kwargs, hnsw_ef_search=self.hnsw_ef_search, + language=self.language, ) @classmethod @@ -231,6 +247,29 @@ def delete_table(self): self._execute_sql(delete_sql, error_msg=f"Could not delete table {self.table_name} in PgvectorDocumentStore") + def _create_keyword_index_if_not_exists(self): + """ + Internal method to create the keyword index if not exists. + """ + index_exists = bool( + self._execute_sql( + "SELECT 1 FROM pg_indexes WHERE tablename = %s AND indexname = %s", + (self.table_name, 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))" + ).format( + index_name=Identifier(KEYWORD_INDEX_NAME), + table_name=Identifier(self.table_name), + language=SQLLiteral(self.language), + ) + + if not index_exists: + self._execute_sql(sql_create_index, error_msg="Could not create keyword index on table") + def _handle_hnsw(self): """ Internal method to handle the HNSW index creation. @@ -475,6 +514,54 @@ def delete_documents(self, document_ids: List[str]) -> None: self._execute_sql(delete_sql, error_msg="Could not delete documents from PgvectorDocumentStore") + def _keyword_retrieval( + self, + query: str, + *, + filters: Optional[Dict[str, Any]] = None, + top_k: int = 10, + ) -> List[Document]: + """ + Retrieves documents that are most similar to the query using a full-text search. + + This method is not meant to be part of the public interface of + `PgvectorDocumentStore` and it should not be called directly. + `PgvectorKeywordRetriever` uses this method directly and is the public interface for it. + + :returns: List of Documents that are most similar to `query` + """ + if not query: + msg = "query must be a non-empty string" + raise ValueError(msg) + + sql_select = SQL(KEYWORD_QUERY).format( + table_name=Identifier(self.table_name), + language=SQLLiteral(self.language), + query=SQLLiteral(query), + ) + + where_params = () + sql_where_clause = SQL("") + if filters: + sql_where_clause, where_params = _convert_filters_to_where_clause_and_params( + filters=filters, operator="AND" + ) + + sql_sort = SQL(" ORDER BY score DESC LIMIT {top_k}").format(top_k=SQLLiteral(top_k)) + + sql_query = sql_select + sql_where_clause + sql_sort + + result = self._execute_sql( + sql_query, + (query, *where_params), + error_msg="Could not retrieve documents from PgvectorDocumentStore.", + cursor=self._dict_cursor, + ) + + records = result.fetchall() + docs = self._from_pg_to_haystack_documents(records) + return docs + def _embedding_retrieval( self, query_embedding: List[float], @@ -489,6 +576,7 @@ def _embedding_retrieval( This method is not meant to be part of the public interface of `PgvectorDocumentStore` and it should not be called directly. `PgvectorEmbeddingRetriever` uses this method directly and is the public interface for it. + :returns: List of Documents that are most similar to `query_embedding` """ diff --git a/integrations/pgvector/src/haystack_integrations/document_stores/pgvector/filters.py b/integrations/pgvector/src/haystack_integrations/document_stores/pgvector/filters.py index daa90f502..d3604cfb3 100644 --- a/integrations/pgvector/src/haystack_integrations/document_stores/pgvector/filters.py +++ b/integrations/pgvector/src/haystack_integrations/document_stores/pgvector/filters.py @@ -3,7 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 from datetime import datetime from itertools import chain -from typing import Any, Dict, List +from typing import Any, Dict, List, Literal, Tuple from haystack.errors import FilterError from pandas import DataFrame @@ -22,7 +22,9 @@ NO_VALUE = "no_value" -def _convert_filters_to_where_clause_and_params(filters: Dict[str, Any]) -> tuple[SQL, tuple]: +def _convert_filters_to_where_clause_and_params( + filters: Dict[str, Any], operator: Literal["WHERE", "AND"] = "WHERE" +) -> Tuple[SQL, Tuple]: """ Convert Haystack filters to a WHERE clause and a tuple of params to query PostgreSQL. """ @@ -31,13 +33,13 @@ def _convert_filters_to_where_clause_and_params(filters: Dict[str, Any]) -> tupl else: query, values = _parse_logical_condition(filters) - where_clause = SQL(" WHERE ") + SQL(query) + where_clause = SQL(f" {operator} ") + SQL(query) params = tuple(value for value in values if value != NO_VALUE) return where_clause, params -def _parse_logical_condition(condition: Dict[str, Any]) -> tuple[str, List[Any]]: +def _parse_logical_condition(condition: Dict[str, Any]) -> Tuple[str, List[Any]]: if "operator" not in condition: msg = f"'operator' key missing in {condition}" raise FilterError(msg) @@ -77,7 +79,7 @@ def _parse_logical_condition(condition: Dict[str, Any]) -> tuple[str, List[Any]] return sql_query, values -def _parse_comparison_condition(condition: Dict[str, Any]) -> tuple[str, List[Any]]: +def _parse_comparison_condition(condition: Dict[str, Any]) -> Tuple[str, List[Any]]: field: str = condition["field"] if "operator" not in condition: msg = f"'operator' key missing in {condition}" @@ -132,20 +134,20 @@ def _treat_meta_field(field: str, value: Any) -> str: return field -def _equal(field: str, value: Any) -> tuple[str, Any]: +def _equal(field: str, value: Any) -> Tuple[str, Any]: if value is None: # NO_VALUE is a placeholder that will be removed in _convert_filters_to_where_clause_and_params return f"{field} IS NULL", NO_VALUE return f"{field} = %s", value -def _not_equal(field: str, value: Any) -> tuple[str, Any]: +def _not_equal(field: str, value: Any) -> Tuple[str, Any]: # we use IS DISTINCT FROM to correctly handle NULL values # (not handled by !=) return f"{field} IS DISTINCT FROM %s", value -def _greater_than(field: str, value: Any) -> tuple[str, Any]: +def _greater_than(field: str, value: Any) -> Tuple[str, Any]: if isinstance(value, str): try: datetime.fromisoformat(value) @@ -162,7 +164,7 @@ def _greater_than(field: str, value: Any) -> tuple[str, Any]: return f"{field} > %s", value -def _greater_than_equal(field: str, value: Any) -> tuple[str, Any]: +def _greater_than_equal(field: str, value: Any) -> Tuple[str, Any]: if isinstance(value, str): try: datetime.fromisoformat(value) @@ -179,7 +181,7 @@ def _greater_than_equal(field: str, value: Any) -> tuple[str, Any]: return f"{field} >= %s", value -def _less_than(field: str, value: Any) -> tuple[str, Any]: +def _less_than(field: str, value: Any) -> Tuple[str, Any]: if isinstance(value, str): try: datetime.fromisoformat(value) @@ -196,7 +198,7 @@ def _less_than(field: str, value: Any) -> tuple[str, Any]: return f"{field} < %s", value -def _less_than_equal(field: str, value: Any) -> tuple[str, Any]: +def _less_than_equal(field: str, value: Any) -> Tuple[str, Any]: if isinstance(value, str): try: datetime.fromisoformat(value) @@ -213,7 +215,7 @@ def _less_than_equal(field: str, value: Any) -> tuple[str, Any]: return f"{field} <= %s", value -def _not_in(field: str, value: Any) -> tuple[str, List]: +def _not_in(field: str, value: Any) -> Tuple[str, List]: if not isinstance(value, list): msg = f"{field}'s value must be a list when using 'not in' comparator in Pinecone" raise FilterError(msg) @@ -221,7 +223,7 @@ def _not_in(field: str, value: Any) -> tuple[str, List]: return f"{field} IS NULL OR {field} != ALL(%s)", [value] -def _in(field: str, value: Any) -> tuple[str, List]: +def _in(field: str, value: Any) -> Tuple[str, List]: if not isinstance(value, list): msg = f"{field}'s value must be a list when using 'in' comparator in Pinecone" raise FilterError(msg) diff --git a/integrations/pgvector/tests/conftest.py b/integrations/pgvector/tests/conftest.py index 94b35a04d..6547db9eb 100644 --- a/integrations/pgvector/tests/conftest.py +++ b/integrations/pgvector/tests/conftest.py @@ -36,10 +36,12 @@ def patches_for_unit_tests(): ) as mock_delete, patch( "haystack_integrations.document_stores.pgvector.document_store.PgvectorDocumentStore._create_table_if_not_exists" ) as mock_create, patch( + "haystack_integrations.document_stores.pgvector.document_store.PgvectorDocumentStore._create_keyword_index_if_not_exists" + ) 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_hnsw + yield mock_connect, mock_register, mock_delete, mock_create, mock_create_kw_index, mock_hnsw @pytest.fixture diff --git a/integrations/pgvector/tests/test_document_store.py b/integrations/pgvector/tests/test_document_store.py index bf5ccd5d4..6fd7e0dc0 100644 --- a/integrations/pgvector/tests/test_document_store.py +++ b/integrations/pgvector/tests/test_document_store.py @@ -89,6 +89,7 @@ def test_to_dict(monkeypatch): "recreate_table": True, "search_strategy": "hnsw", "hnsw_recreate_index_if_exists": True, + "language": "english", "hnsw_index_creation_kwargs": {"m": 32, "ef_construction": 128}, "hnsw_ef_search": 50, }, diff --git a/integrations/pgvector/tests/test_keyword_retrieval.py b/integrations/pgvector/tests/test_keyword_retrieval.py new file mode 100644 index 000000000..4a5614165 --- /dev/null +++ b/integrations/pgvector/tests/test_keyword_retrieval.py @@ -0,0 +1,50 @@ +import pytest +from haystack.dataclasses.document import Document +from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore + + +@pytest.mark.integration +class TestKeywordRetrieval: + def test_keyword_retrieval(self, document_store: PgvectorDocumentStore): + docs = [ + Document(content="The quick brown fox chased the dog", embedding=[0.1] * 768), + Document(content="The fox was brown", embedding=[0.1] * 768), + Document(content="The lazy dog", embedding=[0.1] * 768), + Document(content="fox fox fox", embedding=[0.1] * 768), + ] + + document_store.write_documents(docs) + + results = document_store._keyword_retrieval(query="fox", top_k=2) + + assert len(results) == 2 + for doc in results: + assert "fox" in doc.content + assert results[0].id == docs[-1].id + assert results[0].score > results[1].score + + def test_keyword_retrieval_with_filters(self, document_store: PgvectorDocumentStore): + docs = [ + Document( + content="The quick brown fox chased the dog", + embedding=([0.1] * 768), + meta={"meta_field": "right_value"}, + ), + Document(content="The fox was brown", embedding=([0.1] * 768), meta={"meta_field": "right_value"}), + Document(content="The lazy dog", embedding=([0.1] * 768), meta={"meta_field": "right_value"}), + Document(content="fox fox fox", embedding=([0.1] * 768), meta={"meta_field": "wrong_value"}), + ] + + document_store.write_documents(docs) + + filters = {"field": "meta.meta_field", "operator": "==", "value": "right_value"} + + results = document_store._keyword_retrieval(query="fox", top_k=3, filters=filters) + assert len(results) == 2 + for doc in results: + assert "fox" in doc.content + assert doc.meta["meta_field"] == "right_value" + + def test_empty_query(self, document_store: PgvectorDocumentStore): + with pytest.raises(ValueError): + document_store._keyword_retrieval(query="") diff --git a/integrations/pgvector/tests/test_retriever.py b/integrations/pgvector/tests/test_retrievers.py similarity index 53% rename from integrations/pgvector/tests/test_retriever.py rename to integrations/pgvector/tests/test_retrievers.py index 61381c24e..ef6f918ed 100644 --- a/integrations/pgvector/tests/test_retriever.py +++ b/integrations/pgvector/tests/test_retrievers.py @@ -6,11 +6,11 @@ import pytest from haystack.dataclasses import Document from haystack.utils.auth import EnvVarSecret -from haystack_integrations.components.retrievers.pgvector import PgvectorEmbeddingRetriever +from haystack_integrations.components.retrievers.pgvector import PgvectorEmbeddingRetriever, PgvectorKeywordRetriever from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore -class TestRetriever: +class TestEmbeddingRetriever: def test_init_default(self, mock_store): retriever = PgvectorEmbeddingRetriever(document_store=mock_store) assert retriever.document_store == mock_store @@ -46,6 +46,7 @@ def test_to_dict(self, mock_store): "recreate_table": True, "search_strategy": "exact_nearest_neighbor", "hnsw_recreate_index_if_exists": False, + "language": "english", "hnsw_index_creation_kwargs": {}, "hnsw_ef_search": None, }, @@ -114,3 +115,99 @@ def test_run(self): ) assert res == {"documents": [doc]} + + +class TestKeywordRetriever: + def test_init_default(self, mock_store): + retriever = PgvectorKeywordRetriever(document_store=mock_store) + assert retriever.document_store == mock_store + assert retriever.filters == {} + assert retriever.top_k == 10 + + def test_init(self, mock_store): + retriever = PgvectorKeywordRetriever(document_store=mock_store, filters={"field": "value"}, top_k=5) + assert retriever.document_store == mock_store + assert retriever.filters == {"field": "value"} + assert retriever.top_k == 5 + + def test_to_dict(self, mock_store): + retriever = PgvectorKeywordRetriever(document_store=mock_store, filters={"field": "value"}, top_k=5) + res = retriever.to_dict() + t = "haystack_integrations.components.retrievers.pgvector.keyword_retriever.PgvectorKeywordRetriever" + assert res == { + "type": t, + "init_parameters": { + "document_store": { + "type": "haystack_integrations.document_stores.pgvector.document_store.PgvectorDocumentStore", + "init_parameters": { + "connection_string": {"env_vars": ["PG_CONN_STR"], "strict": True, "type": "env_var"}, + "table_name": "haystack", + "embedding_dimension": 768, + "vector_function": "cosine_similarity", + "recreate_table": True, + "search_strategy": "exact_nearest_neighbor", + "hnsw_recreate_index_if_exists": False, + "language": "english", + "hnsw_index_creation_kwargs": {}, + "hnsw_ef_search": None, + }, + }, + "filters": {"field": "value"}, + "top_k": 5, + }, + } + + @pytest.mark.usefixtures("patches_for_unit_tests") + def test_from_dict(self, monkeypatch): + monkeypatch.setenv("PG_CONN_STR", "some-connection-string") + t = "haystack_integrations.components.retrievers.pgvector.keyword_retriever.PgvectorKeywordRetriever" + data = { + "type": t, + "init_parameters": { + "document_store": { + "type": "haystack_integrations.document_stores.pgvector.document_store.PgvectorDocumentStore", + "init_parameters": { + "connection_string": {"env_vars": ["PG_CONN_STR"], "strict": True, "type": "env_var"}, + "table_name": "haystack_test_to_dict", + "embedding_dimension": 768, + "vector_function": "cosine_similarity", + "recreate_table": True, + "search_strategy": "exact_nearest_neighbor", + "hnsw_recreate_index_if_exists": False, + "hnsw_index_creation_kwargs": {}, + "hnsw_ef_search": None, + }, + }, + "filters": {"field": "value"}, + "top_k": 5, + }, + } + + retriever = PgvectorKeywordRetriever.from_dict(data) + document_store = retriever.document_store + + assert isinstance(document_store, PgvectorDocumentStore) + assert isinstance(document_store.connection_string, EnvVarSecret) + assert document_store.table_name == "haystack_test_to_dict" + assert document_store.embedding_dimension == 768 + assert document_store.vector_function == "cosine_similarity" + assert document_store.recreate_table + assert document_store.search_strategy == "exact_nearest_neighbor" + assert not document_store.hnsw_recreate_index_if_exists + assert document_store.hnsw_index_creation_kwargs == {} + assert document_store.hnsw_ef_search is None + + assert retriever.filters == {"field": "value"} + assert retriever.top_k == 5 + + def test_run(self): + mock_store = Mock(spec=PgvectorDocumentStore) + doc = Document(content="Test doc", embedding=[0.1, 0.2]) + mock_store._keyword_retrieval.return_value = [doc] + + retriever = PgvectorKeywordRetriever(document_store=mock_store) + res = retriever.run(query="test query") + + mock_store._keyword_retrieval.assert_called_once_with(query="test query", filters={}, top_k=10) + + assert res == {"documents": [doc]} diff --git a/integrations/pinecone/pydoc/config.yml b/integrations/pinecone/pydoc/config.yml index 4265eeecc..f49ec4ab4 100644 --- a/integrations/pinecone/pydoc/config.yml +++ b/integrations/pinecone/pydoc/config.yml @@ -16,7 +16,7 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: Pinecone integration for Haystack category_slug: integrations-api title: Pinecone diff --git a/integrations/qdrant/pydoc/config.yml b/integrations/qdrant/pydoc/config.yml index 835eeb2e9..58ededdb5 100644 --- a/integrations/qdrant/pydoc/config.yml +++ b/integrations/qdrant/pydoc/config.yml @@ -17,7 +17,7 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: Qdrant integration for Haystack category_slug: integrations-api title: Qdrant diff --git a/integrations/ragas/pydoc/config.yml b/integrations/ragas/pydoc/config.yml index 3a8e843fe..97d8d808e 100644 --- a/integrations/ragas/pydoc/config.yml +++ b/integrations/ragas/pydoc/config.yml @@ -18,7 +18,7 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: Ragas integration for Haystack category_slug: integrations-api title: Ragas diff --git a/integrations/unstructured/pydoc/config.yml b/integrations/unstructured/pydoc/config.yml index 7179a2607..f2b4061a4 100644 --- a/integrations/unstructured/pydoc/config.yml +++ b/integrations/unstructured/pydoc/config.yml @@ -14,7 +14,7 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: Unstructured integration for Haystack category_slug: integrations-api title: Unstructured diff --git a/integrations/weaviate/pydoc/config.yml b/integrations/weaviate/pydoc/config.yml index e62b21591..ab585ebb7 100644 --- a/integrations/weaviate/pydoc/config.yml +++ b/integrations/weaviate/pydoc/config.yml @@ -18,7 +18,7 @@ processors: - type: smart - type: crossref renderer: - type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer + type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer excerpt: Weaviate integration for Haystack category_slug: integrations-api title: Weaviate diff --git a/integrations/weaviate/src/haystack_integrations/document_stores/weaviate/document_store.py b/integrations/weaviate/src/haystack_integrations/document_stores/weaviate/document_store.py index ec66e07c3..1b6e95155 100644 --- a/integrations/weaviate/src/haystack_integrations/document_stores/weaviate/document_store.py +++ b/integrations/weaviate/src/haystack_integrations/document_stores/weaviate/document_store.py @@ -166,7 +166,9 @@ def __init__( } else: # Set the class if not set - collection_settings["class"] = collection_settings.get("class", "default").capitalize() + _class_name = collection_settings.get("class", "Default") + _class_name = _class_name[0].upper() + _class_name[1:] + collection_settings["class"] = _class_name # Set the properties if they're not set collection_settings["properties"] = collection_settings.get("properties", DOCUMENT_COLLECTION_PROPERTIES) diff --git a/integrations/weaviate/tests/test_document_store.py b/integrations/weaviate/tests/test_document_store.py index cc76923f6..59a6ed2e3 100644 --- a/integrations/weaviate/tests/test_document_store.py +++ b/integrations/weaviate/tests/test_document_store.py @@ -664,3 +664,18 @@ def test_filter_documents_over_default_limit(self, document_store): document_store.write_documents(docs) with pytest.raises(DocumentStoreError): document_store.filter_documents({"field": "content", "operator": "==", "value": "This is some content"}) + + def test_schema_class_name_conversion_preserves_pascal_case(self): + collection_settings = {"class": "CaseDocument"} + doc_score = WeaviateDocumentStore( + url="http://localhost:8080", + collection_settings=collection_settings, + ) + assert doc_score._collection_settings["class"] == "CaseDocument" + + collection_settings = {"class": "lower_case_name"} + doc_score = WeaviateDocumentStore( + url="http://localhost:8080", + collection_settings=collection_settings, + ) + assert doc_score._collection_settings["class"] == "Lower_case_name"