feat: Add AWS bedrock embeddings to embedding encoder

haystack/nodes/retriever/_embedding_encoder.py

@@ -42,6 +42,8 @@
     from haystack.modeling.infer import Inferencer
     from haystack.nodes.retriever._losses import _TRAINING_LOSSES
 
+with LazyImport(message="Run 'pip install boto3'") as boto3_import:
+    import boto3
 
 COHERE_TIMEOUT = float(os.environ.get(HAYSTACK_REMOTE_API_TIMEOUT_SEC, 30))
 COHERE_BACKOFF = int(os.environ.get(HAYSTACK_REMOTE_API_BACKOFF_SEC, 10))
@@ -434,11 +436,79 @@ def save(self, save_dir: Union[Path, str]):
         raise NotImplementedError(f"Saving is not implemented for {self.__class__}")
 
 
+class _BedrockEmbeddingEncoder(_BaseEmbeddingEncoder):
+    def __init__(self, retriever: "EmbeddingRetriever"):
+        boto3_import.check()
+
+        # See https://docs.aws.amazon.com/bedrock/latest/userguide/embeddings.html for more details
+        # The maximum input text is 8K tokens and the maximum output vector length is 1536
+        # Bedrock embeddings do not support batch operations
+        self.model: str = "amazon.titan-embed-text-v1"
+        self.aws_config = retriever.aws_config
+        self.client = self.initialize_boto3_client()
+
+    def initialize_boto3_client(self):
+        if self.aws_config:
+            access_key_id = self.aws_config.get("aws_access_key_id")
+            secret_access_key = self.aws_config.get("aws_secret_access_key")
+            region = self.aws_config.get("region")
+            try:
+                return boto3.client(
+                    "bedrock-runtime",
+                    aws_access_key_id=access_key_id,
+                    aws_secret_access_key=secret_access_key,
+                    region_name=region,
+                )
+            except Exception as e:
+                raise ValueError(f"AWS client error {e}")
+        else:
+            raise ValueError("Please pass boto3.client(bedrock-runtime) credentials configuration")
+
+    def embed(self, text: str) -> np.ndarray:
+        input_body = {}
+        input_body["inputText"] = text
+        body = json.dumps(input_body)
+        response = self.client.invoke_model(
+            body=body, modelId=self.model, accept="application/json", contentType="application/json"
+        )
+
+        response_body = json.loads(response.get("body").read())
+        return np.array(response_body.get("embedding"))
+
+    def embed_queries(self, queries: List[str]) -> np.ndarray:
+        all_embeddings = []
+        for q in queries:
+            generated_embeddings = self.embed(q)
+            all_embeddings.append(generated_embeddings)
+        return np.concatenate(all_embeddings)
+
+    def embed_documents(self, docs: List[Document]) -> np.ndarray:
+        return self.embed_queries([d.content for d in docs])
+
+    def train(
+        self,
+        training_data: List[Dict[str, Any]],
+        learning_rate: float = 2e-5,
+        n_epochs: int = 1,
+        num_warmup_steps: Optional[int] = None,
+        batch_size: int = 16,
+        train_loss: Literal["mnrl", "margin_mse"] = "mnrl",
+        num_workers: int = 0,
+        use_amp: bool = False,
+        **kwargs,
+    ):
+        raise NotImplementedError(f"Training is not implemented for {self.__class__}")
+
+    def save(self, save_dir: Union[Path, str]):
+        raise NotImplementedError(f"Saving is not implemented for {self.__class__}")
+
+
 _EMBEDDING_ENCODERS: Dict[str, Callable] = {
     "farm": _DefaultEmbeddingEncoder,
     "transformers": _DefaultEmbeddingEncoder,
     "sentence_transformers": _SentenceTransformersEmbeddingEncoder,
     "retribert": _RetribertEmbeddingEncoder,
     "openai": _OpenAIEmbeddingEncoder,
     "cohere": _CohereEmbeddingEncoder,
+    "bedrock": _BedrockEmbeddingEncoder,
 }

haystack/nodes/retriever/dense.py

@@ -1468,6 +1468,7 @@ def __init__(
         azure_deployment_name: Optional[str] = None,
         api_base: str = "https://api.openai.com/v1",
         openai_organization: Optional[str] = None,
+        aws_config: Optional[Dict[str, Any]] = None,
     ):
         """
         :param document_store: An instance of DocumentStore from which to retrieve documents.
@@ -1532,6 +1533,7 @@ def __init__(
                                      will not be used.
         :param api_base: The OpenAI API base URL, defaults to `"https://api.openai.com/v1"`.
         :param openai_organization: The OpenAI-Organization ID, defaults to `None`. For more details, see OpenAI
+        :param aws_config: The aws_config contains {aws_access_key, aws_secret_key, aws_region } to use with the boto3 client for an AWS Bedrock retriever. Defaults to 'None'.
         [documentation](https://platform.openai.com/docs/api-reference/requesting-organization).
         """
         torch_and_transformers_import.check()
@@ -1565,6 +1567,7 @@ def __init__(
         self.azure_base_url = azure_base_url
         self.azure_deployment_name = azure_deployment_name
         self.openai_organization = openai_organization
+        self.aws_config = aws_config
         self.model_format = (
             self._infer_model_format(model_name_or_path=embedding_model, use_auth_token=use_auth_token)
             if model_format is None
@@ -1892,6 +1895,8 @@ def _infer_model_format(model_name_or_path: str, use_auth_token: Optional[Union[
             return "openai"
         if model_name_or_path in COHERE_EMBEDDING_MODELS:
             return "cohere"
+        if model_name_or_path == "bedrock":
+            return "bedrock"
         # Check if model name is a local directory with sentence transformers config file in it
         if Path(model_name_or_path).exists():
             if Path(f"{model_name_or_path}/config_sentence_transformers.json").exists():

releasenotes/notes/aws-bedrock-embedding-encoder-a978884c1a2c8237.yaml

@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    Adding Bedrock Embeddings Encoder to use as a retriever.