diff --git a/open_notebook/domain/base.py b/open_notebook/domain/base.py index da9b0ac..4fca352 100644 --- a/open_notebook/domain/base.py +++ b/open_notebook/domain/base.py @@ -48,7 +48,7 @@ class ObjectModel(BaseModel): raise DatabaseOperationError(e) @classmethod - def get(cls: Type[T], id: str) -> Optional[T]: + def get(cls: Type[T], id: str) -> T: if not id: raise InvalidInputError("ID cannot be empty") try: diff --git a/open_notebook/models/__init__.py b/open_notebook/models/__init__.py index 84907bc..dc2cf36 100644 --- a/open_notebook/models/__init__.py +++ b/open_notebook/models/__init__.py @@ -43,7 +43,7 @@ def get_model(model_id, model_type="language", **kwargs): **kwargs: Additional arguments to pass to the model constructor """ assert model_id, "Model ID cannot be empty" - model = Model.get(model_id) + model: Model = Model.get(model_id) if not model: raise ValueError(f"Model with ID {model_id} not found") @@ -65,77 +65,3 @@ def get_model(model_id, model_type="language", **kwargs): return model_instance.to_langchain() return model_instance - - -# from open_notebook.domain.models import Model -# from open_notebook.models.embedding_models import OpenAIEmbeddingModel -# from open_notebook.models.llms import ( -# AnthropicLanguageModel, -# GeminiLanguageModel, -# LiteLLMLanguageModel, -# OllamaLanguageModel, -# OpenAILanguageModel, -# OpenRouterLanguageModel, -# VertexAILanguageModel, -# VertexAnthropicLanguageModel, -# ) -# from open_notebook.models.speech_to_text_models import OpenAISpeechToTextModel - -# SPEECH_TO_TEXT_CLASS_MAP = { -# "openai": OpenAISpeechToTextModel, -# } - - -# # todo: acho que dá pra juntar todos os get models em uma coisa só -# def get_speech_to_text_model(model_id): -# assert model_id, "Model ID cannot be empty" -# model = Model.get(model_id) -# if not model: -# raise ValueError(f"Model with ID {model_id} not found") -# if model.provider not in SPEECH_TO_TEXT_CLASS_MAP.keys(): -# raise ValueError( -# f"Provider {model.provider} not compatible with Embedding Models" -# ) -# return SPEECH_TO_TEXT_CLASS_MAP[model.provider](model_name=model.name) - - -# # Map provider names to classes -# PROVIDER_CLASS_MAP = { -# "ollama": OllamaLanguageModel, -# "openrouter": OpenRouterLanguageModel, -# "vertexai-anthropic": VertexAnthropicLanguageModel, -# "litellm": LiteLLMLanguageModel, -# "vertexai": VertexAILanguageModel, -# "anthropic": AnthropicLanguageModel, -# "openai": OpenAILanguageModel, -# "gemini": GeminiLanguageModel, -# } - - -# # todo: make the provider check type specific -# def get_langchain_model(model_id, json=False): -# model = Model.get(model_id) -# if not model: -# raise ValueError(f"Model {model_id} not found") -# if model.provider not in PROVIDER_CLASS_MAP.keys(): -# raise ValueError(f"Provider {model.provider} not found") -# return PROVIDER_CLASS_MAP[model.provider]( -# model_name=model.name, json=json -# ).to_langchain() - - -# EMBEDDING_CLASS_MAP = { -# "openai": OpenAIEmbeddingModel, -# } - - -# def get_embedding_model(model_id): -# assert model_id, "Model ID cannot be empty" -# model = Model.get(model_id) -# if not model: -# raise ValueError(f"Model with ID {model_id} not found") -# if model.provider not in EMBEDDING_CLASS_MAP.keys(): -# raise ValueError( -# f"Provider {model.provider} not compatible with Embedding Models" -# ) -# return EMBEDDING_CLASS_MAP[model.provider](model_name=model.name)