admin-valentin/open-notebook
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

improve object typing

Browse source
This commit is contained in:
LUIS NOVO 2024-11-01 22:29:59 -03:00
parent b616d1ad17
commit 212d3a33b0
7 changed files with 151 additions and 140 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
open_notebook/domain/base.py
View file

@ -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:

125
open_notebook/domain/models.py
View file

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

8
open_notebook/domain/notebook.py
View file

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

3
open_notebook/graphs/content_processing/audio.py
View file

@ -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

15
open_notebook/graphs/tools.py
View file

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

2
open_notebook/graphs/utils.py
View file

@ -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

125
open_notebook/models/__init__.py
View file

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