feat: implement the new model management based on esperanto framework

open_notebook/domain/base.py

@@ -1,22 +1,8 @@
 from datetime import datetime
-from typing import (
+from typing import Any, ClassVar, Dict, List, Optional, Type, TypeVar, cast
-    Any,
-    ClassVar,
-    Dict,
-    List,
-    Optional,
-    Type,
-    TypeVar,
-    cast,
-)
 
 from loguru import logger
-from pydantic import (
+from pydantic import BaseModel, ValidationError, field_validator, model_validator
-    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 []
                     )

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 esperanto import (
-from open_notebook.domain.base import ObjectModel, RecordModel
+    AIFactory,
-from open_notebook.models import (
-    MODEL_CLASS_MAP,
     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]
+        # todo: change to providers in the future
-        if model.provider not in provider_map:
+        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

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 (
+from open_notebook.database.repository import repo_query
-    repo_query,
-)
 from open_notebook.domain.base import ObjectModel
 from open_notebook.domain.models import model_manager
-from open_notebook.exceptions import (
+from open_notebook.exceptions import DatabaseOperationError, InvalidInputError
-    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);

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