From 744e65118fa763c1968793718ac5d59fba18a3a5 Mon Sep 17 00:00:00 2001 From: Piyush Jain Date: Thu, 2 Nov 2023 13:30:42 -0700 Subject: [PATCH 1/8] Endpoint args for SM endpoints --- .../jupyter_ai_magics/embedding_providers.py | 3 +- .../jupyter_ai_magics/parsers.py | 22 ++++++++++++ .../jupyter_ai_magics/providers.py | 35 +++++++++++++++++-- 3 files changed, 56 insertions(+), 4 deletions(-) diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py index 5fe522beb..d09907fb0 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py @@ -1,4 +1,4 @@ -from typing import ClassVar, List, Type +from typing import ClassVar, List from jupyter_ai_magics.providers import ( AuthStrategy, @@ -12,7 +12,6 @@ HuggingFaceHubEmbeddings, OpenAIEmbeddings, ) -from langchain.embeddings.base import Embeddings from pydantic import BaseModel, Extra diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py index a6acf3525..deffb3176 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py @@ -32,6 +32,22 @@ + "does nothing with other providers." ) +ENDPOINT_ARGS_SHORT_OPTION = "-e" +ENDPOINT_ARGS_LONG_OPTION = "--endpoint-args" +ENDPOINT_ARGS_HELP = ( + "A JSON value that specifies extra values that will be passed " + "to the SageMaker Endpoint invoke function." +) + +MODEL_ARGS_SHORT_OPTION = "-m" +MODEL_ARGS_LONG_OPTION = "--model-args" +MODEL_ARGS_HELP = ( + "A JSON value that specifies extra values that will be passed to" + "the payload body of the invoke function. This can be useful to" + "pass model tuning parameters such as token count, temperature " + "etc., that affects the response generated by of a model." +) + class CellArgs(BaseModel): type: Literal["root"] = "root" @@ -127,6 +143,12 @@ def get_help(self, ctx): required=False, help=RESPONSE_PATH_HELP, ) +@click.option( + ENDPOINT_ARGS_SHORT_OPTION, + ENDPOINT_ARGS_LONG_OPTION, + required=False, + help=ENDPOINT_ARGS_HELP, +) def cell_magic_parser(**kwargs): """ Invokes a language model identified by MODEL_ID, with the prompt being diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py index 9fdbffa7a..933d834ea 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py @@ -5,7 +5,17 @@ import io import json from concurrent.futures import ThreadPoolExecutor -from typing import Any, ClassVar, Coroutine, Dict, List, Literal, Optional, Union +from typing import ( + Any, + ClassVar, + Coroutine, + Dict, + List, + Literal, + Mapping, + Optional, + Union, +) from jsonpath_ng import parse from langchain.chat_models import ( @@ -232,6 +242,15 @@ def allows_concurrency(self): return True +def pop_with_default(model: Mapping[str, Any], name: str, default: Any) -> Any: + try: + value = model.pop(name) + except KeyError as e: + return default + + return value + + class AI21Provider(BaseProvider, AI21): id = "ai21" name = "AI21" @@ -613,6 +632,9 @@ class SmEndpointProvider(BaseProvider, SagemakerEndpoint): TextField( key="response_path", label="Response path (required)", format="jsonpath" ), + MultilineTextField( + key="endpoint_kwargs", label="Endpoint arguments", format="json" + ), ] def __init__(self, *args, **kwargs): @@ -621,7 +643,16 @@ def __init__(self, *args, **kwargs): content_handler = JsonContentHandler( request_schema=request_schema, response_path=response_path ) - super().__init__(*args, **kwargs, content_handler=content_handler) + + endpoint_kwargs = pop_with_default(kwargs, "endpoint_kwargs", "{}") + endpoint_kwargs = json.loads(endpoint_kwargs) + + super().__init__( + *args, + **kwargs, + content_handler=content_handler, + endpoint_kwargs=endpoint_kwargs, + ) async def _acall(self, *args, **kwargs) -> Coroutine[Any, Any, str]: return await self._call_in_executor(*args, **kwargs) From bde20fb52db71de9e8b1f1868cf49700765d86f3 Mon Sep 17 00:00:00 2001 From: Piyush Jain Date: Fri, 3 Nov 2023 11:36:35 -0700 Subject: [PATCH 2/8] Added model and endpoints kwargs options. --- .../jupyter_ai_magics/magics.py | 13 +++--- .../jupyter_ai_magics/parsers.py | 44 ++++++++++++++++++- .../jupyter_ai_magics/providers.py | 12 +++++ 3 files changed, 59 insertions(+), 10 deletions(-) diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py b/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py index b5741a19b..dc5675b4f 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py @@ -518,14 +518,11 @@ def run_ai_cell(self, args: CellArgs, prompt: str): provider_params["request_schema"] = args.request_schema provider_params["response_path"] = args.response_path - # Validate that the request schema is well-formed JSON - try: - json.loads(args.request_schema) - except json.JSONDecodeError as e: - raise ValueError( - "request-schema must be valid JSON. " - f"Error at line {e.lineno}, column {e.colno}: {e.msg}" - ) from None + if args.model_kwargs: + provider_params["model_kwargs"] = args.model_kwargs + + if args.endpoint_kwargs: + provider_params["endpoint_kwargs"] = args.endpoint_kwargs provider = Provider(**provider_params) diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py index deffb3176..da4e889b1 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py @@ -1,3 +1,4 @@ +import json from typing import Literal, Optional, get_args import click @@ -33,14 +34,14 @@ ) ENDPOINT_ARGS_SHORT_OPTION = "-e" -ENDPOINT_ARGS_LONG_OPTION = "--endpoint-args" +ENDPOINT_ARGS_LONG_OPTION = "--endpoint-kwargs" ENDPOINT_ARGS_HELP = ( "A JSON value that specifies extra values that will be passed " "to the SageMaker Endpoint invoke function." ) MODEL_ARGS_SHORT_OPTION = "-m" -MODEL_ARGS_LONG_OPTION = "--model-args" +MODEL_ARGS_LONG_OPTION = "--model-kwargs" MODEL_ARGS_HELP = ( "A JSON value that specifies extra values that will be passed to" "the payload body of the invoke function. This can be useful to" @@ -58,6 +59,8 @@ class CellArgs(BaseModel): region_name: Optional[str] request_schema: Optional[str] response_path: Optional[str] + model_kwargs: Optional[str] + endpoint_kwargs: Optional[str] # Should match CellArgs, but without "reset" @@ -109,6 +112,19 @@ def get_help(self, ctx): click.echo(super().get_help(ctx)) +def verify_json_value(ctx, param, value): + if not value: + return value + try: + json.loads(value) + except json.JSONDecodeError as e: + raise ValueError( + f"{param.get_error_hint(ctx)} must be valid JSON. " + f"Error at line {e.lineno}, column {e.colno}: {e.msg}" + ) + return value + + @click.command() @click.argument("model_id") @click.option( @@ -136,6 +152,7 @@ def get_help(self, ctx): REQUEST_SCHEMA_LONG_OPTION, required=False, help=REQUEST_SCHEMA_HELP, + callback=verify_json_value, ) @click.option( RESPONSE_PATH_SHORT_OPTION, @@ -148,6 +165,14 @@ def get_help(self, ctx): ENDPOINT_ARGS_LONG_OPTION, required=False, help=ENDPOINT_ARGS_HELP, + callback=verify_json_value, +) +@click.option( + MODEL_ARGS_SHORT_OPTION, + MODEL_ARGS_LONG_OPTION, + required=False, + help=MODEL_ARGS_HELP, + callback=verify_json_value, ) def cell_magic_parser(**kwargs): """ @@ -188,6 +213,7 @@ def line_magic_parser(): REQUEST_SCHEMA_LONG_OPTION, required=False, help=REQUEST_SCHEMA_HELP, + callback=verify_json_value, ) @click.option( RESPONSE_PATH_SHORT_OPTION, @@ -195,6 +221,20 @@ def line_magic_parser(): required=False, help=RESPONSE_PATH_HELP, ) +@click.option( + ENDPOINT_ARGS_SHORT_OPTION, + ENDPOINT_ARGS_LONG_OPTION, + required=False, + help=ENDPOINT_ARGS_HELP, + callback=verify_json_value, +) +@click.option( + MODEL_ARGS_SHORT_OPTION, + MODEL_ARGS_LONG_OPTION, + required=False, + help=MODEL_ARGS_HELP, + callback=verify_json_value, +) def error_subparser(**kwargs): """ Explains the most recent error. Takes the same options (except -r) as diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py index 933d834ea..edc139d5e 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py @@ -677,8 +677,14 @@ class BedrockProvider(BaseProvider, Bedrock): format="text", ), TextField(key="region_name", label="Region name (optional)", format="text"), + MultilineTextField(key="model_kwargs", label="Model Arguments", format="json"), ] + def __init__(self, *args, **kwargs): + model_kwargs = pop_with_default(kwargs, "model_kwargs", "{}") + model_kwargs = json.loads(model_kwargs) + super().__init__(*args, **kwargs, model_kwargs=model_kwargs) + async def _acall(self, *args, **kwargs) -> Coroutine[Any, Any, str]: return await self._call_in_executor(*args, **kwargs) @@ -701,8 +707,14 @@ class BedrockChatProvider(BaseProvider, BedrockChat): format="text", ), TextField(key="region_name", label="Region name (optional)", format="text"), + MultilineTextField(key="model_kwargs", label="Model Arguments", format="json"), ] + def __init__(self, *args, **kwargs): + model_kwargs = pop_with_default(kwargs, "model_kwargs", "{}") + model_kwargs = json.loads(model_kwargs) + super().__init__(*args, **kwargs, model_kwargs=model_kwargs) + async def _acall(self, *args, **kwargs) -> Coroutine[Any, Any, str]: return await self._call_in_executor(*args, **kwargs) From 628641f4188d88a14ec42330bf6c62defc92acc4 Mon Sep 17 00:00:00 2001 From: Piyush Jain Date: Tue, 7 Nov 2023 18:57:52 -0800 Subject: [PATCH 3/8] Added configurable option for model parameters. --- docs/source/users/index.md | 69 +++++++++++++++++++ .../jupyter_ai_magics/providers.py | 34 +-------- .../jupyter_ai/chat_handlers/ask.py | 3 +- .../jupyter_ai/chat_handlers/base.py | 11 ++- .../jupyter_ai/chat_handlers/default.py | 3 +- .../jupyter_ai/chat_handlers/generate.py | 4 +- packages/jupyter-ai/jupyter_ai/extension.py | 17 ++++- 7 files changed, 103 insertions(+), 38 deletions(-) diff --git a/docs/source/users/index.md b/docs/source/users/index.md index b52fb66bb..4226df56b 100644 --- a/docs/source/users/index.md +++ b/docs/source/users/index.md @@ -882,3 +882,72 @@ To allow more than one provider in the allowlist, repeat the runtime configurati ``` jupyter lab --AiExtension.allowed_providers=openai --AiExtension.allowed_providers=ai21 ``` + +### Model parameters +This configuration allows specifying arbitrary parameters that are unpacked and passed to the provider class. +This is useful for passing parameters such as model tuning that affect the response generation by the model. +This is also an appropriate place to pass in custom attributes required by certain providers/models. + +The accepted value should be a dictionary, with top level keys as the model id (provider:model_id), and value +should be any arbitrary dictionary which is unpacked and passed as is to the provider class. + +#### Configuring as a startup option +In this sample, the `bedrock` provider will be created with the value for `model_kwargs` when `ai21.j2-mid-v1` model is selected. + +```bash +jupyter lab --AiExtension.model_parameters {"bedrock:ai21.j2-mid-v1":{"model_kwargs":{"maxTokens":200}}} +``` +The above will result in the following LLM class to be generated. + +```python +BedrockProvider(model_kwargs={"maxTokens":200}, ...) +``` + +Here is another example, where `anthropic` provider will be created with the values for `max_tokens` and `temperature`, when `claude-2` model is selected. + + +```bash +jupyter lab --AiExtension.model_parameters {"anthropic:claude-2":{"max_tokens":1024,"temperature":0.9}} +``` +The above will result in the following LLM class to be generated. + +```python +AnthropicProvider(max_tokens=1024, temperature=0.9, ...) +``` + +#### Configuring as a config file +This configuration can also be specified in a config file in json format. The file should be named `jupyter_jupyter_ai_config.json` and saved in a path that JupyterLab can pick from. You can find this +path by running `jupyter --paths` command, and picking one of the paths from the `config` section. + +Here is an example of running the `jupyter --paths` command. + +```bash +(jupyter-ai-lab4) ➜ jupyter --paths +config: + /opt/anaconda3/envs/jupyter-ai-lab4/etc/jupyter + /Users/3coins/.jupyter + /Users/3coins/.local/etc/jupyter + /usr/3coins/etc/jupyter + /etc/jupyter +data: + /opt/anaconda3/envs/jupyter-ai-lab4/share/jupyter + /Users/3coins/Library/Jupyter + /Users/3coins/.local/share/jupyter + /usr/local/share/jupyter + /usr/share/jupyter +runtime: + /Users/3coins/Library/Jupyter/runtime +``` + +Here is an example for configuring the `bedrock` provider for `ai21.j2-mid-v1` model. +```json +{ + "AiExtension": { + "bedrock:ai21.j2-mid-v1": { + "model_kwargs": { + "maxTokens": 200 + } + } + } +} +``` diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py index edc139d5e..a17067921 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py @@ -242,15 +242,6 @@ def allows_concurrency(self): return True -def pop_with_default(model: Mapping[str, Any], name: str, default: Any) -> Any: - try: - value = model.pop(name) - except KeyError as e: - return default - - return value - - class AI21Provider(BaseProvider, AI21): id = "ai21" name = "AI21" @@ -632,9 +623,6 @@ class SmEndpointProvider(BaseProvider, SagemakerEndpoint): TextField( key="response_path", label="Response path (required)", format="jsonpath" ), - MultilineTextField( - key="endpoint_kwargs", label="Endpoint arguments", format="json" - ), ] def __init__(self, *args, **kwargs): @@ -644,15 +632,7 @@ def __init__(self, *args, **kwargs): request_schema=request_schema, response_path=response_path ) - endpoint_kwargs = pop_with_default(kwargs, "endpoint_kwargs", "{}") - endpoint_kwargs = json.loads(endpoint_kwargs) - - super().__init__( - *args, - **kwargs, - content_handler=content_handler, - endpoint_kwargs=endpoint_kwargs, - ) + super().__init__(*args, **kwargs, content_handler=content_handler) async def _acall(self, *args, **kwargs) -> Coroutine[Any, Any, str]: return await self._call_in_executor(*args, **kwargs) @@ -677,14 +657,8 @@ class BedrockProvider(BaseProvider, Bedrock): format="text", ), TextField(key="region_name", label="Region name (optional)", format="text"), - MultilineTextField(key="model_kwargs", label="Model Arguments", format="json"), ] - def __init__(self, *args, **kwargs): - model_kwargs = pop_with_default(kwargs, "model_kwargs", "{}") - model_kwargs = json.loads(model_kwargs) - super().__init__(*args, **kwargs, model_kwargs=model_kwargs) - async def _acall(self, *args, **kwargs) -> Coroutine[Any, Any, str]: return await self._call_in_executor(*args, **kwargs) @@ -707,14 +681,8 @@ class BedrockChatProvider(BaseProvider, BedrockChat): format="text", ), TextField(key="region_name", label="Region name (optional)", format="text"), - MultilineTextField(key="model_kwargs", label="Model Arguments", format="json"), ] - def __init__(self, *args, **kwargs): - model_kwargs = pop_with_default(kwargs, "model_kwargs", "{}") - model_kwargs = json.loads(model_kwargs) - super().__init__(*args, **kwargs, model_kwargs=model_kwargs) - async def _acall(self, *args, **kwargs) -> Coroutine[Any, Any, str]: return await self._call_in_executor(*args, **kwargs) diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py index 2f3f1388a..e5c852051 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py +++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py @@ -36,7 +36,8 @@ def __init__(self, retriever, *args, **kwargs): def create_llm_chain( self, provider: Type[BaseProvider], provider_params: Dict[str, str] ): - self.llm = provider(**provider_params) + model_parameters = self.get_model_parameters(provider, provider_params) + self.llm = provider(**provider_params, **model_parameters) memory = ConversationBufferWindowMemory( memory_key="chat_history", return_messages=True, k=2 ) diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py index 7894dcfa5..5ffe65c7c 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py +++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py @@ -3,7 +3,7 @@ import traceback # necessary to prevent circular import -from typing import TYPE_CHECKING, Dict, Optional, Type +from typing import TYPE_CHECKING, Any, Dict, Optional, Type from uuid import uuid4 from jupyter_ai.config_manager import ConfigManager, Logger @@ -23,10 +23,12 @@ def __init__( log: Logger, config_manager: ConfigManager, root_chat_handlers: Dict[str, "RootChatHandler"], + model_parameters: Dict[str, Dict], ): self.log = log self.config_manager = config_manager self._root_chat_handlers = root_chat_handlers + self.model_parameters = model_parameters self.parser = argparse.ArgumentParser() self.llm = None self.llm_params = None @@ -122,6 +124,13 @@ def get_llm_chain(self): self.llm_params = lm_provider_params return self.llm_chain + def get_model_parameters( + self, provider: Type[BaseProvider], provider_params: Dict[str, str] + ): + return self.model_parameters.get( + f"{provider.id}:{provider_params['model_id']}", {} + ) + def create_llm_chain( self, provider: Type[BaseProvider], provider_params: Dict[str, str] ): diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/default.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/default.py index 3638b6cfb..d329e05e2 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/default.py +++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/default.py @@ -40,7 +40,8 @@ def __init__(self, chat_history: List[ChatMessage], *args, **kwargs): def create_llm_chain( self, provider: Type[BaseProvider], provider_params: Dict[str, str] ): - llm = provider(**provider_params) + model_parameters = self.get_model_parameters(provider, provider_params) + llm = provider(**provider_params, **model_parameters) if llm.is_chat_provider: prompt_template = ChatPromptTemplate.from_messages( diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py index 5036c8dfb..cdddee92d 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py +++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py @@ -226,7 +226,9 @@ def __init__(self, root_dir: str, *args, **kwargs): def create_llm_chain( self, provider: Type[BaseProvider], provider_params: Dict[str, str] ): - llm = provider(**provider_params) + model_parameters = self.get_model_parameters(provider, provider_params) + llm = provider(**provider_params, **model_parameters) + self.llm = llm return llm diff --git a/packages/jupyter-ai/jupyter_ai/extension.py b/packages/jupyter-ai/jupyter_ai/extension.py index a2ecd5245..500bc9649 100644 --- a/packages/jupyter-ai/jupyter_ai/extension.py +++ b/packages/jupyter-ai/jupyter_ai/extension.py @@ -4,7 +4,7 @@ from jupyter_ai.chat_handlers.learn import Retriever from jupyter_ai_magics.utils import get_em_providers, get_lm_providers from jupyter_server.extension.application import ExtensionApp -from traitlets import List, Unicode +from traitlets import Dict, List, Unicode from .chat_handlers import ( AskChatHandler, @@ -81,6 +81,17 @@ class AiExtension(ExtensionApp): config=True, ) + model_parameters = Dict( + key_trait=Unicode(), + value_trait=Dict(), + default_value={}, + help="""Key-value pairs for model id and corresponding parameters that + are passed to the provider class. The values are unpacked and passed to + the provider class as-is.""", + allow_none=True, + config=True, + ) + def initialize_settings(self): start = time.time() @@ -96,6 +107,9 @@ def initialize_settings(self): self.log.info(f"Configured model allowlist: {self.allowed_models}") self.log.info(f"Configured model blocklist: {self.blocked_models}") + + self.settings["model_parameters"] = self.model_parameters + # Fetch LM & EM providers self.settings["lm_providers"] = get_lm_providers( log=self.log, restrictions=restrictions @@ -147,6 +161,7 @@ def initialize_settings(self): "log": self.log, "config_manager": self.settings["jai_config_manager"], "root_chat_handlers": self.settings["jai_root_chat_handlers"], + "model_parameters": self.settings["model_parameters"], } default_chat_handler = DefaultChatHandler( **chat_handler_kwargs, chat_history=self.settings["chat_history"] From 400e3cf0b52ad941a002a418cdc57bf2c442a30a Mon Sep 17 00:00:00 2001 From: Piyush Jain Date: Tue, 7 Nov 2023 19:36:32 -0800 Subject: [PATCH 4/8] Updated magics, added model_parameters, removed model_kwargs and endpoint_kwargs. --- docs/source/users/index.md | 4 +- .../jupyter_ai_magics/magics.py | 8 +-- .../jupyter_ai_magics/parsers.py | 49 ++++++------------- 3 files changed, 18 insertions(+), 43 deletions(-) diff --git a/docs/source/users/index.md b/docs/source/users/index.md index 4226df56b..b5e7066b5 100644 --- a/docs/source/users/index.md +++ b/docs/source/users/index.md @@ -888,8 +888,8 @@ This configuration allows specifying arbitrary parameters that are unpacked and This is useful for passing parameters such as model tuning that affect the response generation by the model. This is also an appropriate place to pass in custom attributes required by certain providers/models. -The accepted value should be a dictionary, with top level keys as the model id (provider:model_id), and value -should be any arbitrary dictionary which is unpacked and passed as is to the provider class. +The accepted value is a dictionary, with top level keys as the model id (provider:model_id), and value +should be any arbitrary dictionary which is unpacked and passed as-is to the provider class. #### Configuring as a startup option In this sample, the `bedrock` provider will be created with the value for `model_kwargs` when `ai21.j2-mid-v1` model is selected. diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py b/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py index dc5675b4f..5266b071d 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py @@ -518,13 +518,9 @@ def run_ai_cell(self, args: CellArgs, prompt: str): provider_params["request_schema"] = args.request_schema provider_params["response_path"] = args.response_path - if args.model_kwargs: - provider_params["model_kwargs"] = args.model_kwargs + model_parameters = json.loads(args.model_parameters) - if args.endpoint_kwargs: - provider_params["endpoint_kwargs"] = args.endpoint_kwargs - - provider = Provider(**provider_params) + provider = Provider(**provider_params, **model_parameters) # Apply a prompt template. prompt = provider.get_prompt_template(args.format).format(prompt=prompt) diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py index da4e889b1..a3b14bbb0 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py @@ -33,20 +33,12 @@ + "does nothing with other providers." ) -ENDPOINT_ARGS_SHORT_OPTION = "-e" -ENDPOINT_ARGS_LONG_OPTION = "--endpoint-kwargs" -ENDPOINT_ARGS_HELP = ( +MODEL_PARAMETERS_SHORT_OPTION = "-m" +MODEL_PARAMETERS_LONG_OPTION = "--model-parameters" +MODEL_PARAMETERS_HELP = ( "A JSON value that specifies extra values that will be passed " - "to the SageMaker Endpoint invoke function." -) - -MODEL_ARGS_SHORT_OPTION = "-m" -MODEL_ARGS_LONG_OPTION = "--model-kwargs" -MODEL_ARGS_HELP = ( - "A JSON value that specifies extra values that will be passed to" - "the payload body of the invoke function. This can be useful to" - "pass model tuning parameters such as token count, temperature " - "etc., that affects the response generated by of a model." + "to the model. The accepted value parsed to a dict, unpacked " + "and passed as-is to the provider class." ) @@ -59,8 +51,7 @@ class CellArgs(BaseModel): region_name: Optional[str] request_schema: Optional[str] response_path: Optional[str] - model_kwargs: Optional[str] - endpoint_kwargs: Optional[str] + model_parameters: Optional[str] # Should match CellArgs, but without "reset" @@ -161,18 +152,12 @@ def verify_json_value(ctx, param, value): help=RESPONSE_PATH_HELP, ) @click.option( - ENDPOINT_ARGS_SHORT_OPTION, - ENDPOINT_ARGS_LONG_OPTION, - required=False, - help=ENDPOINT_ARGS_HELP, - callback=verify_json_value, -) -@click.option( - MODEL_ARGS_SHORT_OPTION, - MODEL_ARGS_LONG_OPTION, + MODEL_PARAMETERS_SHORT_OPTION, + MODEL_PARAMETERS_LONG_OPTION, required=False, - help=MODEL_ARGS_HELP, + help=MODEL_PARAMETERS_HELP, callback=verify_json_value, + default="{}", ) def cell_magic_parser(**kwargs): """ @@ -222,18 +207,12 @@ def line_magic_parser(): help=RESPONSE_PATH_HELP, ) @click.option( - ENDPOINT_ARGS_SHORT_OPTION, - ENDPOINT_ARGS_LONG_OPTION, - required=False, - help=ENDPOINT_ARGS_HELP, - callback=verify_json_value, -) -@click.option( - MODEL_ARGS_SHORT_OPTION, - MODEL_ARGS_LONG_OPTION, + MODEL_PARAMETERS_SHORT_OPTION, + MODEL_PARAMETERS_LONG_OPTION, required=False, - help=MODEL_ARGS_HELP, + help=MODEL_PARAMETERS_HELP, callback=verify_json_value, + default="{}", ) def error_subparser(**kwargs): """ From 486b2fee678f1baaf40a40ebdb84cafe31edf110 Mon Sep 17 00:00:00 2001 From: Piyush Jain Date: Tue, 7 Nov 2023 20:46:13 -0800 Subject: [PATCH 5/8] Fixes %ai error for SM endpoints. --- packages/jupyter-ai-magics/jupyter_ai_magics/magics.py | 8 +++++--- packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py | 3 ++- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py b/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py index 5266b071d..8763566f4 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py @@ -392,9 +392,11 @@ def handle_error(self, args: ErrorArgs): prompt = f"Explain the following error:\n\n{last_error}" # Set CellArgs based on ErrorArgs - cell_args = CellArgs( - type="root", model_id=args.model_id, format=args.format, reset=False - ) + values = args.dict() + values["type"] = "root" + values["reset"] = False + cell_args = CellArgs(**values) + return self.run_ai_cell(cell_args, prompt) def _append_exchange_openai(self, prompt: str, output: str): diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py index a3b14bbb0..cadd41f4a 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/parsers.py @@ -47,11 +47,11 @@ class CellArgs(BaseModel): model_id: str format: FORMAT_CHOICES_TYPE reset: bool + model_parameters: Optional[str] # The following parameters are required only for SageMaker models region_name: Optional[str] request_schema: Optional[str] response_path: Optional[str] - model_parameters: Optional[str] # Should match CellArgs, but without "reset" @@ -59,6 +59,7 @@ class ErrorArgs(BaseModel): type: Literal["error"] = "error" model_id: str format: FORMAT_CHOICES_TYPE + model_parameters: Optional[str] # The following parameters are required only for SageMaker models region_name: Optional[str] request_schema: Optional[str] From 35bd907bc664322a6b0bff81b96f7614d06164fd Mon Sep 17 00:00:00 2001 From: Piyush Jain Date: Wed, 8 Nov 2023 09:35:19 -0800 Subject: [PATCH 6/8] Fixed docs --- docs/source/users/index.md | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/docs/source/users/index.md b/docs/source/users/index.md index b5e7066b5..cf5e1b7b6 100644 --- a/docs/source/users/index.md +++ b/docs/source/users/index.md @@ -943,9 +943,11 @@ Here is an example for configuring the `bedrock` provider for `ai21.j2-mid-v1` m ```json { "AiExtension": { - "bedrock:ai21.j2-mid-v1": { - "model_kwargs": { - "maxTokens": 200 + "model_parameters": { + "bedrock:ai21.j2-mid-v1": { + "model_kwargs": { + "maxTokens": 200 + } } } } From 9f6f23879a19559ccf0bb7201fdb476d75e60c84 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 8 Nov 2023 17:35:36 +0000 Subject: [PATCH 7/8] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- packages/jupyter-ai/jupyter_ai/extension.py | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/jupyter-ai/jupyter_ai/extension.py b/packages/jupyter-ai/jupyter_ai/extension.py index 500bc9649..dad3e2042 100644 --- a/packages/jupyter-ai/jupyter_ai/extension.py +++ b/packages/jupyter-ai/jupyter_ai/extension.py @@ -107,7 +107,6 @@ def initialize_settings(self): self.log.info(f"Configured model allowlist: {self.allowed_models}") self.log.info(f"Configured model blocklist: {self.blocked_models}") - self.settings["model_parameters"] = self.model_parameters # Fetch LM & EM providers From 71e2f13687e95d389192392fceabdab2333355a1 Mon Sep 17 00:00:00 2001 From: david qiu Date: Wed, 8 Nov 2023 10:41:20 -0800 Subject: [PATCH 8/8] 430 fixes (#2) * log configured model_parameters * fix markdown formatting in docs * fix single quotes and use preferred traitlets CLI syntax --- docs/source/users/index.md | 73 ++++++++++++++++----- packages/jupyter-ai/jupyter_ai/extension.py | 1 + 2 files changed, 56 insertions(+), 18 deletions(-) diff --git a/docs/source/users/index.md b/docs/source/users/index.md index cf5e1b7b6..fb52e29ba 100644 --- a/docs/source/users/index.md +++ b/docs/source/users/index.md @@ -855,69 +855,104 @@ The `--response-path` option is a [JSONPath](https://goessner.net/articles/JsonP ## Configuration -You can specify an allowlist, to only allow only a certain list of providers, or a blocklist, to block some providers. +You can specify an allowlist, to only allow only a certain list of providers, or +a blocklist, to block some providers. ### Blocklisting providers -This configuration allows for blocking specific providers in the settings panel. This list takes precedence over the allowlist in the next section. + +This configuration allows for blocking specific providers in the settings panel. +This list takes precedence over the allowlist in the next section. ``` jupyter lab --AiExtension.blocked_providers=openai ``` -To block more than one provider in the block-list, repeat the runtime configuration. +To block more than one provider in the block-list, repeat the runtime +configuration. ``` jupyter lab --AiExtension.blocked_providers=openai --AiExtension.blocked_providers=ai21 ``` ### Allowlisting providers -This configuration allows for filtering the list of providers in the settings panel to only an allowlisted set of providers. + +This configuration allows for filtering the list of providers in the settings +panel to only an allowlisted set of providers. ``` jupyter lab --AiExtension.allowed_providers=openai ``` -To allow more than one provider in the allowlist, repeat the runtime configuration. +To allow more than one provider in the allowlist, repeat the runtime +configuration. ``` jupyter lab --AiExtension.allowed_providers=openai --AiExtension.allowed_providers=ai21 ``` ### Model parameters -This configuration allows specifying arbitrary parameters that are unpacked and passed to the provider class. -This is useful for passing parameters such as model tuning that affect the response generation by the model. -This is also an appropriate place to pass in custom attributes required by certain providers/models. -The accepted value is a dictionary, with top level keys as the model id (provider:model_id), and value -should be any arbitrary dictionary which is unpacked and passed as-is to the provider class. +This configuration allows specifying arbitrary parameters that are unpacked and +passed to the provider class. This is useful for passing parameters such as +model tuning that affect the response generation by the model. This is also an +appropriate place to pass in custom attributes required by certain +providers/models. + +The accepted value is a dictionary, with top level keys as the model id +(provider:model_id), and value should be any arbitrary dictionary which is +unpacked and passed as-is to the provider class. #### Configuring as a startup option -In this sample, the `bedrock` provider will be created with the value for `model_kwargs` when `ai21.j2-mid-v1` model is selected. + +In this sample, the `bedrock` provider will be created with the value for +`model_kwargs` when `ai21.j2-mid-v1` model is selected. ```bash -jupyter lab --AiExtension.model_parameters {"bedrock:ai21.j2-mid-v1":{"model_kwargs":{"maxTokens":200}}} +jupyter lab --AiExtension.model_parameters bedrock:ai21.j2-mid-v1='{"model_kwargs":{"maxTokens":200}}' ``` -The above will result in the following LLM class to be generated. + +Note the usage of single quotes surrounding the dictionary to escape the double +quotes. This is required in some shells. The above will result in the following +LLM class to be generated. ```python BedrockProvider(model_kwargs={"maxTokens":200}, ...) ``` -Here is another example, where `anthropic` provider will be created with the values for `max_tokens` and `temperature`, when `claude-2` model is selected. +Here is another example, where `anthropic` provider will be created with the +values for `max_tokens` and `temperature`, when `claude-2` model is selected. ```bash -jupyter lab --AiExtension.model_parameters {"anthropic:claude-2":{"max_tokens":1024,"temperature":0.9}} +jupyter lab --AiExtension.model_parameters anthropic:claude-2='{"max_tokens":1024,"temperature":0.9}' ``` + The above will result in the following LLM class to be generated. ```python AnthropicProvider(max_tokens=1024, temperature=0.9, ...) ``` +To pass multiple sets of model parameters for multiple models in the +command-line, you can append them as additional arguments to +`--AiExtension.model_parameters`, as shown below. + +```bash +jupyter lab \ +--AiExtension.model_parameters bedrock:ai21.j2-mid-v1='{"model_kwargs":{"maxTokens":200}}' \ +--AiExtension.model_parameters anthropic:claude-2='{"max_tokens":1024,"temperature":0.9}' +``` + +However, for more complex configuration, we highly recommend that you specify +this in a dedicated configuration file. We will describe how to do so in the +following section. + #### Configuring as a config file -This configuration can also be specified in a config file in json format. The file should be named `jupyter_jupyter_ai_config.json` and saved in a path that JupyterLab can pick from. You can find this -path by running `jupyter --paths` command, and picking one of the paths from the `config` section. + +This configuration can also be specified in a config file in json format. The +file should be named `jupyter_jupyter_ai_config.json` and saved in a path that +JupyterLab can pick from. You can find this path by running `jupyter --paths` +command, and picking one of the paths from the `config` section. Here is an example of running the `jupyter --paths` command. @@ -939,7 +974,9 @@ runtime: /Users/3coins/Library/Jupyter/runtime ``` -Here is an example for configuring the `bedrock` provider for `ai21.j2-mid-v1` model. +Here is an example for configuring the `bedrock` provider for `ai21.j2-mid-v1` +model. + ```json { "AiExtension": { diff --git a/packages/jupyter-ai/jupyter_ai/extension.py b/packages/jupyter-ai/jupyter_ai/extension.py index dad3e2042..8ab8c0cc6 100644 --- a/packages/jupyter-ai/jupyter_ai/extension.py +++ b/packages/jupyter-ai/jupyter_ai/extension.py @@ -108,6 +108,7 @@ def initialize_settings(self): self.log.info(f"Configured model blocklist: {self.blocked_models}") self.settings["model_parameters"] = self.model_parameters + self.log.info(f"Configured model parameters: {self.model_parameters}") # Fetch LM & EM providers self.settings["lm_providers"] = get_lm_providers(