diff --git a/docs/source/users/index.md b/docs/source/users/index.md index cc518756c..d6912cee7 100644 --- a/docs/source/users/index.md +++ b/docs/source/users/index.md @@ -279,7 +279,7 @@ To clear the local vector database, you can run `/learn -d` and Jupyter AI will To clear the chat panel, use the `/clear` command. This does not reset the AI model; the model may still remember previous messages that you sent it, and it may use them to inform its responses. -## The `%%ai` magic command +## The `%ai` and `%%ai` magic commands Jupyter AI can also be used in notebooks via Jupyter AI magics. This section provides guidance on how to use Jupyter AI magics effectively. The examples in @@ -320,8 +320,11 @@ model you want to use and specify natural language prompts. ### Choosing a provider and model -The `%%ai` magic command enables you to specify a model provider and model with the -syntax `:`. The natural language prompt starts on the second line of the cell. +The `%%ai` cell magic allows you to invoke a language model of your choice with +a given prompt. The model is identified with a **global model ID**, which is a string with the +syntax `:`, where `` is the ID of the +provider and `` is the ID of the model scoped to that provider. +The prompt begins on the second line of the cell. For example, to send a text prompt to the provider `anthropic` and the model ID `claude-v1.2`, enter the following code into a cell and run it: @@ -331,8 +334,7 @@ For example, to send a text prompt to the provider `anthropic` and the model ID Write a poem about C++. ``` -We support the following providers, and all model IDs for each of these -providers, as defined in [`langchain.llms`](https://langchain.readthedocs.io/en/latest/reference/modules/llms.html#module-langchain.llms): +We currently support the following language model providers: - `ai21` - `anthropic` @@ -342,14 +344,29 @@ providers, as defined in [`langchain.llms`](https://langchain.readthedocs.io/en/ - `openai-chat` - `sagemaker-endpoint` -You can find a list of supported providers and models by running `%ai list`. Some providers -define a list of supported models. If a provider does not define a list of supported models, -consult the vendor's documentation. The [Hugging Face web site](https://huggingface.co/) -includes a list of models, for example. +:::{warning} +As of v0.8.0, only the `%%ai` *cell* magic may be used to invoke a language +model, while the `%ai` *line* magic is reserved for invoking subcommands. +::: + +### Listing available models + +Jupyter AI also includes multiple subcommands, which may be invoked via the +`%ai` *line* magic. Jupyter AI uses subcommands to provide additional utilities +in notebooks while keeping the same concise syntax for invoking a language model. + +The `%ai list` subcommand prints a list of available providers and models. Some +providers explicitly define a list of supported models in their API. However, +other providers, like HuggingFace Hub, lack a well-defined list of available +models. In such cases, it's best to consult the provider's upstream +documentation. The [HuggingFace website](https://huggingface.co/) includes a +list of models, for example. + +Optionally, you can specify a provider ID as a positional argument to `%ai list` +to get all models provided by one provider. For example, `%ai list openai` will +display only models provided by the `openai` provider. -Optionally, you can pass a provider ID as a parameter to `%ai list` to get all -models provided by one provider. For example, `%ai list openai` will display only models -provided by the `openai` provider. +### Abbreviated syntax If your model ID is associated with only one provider, you can omit the `provider-id` and the colon from the first line. For example, because `ai21` is the only provider of the diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py b/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py index ccc082c3b..01e47196e 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py @@ -5,12 +5,14 @@ import warnings from typing import Optional +import click from IPython import get_ipython from IPython.core.magic import Magics, magics_class, line_cell_magic -from IPython.core.magic_arguments import magic_arguments, argument, parse_argstring from IPython.display import HTML, JSON, Markdown, Math from jupyter_ai_magics.utils import decompose_model_id, load_providers + from .providers import BaseProvider +from .parsers import cell_magic_parser, line_magic_parser, CellArgs, HelpArgs, ListArgs MODEL_ID_ALIASES = { @@ -96,6 +98,9 @@ def __missing__(self, key): class EnvironmentError(BaseException): pass +class CellMagicError(BaseException): + pass + @magics_class class AiMagics(Magics): def __init__(self, shell): @@ -109,22 +114,6 @@ def __init__(self, shell): self.providers = load_providers() - def _ai_help_command_markdown(self): - table = ("| Command | Description |\n" - "| ------- | ----------- |\n") - - for command in AI_COMMANDS: - table += "| `" + command + "` | " + AI_COMMANDS[command] + "|\n"; - - return table - - def _ai_help_command_text(self): - output = "" - - for command in AI_COMMANDS: - output += command + " - " + AI_COMMANDS[command] + "\n"; - - return output def _ai_bulleted_list_models_for_provider(self, provider_id, Provider): output = "" @@ -223,30 +212,6 @@ def _ai_list_command_text(self, single_provider=None): return output - # Run an AI command using the arguments provided as a space-delimited value - def _ai_command(self, command, args_string): - args = args_string.split() # Split by whitespace - - # When we can use Python 3.10+, replace this with a 'match' command - if (command == 'help'): - return TextOrMarkdown(self._ai_help_command_text(), self._ai_help_command_markdown()) - elif (command == 'list'): - # Optional parameter: model provider ID - provider_id = None - if (len(args) >= 1): - provider_id = args[0] - - return TextOrMarkdown( - self._ai_list_command_text(provider_id), - self._ai_list_command_markdown(provider_id) - ) - else: - # This should be unreachable, since unhandled commands are treated like model names - return TextOrMarkdown( - f"No handler for command {command}\n", - f"No handler for command `{command}`" - ) - def _append_exchange_openai(self, prompt: str, output: str): """Appends a conversational exchange between user and an OpenAI Chat model to a transcript that will be included in future exchanges.""" @@ -269,34 +234,48 @@ def _get_provider(self, provider_id: Optional[str]) -> BaseProvider: return self.providers[provider_id] - @magic_arguments() - @argument('model_id', - help="""Model to run, specified as a model ID that may be - optionally prefixed with the ID of the model provider, delimited - by a colon.""") - @argument('-f', '--format', - choices=["code", "html", "image", "json", "markdown", "math", "md", "text"], - nargs="?", - default="markdown", - help="""IPython display to use when rendering output. [default="markdown"]""") - @argument('-r', '--reset', - action="store_true", - help="""Clears the conversation transcript used when interacting - with an OpenAI chat model provider. Does nothing with other - providers.""") - @argument('prompt', - nargs='*', - help="""Prompt for code generation. When used as a line magic, it - runs to the end of the line. In cell mode, the entire cell is - considered the code generation prompt.""") + def handle_help(self, _: HelpArgs): + with click.Context(cell_magic_parser, info_name="%%ai") as ctx: + click.echo(cell_magic_parser.get_help(ctx)) + click.echo('-' * 78) + with click.Context(line_magic_parser, info_name="%ai") as ctx: + click.echo(line_magic_parser.get_help(ctx)) + + def handle_list(self, args: ListArgs): + return TextOrMarkdown( + self._ai_list_command_text(args.provider_id), + self._ai_list_command_markdown(args.provider_id) + ) + @line_cell_magic def ai(self, line, cell=None): - # parse arguments - args = parse_argstring(self.ai, line) - if cell is None: - prompt = ' '.join(args.prompt) + raw_args = line.split(' ') + if cell: + args = cell_magic_parser(raw_args, prog_name="%%ai", standalone_mode=False) else: - prompt = cell + args = line_magic_parser(raw_args, prog_name="%ai", standalone_mode=False) + + if args == 0: + # this happens when `--help` is called on the root command, in which + # case we want to exit early. + return + + if args.type == "help": + return self.handle_help(args) + if args.type == "list": + return self.handle_list(args) + + # hint to the IDE that this object must be of type `RootArgs` + args: CellArgs = args + + if not cell: + raise CellMagicError( + """[0.8+]: To invoke a language model, you must use the `%%ai` + cell magic. The `%ai` line magic is only for use with + subcommands.""" + ) + + prompt = cell.strip() # If the user is attempting to run a command, run the command separately. if (args.model_id in AI_COMMANDS): diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py new file mode 100644 index 000000000..5847cb5a3 --- /dev/null +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py @@ -0,0 +1,70 @@ +import click +from pydantic import BaseModel +from typing import Optional, Literal, get_args + +FORMAT_CHOICES_TYPE = Literal["code", "html", "image", "json", "markdown", "math", "md", "text"] +FORMAT_CHOICES = list(get_args(FORMAT_CHOICES_TYPE)) + +class CellArgs(BaseModel): + type: Literal["root"] = "root" + model_id: str + format: FORMAT_CHOICES_TYPE + reset: bool + +class HelpArgs(BaseModel): + type: Literal["help"] = "help" + +class ListArgs(BaseModel): + type: Literal["list"] = "list" + provider_id: Optional[str] + +class LineMagicGroup(click.Group): + """Helper class to print the help string for cell magics as well when + `%ai --help` is called.""" + def get_help(self, ctx): + with click.Context(cell_magic_parser, info_name="%%ai") as ctx: + click.echo(cell_magic_parser.get_help(ctx)) + click.echo('-' * 78) + with click.Context(line_magic_parser, info_name="%ai") as ctx: + click.echo(super().get_help(ctx)) + +@click.command() +@click.argument('model_id') +@click.option('-f', '--format', + type=click.Choice(FORMAT_CHOICES, case_sensitive=False), + default="markdown", + help="""IPython display to use when rendering output. [default="markdown"]""" +) +@click.option('-r', '--reset', is_flag=True, + help="""Clears the conversation transcript used when interacting with an + OpenAI chat model provider. Does nothing with other providers.""" +) +def cell_magic_parser(**kwargs): + """ + Invokes a language model identified by MODEL_ID, with the prompt being + contained in all lines after the first. Both local model IDs and global + model IDs (with the provider ID explicitly prefixed, followed by a colon) + are accepted. + + To view available language models, please run `%ai list`. + """ + return CellArgs(**kwargs) + +@click.group(cls=LineMagicGroup) +def line_magic_parser(): + """ + Invokes a subcommand. + """ + +@line_magic_parser.command(name='help') +def help_subparser(): + """Show this message and exit.""" + return HelpArgs() + +@line_magic_parser.command(name='list', + short_help="List language models. See `%ai list --help` for options." +) +@click.argument('provider_id', required=False) +def list_subparser(**kwargs): + """List language models, optionally scoped to PROVIDER_ID.""" + return ListArgs(**kwargs) diff --git a/packages/jupyter-ai-magics/pyproject.toml b/packages/jupyter-ai-magics/pyproject.toml index dfc455ba6..1e36bd238 100644 --- a/packages/jupyter-ai-magics/pyproject.toml +++ b/packages/jupyter-ai-magics/pyproject.toml @@ -25,7 +25,8 @@ dependencies = [ "pydantic", "importlib_metadata~=5.2.0", "langchain==0.0.159", - "typing_extensions==4.5.0" + "typing_extensions==4.5.0", + "click~=8.0", ] [project.optional-dependencies]