mirror of
https://github.com/lfnovo/open-notebook.git
synced 2026-04-28 19:40:50 +00:00
46 lines
1.3 KiB
Python
46 lines
1.3 KiB
Python
from typing import ClassVar, Optional
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from open_notebook.database.repository import (
|
|
repo_query,
|
|
repo_update,
|
|
)
|
|
from open_notebook.domain.base import ObjectModel
|
|
|
|
|
|
class Model(ObjectModel):
|
|
table_name: ClassVar[str] = "model"
|
|
name: str
|
|
provider: str
|
|
type: str
|
|
|
|
@classmethod
|
|
def get_models_by_type(cls, model_type):
|
|
models = repo_query(
|
|
"SELECT * FROM model WHERE type=$model_type;", {"model_type": model_type}
|
|
)
|
|
return [Model(**model) for model in models]
|
|
|
|
|
|
class DefaultModels(BaseModel):
|
|
default_chat_model: Optional[str] = None
|
|
default_transformation_model: Optional[str] = None
|
|
large_context_model: Optional[str] = None
|
|
default_text_to_speech_model: Optional[str] = None
|
|
default_speech_to_text_model: Optional[str] = None
|
|
# default_vision_model: Optional[str] = None
|
|
default_embedding_model: Optional[str] = None
|
|
|
|
@classmethod
|
|
def load(self):
|
|
result = repo_query("SELECT * FROM open_notebook:default_models;")
|
|
if result:
|
|
result = result[0]
|
|
dm = DefaultModels(**result)
|
|
return dm
|
|
return DefaultModels()
|
|
|
|
@classmethod
|
|
def update(self, data):
|
|
repo_update("open_notebook:default_models", data)
|