feat: support bedrock embedding model provider (#511)

Close #379 --------- Co-authored-by: Mini256 <[email protected]>
pingcap · Dec 31, 2024 · 85872a9 · 85872a9
1 parent 11852d1
commit 85872a9
Show file tree

Hide file tree

Showing 8 changed files with 69 additions and 0 deletions.
diff --git a/backend/app/rag/chat_config.py b/backend/app/rag/chat_config.py
@@ -17,6 +17,7 @@
 from llama_index.embeddings.openai import OpenAIEmbedding
 from llama_index.embeddings.jinaai import JinaEmbedding
 from llama_index.embeddings.cohere import CohereEmbedding
+from llama_index.embeddings.bedrock import BedrockEmbedding
 from llama_index.embeddings.ollama import OllamaEmbedding
 from llama_index.postprocessor.jinaai_rerank import JinaRerank
 from llama_index.postprocessor.cohere_rerank import CohereRerank
@@ -331,6 +332,14 @@ def get_embed_model(
                 model_name=model,
                 cohere_api_key=credentials,
             )
+        case EmbeddingProvider.BEDROCK:
+            return BedrockEmbedding(
+                model_name=model,
+                aws_access_key_id=credentials["aws_access_key_id"],
+                aws_secret_access_key=credentials["aws_secret_access_key"],
+                region_name=credentials["aws_region_name"],
+                **config,
+            )
         case EmbeddingProvider.OLLAMA:
             return OllamaEmbedding(
                 model_name=model,

diff --git a/backend/app/rag/embed_model_option.py b/backend/app/rag/embed_model_option.py
@@ -57,6 +57,22 @@ class EmbeddingModelOption(BaseModel):
         credentials_type="str",
         default_credentials="*****",
     ),
+    EmbeddingModelOption(
+        provider=EmbeddingProvider.BEDROCK,
+        provider_display_name="Bedrock",
+        provider_description="Amazon Bedrock is a fully managed foundation models service.",
+        provider_url="https://docs.aws.amazon.com/bedrock/",
+        default_embedding_model="amazon.titan-embed-text-v2:0",
+        embedding_model_description="",
+        credentials_display_name="AWS Bedrock Credentials JSON",
+        credentials_description="The JSON Object of AWS Credentials, refer to https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-files.html#cli-configure-files-global",
+        credentials_type="dict",
+        default_credentials={
+            "aws_access_key_id": "****",
+            "aws_secret_access_key": "****",
+            "aws_region_name": "us-west-2",
+        },
+    ),
     EmbeddingModelOption(
         provider=EmbeddingProvider.OLLAMA,
         provider_display_name="Ollama",

diff --git a/backend/app/types.py b/backend/app/types.py
@@ -15,6 +15,7 @@ class EmbeddingProvider(str, enum.Enum):
     OPENAI = "openai"
     JINA = "jina"
     COHERE = "cohere"
+    BEDROCK = "bedrock"
     OLLAMA = "ollama"
     GITEEAI = "giteeai"
     LOCAL = "local"

diff --git a/backend/pyproject.toml b/backend/pyproject.toml
@@ -58,6 +58,7 @@ dependencies = [
     "retry>=0.9.2",
     "langchain-openai>=0.2.9",
     "ragas>=0.2.6",
+    "llama-index-embeddings-bedrock<=0.3.0",
 ]
 readme = "README.md"
 requires-python = ">= 3.8"

diff --git a/backend/requirements-dev.lock b/backend/requirements-dev.lock
@@ -54,6 +54,7 @@ billiard==4.2.0
     # via celery
 boto3==1.34.156
     # via cohere
+    # via llama-index-embeddings-bedrock
     # via llama-index-llms-bedrock
 botocore==1.34.156
     # via boto3
@@ -319,6 +320,7 @@ llama-index-core==0.11.10
     # via llama-index
     # via llama-index-agent-openai
     # via llama-index-cli
+    # via llama-index-embeddings-bedrock
     # via llama-index-embeddings-cohere
     # via llama-index-embeddings-jinaai
     # via llama-index-embeddings-ollama
@@ -338,6 +340,7 @@ llama-index-core==0.11.10
     # via llama-index-readers-file
     # via llama-index-readers-llama-parse
     # via llama-parse
+llama-index-embeddings-bedrock==0.3.0
 llama-index-embeddings-cohere==0.2.0
 llama-index-embeddings-jinaai==0.3.0
 llama-index-embeddings-ollama==0.3.0

diff --git a/backend/requirements.lock b/backend/requirements.lock
@@ -54,6 +54,7 @@ billiard==4.2.0
     # via celery
 boto3==1.34.156
     # via cohere
+    # via llama-index-embeddings-bedrock
     # via llama-index-llms-bedrock
 botocore==1.34.156
     # via boto3
@@ -319,6 +320,7 @@ llama-index-core==0.11.10
     # via llama-index
     # via llama-index-agent-openai
     # via llama-index-cli
+    # via llama-index-embeddings-bedrock
     # via llama-index-embeddings-cohere
     # via llama-index-embeddings-jinaai
     # via llama-index-embeddings-ollama
@@ -338,6 +340,7 @@ llama-index-core==0.11.10
     # via llama-index-readers-file
     # via llama-index-readers-llama-parse
     # via llama-parse
+llama-index-embeddings-bedrock==0.3.0
 llama-index-embeddings-cohere==0.2.0
 llama-index-embeddings-jinaai==0.3.0
 llama-index-embeddings-ollama==0.3.0

diff --git a/frontend/app/src/pages/docs/embedding-model.mdx b/frontend/app/src/pages/docs/embedding-model.mdx
@@ -109,6 +109,34 @@ Cohere provides industry-leading large language models (LLMs) and RAG capabiliti
 
 For more information, see the [Cohere Embed documentation](https://docs.cohere.com/docs/cohere-embed).
 
+### Amazon Bedrock
+
+Amazon Bedrock is a fully managed foundation models service that provides a range of large language models and embedding models.
+
+**Featured Models**:
+
+| Embedding Model                 | Vector Dimensions | Max Tokens |
+| ------------------------------- | ----------------- | ---------- |
+| `amazon.titan-embed-text-v2:0`  | 1024              | 8192       |
+| `amazon.titan-embed-text-v1`    | 1536              | 8192       |
+| `amazon.titan-embed-g1-text-02` | 1536              | 8192       |
+| `cohere.embed-english-v3`       | 1024              | 512        |
+| `cohere.embed-multilingual-v3`  | 1024              | 512        |
+
+To check all embbeding models supported by Bedrock, go to [Bedrock console](https://console.aws.amazon.com/bedrock).
+
+To use Amazon Bedrock, you'll need to provide a JSON Object of your AWS Credentials, as described in the [AWS CLI config global settings](https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-files.html#cli-configure-files-global):
+
+```json
+{
+    "aws_access_key_id": "****",
+    "aws_secret_access_key": "****",
+    "aws_region_name": "us-west-2"
+}
+```
+
+For more information, see the [Amazon Bedrock documentation](https://docs.aws.amazon.com/bedrock/).
+
 ### Ollama
 
 Ollama is a lightweight framework for building and running large language models and embedding models locally.

diff --git a/frontend/app/src/pages/docs/llm.mdx b/frontend/app/src/pages/docs/llm.mdx
@@ -30,6 +30,14 @@ Currently Autoflow supports the following LLM providers:
 - [Google Gemini](https://gemini.google.com/)
 - [Anthropic Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-claude)
 - [Amazon Bedrock](https://aws.amazon.com/bedrock/)
+    - To use Amazon Bedrock, you'll need to provide a JSON Object of your AWS Credentials, as described in the [AWS CLI config global settings](https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-files.html#cli-configure-files-global):
+    ```json
+    {
+        "aws_access_key_id": "****",
+        "aws_secret_access_key": "****",
+        "aws_region_name": "us-west-2"
+    }
+    ```
 - [Gitee AI](https://ai.gitee.com/serverless-api)
 - And all OpenAI-Like models:
     - [OpenRouter](https://openrouter.ai/)