improve object typing

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:

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()

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")

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
 
 

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"]

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
 

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
+ModelType = Union[LanguageModel, EmbeddingModel, SpeechToTextModel, TextToSpeechModel]
-MODEL_CLASS_MAP = {
+
+
+ProviderMap = Dict[str, Type[ModelType]]
+
+MODEL_CLASS_MAP: Dict[str, ProviderMap] = {
     "language": {
         "ollama": OllamaLanguageModel,
         "openrouter": OpenRouterLanguageModel,
@@ -56,109 +59,11 @@ MODEL_CLASS_MAP = {
     },
 }
 
-
+__all__ = [
-class ModelManager:
+    "MODEL_CLASS_MAP",
-    _instance = None
+    "EmbeddingModel",
-    _model_cache: Dict[str, object] = {}
+    "LanguageModel",
-    _default_models: Optional[DefaultModels] = None
+    "SpeechToTextModel",
-
+    "TextToSpeechModel",
-    def __new__(cls):
+    "ModelType",
-        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()