From 0676fb4af162fcdcf1c921edf3459e2324c9ada3 Mon Sep 17 00:00:00 2001 From: Tobias Wochinger Date: Wed, 6 Mar 2024 10:06:24 +0100 Subject: [PATCH 1/8] docs: small consistency improvements (#536) * docs: backticks everywhere * chore: ignore `.idea` directory * docs: consistently use `returns` over `return` --- .gitignore | 1 + .../common/amazon_bedrock/utils.py | 4 +-- .../amazon_bedrock/document_embedder.py | 2 +- .../embedders/amazon_bedrock/text_embedder.py | 4 +-- .../generators/amazon_bedrock/adapters.py | 34 +++++++++---------- .../amazon_bedrock/chat/adapters.py | 34 +++++++++---------- .../amazon_bedrock/chat/chat_generator.py | 6 ++-- .../generators/amazon_bedrock/generator.py | 8 ++--- .../generators/amazon_bedrock/handlers.py | 6 ++-- .../generators/amazon_sagemaker/sagemaker.py | 4 +-- .../components/retrievers/astra/retriever.py | 2 +- .../components/retrievers/chroma/retriever.py | 8 ++--- .../document_stores/chroma/document_store.py | 2 +- .../embedders/cohere/text_embedder.py | 4 +-- .../generators/cohere/chat/chat_generator.py | 6 ++-- .../components/generators/cohere/generator.py | 6 ++-- .../fastembed/fastembed_document_embedder.py | 2 +- .../fastembed/fastembed_text_embedder.py | 2 +- .../google_vertex/image_generator.py | 2 +- .../google_vertex/question_answering.py | 2 +- .../google_vertex/text_generator.py | 6 ++-- .../generators/llama_cpp/generator.py | 4 +-- .../opensearch/document_store.py | 4 +-- .../document_stores/qdrant/document_store.py | 4 +-- .../document_stores/qdrant/filters.py | 2 +- .../converters/unstructured/converter.py | 2 +- .../text2speech/utils/text_to_speech.py | 4 +-- 27 files changed, 83 insertions(+), 82 deletions(-) diff --git a/.gitignore b/.gitignore index 8634bc259..c3a7cf863 100644 --- a/.gitignore +++ b/.gitignore @@ -134,3 +134,4 @@ dmypy.json # Docs generation artifacts _readme_*.md +.idea diff --git a/integrations/amazon_bedrock/src/haystack_integrations/common/amazon_bedrock/utils.py b/integrations/amazon_bedrock/src/haystack_integrations/common/amazon_bedrock/utils.py index e1683e3b3..3148818c1 100644 --- a/integrations/amazon_bedrock/src/haystack_integrations/common/amazon_bedrock/utils.py +++ b/integrations/amazon_bedrock/src/haystack_integrations/common/amazon_bedrock/utils.py @@ -34,7 +34,7 @@ def get_aws_session( :param kwargs: The kwargs passed down to the service client. Supported kwargs depend on the model chosen. See https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html. :raises AWSConfigurationError: If the provided AWS credentials are invalid. - :return: The created AWS session. + :returns: The created AWS session. """ try: return boto3.Session( @@ -54,7 +54,7 @@ def aws_configured(**kwargs) -> bool: """ Checks whether AWS configuration is provided. :param kwargs: The kwargs passed down to the generator. - :return: True if AWS configuration is provided, False otherwise. + :returns: True if AWS configuration is provided, False otherwise. """ aws_config_provided = any(key in kwargs for key in AWS_CONFIGURATION_KEYS) return aws_config_provided diff --git a/integrations/amazon_bedrock/src/haystack_integrations/components/embedders/amazon_bedrock/document_embedder.py b/integrations/amazon_bedrock/src/haystack_integrations/components/embedders/amazon_bedrock/document_embedder.py index 5a82821e3..8cf98cd45 100644 --- a/integrations/amazon_bedrock/src/haystack_integrations/components/embedders/amazon_bedrock/document_embedder.py +++ b/integrations/amazon_bedrock/src/haystack_integrations/components/embedders/amazon_bedrock/document_embedder.py @@ -235,7 +235,7 @@ def run(self, documents: List[Document]): def to_dict(self) -> Dict[str, Any]: """ Serialize this component to a dictionary. - :return: The serialized component as a dictionary. + :returns: The serialized component as a dictionary. """ return default_to_dict( self, diff --git a/integrations/amazon_bedrock/src/haystack_integrations/components/embedders/amazon_bedrock/text_embedder.py b/integrations/amazon_bedrock/src/haystack_integrations/components/embedders/amazon_bedrock/text_embedder.py index 8804702a0..ed6768737 100644 --- a/integrations/amazon_bedrock/src/haystack_integrations/components/embedders/amazon_bedrock/text_embedder.py +++ b/integrations/amazon_bedrock/src/haystack_integrations/components/embedders/amazon_bedrock/text_embedder.py @@ -154,7 +154,7 @@ def run(self, text: str): def to_dict(self) -> Dict[str, Any]: """ Serialize this component to a dictionary. - :return: The serialized component as a dictionary. + :returns: The serialized component as a dictionary. """ return default_to_dict( self, @@ -172,7 +172,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "AmazonBedrockTextEmbedder": """ Deserialize this component from a dictionary. :param data: The dictionary representation of this component. - :return: The deserialized component instance. + :returns: The deserialized component instance. """ deserialize_secrets_inplace( data["init_parameters"], diff --git a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/adapters.py b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/adapters.py index a1704ef13..f842f0ef5 100644 --- a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/adapters.py +++ b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/adapters.py @@ -25,7 +25,7 @@ def prepare_body(self, prompt: str, **inference_kwargs) -> Dict[str, Any]: :param prompt: The prompt to be sent to the model. :param inference_kwargs: Additional keyword arguments passed to the handler. - :return: A dictionary containing the body for the request. + :returns: A dictionary containing the body for the request. """ def get_responses(self, response_body: Dict[str, Any]) -> List[str]: @@ -33,7 +33,7 @@ def get_responses(self, response_body: Dict[str, Any]) -> List[str]: Extracts the responses from the Amazon Bedrock response. :param response_body: The response body from the Amazon Bedrock request. - :return: A list of responses. + :returns: A list of responses. """ completions = self._extract_completions_from_response(response_body) responses = [completion.lstrip() for completion in completions] @@ -45,7 +45,7 @@ def get_stream_responses(self, stream, stream_handler: TokenStreamingHandler) -> :param stream: The streaming response from the Amazon Bedrock request. :param stream_handler: The handler for the streaming response. - :return: A list of string responses. + :returns: A list of string responses. """ tokens: List[str] = [] for event in stream: @@ -64,7 +64,7 @@ def _get_params(self, inference_kwargs: Dict[str, Any], default_params: Dict[str Includes param if it's in kwargs or its default is not None (i.e. it is actually defined). :param inference_kwargs: The inference kwargs. :param default_params: The default params. - :return: A dictionary containing the merged params. + :returns: A dictionary containing the merged params. """ kwargs = self.model_kwargs.copy() kwargs.update(inference_kwargs) @@ -80,7 +80,7 @@ def _extract_completions_from_response(self, response_body: Dict[str, Any]) -> L Extracts the responses from the Amazon Bedrock response. :param response_body: The response body from the Amazon Bedrock request. - :return: A list of string responses. + :returns: A list of string responses. """ @abstractmethod @@ -89,7 +89,7 @@ def _extract_token_from_stream(self, chunk: Dict[str, Any]) -> str: Extracts the token from a streaming chunk. :param chunk: The streaming chunk. - :return: A string token. + :returns: A string token. """ @@ -121,7 +121,7 @@ def _extract_completions_from_response(self, response_body: Dict[str, Any]) -> L Extracts the responses from the Amazon Bedrock response. :param response_body: The response body from the Amazon Bedrock request. - :return: A list of string responses. + :returns: A list of string responses. """ return [response_body["completion"]] @@ -130,7 +130,7 @@ def _extract_token_from_stream(self, chunk: Dict[str, Any]) -> str: Extracts the token from a streaming chunk. :param chunk: The streaming chunk. - :return: A string token. + :returns: A string token. """ return chunk.get("completion", "") @@ -146,7 +146,7 @@ def prepare_body(self, prompt: str, **inference_kwargs) -> Dict[str, Any]: :param prompt: The prompt to be sent to the model. :param inference_kwargs: Additional keyword arguments passed to the handler. - :return: A dictionary containing the body for the request. + :returns: A dictionary containing the body for the request. """ default_params = { "max_tokens": self.max_length, @@ -170,7 +170,7 @@ def _extract_completions_from_response(self, response_body: Dict[str, Any]) -> L Extracts the responses from the Cohere Command model response. :param response_body: The response body from the Amazon Bedrock request. - :return: A list of string responses. + :returns: A list of string responses. """ responses = [generation["text"] for generation in response_body["generations"]] return responses @@ -180,7 +180,7 @@ def _extract_token_from_stream(self, chunk: Dict[str, Any]) -> str: Extracts the token from a streaming chunk. :param chunk: The streaming chunk. - :return: A string token. + :returns: A string token. """ return chunk.get("text", "") @@ -226,7 +226,7 @@ def prepare_body(self, prompt: str, **inference_kwargs) -> Dict[str, Any]: :param prompt: The prompt to be sent to the model. :param inference_kwargs: Additional keyword arguments passed to the handler. - :return: A dictionary containing the body for the request. + :returns: A dictionary containing the body for the request. """ default_params = { "maxTokenCount": self.max_length, @@ -244,7 +244,7 @@ def _extract_completions_from_response(self, response_body: Dict[str, Any]) -> L Extracts the responses from the Titan model response. :param response_body: The response body for Titan model response. - :return: A list of string responses. + :returns: A list of string responses. """ responses = [result["outputText"] for result in response_body["results"]] return responses @@ -254,7 +254,7 @@ def _extract_token_from_stream(self, chunk: Dict[str, Any]) -> str: Extracts the token from a streaming chunk. :param chunk: The streaming chunk. - :return: A string token. + :returns: A string token. """ return chunk.get("outputText", "") @@ -270,7 +270,7 @@ def prepare_body(self, prompt: str, **inference_kwargs) -> Dict[str, Any]: :param prompt: The prompt to be sent to the model. :param inference_kwargs: Additional keyword arguments passed to the handler. - :return: A dictionary containing the body for the request. + :returns: A dictionary containing the body for the request. """ default_params = { "max_gen_len": self.max_length, @@ -287,7 +287,7 @@ def _extract_completions_from_response(self, response_body: Dict[str, Any]) -> L Extracts the responses from the Llama2 model response. :param response_body: The response body from the Llama2 model request. - :return: A list of string responses. + :returns: A list of string responses. """ return [response_body["generation"]] @@ -296,6 +296,6 @@ def _extract_token_from_stream(self, chunk: Dict[str, Any]) -> str: Extracts the token from a streaming chunk. :param chunk: The streaming chunk. - :return: A string token. + :returns: A string token. """ return chunk.get("generation", "") diff --git a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/adapters.py b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/adapters.py index d5dc100f9..196a55743 100644 --- a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/adapters.py +++ b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/adapters.py @@ -34,7 +34,7 @@ def prepare_body(self, messages: List[ChatMessage], **inference_kwargs) -> Dict[ :param messages: The chat messages to package into the request. :param inference_kwargs: Additional inference kwargs to use. - :return: The prepared body. + :returns: The prepared body. """ def get_responses(self, response_body: Dict[str, Any]) -> List[ChatMessage]: @@ -42,7 +42,7 @@ def get_responses(self, response_body: Dict[str, Any]) -> List[ChatMessage]: Extracts the responses from the Amazon Bedrock response. :param response_body: The response body. - :return: The extracted responses. + :returns: The extracted responses. """ return self._extract_messages_from_response(self.response_body_message_key(), response_body) @@ -85,7 +85,7 @@ def _get_params(self, inference_kwargs: Dict[str, Any], default_params: Dict[str :param inference_kwargs: The inference kwargs to merge. :param default_params: The default params to start with. - :return: The merged params. + :returns: The merged params. """ # Start with a copy of default_params kwargs = default_params.copy() @@ -100,7 +100,7 @@ def _ensure_token_limit(self, prompt: str) -> str: """ Ensures that the prompt is within the token limit for the model. :param prompt: The prompt to check. - :return: The resized prompt. + :returns: The resized prompt. """ resize_info = self.check_prompt(prompt) if resize_info["prompt_length"] != resize_info["new_prompt_length"]: @@ -121,7 +121,7 @@ def check_prompt(self, prompt: str) -> Dict[str, Any]: Checks the prompt length and resizes it if necessary. If the prompt is too long, it will be truncated. :param prompt: The prompt to check. - :return: A dictionary containing the resized prompt and additional information. + :returns: A dictionary containing the resized prompt and additional information. """ def _extract_messages_from_response(self, message_tag: str, response_body: Dict[str, Any]) -> List[ChatMessage]: @@ -130,7 +130,7 @@ def _extract_messages_from_response(self, message_tag: str, response_body: Dict[ :param message_tag: The key for the message in the response body. :param response_body: The response body. - :return: The extracted ChatMessage list. + :returns: The extracted ChatMessage list. """ metadata = {k: v for (k, v) in response_body.items() if k != message_tag} return [ChatMessage.from_assistant(response_body[message_tag], meta=metadata)] @@ -141,7 +141,7 @@ def response_body_message_key(self) -> str: Returns the key for the message in the response body. Subclasses should override this method to return the correct message key - where the response is located. - :return: The key for the message in the response body. + :returns: The key for the message in the response body. """ @abstractmethod @@ -150,7 +150,7 @@ def _extract_token_from_stream(self, chunk: Dict[str, Any]) -> str: Extracts the token from a streaming chunk. :param chunk: The streaming chunk. - :return: The extracted token. + :returns: The extracted token. """ @@ -192,7 +192,7 @@ def prepare_body(self, messages: List[ChatMessage], **inference_kwargs) -> Dict[ :param messages: The chat messages to package into the request. :param inference_kwargs: Additional inference kwargs to use. - :return: The prepared body. + :returns: The prepared body. """ default_params = { "max_tokens_to_sample": self.generation_kwargs.get("max_tokens_to_sample") or 512, @@ -212,7 +212,7 @@ def prepare_chat_messages(self, messages: List[ChatMessage]) -> str: Prepares the chat messages for the Anthropic Claude request. :param messages: The chat messages to prepare. - :return: The prepared chat messages as a string. + :returns: The prepared chat messages as a string. """ conversation = [] for index, message in enumerate(messages): @@ -241,7 +241,7 @@ def check_prompt(self, prompt: str) -> Dict[str, Any]: Checks the prompt length and resizes it if necessary. If the prompt is too long, it will be truncated. :param prompt: The prompt to check. - :return: A dictionary containing the resized prompt and additional information. + :returns: A dictionary containing the resized prompt and additional information. """ return self.prompt_handler(prompt) @@ -249,7 +249,7 @@ def response_body_message_key(self) -> str: """ Returns the key for the message in the response body for Anthropic Claude i.e. "completion". - :return: The key for the message in the response body. + :returns: The key for the message in the response body. """ return "completion" @@ -258,7 +258,7 @@ def _extract_token_from_stream(self, chunk: Dict[str, Any]) -> str: Extracts the token from a streaming chunk. :param chunk: The streaming chunk. - :return: The extracted token. + :returns: The extracted token. """ return chunk.get("completion", "") @@ -340,7 +340,7 @@ def prepare_chat_messages(self, messages: List[ChatMessage]) -> str: Prepares the chat messages for the Meta Llama 2 request. :param messages: The chat messages to prepare. - :return: The prepared chat messages as a string ready for the model. + :returns: The prepared chat messages as a string ready for the model. """ prepared_prompt: str = self.prompt_handler.tokenizer.apply_chat_template( conversation=messages, tokenize=False, chat_template=self.chat_template @@ -352,7 +352,7 @@ def check_prompt(self, prompt: str) -> Dict[str, Any]: Checks the prompt length and resizes it if necessary. If the prompt is too long, it will be truncated. :param prompt: The prompt to check. - :return: A dictionary containing the resized prompt and additional information. + :returns: A dictionary containing the resized prompt and additional information. """ return self.prompt_handler(prompt) @@ -361,7 +361,7 @@ def response_body_message_key(self) -> str: """ Returns the key for the message in the response body for Meta Llama 2 i.e. "generation". - :return: The key for the message in the response body. + :returns: The key for the message in the response body. """ return "generation" @@ -370,6 +370,6 @@ def _extract_token_from_stream(self, chunk: Dict[str, Any]) -> str: Extracts the token from a streaming chunk. :param chunk: The streaming chunk. - :return: The extracted token. + :returns: The extracted token. """ return chunk.get("generation", "") diff --git a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py index 3b5a8f6cc..bea6924f6 100644 --- a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py +++ b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py @@ -140,7 +140,7 @@ def invoke(self, *args, **kwargs): :param args: The positional arguments passed to the generator. :param kwargs: The keyword arguments passed to the generator. - :return: List of `ChatMessage` generated by LLM. + :returns: List of `ChatMessage` generated by LLM. """ kwargs = kwargs.copy() @@ -183,7 +183,7 @@ def run(self, messages: List[ChatMessage], generation_kwargs: Optional[Dict[str, :param messages: The messages to generate a response to. :param generation_kwargs: Additional generation keyword arguments passed to the model. - :return: A dictionary with the following keys: + :returns: A dictionary with the following keys: - `replies`: The generated List of `ChatMessage` objects. """ return {"replies": self.invoke(messages=messages, **(generation_kwargs or {}))} @@ -194,7 +194,7 @@ def get_model_adapter(cls, model: str) -> Optional[Type[BedrockModelChatAdapter] Returns the model adapter for the given model. :param model: The model to get the adapter for. - :return: The model adapter for the given model, or None if the model is not supported. + :returns: The model adapter for the given model, or None if the model is not supported. """ for pattern, adapter in cls.SUPPORTED_MODEL_PATTERNS.items(): if re.fullmatch(pattern, model): diff --git a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py index 706d29c98..f6af48ae1 100644 --- a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py +++ b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py @@ -142,7 +142,7 @@ def _ensure_token_limit(self, prompt: Union[str, List[Dict[str, str]]]) -> Union the initialization of the component. :param prompt: The prompt to be sent to the model. - :return: The resized prompt. + :returns: The resized prompt. """ # the prompt for this model will be of the type str if isinstance(prompt, List): @@ -171,7 +171,7 @@ def invoke(self, *args, **kwargs): :param args: Additional positional arguments passed to the generator. :param kwargs: Additional keyword arguments passed to the generator. - :return: A list of generated responses (strings). + :returns: A list of generated responses (strings). """ kwargs = kwargs.copy() prompt: str = kwargs.pop("prompt", None) @@ -225,7 +225,7 @@ def run(self, prompt: str, generation_kwargs: Optional[Dict[str, Any]] = None): :param prompt: The prompt to generate a response for. :param generation_kwargs: Additional keyword arguments passed to the generator. - :return: A dictionary with the following keys: + :returns: A dictionary with the following keys: - `replies`: A list of generated responses (strings). """ return {"replies": self.invoke(prompt=prompt, **(generation_kwargs or {}))} @@ -236,7 +236,7 @@ def get_model_adapter(cls, model: str) -> Optional[Type[BedrockModelAdapter]]: Gets the model adapter for the given model. :param model: The model name. - :return: The model adapter class, or None if no adapter is found. + :returns: The model adapter class, or None if no adapter is found. """ for pattern, adapter in cls.SUPPORTED_MODEL_PATTERNS.items(): if re.fullmatch(pattern, model): diff --git a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/handlers.py b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/handlers.py index b7b555ec0..ddc276264 100644 --- a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/handlers.py +++ b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/handlers.py @@ -34,7 +34,7 @@ def __call__(self, prompt: str, **kwargs) -> Dict[str, Union[str, int]]: :param prompt: the prompt to be sent to the model. :param kwargs: Additional keyword arguments passed to the handler. - :return: A dictionary containing the resized prompt and additional information. + :returns: A dictionary containing the resized prompt and additional information. """ resized_prompt = prompt prompt_length = 0 @@ -75,7 +75,7 @@ def __call__(self, token_received: str, **kwargs) -> str: :param token_received: The token received from the stream. :param kwargs: Additional keyword arguments passed to the handler. - :return: The token to be sent to the stream. + :returns: The token to be sent to the stream. """ pass @@ -87,7 +87,7 @@ def __call__(self, token_received, **kwargs) -> str: :param token_received: The token received from the stream. :param kwargs: Additional keyword arguments passed to the handler. - :return: The token to be sent to the stream. + :returns: The token to be sent to the stream. """ print(token_received, flush=True, end="") # noqa: T201 return token_received diff --git a/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/sagemaker.py b/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/sagemaker.py index 106698558..c171ccdf6 100644 --- a/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/sagemaker.py +++ b/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/sagemaker.py @@ -179,7 +179,7 @@ def _get_aws_session( :param aws_profile_name: AWS profile name. :raises AWSConfigurationError: If the provided AWS credentials are invalid. - :return: The created AWS session. + :returns: The created AWS session. """ try: return boto3.Session( @@ -202,7 +202,7 @@ def run(self, prompt: str, generation_kwargs: Optional[Dict[str, Any]] = None): :param generation_kwargs: Additional keyword arguments for text generation. These parameters will potentially override the parameters passed in the `__init__` method. - :return: A dictionary with the following keys: + :returns: A dictionary with the following keys: - `replies`: A list of strings containing the generated responses - `meta`: A list of dictionaries containing the metadata for each response. """ diff --git a/integrations/astra/src/haystack_integrations/components/retrievers/astra/retriever.py b/integrations/astra/src/haystack_integrations/components/retrievers/astra/retriever.py index 2b9ac7d28..80e436e0a 100644 --- a/integrations/astra/src/haystack_integrations/components/retrievers/astra/retriever.py +++ b/integrations/astra/src/haystack_integrations/components/retrievers/astra/retriever.py @@ -52,7 +52,7 @@ def run(self, query_embedding: List[float], filters: Optional[Dict[str, Any]] = :param filters: filters to narrow down the search space. :param top_k: the maximum number of documents to retrieve. :returns: a dictionary with the following keys: - - documents: A list of documents retrieved from the AstraDocumentStore. + - `documents`: A list of documents retrieved from the AstraDocumentStore. """ if not top_k: diff --git a/integrations/chroma/src/haystack_integrations/components/retrievers/chroma/retriever.py b/integrations/chroma/src/haystack_integrations/components/retrievers/chroma/retriever.py index 10f97f01f..7138eff88 100644 --- a/integrations/chroma/src/haystack_integrations/components/retrievers/chroma/retriever.py +++ b/integrations/chroma/src/haystack_integrations/components/retrievers/chroma/retriever.py @@ -64,8 +64,8 @@ def run( :param query: The input data for the retriever. In this case, a plain-text query. :param top_k: The maximum number of documents to retrieve. If not specified, the default value from the constructor is used. - :return: A dictionary with the following keys: - - "documents": List of documents returned by the search engine. + :returns: A dictionary with the following keys: + - `documents`: List of documents returned by the search engine. :raises ValueError: If the specified document store is not found or is not a MemoryDocumentStore instance. """ @@ -119,8 +119,8 @@ def run( Run the retriever on the given input data. :param query_embedding: the query embeddings. - :return: a dictionary with the following keys: - - "documents": List of documents returned by the search engine. + :returns: a dictionary with the following keys: + - `documents`: List of documents returned by the search engine. """ top_k = top_k or self.top_k diff --git a/integrations/chroma/src/haystack_integrations/document_stores/chroma/document_store.py b/integrations/chroma/src/haystack_integrations/document_stores/chroma/document_store.py index 4201de23b..0db9f832a 100644 --- a/integrations/chroma/src/haystack_integrations/document_stores/chroma/document_store.py +++ b/integrations/chroma/src/haystack_integrations/document_stores/chroma/document_store.py @@ -206,7 +206,7 @@ def search(self, queries: List[str], top_k: int) -> List[List[Document]]: :param queries: the list of queries to search for. :param top_k: top_k documents to return for each query. - :return: matching documents for each query. + :returns: matching documents for each query. """ results = self._collection.query( query_texts=queries, n_results=top_k, include=["embeddings", "documents", "metadatas", "distances"] 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 305743126..87b4f8834 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 @@ -112,8 +112,8 @@ def run(self, text: str): :param text: the text to embed. :returns: A dictionary with the following keys: - - "embedding": the embedding of the text. - - "meta": metadata about the request. + - `embedding`: the embedding of the text. + - `meta`: metadata about the request. :raises TypeError: If the input is not a string. """ if not isinstance(text, str): 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 dcee45d10..980441009 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 @@ -149,7 +149,7 @@ def run(self, messages: List[ChatMessage], generation_kwargs: Optional[Dict[str, For more details on the parameters supported by the Cohere API, refer to the Cohere [documentation](https://docs.cohere.com/reference/chat). :returns: A dictionary with the following keys: - - "replies": a list of `ChatMessage` instances representing the generated responses. + - `replies`: a list of `ChatMessage` instances representing the generated responses. """ # update generation kwargs by merging with the generation kwargs passed to the run method generation_kwargs = {**self.generation_kwargs, **(generation_kwargs or {})} @@ -186,7 +186,7 @@ def _build_chunk(self, chunk) -> StreamingChunk: Converts the response from the Cohere API to a StreamingChunk. :param chunk: The chunk returned by the OpenAI API. :param choice: The choice returned by the OpenAI API. - :return: The StreamingChunk. + :returns: The StreamingChunk. """ chat_message = StreamingChunk(content=chunk.text, meta={"index": chunk.index, "event_type": chunk.event_type}) return chat_message @@ -195,7 +195,7 @@ def _build_message(self, cohere_response): """ Converts the non-streaming response from the Cohere API to a ChatMessage. :param cohere_response: The completion returned by the Cohere API. - :return: The ChatMessage. + :returns: The ChatMessage. """ content = cohere_response.text message = ChatMessage.from_assistant(content=content) 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 6a25727dd..2c9a97478 100644 --- a/integrations/cohere/src/haystack_integrations/components/generators/cohere/generator.py +++ b/integrations/cohere/src/haystack_integrations/components/generators/cohere/generator.py @@ -127,8 +127,8 @@ def run(self, prompt: str): :param prompt: the prompt to be sent to the generative model. :returns: A dictionary with the following keys: - - "replies": the list of replies generated by the model. - - "meta": metadata about the request. + - `replies`: the list of replies generated by the model. + - `meta`: metadata about the request. """ response = self.client.generate( model=self.model, prompt=prompt, stream=self.streaming_callback is not None, **self.model_parameters @@ -153,7 +153,7 @@ def _build_chunk(self, chunk) -> StreamingChunk: """ Converts the response from the Cohere API to a StreamingChunk. :param chunk: The chunk returned by the OpenAI API. - :return: The StreamingChunk. + :returns: The StreamingChunk. """ streaming_chunk = StreamingChunk(content=chunk.text, meta={"index": chunk.index}) return streaming_chunk 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 f08ff1adc..4af8e1bbe 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 @@ -144,7 +144,7 @@ def run(self, documents: List[Document]): Embeds a list of Documents. :param documents: List of Documents to embed. - :return: A dictionary with the following keys: + :returns: A dictionary with the following keys: - `documents`: List of Documents with each Document's `embedding` field set to the computed embeddings. """ if not isinstance(documents, list) or documents and not isinstance(documents[0], Document): 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 ffba6a902..13a89d1ce 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 @@ -101,7 +101,7 @@ def run(self, text: str): Embeds text using the Fastembed model. :param text: A string to embed. - :return: A dictionary with the following keys: + :returns: A dictionary with the following keys: - `embedding`: A list of floats representing the embedding of the input text. :raises TypeError: If the input is not a string. :raises RuntimeError: If the embedding model has not been loaded. diff --git a/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/image_generator.py b/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/image_generator.py index 422e1cfe6..ae8c4892f 100644 --- a/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/image_generator.py +++ b/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/image_generator.py @@ -85,7 +85,7 @@ def run(self, prompt: str, negative_prompt: Optional[str] = None): :param negative_prompt: A description of what you want to omit in the generated images. :returns: A dictionary with the following keys: - - images: A list of ByteStream objects, each containing an image. + - `images`: A list of ByteStream objects, each containing an image. """ negative_prompt = negative_prompt or self._kwargs.get("negative_prompt") res = self._model.generate_images(prompt=prompt, negative_prompt=negative_prompt, **self._kwargs) diff --git a/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/question_answering.py b/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/question_answering.py index 79c343b02..32cde86ef 100644 --- a/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/question_answering.py +++ b/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/question_answering.py @@ -89,7 +89,7 @@ def run(self, image: ByteStream, question: str): :param image: The image to ask the question about. :param question: The question to ask. :returns: A dictionary with the following keys: - - answers: A list of answers to the question. + - `answers`: A list of answers to the question. """ answers = self._model.ask_question(image=Image(image.data), question=question, **self._kwargs) return {"answers": answers} diff --git a/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/text_generator.py b/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/text_generator.py index e16954f8f..a00678f19 100644 --- a/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/text_generator.py +++ b/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/text_generator.py @@ -115,10 +115,10 @@ def run(self, prompt: str): :param prompt: The prompt to use for text generation. :returns: A dictionary with the following keys: - - answers: A list of generated answers. - - safety_attributes: A dictionary with the [safety scores](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/responsible-ai#safety_attribute_descriptions) + - `answers`: A list of generated answers. + - `safety_attributes`: A dictionary with the [safety scores](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/responsible-ai#safety_attribute_descriptions) of each answer. - - citations: A list of citations for each answer. + - `citations`: A list of citations for each answer. """ res = self._model.predict(prompt=prompt, **self._kwargs) diff --git a/integrations/llama_cpp/src/haystack_integrations/components/generators/llama_cpp/generator.py b/integrations/llama_cpp/src/haystack_integrations/components/generators/llama_cpp/generator.py index c5d37c1a8..1c504b6f3 100644 --- a/integrations/llama_cpp/src/haystack_integrations/components/generators/llama_cpp/generator.py +++ b/integrations/llama_cpp/src/haystack_integrations/components/generators/llama_cpp/generator.py @@ -79,8 +79,8 @@ def run(self, prompt: str, generation_kwargs: Optional[Dict[str, Any]] = None): For more information on the available kwargs, see [llama.cpp documentation](https://llama-cpp-python.readthedocs.io/en/latest/api-reference/#llama_cpp.Llama.create_completion). :returns: A dictionary with the following keys: - - "replies": the list of replies generated by the model. - - "meta": metadata about the request. + - `replies`: the list of replies generated by the model. + - `meta`: metadata about the request. """ if self.model is None: error_msg = "The model has not been loaded. Please call warm_up() before running." diff --git a/integrations/opensearch/src/haystack_integrations/document_stores/opensearch/document_store.py b/integrations/opensearch/src/haystack_integrations/document_stores/opensearch/document_store.py index e91347728..e9c88274c 100644 --- a/integrations/opensearch/src/haystack_integrations/document_stores/opensearch/document_store.py +++ b/integrations/opensearch/src/haystack_integrations/document_stores/opensearch/document_store.py @@ -257,7 +257,7 @@ def _bm25_retrieval( :param scale_score: If `True` scales the Document`s scores between 0 and 1, defaults to False :param all_terms_must_match: If `True` all terms in `query` must be present in the Document, defaults to False :raises ValueError: If `query` is an empty string - :return: List of Document that match `query` + :returns: List of Document that match `query` """ if not query: @@ -314,7 +314,7 @@ def _embedding_retrieval( Filters are applied during the approximate kNN search to ensure that top_k matching documents are returned. :param top_k: Maximum number of Documents to return, defaults to 10 :raises ValueError: If `query_embedding` is an empty list - :return: List of Document that are most similar to `query_embedding` + :returns: List of Document that are most similar to `query_embedding` """ if not query_embedding: diff --git a/integrations/qdrant/src/haystack_integrations/document_stores/qdrant/document_store.py b/integrations/qdrant/src/haystack_integrations/document_stores/qdrant/document_store.py index 4a47bf59e..645db88ae 100644 --- a/integrations/qdrant/src/haystack_integrations/document_stores/qdrant/document_store.py +++ b/integrations/qdrant/src/haystack_integrations/document_stores/qdrant/document_store.py @@ -420,7 +420,7 @@ def _handle_duplicate_documents( overwrite: Update any existing documents with the same ID when adding documents. fail: an error is raised if the document ID of the document being added already exists. - :return: A list of Haystack Document objects. + :returns: A list of Haystack Document objects. """ index = index or self.index @@ -443,7 +443,7 @@ def _drop_duplicate_documents(self, documents: List[Document], index: Optional[s :param documents: A list of Haystack Document objects. :param index: name of the index - :return: A list of Haystack Document objects. + :returns: A list of Haystack Document objects. """ _hash_ids: Set = set() _documents: List[Document] = [] diff --git a/integrations/qdrant/src/haystack_integrations/document_stores/qdrant/filters.py b/integrations/qdrant/src/haystack_integrations/document_stores/qdrant/filters.py index 77d800853..fc18a3ab2 100644 --- a/integrations/qdrant/src/haystack_integrations/document_stores/qdrant/filters.py +++ b/integrations/qdrant/src/haystack_integrations/document_stores/qdrant/filters.py @@ -186,7 +186,7 @@ def _squeeze_filter(self, payload_filter: models.Filter) -> models.Filter: Simplify given payload filter, if the nested structure might be unnested. That happens if there is a single clause in that filter. :param payload_filter: - :return: + :returns: """ filter_parts = { "must": payload_filter.must, diff --git a/integrations/unstructured/src/haystack_integrations/components/converters/unstructured/converter.py b/integrations/unstructured/src/haystack_integrations/components/converters/unstructured/converter.py index 84865cee5..637c0840f 100644 --- a/integrations/unstructured/src/haystack_integrations/components/converters/unstructured/converter.py +++ b/integrations/unstructured/src/haystack_integrations/components/converters/unstructured/converter.py @@ -139,7 +139,7 @@ def run( (same metadata for all files). :returns: A dictionary with the following key: - - "documents": List of Haystack Documents. + - `documents`: List of Haystack Documents. :raises ValueError: If `meta` is a list and `paths` contains directories. """ diff --git a/nodes/text2speech/text2speech/utils/text_to_speech.py b/nodes/text2speech/text2speech/utils/text_to_speech.py index 84c08e90f..e7d22cc11 100644 --- a/nodes/text2speech/text2speech/utils/text_to_speech.py +++ b/nodes/text2speech/text2speech/utils/text_to_speech.py @@ -104,7 +104,7 @@ def text_to_audio_file( leaves it untouched. :param audio_naming_function: A function mapping the input text into the audio file name. By default, the audio file gets the name from the MD5 sum of the input text. - :return: The path to the generated file. + :returns: The path to the generated file. """ if not os.path.exists(generated_audio_dir): os.mkdir(generated_audio_dir) @@ -140,7 +140,7 @@ def text_to_audio_data(self, text: str, _models_output_key: str = "wav") -> np.a :param text: The text to convert into audio. :param _models_output_key: The key in the prediction dictionary that contains the audio data. Defaults to 'wav'. - :return: A numpy array representing the audio generated by the model. + :returns: A numpy array representing the audio generated by the model. """ prediction = self.model(text) if not prediction: From d4a2ed56377c5d440665cc6c4bbf0721938a7ea4 Mon Sep 17 00:00:00 2001 From: Madeesh Kannan Date: Wed, 6 Mar 2024 10:15:51 +0100 Subject: [PATCH 2/8] feat: Add `NvidiaTextEmbedder`, `NvidiaDocumentEmbedder` and co. (#537) --- integrations/nvidia/pydoc/config.yml | 7 +- integrations/nvidia/pyproject.toml | 1 + .../components/embedders/nvidia/__init__.py | 9 + .../components/embedders/nvidia/_schema.py | 91 ++++++ .../embedders/nvidia/document_embedder.py | 208 +++++++++++++ .../components/embedders/nvidia/models.py | 31 ++ .../embedders/nvidia/text_embedder.py | 144 +++++++++ .../utils/nvidia/__init__.py | 3 + .../utils/nvidia/client.py | 61 ++++ .../nvidia/tests/test_document_embedder.py | 287 ++++++++++++++++++ integrations/nvidia/tests/test_placeholder.py | 2 - .../nvidia/tests/test_text_embedder.py | 119 ++++++++ 12 files changed, 960 insertions(+), 3 deletions(-) create mode 100644 integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/__init__.py create mode 100644 integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/_schema.py create mode 100644 integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/document_embedder.py create mode 100644 integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/models.py create mode 100644 integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/text_embedder.py create mode 100644 integrations/nvidia/src/haystack_integrations/utils/nvidia/__init__.py create mode 100644 integrations/nvidia/src/haystack_integrations/utils/nvidia/client.py create mode 100644 integrations/nvidia/tests/test_document_embedder.py delete mode 100644 integrations/nvidia/tests/test_placeholder.py create mode 100644 integrations/nvidia/tests/test_text_embedder.py diff --git a/integrations/nvidia/pydoc/config.yml b/integrations/nvidia/pydoc/config.yml index 65bedf8b4..675db0335 100644 --- a/integrations/nvidia/pydoc/config.yml +++ b/integrations/nvidia/pydoc/config.yml @@ -1,7 +1,12 @@ loaders: - type: haystack_pydoc_tools.loaders.CustomPythonLoader search_path: [../src] - modules: [] + modules: + [ + "haystack_integrations.components.embedders.nvidia.document_embedder", + "haystack_integrations.components.embedders.nvidia.text_embedder", + "haystack_integrations.components.embedders.nvidia.models", + ] ignore_when_discovered: ["__init__"] processors: - type: filter diff --git a/integrations/nvidia/pyproject.toml b/integrations/nvidia/pyproject.toml index cded4787b..ba25812a8 100644 --- a/integrations/nvidia/pyproject.toml +++ b/integrations/nvidia/pyproject.toml @@ -116,6 +116,7 @@ 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 new file mode 100644 index 000000000..6ad2f9f6b --- /dev/null +++ b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/__init__.py @@ -0,0 +1,9 @@ +from .document_embedder import NvidiaDocumentEmbedder +from .models import NvidiaEmbeddingModel +from .text_embedder import NvidiaTextEmbedder + +__all__ = [ + "NvidiaDocumentEmbedder", + "NvidiaEmbeddingModel", + "NvidiaTextEmbedder", +] diff --git a/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/_schema.py b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/_schema.py new file mode 100644 index 000000000..a0598be86 --- /dev/null +++ b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/_schema.py @@ -0,0 +1,91 @@ +from dataclasses import asdict, dataclass +from typing import Any, Dict, List, Literal, Union + +from haystack_integrations.utils.nvidia import NvidiaCloudFunctionsClient + +from .models import NvidiaEmbeddingModel + +MAX_INPUT_STRING_LENGTH = 2048 +MAX_INPUTS = 50 + + +def get_model_nvcf_id(model: NvidiaEmbeddingModel, client: NvidiaCloudFunctionsClient) -> str: + """ + Returns the Nvidia Cloud Functions UUID for the given model. + """ + + available_functions = client.available_functions() + func = available_functions.get(str(model)) + if func is None: + msg = f"Model '{model}' was not found on the Nvidia Cloud Functions backend" + raise ValueError(msg) + elif func.status != "ACTIVE": + msg = f"Model '{model}' is not currently active/usable on the Nvidia Cloud Functions backend" + raise ValueError(msg) + + return func.id + + +@dataclass +class EmbeddingsRequest: + input: Union[str, List[str]] + model: Literal["query", "passage"] + encoding_format: Literal["float", "base64"] = "float" + + def __post_init__(self): + if isinstance(self.input, list): + if len(self.input) > MAX_INPUTS: + msg = f"The number of inputs should not exceed {MAX_INPUTS}" + raise ValueError(msg) + else: + self.input = [self.input] + + if len(self.input) == 0: + msg = "The number of inputs should not be 0" + raise ValueError(msg) + + if any(len(x) > MAX_INPUT_STRING_LENGTH for x in self.input): + msg = f"The length of each input should not exceed {MAX_INPUT_STRING_LENGTH} characters" + raise ValueError(msg) + + if self.encoding_format not in ["float", "base64"]: + msg = "encoding_format should be either 'float' or 'base64'" + raise ValueError(msg) + + if self.model not in ["query", "passage"]: + msg = "model should be either 'query' or 'passage'" + raise ValueError(msg) + + def to_dict(self) -> Dict[str, Any]: + return asdict(self) + + +@dataclass +class Usage: + prompt_tokens: int + total_tokens: int + + def to_dict(self) -> Dict[str, Any]: + return asdict(self) + + +@dataclass +class Embeddings: + index: int + embedding: Union[List[float], str] + + +@dataclass +class EmbeddingsResponse: + data: List[Embeddings] + usage: Usage + + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> "EmbeddingsResponse": + try: + embeddings = [Embeddings(**x) for x in data["data"]] + usage = Usage(**data["usage"]) + return cls(data=embeddings, usage=usage) + except (KeyError, TypeError) as e: + msg = f"Failed to parse EmbeddingsResponse from data: {data}" + raise ValueError(msg) from e 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 new file mode 100644 index 000000000..139a184b7 --- /dev/null +++ b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/document_embedder.py @@ -0,0 +1,208 @@ +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 +from haystack_integrations.utils.nvidia import NvidiaCloudFunctionsClient +from tqdm import tqdm + +from ._schema import MAX_INPUTS, EmbeddingsRequest, EmbeddingsResponse, Usage, get_model_nvcf_id +from .models import NvidiaEmbeddingModel + + +@component +class NvidiaDocumentEmbedder: + """ + A component for embedding documents using embedding models provided by + [NVIDIA AI Foundation Endpoints](https://www.nvidia.com/en-us/ai-data-science/foundation-models/). + + Usage example: + ```python + from haystack_integrations.components.embedders.nvidia import NvidiaDocumentEmbedder, NvidiaEmbeddingModel + + doc = Document(content="I love pizza!") + + text_embedder = NvidiaDocumentEmbedder(model=NvidiaEmbeddingModel.NVOLVE_40K) + text_embedder.warm_up() + + result = document_embedder.run([doc]) + print(result["documents"][0].embedding) + ``` + """ + + def __init__( + self, + model: Union[str, NvidiaEmbeddingModel], + api_key: Secret = Secret.from_env_var("NVIDIA_API_KEY"), + prefix: str = "", + suffix: str = "", + batch_size: int = 32, + progress_bar: bool = True, + meta_fields_to_embed: Optional[List[str]] = None, + embedding_separator: str = "\n", + ): + """ + Create a NvidiaTextEmbedder component. + + :param model: + Embedding model to use. + :param api_key: + API key for the NVIDIA AI Foundation Endpoints. + :param prefix: + A string to add to the beginning of each text. + :param suffix: + A string to add to the end of each text. + :param batch_size: + Number of Documents to encode at once. + Cannot be greater than 50. + :param progress_bar: + Whether to show a progress bar or not. + :param meta_fields_to_embed: + 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. + """ + + if isinstance(model, str): + model = NvidiaEmbeddingModel.from_str(model) + + resolved_api_key = api_key.resolve_value() + assert resolved_api_key is not None + + # Upper-limit for the endpoint. + if batch_size > MAX_INPUTS: + msg = f"NVIDIA Cloud Functions currently support a maximum batch size of {MAX_INPUTS}." + raise ValueError(msg) + + self.api_key = api_key + self.model = model + self.prefix = prefix + self.suffix = suffix + self.batch_size = batch_size + self.progress_bar = progress_bar + self.meta_fields_to_embed = meta_fields_to_embed or [] + self.embedding_separator = embedding_separator + + self.client = NvidiaCloudFunctionsClient( + api_key=resolved_api_key, + headers={ + "Content-Type": "application/json", + "Accept": "application/json", + }, + ) + self.nvcf_id = None + self._initialized = False + + def warm_up(self): + """ + Initializes the component. + """ + if self._initialized: + return + + self.nvcf_id = get_model_nvcf_id(self.model, self.client) + self._initialized = True + + def to_dict(self) -> Dict[str, Any]: + """ + Serializes the component to a dictionary. + + :returns: + Dictionary with serialized data. + """ + return default_to_dict( + self, + api_key=self.api_key.to_dict(), + model=str(self.model), + prefix=self.prefix, + suffix=self.suffix, + batch_size=self.batch_size, + progress_bar=self.progress_bar, + meta_fields_to_embed=self.meta_fields_to_embed, + embedding_separator=self.embedding_separator, + ) + + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> "NvidiaDocumentEmbedder": + """ + Deserializes the component from a dictionary. + + :param data: + The dictionary to deserialize from. + :returns: + The deserialized component. + """ + data["init_parameters"]["model"] = NvidiaEmbeddingModel.from_str(data["init_parameters"]["model"]) + deserialize_secrets_inplace(data["init_parameters"], keys=["api_key"]) + return default_from_dict(cls, data) + + def _prepare_texts_to_embed(self, documents: List[Document]) -> List[str]: + texts_to_embed = [] + for doc in documents: + meta_values_to_embed = [ + str(doc.meta[key]) for key in self.meta_fields_to_embed if key in doc.meta and doc.meta[key] is not None + ] + text_to_embed = ( + self.prefix + self.embedding_separator.join([*meta_values_to_embed, doc.content or ""]) + self.suffix + ) + texts_to_embed.append(text_to_embed) + + return texts_to_embed + + def _embed_batch(self, texts_to_embed: List[str], batch_size: int) -> Tuple[List[List[float]], Dict[str, Any]]: + all_embeddings: List[List[float]] = [] + usage = Usage(prompt_tokens=0, total_tokens=0) + assert self.nvcf_id is not None + + for i in tqdm( + range(0, len(texts_to_embed), batch_size), disable=not self.progress_bar, desc="Calculating embeddings" + ): + batch = texts_to_embed[i : i + batch_size] + + request = EmbeddingsRequest(input=batch, model="passage").to_dict() + json_response = self.client.query_function(self.nvcf_id, request) + response = EmbeddingsResponse.from_dict(json_response) + + # Sort resulting embeddings by index + assert all(isinstance(r.embedding, list) for r in response.data) + sorted_embeddings: List[List[float]] = [r.embedding for r in sorted(response.data, key=lambda e: e.index)] # type: ignore + all_embeddings.extend(sorted_embeddings) + + usage.prompt_tokens += response.usage.prompt_tokens + usage.total_tokens += response.usage.total_tokens + + return all_embeddings, {"usage": usage.to_dict()} + + @component.output_types(documents=List[Document], meta=Dict[str, Any]) + def run(self, documents: List[Document]): + """ + Embed a list of Documents. + + The embedding of each Document is stored in the `embedding` field of the Document. + + :param documents: + A list of Documents to embed. + :returns: + A dictionary with the following keys and values: + - `documents` - List of processed Documents with embeddings. + - `meta` - Metadata on usage statistics, etc. + :raises RuntimeError: + If the component was not initialized. + :raises TypeError: + If the input is not a string. + """ + if not self._initialized: + msg = "The embedding model has not been loaded. Please call warm_up() before running." + raise RuntimeError(msg) + if not isinstance(documents, list) or documents and not isinstance(documents[0], Document): + msg = ( + "NvidiaDocumentEmbedder expects a list of Documents as input." + "In case you want to embed a string, please use the NvidiaTextEmbedder." + ) + raise TypeError(msg) + + texts_to_embed = self._prepare_texts_to_embed(documents) + embeddings, metadata = self._embed_batch(texts_to_embed, self.batch_size) + for doc, emb in zip(documents, embeddings): + doc.embedding = emb + + return {"documents": documents, "meta": metadata} diff --git a/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/models.py b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/models.py new file mode 100644 index 000000000..dd11ac727 --- /dev/null +++ b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/models.py @@ -0,0 +1,31 @@ +from enum import Enum + + +class NvidiaEmbeddingModel(Enum): + """ + [NVIDIA AI Foundation models](https://catalog.ngc.nvidia.com/ai-foundation-models) + used for generating embeddings. + """ + + #: [Retrieval QA Embedding Model](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/ai-foundation/models/nvolve-40k). + NVOLVE_40K = "playground_nvolveqa_40k" + + def __str__(self): + return self.value + + @classmethod + def from_str(cls, string: str) -> "NvidiaEmbeddingModel": + """ + Create an embedding model from a string. + + :param string: + String to convert. + :returns: + Embedding model. + """ + enum_map = {e.value: e for e in NvidiaEmbeddingModel} + emb_model = enum_map.get(string) + if emb_model is None: + msg = f"Unknown embedding model '{string}'. Supported modes are: {list(enum_map.keys())}" + raise ValueError(msg) + return emb_model 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 new file mode 100644 index 000000000..43d62ed92 --- /dev/null +++ b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/text_embedder.py @@ -0,0 +1,144 @@ +from typing import Any, Dict, List, Union + +from haystack import component, default_from_dict, default_to_dict +from haystack.utils import Secret, deserialize_secrets_inplace +from haystack_integrations.utils.nvidia import NvidiaCloudFunctionsClient + +from ._schema import EmbeddingsRequest, EmbeddingsResponse, get_model_nvcf_id +from .models import NvidiaEmbeddingModel + + +@component +class NvidiaTextEmbedder: + """ + A component for embedding strings using embedding models provided by + [NVIDIA AI Foundation Endpoints](https://www.nvidia.com/en-us/ai-data-science/foundation-models/). + + For models that differentiate between query and document inputs, + this component embeds the input string as a query. + + Usage example: + ```python + from haystack_integrations.components.embedders.nvidia import NvidiaTextEmbedder, NvidiaEmbeddingModel + + text_to_embed = "I love pizza!" + + text_embedder = NvidiaTextEmbedder(model=NvidiaEmbeddingModel.NVOLVE_40K) + text_embedder.warm_up() + + print(text_embedder.run(text_to_embed)) + ``` + """ + + def __init__( + self, + model: Union[str, NvidiaEmbeddingModel], + api_key: Secret = Secret.from_env_var("NVIDIA_API_KEY"), + prefix: str = "", + suffix: str = "", + ): + """ + Create a NvidiaTextEmbedder component. + + :param model: + Embedding model to use. + :param api_key: + API key for the NVIDIA AI Foundation Endpoints. + :param prefix: + A string to add to the beginning of each text. + :param suffix: + A string to add to the end of each text. + """ + + if isinstance(model, str): + model = NvidiaEmbeddingModel.from_str(model) + + resolved_api_key = api_key.resolve_value() + assert resolved_api_key is not None + + self.api_key = api_key + self.model = model + self.prefix = prefix + self.suffix = suffix + self.client = NvidiaCloudFunctionsClient( + api_key=resolved_api_key, + headers={ + "Content-Type": "application/json", + "Accept": "application/json", + }, + ) + self.nvcf_id = None + self._initialized = False + + def warm_up(self): + """ + Initializes the component. + """ + if self._initialized: + return + + self.nvcf_id = get_model_nvcf_id(self.model, self.client) + self._initialized = True + + def to_dict(self) -> Dict[str, Any]: + """ + Serializes the component to a dictionary. + + :returns: + Dictionary with serialized data. + """ + return default_to_dict( + self, + api_key=self.api_key.to_dict(), + model=str(self.model), + prefix=self.prefix, + suffix=self.suffix, + ) + + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> "NvidiaTextEmbedder": + """ + Deserializes the component from a dictionary. + + :param data: + The dictionary to deserialize from. + :returns: + The deserialized component. + """ + data["init_parameters"]["model"] = NvidiaEmbeddingModel.from_str(data["init_parameters"]["model"]) + deserialize_secrets_inplace(data["init_parameters"], keys=["api_key"]) + return default_from_dict(cls, data) + + @component.output_types(embedding=List[float], meta=Dict[str, Any]) + def run(self, text: str): + """ + Embed a string. + + :param text: + The text to embed. + :returns: + A dictionary with the following keys and values: + - `embedding` - Embeddng of the text. + - `meta` - Metadata on usage statistics, etc. + :raises RuntimeError: + If the component was not initialized. + :raises TypeError: + If the input is not a string. + """ + if not self._initialized: + msg = "The embedding model has not been loaded. Please call warm_up() before running." + raise RuntimeError(msg) + if not isinstance(text, str): + msg = ( + "NvidiaTextEmbedder expects a string as an input." + "In case you want to embed a list of Documents, please use the NvidiaDocumentEmbedder." + ) + raise TypeError(msg) + + assert self.nvcf_id is not None + text_to_embed = self.prefix + text + self.suffix + request = EmbeddingsRequest(input=text_to_embed, model="query").to_dict() + json_response = self.client.query_function(self.nvcf_id, request) + response = EmbeddingsResponse.from_dict(json_response) + + return {"embedding": response.data[0].embedding, "meta": {"usage": response.usage.to_dict()}} diff --git a/integrations/nvidia/src/haystack_integrations/utils/nvidia/__init__.py b/integrations/nvidia/src/haystack_integrations/utils/nvidia/__init__.py new file mode 100644 index 000000000..b8015cfda --- /dev/null +++ b/integrations/nvidia/src/haystack_integrations/utils/nvidia/__init__.py @@ -0,0 +1,3 @@ +from .client import NvidiaCloudFunctionsClient + +__all__ = ["NvidiaCloudFunctionsClient"] diff --git a/integrations/nvidia/src/haystack_integrations/utils/nvidia/client.py b/integrations/nvidia/src/haystack_integrations/utils/nvidia/client.py new file mode 100644 index 000000000..5227e8c45 --- /dev/null +++ b/integrations/nvidia/src/haystack_integrations/utils/nvidia/client.py @@ -0,0 +1,61 @@ +import copy +from dataclasses import dataclass +from typing import Dict, Optional + +import requests + +FUNCTIONS_ENDPOINT = "https://api.nvcf.nvidia.com/v2/nvcf/functions" +INVOKE_ENDPOINT = "https://api.nvcf.nvidia.com/v2/nvcf/pexec/functions" +STATUS_ENDPOINT = "https://api.nvcf.nvidia.com/v2/nvcf/pexec/status" + +ACCEPTED_STATUS_CODE = 202 + + +@dataclass +class AvailableNvidiaCloudFunctions: + name: str + id: str + status: Optional[str] = None + + +class NvidiaCloudFunctionsClient: + def __init__(self, *, api_key: str, headers: Dict[str, str], timeout: int = 60): + self.api_key = api_key + self.fetch_url_format = STATUS_ENDPOINT + self.headers = copy.deepcopy(headers) + self.headers.update( + { + "Authorization": f"Bearer {api_key}", + } + ) + self.timeout = timeout + self.session = requests.Session() + + def query_function(self, func_id: str, payload: Dict[str, str]) -> Dict[str, str]: + invoke_url = f"{INVOKE_ENDPOINT}/{func_id}" + + response = self.session.post(invoke_url, headers=self.headers, json=payload, timeout=self.timeout) + request_id = response.headers.get("NVCF-REQID") + if request_id is None: + msg = "NVCF-REQID header not found in response" + raise ValueError(msg) + + while response.status_code == ACCEPTED_STATUS_CODE: + fetch_url = f"{self.fetch_url_format}/{request_id}" + response = self.session.get(fetch_url, headers=self.headers, timeout=self.timeout) + + response.raise_for_status() + return response.json() + + def available_functions(self) -> Dict[str, AvailableNvidiaCloudFunctions]: + response = self.session.get(FUNCTIONS_ENDPOINT, headers=self.headers, timeout=self.timeout) + response.raise_for_status() + + return { + f["name"]: AvailableNvidiaCloudFunctions( + name=f["name"], + id=f["id"], + status=f.get("status"), + ) + for f in response.json()["functions"] + } diff --git a/integrations/nvidia/tests/test_document_embedder.py b/integrations/nvidia/tests/test_document_embedder.py new file mode 100644 index 000000000..4f19633e8 --- /dev/null +++ b/integrations/nvidia/tests/test_document_embedder.py @@ -0,0 +1,287 @@ +import os + +import pytest +from haystack import Document +from haystack.utils import Secret +from haystack_integrations.components.embedders.nvidia import NvidiaDocumentEmbedder, NvidiaEmbeddingModel +from haystack_integrations.utils.nvidia.client import AvailableNvidiaCloudFunctions + + +class MockClient: + def query_function(self, func_id, payload): + inputs = payload["input"] + data = [{"index": i, "embedding": [0.1, 0.2, 0.3]} for i in range(len(inputs))] + return {"data": data, "usage": {"total_tokens": 4, "prompt_tokens": 4}} + + def available_functions(self): + return { + NvidiaEmbeddingModel.NVOLVE_40K.value: AvailableNvidiaCloudFunctions( + name=NvidiaEmbeddingModel.NVOLVE_40K.value, id="fake-id", status="ACTIVE" + ) + } + + +class TestNvidiaDocumentEmbedder: + def test_init_default(self, monkeypatch): + monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key") + embedder = NvidiaDocumentEmbedder(NvidiaEmbeddingModel.NVOLVE_40K) + + assert embedder.api_key == Secret.from_env_var("NVIDIA_API_KEY") + assert embedder.model == NvidiaEmbeddingModel.NVOLVE_40K + assert embedder.prefix == "" + assert embedder.suffix == "" + assert embedder.batch_size == 32 + assert embedder.progress_bar is True + assert embedder.meta_fields_to_embed == [] + assert embedder.embedding_separator == "\n" + + def test_init_with_parameters(self): + embedder = NvidiaDocumentEmbedder( + api_key=Secret.from_token("fake-api-key"), + model="playground_nvolveqa_40k", + prefix="prefix", + suffix="suffix", + batch_size=30, + progress_bar=False, + meta_fields_to_embed=["test_field"], + embedding_separator=" | ", + ) + + assert embedder.api_key == Secret.from_token("fake-api-key") + assert embedder.model == NvidiaEmbeddingModel.NVOLVE_40K + assert embedder.prefix == "prefix" + assert embedder.suffix == "suffix" + assert embedder.batch_size == 30 + assert embedder.progress_bar is False + assert embedder.meta_fields_to_embed == ["test_field"] + assert embedder.embedding_separator == " | " + + def test_init_fail_wo_api_key(self, monkeypatch): + monkeypatch.delenv("NVIDIA_API_KEY", raising=False) + with pytest.raises(ValueError): + NvidiaDocumentEmbedder(NvidiaEmbeddingModel.NVOLVE_40K) + + def test_init_fail_batch_size(self, monkeypatch): + with pytest.raises(ValueError): + NvidiaDocumentEmbedder(model="playground_nvolveqa_40k", batch_size=55) + + def test_to_dict(self, monkeypatch): + monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key") + component = NvidiaDocumentEmbedder("playground_nvolveqa_40k") + data = component.to_dict() + assert data == { + "type": "haystack_integrations.components.embedders.nvidia.document_embedder.NvidiaDocumentEmbedder", + "init_parameters": { + "api_key": {"env_vars": ["NVIDIA_API_KEY"], "strict": True, "type": "env_var"}, + "model": "playground_nvolveqa_40k", + "prefix": "", + "suffix": "", + "batch_size": 32, + "progress_bar": True, + "meta_fields_to_embed": [], + "embedding_separator": "\n", + }, + } + + def test_to_dict_with_custom_init_parameters(self, monkeypatch): + monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key") + component = NvidiaDocumentEmbedder( + model="playground_nvolveqa_40k", + prefix="prefix", + suffix="suffix", + batch_size=10, + progress_bar=False, + meta_fields_to_embed=["test_field"], + embedding_separator=" | ", + ) + data = component.to_dict() + assert data == { + "type": "haystack_integrations.components.embedders.nvidia.document_embedder.NvidiaDocumentEmbedder", + "init_parameters": { + "api_key": {"env_vars": ["NVIDIA_API_KEY"], "strict": True, "type": "env_var"}, + "model": "playground_nvolveqa_40k", + "prefix": "prefix", + "suffix": "suffix", + "batch_size": 10, + "progress_bar": False, + "meta_fields_to_embed": ["test_field"], + "embedding_separator": " | ", + }, + } + + 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) + ] + + embedder = NvidiaDocumentEmbedder( + "playground_nvolveqa_40k", + api_key=Secret.from_token("fake-api-key"), + meta_fields_to_embed=["meta_field"], + embedding_separator=" | ", + ) + + prepared_texts = embedder._prepare_texts_to_embed(documents) + + # note that newline is replaced by space + assert prepared_texts == [ + "meta_value 0 | document number 0:\ncontent", + "meta_value 1 | document number 1:\ncontent", + "meta_value 2 | document number 2:\ncontent", + "meta_value 3 | document number 3:\ncontent", + "meta_value 4 | document number 4:\ncontent", + ] + + def test_prepare_texts_to_embed_w_suffix(self): + documents = [Document(content=f"document number {i}") for i in range(5)] + + embedder = NvidiaDocumentEmbedder( + "playground_nvolveqa_40k", + api_key=Secret.from_token("fake-api-key"), + prefix="my_prefix ", + suffix=" my_suffix", + ) + + prepared_texts = embedder._prepare_texts_to_embed(documents) + + assert prepared_texts == [ + "my_prefix document number 0 my_suffix", + "my_prefix document number 1 my_suffix", + "my_prefix document number 2 my_suffix", + "my_prefix document number 3 my_suffix", + "my_prefix document number 4 my_suffix", + ] + + def test_embed_batch(self): + texts = ["text 1", "text 2", "text 3", "text 4", "text 5"] + + embedder = NvidiaDocumentEmbedder( + "playground_nvolveqa_40k", + api_key=Secret.from_token("fake-api-key"), + ) + embedder.client = MockClient() + embedder.warm_up() + + embeddings, metadata = embedder._embed_batch(texts_to_embed=texts, batch_size=2) + + assert isinstance(embeddings, list) + assert len(embeddings) == len(texts) + for embedding in embeddings: + assert isinstance(embedding, list) + assert len(embedding) == 3 + assert all(isinstance(x, float) for x in embedding) + + assert metadata == {"usage": {"prompt_tokens": 3 * 4, "total_tokens": 3 * 4}} + + def test_run(self): + docs = [ + Document(content="I love cheese", meta={"topic": "Cuisine"}), + Document(content="A transformer is a deep learning architecture", meta={"topic": "ML"}), + ] + + model = "playground_nvolveqa_40k" + embedder = NvidiaDocumentEmbedder( + api_key=Secret.from_token("fake-api-key"), + model=model, + prefix="prefix ", + suffix=" suffix", + meta_fields_to_embed=["topic"], + embedding_separator=" | ", + ) + embedder.client = MockClient() + embedder.warm_up() + + result = embedder.run(documents=docs) + + documents_with_embeddings = result["documents"] + metadata = result["meta"] + + assert isinstance(documents_with_embeddings, list) + assert len(documents_with_embeddings) == len(docs) + for doc in documents_with_embeddings: + assert isinstance(doc, Document) + assert isinstance(doc.embedding, list) + assert len(doc.embedding) == 3 + assert all(isinstance(x, float) for x in doc.embedding) + assert metadata == {"usage": {"prompt_tokens": 4, "total_tokens": 4}} + + def test_run_custom_batch_size(self): + docs = [ + Document(content="I love cheese", meta={"topic": "Cuisine"}), + Document(content="A transformer is a deep learning architecture", meta={"topic": "ML"}), + ] + model = "playground_nvolveqa_40k" + embedder = NvidiaDocumentEmbedder( + api_key=Secret.from_token("fake-api-key"), + model=model, + prefix="prefix ", + suffix=" suffix", + meta_fields_to_embed=["topic"], + embedding_separator=" | ", + batch_size=1, + ) + embedder.client = MockClient() + embedder.warm_up() + + result = embedder.run(documents=docs) + + documents_with_embeddings = result["documents"] + metadata = result["meta"] + + assert isinstance(documents_with_embeddings, list) + assert len(documents_with_embeddings) == len(docs) + for doc in documents_with_embeddings: + assert isinstance(doc, Document) + assert isinstance(doc.embedding, list) + assert len(doc.embedding) == 3 + assert all(isinstance(x, float) for x in doc.embedding) + + assert metadata == {"usage": {"prompt_tokens": 2 * 4, "total_tokens": 2 * 4}} + + def test_run_wrong_input_format(self): + embedder = NvidiaDocumentEmbedder("playground_nvolveqa_40k", api_key=Secret.from_token("fake-api-key")) + embedder.client = MockClient() + embedder.warm_up() + + string_input = "text" + list_integers_input = [1, 2, 3] + + with pytest.raises(TypeError, match="NvidiaDocumentEmbedder expects a list of Documents as input"): + embedder.run(documents=string_input) + + with pytest.raises(TypeError, match="NvidiaDocumentEmbedder expects a list of Documents as input"): + embedder.run(documents=list_integers_input) + + def test_run_on_empty_list(self): + embedder = NvidiaDocumentEmbedder("playground_nvolveqa_40k", api_key=Secret.from_token("fake-api-key")) + embedder.client = MockClient() + embedder.warm_up() + + empty_list_input = [] + result = embedder.run(documents=empty_list_input) + + assert result["documents"] is not None + assert not result["documents"] # empty list + + @pytest.mark.skipif( + not os.environ.get("NVIDIA_API_KEY", None), + reason="Export an env var called NVIDIA_API_KEY containing the Nvidia API key to run this test.", + ) + @pytest.mark.integration + def test_run_integration(self): + embedder = NvidiaDocumentEmbedder("playground_nvolveqa_40k") + 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_placeholder.py b/integrations/nvidia/tests/test_placeholder.py deleted file mode 100644 index 3ada1ee4e..000000000 --- a/integrations/nvidia/tests/test_placeholder.py +++ /dev/null @@ -1,2 +0,0 @@ -def test_placeholder(): - assert True diff --git a/integrations/nvidia/tests/test_text_embedder.py b/integrations/nvidia/tests/test_text_embedder.py new file mode 100644 index 000000000..b4239308b --- /dev/null +++ b/integrations/nvidia/tests/test_text_embedder.py @@ -0,0 +1,119 @@ +import os + +import pytest +from haystack.utils import Secret +from haystack_integrations.components.embedders.nvidia import NvidiaEmbeddingModel, NvidiaTextEmbedder +from haystack_integrations.utils.nvidia.client import AvailableNvidiaCloudFunctions + + +class MockClient: + def query_function(self, func_id, payload): + data = [{"index": 0, "embedding": [0.1, 0.2, 0.3]}] + return {"data": data, "usage": {"total_tokens": 4, "prompt_tokens": 4}} + + def available_functions(self): + return { + NvidiaEmbeddingModel.NVOLVE_40K.value: AvailableNvidiaCloudFunctions( + name=NvidiaEmbeddingModel.NVOLVE_40K.value, id="fake-id", status="ACTIVE" + ) + } + + +class TestNvidiaTextEmbedder: + def test_init_default(self, monkeypatch): + monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key") + embedder = NvidiaTextEmbedder(NvidiaEmbeddingModel.NVOLVE_40K) + + assert embedder.api_key == Secret.from_env_var("NVIDIA_API_KEY") + assert embedder.model == NvidiaEmbeddingModel.NVOLVE_40K + assert embedder.prefix == "" + assert embedder.suffix == "" + + def test_init_with_parameters(self): + embedder = NvidiaTextEmbedder( + api_key=Secret.from_token("fake-api-key"), + model="playground_nvolveqa_40k", + prefix="prefix", + suffix="suffix", + ) + assert embedder.api_key == Secret.from_token("fake-api-key") + assert embedder.model == NvidiaEmbeddingModel.NVOLVE_40K + assert embedder.prefix == "prefix" + assert embedder.suffix == "suffix" + + def test_init_fail_wo_api_key(self, monkeypatch): + monkeypatch.delenv("NVIDIA_API_KEY", raising=False) + with pytest.raises(ValueError): + NvidiaTextEmbedder(NvidiaEmbeddingModel.NVOLVE_40K) + + def test_to_dict(self, monkeypatch): + monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key") + component = NvidiaTextEmbedder(NvidiaEmbeddingModel.NVOLVE_40K) + data = component.to_dict() + assert data == { + "type": "haystack_integrations.components.embedders.nvidia.text_embedder.NvidiaTextEmbedder", + "init_parameters": { + "api_key": {"env_vars": ["NVIDIA_API_KEY"], "strict": True, "type": "env_var"}, + "model": "playground_nvolveqa_40k", + "prefix": "", + "suffix": "", + }, + } + + def test_to_dict_with_custom_init_parameters(self, monkeypatch): + monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key") + component = NvidiaTextEmbedder( + model=NvidiaEmbeddingModel.NVOLVE_40K, + prefix="prefix", + suffix="suffix", + ) + data = component.to_dict() + assert data == { + "type": "haystack_integrations.components.embedders.nvidia.text_embedder.NvidiaTextEmbedder", + "init_parameters": { + "api_key": {"env_vars": ["NVIDIA_API_KEY"], "strict": True, "type": "env_var"}, + "model": "playground_nvolveqa_40k", + "prefix": "prefix", + "suffix": "suffix", + }, + } + + def test_run(self): + embedder = NvidiaTextEmbedder( + "playground_nvolveqa_40k", api_key=Secret.from_token("fake-api-key"), prefix="prefix ", suffix=" suffix" + ) + embedder.client = MockClient() + embedder.warm_up() + result = embedder.run(text="The food was delicious") + + assert len(result["embedding"]) == 3 + assert all(isinstance(x, float) for x in result["embedding"]) + assert result["meta"] == { + "usage": {"prompt_tokens": 4, "total_tokens": 4}, + } + + def test_run_wrong_input_format(self): + embedder = NvidiaTextEmbedder("playground_nvolveqa_40k", api_key=Secret.from_token("fake-api-key")) + embedder.client = MockClient() + embedder.warm_up() + + list_integers_input = [1, 2, 3] + + with pytest.raises(TypeError, match="NvidiaTextEmbedder expects a string as an input"): + embedder.run(text=list_integers_input) + + @pytest.mark.skipif( + not os.environ.get("NVIDIA_API_KEY", None), + reason="Export an env var called NVIDIA_API_KEY containing the Nvidia API key to run this test.", + ) + @pytest.mark.integration + def test_run_integration(self): + embedder = NvidiaTextEmbedder("playground_nvolveqa_40k") + 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 From 5a339d44f73eed13908db8046e03b9079d51f24c Mon Sep 17 00:00:00 2001 From: Vladimir Blagojevic Date: Wed, 6 Mar 2024 10:25:22 +0100 Subject: [PATCH 3/8] docs: Final API docs touches (#538) * Final API docs touches * small fixes * remove error from API ref: MongoDBAtlasDocumentStoreError is never used --------- Co-authored-by: Stefano Fiorucci --- integrations/mongodb_atlas/pydoc/config.yml | 1 + .../document_stores/mongodb_atlas/document_store.py | 12 ++++++------ 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/integrations/mongodb_atlas/pydoc/config.yml b/integrations/mongodb_atlas/pydoc/config.yml index 782d662fa..e59993881 100644 --- a/integrations/mongodb_atlas/pydoc/config.yml +++ b/integrations/mongodb_atlas/pydoc/config.yml @@ -4,6 +4,7 @@ loaders: modules: [ "haystack_integrations.document_stores.mongodb_atlas.document_store", "haystack_integrations.document_stores.mongodb_atlas.filters", + "haystack_integrations.components.retrievers.mongodb_atlas.embedding_retriever", ] ignore_when_discovered: ["__init__"] processors: diff --git a/integrations/mongodb_atlas/src/haystack_integrations/document_stores/mongodb_atlas/document_store.py b/integrations/mongodb_atlas/src/haystack_integrations/document_stores/mongodb_atlas/document_store.py index 0d7116b3a..27cb853db 100644 --- a/integrations/mongodb_atlas/src/haystack_integrations/document_stores/mongodb_atlas/document_store.py +++ b/integrations/mongodb_atlas/src/haystack_integrations/document_stores/mongodb_atlas/document_store.py @@ -20,11 +20,11 @@ class MongoDBAtlasDocumentStore: """ - MongoDBAtlasDocumentStore is a DocumentStore implementation that uses [MongoDB Atlas](https://www.mongodb.com/atlas/database). - service that is easy to deploy, operate, and scale. + MongoDBAtlasDocumentStore is a DocumentStore implementation that uses + [MongoDB Atlas](https://www.mongodb.com/atlas/database) service that is easy to deploy, operate, and scale. To connect to MongoDB Atlas, you need to provide a connection string in the format: - "mongodb+srv://{mongo_atlas_username}:{mongo_atlas_password}@{mongo_atlas_host}/?{mongo_atlas_params_string}". + `"mongodb+srv://{mongo_atlas_username}:{mongo_atlas_password}@{mongo_atlas_host}/?{mongo_atlas_params_string}"`. This connection string can be obtained on the MongoDB Atlas Dashboard by clicking on the `CONNECT` button, selecting Python as the driver, and copying the connection string. The connection string can be provided as an environment @@ -39,7 +39,7 @@ class MongoDBAtlasDocumentStore: can support a chosen metric (i.e. cosine, dot product, or euclidean) and can be created in the Atlas web UI. For more details on MongoDB Atlas, see the official - MongoDB Atlas [documentation](https://www.mongodb.com/docs/atlas/getting-started/) + MongoDB Atlas [documentation](https://www.mongodb.com/docs/atlas/getting-started/). Usage example: ```python @@ -64,7 +64,7 @@ def __init__( Creates a new MongoDBAtlasDocumentStore instance. :param mongo_connection_string: MongoDB Atlas connection string in the format: - "mongodb+srv://{mongo_atlas_username}:{mongo_atlas_password}@{mongo_atlas_host}/?{mongo_atlas_params_string}". + `"mongodb+srv://{mongo_atlas_username}:{mongo_atlas_password}@{mongo_atlas_host}/?{mongo_atlas_params_string}"`. This can be obtained on the MongoDB Atlas Dashboard by clicking on the `CONNECT` button. This value will be read automatically from the env var "MONGO_CONNECTION_STRING". :param database_name: Name of the database to use. @@ -73,7 +73,7 @@ def __init__( :param vector_search_index: The name of the vector search index to use for vector search operations. Create a vector_search_index in the Atlas web UI and specify the init params of MongoDBAtlasDocumentStore. \ For more details refer to MongoDB - Atlas [documentation](https://www.mongodb.com/docs/atlas/atlas-vector-search/create-index/#std-label-avs-create-index) + Atlas [documentation](https://www.mongodb.com/docs/atlas/atlas-vector-search/create-index/#std-label-avs-create-index). :raises ValueError: If the collection name contains invalid characters. """ From 4998a7a6d840070d87f838c73cdf9010aba16865 Mon Sep 17 00:00:00 2001 From: Madeesh Kannan Date: Wed, 6 Mar 2024 11:44:08 +0100 Subject: [PATCH 4/8] fix: `nvidia-haystack`- Handle non-strict env var secrets correctly (#543) --- .../components/embedders/nvidia/document_embedder.py | 7 ++----- .../components/embedders/nvidia/text_embedder.py | 7 ++----- .../src/haystack_integrations/utils/nvidia/client.py | 11 ++++++++--- 3 files changed, 12 insertions(+), 13 deletions(-) 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 139a184b7..bbc68b492 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 @@ -65,9 +65,6 @@ def __init__( if isinstance(model, str): model = NvidiaEmbeddingModel.from_str(model) - resolved_api_key = api_key.resolve_value() - assert resolved_api_key is not None - # Upper-limit for the endpoint. if batch_size > MAX_INPUTS: msg = f"NVIDIA Cloud Functions currently support a maximum batch size of {MAX_INPUTS}." @@ -83,7 +80,7 @@ def __init__( self.embedding_separator = embedding_separator self.client = NvidiaCloudFunctionsClient( - api_key=resolved_api_key, + api_key=api_key, headers={ "Content-Type": "application/json", "Accept": "application/json", @@ -193,7 +190,7 @@ def run(self, documents: List[Document]): if not self._initialized: msg = "The embedding model has not been loaded. Please call warm_up() before running." raise RuntimeError(msg) - if not isinstance(documents, list) or documents and not isinstance(documents[0], Document): + elif not isinstance(documents, list) or documents and not isinstance(documents[0], Document): msg = ( "NvidiaDocumentEmbedder expects a list of Documents as input." "In case you want to embed a string, please use the NvidiaTextEmbedder." 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 43d62ed92..a2636b4b8 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 @@ -53,15 +53,12 @@ def __init__( if isinstance(model, str): model = NvidiaEmbeddingModel.from_str(model) - resolved_api_key = api_key.resolve_value() - assert resolved_api_key is not None - self.api_key = api_key self.model = model self.prefix = prefix self.suffix = suffix self.client = NvidiaCloudFunctionsClient( - api_key=resolved_api_key, + api_key=api_key, headers={ "Content-Type": "application/json", "Accept": "application/json", @@ -128,7 +125,7 @@ def run(self, text: str): if not self._initialized: msg = "The embedding model has not been loaded. Please call warm_up() before running." raise RuntimeError(msg) - if not isinstance(text, str): + elif not isinstance(text, str): msg = ( "NvidiaTextEmbedder expects a string as an input." "In case you want to embed a list of Documents, please use the NvidiaDocumentEmbedder." diff --git a/integrations/nvidia/src/haystack_integrations/utils/nvidia/client.py b/integrations/nvidia/src/haystack_integrations/utils/nvidia/client.py index 5227e8c45..e582b09ba 100644 --- a/integrations/nvidia/src/haystack_integrations/utils/nvidia/client.py +++ b/integrations/nvidia/src/haystack_integrations/utils/nvidia/client.py @@ -3,6 +3,7 @@ from typing import Dict, Optional import requests +from haystack.utils import Secret FUNCTIONS_ENDPOINT = "https://api.nvcf.nvidia.com/v2/nvcf/functions" INVOKE_ENDPOINT = "https://api.nvcf.nvidia.com/v2/nvcf/pexec/functions" @@ -19,13 +20,17 @@ class AvailableNvidiaCloudFunctions: class NvidiaCloudFunctionsClient: - def __init__(self, *, api_key: str, headers: Dict[str, str], timeout: int = 60): - self.api_key = api_key + def __init__(self, *, api_key: Secret, headers: Dict[str, str], timeout: int = 60): + self.api_key = api_key.resolve_value() + if self.api_key is None: + msg = "Nvidia Cloud Functions API key is not set." + raise ValueError(msg) + self.fetch_url_format = STATUS_ENDPOINT self.headers = copy.deepcopy(headers) self.headers.update( { - "Authorization": f"Bearer {api_key}", + "Authorization": f"Bearer {self.api_key}", } ) self.timeout = timeout From 6d340795137a3315cb58b21c1098260610ea548e Mon Sep 17 00:00:00 2001 From: Stefano Fiorucci Date: Wed, 6 Mar 2024 12:23:02 +0100 Subject: [PATCH 5/8] mongodb: improve example (#546) --- .../mongodb_atlas/examples/example.py | 25 ++++++++++++++++--- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/integrations/mongodb_atlas/examples/example.py b/integrations/mongodb_atlas/examples/example.py index 4b02bfd59..4cd3edc21 100644 --- a/integrations/mongodb_atlas/examples/example.py +++ b/integrations/mongodb_atlas/examples/example.py @@ -10,18 +10,24 @@ from haystack import Pipeline from haystack.components.converters import MarkdownToDocument -from haystack.components.embedders import SentenceTransformersDocumentEmbedder +from haystack.components.embedders import SentenceTransformersDocumentEmbedder, SentenceTransformersTextEmbedder from haystack.components.preprocessors import DocumentSplitter from haystack.components.writers import DocumentWriter +from haystack_integrations.components.retrievers.mongodb_atlas import MongoDBAtlasEmbeddingRetriever from haystack_integrations.document_stores.mongodb_atlas import MongoDBAtlasDocumentStore -# Provide your connection string -connection_string = input("Enter your MongoDB Atlas connection string: ") +# To use the MongoDBAtlasDocumentStore, you must have a running MongoDB Atlas database. +# For details, see https://www.mongodb.com/docs/atlas/getting-started/ + +# Once your database is set, set the environment variable `MONGO_CONNECTION_STRING` +# with the connection string to your MongoDB Atlas database. +# format: "mongodb+srv://{mongo_atlas_username}:{mongo_atlas_password}@{mongo_atlas_host}/?{mongo_atlas_params_string}". # Initialize the document store document_store = MongoDBAtlasDocumentStore( database_name="haystack_test", collection_name="test_collection", + vector_search_index="test_vector_search_index", ) # Create the indexing Pipeline and index some documents @@ -39,4 +45,15 @@ indexing.run({"converter": {"sources": file_paths}}) -print("Indexed documents:" + document_store.count_documents() + "\n - ".join(document_store.filter_documents())) + +# Create the querying Pipeline and try a query +querying = Pipeline() +querying.add_component("embedder", SentenceTransformersTextEmbedder()) +querying.add_component("retriever", MongoDBAtlasEmbeddingRetriever(document_store=document_store, top_k=3)) +querying.connect("embedder", "retriever") + +results = querying.run({"embedder": {"text": "What is a cross-encoder?"}}) + +for doc in results["retriever"]["documents"]: + print(doc) + print("-" * 10) From 2c6b218703211c32d7f29a0031494babdb831629 Mon Sep 17 00:00:00 2001 From: Tobias Wochinger Date: Wed, 6 Mar 2024 13:03:19 +0100 Subject: [PATCH 6/8] docs: review integrations sagemaker (#544) * refactor: remove reimplementing exceptions * docs: review docs * style: reformat * style: shorten line --- .../generators/amazon_sagemaker/errors.py | 32 +---------------- .../generators/amazon_sagemaker/sagemaker.py | 35 ++++++++++--------- 2 files changed, 19 insertions(+), 48 deletions(-) diff --git a/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/errors.py b/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/errors.py index 6c13d0fcb..e518c3c6d 100644 --- a/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/errors.py +++ b/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/errors.py @@ -1,46 +1,16 @@ -from typing import Optional - - class SagemakerError(Exception): """ - Error generated by the Amazon Sagemaker integration. + Parent class for all exceptions raised by the Sagemaker component """ - def __init__( - self, - message: Optional[str] = None, - ): - super().__init__() - if message: - self.message = message - - def __getattr__(self, attr): - # If self.__cause__ is None, it will raise the expected AttributeError - getattr(self.__cause__, attr) - - def __str__(self): - return self.message - - def __repr__(self): - return str(self) - class AWSConfigurationError(SagemakerError): """Exception raised when AWS is not configured correctly""" - def __init__(self, message: Optional[str] = None): - super().__init__(message=message) - class SagemakerNotReadyError(SagemakerError): """Exception for issues that occur during Sagemaker inference""" - def __init__(self, message: Optional[str] = None): - super().__init__(message=message) - class SagemakerInferenceError(SagemakerError): """Exception for issues that occur during Sagemaker inference""" - - def __init__(self, message: Optional[str] = None): - super().__init__(message=message) diff --git a/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/sagemaker.py b/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/sagemaker.py index c171ccdf6..2a04d6a2a 100644 --- a/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/sagemaker.py +++ b/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/sagemaker.py @@ -31,19 +31,17 @@ class SagemakerGenerator: [SageMaker JumpStart foundation models documentation](https://docs.aws.amazon.com/sagemaker/latest/dg/jumpstart-foundation-models-use.html). Usage example: - - Make sure your AWS credentials are set up correctly. You can use environment variables or a shared credentials file. - Then you can use the generator as follows: ```python + # Make sure your AWS credentials are set up correctly. You can use environment variables or a shared credentials + # file. Then you can use the generator as follows: from haystack_integrations.components.generators.amazon_sagemaker import SagemakerGenerator + generator = SagemakerGenerator(model="jumpstart-dft-hf-llm-falcon-7b-bf16") response = generator.run("What's Natural Language Processing? Be brief.") print(response) - ``` - ``` - >> {'replies': ['Natural Language Processing (NLP) is a branch of artificial intelligence that focuses on - >> the interaction between computers and human language. It involves enabling computers to understand, interpret, - >> and respond to natural human language in a way that is both meaningful and useful.'], 'meta': [{}]} + >>> {'replies': ['Natural Language Processing (NLP) is a branch of artificial intelligence that focuses on + >>> the interaction between computers and human language. It involves enabling computers to understand, interpret, + >>> and respond to natural human language in a way that is both meaningful and useful.'], 'meta': [{}]} ``` """ @@ -73,7 +71,6 @@ def __init__( :param model: The name for SageMaker Model Endpoint. :param aws_custom_attributes: Custom attributes to be passed to SageMaker, for example `{"accept_eula": True}` in case of Llama-2 models. - :param generation_kwargs: Additional keyword arguments for text generation. For a list of supported parameters see your model's documentation page, for example here for HuggingFace models: https://huggingface.co/blog/sagemaker-huggingface-llm#4-run-inference-and-chat-with-our-model @@ -121,15 +118,15 @@ def resolve_secret(secret: Optional[Secret]) -> Optional[str]: def _get_telemetry_data(self) -> Dict[str, Any]: """ Returns data that is sent to Posthog for usage analytics. - :returns: a dictionary with following keys: - - model: The name of the model. - + :returns: A dictionary with the following keys: + - `model`: The name of the model. """ return {"model": self.model} def to_dict(self) -> Dict[str, Any]: """ Serializes the component to a dictionary. + :returns: Dictionary with serialized data. """ @@ -149,10 +146,11 @@ def to_dict(self) -> Dict[str, Any]: def from_dict(cls, data) -> "SagemakerGenerator": """ Deserializes the component from a dictionary. + :param data: Dictionary to deserialize from. :returns: - Deserialized component. + Deserialized component. """ deserialize_secrets_inplace( data["init_parameters"], @@ -170,6 +168,7 @@ def _get_aws_session( ): """ Creates an AWS Session with the given parameters. + Checks if the provided AWS credentials are valid and can be used to connect to AWS. :param aws_access_key_id: AWS access key ID. @@ -200,8 +199,10 @@ def run(self, prompt: str, generation_kwargs: Optional[Dict[str, Any]] = None): :param prompt: The string prompt to use for text generation. :param generation_kwargs: Additional keyword arguments for text generation. These parameters will - potentially override the parameters passed in the `__init__` method. - + potentially override the parameters passed in the `__init__` method. + :raises ValueError: If the model response type is not a list of dictionaries or a single dictionary. + :raises SagemakerNotReadyError: If the SageMaker model is not ready to accept requests. + :raises SagemakerInferenceError: If the SageMaker Inference returns an error. :returns: A dictionary with the following keys: - `replies`: A list of strings containing the generated responses - `meta`: A list of dictionaries containing the metadata for each response. @@ -249,5 +250,5 @@ def run(self, prompt: str, generation_kwargs: Optional[Dict[str, Any]] = None): msg = f"Sagemaker model not ready: {res.text}" raise SagemakerNotReadyError(msg) from err - msg = f"SageMaker Inference returned an error. Status code: {res.status_code} Response body: {res.text}" - raise SagemakerInferenceError(msg, status_code=res.status_code) from err + msg = f"SageMaker Inference returned an error. Status code: {res.status_code}. Response body: {res.text}" + raise SagemakerInferenceError(msg) from err From d7ad329dac59b667de1726f21ed91db3580fd28b Mon Sep 17 00:00:00 2001 From: Tobias Wochinger Date: Wed, 6 Mar 2024 13:28:56 +0100 Subject: [PATCH 7/8] docs: review Elastic (#541) * docs: review Elastic * docs: correctly describe `DocumentStoreError` Co-authored-by: Stefano Fiorucci --------- Co-authored-by: Stefano Fiorucci --- .../elasticsearch/bm25_retriever.py | 24 +++--- .../elasticsearch/embedding_retriever.py | 16 ++-- .../elasticsearch/document_store.py | 81 ++++++++++--------- 3 files changed, 61 insertions(+), 60 deletions(-) diff --git a/integrations/elasticsearch/src/haystack_integrations/components/retrievers/elasticsearch/bm25_retriever.py b/integrations/elasticsearch/src/haystack_integrations/components/retrievers/elasticsearch/bm25_retriever.py index df1cb4a26..867d49c0e 100644 --- a/integrations/elasticsearch/src/haystack_integrations/components/retrievers/elasticsearch/bm25_retriever.py +++ b/integrations/elasticsearch/src/haystack_integrations/components/retrievers/elasticsearch/bm25_retriever.py @@ -54,13 +54,13 @@ def __init__( :param document_store: An instance of ElasticsearchDocumentStore. :param filters: Filters applied to the retrieved Documents, for more info - see `ElasticsearchDocumentStore.filter_documents`, defaults to None - :param fuzziness: Fuzziness parameter passed to Elasticsearch, defaults to "AUTO". - See the official - [documentation](https://www.elastic.co/guide/en/elasticsearch/reference/current/common-options.html#fuzziness) - for more details. - :param top_k: Maximum number of Documents to return, defaults to 10 - :param scale_score: If `True` scales the Document`s scores between 0 and 1, defaults to False + see `ElasticsearchDocumentStore.filter_documents`. + :param fuzziness: Fuzziness parameter passed to Elasticsearch. See the official + [documentation](https://www.elastic.co/guide/en/elasticsearch/reference/current/common-options.html#fuzziness) + for more details. + :param top_k: Maximum number of Documents to return. + :param scale_score: If `True` scales the Document`s scores between 0 and 1. + :raises ValueError: If `document_store` is not an instance of `ElasticsearchDocumentStore`. """ if not isinstance(document_store, ElasticsearchDocumentStore): @@ -97,7 +97,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "ElasticsearchBM25Retriever": :param data: Dictionary to deserialize from. :returns: - Deserialized component. + Deserialized component. """ data["init_parameters"]["document_store"] = ElasticsearchDocumentStore.from_dict( data["init_parameters"]["document_store"] @@ -109,11 +109,11 @@ def run(self, query: str, filters: Optional[Dict[str, Any]] = None, top_k: Optio """ Retrieve documents using the BM25 keyword-based algorithm. - :param query: String to search in Documents' text. - :param filters: Filters applied to the retrieved Documents. - :param top_k: Maximum number of Documents to return. + :param query: String to search in `Document`s' text. + :param filters: Filters applied to the retrieved `Document`s. + :param top_k: Maximum number of `Document` to return. :returns: A dictionary with the following keys: - - `documents`: List of Documents that match the query. + - `documents`: List of `Document`s that match the query. """ docs = self._document_store._bm25_retrieval( query=query, diff --git a/integrations/elasticsearch/src/haystack_integrations/components/retrievers/elasticsearch/embedding_retriever.py b/integrations/elasticsearch/src/haystack_integrations/components/retrievers/elasticsearch/embedding_retriever.py index d9f7f1fe6..fa292fe63 100644 --- a/integrations/elasticsearch/src/haystack_integrations/components/retrievers/elasticsearch/embedding_retriever.py +++ b/integrations/elasticsearch/src/haystack_integrations/components/retrievers/elasticsearch/embedding_retriever.py @@ -38,7 +38,7 @@ class ElasticsearchEmbeddingRetriever: result = retriever.run(query=query_embeddings) for doc in result["documents"]: - print(doc.content) + print(doc.content) ``` """ @@ -54,9 +54,9 @@ def __init__( Create the ElasticsearchEmbeddingRetriever component. :param document_store: An instance of ElasticsearchDocumentStore. - :param filters: Filters applied to the retrieved Documents. Defaults to None. - Filters are applied during the approximate kNN search to ensure that top_k matching documents are returned. - :param top_k: Maximum number of Documents to return, defaults to 10 + :param filters: Filters applied to the retrieved Documents. + Filters are applied during the approximate KNN search to ensure that top_k matching documents are returned. + :param top_k: Maximum number of Documents to return. :param num_candidates: Number of approximate nearest neighbor candidates on each shard. Defaults to top_k * 10. Increasing this value will improve search accuracy at the cost of slower search speeds. You can read more about it in the Elasticsearch @@ -95,7 +95,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "ElasticsearchEmbeddingRetriever": :param data: Dictionary to deserialize from. :returns: - Deserialized component. + Deserialized component. """ data["init_parameters"]["document_store"] = ElasticsearchDocumentStore.from_dict( data["init_parameters"]["document_store"] @@ -108,10 +108,10 @@ def run(self, query_embedding: List[float], filters: Optional[Dict[str, Any]] = Retrieve documents using a vector similarity metric. :param query_embedding: Embedding of the query. - :param filters: Filters applied to the retrieved Documents. - :param top_k: Maximum number of Documents to return. + :param filters: Filters applied to the retrieved `Document`s. + :param top_k: Maximum number of `Document`s to return. :returns: A dictionary with the following keys: - - `documents`: List of Documents most similar to the given query_embedding + - `documents`: List of `Document`s most similar to the given `query_embedding` """ docs = self._document_store._embedding_retrieval( query_embedding=query_embedding, diff --git a/integrations/elasticsearch/src/haystack_integrations/document_stores/elasticsearch/document_store.py b/integrations/elasticsearch/src/haystack_integrations/document_stores/elasticsearch/document_store.py index f50e2b1b3..0429f8811 100644 --- a/integrations/elasticsearch/src/haystack_integrations/document_stores/elasticsearch/document_store.py +++ b/integrations/elasticsearch/src/haystack_integrations/document_stores/elasticsearch/document_store.py @@ -38,13 +38,13 @@ class ElasticsearchDocumentStore: ElasticsearchDocumentStore is a Document Store for Elasticsearch. It can be used with Elastic Cloud or your own Elasticsearch cluster. - Usage example with Elastic Cloud: + Usage example (Elastic Cloud): ```python from haystack.document_store.elasticsearch import ElasticsearchDocumentStore document_store = ElasticsearchDocumentStore(cloud_id="YOUR_CLOUD_ID", api_key="YOUR_API_KEY") ``` - Usage example with a self-hosted Elasticsearch instance: + Usage example (self-hosted Elasticsearch instance): ```python from haystack.document_store.elasticsearch import ElasticsearchDocumentStore document_store = ElasticsearchDocumentStore(hosts="http://localhost:9200") @@ -69,8 +69,8 @@ def __init__( ): """ Creates a new ElasticsearchDocumentStore instance. - When no index is explicitly specified, it will use the default index "default". - It will also try to create that index if it doesn't exist yet. Otherwise it will use the existing one. + + It will also try to create that index if it doesn't exist yet. Otherwise, it will use the existing one. One can also set the similarity function used to compare Documents embeddings. This is mostly useful when using the `ElasticsearchDocumentStore` in a Pipeline with an `ElasticsearchEmbeddingRetriever`. @@ -81,14 +81,14 @@ def __init__( For the full list of supported kwargs, see the official Elasticsearch [reference](https://elasticsearch-py.readthedocs.io/en/stable/api.html#module-elasticsearch) - :param hosts: List of hosts running the Elasticsearch client. Defaults to None - :param index: Name of index in Elasticsearch, if it doesn't exist it will be created. Defaults to "default" + :param hosts: List of hosts running the Elasticsearch client. + :param index: Name of index in Elasticsearch. :param embedding_similarity_function: The similarity function used to compare Documents embeddings. - Defaults to "cosine". This parameter only takes effect if the index does not yet exist and is created. + This parameter only takes effect if the index does not yet exist and is created. To choose the most appropriate function, look for information about your embedding model. To understand how document scores are computed, see the Elasticsearch [documentation](https://www.elastic.co/guide/en/elasticsearch/reference/current/dense-vector.html#dense-vector-params) - :param **kwargs: Optional arguments that ``Elasticsearch`` takes. + :param **kwargs: Optional arguments that `Elasticsearch` takes. """ self._hosts = hosts self._client = Elasticsearch( @@ -140,7 +140,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "ElasticsearchDocumentStore": :param data: Dictionary to deserialize from. :returns: - Deserialized component. + Deserialized component. """ return default_from_dict(cls, data) @@ -186,7 +186,7 @@ def filter_documents(self, filters: Optional[Dict[str, Any]] = None) -> List[Doc :param filters: A dictionary of filters to apply. For more information on the structure of the filters, see the official Elasticsearch [documentation](https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl.html) - :returns: List of Documents that match the filters. + :returns: List of `Document`s that match the filters. """ if filters and "operator" not in filters and "conditions" not in filters: filters = convert(filters) @@ -197,13 +197,14 @@ def filter_documents(self, filters: Optional[Dict[str, Any]] = None) -> List[Doc def write_documents(self, documents: List[Document], policy: DuplicatePolicy = DuplicatePolicy.NONE) -> int: """ - Writes Documents to Elasticsearch. - - If policy is not specified or set to DuplicatePolicy.NONE, it will raise an exception if a document with the - same ID already exists in the document store. + Writes `Document`s to Elasticsearch. :param documents: List of Documents to write to the document store. :param policy: DuplicatePolicy to apply when a document with the same ID already exists in the document store. + :raises ValueError: If `documents` is not a list of `Document`s. + :raises DuplicateDocumentError: If a document with the same ID already exists in the document store and + `policy` is set to `DuplicatePolicy.FAIL` or `DuplicatePolicy.NONE`. + :raises DocumentStoreError: If an error occurs while writing the documents to the document store. :returns: Number of documents written to the document store. """ if len(documents) > 0: @@ -253,13 +254,15 @@ def write_documents(self, documents: List[Document], policy: DuplicatePolicy = D return documents_written - def _deserialize_document(self, hit: Dict[str, Any]) -> Document: + @staticmethod + def _deserialize_document(hit: Dict[str, Any]) -> Document: """ - Creates a Document from the search hit provided. + Creates a `Document` from the search hit provided. + This is mostly useful in self.filter_documents(). :param hit: A search hit from Elasticsearch. - :returns: Document created from the search hit. + :returns: `Document` created from the search hit. """ data = hit["_source"] @@ -271,12 +274,11 @@ def _deserialize_document(self, hit: Dict[str, Any]) -> Document: def delete_documents(self, document_ids: List[str]) -> None: """ - Deletes all documents with a matching document_ids from the document store. + Deletes all `Document`s with a matching `document_ids` from the document store. - :param document_ids: the object_ids to delete + :param document_ids: the object IDs to delete """ - # helpers.bulk( client=self._client, actions=({"_op_type": "delete", "_id": id_} for id_ in document_ids), @@ -295,27 +297,25 @@ def _bm25_retrieval( scale_score: bool = False, ) -> List[Document]: """ - Elasticsearch by defaults uses BM25 search algorithm. + Retrieves `Document`s from Elasticsearch using the BM25 search algorithm. + Even though this method is called `bm25_retrieval` it searches for `query` using the search algorithm `_client` was configured with. - This method is not mean to be part of the public interface of + This method is not meant to be part of the public interface of `ElasticsearchDocumentStore` nor called directly. `ElasticsearchBM25Retriever` uses this method directly and is the public interface for it. - `query` must be a non-empty string, otherwise a `ValueError` will be raised. - - :param query: String to search in saved Documents' text. - :param filters: Filters applied to the retrieved Documents, for more info - see `ElasticsearchDocumentStore.filter_documents`, defaults to None - :param fuzziness: Fuzziness parameter passed to Elasticsearch, defaults to "AUTO". - see the official - [documentation](https://www.elastic.co/guide/en/elasticsearch/reference/current/common-options.html#fuzziness) - for valid values. - :param top_k: Maximum number of Documents to return, defaults to 10 - :param scale_score: If `True` scales the Document`s scores between 0 and 1, defaults to False + :param query: String to search in saved `Document`s' text. + :param filters: Filters applied to the retrieved `Document`s, for more info + see `ElasticsearchDocumentStore.filter_documents`. + :param fuzziness: Fuzziness parameter passed to Elasticsearch. See the official + [documentation](https://www.elastic.co/guide/en/elasticsearch/reference/current/common-options.html#fuzziness) + for valid values. + :param top_k: Maximum number of `Document`s to return. + :param scale_score: If `True` scales the `Document``s scores between 0 and 1. :raises ValueError: If `query` is an empty string - :returns: List of Document that match `query` + :returns: List of `Document` that match `query` """ if not query: @@ -361,22 +361,23 @@ def _embedding_retrieval( ) -> List[Document]: """ Retrieves documents that are most similar to the query embedding using a vector similarity metric. + It uses the Elasticsearch's Approximate k-Nearest Neighbors search algorithm. - This method is not mean to be part of the public interface of + This method is not meant to be part of the public interface of `ElasticsearchDocumentStore` nor called directly. `ElasticsearchEmbeddingRetriever` uses this method directly and is the public interface for it. :param query_embedding: Embedding of the query. - :param filters: Filters applied to the retrieved Documents. Defaults to None. + :param filters: Filters applied to the retrieved `Document`s. Filters are applied during the approximate kNN search to ensure that top_k matching documents are returned. - :param top_k: Maximum number of Documents to return, defaults to 10 + :param top_k: Maximum number of `Document`s to return. :param num_candidates: Number of approximate nearest neighbor candidates on each shard. Defaults to top_k * 10. Increasing this value will improve search accuracy at the cost of slower search speeds. You can read more about it in the Elasticsearch - [documentation](https://www.elastic.co/guide/en/elasticsearch/reference/current/knn-search.html#tune-approximate-knn-for-speed-accuracy) - :raises ValueError: If `query_embedding` is an empty list - :returns: List of Document that are most similar to `query_embedding` + [documentation](https://www.elastic.co/guide/en/elasticsearch/reference/current/knn-search.html#tune-approximate-knn-for-speed-accuracy) + :raises ValueError: If `query_embedding` is an empty list. + :returns: List of `Document` that are most similar to `query_embedding`. """ if not query_embedding: From 67decea0523dc15b2268bc7953694a4764f7515a Mon Sep 17 00:00:00 2001 From: Vladimir Blagojevic Date: Wed, 6 Mar 2024 14:08:55 +0100 Subject: [PATCH 8/8] Fix API docs (#540) --- .../embedders/fastembed/fastembed_document_embedder.py | 6 +++--- .../embedders/fastembed/fastembed_text_embedder.py | 7 +++---- 2 files changed, 6 insertions(+), 7 deletions(-) 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 4af8e1bbe..b5dd71231 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 @@ -68,11 +68,11 @@ def __init__( Create an FastembedDocumentEmbedder component. :param model: Local path or name of the model in Hugging Face's model hub, - such as ``'BAAI/bge-small-en-v1.5'``. - :param cache_dir (str, optional): The path to the cache directory. + such as `BAAI/bge-small-en-v1.5`. + :param cache_dir: The path to the cache directory. Can be set using the `FASTEMBED_CACHE_PATH` env variable. Defaults to `fastembed_cache` in the system's temp directory. - :param threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None. + :param threads: The number of threads single onnxruntime session can use. Defaults to None. :param prefix: A string to add to the beginning of each text. :param suffix: A string to add to the end of each text. :param batch_size: Number of strings to encode at once. 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 13a89d1ce..743884ec1 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 @@ -42,12 +42,11 @@ def __init__( """ Create a FastembedTextEmbedder component. - :param model: Local path or name of the model in Fastembed's model hub, - such as ``'BAAI/bge-small-en-v1.5'``. - :param cache_dir (str, optional): The path to the cache directory. + :param model: Local path or name of the model in Fastembed's model hub, such as `BAAI/bge-small-en-v1.5` + :param cache_dir: The path to the cache directory. Can be set using the `FASTEMBED_CACHE_PATH` env variable. Defaults to `fastembed_cache` in the system's temp directory. - :param threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None. + :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 prefix: A string to add to the beginning of each text. :param suffix: A string to add to the end of each text.