diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/gemini.py b/packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/gemini.py new file mode 100644 index 000000000..d56d30e91 --- /dev/null +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/gemini.py @@ -0,0 +1,13 @@ +from jupyter_ai_magics.providers import BaseProvider, EnvAuthStrategy +from langchain_google_genai import GoogleGenerativeAI + + +class GeminiProvider(BaseProvider, GoogleGenerativeAI): + id = "gemini" + name = "Gemini" + models = [ + "gemini-pro", + ] + model_id_key = "model" + auth_strategy = EnvAuthStrategy(name="GOOGLE_API_KEY") + pypi_package_deps = ["langchain-google-genai"] diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py index f787067fd..dfaaf8356 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py @@ -39,7 +39,7 @@ SagemakerEndpoint, Together, ) -from langchain_google_genai import GoogleGenerativeAI + # this is necessary because `langchain.pydantic_v1.main` does not include # `ModelMetaclass`, as it is not listed in `__all__` by the `pydantic.main` @@ -901,18 +901,3 @@ class QianfanProvider(BaseProvider, QianfanChatEndpoint): model_id_key = "model_name" pypi_package_deps = ["qianfan"] auth_strategy = MultiEnvAuthStrategy(names=["QIANFAN_AK", "QIANFAN_SK"]) - - -class GeminiProvider(BaseProvider, GoogleGenerativeAI): - id = "gemini" - name = "Gemini" - models = [ - "gemini-pro", - ] - model_id_key = "model_name" - pypi_package_deps = ["langchain-google-genai"] - auth_strategy = EnvAuthStrategy(name="GOOGLE_API_KEY") - - @classmethod - def is_api_key_exc(cls, e: Exception): - return False diff --git a/packages/jupyter-ai-magics/pyproject.toml b/packages/jupyter-ai-magics/pyproject.toml index 11d14c8b1..64963657e 100644 --- a/packages/jupyter-ai-magics/pyproject.toml +++ b/packages/jupyter-ai-magics/pyproject.toml @@ -65,6 +65,7 @@ amazon-bedrock-chat = "jupyter_ai_magics:BedrockChatProvider" qianfan = "jupyter_ai_magics:QianfanProvider" nvidia-chat = "jupyter_ai_magics.partner_providers.nvidia:ChatNVIDIAProvider" together-ai = "jupyter_ai_magics:TogetherAIProvider" +gemini = "jupyter_ai_magics.partner_providers.gemini:GeminiProvider" [project.entry-points."jupyter_ai.embeddings_model_providers"] bedrock = "jupyter_ai_magics:BedrockEmbeddingsProvider"