From bea43f3ce7de3c492e1a9f478713e7ce00cfc297 Mon Sep 17 00:00:00 2001 From: LUIS NOVO Date: Sun, 8 Jun 2025 19:38:43 -0300 Subject: [PATCH] feat: implement the new model management based on esperanto framework --- open_notebook/domain/base.py | 20 ++--------- open_notebook/domain/models.py | 62 ++++++++++++++++++++++++-------- open_notebook/domain/notebook.py | 15 +++----- open_notebook/graphs/utils.py | 3 +- 4 files changed, 58 insertions(+), 42 deletions(-) diff --git a/open_notebook/domain/base.py b/open_notebook/domain/base.py index 2813e8c..eb4f000 100644 --- a/open_notebook/domain/base.py +++ b/open_notebook/domain/base.py @@ -1,22 +1,8 @@ from datetime import datetime -from typing import ( - Any, - ClassVar, - Dict, - List, - Optional, - Type, - TypeVar, - cast, -) +from typing import Any, ClassVar, Dict, List, Optional, Type, TypeVar, cast from loguru import logger -from pydantic import ( - BaseModel, - ValidationError, - field_validator, - model_validator, -) +from pydantic import BaseModel, ValidationError, field_validator, model_validator from open_notebook.database.repository import ( repo_create, @@ -140,7 +126,7 @@ class ObjectModel(BaseModel): "No embedding model found. Content will not be searchable." ) data["embedding"] = ( - EMBEDDING_MODEL.embed(embedding_content) + EMBEDDING_MODEL.embed([embedding_content])[0] if EMBEDDING_MODEL else [] ) diff --git a/open_notebook/domain/models.py b/open_notebook/domain/models.py index 645aa94..3c273a2 100644 --- a/open_notebook/domain/models.py +++ b/open_notebook/domain/models.py @@ -1,16 +1,18 @@ -from typing import ClassVar, Dict, Optional +from typing import ClassVar, Dict, Optional, Union -from open_notebook.database.repository import repo_query -from open_notebook.domain.base import ObjectModel, RecordModel -from open_notebook.models import ( - MODEL_CLASS_MAP, +from esperanto import ( + AIFactory, EmbeddingModel, LanguageModel, - ModelType, SpeechToTextModel, TextToSpeechModel, ) +from open_notebook.database.repository import repo_query +from open_notebook.domain.base import ObjectModel, RecordModel + +ModelType = Union[LanguageModel, EmbeddingModel, SpeechToTextModel, TextToSpeechModel] + class Model(ObjectModel): table_name: ClassVar[str] = "model" @@ -75,21 +77,53 @@ class ModelManager: if not model: raise ValueError(f"Model with ID {model_id} not found") - if not model.type or model.type not in MODEL_CLASS_MAP: + if not model.type or model.type not in [ + "language", + "embedding", + "speech_to_text", + "text_to_speech", + ]: raise ValueError(f"Invalid model type: {model.type}") - provider_map = MODEL_CLASS_MAP[model.type] - if model.provider not in provider_map: + # todo: change to providers in the future + if model.provider not in [ + "ollama", + "openrouter", + "vertexai-anthropic", + "litellm", + "vertexai", + "anthropic", + "openai", + "xai", + ]: raise ValueError( f"Provider {model.provider} not compatible with {model.type} models" ) - model_class = provider_map[model.provider] - model_instance = model_class(model_name=model.name, **kwargs) - - # Special handling for language models that need langchain conversion if model.type == "language": - model_instance = model_instance + model_instance: LanguageModel = AIFactory.create_language( + model_name=model.name, + provider=model.provider, + config=kwargs, + ) + elif model.type == "embedding": + model_instance: EmbeddingModel = AIFactory.create_embedding( + model_name=model.name, + provider=model.provider, + config=kwargs, + ) + elif model.type == "speech_to_text": + model_instance: SpeechToTextModel = AIFactory.create_speech_to_text( + model_name=model.name, + provider=model.provider, + config=kwargs, + ) + elif model.type == "text_to_speech": + model_instance: TextToSpeechModel = AIFactory.create_text_to_speech( + model_name=model.name, + provider=model.provider, + config=kwargs, + ) self._model_cache[cache_key] = model_instance return model_instance diff --git a/open_notebook/domain/notebook.py b/open_notebook/domain/notebook.py index b2a8833..5387c52 100644 --- a/open_notebook/domain/notebook.py +++ b/open_notebook/domain/notebook.py @@ -4,15 +4,10 @@ from typing import Any, ClassVar, Dict, List, Literal, Optional, Tuple from loguru import logger from pydantic import BaseModel, Field, field_validator -from open_notebook.database.repository import ( - repo_query, -) +from open_notebook.database.repository import repo_query from open_notebook.domain.base import ObjectModel from open_notebook.domain.models import model_manager -from open_notebook.exceptions import ( - DatabaseOperationError, - InvalidInputError, -) +from open_notebook.exceptions import DatabaseOperationError, InvalidInputError from open_notebook.utils import split_text, surreal_clean @@ -212,7 +207,7 @@ class Source(ObjectModel): idx, chunk = args logger.debug(f"Processing chunk {idx}/{chunk_count}") try: - embedding = EMBEDDING_MODEL.embed(chunk) + embedding = EMBEDDING_MODEL.embed([chunk])[0] cleaned_content = surreal_clean(chunk) logger.debug(f"Successfully processed chunk {idx}") return (idx, embedding, cleaned_content) @@ -259,7 +254,7 @@ class Source(ObjectModel): if not insight_type or not content: raise InvalidInputError("Insight type and content must be provided") try: - embedding = EMBEDDING_MODEL.embed(content) if EMBEDDING_MODEL else [] + embedding = EMBEDDING_MODEL.embed([content])[0] if EMBEDDING_MODEL else [] return repo_query( f""" CREATE source_insight CONTENT {{ @@ -351,7 +346,7 @@ def vector_search( raise InvalidInputError("Search keyword cannot be empty") try: EMBEDDING_MODEL = model_manager.embedding_model - embed = EMBEDDING_MODEL.embed(keyword) + embed = EMBEDDING_MODEL.embed([keyword])[0] results = repo_query( """ SELECT * FROM fn::vector_search($embed, $results, $source, $note, $minimum_score); diff --git a/open_notebook/graphs/utils.py b/open_notebook/graphs/utils.py index ced6ed0..e6a256b 100644 --- a/open_notebook/graphs/utils.py +++ b/open_notebook/graphs/utils.py @@ -1,8 +1,8 @@ +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.models.llms import LanguageModel from open_notebook.utils import token_count @@ -27,5 +27,6 @@ def provision_langchain_model( else: model = model_manager.get_default_model(default_type, **kwargs) + logger.debug(f"Using model: {model}") assert isinstance(model, LanguageModel), f"Model is not a LanguageModel: {model}" return model.to_langchain()