Added Bedrock provider (jupyterlab#263)

* Added Bedrock provider * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Added Bedrock info to docs * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update docs/source/users/index.md Co-authored-by: Jason Weill <93281816+JasonWeill@users.noreply.github.com> --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Jason Weill <93281816+JasonWeill@users.noreply.github.com>
dbelgrod · Jul 13, 2023 · 11eff22 · 11eff22
1 parent d9adbdb
commit 11eff22
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 1 deletion.
diff --git a/docs/source/users/index.md b/docs/source/users/index.md
@@ -39,6 +39,7 @@ Jupyter AI supports the following model providers:
 |---------------------|----------------------|----------------------------|---------------------------------|
 | AI21                | `ai21`               | `AI21_API_KEY`             | `ai21`                          |
 | Anthropic           | `anthropic`          | `ANTHROPIC_API_KEY`        | `anthropic`                     |
+| Bedrock             | `amazon-bedrock`     | N/A                        | `boto3`                         |
 | Cohere              | `cohere`             | `COHERE_API_KEY`           | `cohere`                        |
 | Hugging Face Hub    | `huggingface_hub`    | `HUGGINGFACEHUB_API_TOKEN` | `huggingface_hub`, `ipywidgets`, `pillow` |
 | OpenAI              | `openai`             | `OPENAI_API_KEY`           | `openai`                        |
@@ -47,6 +48,12 @@ Jupyter AI supports the following model providers:
 
 The environment variable names shown above are also the names of the settings keys used when setting up the chat interface.
 
+To use the Bedrock models, you need access to the Bedrock service. For more information, see the
+[Amazon Bedrock Homepage](https://aws.amazon.com/bedrock/).
+
+To use Bedrock models, you will need to authenticate via
+[boto3](https://github.com/boto/boto3).
+
 You need the `pillow` Python package to use Hugging Face Hub's text-to-image models.
 
 You can find a list of Hugging Face's models at https://huggingface.co/models.

diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/__init__.py b/packages/jupyter-ai-magics/jupyter_ai_magics/__init__.py
@@ -14,6 +14,7 @@
     AI21Provider,
     AnthropicProvider,
     BaseProvider,
+    BedrockProvider,
     ChatOpenAINewProvider,
     ChatOpenAIProvider,
     CohereProvider,

diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py
@@ -12,6 +12,7 @@
 from langchain.llms import (
     AI21,
     Anthropic,
+    Bedrock,
     Cohere,
     HuggingFaceHub,
     OpenAI,
@@ -123,7 +124,8 @@ def __init__(self, *args, **kwargs):
             )
 
         model_kwargs = {}
-        model_kwargs[self.__class__.model_id_key] = kwargs["model_id"]
+        if self.__class__.model_id_key != "model_id":
+            model_kwargs[self.__class__.model_id_key] = kwargs["model_id"]
 
         super().__init__(*args, **kwargs, **model_kwargs)
 
@@ -430,3 +432,21 @@ def __init__(self, *args, **kwargs):
 
     async def _acall(self, *args, **kwargs) -> Coroutine[Any, Any, str]:
         return await self._call_in_executor(*args, **kwargs)
+
+
+class BedrockProvider(BaseProvider, Bedrock):
+    id = "bedrock"
+    name = "Amazon Bedrock"
+    models = [
+        "amazon.titan-tg1-large",
+        "anthropic.claude-v1",
+        "anthropic.claude-instant-v1",
+        "ai21.j2-jumbo-instruct",
+        "ai21.j2-grande-instruct",
+    ]
+    model_id_key = "model_id"
+    pypi_package_deps = ["boto3"]
+    auth_strategy = AwsAuthStrategy()
+
+    async def _acall(self, *args, **kwargs) -> Coroutine[Any, Any, str]:
+        return await self._call_in_executor(*args, **kwargs)
diff --git a/packages/jupyter-ai-magics/pyproject.toml b/packages/jupyter-ai-magics/pyproject.toml
@@ -62,6 +62,7 @@ openai = "jupyter_ai_magics:OpenAIProvider"
 openai-chat = "jupyter_ai_magics:ChatOpenAIProvider"
 openai-chat-new = "jupyter_ai_magics:ChatOpenAINewProvider"
 sagemaker-endpoint = "jupyter_ai_magics:SmEndpointProvider"
+amazon-bedrock = "jupyter_ai_magics:BedrockProvider"
 
 [project.entry-points."jupyter_ai.embeddings_model_providers"]
 cohere = "jupyter_ai_magics:CohereEmbeddingsProvider"