Add Amazon Bedrock Embedding function

chromadb/utils/embedding_functions.py

@@ -581,6 +581,105 @@ def __call__(self, input: Union[Documents, Images]) -> Embeddings:
         return embeddings
 
 
+class AmazonBedrockEmbeddingFunction(EmbeddingFunction[Documents]):
+    def __init__(
+        self,
+        profile_name: Optional[str] = None,
+        region: Optional[str] = None,
+        model_name: str = "amazon.titan-embed-text-v1",
+    ):
+        """Initialize AmazonBedrockEmbeddingFucntion.
+
+        Args:
+            profile_name (str, optional): The name of a profile to use. If not given, then the default profile is used, defaults to None
+            region (str, optional): Default region when creating new connections, defaults to None
+            model_name (str, optional): Identifier of the model, defaults to "amazon.titan-embed-text-v1"
+        """
+
+        self._model_name = model_name
+
+        try:
+            import boto3
+            from botocore.config import Config
+        except ImportError:
+            raise ValueError(
+                "The boto3 python package is not installed. Please install it with `pip install boto3`"
+            )
+
+        if not region:
+            target_region = os.environ.get(
+                "AWS_REGION", os.environ.get("AWS_DEFAULT_REGION")
+            )
+        else:
+            target_region = region
+
+        session_kwargs = {"region_name": target_region}
+        client_kwargs = {**session_kwargs}
+
+        if profile_name:
+            target_profile = profile_name
+        else:
+            target_profile = os.environ.get("AWS_PROFILE")
+
+        if target_profile:
+            session_kwargs["profile_name"] = target_profile
+
+        retry_config = Config(
+            region_name=target_region,
+            retries={
+                "max_attempts": 10,
+                "mode": "standard",
+            },
+        )
+
+        session = boto3.Session(**session_kwargs)
+        self._client = session.client(
+            service_name="bedrock-runtime", config=retry_config, **client_kwargs
+        )
+
+    def __call__(self, input: Documents) -> Embeddings:
+        """Get the embeddings for a list of texts.
+
+        Args:
+            input (Documents): A list of texts to get embeddings for.
+
+        Returns:
+            Embeddings: The embeddings for the texts.
+
+        Example:
+            >>> bedrock = AmazonBedrockEmbeddingFunction(profile_name="profile")
+            >>> texts = ["Hello, world!", "How are you?"]
+            >>> embeddings = bedrock(texts)
+        """
+        embeddings = []
+        for text in input:
+            embeddings.append(self._invoke(text))
+        return embeddings
+
+    def _invoke(
+        self,
+        text: str,
+    ) -> Embedding:
+        import json
+
+        input_body = {"inputText": text}
+        body = json.dumps(input_body)
+        accept = "application/json"
+        content_type = "application/json"
+        try:
+            response = self._client.invoke_model(
+                body=body,
+                modelId=self._model_name,
+                accept=accept,
+                contentType=content_type,
+            )
+            embedding = json.load(response.get("body")).get("embedding")
+        except Exception as e:
+            raise ValueError(f"Error raised by bedrock service: {e}")
+
+        return embedding
+
+
 # List of all classes in this module
 _classes = [
     name