diff --git a/docs/source/users/index.md b/docs/source/users/index.md index 74227af10..e24c21f64 100644 --- 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 index c3dce919a..c66891c0f 100644 --- 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 index a70180339..57a38e8d5 100644 --- 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 index 814b04d6e..bb5e5a6ca 100644 --- 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"