diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/aliases.py b/packages/jupyter-ai-magics/jupyter_ai_magics/aliases.py index f34826428..e34a31aa0 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/aliases.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/aliases.py @@ -6,4 +6,5 @@ "ernie-bot": "qianfan:ERNIE-Bot", "ernie-bot-4": "qianfan:ERNIE-Bot-4", "titan": "bedrock:amazon.titan-tg1-large", + "openrouter-claude": "openrouter:anthropic/claude-3.5-sonnet:beta", } diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/openrouter.py b/packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/openrouter.py new file mode 100644 index 000000000..81c2d7ab1 --- /dev/null +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/openrouter.py @@ -0,0 +1,60 @@ +from typing import Dict + +from jupyter_ai_magics import BaseProvider +from jupyter_ai_magics.providers import EnvAuthStrategy, TextField +from langchain_core.pydantic_v1 import root_validator +from langchain_core.utils import convert_to_secret_str, get_from_dict_or_env +from langchain_openai import ChatOpenAI + + +class ChatOpenRouter(ChatOpenAI): + @property + def lc_secrets(self) -> Dict[str, str]: + return {"openai_api_key": "OPENROUTER_API_KEY"} + + +class OpenRouterProvider(BaseProvider, ChatOpenRouter): + id = "openrouter" + name = "OpenRouter" + models = [ + "*" + ] # OpenRouter supports multiple models, so we use "*" to indicate it's a registry + model_id_key = "model_name" + pypi_package_deps = ["langchain_openai"] + auth_strategy = EnvAuthStrategy(name="OPENROUTER_API_KEY") + registry = True + + fields = [ + TextField( + key="openai_api_base", label="API Base URL (optional)", format="text" + ), + ] + + def __init__(self, **kwargs): + openrouter_api_key = kwargs.pop("openrouter_api_key", None) + openrouter_api_base = kwargs.pop( + "openai_api_base", "https://openrouter.ai/api/v1" + ) + + super().__init__( + openai_api_key=openrouter_api_key, + openai_api_base=openrouter_api_base, + **kwargs, + ) + + @root_validator(pre=False, skip_on_failure=True, allow_reuse=True) + def validate_environment(cls, values: Dict) -> Dict: + """Validate that api key and python package exists in environment.""" + values["openai_api_key"] = convert_to_secret_str( + get_from_dict_or_env(values, "openai_api_key", "OPENROUTER_API_KEY") + ) + return super().validate_environment(values) + + @classmethod + def is_api_key_exc(cls, e: Exception): + import openai + + if isinstance(e, openai.AuthenticationError): + error_details = e.json_body.get("error", {}) + return error_details.get("code") == "invalid_api_key" + return False diff --git a/packages/jupyter-ai-magics/pyproject.toml b/packages/jupyter-ai-magics/pyproject.toml index d42a9bc7b..9273c2d9a 100644 --- a/packages/jupyter-ai-magics/pyproject.toml +++ b/packages/jupyter-ai-magics/pyproject.toml @@ -74,6 +74,7 @@ nvidia-chat = "jupyter_ai_magics.partner_providers.nvidia:ChatNVIDIAProvider" together-ai = "jupyter_ai_magics:TogetherAIProvider" gemini = "jupyter_ai_magics.partner_providers.gemini:GeminiProvider" mistralai = "jupyter_ai_magics.partner_providers.mistralai:MistralAIProvider" +openrouter = "jupyter_ai_magics.partner_providers.openrouter:OpenRouterProvider" [project.entry-points."jupyter_ai.embeddings_model_providers"] azure = "jupyter_ai_magics.partner_providers.openai:AzureOpenAIEmbeddingsProvider"