add baidu's ERNIE-Bot and ERNIE-Bot-4 support

.gitignore

@@ -132,3 +132,5 @@ dev.sh
 .jupyter_ystore.db
 
 .yarn
+
+.conda/

examples/commands.ipynb

@@ -1117,7 +1117,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.4"
+   "version": "3.11.5"
   }
  },
  "nbformat": 4,

examples/errors.ipynb

@@ -259,7 +259,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.8"
+   "version": "3.11.5"
   }
  },
  "nbformat": 4,

examples/images.ipynb

@@ -136,7 +136,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.16"
+   "version": "3.11.5"
   }
  },
  "nbformat": 4,

packages/jupyter-ai-magics/jupyter_ai_magics/__init__.py

@@ -7,6 +7,7 @@
     GPT4AllEmbeddingsProvider,
     HfHubEmbeddingsProvider,
     OpenAIEmbeddingsProvider,
+    QianfanEmbeddingsEndpointProvider,
 )
 from .exception import store_exception
 from .magics import AiMagics
@@ -26,6 +27,7 @@
     GPT4AllProvider,
     HfHubProvider,
     OpenAIProvider,
+    QianfanProvider,
     SmEndpointProvider,
 )
 

packages/jupyter-ai-magics/jupyter_ai_magics/aliases.py

@@ -3,4 +3,7 @@
     "gpt3": "openai:text-davinci-003",
     "chatgpt": "openai-chat:gpt-3.5-turbo",
     "gpt4": "openai-chat:gpt-4",
+    "ernie-bot": "qianfan:ERNIE-Bot",
+    "ernie-bot-4": "qianfan:ERNIE-Bot-4",
+    "titan": "bedrock:amazon.titan-tg1-large",
 }

packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py

@@ -6,13 +6,15 @@
     AwsAuthStrategy,
     EnvAuthStrategy,
     Field,
+    MultiEnvAuthStrategy,
 )
 from langchain.embeddings import (
     BedrockEmbeddings,
     CohereEmbeddings,
     GPT4AllEmbeddings,
     HuggingFaceHubEmbeddings,
     OpenAIEmbeddings,
+    QianfanEmbeddingsEndpoint,
 )
 from langchain.pydantic_v1 import BaseModel, Extra
 
@@ -127,3 +129,14 @@ def __init__(self, **kwargs):
     models = ["all-MiniLM-L6-v2-f16"]
     model_id_key = "model_id"
     pypi_package_deps = ["gpt4all"]
+
+
+class QianfanEmbeddingsEndpointProvider(
+    BaseEmbeddingsProvider, QianfanEmbeddingsEndpoint
+):
+    id = "qianfan"
+    name = "ERNIE-Bot"
+    models = ["ERNIE-Bot", "ERNIE-Bot-4"]
+    model_id_key = "model"
+    pypi_package_deps = ["qianfan"]
+    auth_strategy = MultiEnvAuthStrategy(names=["QIANFAN_AK", "QIANFAN_SK"])

packages/jupyter-ai-magics/jupyter_ai_magics/magics.py

@@ -11,6 +11,7 @@
 from IPython import get_ipython
 from IPython.core.magic import Magics, line_cell_magic, magics_class
 from IPython.display import HTML, JSON, Markdown, Math
+from jupyter_ai_magics.aliases import MODEL_ID_ALIASES
 from jupyter_ai_magics.utils import decompose_model_id, get_lm_providers
 from langchain.chains import LLMChain
 from langchain.schema import HumanMessage
@@ -28,14 +29,6 @@
 )
 from .providers import BaseProvider
 
-MODEL_ID_ALIASES = {
-    "gpt2": "huggingface_hub:gpt2",
-    "gpt3": "openai:text-davinci-003",
-    "chatgpt": "openai-chat:gpt-3.5-turbo",
-    "gpt4": "openai-chat:gpt-4",
-    "titan": "bedrock:amazon.titan-tg1-large",
-}
-
 
 class TextOrMarkdown:
     def __init__(self, text, markdown):

packages/jupyter-ai-magics/jupyter_ai_magics/providers.py

@@ -23,6 +23,7 @@
     BedrockChat,
     ChatAnthropic,
     ChatOpenAI,
+    QianfanChatEndpoint,
 )
 from langchain.chat_models.base import BaseChatModel
 from langchain.llms import (
@@ -34,6 +35,7 @@
     HuggingFaceHub,
     OpenAI,
     OpenAIChat,
+    QianfanLLMEndpoint,
     SagemakerEndpoint,
 )
 from langchain.llms.sagemaker_endpoint import LLMContentHandler
@@ -703,3 +705,13 @@ async def _agenerate(self, *args, **kwargs) -> Coroutine[Any, Any, LLMResult]:
     @property
     def allows_concurrency(self):
         return not "anthropic" in self.model_id
+
+
+# Baidu QianfanChat provider. temporarily living as a separate class until
+class QianfanProvider(BaseProvider, QianfanChatEndpoint):
+    id = "qianfan"
+    name = "ERNIE-Bot"
+    models = ["ERNIE-Bot", "ERNIE-Bot-4"]
+    model_id_key = "model_name"
+    pypi_package_deps = ["qianfan"]
+    auth_strategy = MultiEnvAuthStrategy(names=["QIANFAN_AK", "QIANFAN_SK"])

packages/jupyter-ai-magics/pyproject.toml

@@ -50,7 +50,8 @@ all = [
     "ipywidgets",
     "pillow",
     "openai",
-    "boto3"
+    "boto3",
+    "qianfan"
 ]
 
 [project.entry-points."jupyter_ai.model_providers"]
@@ -67,13 +68,15 @@ 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"
 
 [project.entry-points."jupyter_ai.embeddings_model_providers"]
 bedrock = "jupyter_ai_magics:BedrockEmbeddingsProvider"
 cohere = "jupyter_ai_magics:CohereEmbeddingsProvider"
 gpt4all = "jupyter_ai_magics:GPT4AllEmbeddingsProvider"
 huggingface_hub = "jupyter_ai_magics:HfHubEmbeddingsProvider"
 openai = "jupyter_ai_magics:OpenAIEmbeddingsProvider"
+qianfan = "jupyter_ai_magics:QianfanEmbeddingsEndpointProvider"
 
 [tool.hatch.version]
 source = "nodejs"