[magics] Add --api-base and --api-key-name arguments (#1471)

* [magics] Enable passing API base url and key for proxy service access in magicsfor

* Update magics.py

* updated api-key-name

Author: srdas

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

@@ -10,6 +10,7 @@
 import click
 import litellm
 import traitlets
+from typing import Optional
 from IPython.core.magic import Magics, line_cell_magic, magics_class
 from IPython.display import HTML, JSON, Markdown, Math
 from jupyter_ai.model_providers.model_list import CHAT_MODELS
@@ -314,10 +315,29 @@ def run_ai_cell(self, args: CellArgs, prompt: str):
                 print(error_msg, file=sys.stderr)  # Log to stderr
                 return
         try:
+            # Prepare litellm completion arguments
+            completion_args = {
+                "model": model_id, 
+                "messages": messages, 
+                "stream": False
+            }
+
+            # Add api_base if provided
+            if args.api_base:
+                completion_args["api_base"] = args.api_base
+
+            # Add API key from .env if api_key_name is provided
+            if args.api_key_name:
+                # Retrieve the actual API key from the .env file
+                api_key_name_value = os.getenv(args.api_key_name)
+                if not api_key_name_value:
+                    error_msg = f"API key '{args.api_key_name}' not found in .env file."
+                    print(error_msg, file=sys.stderr)
+                    return
+                completion_args["api_key_name"] = api_key_name_value
+
             # Call litellm completion
-            response = litellm.completion(
-                model=model_id, messages=messages, stream=False
-            )
+            response = litellm.completion(**completion_args)
 
             # Extract output text from response
             output = response.choices[0].message.content

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

@@ -51,6 +51,9 @@ class CellArgs(BaseModel):
     region_name: Optional[str] = None
     request_schema: Optional[str] = None
     response_path: Optional[str] = None
+    # Parameters for custom API endpoints
+    api_base: Optional[str] = None
+    api_key_name: Optional[str] = None
 
 
 # Should match CellArgs
@@ -63,6 +66,9 @@ class FixArgs(BaseModel):
     region_name: Optional[str] = None
     request_schema: Optional[str] = None
     response_path: Optional[str] = None
+    # Parameters for custom API endpoints
+    api_base: Optional[str] = None
+    api_key_name: Optional[str] = None
 
 
 class HelpArgs(BaseModel):
@@ -133,6 +139,16 @@ def verify_json_value(ctx, param, value):
     default="markdown",
     help=FORMAT_HELP,
 )
+@click.option(
+    "--api-base",
+    required=False,
+    help="Base URL for the API endpoint.",
+)
+@click.option(
+    "--api-key-name",
+    required=False,
+    help="Name of the environment variable containing the API key.",
+)
 @click.option(
     REGION_NAME_SHORT_OPTION,
     REGION_NAME_LONG_OPTION,
@@ -168,6 +184,11 @@ def cell_magic_parser(context: click.Context, **kwargs):
     model IDs (with the provider ID explicitly prefixed, followed by a colon)
     are accepted.
 
+    Optional parameters:
+    --api-base: Base URL for the API endpoint
+    --api-key-name: Name of the environment variable containing the API key
+               (the actual key should be stored in the .env file)
+
     To view available language models, please run `%ai list`.
     """
     if not kwargs["model_id"] and context.default_map:
@@ -191,6 +212,16 @@ def line_magic_parser():
     default="markdown",
     help=FORMAT_HELP,
 )
+@click.option(
+    "--api-base",
+    required=False,
+    help="Base URL for the API endpoint.",
+)
+@click.option(
+    "--api-key-name",
+    required=False,
+    help="Name of the environment variable containing the API key.",
+)
 @click.option(
     REGION_NAME_SHORT_OPTION,
     REGION_NAME_LONG_OPTION,