From 212d3a33b070e2c2da6b509216762bf0a11919e7 Mon Sep 17 00:00:00 2001 From: LUIS NOVO Date: Fri, 1 Nov 2024 22:29:59 -0300 Subject: [PATCH] improve object typing --- open_notebook/domain/base.py | 13 +- open_notebook/domain/models.py | 125 +++++++++++++++++- open_notebook/domain/notebook.py | 8 +- .../graphs/content_processing/audio.py | 3 +- open_notebook/graphs/tools.py | 15 +-- open_notebook/graphs/utils.py | 2 +- open_notebook/models/__init__.py | 125 +++--------------- 7 files changed, 151 insertions(+), 140 deletions(-) diff --git a/open_notebook/domain/base.py b/open_notebook/domain/base.py index 84a57c9..c33a292 100644 --- a/open_notebook/domain/base.py +++ b/open_notebook/domain/base.py @@ -55,7 +55,8 @@ class ObjectModel(BaseModel): result = repo_query(f"SELECT * FROM {id}") if result: return cls(**result[0]) - return None + else: + raise NotFoundError(f"{cls.table_name} with id {id} not found") except Exception as e: logger.error(f"Error fetching {cls.table_name} with id {id}: {str(e)}") logger.exception(e) @@ -68,12 +69,12 @@ class ObjectModel(BaseModel): return None def save(self) -> None: - from open_notebook.models import model_manager + from open_notebook.domain.models import model_manager + from open_notebook.models import EmbeddingModel - EMBEDDING_MODEL = model_manager.get_default_model("embedding") + EMBEDDING_MODEL: EmbeddingModel = model_manager.embedding_model try: - logger.debug(f"Validating {self.__class__.__name__}") self.model_validate(self.model_dump(), strict=True) data = self._prepare_save_data() data["updated"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S") @@ -90,7 +91,7 @@ class ObjectModel(BaseModel): else: data["created"] = ( self.created.strftime("%Y-%m-%d %H:%M:%S") - if type(self.created) == datetime + if isinstance(self.created, datetime) else self.created ) logger.debug(f"Updating record with id {self.id}") @@ -118,8 +119,6 @@ class ObjectModel(BaseModel): def _prepare_save_data(self) -> Dict[str, Any]: data = self.model_dump() - # del data["created"] - # del data["updated"] return {key: value for key, value in data.items() if value is not None} def delete(self) -> bool: diff --git a/open_notebook/domain/models.py b/open_notebook/domain/models.py index 94549b1..cf6c8f9 100644 --- a/open_notebook/domain/models.py +++ b/open_notebook/domain/models.py @@ -1,7 +1,15 @@ -from typing import ClassVar, Optional +from typing import ClassVar, Dict, Optional from open_notebook.database.repository import repo_query from open_notebook.domain.base import ObjectModel, RecordModel +from open_notebook.models import ( + MODEL_CLASS_MAP, + EmbeddingModel, + LanguageModel, + ModelType, + SpeechToTextModel, + TextToSpeechModel, +) class Model(ObjectModel): @@ -18,7 +26,6 @@ class Model(ObjectModel): return [Model(**model) for model in models] -# todo: future: colocar um cache aqui class DefaultModels(RecordModel): record_id: ClassVar[str] = "open_notebook:default_models" @@ -29,3 +36,117 @@ class DefaultModels(RecordModel): default_speech_to_text_model: Optional[str] = None # default_vision_model: Optional[str] = None default_embedding_model: Optional[str] = None + + +class ModelManager: + _instance = None + + def __new__(cls): + if cls._instance is None: + cls._instance = super(ModelManager, cls).__new__(cls) + return cls._instance + + def __init__(self): + if not hasattr(self, "_initialized"): + self._initialized = True + self._model_cache: Dict[str, ModelType] = {} + self._default_models = None + self.refresh_defaults() + + def get_model(self, model_id: str, **kwargs) -> ModelType: + cache_key = f"{model_id}:{str(kwargs)}" + + if cache_key in self._model_cache: + cached_model = self._model_cache[cache_key] + if not isinstance( + cached_model, + (LanguageModel, EmbeddingModel, SpeechToTextModel, TextToSpeechModel), + ): + raise TypeError( + f"Cached model is of unexpected type: {type(cached_model)}" + ) + return cached_model + + assert model_id, "Model ID cannot be empty" + model: Model = Model.get(model_id) + + 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: + raise ValueError(f"Invalid model type: {model.type}") + + provider_map = MODEL_CLASS_MAP[model.type] + if model.provider not in provider_map: + 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 + + self._model_cache[cache_key] = model_instance + return model_instance + + def refresh_defaults(self): + """Refresh the default models from the database""" + self._default_models = DefaultModels.load() + + @property + def defaults(self) -> DefaultModels: + """Get the default models configuration""" + if not self._default_models: + self.refresh_defaults() + if not self._default_models: + raise RuntimeError("Failed to initialize default models configuration") + return self._default_models + + @property + def embedding_model(self, **kwargs) -> EmbeddingModel: + """Get the default embedding model""" + model = self.get_default_model("embedding", **kwargs) + if not isinstance(model, EmbeddingModel): + raise TypeError(f"Expected EmbeddingModel but got {type(model)}") + return model + + def get_default_model(self, model_type: str, **kwargs) -> ModelType: + """ + Get the default model for a specific type. + + Args: + model_type: The type of model to retrieve (e.g., 'chat', 'embedding', etc.) + **kwargs: Additional arguments to pass to the model constructor + """ + model_id = None + + if model_type == "chat": + model_id = self.defaults.default_chat_model + elif model_type == "transformation": + model_id = ( + self.defaults.default_transformation_model + or self.defaults.default_chat_model + ) + elif model_type == "embedding": + model_id = self.defaults.default_embedding_model + elif model_type == "text_to_speech": + model_id = self.defaults.default_text_to_speech_model + elif model_type == "speech_to_text": + model_id = self.defaults.default_speech_to_text_model + elif model_type == "large_context": + model_id = self.defaults.large_context_model + + if not model_id: + raise ValueError(f"No default model configured for type: {model_type}") + + return self.get_model(model_id, **kwargs) + + def clear_cache(self): + """Clear the model cache""" + self._model_cache.clear() + + +model_manager = ModelManager() diff --git a/open_notebook/domain/notebook.py b/open_notebook/domain/notebook.py index 207f5ca..1eb2edd 100644 --- a/open_notebook/domain/notebook.py +++ b/open_notebook/domain/notebook.py @@ -9,13 +9,11 @@ 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 temp.recursive_toc import graph as toc_graph -from open_notebook.models import model_manager from open_notebook.utils import split_text, surreal_clean @@ -139,7 +137,7 @@ class Source(ObjectModel): raise DatabaseOperationError(e) def vectorize(self) -> None: - EMBEDDING_MODEL = model_manager.get_default_model("embedding") + EMBEDDING_MODEL = model_manager.embedding_model try: if not self.full_text: @@ -190,7 +188,7 @@ class Source(ObjectModel): raise DatabaseOperationError("Failed to search sources") def add_insight(self, insight_type: str, content: str) -> Any: - EMBEDDING_MODEL = model_manager.get_default_model("embedding") + EMBEDDING_MODEL = model_manager.embedding_model if not insight_type or not content: raise InvalidInputError("Insight type and content must be provided") diff --git a/open_notebook/graphs/content_processing/audio.py b/open_notebook/graphs/content_processing/audio.py index ac81481..be8c441 100644 --- a/open_notebook/graphs/content_processing/audio.py +++ b/open_notebook/graphs/content_processing/audio.py @@ -4,9 +4,10 @@ from math import ceil from loguru import logger from pydub import AudioSegment +from open_notebook.domain.models import model_manager from open_notebook.graphs.content_processing.state import SourceState -from open_notebook.models import model_manager +# todo: remove reference to model_manager # future: parallelize the transcription process diff --git a/open_notebook/graphs/tools.py b/open_notebook/graphs/tools.py index 96aeacc..9c3df13 100644 --- a/open_notebook/graphs/tools.py +++ b/open_notebook/graphs/tools.py @@ -3,6 +3,7 @@ from datetime import datetime from langchain.tools import tool +# todo: turn this into a system prompt variable @tool def get_current_timestamp() -> str: """ @@ -10,17 +11,3 @@ def get_current_timestamp() -> str: Returns the current timestamp in the format YYYYMMDDHHmmss. """ return datetime.now().strftime("%Y%m%d%H%M%S") - - -# @tool -# def doc_query(doc_id: str, question: str): -# """ -# name: doc_query -# Use this tool if you need to investigate into a particular document. -# Another LLM will read the document and answer the question that you might have. -# Use this when the user question cannot be answered with the content you have in context. -# """ -# from temp.doc_query import graph - -# result = graph.invoke({"doc_id": doc_id, "question": question}) -# return result["answer"] diff --git a/open_notebook/graphs/utils.py b/open_notebook/graphs/utils.py index 5d8339c..d67d5a6 100644 --- a/open_notebook/graphs/utils.py +++ b/open_notebook/graphs/utils.py @@ -2,7 +2,7 @@ from langchain.output_parsers import OutputFixingParser from langchain_core.messages import AIMessage from loguru import logger -from open_notebook.models import model_manager +from open_notebook.domain.models import model_manager from open_notebook.prompter import Prompter from open_notebook.utils import token_count diff --git a/open_notebook/models/__init__.py b/open_notebook/models/__init__.py index 679f262..7e95d9c 100644 --- a/open_notebook/models/__init__.py +++ b/open_notebook/models/__init__.py @@ -1,6 +1,5 @@ -from typing import Dict, Optional, Union +from typing import Dict, Type, Union -from open_notebook.domain.models import DefaultModels, Model from open_notebook.models.embedding_models import ( EmbeddingModel, GeminiEmbeddingModel, @@ -29,8 +28,12 @@ from open_notebook.models.text_to_speech_models import ( TextToSpeechModel, ) -# Unified model class map with type information -MODEL_CLASS_MAP = { +ModelType = Union[LanguageModel, EmbeddingModel, SpeechToTextModel, TextToSpeechModel] + + +ProviderMap = Dict[str, Type[ModelType]] + +MODEL_CLASS_MAP: Dict[str, ProviderMap] = { "language": { "ollama": OllamaLanguageModel, "openrouter": OpenRouterLanguageModel, @@ -56,109 +59,11 @@ MODEL_CLASS_MAP = { }, } - -class ModelManager: - _instance = None - _model_cache: Dict[str, object] = {} - _default_models: Optional[DefaultModels] = None - - def __new__(cls): - if cls._instance is None: - cls._instance = super(ModelManager, cls).__new__(cls) - return cls._instance - - def __init__(self): - if not hasattr(self, "_initialized"): - self._initialized = True - self.refresh_defaults() - - def refresh_defaults(self): - """Refresh the default models from the database""" - self._default_models = DefaultModels.load() - - @property - def defaults(self) -> DefaultModels: - """Get the default models configuration""" - if not self._default_models: - self.refresh_defaults() - return self._default_models - - def get_model( - self, model_id: str, **kwargs - ) -> Union[LanguageModel, EmbeddingModel, SpeechToTextModel, TextToSpeechModel]: - """ - Get a model instance based on model_id. Uses caching to avoid recreating instances. - - Args: - model_id: The ID of the model to retrieve - **kwargs: Additional arguments to pass to the model constructor - """ - cache_key = f"{model_id}:{str(kwargs)}" - - if cache_key in self._model_cache: - return self._model_cache[cache_key] - - assert model_id, "Model ID cannot be empty" - model: Model = Model.get(model_id) - - 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: - raise ValueError(f"Invalid model type: {model.type}") - - provider_map = MODEL_CLASS_MAP[model.type] - if model.provider not in provider_map: - 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 - - self._model_cache[cache_key] = model_instance - return model_instance - - def get_default_model( - self, model_type: str, **kwargs - ) -> Union[LanguageModel, EmbeddingModel, SpeechToTextModel, TextToSpeechModel]: - """ - Get the default model for a specific type. - - Args: - model_type: The type of model to retrieve (e.g., 'chat', 'embedding', etc.) - **kwargs: Additional arguments to pass to the model constructor - """ - model_id = None - - if model_type == "chat": - model_id = self.defaults.default_chat_model - elif model_type == "transformation": - model_id = ( - self.defaults.default_transformation_model - or self.defaults.default_chat_model - ) - elif model_type == "embedding": - model_id = self.defaults.default_embedding_model - elif model_type == "text_to_speech": - model_id = self.defaults.default_text_to_speech_model - elif model_type == "speech_to_text": - model_id = self.defaults.default_speech_to_text_model - elif model_type == "large_context": - model_id = self.defaults.large_context_model - - if not model_id: - raise ValueError(f"No default model configured for type: {model_type}") - - return self.get_model(model_id, **kwargs) - - def clear_cache(self): - """Clear the model cache""" - self._model_cache.clear() - - -model_manager = ModelManager() +__all__ = [ + "MODEL_CLASS_MAP", + "EmbeddingModel", + "LanguageModel", + "SpeechToTextModel", + "TextToSpeechModel", + "ModelType", +]