change model provisioning strategy

This commit is contained in:
LUIS NOVO 2024-11-04 09:49:11 -03:00
parent d9c0c93deb
commit b4ba3ef4c8
2 changed files with 7 additions and 5 deletions

View file

@ -8,7 +8,7 @@ from open_notebook.prompter import Prompter
from open_notebook.utils import token_count
def provision_langchain_model(content, config, default_type) -> BaseChatModel:
def provision_langchain_model(content, config, default_type, **kwargs) -> BaseChatModel:
"""
Returns the best model to use based on the context size and on whether there is a specific model being requested in Config.
If context > 105_000, returns the large_context_model
@ -21,11 +21,13 @@ def provision_langchain_model(content, config, default_type) -> BaseChatModel:
logger.debug(
f"Using large context model because the content has {tokens} tokens"
)
model = model_manager.get_default_model("large_context")
model = model_manager.get_default_model("large_context", **kwargs)
elif config.get("configurable", {}).get("model_id"):
model = model_manager.get_model(config.get("configurable", {}).get("model_id"))
model = model_manager.get_model(
config.get("configurable", {}).get("model_id"), **kwargs
)
else:
model = model_manager.get_default_model(default_type)
model = model_manager.get_default_model(default_type, **kwargs)
assert isinstance(model, LanguageModel), f"Model is not a LanguageModel: {model}"
return model.to_langchain()