mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-01 05:23:47 +08:00
Feat: tenant llm provider (#14595)
### What problem does this PR solve? Python implementation of the Go-based model_provider API suite. ### Type of change - [x] New Feature (non-breaking change which adds functionality) --------- Co-authored-by: bill <yibie_jingnian@163.com>
This commit is contained in:
Lynn
@@ -763,6 +763,7 @@ class Tenant(DataBaseModel):
|
||||
tenant_rerank_id = IntegerField(null=True, help_text="id in tenant_llm", index=True)
|
||||
tts_id = CharField(max_length=256, null=True, help_text="default tts model ID", index=True)
|
||||
tenant_tts_id = IntegerField(null=True, help_text="id in tenant_llm", index=True)
|
||||
ocr_id = CharField(max_length=256, null=True, help_text="default OCR model ID", index=True)
|
||||
parser_ids = CharField(max_length=256, null=False, help_text="document processors", index=True)
|
||||
credit = IntegerField(default=512, index=True)
|
||||
status = CharField(max_length=1, null=True, help_text="is it validate(0: wasted, 1: validate)", default="1", index=True)
|
||||
@@ -1350,6 +1351,67 @@ class SystemSettings(DataBaseModel):
|
||||
class Meta:
|
||||
db_table = "system_settings"
|
||||
|
||||
class TenantModelProvider(DataBaseModel):
|
||||
id = CharField(max_length=32, primary_key=True)
|
||||
provider_name = CharField(max_length=128, null=False, index=False, help_text="LLM provider name")
|
||||
tenant_id = CharField(max_length=32, null=False, index=True)
|
||||
|
||||
class Meta:
|
||||
db_table = "tenant_model_provider"
|
||||
indexes = (
|
||||
(("tenant_id", "provider_name"), True),
|
||||
)
|
||||
|
||||
class TenantModelInstance(DataBaseModel):
|
||||
id = CharField(max_length=32, primary_key=True)
|
||||
instance_name = CharField(max_length=128, null=False, index=False, help_text="Model instance name")
|
||||
provider_id = CharField(max_length=32, null=False, index=False)
|
||||
api_key = CharField(max_length=512, null=False, index=False, help_text="API key")
|
||||
status = CharField(max_length=32, default="active", index=False)
|
||||
extra = CharField(max_length=512, default="{}", index=False)
|
||||
|
||||
class Meta:
|
||||
db_table = "tenant_model_instance"
|
||||
indexes = (
|
||||
(("api_key", "provider_id"), True),
|
||||
)
|
||||
|
||||
|
||||
class TenantModel(DataBaseModel):
|
||||
id = CharField(max_length=32, primary_key=True)
|
||||
model_name = CharField(max_length=128, null=True, index=False, help_text="Model name")
|
||||
provider_id = CharField(max_length=32, null=False, index=False)
|
||||
instance_id = CharField(max_length=32, null=False, index=True)
|
||||
model_type = CharField(max_length=32, null=False, index=False, help_text="Model type")
|
||||
status = CharField(max_length=32, default="active", index=False)
|
||||
extra = CharField(max_length=1024, default="{}", index=False)
|
||||
|
||||
class Meta:
|
||||
db_table = "tenant_model"
|
||||
|
||||
|
||||
class TenantModelGroup(DataBaseModel):
|
||||
id = CharField(max_length=32, primary_key=True)
|
||||
group_type = CharField(max_length=32, null=False, index=False, help_text="Group type")
|
||||
model_name = CharField(max_length=128, null=True, index=False, help_text="Model name")
|
||||
strategy = CharField(max_length=32, default="weighted", index=False, help_text="Routing strategy")
|
||||
|
||||
class Meta:
|
||||
db_table = "tenant_model_group"
|
||||
|
||||
class TenantModelGroupMapping(DataBaseModel):
|
||||
group_id = CharField(max_length=32, null=False, index=True, help_text="Group ID")
|
||||
provider_id = CharField(max_length=32, null=False, index=False)
|
||||
instance_id = CharField(max_length=32, null=False, index=False)
|
||||
model_id = CharField(max_length=32, null=False, index=True)
|
||||
weight = IntegerField(default=100, index=False, help_text="Routing weight")
|
||||
status = CharField(max_length=32, default="active", index=False)
|
||||
|
||||
class Meta:
|
||||
db_table = "tenant_model_group_mapping"
|
||||
primary_key = CompositeKey("group_id", "provider_id", "instance_id", "model_id")
|
||||
|
||||
|
||||
def alter_db_add_column(migrator, table_name, column_name, column_type):
|
||||
try:
|
||||
migrate(migrator.add_column(table_name, column_name, column_type))
|
||||
@@ -1668,6 +1730,7 @@ def migrate_db():
|
||||
alter_db_add_column(migrator, "api_4_conversation", "version_title", CharField(max_length=255, null=True, help_text="canvas version title when session created", index=False))
|
||||
alter_db_column_type(migrator, "document", "size", BigIntegerField(default=0, index=True))
|
||||
alter_db_column_type(migrator, "file", "size", BigIntegerField(default=0, index=True))
|
||||
alter_db_add_column(migrator, "tenant", "ocr_id", CharField(max_length=128, null=True, help_text="default ocr model ID", index=True))
|
||||
logging.disable(logging.NOTSET)
|
||||
# this is after re-enabling logging to allow logging changed user emails
|
||||
migrate_add_unique_email(migrator)
|
||||
|
||||
@@ -24,17 +24,15 @@ from copy import deepcopy
|
||||
from peewee import IntegrityError
|
||||
|
||||
from api.db import UserTenantRole
|
||||
from api.db.db_models import init_database_tables as init_web_db, LLMFactories, LLM, TenantLLM, Knowledgebase, Dialog, Memory
|
||||
from api.db.db_models import init_database_tables as init_web_db, LLMFactories, LLM, TenantLLM
|
||||
from api.db.services import UserService
|
||||
from api.db.services.canvas_service import CanvasTemplateService
|
||||
from api.db.services.document_service import DocumentService
|
||||
from api.db.services.knowledgebase_service import KnowledgebaseService
|
||||
from api.db.services.memory_service import MemoryService
|
||||
from api.db.services.tenant_llm_service import LLMFactoriesService, TenantLLMService
|
||||
from api.db.services.llm_service import LLMService, LLMBundle, get_init_tenant_llm
|
||||
from api.db.services.user_service import TenantService, UserTenantService
|
||||
from api.db.services.system_settings_service import SystemSettingsService
|
||||
from api.db.services.dialog_service import DialogService
|
||||
from api.db.template_utils import normalize_canvas_template_categories
|
||||
from api.db.joint_services.memory_message_service import init_message_id_sequence, init_memory_size_cache, fix_missing_tokenized_memory
|
||||
from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type
|
||||
@@ -154,6 +152,8 @@ def init_llm_factory():
|
||||
except Exception:
|
||||
pass
|
||||
break
|
||||
|
||||
def update_document_number_in_init():
|
||||
doc_count = DocumentService.get_all_kb_doc_count()
|
||||
for kb_id in KnowledgebaseService.get_all_ids():
|
||||
KnowledgebaseService.update_document_number_in_init(kb_id=kb_id, doc_num=doc_count.get(kb_id, 0))
|
||||
@@ -189,7 +189,8 @@ def init_web_data():
|
||||
|
||||
init_table()
|
||||
|
||||
init_llm_factory()
|
||||
# init_llm_factory()
|
||||
update_document_number_in_init()
|
||||
# if not UserService.get_all().count():
|
||||
# init_superuser()
|
||||
|
||||
@@ -197,7 +198,6 @@ def init_web_data():
|
||||
init_message_id_sequence()
|
||||
init_memory_size_cache()
|
||||
fix_missing_tokenized_memory()
|
||||
fix_empty_tenant_model_id()
|
||||
logging.info("init web data success:{}".format(time.time() - start_time))
|
||||
|
||||
def init_table():
|
||||
@@ -226,105 +226,6 @@ def init_table():
|
||||
raise e
|
||||
|
||||
|
||||
def fix_empty_tenant_model_id():
|
||||
# knowledgebase
|
||||
empty_tenant_embd_id_kbs = KnowledgebaseService.get_null_tenant_embd_id_row()
|
||||
if empty_tenant_embd_id_kbs:
|
||||
logging.info(f"Found {len(empty_tenant_embd_id_kbs)} empty tenant_embd_id knowledgebase.")
|
||||
kb_groups: dict = {}
|
||||
for obj in empty_tenant_embd_id_kbs:
|
||||
if kb_groups.get((obj.tenant_id, obj.embd_id)):
|
||||
kb_groups[(obj.tenant_id, obj.embd_id)].append(obj.id)
|
||||
else:
|
||||
kb_groups[(obj.tenant_id, obj.embd_id)] = [obj.id]
|
||||
update_cnt = 0
|
||||
for k, v in kb_groups.items():
|
||||
tenant_llm = TenantLLMService.get_api_key(k[0], k[1])
|
||||
if tenant_llm:
|
||||
update_cnt += KnowledgebaseService.filter_update([Knowledgebase.id.in_(v)], {"tenant_embd_id": tenant_llm.id})
|
||||
logging.info(f"Update {update_cnt} tenant_embd_id in table knowledgebase.")
|
||||
# dialog
|
||||
empty_tenant_llm_id_dialog = DialogService.get_null_tenant_llm_id_row()
|
||||
if empty_tenant_llm_id_dialog:
|
||||
logging.info(f"Found {len(empty_tenant_llm_id_dialog)} empty tenant_llm_id dialogs.")
|
||||
dialog_groups: dict = {}
|
||||
for obj in empty_tenant_llm_id_dialog:
|
||||
if dialog_groups.get((obj.tenant_id, obj.llm_id)):
|
||||
dialog_groups[(obj.tenant_id, obj.llm_id)].append(obj.id)
|
||||
else:
|
||||
dialog_groups[(obj.tenant_id, obj.llm_id)] = [obj.id]
|
||||
update_cnt = 0
|
||||
for k, v in dialog_groups.items():
|
||||
tenant_llm = TenantLLMService.get_api_key(k[0], k[1])
|
||||
if tenant_llm:
|
||||
update_cnt += DialogService.filter_update([Dialog.id.in_(v)], {"tenant_llm_id": tenant_llm.id})
|
||||
logging.info(f"Update {update_cnt} tenant_llm_id in table dialog.")
|
||||
|
||||
empty_tenant_rerank_id_dialog = DialogService.get_null_tenant_rerank_id_row()
|
||||
if empty_tenant_rerank_id_dialog:
|
||||
logging.info(f"Found {len(empty_tenant_rerank_id_dialog)} empty tenant_rerank_id dialogs.")
|
||||
dialog_groups: dict = {}
|
||||
for obj in empty_tenant_rerank_id_dialog:
|
||||
if dialog_groups.get((obj.tenant_id, obj.rerank_id)):
|
||||
dialog_groups[(obj.tenant_id, obj.rerank_id)].append(obj.id)
|
||||
else:
|
||||
dialog_groups[(obj.tenant_id, obj.rerank_id)] = [obj.id]
|
||||
update_cnt = 0
|
||||
for k, v in dialog_groups.items():
|
||||
tenant_llm = TenantLLMService.get_api_key(k[0], k[1])
|
||||
if tenant_llm:
|
||||
update_cnt += DialogService.filter_update([Dialog.id.in_(v)], {"tenant_rerank_id": tenant_llm.id})
|
||||
logging.info(f"Update {update_cnt} tenant_rerank_id in table dialog.")
|
||||
# memory
|
||||
empty_tenant_embd_id_memories = MemoryService.get_null_tenant_embd_id_row()
|
||||
if empty_tenant_embd_id_memories:
|
||||
logging.info(f"Found {len(empty_tenant_embd_id_memories)} empty tenant_embd_id memories.")
|
||||
memory_groups: dict = {}
|
||||
for obj in empty_tenant_embd_id_memories:
|
||||
if memory_groups.get((obj.tenant_id, obj.embd_id)):
|
||||
memory_groups[(obj.tenant_id, obj.embd_id)].append(obj.id)
|
||||
else:
|
||||
memory_groups[(obj.tenant_id, obj.embd_id)] = [obj.id]
|
||||
update_cnt = 0
|
||||
for k, v in memory_groups.items():
|
||||
tenant_llm = TenantLLMService.get_api_key(k[0], k[1])
|
||||
if tenant_llm:
|
||||
update_cnt += MemoryService.filter_update([Memory.id.in_(v)], {"tenant_embd_id": tenant_llm.id})
|
||||
logging.info(f"Update {update_cnt} tenant_embd_id in table memory.")
|
||||
|
||||
empty_tenant_llm_id_memories = MemoryService.get_null_tenant_llm_id_row()
|
||||
if empty_tenant_llm_id_memories:
|
||||
logging.info(f"Found {len(empty_tenant_llm_id_memories)} empty tenant_llm_id memories.")
|
||||
memory_groups: dict = {}
|
||||
for obj in empty_tenant_llm_id_memories:
|
||||
if memory_groups.get((obj.tenant_id, obj.llm_id)):
|
||||
memory_groups[(obj.tenant_id, obj.llm_id)].append(obj.id)
|
||||
else:
|
||||
memory_groups[(obj.tenant_id, obj.llm_id)] = [obj.id]
|
||||
update_cnt = 0
|
||||
for k, v in memory_groups.items():
|
||||
tenant_llm = TenantLLMService.get_api_key(k[0], k[1])
|
||||
if tenant_llm:
|
||||
update_cnt += MemoryService.filter_update([Memory.id.in_(v)], {"tenant_llm_id": tenant_llm.id})
|
||||
logging.info(f"Update {update_cnt} tenant_llm_id in table memory.")
|
||||
# tenant
|
||||
empty_tenant_model_id_tenants = TenantService.get_null_tenant_model_id_rows()
|
||||
if empty_tenant_model_id_tenants:
|
||||
logging.info(f"Found {len(empty_tenant_model_id_tenants)} empty tenant_model_id tenants.")
|
||||
update_cnt = 0
|
||||
for obj in empty_tenant_model_id_tenants:
|
||||
tenant_dict = obj.to_dict()
|
||||
update_dict = {}
|
||||
for key in ["llm_id", "embd_id", "asr_id", "img2txt_id", "rerank_id", "tts_id"]:
|
||||
if tenant_dict.get(key) and not tenant_dict.get(f"tenant_{key}"):
|
||||
tenant_model = TenantLLMService.get_api_key(tenant_dict["id"], tenant_dict[key])
|
||||
if tenant_model:
|
||||
update_dict.update({f"tenant_{key}": tenant_model.id})
|
||||
if update_dict:
|
||||
update_cnt += TenantService.update_by_id(tenant_dict["id"], update_dict)
|
||||
logging.info(f"Update {update_cnt} tenant_model_id in table tenant.")
|
||||
logging.info("Fix empty tenant_model_id done.")
|
||||
|
||||
if __name__ == '__main__':
|
||||
init_web_db()
|
||||
init_web_data()
|
||||
|
||||
@@ -26,7 +26,7 @@ from api.db.db_models import Task
|
||||
from api.db.services.task_service import TaskService
|
||||
from api.db.services.memory_service import MemoryService
|
||||
from api.db.services.llm_service import LLMBundle
|
||||
from api.db.joint_services.tenant_model_service import get_model_config_by_id, get_model_config_by_type_and_name
|
||||
from api.db.joint_services.tenant_model_service import get_model_config_from_provider_instance
|
||||
from api.utils.memory_utils import get_memory_type_human
|
||||
from memory.services.messages import MessageService
|
||||
from memory.services.query import MsgTextQuery, get_vector
|
||||
@@ -153,14 +153,7 @@ async def extract_by_llm(tenant_id: str, tenant_llm_id: int, extract_conf: dict,
|
||||
user_prompts.append({"role": "user", "content": f"Conversation: {conversation_content}\nConversation Time: {conversation_time}\nCurrent Time: {conversation_time}"})
|
||||
else:
|
||||
user_prompts.append({"role": "user", "content": PromptAssembler.assemble_user_prompt(conversation_content, conversation_time, conversation_time)})
|
||||
if tenant_llm_id:
|
||||
llm_config = get_model_config_by_id(
|
||||
tenant_llm_id,
|
||||
allowed_tenant_ids=tenant_id,
|
||||
requester_tenant_id=tenant_id,
|
||||
)
|
||||
else:
|
||||
llm_config = get_model_config_by_type_and_name(tenant_id, LLMType.CHAT, llm_id)
|
||||
llm_config = get_model_config_from_provider_instance(tenant_id, LLMType.CHAT, llm_id)
|
||||
with LLMBundle(tenant_id, llm_config) as llm:
|
||||
if task_id:
|
||||
TaskService.update_progress(task_id, {"progress": 0.15, "progress_msg": timestamp_to_date(current_timestamp())+ " " + "Prepared prompts and LLM."})
|
||||
@@ -177,14 +170,7 @@ async def extract_by_llm(tenant_id: str, tenant_llm_id: int, extract_conf: dict,
|
||||
|
||||
|
||||
async def embed_and_save(memory, message_list: list[dict], task_id: str=None):
|
||||
if memory.tenant_embd_id:
|
||||
embd_model_config = get_model_config_by_id(
|
||||
memory.tenant_embd_id,
|
||||
allowed_tenant_ids=memory.tenant_id,
|
||||
requester_tenant_id=memory.tenant_id,
|
||||
)
|
||||
else:
|
||||
embd_model_config = get_model_config_by_type_and_name(memory.tenant_id, LLMType.EMBEDDING, memory.embd_id)
|
||||
embd_model_config = get_model_config_from_provider_instance(memory.tenant_id, LLMType.EMBEDDING, memory.embd_id)
|
||||
with LLMBundle(memory.tenant_id, embd_model_config) as embedding_model:
|
||||
if task_id:
|
||||
TaskService.update_progress(task_id, {"progress": 0.65, "progress_msg": timestamp_to_date(current_timestamp())+ " " + "Prepared embedding model."})
|
||||
@@ -255,14 +241,7 @@ def query_message(filter_dict: dict, params: dict):
|
||||
question = params["query"]
|
||||
question = question.strip()
|
||||
memory = memory_list[0]
|
||||
if memory.tenant_embd_id:
|
||||
embd_model_config = get_model_config_by_id(
|
||||
memory.tenant_embd_id,
|
||||
allowed_tenant_ids=memory.tenant_id,
|
||||
requester_tenant_id=memory.tenant_id,
|
||||
)
|
||||
else:
|
||||
embd_model_config = get_model_config_by_type_and_name(memory.tenant_id, LLMType.EMBEDDING, memory.embd_id)
|
||||
embd_model_config = get_model_config_from_provider_instance(memory.tenant_id, LLMType.EMBEDDING, memory.embd_id)
|
||||
embd_model = LLMBundle(memory.tenant_id, embd_model_config)
|
||||
match_dense = get_vector(question, embd_model, similarity=params["similarity_threshold"])
|
||||
match_text, _ = MsgTextQuery().question(question, min_match=params["similarity_threshold"])
|
||||
|
||||
@@ -17,131 +17,15 @@ import logging
|
||||
import os
|
||||
import enum
|
||||
from common import settings
|
||||
from common.constants import LLMType
|
||||
from api.db.services.llm_service import LLMService
|
||||
from common.constants import LLMType, ActiveStatusEnum
|
||||
from api.db.services.tenant_llm_service import TenantLLMService, TenantService
|
||||
from api.db.services.tenant_model_provider_service import TenantModelProviderService
|
||||
from api.db.services.tenant_model_instance_service import TenantModelInstanceService
|
||||
from api.db.services.tenant_model_service import TenantModelService
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_model_config_by_id(
|
||||
tenant_model_id: int,
|
||||
allowed_tenant_ids: str | list[str] | set[str] | tuple[str, ...] | None = None,
|
||||
requester_tenant_id: str | None = None,
|
||||
) -> dict:
|
||||
found, model_config = TenantLLMService.get_by_id(tenant_model_id)
|
||||
if not found:
|
||||
raise LookupError(f"Tenant Model with id {tenant_model_id} not found")
|
||||
if allowed_tenant_ids is not None:
|
||||
if isinstance(allowed_tenant_ids, str):
|
||||
allowed_tenant_ids = {allowed_tenant_ids}
|
||||
else:
|
||||
allowed_tenant_ids = {str(tenant_id) for tenant_id in allowed_tenant_ids if tenant_id}
|
||||
if str(model_config.tenant_id) not in allowed_tenant_ids:
|
||||
logger.warning(
|
||||
"Denied tenant model access: tenant_model_id=%s model_tenant_id=%s "
|
||||
"allowed_tenant_ids=%s requester_tenant_id=%s",
|
||||
tenant_model_id,
|
||||
model_config.tenant_id,
|
||||
sorted(allowed_tenant_ids),
|
||||
requester_tenant_id,
|
||||
)
|
||||
raise LookupError(f"Tenant Model with id {tenant_model_id} not authorized")
|
||||
config_dict = model_config.to_dict()
|
||||
api_key, is_tools, api_key_payload = TenantLLMService._decode_api_key_config(config_dict.get("api_key", ""))
|
||||
config_dict["api_key"] = api_key
|
||||
if api_key_payload is not None:
|
||||
config_dict["api_key_payload"] = api_key_payload
|
||||
if is_tools is not None:
|
||||
config_dict["is_tools"] = is_tools
|
||||
llm = LLMService.query(llm_name=config_dict["llm_name"])
|
||||
if "is_tools" not in config_dict and llm:
|
||||
config_dict["is_tools"] = llm[0].is_tools
|
||||
return config_dict
|
||||
|
||||
|
||||
def get_model_config_by_type_and_name(tenant_id: str, model_type: str, model_name: str):
|
||||
if not model_name:
|
||||
raise Exception("Model Name is required")
|
||||
model_type_val = model_type.value if hasattr(model_type, "value") else model_type
|
||||
model_config = TenantLLMService.get_api_key(tenant_id, model_name, model_type_val)
|
||||
if not model_config:
|
||||
# model_name in format 'name@factory', split model_name and try again
|
||||
pure_model_name, fid = TenantLLMService.split_model_name_and_factory(model_name)
|
||||
compose_profiles = os.getenv("COMPOSE_PROFILES", "")
|
||||
is_tei_builtin_embedding = (
|
||||
model_type_val == LLMType.EMBEDDING.value
|
||||
and "tei-" in compose_profiles
|
||||
and pure_model_name == os.getenv("TEI_MODEL", "")
|
||||
and (fid == "Builtin" or fid is None)
|
||||
)
|
||||
if is_tei_builtin_embedding:
|
||||
# configured local embedding model
|
||||
embedding_cfg = settings.EMBEDDING_CFG
|
||||
config_dict = {
|
||||
"llm_factory": "Builtin",
|
||||
"api_key": embedding_cfg["api_key"],
|
||||
"llm_name": pure_model_name,
|
||||
"api_base": embedding_cfg["base_url"],
|
||||
"model_type": LLMType.EMBEDDING.value,
|
||||
}
|
||||
elif model_type_val == LLMType.CHAT.value:
|
||||
# Retry as CHAT with pure_model_name first; then fall back to a multimodal model registered under IMAGE2TEXT.
|
||||
model_config = TenantLLMService.get_api_key(tenant_id, pure_model_name, LLMType.CHAT.value)
|
||||
if not model_config:
|
||||
model_config = TenantLLMService.get_api_key(tenant_id, pure_model_name, LLMType.IMAGE2TEXT.value)
|
||||
if not model_config:
|
||||
raise LookupError(f"Tenant Model with name {model_name} and type {model_type_val} not found")
|
||||
config_dict = model_config.to_dict()
|
||||
elif model_type_val == LLMType.IMAGE2TEXT.value:
|
||||
model_config = TenantLLMService.get_api_key(tenant_id, pure_model_name, LLMType.IMAGE2TEXT.value)
|
||||
if not model_config:
|
||||
# Fall back to a chat model only if it has declared IMAGE2TEXT capability (tag check via llm table)
|
||||
chat_config = TenantLLMService.get_api_key(tenant_id, pure_model_name, LLMType.CHAT.value)
|
||||
logger.debug("IMAGE2TEXT config not found for %s; chat_config found: %s", pure_model_name, chat_config is not None)
|
||||
if chat_config:
|
||||
llm_entry = LLMService.query(fid=chat_config.llm_factory, llm_name=chat_config.llm_name)
|
||||
tags = [t.strip() for t in (llm_entry[0].tags or "").split(",")] if llm_entry else []
|
||||
logger.debug("LLM tags for %s/%s: %s", chat_config.llm_factory, chat_config.llm_name, tags)
|
||||
if "IMAGE2TEXT" in tags:
|
||||
logger.debug("Promoting chat config to IMAGE2TEXT for %s", pure_model_name)
|
||||
model_config = chat_config
|
||||
if not model_config:
|
||||
raise LookupError(f"Tenant Model with name {model_name} and type {model_type_val} not found")
|
||||
config_dict = model_config.to_dict()
|
||||
config_dict["model_type"] = LLMType.IMAGE2TEXT.value
|
||||
else:
|
||||
model_config = TenantLLMService.get_api_key(tenant_id, pure_model_name, model_type_val)
|
||||
if not model_config:
|
||||
raise LookupError(f"Tenant Model with name {model_name} and type {model_type_val} not found")
|
||||
config_dict = model_config.to_dict()
|
||||
else:
|
||||
# model_name without @factory
|
||||
config_dict = model_config.to_dict()
|
||||
api_key, is_tools, api_key_payload = TenantLLMService._decode_api_key_config(config_dict.get("api_key", ""))
|
||||
config_dict["api_key"] = api_key
|
||||
if api_key_payload is not None:
|
||||
config_dict["api_key_payload"] = api_key_payload
|
||||
if is_tools is not None:
|
||||
config_dict["is_tools"] = is_tools
|
||||
config_model_type = config_dict.get("model_type")
|
||||
config_model_type = config_model_type.value if hasattr(config_model_type, "value") else config_model_type
|
||||
if config_model_type != model_type_val and not (
|
||||
model_type_val == LLMType.CHAT.value
|
||||
and config_model_type == LLMType.IMAGE2TEXT.value
|
||||
) and not (
|
||||
model_type_val == LLMType.IMAGE2TEXT.value
|
||||
and config_model_type == LLMType.CHAT.value
|
||||
):
|
||||
raise LookupError(
|
||||
f"Tenant Model with name {model_name} has type {config_model_type}, expected {model_type_val}"
|
||||
)
|
||||
llm = LLMService.query(llm_name=config_dict["llm_name"])
|
||||
if "is_tools" not in config_dict and llm:
|
||||
config_dict["is_tools"] = llm[0].is_tools
|
||||
return config_dict
|
||||
|
||||
|
||||
def get_tenant_default_model_by_type(tenant_id: str, model_type: str|enum.Enum):
|
||||
exist, tenant = TenantService.get_by_id(tenant_id)
|
||||
if not exist:
|
||||
@@ -167,4 +51,135 @@ def get_tenant_default_model_by_type(tenant_id: str, model_type: str|enum.Enum):
|
||||
raise Exception(f"Unknown model type {model_type}")
|
||||
if not model_name:
|
||||
raise Exception(f"No default {model_type} model is set.")
|
||||
return get_model_config_by_type_and_name(tenant_id, model_type, model_name)
|
||||
return get_model_config_from_provider_instance(tenant_id, model_type, model_name)
|
||||
|
||||
|
||||
def split_model_name(model_name: str):
|
||||
# Parse model_name: {model_name} or {model_name}@{factory_name} or {model_name}@{instance_name}@{factory_name}
|
||||
parts = model_name.split("@")
|
||||
if len(parts) == 1:
|
||||
pure_model_name = parts[0]
|
||||
provider_name = ""
|
||||
instance_name = ""
|
||||
elif len(parts) == 2:
|
||||
pure_model_name = parts[0]
|
||||
provider_name = parts[1]
|
||||
instance_name = "default"
|
||||
else:
|
||||
pure_model_name = parts[0]
|
||||
instance_name = parts[1]
|
||||
provider_name = parts[2]
|
||||
return pure_model_name, instance_name, provider_name
|
||||
|
||||
|
||||
def get_model_config_from_provider_instance(tenant_id, model_type: str|enum.Enum, model_name: str):
|
||||
pure_model_name, instance_name, provider_name = split_model_name(model_name)
|
||||
model_type_val = model_type if isinstance(model_type, str) else model_type.value
|
||||
# Builtin embedding model
|
||||
compose_profiles = os.getenv("COMPOSE_PROFILES", "")
|
||||
is_tei_builtin_embedding = (
|
||||
model_type_val == LLMType.EMBEDDING.value
|
||||
and "tei-" in compose_profiles
|
||||
and pure_model_name == os.getenv("TEI_MODEL", "")
|
||||
and (provider_name == "Builtin" or provider_name is None)
|
||||
)
|
||||
if is_tei_builtin_embedding:
|
||||
# configured local embedding model
|
||||
embedding_cfg = settings.EMBEDDING_CFG
|
||||
return {
|
||||
"llm_factory": "Builtin",
|
||||
"api_key": embedding_cfg["api_key"],
|
||||
"llm_name": pure_model_name,
|
||||
"api_base": embedding_cfg["base_url"],
|
||||
"model_type": LLMType.EMBEDDING.value,
|
||||
}
|
||||
|
||||
provider_obj = TenantModelProviderService.get_by_tenant_id_and_provider_name(tenant_id, provider_name)
|
||||
if not provider_obj:
|
||||
raise LookupError(f"Provider {provider_name} not found for model {model_name}.")
|
||||
instance_obj = TenantModelInstanceService.get_by_provider_id_and_instance_name(provider_obj.id, instance_name)
|
||||
if not instance_obj:
|
||||
raise LookupError(f"Instance {instance_name} not found for model {model_name}.")
|
||||
model_obj = TenantModelService.get_by_provider_id_and_instance_id_and_model_type_and_model_name(provider_obj.id, instance_obj.id, model_type_val, pure_model_name)
|
||||
|
||||
import json
|
||||
api_key, is_tool, api_key_payload = TenantLLMService._decode_api_key_config(instance_obj.api_key)
|
||||
extra_fields = json.loads(instance_obj.extra) if instance_obj.extra else {}
|
||||
|
||||
if model_obj:
|
||||
if model_obj.status == ActiveStatusEnum.INACTIVE.value:
|
||||
raise LookupError(f"Model {model_name} is disabled.")
|
||||
|
||||
model_config = {
|
||||
"llm_factory": provider_obj.provider_name,
|
||||
"api_key": api_key,
|
||||
"llm_name": model_obj.model_name,
|
||||
"api_base": extra_fields.get("base_url", ""),
|
||||
"model_type": model_obj.model_type,
|
||||
"is_tool": extra_fields.get("is_tool", is_tool)
|
||||
}
|
||||
if api_key_payload is not None:
|
||||
model_config["api_key_payload"] = api_key_payload
|
||||
|
||||
return model_config
|
||||
else:
|
||||
fac_list = [f for f in settings.FACTORY_LLM_INFOS if f["name"] == provider_name]
|
||||
if not fac_list:
|
||||
raise LookupError(f"Model provider config not found: {provider_name}")
|
||||
llm_list = [llm for llm in fac_list[0]["llm"] if llm["llm_name"] == pure_model_name]
|
||||
if not llm_list:
|
||||
raise LookupError(f"Model config not found: {model_name}")
|
||||
llm_info = llm_list[0]
|
||||
model_config = {
|
||||
"llm_factory": provider_obj.provider_name,
|
||||
"api_key": api_key,
|
||||
"llm_name": llm_info["llm_name"],
|
||||
"api_base": extra_fields.get("base_url", ""),
|
||||
"model_type": llm_info["model_type"],
|
||||
"is_tool": llm_info.get("is_tool", is_tool)
|
||||
}
|
||||
if api_key_payload is not None:
|
||||
model_config["api_key_payload"] = api_key_payload
|
||||
return model_config
|
||||
|
||||
|
||||
def get_api_key(tenant_id: str, model_name: str):
|
||||
_, instance_name, provider_name = split_model_name(model_name)
|
||||
|
||||
if not provider_name:
|
||||
raise LookupError("Provider name is required.")
|
||||
provider_obj = TenantModelProviderService.get_by_tenant_id_and_provider_name(tenant_id, provider_name)
|
||||
if not provider_obj:
|
||||
raise LookupError(f"Provider {provider_name} not found.")
|
||||
instance_obj = TenantModelInstanceService.get_by_provider_id_and_instance_name(provider_obj.id, instance_name)
|
||||
if not instance_obj:
|
||||
raise LookupError(f"Instance {instance_name} not found.")
|
||||
return instance_obj.api_key
|
||||
|
||||
|
||||
def get_model_type_by_name(tenant_id: str, model_name: str):
|
||||
pure_model_name, instance_name, provider_name = split_model_name(model_name)
|
||||
provider_obj = TenantModelProviderService.get_by_tenant_id_and_provider_name(tenant_id, provider_name)
|
||||
if not provider_obj:
|
||||
raise LookupError(f"Provider {provider_name} not found for model {model_name}.")
|
||||
instance_obj = TenantModelInstanceService.get_by_provider_id_and_instance_name(provider_obj.id, instance_name)
|
||||
if not instance_obj:
|
||||
raise LookupError(f"Instance {instance_name} not found for model {model_name}.")
|
||||
model_objs = TenantModelService.get_by_provider_id_and_instance_id_and_model_name(provider_obj.id, instance_obj.id, pure_model_name)
|
||||
if not model_objs:
|
||||
fac_list = [f for f in settings.FACTORY_LLM_INFOS if f["name"] == provider_name]
|
||||
if not fac_list:
|
||||
raise LookupError(f"Model provider config not found: {provider_name}")
|
||||
llm_list = [llm for llm in fac_list[0]["llm"] if llm["llm_name"] == pure_model_name]
|
||||
if not llm_list:
|
||||
raise LookupError(f"Model {pure_model_name} not found for model {model_name}.")
|
||||
return [llm_list[0]["model_type"]]
|
||||
return [model_obj.model_type for model_obj in model_objs]
|
||||
|
||||
|
||||
def delete_models_by_instance_ids(instance_ids: list[str]):
|
||||
return TenantModelService.delete_by_instance_ids(instance_ids)
|
||||
|
||||
|
||||
def delete_instances_by_provider_ids(provider_ids: list[str]):
|
||||
return TenantModelInstanceService.delete_by_provider_ids(provider_ids)
|
||||
|
||||
@@ -42,7 +42,7 @@ def _split_name_counter(filename: str) -> tuple[str, int | None]:
|
||||
return filename, None
|
||||
|
||||
|
||||
def duplicate_name(query_func, **kwargs) -> str:
|
||||
def duplicate_name(query_func, name_field: str="name", **kwargs) -> str:
|
||||
"""
|
||||
Generates a unique filename by appending/incrementing a counter when duplicates exist.
|
||||
|
||||
@@ -54,6 +54,7 @@ def duplicate_name(query_func, **kwargs) -> str:
|
||||
query_func: Callable that accepts keyword arguments and returns:
|
||||
- True if name exists (should be modified)
|
||||
- False if name is available
|
||||
name_field: the field name of name in db. default to 'name'
|
||||
**kwargs: Must contain 'name' key with original filename to check
|
||||
|
||||
Returns:
|
||||
@@ -72,10 +73,10 @@ def duplicate_name(query_func, **kwargs) -> str:
|
||||
"""
|
||||
MAX_RETRIES = 1000
|
||||
|
||||
if "name" not in kwargs:
|
||||
raise KeyError("Arguments must contain 'name' key")
|
||||
if name_field not in kwargs:
|
||||
raise KeyError(f"Arguments must contain '{name_field}' key")
|
||||
|
||||
original_name = kwargs["name"]
|
||||
original_name = kwargs[name_field]
|
||||
current_name = original_name
|
||||
retries = 0
|
||||
|
||||
@@ -92,7 +93,7 @@ def duplicate_name(query_func, **kwargs) -> str:
|
||||
|
||||
new_name = f"{main_part}({counter}){suffix}"
|
||||
|
||||
kwargs["name"] = new_name
|
||||
kwargs[name_field] = new_name
|
||||
current_name = new_name
|
||||
retries += 1
|
||||
|
||||
|
||||
@@ -39,8 +39,7 @@ from api.utils.reference_metadata_utils import (
|
||||
enrich_chunks_with_document_metadata,
|
||||
resolve_reference_metadata_preferences,
|
||||
)
|
||||
from api.db.services.tenant_llm_service import TenantLLMService
|
||||
from api.db.joint_services.tenant_model_service import get_model_config_by_id, get_model_config_by_type_and_name, get_tenant_default_model_by_type
|
||||
from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, get_model_config_from_provider_instance, get_model_type_by_name
|
||||
from common.time_utils import current_timestamp, datetime_format
|
||||
from common.text_utils import normalize_arabic_digits
|
||||
from rag.graphrag.general.mind_map_extractor import MindMapExtractor
|
||||
@@ -288,21 +287,19 @@ class DialogService(CommonService):
|
||||
|
||||
|
||||
async def async_chat_solo(dialog, messages, stream=True):
|
||||
llm_type = TenantLLMService.llm_id2llm_type(dialog.llm_id)
|
||||
llm_types = get_model_type_by_name(dialog.tenant_id, dialog.llm_id)
|
||||
attachments = ""
|
||||
image_attachments = []
|
||||
image_files = []
|
||||
if "files" in messages[-1]:
|
||||
if llm_type == "chat":
|
||||
if "chat" in llm_types:
|
||||
text_attachments, image_attachments = split_file_attachments(messages[-1]["files"])
|
||||
else:
|
||||
text_attachments, image_files = split_file_attachments(messages[-1]["files"], raw=True)
|
||||
attachments = "\n\n".join(text_attachments)
|
||||
|
||||
if dialog.llm_id:
|
||||
model_config = get_model_config_by_type_and_name(dialog.tenant_id, LLMType.CHAT, dialog.llm_id)
|
||||
elif dialog.tenant_llm_id:
|
||||
model_config = get_model_config_by_id(dialog.tenant_llm_id)
|
||||
model_config = get_model_config_from_provider_instance(dialog.tenant_id, LLMType.CHAT, dialog.llm_id)
|
||||
else:
|
||||
model_config = get_tenant_default_model_by_type(dialog.tenant_id, LLMType.CHAT)
|
||||
|
||||
@@ -317,10 +314,10 @@ async def async_chat_solo(dialog, messages, stream=True):
|
||||
msg = [{"role": m["role"], "content": re.sub(r"##\d+\$\$", "", m["content"])} for m in messages if m["role"] != "system"]
|
||||
if attachments and msg:
|
||||
msg[-1]["content"] += attachments
|
||||
if llm_type == "chat" and image_attachments:
|
||||
if "chat" in llm_types and image_attachments:
|
||||
convert_last_user_msg_to_multimodal(msg, image_attachments, factory)
|
||||
if stream:
|
||||
if llm_type == "chat":
|
||||
if "chat" in llm_types:
|
||||
stream_iter = chat_mdl.async_chat_streamly_delta(prompt_config.get("system", ""), msg, dialog.llm_setting)
|
||||
else:
|
||||
stream_iter = chat_mdl.async_chat_streamly_delta(prompt_config.get("system", ""), msg, dialog.llm_setting, images=image_files)
|
||||
@@ -331,7 +328,7 @@ async def async_chat_solo(dialog, messages, stream=True):
|
||||
continue
|
||||
yield {"answer": value, "reference": {}, "audio_binary": tts(tts_mdl, value), "prompt": "", "created_at": time.time(), "final": False}
|
||||
else:
|
||||
if llm_type == "chat":
|
||||
if "chat" in llm_types:
|
||||
answer = await chat_mdl.async_chat(prompt_config.get("system", ""), msg, dialog.llm_setting)
|
||||
else:
|
||||
answer = await chat_mdl.async_chat(prompt_config.get("system", ""), msg, dialog.llm_setting, images=image_files)
|
||||
@@ -349,22 +346,20 @@ def get_models(dialog):
|
||||
|
||||
if embedding_list:
|
||||
embd_owner_tenant_id = kbs[0].tenant_id
|
||||
embd_model_config = get_model_config_by_type_and_name(embd_owner_tenant_id, LLMType.EMBEDDING, embedding_list[0])
|
||||
embd_model_config = get_model_config_from_provider_instance(embd_owner_tenant_id, LLMType.EMBEDDING, embedding_list[0])
|
||||
embd_mdl = LLMBundle(embd_owner_tenant_id, embd_model_config)
|
||||
if not embd_mdl:
|
||||
raise LookupError("Embedding model(%s) not found" % embedding_list[0])
|
||||
|
||||
if dialog.llm_id:
|
||||
chat_model_config = get_model_config_by_type_and_name(dialog.tenant_id, LLMType.CHAT, dialog.llm_id)
|
||||
elif dialog.tenant_llm_id:
|
||||
chat_model_config = get_model_config_by_id(dialog.tenant_llm_id)
|
||||
chat_model_config = get_model_config_from_provider_instance(dialog.tenant_id, LLMType.CHAT, dialog.llm_id)
|
||||
else:
|
||||
chat_model_config = get_tenant_default_model_by_type(dialog.tenant_id, LLMType.CHAT)
|
||||
|
||||
chat_mdl = LLMBundle(dialog.tenant_id, chat_model_config)
|
||||
|
||||
if dialog.rerank_id:
|
||||
rerank_model_config = get_model_config_by_type_and_name(dialog.tenant_id, LLMType.RERANK, dialog.rerank_id)
|
||||
rerank_model_config = get_model_config_from_provider_instance(dialog.tenant_id, LLMType.RERANK, dialog.rerank_id)
|
||||
rerank_mdl = LLMBundle(dialog.tenant_id, rerank_model_config)
|
||||
|
||||
if dialog.prompt_config.get("tts"):
|
||||
@@ -554,11 +549,14 @@ async def async_chat(dialog, messages, stream=True, **kwargs):
|
||||
return
|
||||
|
||||
chat_start_ts = timer()
|
||||
llm_type = TenantLLMService.llm_id2llm_type(dialog.llm_id)
|
||||
if llm_type == "image2text":
|
||||
llm_model_config = TenantLLMService.get_model_config(dialog.tenant_id, LLMType.IMAGE2TEXT, dialog.llm_id)
|
||||
if dialog.llm_id:
|
||||
llm_types = get_model_type_by_name(dialog.tenant_id, dialog.llm_id)
|
||||
if "image2text" in llm_types:
|
||||
llm_model_config = get_model_config_from_provider_instance(dialog.tenant_id, LLMType.IMAGE2TEXT, dialog.llm_id)
|
||||
else:
|
||||
llm_model_config = get_model_config_from_provider_instance(dialog.tenant_id, LLMType.CHAT, dialog.llm_id)
|
||||
else:
|
||||
llm_model_config = TenantLLMService.get_model_config(dialog.tenant_id, LLMType.CHAT, dialog.llm_id)
|
||||
llm_model_config = get_tenant_default_model_by_type(dialog.tenant_id, LLMType.CHAT)
|
||||
|
||||
factory = llm_model_config.get("llm_factory", "") if llm_model_config else ""
|
||||
max_tokens = llm_model_config.get("max_tokens", 8192)
|
||||
@@ -598,7 +596,7 @@ async def async_chat(dialog, messages, stream=True, **kwargs):
|
||||
if "doc_ids" in messages[-1]:
|
||||
attachments = [doc_id for doc_id in messages[-1]["doc_ids"] if doc_id]
|
||||
if "files" in messages[-1]:
|
||||
if llm_type == "chat":
|
||||
if llm_model_config["model_type"] == "chat":
|
||||
text_attachments, image_attachments = split_file_attachments(messages[-1]["files"])
|
||||
else:
|
||||
text_attachments, image_files = split_file_attachments(messages[-1]["files"], raw=True)
|
||||
@@ -769,7 +767,7 @@ async def async_chat(dialog, messages, stream=True, **kwargs):
|
||||
prompt4citation = citation_prompt()
|
||||
msg.extend([{"role": m["role"], "content": re.sub(r"##\d+\$\$", "", m["content"])} for m in messages if m["role"] != "system"])
|
||||
used_token_count, msg = message_fit_in(msg, int(max_tokens * 0.95))
|
||||
if llm_type == "chat" and image_attachments:
|
||||
if llm_model_config["model_type"] == "chat" and image_attachments:
|
||||
convert_last_user_msg_to_multimodal(msg, image_attachments, factory)
|
||||
assert len(msg) >= 2, f"message_fit_in has bug: {msg}"
|
||||
prompt = msg[0]["content"]
|
||||
@@ -881,7 +879,7 @@ async def async_chat(dialog, messages, stream=True, **kwargs):
|
||||
langfuse_generation = None
|
||||
|
||||
if stream:
|
||||
if llm_type == "chat":
|
||||
if llm_model_config["model_type"] == "chat":
|
||||
stream_iter = chat_mdl.async_chat_streamly_delta(prompt + prompt4citation, msg[1:], gen_conf)
|
||||
else:
|
||||
stream_iter = chat_mdl.async_chat_streamly_delta(prompt + prompt4citation, msg[1:], gen_conf, images=image_files)
|
||||
@@ -900,7 +898,7 @@ async def async_chat(dialog, messages, stream=True, **kwargs):
|
||||
final["audio_binary"] = None
|
||||
yield final
|
||||
else:
|
||||
if llm_type == "chat":
|
||||
if llm_model_config["model_type"] == "chat":
|
||||
answer = await chat_mdl.async_chat(prompt + prompt4citation, msg[1:], gen_conf)
|
||||
else:
|
||||
answer = await chat_mdl.async_chat(prompt + prompt4citation, msg[1:], gen_conf, images=image_files)
|
||||
@@ -1540,12 +1538,12 @@ async def async_ask(question, kb_ids, tenant_id, chat_llm_name=None, search_conf
|
||||
is_knowledge_graph = all([kb.parser_id == ParserType.KG for kb in kbs])
|
||||
retriever = settings.retriever if not is_knowledge_graph else settings.kg_retriever
|
||||
embd_owner_tenant_id = kbs[0].tenant_id
|
||||
embd_model_config = get_model_config_by_type_and_name(embd_owner_tenant_id, LLMType.EMBEDDING, embedding_list[0])
|
||||
embd_model_config = get_model_config_from_provider_instance(embd_owner_tenant_id, LLMType.EMBEDDING, embedding_list[0])
|
||||
embd_mdl = LLMBundle(embd_owner_tenant_id, embd_model_config)
|
||||
chat_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.CHAT, chat_llm_name)
|
||||
chat_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.CHAT, chat_llm_name)
|
||||
chat_mdl = LLMBundle(tenant_id, chat_model_config)
|
||||
if rerank_id:
|
||||
rerank_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.RERANK, rerank_id)
|
||||
rerank_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.RERANK, rerank_id)
|
||||
rerank_mdl = LLMBundle(tenant_id, rerank_model_config)
|
||||
max_tokens = chat_mdl.max_length
|
||||
tenant_ids = list(set([kb.tenant_id for kb in kbs]))
|
||||
@@ -1649,23 +1647,18 @@ async def gen_mindmap(question, kb_ids, tenant_id, search_config={}):
|
||||
kbs = KnowledgebaseService.get_by_ids(kb_ids)
|
||||
if not kbs:
|
||||
return {"error": "No KB selected"}
|
||||
tenant_embedding_list = list(set([kb.tenant_embd_id for kb in kbs]))
|
||||
tenant_ids = list(set([kb.tenant_id for kb in kbs]))
|
||||
if tenant_embedding_list[0]:
|
||||
embd_model_config = get_model_config_by_id(tenant_embedding_list[0])
|
||||
embd_owner_tenant_id = kbs[0].tenant_id
|
||||
else:
|
||||
embd_owner_tenant_id = kbs[0].tenant_id
|
||||
embd_model_config = get_model_config_by_type_and_name(embd_owner_tenant_id, LLMType.EMBEDDING, kbs[0].embd_id)
|
||||
embd_owner_tenant_id = kbs[0].tenant_id
|
||||
embd_model_config = get_model_config_from_provider_instance(embd_owner_tenant_id, LLMType.EMBEDDING, kbs[0].embd_id)
|
||||
embd_mdl = LLMBundle(embd_owner_tenant_id, embd_model_config)
|
||||
chat_id = search_config.get("chat_id", "")
|
||||
if chat_id:
|
||||
chat_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.CHAT, chat_id)
|
||||
chat_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.CHAT, chat_id)
|
||||
else:
|
||||
chat_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.CHAT)
|
||||
chat_mdl = LLMBundle(tenant_id, chat_model_config)
|
||||
if rerank_id:
|
||||
rerank_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.RERANK, rerank_id)
|
||||
rerank_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.RERANK, rerank_id)
|
||||
rerank_mdl = LLMBundle(tenant_id, rerank_model_config)
|
||||
|
||||
if meta_data_filter:
|
||||
|
||||
@@ -24,7 +24,7 @@ from typing import Generator
|
||||
|
||||
from api.db.db_models import LLM
|
||||
from api.db.services.common_service import CommonService
|
||||
from api.db.services.tenant_llm_service import LLM4Tenant, TenantLLMService
|
||||
from api.db.services.tenant_llm_service import LLM4Tenant
|
||||
from common.constants import LLMType
|
||||
from common.token_utils import num_tokens_from_string
|
||||
|
||||
@@ -137,9 +137,9 @@ class LLMBundle(LLM4Tenant):
|
||||
|
||||
embeddings, used_tokens = self.mdl.encode(safe_texts)
|
||||
if self.model_config["llm_factory"] == "Builtin":
|
||||
logging.debug("LLMBundle.encode_queries query: {}, emd len: {}, used_tokens: {}. Builtin model don't need to update token usage".format(texts, len(embeddings), used_tokens))
|
||||
elif not TenantLLMService.increase_usage_by_id(self.model_config["id"], used_tokens):
|
||||
logging.error("LLMBundle.encode can't update token usage for <tenant redacted>/EMBEDDING used_tokens: {}".format(used_tokens))
|
||||
logging.debug("LLMBundle.encode query: {}, emd len: {}, used_tokens: {}. Builtin model don't need to update token usage".format(texts, len(embeddings), used_tokens))
|
||||
else:
|
||||
logging.info("LLMBundle.encode used_tokens: {}, llm_name: {}".format(used_tokens, self.model_config["llm_name"]))
|
||||
|
||||
if self.langfuse:
|
||||
generation.update(usage_details={"total_tokens": used_tokens})
|
||||
@@ -162,8 +162,8 @@ class LLMBundle(LLM4Tenant):
|
||||
emd, used_tokens = self.mdl.encode_queries(query)
|
||||
if self.model_config["llm_factory"] == "Builtin":
|
||||
logging.info("LLMBundle.encode_queries query: {}, emd len: {}, used_tokens: {}. Builtin model don't need to update token usage".format(query, len(emd), used_tokens))
|
||||
elif not TenantLLMService.increase_usage_by_id(self.model_config["id"], used_tokens):
|
||||
logging.error("LLMBundle.encode_queries can't update token usage for <tenant redacted>/EMBEDDING used_tokens: {}".format(used_tokens))
|
||||
else:
|
||||
logging.info("LLMBundle.encode_queries used_tokens: {}, llm_name: {}".format(used_tokens, self.model_config["llm_name"]))
|
||||
|
||||
if self.langfuse:
|
||||
generation.update(usage_details={"total_tokens": used_tokens})
|
||||
@@ -176,8 +176,7 @@ class LLMBundle(LLM4Tenant):
|
||||
generation = self.langfuse.start_observation(trace_context=self.trace_context, as_type="generation", name="similarity", model=self.model_config["llm_name"], input={"query": query, "texts": texts})
|
||||
|
||||
sim, used_tokens = self.mdl.similarity(query, texts)
|
||||
if not TenantLLMService.increase_usage_by_id(self.model_config["id"], used_tokens):
|
||||
logging.error("LLMBundle.similarity can't update token usage for {}/RERANK used_tokens: {}".format(self.tenant_id, used_tokens))
|
||||
logging.info("LLMBundle.similarity used_tokens: {}, llm_name: {}".format(used_tokens, self.model_config["llm_name"]))
|
||||
|
||||
if self.langfuse:
|
||||
generation.update(usage_details={"total_tokens": used_tokens})
|
||||
@@ -190,8 +189,7 @@ class LLMBundle(LLM4Tenant):
|
||||
generation = self.langfuse.start_observation(trace_context=self.trace_context, as_type="generation", name="describe", metadata={"model": self.model_config["llm_name"]})
|
||||
|
||||
txt, used_tokens = self.mdl.describe(image)
|
||||
if not TenantLLMService.increase_usage_by_id(self.model_config["id"], used_tokens):
|
||||
logging.error("LLMBundle.describe can't update token usage for {}/IMAGE2TEXT used_tokens: {}".format(self.tenant_id, used_tokens))
|
||||
logging.info("LLMBundle.describe used_tokens: {}, llm_name: {}".format(used_tokens, self.model_config["llm_name"]))
|
||||
|
||||
if self.langfuse:
|
||||
generation.update(output={"output": txt}, usage_details={"total_tokens": used_tokens})
|
||||
@@ -204,8 +202,7 @@ class LLMBundle(LLM4Tenant):
|
||||
generation = self.langfuse.start_observation(trace_context=self.trace_context, as_type="generation", name="describe_with_prompt", metadata={"model": self.model_config["llm_name"], "prompt": prompt})
|
||||
|
||||
txt, used_tokens = self.mdl.describe_with_prompt(image, prompt)
|
||||
if not TenantLLMService.increase_usage_by_id(self.model_config["id"], used_tokens):
|
||||
logging.error("LLMBundle.describe can't update token usage for {}/IMAGE2TEXT used_tokens: {}".format(self.tenant_id, used_tokens))
|
||||
logging.info("LLMBundle.describe_with_prompt used_tokens: {}, llm_name: {}".format(used_tokens, self.model_config["llm_name"]))
|
||||
|
||||
if self.langfuse:
|
||||
generation.update(output={"output": txt}, usage_details={"total_tokens": used_tokens})
|
||||
@@ -218,8 +215,7 @@ class LLMBundle(LLM4Tenant):
|
||||
generation = self.langfuse.start_observation(trace_context=self.trace_context, as_type="generation", name="transcription", metadata={"model": self.model_config["llm_name"]})
|
||||
|
||||
txt, used_tokens = self.mdl.transcription(audio)
|
||||
if not TenantLLMService.increase_usage_by_id(self.model_config["id"], used_tokens):
|
||||
logging.error("LLMBundle.transcription can't update token usage for {}/SEQUENCE2TXT used_tokens: {}".format(self.tenant_id, used_tokens))
|
||||
logging.info("LLMBundle.transcription used_tokens: {}, llm_name: {}".format(used_tokens, self.model_config["llm_name"]))
|
||||
|
||||
if self.langfuse:
|
||||
generation.update(output={"output": txt}, usage_details={"total_tokens": used_tokens})
|
||||
@@ -254,7 +250,7 @@ class LLMBundle(LLM4Tenant):
|
||||
finally:
|
||||
if final_text:
|
||||
used_tokens = num_tokens_from_string(final_text)
|
||||
TenantLLMService.increase_usage_by_id(self.model_config["id"], used_tokens)
|
||||
logging.info("LLMBundle.stream_transcription used_tokens: {}, llm_name: {}".format(used_tokens, self.model_config["llm_name"]))
|
||||
|
||||
if self.langfuse:
|
||||
generation.update(
|
||||
@@ -273,8 +269,7 @@ class LLMBundle(LLM4Tenant):
|
||||
)
|
||||
|
||||
full_text, used_tokens = mdl.transcription(audio)
|
||||
if not TenantLLMService.increase_usage_by_id(self.model_config["id"], used_tokens):
|
||||
logging.error(f"LLMBundle.stream_transcription can't update token usage for {self.tenant_id}/SEQUENCE2TXT used_tokens: {used_tokens}")
|
||||
logging.info("LLMBundle.stream_transcription used_tokens: {}, llm_name: {}".format(used_tokens, self.model_config["llm_name"]))
|
||||
|
||||
if self.langfuse:
|
||||
generation.update(
|
||||
@@ -295,8 +290,7 @@ class LLMBundle(LLM4Tenant):
|
||||
|
||||
for chunk in self.mdl.tts(text):
|
||||
if isinstance(chunk, int):
|
||||
if not TenantLLMService.increase_usage_by_id(self.model_config["id"], chunk):
|
||||
logging.error("LLMBundle.tts can't update token usage for {}/TTS".format(self.tenant_id))
|
||||
logging.info("LLMBundle.tts used_tokens: {}, model_name: {}".format(chunk, self.model_config["llm_name"]))
|
||||
return
|
||||
yield chunk
|
||||
|
||||
@@ -431,8 +425,8 @@ class LLMBundle(LLM4Tenant):
|
||||
if not self.verbose_tool_use:
|
||||
txt = re.sub(r"<tool_call>.*?</tool_call>", "", txt, flags=re.DOTALL)
|
||||
|
||||
if used_tokens and not TenantLLMService.increase_usage_by_id(self.model_config["id"], used_tokens):
|
||||
logging.error("LLMBundle.async_chat can't update token usage for {}/CHAT llm_name: {}, used_tokens: {}".format(self.tenant_id, self.model_config["llm_name"], used_tokens))
|
||||
if used_tokens:
|
||||
logging.info("LLMBundle.async_chat used_tokens: {}, llm_name: {}".format(used_tokens, self.model_config["llm_name"]))
|
||||
|
||||
if generation:
|
||||
generation.update(output={"output": txt}, usage_details={"total_tokens": used_tokens})
|
||||
@@ -479,8 +473,8 @@ class LLMBundle(LLM4Tenant):
|
||||
generation.update(output={"error": str(e)})
|
||||
generation.end()
|
||||
raise
|
||||
if total_tokens and not TenantLLMService.increase_usage_by_id(self.model_config["id"], total_tokens):
|
||||
logging.error("LLMBundle.async_chat_streamly can't update token usage for {}/CHAT llm_name: {}, used_tokens: {}".format(self.tenant_id, self.model_config["llm_name"], total_tokens))
|
||||
if total_tokens:
|
||||
logging.info("LLMBundle.async_chat_streamly used_tokens: {}, llm_name: {}".format(total_tokens, self.model_config["llm_name"]))
|
||||
if generation:
|
||||
generation.update(output={"output": ans}, usage_details={"total_tokens": total_tokens})
|
||||
generation.end()
|
||||
@@ -522,8 +516,8 @@ class LLMBundle(LLM4Tenant):
|
||||
generation.update(output={"error": str(e)})
|
||||
generation.end()
|
||||
raise
|
||||
if total_tokens and not TenantLLMService.increase_usage_by_id(self.model_config["id"], total_tokens):
|
||||
logging.error("LLMBundle.async_chat_streamly can't update token usage for {}/CHAT llm_name: {}, used_tokens: {}".format(self.tenant_id, self.model_config["llm_name"], total_tokens))
|
||||
if total_tokens:
|
||||
logging.info("LLMBundle.async_chat_streamly_delta used_tokens: {}, llm_name: {}".format(total_tokens, self.model_config["llm_name"]))
|
||||
if generation:
|
||||
generation.update(output={"output": ans}, usage_details={"total_tokens": total_tokens})
|
||||
generation.end()
|
||||
|
||||
@@ -112,7 +112,7 @@ class MemoryService(CommonService):
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def create_memory(cls, tenant_id: str, name: str, memory_type: List[str], embd_id: str, tenant_embd_id: int, llm_id: str, tenant_llm_id: int):
|
||||
def create_memory(cls, tenant_id: str, name: str, memory_type: List[str], embd_id: str, llm_id: str):
|
||||
# Deduplicate name within tenant
|
||||
memory_name = duplicate_name(
|
||||
cls.query,
|
||||
@@ -131,9 +131,7 @@ class MemoryService(CommonService):
|
||||
"memory_type": calculate_memory_type(memory_type),
|
||||
"tenant_id": tenant_id,
|
||||
"embd_id": embd_id,
|
||||
"tenant_embd_id": tenant_embd_id,
|
||||
"llm_id": llm_id,
|
||||
"tenant_llm_id": tenant_llm_id,
|
||||
"system_prompt": PromptAssembler.assemble_system_prompt({"memory_type": memory_type}),
|
||||
"create_time": timestamp,
|
||||
"create_date": format_time,
|
||||
|
||||
@@ -188,30 +188,36 @@ class TenantLLMService(CommonService):
|
||||
api_key = model_config.get("api_key_payload", model_config["api_key"])
|
||||
if model_config["model_type"] == LLMType.EMBEDDING.value:
|
||||
if model_config["llm_factory"] not in EmbeddingModel:
|
||||
logging.error(f"Factory {model_config['llm_factory']} not in embedding model. Supported factories: {EmbeddingModel.keys()}")
|
||||
return None
|
||||
return EmbeddingModel[model_config["llm_factory"]](api_key, model_config["llm_name"], base_url=model_config["api_base"])
|
||||
|
||||
elif model_config["model_type"] == LLMType.RERANK:
|
||||
elif model_config["model_type"] == LLMType.RERANK.value:
|
||||
if model_config["llm_factory"] not in RerankModel:
|
||||
logging.error(f"Factory {model_config['llm_factory']} not in rerank model. Supported factories: {RerankModel.keys()}")
|
||||
return None
|
||||
return RerankModel[model_config["llm_factory"]](api_key, model_config["llm_name"], base_url=model_config["api_base"])
|
||||
|
||||
elif model_config["model_type"] == LLMType.IMAGE2TEXT.value:
|
||||
if model_config["llm_factory"] not in CvModel:
|
||||
logging.error(f"Factory {model_config['llm_factory']} not in cv model. Supported factories: {CvModel.keys()}")
|
||||
return None
|
||||
return CvModel[model_config["llm_factory"]](api_key, model_config["llm_name"], lang, base_url=model_config["api_base"], **kwargs)
|
||||
|
||||
elif model_config["model_type"] == LLMType.CHAT.value:
|
||||
if model_config["llm_factory"] not in ChatModel:
|
||||
logging.error(f"Factory {model_config['llm_factory']} not in chat model. Supported factories: {ChatModel.keys()}")
|
||||
return None
|
||||
return ChatModel[model_config["llm_factory"]](api_key, model_config["llm_name"], base_url=model_config["api_base"], **kwargs)
|
||||
|
||||
elif model_config["model_type"] == LLMType.SPEECH2TEXT:
|
||||
elif model_config["model_type"] == LLMType.SPEECH2TEXT.value:
|
||||
if model_config["llm_factory"] not in Seq2txtModel:
|
||||
logging.error(f"Factory {model_config['llm_factory']} not in speech2text model. Supported factories: {Seq2txtModel.keys()}")
|
||||
return None
|
||||
return Seq2txtModel[model_config["llm_factory"]](key=api_key, model_name=model_config["llm_name"], lang=lang, base_url=model_config["api_base"])
|
||||
elif model_config["model_type"] == LLMType.TTS:
|
||||
elif model_config["model_type"] == LLMType.TTS.value:
|
||||
if model_config["llm_factory"] not in TTSModel:
|
||||
logging.error(f"Factory {model_config['llm_factory']} not in tts model. Supported factories: {TTSModel.keys()}")
|
||||
return None
|
||||
return TTSModel[model_config["llm_factory"]](
|
||||
api_key,
|
||||
@@ -219,8 +225,9 @@ class TenantLLMService(CommonService):
|
||||
base_url=model_config["api_base"],
|
||||
)
|
||||
|
||||
elif model_config["model_type"] == LLMType.OCR:
|
||||
elif model_config["model_type"] == LLMType.OCR.value:
|
||||
if model_config["llm_factory"] not in OcrModel:
|
||||
logging.error(f"Factory {model_config['llm_factory']} not in ocr model. Supported factories: {OcrModel.keys()}")
|
||||
return None
|
||||
return OcrModel[model_config["llm_factory"]](
|
||||
key=api_key,
|
||||
|
||||
31
api/db/services/tenant_model_group_mapping_service.py
Normal file
31
api/db/services/tenant_model_group_mapping_service.py
Normal file
@@ -0,0 +1,31 @@
|
||||
#
|
||||
# Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
from api.db.db_models import DB, TenantModelGroupMapping
|
||||
from api.db.services.common_service import CommonService
|
||||
|
||||
|
||||
class TenantModelGroupMappingService(CommonService):
|
||||
model = TenantModelGroupMapping
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def get_by_composite_id(cls, group_id, provider_id, instance_id, model_id):
|
||||
return cls.model.get_or_none(
|
||||
cls.model.group_id == group_id,
|
||||
cls.model.provider_id == provider_id,
|
||||
cls.model.instance_id == instance_id,
|
||||
cls.model.model_id == model_id,
|
||||
)
|
||||
21
api/db/services/tenant_model_group_service.py
Normal file
21
api/db/services/tenant_model_group_service.py
Normal file
@@ -0,0 +1,21 @@
|
||||
#
|
||||
# Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
from api.db.db_models import TenantModelGroup
|
||||
from api.db.services.common_service import CommonService
|
||||
|
||||
|
||||
class TenantModelGroupService(CommonService):
|
||||
model = TenantModelGroup
|
||||
69
api/db/services/tenant_model_instance_service.py
Normal file
69
api/db/services/tenant_model_instance_service.py
Normal file
@@ -0,0 +1,69 @@
|
||||
#
|
||||
# Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
from common.misc_utils import get_uuid
|
||||
from api.db.db_models import DB, TenantModelInstance
|
||||
from api.db.services.common_service import CommonService
|
||||
from api.db.services import duplicate_name
|
||||
|
||||
class TenantModelInstanceService(CommonService):
|
||||
model = TenantModelInstance
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def create_instance(cls, provider_id: str, instance_name: str, api_key: str, extra: str):
|
||||
unique_instance_name = duplicate_name(cls.query, name_field="instance_name", provider_id=provider_id, instance_name=instance_name)
|
||||
return cls.insert(id=get_uuid(), provider_id=provider_id, instance_name=unique_instance_name, api_key=api_key, extra=extra)
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def get_all_by_provider_id(cls, provider_id):
|
||||
return list(cls.model.select().where(cls.model.provider_id == provider_id))
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def get_by_provider_ids(cls, provider_ids):
|
||||
return list(cls.model.select().where(cls.model.provider_id.in_(provider_ids)))
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def get_by_provider_id_and_instance_name(cls, provider_id, instance_name):
|
||||
return cls.model.get_or_none(
|
||||
cls.model.provider_id == provider_id,
|
||||
cls.model.instance_name == instance_name,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def get_by_provider_id_and_api_key(cls, provider_id, api_key):
|
||||
return cls.model.get_or_none(
|
||||
cls.model.provider_id == provider_id,
|
||||
cls.model.api_key == api_key
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def delete_by_provider_id_and_instance_name(cls, provider_id, instance_name):
|
||||
return cls.model.delete().where(
|
||||
cls.model.provider_id == provider_id,
|
||||
cls.model.instance_name == instance_name,
|
||||
).execute()
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def delete_by_provider_ids(cls, provider_ids):
|
||||
return cls.model.delete().where(
|
||||
cls.model.provider_id.in_(provider_ids)
|
||||
).execute()
|
||||
52
api/db/services/tenant_model_provider_service.py
Normal file
52
api/db/services/tenant_model_provider_service.py
Normal file
@@ -0,0 +1,52 @@
|
||||
#
|
||||
# Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
from api.db.db_models import DB, TenantModelProvider
|
||||
from api.db.services.common_service import CommonService
|
||||
|
||||
|
||||
class TenantModelProviderService(CommonService):
|
||||
model = TenantModelProvider
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def get_by_tenant_id_and_provider_name(cls, tenant_id, provider_name):
|
||||
return cls.model.get_or_none(
|
||||
cls.model.tenant_id == tenant_id,
|
||||
cls.model.provider_name == provider_name,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def get_by_tenant_id(cls, tenant_id):
|
||||
return list(cls.model.select().where(cls.model.tenant_id == tenant_id))
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def delete_by_tenant_id(cls, tenant_id):
|
||||
return cls.model.delete().where(cls.model.tenant_id == tenant_id).execute()
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def delete_by_tenant_id_and_provider_name(cls, tenant_id, provider_name):
|
||||
return cls.model.delete().where(
|
||||
cls.model.tenant_id == tenant_id,
|
||||
cls.model.provider_name == provider_name,
|
||||
).execute()
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def list_provider_names_by_tenant_id(cls, tenant_id):
|
||||
return [row.provider_name for row in cls.model.select(cls.model.provider_name).where(cls.model.tenant_id == tenant_id)]
|
||||
61
api/db/services/tenant_model_service.py
Normal file
61
api/db/services/tenant_model_service.py
Normal file
@@ -0,0 +1,61 @@
|
||||
#
|
||||
# Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
from api.db.db_models import DB, TenantModel
|
||||
from api.db.services.common_service import CommonService
|
||||
|
||||
|
||||
class TenantModelService(CommonService):
|
||||
model = TenantModel
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def get_by_provider_id_and_instance_id_and_model_name(cls, provider_id, instance_id, model_name):
|
||||
return list(cls.model.select().where(cls.model.provider_id == provider_id, cls.model.instance_id == instance_id, cls.model.model_name == model_name))
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def get_by_provider_id_and_instance_id_and_model_type_and_model_name(cls, provider_id, instance_id, model_type, model_name):
|
||||
return cls.model.get_or_none(
|
||||
cls.model.provider_id == provider_id,
|
||||
cls.model.instance_id == instance_id,
|
||||
cls.model.model_type == model_type,
|
||||
cls.model.model_name == model_name
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def get_models_by_instance_id(cls, instance_id):
|
||||
return list(cls.model.select().where(cls.model.instance_id == instance_id))
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def get_models_by_provider_ids_and_instance_ids(cls, provider_ids, instance_ids):
|
||||
return list(cls.model.select().where(cls.model.provider_id.in_(provider_ids), cls.model.instance_id.in_(instance_ids)))
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def batch_update_model_status(cls, model_ids, status):
|
||||
return cls.model.update(status=status).where(cls.model.id.in_(model_ids)).execute()
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def delete_by_id(cls, model_id):
|
||||
return cls.model.delete().where(cls.model.id == model_id).execute()
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def delete_by_instance_ids(cls, instance_ids):
|
||||
return cls.model.delete().where(cls.model.instance_id.in_(instance_ids)).execute()
|
||||
@@ -191,6 +191,7 @@ class TenantService(CommonService):
|
||||
cls.model.asr_id,
|
||||
cls.model.img2txt_id,
|
||||
cls.model.tts_id,
|
||||
cls.model.ocr_id,
|
||||
cls.model.parser_ids,
|
||||
UserTenant.role]
|
||||
return list(cls.model.select(*fields)
|
||||
|
||||
Reference in New Issue
Block a user