Added support for Amazon Bedrock embeddings. (#41)

* added support for Bedrock embedding models

* added support for Bedrock embedding models

---------

Co-authored-by: ivantang <zilliz@ivantang-work.lan>

Author: hiimivantang

deepsearcher/embedding/__init__.py

@@ -1,9 +1,11 @@
 from .milvus_embedding import MilvusEmbedding
 from .openai_embedding import OpenAIEmbedding
 from .voyage_embedding import VoyageEmbedding
+from .bedrock_embedding import BedrockEmbedding
 
 __all__ = [
     "MilvusEmbedding",
     "OpenAIEmbedding",
     "VoyageEmbedding",
-]
+    "BedrockEmbedding"
+]

deepsearcher/embedding/bedrock_embedding.py

@@ -0,0 +1,57 @@
+import os
+from typing import List
+import boto3
+import json
+from deepsearcher.embedding.base import BaseEmbedding
+
+MODEL_ID_TITAN_TEXT_G1 = "amazon.titan-embed-text-v1"
+MODEL_ID_TITAN_TEXT_V2 = "amazon.titan-embed-text-v2:0"
+MODEL_ID_TITAN_MULTIMODAL_G1 = "amazon.titan-embed-image-v1"
+MODEL_ID_COHERE_ENGLISH_V3 = "cohere.embed-english-v3"
+MODEL_ID_COHERE_MULTILINGUAL_V3 = "cohere.embed-multilingual-v3"
+
+BEDROCK_MODEL_DIM_MAP = {
+    MODEL_ID_TITAN_TEXT_G1: 1536,
+    MODEL_ID_TITAN_TEXT_V2: 1024,
+    MODEL_ID_TITAN_MULTIMODAL_G1: 1024,
+    MODEL_ID_COHERE_ENGLISH_V3: 1024,
+    MODEL_ID_COHERE_MULTILINGUAL_V3: 1024
+}
+
+DEFAULT_MODEL_ID = MODEL_ID_TITAN_TEXT_V2
+
+class BedrockEmbedding(BaseEmbedding):
+    def __init__(self, model: str = DEFAULT_MODEL_ID, **kwargs):
+        """
+        Args:
+            model_name (`str`):
+                Can be one of the following:
+                'amazon.titan-embed-text-v2:0': dimensions include 256, 512, 1024, default is 1024,
+        """
+
+        aws_access_key_id = kwargs.pop("aws_access_key_id", os.getenv("AWS_ACCESS_KEY_ID"))
+        aws_secret_access_key = kwargs.pop("aws_secret_access_key", os.getenv("AWS_SECRET_ACCESS_KEY"))
+
+        if model in {None, DEFAULT_MODEL_ID} and "model_name" in kwargs:
+            model = kwargs.pop("model_name") #overwrites `model` with `model_name`
+        
+        self.model = model
+
+        #TODO: initiate boto3 client
+        self.client = boto3.client("bedrock-runtime", 
+                                    region_name="us-east-1", #FIXME: allow users to specify
+                                    aws_access_key_id=aws_access_key_id,
+                                    aws_secret_access_key=aws_secret_access_key)
+        
+    def embed_query(self, text: str) -> List[float]:
+        response = self.client.invoke_model(modelId=self.model, body=json.dumps({"inputText": text}))
+        model_response = json.loads(response["body"].read())
+        embedding = model_response["embedding"]
+        return embedding
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        return [self.embed_query(text) for text in texts]
+
+    @property
+    def dimension(self) -> int:
+        return BEDROCK_MODEL_DIM_MAP[self.model]

deepsearcher/online_query.py

@@ -20,7 +20,6 @@ def query(
 ) -> Tuple[str, List[RetrievalResult], int]:
     return asyncio.run(async_query(original_query, max_iter))
 
-
 async def async_query(
     original_query: str, max_iter: int = 3
 ) -> Tuple[str, List[RetrievalResult], int]:

requirements.txt

@@ -10,3 +10,4 @@ termcolor
 fastapi
 uvicorn
 pydantic-settings
+boto3