Update Anthropic providers to use langchain_anthropic partner packa…

…ge (jupyterlab#700) * update anthropic providers to use langchain_anthropic partner package * pre-commit
Marchlak · Oct 28, 2024 · 5c0388b · 5c0388b
1 parent 3dbf3c4
commit 5c0388b
Show file tree

Hide file tree

Showing 5 changed files with 63 additions and 62 deletions.
diff --git a/docs/source/users/index.md b/docs/source/users/index.md
@@ -134,8 +134,8 @@ Jupyter AI supports the following model providers:
 | Provider            | Provider ID          | Environment variable(s)    | Python package(s)               |
 |---------------------|----------------------|----------------------------|---------------------------------|
 | AI21                | `ai21`               | `AI21_API_KEY`             | `ai21`                          |
-| Anthropic           | `anthropic`          | `ANTHROPIC_API_KEY`        | `anthropic`                     |
-| Anthropic (chat)    | `anthropic-chat`     | `ANTHROPIC_API_KEY`        | `anthropic`                     |
+| Anthropic           | `anthropic`          | `ANTHROPIC_API_KEY`        | `langchain-anthropic`           |
+| Anthropic (chat)    | `anthropic-chat`     | `ANTHROPIC_API_KEY`        | `langchain-anthropic`           |
 | Bedrock             | `bedrock`            | N/A                        | `boto3`                         |
 | Bedrock (chat)      | `bedrock-chat`       | N/A                        | `boto3`                         |
 | Cohere              | `cohere`             | `COHERE_API_KEY`           | `cohere`                        |

diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/__init__.py b/packages/jupyter-ai-magics/jupyter_ai_magics/__init__.py
@@ -15,11 +15,9 @@
 # expose model providers on the package root
 from .providers import (
     AI21Provider,
-    AnthropicProvider,
     BaseProvider,
     BedrockChatProvider,
     BedrockProvider,
-    ChatAnthropicProvider,
     CohereProvider,
     GPT4AllProvider,
     HfHubProvider,

diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/anthropic.py b/packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/anthropic.py
@@ -0,0 +1,58 @@
+from langchain_anthropic import AnthropicLLM, ChatAnthropic
+
+from ..providers import BaseProvider, EnvAuthStrategy
+
+
+class AnthropicProvider(BaseProvider, AnthropicLLM):
+    id = "anthropic"
+    name = "Anthropic"
+    models = [
+        "claude-v1",
+        "claude-v1.0",
+        "claude-v1.2",
+        "claude-2",
+        "claude-2.0",
+        "claude-instant-v1",
+        "claude-instant-v1.0",
+        "claude-instant-v1.2",
+    ]
+    model_id_key = "model"
+    pypi_package_deps = ["anthropic"]
+    auth_strategy = EnvAuthStrategy(name="ANTHROPIC_API_KEY")
+
+    @property
+    def allows_concurrency(self):
+        return False
+
+    @classmethod
+    def is_api_key_exc(cls, e: Exception):
+        """
+        Determine if the exception is an Anthropic API key error.
+        """
+        import anthropic
+
+        if isinstance(e, anthropic.AuthenticationError):
+            return e.status_code == 401 and "Invalid API Key" in str(e)
+        return False
+
+
+class ChatAnthropicProvider(
+    BaseProvider, ChatAnthropic
+):  # https://docs.anthropic.com/claude/docs/models-overview
+    id = "anthropic-chat"
+    name = "ChatAnthropic"
+    models = [
+        "claude-2.0",
+        "claude-2.1",
+        "claude-instant-1.2",
+        "claude-3-opus-20240229",
+        "claude-3-sonnet-20240229",
+        "claude-3-haiku-20240307",
+    ]
+    model_id_key = "model"
+    pypi_package_deps = ["anthropic"]
+    auth_strategy = EnvAuthStrategy(name="ANTHROPIC_API_KEY")
+
+    @property
+    def allows_concurrency(self):
+        return False
diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py
@@ -402,61 +402,6 @@ def is_api_key_exc(cls, e: Exception):
         return False
 
 
-class AnthropicProvider(BaseProvider, Anthropic):
-    id = "anthropic"
-    name = "Anthropic"
-    models = [
-        "claude-v1",
-        "claude-v1.0",
-        "claude-v1.2",
-        "claude-2",
-        "claude-2.0",
-        "claude-instant-v1",
-        "claude-instant-v1.0",
-        "claude-instant-v1.2",
-    ]
-    model_id_key = "model"
-    pypi_package_deps = ["anthropic"]
-    auth_strategy = EnvAuthStrategy(name="ANTHROPIC_API_KEY")
-
-    @property
-    def allows_concurrency(self):
-        return False
-
-    @classmethod
-    def is_api_key_exc(cls, e: Exception):
-        """
-        Determine if the exception is an Anthropic API key error.
-        """
-        import anthropic
-
-        if isinstance(e, anthropic.AuthenticationError):
-            return e.status_code == 401 and "Invalid API Key" in str(e)
-        return False
-
-
-class ChatAnthropicProvider(
-    BaseProvider, ChatAnthropic
-):  # https://docs.anthropic.com/claude/docs/models-overview
-    id = "anthropic-chat"
-    name = "ChatAnthropic"
-    models = [
-        "claude-2.0",
-        "claude-2.1",
-        "claude-instant-1.2",
-        "claude-3-opus-20240229",
-        "claude-3-sonnet-20240229",
-        "claude-3-haiku-20240307",
-    ]
-    model_id_key = "model"
-    pypi_package_deps = ["anthropic"]
-    auth_strategy = EnvAuthStrategy(name="ANTHROPIC_API_KEY")
-
-    @property
-    def allows_concurrency(self):
-        return False
-
-
 class CohereProvider(BaseProvider, Cohere):
     id = "cohere"
     name = "Cohere"

diff --git a/packages/jupyter-ai-magics/pyproject.toml b/packages/jupyter-ai-magics/pyproject.toml
@@ -36,11 +36,11 @@ test = ["coverage", "pytest", "pytest-asyncio", "pytest-cov"]
 
 all = [
     "ai21",
-    "anthropic~=0.3.0",
     "cohere>4.40,<5",
     "gpt4all",
     "huggingface_hub",
     "ipywidgets",
+    "langchain_anthropic",
     "langchain_nvidia_ai_endpoints",
     "langchain-openai",
     "pillow",
@@ -52,7 +52,8 @@ all = [
 
 [project.entry-points."jupyter_ai.model_providers"]
 ai21 = "jupyter_ai_magics:AI21Provider"
-anthropic = "jupyter_ai_magics:AnthropicProvider"
+anthropic = "jupyter_ai_magics.partner_providers.anthropic:AnthropicProvider"
+anthropic-chat = "jupyter_ai_magics.partner_providers.anthropic:ChatAnthropicProvider"
 cohere = "jupyter_ai_magics:CohereProvider"
 gpt4all = "jupyter_ai_magics:GPT4AllProvider"
 huggingface_hub = "jupyter_ai_magics:HfHubProvider"
@@ -61,7 +62,6 @@ openai-chat = "jupyter_ai_magics.partner_providers.openai:ChatOpenAIProvider"
 azure-chat-openai = "jupyter_ai_magics.partner_providers.openai:AzureChatOpenAIProvider"
 sagemaker-endpoint = "jupyter_ai_magics:SmEndpointProvider"
 amazon-bedrock = "jupyter_ai_magics:BedrockProvider"
-anthropic-chat = "jupyter_ai_magics:ChatAnthropicProvider"
 amazon-bedrock-chat = "jupyter_ai_magics:BedrockChatProvider"
 qianfan = "jupyter_ai_magics:QianfanProvider"
 nvidia-chat = "jupyter_ai_magics.partner_providers.nvidia:ChatNVIDIAProvider"