feat: AmazonBedrockEmbeddingFunction support for Cohere Embed models

chromadb/utils/embedding_functions.py

@@ -742,27 +742,54 @@ def __init__(
         """
 
         self._model_name = model_name
+        self._model_provider = self._model_name.split('.')[0]
+
+        self._bedrock_client = session.client(
+            service_name="bedrock",
+            **kwargs
+        )
+
+        self._model_details = self._bedrock_client.get_foundation_model(modelIdentifier=self._model_name)['modelDetails']
+        assert "EMBEDDING" in self._model_details['outputModalities'], f"{self._model_name} doesn't have embedding modality output!"
 
-        self._client = session.client(
+        self._bedrock_runtime_client = session.client(
             service_name="bedrock-runtime",
             **kwargs,
         )
+
+    def call_model(self, body) -> dict:
+        body = json.dumps(body)
+        response = self._bedrock_runtime_client.invoke_model(
+            body=body,
+            modelId=self._model_name,
+            accept="application/json",
+            contentType="application/json",
+        )
+        return response
 
     def __call__(self, input: Documents) -> Embeddings:
-        accept = "application/json"
-        content_type = "application/json"
-        embeddings = []
-        for text in input:
-            input_body = {"inputText": text}
-            body = json.dumps(input_body)
-            response = self._client.invoke_model(
-                body=body,
-                modelId=self._model_name,
-                accept=accept,
-                contentType=content_type,
-            )
-            embedding = json.load(response.get("body")).get("embedding")
-            embeddings.append(embedding)
+        if self._model_provider == "amazon":
+            embeddings = []
+            for text in input:
+                input_body = {
+                    "inputText": text
+                }
+                response = self.call_model(input_body)
+                embedding = json.load(response.get("body")).get("embedding")
+                embeddings.append(embedding)
+        elif self._model_provider == "cohere":
+            # See Amazon Bedrock User Guide > Cohere Embed models for more information
+            # https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-embed.html
+            assert len(input) <= 128, f"Input texts exceeds max size (Got: {len(input)}, Expected: <=128)"
+            assert all(len(text) <= 2048 for text in input), f"Input contains texts exceeding max length (2048)"
+            input_body = {
+                "texts": input,
+                "input_type": "search_document"
+            }
+            response = self.call_model(input_body)
+            embeddings = json.load(response.get("body")).get("embeddings")
+        else:
+            raise NotImplementedError(f"Model {self._model_name} is not supported!")
         return embeddings