From d7ad329dac59b667de1726f21ed91db3580fd28b Mon Sep 17 00:00:00 2001 From: Tobias Wochinger Date: Wed, 6 Mar 2024 13:28:56 +0100 Subject: [PATCH 1/4] 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 2/4] 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. From 2a91820f1fcc319d482bb0d3558a72e32ecfe82f Mon Sep 17 00:00:00 2001 From: Tobias Wochinger Date: Wed, 6 Mar 2024 15:27:59 +0100 Subject: [PATCH 3/4] docs: review integrations bedrock (#550) * docs: review bedrock * refactor: drop duplicate code * refactor: make invoke public again * docs: add missing colon Co-authored-by: Stefano Fiorucci * docs: add missing backticks Co-authored-by: Stefano Fiorucci --------- Co-authored-by: Stefano Fiorucci --- .../common/amazon_bedrock/errors.py | 30 ------------------- .../amazon_bedrock/document_embedder.py | 25 ++++++++++++---- .../embedders/amazon_bedrock/text_embedder.py | 26 ++++++++++++---- .../generators/amazon_bedrock/adapters.py | 24 +++++++++++++-- .../generators/amazon_bedrock/generator.py | 10 +++++-- .../generators/amazon_bedrock/handlers.py | 2 ++ 6 files changed, 70 insertions(+), 47 deletions(-) diff --git a/integrations/amazon_bedrock/src/haystack_integrations/common/amazon_bedrock/errors.py b/integrations/amazon_bedrock/src/haystack_integrations/common/amazon_bedrock/errors.py index aa8a3f6e4..130c7d5e6 100644 --- a/integrations/amazon_bedrock/src/haystack_integrations/common/amazon_bedrock/errors.py +++ b/integrations/amazon_bedrock/src/haystack_integrations/common/amazon_bedrock/errors.py @@ -1,6 +1,3 @@ -from typing import Optional - - class AmazonBedrockError(Exception): """ Any error generated by the Amazon Bedrock integration. @@ -10,41 +7,14 @@ class AmazonBedrockError(Exception): `AmazonBedrockError.message` will exist and have the expected content. """ - 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(AmazonBedrockError): """Exception raised when AWS is not configured correctly""" - def __init__(self, message: Optional[str] = None): - super().__init__(message=message) - class AmazonBedrockConfigurationError(AmazonBedrockError): """Exception raised when AmazonBedrock node is not configured correctly""" - def __init__(self, message: Optional[str] = None): - super().__init__(message=message) - class AmazonBedrockInferenceError(AmazonBedrockError): """Exception for issues that occur in the Bedrock inference node""" - - def __init__(self, message: Optional[str] = None): - super().__init__(message=message) 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 8cf98cd45..a5621cbd2 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 @@ -77,13 +77,12 @@ def __init__( :param model: The embedding model to use. The model has to be specified in the format outlined in the Amazon Bedrock [documentation](https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids.html). - :type model: Literal["amazon.titan-embed-text-v1", "cohere.embed-english-v3", "cohere.embed-multilingual-v3"] :param aws_access_key_id: AWS access key ID. :param aws_secret_access_key: AWS secret access key. :param aws_session_token: AWS session token. :param aws_region_name: AWS region name. :param aws_profile_name: AWS profile name. - :param batch_size: Number of Documents to encode at once. Default is 32. + :param batch_size: Number of Documents to encode at once. Only Cohere models support batch inference. This parameter is ignored for Amazon Titan models. :param progress_bar: Whether to show a progress bar or not. Can be helpful to disable in production deployments to keep the logs clean. @@ -91,6 +90,8 @@ def __init__( :param embedding_separator: Separator used to concatenate the meta fields to the Document text. :param kwargs: Additional parameters to pass for model inference. For example, `input_type` and `truncate` for Cohere models. + :raises ValueError: If the model is not supported. + :raises AmazonBedrockConfigurationError: If the AWS environment is not configured correctly. """ if not model or model not in SUPPORTED_EMBEDDING_MODELS: @@ -218,6 +219,13 @@ def _embed_titan(self, documents: List[Document]) -> List[Document]: @component.output_types(documents=List[Document]) def run(self, documents: List[Document]): + """Embed the provided `Document`s using the specified model. + + :param documents: The `Document`s to embed. + :returns: A dictionary with the following keys: + - `documents`: The `Document`s with the `embedding` field populated. + :raises AmazonBedrockInferenceError: If the inference fails. + """ if not isinstance(documents, list) or documents and not isinstance(documents[0], Document): msg = ( "AmazonBedrockDocumentEmbedder expects a list of Documents as input." @@ -234,8 +242,10 @@ def run(self, documents: List[Document]): def to_dict(self) -> Dict[str, Any]: """ - Serialize this component to a dictionary. - :returns: The serialized component as a dictionary. + Serializes the component to a dictionary. + + :returns: + Dictionary with serialized data. """ return default_to_dict( self, @@ -255,7 +265,12 @@ def to_dict(self) -> Dict[str, Any]: @classmethod def from_dict(cls, data: Dict[str, Any]) -> "AmazonBedrockDocumentEmbedder": """ - Deserialize this component from a dictionary. + Deserializes the component from a dictionary. + + :param data: + Dictionary to deserialize from. + :returns: + Deserialized component. """ deserialize_secrets_inplace( data["init_parameters"], 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 ed6768737..91a9e3b72 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 @@ -66,7 +66,6 @@ def __init__( :param model: The embedding model to use. The model has to be specified in the format outlined in the Amazon Bedrock [documentation](https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids.html). - :type model: Literal["amazon.titan-embed-text-v1", "cohere.embed-english-v3", "cohere.embed-multilingual-v3"] :param aws_access_key_id: AWS access key ID. :param aws_secret_access_key: AWS secret access key. :param aws_session_token: AWS session token. @@ -74,6 +73,8 @@ def __init__( :param aws_profile_name: AWS profile name. :param kwargs: Additional parameters to pass for model inference. For example, `input_type` and `truncate` for Cohere models. + :raises ValueError: If the model is not supported. + :raises AmazonBedrockConfigurationError: If the AWS environment is not configured correctly. """ if not model or model not in SUPPORTED_EMBEDDING_MODELS: msg = "Please provide a valid model from the list of supported models: " + ", ".join( @@ -110,6 +111,14 @@ def resolve_secret(secret: Optional[Secret]) -> Optional[str]: @component.output_types(embedding=List[float]) def run(self, text: str): + """Embeds the input text using the Amazon Bedrock model. + + :param text: The input text to embed. + :returns: A dictionary with the following keys: + - `embedding`: The embedding of the input text. + :raises TypeError: If the input text is not a string. + :raises AmazonBedrockInferenceError: If the model inference fails. + """ if not isinstance(text, str): msg = ( "AmazonBedrockTextEmbedder expects a string as an input." @@ -153,8 +162,10 @@ def run(self, text: str): def to_dict(self) -> Dict[str, Any]: """ - Serialize this component to a dictionary. - :returns: The serialized component as a dictionary. + Serializes the component to a dictionary. + + :returns: + Dictionary with serialized data. """ return default_to_dict( self, @@ -170,9 +181,12 @@ def to_dict(self) -> Dict[str, Any]: @classmethod def from_dict(cls, data: Dict[str, Any]) -> "AmazonBedrockTextEmbedder": """ - Deserialize this component from a dictionary. - :param data: The dictionary representation of this component. - :returns: The deserialized component instance. + Deserializes the component from a dictionary. + + :param data: + Dictionary to deserialize from. + :returns: + Deserialized component. """ 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 f842f0ef5..0fbf04de7 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 @@ -103,6 +103,10 @@ def prepare_body(self, prompt: str, **inference_kwargs) -> Dict[str, Any]: Prepares the body for the Claude model :param prompt: The prompt to be sent to the model. + :param inference_kwargs: Additional keyword arguments passed to the handler. + :returns: A dictionary with the following keys: + - `prompt`: The prompt to be sent to the model. + - specified inference parameters. """ default_params = { "max_tokens_to_sample": self.max_length, @@ -146,7 +150,9 @@ 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. - :returns: A dictionary containing the body for the request. + :returns: A dictionary with the following keys: + - `prompt`: The prompt to be sent to the model. + - specified inference parameters. """ default_params = { "max_tokens": self.max_length, @@ -191,6 +197,14 @@ class AI21LabsJurassic2Adapter(BedrockModelAdapter): """ def prepare_body(self, prompt: str, **inference_kwargs) -> Dict[str, Any]: + """Prepares the body for the Jurassic 2 model. + + :param prompt: The prompt to be sent to the model. + :param inference_kwargs: Additional keyword arguments passed to the handler. + :returns: A dictionary with the following keys: + - `prompt`: The prompt to be sent to the model. + - specified inference parameters. + """ default_params = { "maxTokens": self.max_length, "stopSequences": None, @@ -226,7 +240,9 @@ 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. - :returns: A dictionary containing the body for the request. + :returns: A dictionary with the following keys + - `inputText`: The prompt to be sent to the model. + - specified inference parameters. """ default_params = { "maxTokenCount": self.max_length, @@ -270,7 +286,9 @@ 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. - :returns: A dictionary containing the body for the request. + :returns: A dictionary with the following keys: + - `prompt`: The prompt to be sent to the model. + - specified inference parameters. """ default_params = { "max_gen_len": self.max_length, 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 f6af48ae1..0b1b73812 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 @@ -84,7 +84,9 @@ def __init__( :param aws_profile_name: The AWS profile name. :param max_length: The maximum length of the generated text. :param kwargs: Additional keyword arguments to be passed to the model. - + :raises ValueError: If the model name is empty or None. + :raises AmazonBedrockConfigurationError: If the AWS environment is not configured correctly or the model is + not supported. """ if not model: msg = "'model' cannot be None or empty string" @@ -226,7 +228,9 @@ 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. :returns: A dictionary with the following keys: - - `replies`: A list of generated responses (strings). + - `replies`: A list of generated responses. + :raises ValueError: If the prompt is empty or None. + :raises AmazonBedrockInferenceError: If the model cannot be invoked. """ return {"replies": self.invoke(prompt=prompt, **(generation_kwargs or {}))} @@ -269,7 +273,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "AmazonBedrockGenerator": :param data: Dictionary to deserialize from. :returns: - Deserialized component. + Deserialized component. """ deserialize_secrets_inplace( data["init_parameters"], 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 ddc276264..f4dc1aa4f 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 @@ -15,6 +15,8 @@ def __init__(self, tokenizer: Union[str, PreTrainedTokenizerBase], model_max_len :param tokenizer: The tokenizer to be used to tokenize the prompt. :param model_max_length: The maximum length of the prompt and answer tokens combined. :param max_length: The maximum length of the answer tokens. + :raises ValueError: If the tokenizer is not a string or a `PreTrainedTokenizer` or `PreTrainedTokenizerFast` + instance. """ if isinstance(tokenizer, str): self.tokenizer = AutoTokenizer.from_pretrained(tokenizer) From 1032b2c7566ffb3d6e8cac0541f1984bdaeb64aa Mon Sep 17 00:00:00 2001 From: Stefano Fiorucci Date: Wed, 6 Mar 2024 16:16:00 +0100 Subject: [PATCH 4/4] Ollama: change testing workflow (#551) * try changing workflow * show coverage * add healtcheck with timeout --- .github/workflows/ollama.yml | 33 +++++++++++++++++++----------- integrations/ollama/pyproject.toml | 17 +++++++-------- 2 files changed, 28 insertions(+), 22 deletions(-) diff --git a/.github/workflows/ollama.yml b/.github/workflows/ollama.yml index 055f63dd0..28b522890 100644 --- a/.github/workflows/ollama.yml +++ b/.github/workflows/ollama.yml @@ -33,16 +33,31 @@ jobs: matrix: os: [ubuntu-latest] python-version: ["3.9","3.10","3.11"] - services: - ollama: - image: ollama/ollama:latest - options: --name ollama - ports: - - 11434:11434 steps: - uses: actions/checkout@v4 + - name: Install Ollama and pull the required models + run: | + curl -fsSL https://ollama.com/install.sh | sh + ollama serve & + + # Check if the service is up and running with a timeout of 60 seconds + timeout=60 + while [ $timeout -gt 0 ] && ! curl -sSf http://localhost:11434/ > /dev/null; do + echo "Waiting for Ollama service to start..." + sleep 5 + ((timeout-=5)) + done + + if [ $timeout -eq 0 ]; then + echo "Timed out waiting for Ollama service to start." + exit 1 + fi + + ollama pull ${{ env.LLM_FOR_TESTS }} + ollama pull ${{ env.EMBEDDER_FOR_TESTS }} + - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v5 with: @@ -55,12 +70,6 @@ jobs: if: matrix.python-version == '3.9' run: hatch run lint:all - - name: Pull the LLM in the Ollama service - run: docker exec ollama ollama pull ${{ env.LLM_FOR_TESTS }} - - - name: Pull the Embedding Model in the Ollama service - run: docker exec ollama ollama pull ${{ env.EMBEDDER_FOR_TESTS }} - - name: Generate docs if: matrix.python-version == '3.9' && runner.os == 'Linux' run: hatch run docs diff --git a/integrations/ollama/pyproject.toml b/integrations/ollama/pyproject.toml index e3bb738b6..17c9d46b2 100644 --- a/integrations/ollama/pyproject.toml +++ b/integrations/ollama/pyproject.toml @@ -157,20 +157,17 @@ ban-relative-imports = "parents" [tool.coverage.run] -source_pkgs = ["src", "tests"] +source = ["haystack_integrations"] branch = true -parallel = true - - -[tool.coverage.paths] -ollama_haystack = ["src/haystack_integrations", "*/ollama-haystack/src"] -tests = ["tests", "*/ollama-haystack/tests"] +parallel = false [tool.coverage.report] +omit = ["*/tests/*", "*/__init__.py"] +show_missing=true exclude_lines = [ - "no cov", - "if __name__ == .__main__.:", - "if TYPE_CHECKING:", + "no cov", + "if __name__ == .__main__.:", + "if TYPE_CHECKING:", ] [tool.pytest.ini_options]