diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/ollama.py b/packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/ollama.py index 5babc5adb..bf7d8474a 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/ollama.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/ollama.py @@ -1,7 +1,7 @@ from langchain_ollama import ChatOllama, OllamaEmbeddings from ..embedding_providers import BaseEmbeddingsProvider -from ..providers import BaseProvider, EnvAuthStrategy, TextField +from ..providers import BaseProvider, TextField class OllamaProvider(BaseProvider, ChatOllama): @@ -23,10 +23,14 @@ class OllamaEmbeddingsProvider(BaseEmbeddingsProvider, OllamaEmbeddings): id = "ollama" name = "Ollama" # source: https://ollama.com/library + model_id_key = "model" models = [ "nomic-embed-text", "mxbai-embed-large", "all-minilm", "snowflake-arctic-embed", ] - model_id_key = "model" + registry = True + fields = [ + TextField(key="base_url", label="Base API URL (optional)", format="text"), + ] diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/openai.py b/packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/openai.py index 7e97995d4..34ca76a8e 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/openai.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/openai.py @@ -107,6 +107,12 @@ class OpenAIEmbeddingsProvider(BaseEmbeddingsProvider, OpenAIEmbeddings): model_id_key = "model" pypi_package_deps = ["langchain_openai"] auth_strategy = EnvAuthStrategy(name="OPENAI_API_KEY") + registry = True + fields = [ + TextField( + key="openai_api_base", label="Base API URL (optional)", format="text" + ), + ] class AzureOpenAIEmbeddingsProvider(BaseEmbeddingsProvider, AzureOpenAIEmbeddings): @@ -122,5 +128,7 @@ class AzureOpenAIEmbeddingsProvider(BaseEmbeddingsProvider, AzureOpenAIEmbedding auth_strategy = EnvAuthStrategy( name="AZURE_OPENAI_API_KEY", keyword_param="openai_api_key" ) - registry = True + fields = [ + TextField(key="azure_endpoint", label="Base API URL (optional)", format="text"), + ] diff --git a/packages/jupyter-ai/jupyter_ai/config_manager.py b/packages/jupyter-ai/jupyter_ai/config_manager.py index 7b309faae..4732e152c 100644 --- a/packages/jupyter-ai/jupyter_ai/config_manager.py +++ b/packages/jupyter-ai/jupyter_ai/config_manager.py @@ -462,6 +462,13 @@ def _provider_params(self, key, listing, completions: bool = False): else: fields = config.fields.get(model_uid, {}) + # exclude empty fields + # TODO: modify the config manager to never save empty fields in the + # first place. + for field_key in fields: + if isinstance(fields[field_key], str) and not len(fields[field_key]): + fields[field_key] = None + # get authn fields _, Provider = get_em_provider(model_uid, listing) authn_fields = {} diff --git a/packages/jupyter-ai/src/components/chat-settings.tsx b/packages/jupyter-ai/src/components/chat-settings.tsx index a1ad0a9b6..c32eb46fd 100644 --- a/packages/jupyter-ai/src/components/chat-settings.tsx +++ b/packages/jupyter-ai/src/components/chat-settings.tsx @@ -88,6 +88,9 @@ export function ChatSettings(props: ChatSettingsProps): JSX.Element { const [apiKeys, setApiKeys] = useState>({}); const [sendWse, setSendWse] = useState(false); const [fields, setFields] = useState>({}); + const [embeddingModelFields, setEmbeddingModelFields] = useState< + Record + >({}); const [isCompleterEnabled, setIsCompleterEnabled] = useState( props.completionProvider && props.completionProvider.isEnabled() @@ -188,7 +191,15 @@ export function ChatSettings(props: ChatSettingsProps): JSX.Element { const currFields: Record = server.config.fields?.[lmGlobalId] ?? {}; setFields(currFields); - }, [server, lmProvider]); + + if (!emGlobalId) { + return; + } + + const initEmbeddingModelFields: Record = + server.config.fields?.[emGlobalId] ?? {}; + setEmbeddingModelFields(initEmbeddingModelFields); + }, [server, lmGlobalId, emGlobalId]); const handleSave = async () => { // compress fields with JSON values @@ -222,6 +233,9 @@ export function ChatSettings(props: ChatSettingsProps): JSX.Element { }), ...(clmGlobalId && { [clmGlobalId]: fields + }), + ...(emGlobalId && { + [emGlobalId]: embeddingModelFields }) } }), @@ -376,26 +390,35 @@ export function ChatSettings(props: ChatSettingsProps): JSX.Element { {/* Embedding model section */}

Embedding model

{server.emProviders.providers.length > 0 ? ( - { + const emGid = e.target.value === 'null' ? null : e.target.value; + setEmGlobalId(emGid); + }} + MenuProps={{ sx: { maxHeight: '50%', minHeight: 400 } }} + > + None + {server.emProviders.providers.map(emp => + emp.models + .filter(em => em !== '*') // TODO: support registry providers + .map(em => ( + + {emp.name} :: {em} + + )) + )} + + {emGlobalId && ( + )} - + ) : (

No embedding models available.

)}