baveshraam's picture
FIX: SurrealDB 2.0 migration syntax and Frontend/CORS link
f871fed
from esperanto import LanguageModel
from langchain_core.language_models.chat_models import BaseChatModel
from loguru import logger
from open_notebook.domain.models import model_manager
from open_notebook.utils import token_count
async def provision_langchain_model(
content, model_id, 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
If model_id is specified in Config, returns that model
Otherwise, returns the default model for the given type
"""
tokens = token_count(content)
model = None
try:
if tokens > 105_000:
logger.debug(
f"Using large context model because the content has {tokens} tokens"
)
model = await model_manager.get_default_model("large_context", **kwargs)
elif model_id:
logger.debug(f"Attempting to use specified model: {model_id}")
model = await model_manager.get_model(model_id, **kwargs)
else:
logger.debug(f"Using default model for type: {default_type}")
model = await model_manager.get_default_model(default_type, **kwargs)
except Exception as e:
logger.warning(f"Failed to get requested model (model_id={model_id}, type={default_type}): {e}")
# Fallback to chat model if specified model fails
if default_type != "chat":
logger.info("Falling back to default chat model")
try:
model = await model_manager.get_default_model("chat", **kwargs)
except Exception as chat_error:
logger.error(f"Failed to get fallback chat model: {chat_error}")
raise RuntimeError(f"No available models. Please configure at least one model in settings.") from chat_error
if model is None:
raise RuntimeError(f"Failed to provision model (model_id={model_id}, type={default_type}). Please configure models in settings.")
logger.debug(f"Using model: {model}")
assert isinstance(model, LanguageModel), f"Model is not a LanguageModel: {model}"
return model.to_langchain()