diff --git a/agent/canvas.py b/agent/canvas.py index 2b219ef162..48c4a6cbbc 100644 --- a/agent/canvas.py +++ b/agent/canvas.py @@ -31,6 +31,7 @@ from agent.component.base import ComponentBase from api.db.services.file_service import FileService from api.db.services.llm_service import LLMBundle from api.db.services.task_service import has_canceled +from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type from common.constants import LLMType from common.misc_utils import get_uuid, hash_str2int from common.exceptions import TaskCanceledException @@ -508,7 +509,8 @@ class Canvas(Graph): cpn_obj = self.get_component_obj(self.path[i]) if cpn_obj.component_name.lower() == "message": if cpn_obj.get_param("auto_play"): - tts_mdl = LLMBundle(self._tenant_id, LLMType.TTS) + tts_model_config = get_tenant_default_model_by_type(self._tenant_id, LLMType.TTS) + tts_mdl = LLMBundle(self._tenant_id, tts_model_config) if isinstance(cpn_obj.output("content"), partial): _m = "" buff_m = "" diff --git a/agent/component/agent_with_tools.py b/agent/component/agent_with_tools.py index 55e865346a..b0006fac46 100644 --- a/agent/component/agent_with_tools.py +++ b/agent/component/agent_with_tools.py @@ -28,6 +28,7 @@ from agent.tools.base import LLMToolPluginCallSession, ToolParamBase, ToolBase, from api.db.services.llm_service import LLMBundle from api.db.services.tenant_llm_service import TenantLLMService from api.db.services.mcp_server_service import MCPServerService +from api.db.joint_services.tenant_model_service import get_model_config_by_type_and_name from common.connection_utils import timeout from rag.prompts.generator import next_step_async, COMPLETE_TASK, \ citation_prompt, kb_prompt, citation_plus, full_question, message_fit_in, structured_output_prompt @@ -90,8 +91,8 @@ class Agent(LLM, ToolBase): original_name = cpn.get_meta()["function"]["name"] indexed_name = f"{original_name}_{idx}" self.tools[indexed_name] = cpn - - self.chat_mdl = LLMBundle(self._canvas.get_tenant_id(), TenantLLMService.llm_id2llm_type(self._param.llm_id), self._param.llm_id, + chat_model_config = get_model_config_by_type_and_name(self._canvas.get_tenant_id(), TenantLLMService.llm_id2llm_type(self._param.llm_id), self._param.llm_id) + self.chat_mdl = LLMBundle(self._canvas.get_tenant_id(), chat_model_config, max_retries=self._param.max_retries, retry_interval=self._param.delay_after_error, max_rounds=self._param.max_rounds, diff --git a/agent/component/categorize.py b/agent/component/categorize.py index b5a6a4b9c6..708ce142fe 100644 --- a/agent/component/categorize.py +++ b/agent/component/categorize.py @@ -21,6 +21,7 @@ from abc import ABC from common.constants import LLMType from api.db.services.llm_service import LLMBundle +from api.db.joint_services.tenant_model_service import get_model_config_by_type_and_name from agent.component.llm import LLMParam, LLM from common.connection_utils import timeout from rag.llm.chat_model import ERROR_PREFIX @@ -122,7 +123,8 @@ class Categorize(LLM, ABC): msg[-1]["content"] = query_value self.set_input_value(query_key, msg[-1]["content"]) self._param.update_prompt() - chat_mdl = LLMBundle(self._canvas.get_tenant_id(), LLMType.CHAT, self._param.llm_id) + chat_model_config = get_model_config_by_type_and_name(self._canvas.get_tenant_id(), LLMType.CHAT, self._param.llm_id) + chat_mdl = LLMBundle(self._canvas.get_tenant_id(), chat_model_config) user_prompt = """ ---- Real Data ---- diff --git a/agent/component/llm.py b/agent/component/llm.py index 7538e0d736..f651fd5abd 100644 --- a/agent/component/llm.py +++ b/agent/component/llm.py @@ -25,6 +25,7 @@ from functools import partial from common.constants import LLMType from api.db.services.llm_service import LLMBundle from api.db.services.tenant_llm_service import TenantLLMService +from api.db.joint_services.tenant_model_service import get_model_config_by_type_and_name from agent.component.base import ComponentBase, ComponentParamBase from common.connection_utils import timeout from rag.prompts.generator import tool_call_summary, message_fit_in, citation_prompt, structured_output_prompt @@ -84,10 +85,10 @@ class LLM(ComponentBase): def __init__(self, canvas, component_id, param: ComponentParamBase): super().__init__(canvas, component_id, param) - self.chat_mdl = LLMBundle(self._canvas.get_tenant_id(), TenantLLMService.llm_id2llm_type(self._param.llm_id), - self._param.llm_id, max_retries=self._param.max_retries, - retry_interval=self._param.delay_after_error - ) + chat_model_config = get_model_config_by_type_and_name(self._canvas.get_tenant_id(), TenantLLMService.llm_id2llm_type(self._param.llm_id)) + self.chat_mdl = LLMBundle(self._canvas.get_tenant_id(), chat_model_config, + max_retries=self._param.max_retries, + retry_interval=self._param.delay_after_error) self.imgs = [] def get_input_form(self) -> dict[str, dict]: diff --git a/agent/tools/retrieval.py b/agent/tools/retrieval.py index 29bddde238..abf82fa1ae 100644 --- a/agent/tools/retrieval.py +++ b/agent/tools/retrieval.py @@ -27,6 +27,7 @@ from api.db.services.knowledgebase_service import KnowledgebaseService from api.db.services.llm_service import LLMBundle from api.db.services.memory_service import MemoryService from api.db.joint_services import memory_message_service +from api.db.joint_services.tenant_model_service import get_model_config_by_type_and_name, get_tenant_default_model_by_type from common import settings from common.connection_utils import timeout from rag.app.tag import label_question @@ -113,11 +114,14 @@ class Retrieval(ToolBase, ABC): embd_mdl = None if embd_nms: - embd_mdl = LLMBundle(self._canvas.get_tenant_id(), LLMType.EMBEDDING, embd_nms[0]) + tenant_id = self._canvas.get_tenant_id() + embd_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.EMBEDDING, embd_nms[0]) + embd_mdl = LLMBundle(tenant_id, embd_model_config) rerank_mdl = None if self._param.rerank_id: - rerank_mdl = LLMBundle(kbs[0].tenant_id, LLMType.RERANK, self._param.rerank_id) + rerank_model_config = get_model_config_by_type_and_name(kbs[0].tenant_id, LLMType.RERANK, self._param.rerank_id) + rerank_mdl = LLMBundle(kbs[0].tenant_id, rerank_model_config) vars = self.get_input_elements_from_text(query_text) vars = {k: o["value"] for k, o in vars.items()} @@ -158,7 +162,9 @@ class Retrieval(ToolBase, ABC): chat_mdl = None if self._param.meta_data_filter.get("method") in ["auto", "semi_auto"]: - chat_mdl = LLMBundle(self._canvas.get_tenant_id(), LLMType.CHAT) + tenant_id = self._canvas.get_tenant_id() + chat_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.CHAT) + chat_mdl = LLMBundle(tenant_id, chat_model_config) doc_ids = await apply_meta_data_filter( self._param.meta_data_filter, @@ -192,7 +198,9 @@ class Retrieval(ToolBase, ABC): return if self._param.toc_enhance: - chat_mdl = LLMBundle(self._canvas._tenant_id, LLMType.CHAT) + tenant_id = self._canvas._tenant_id + chat_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.CHAT) + chat_mdl = LLMBundle(tenant_id, chat_model_config) cks = await settings.retriever.retrieval_by_toc(query, kbinfos["chunks"], [kb.tenant_id for kb in kbs], chat_mdl, self._param.top_n) if self.check_if_canceled("Retrieval processing"): @@ -202,11 +210,13 @@ class Retrieval(ToolBase, ABC): kbinfos["chunks"] = settings.retriever.retrieval_by_children(kbinfos["chunks"], [kb.tenant_id for kb in kbs]) if self._param.use_kg: + tenant_id = self._canvas.get_tenant_id() + chat_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.CHAT) ck = await settings.kg_retriever.retrieval(query, [kb.tenant_id for kb in kbs], kb_ids, embd_mdl, - LLMBundle(self._canvas.get_tenant_id(), LLMType.CHAT)) + LLMBundle(tenant_id, chat_model_config)) if self.check_if_canceled("Retrieval processing"): return if ck["content_with_weight"]: @@ -215,8 +225,9 @@ class Retrieval(ToolBase, ABC): kbinfos = {"chunks": [], "doc_aggs": []} if self._param.use_kg and kbs: + chat_model_config = get_tenant_default_model_by_type(kbs[0].tenant_id, LLMType.CHAT) ck = await settings.kg_retriever.retrieval(query, [kb.tenant_id for kb in kbs], filtered_kb_ids, embd_mdl, - LLMBundle(kbs[0].tenant_id, LLMType.CHAT)) + LLMBundle(kbs[0].tenant_id, chat_model_config)) if self.check_if_canceled("Retrieval processing"): return if ck["content_with_weight"]: diff --git a/api/apps/chunk_app.py b/api/apps/chunk_app.py index c1be1ef88c..3b1c153aa8 100644 --- a/api/apps/chunk_app.py +++ b/api/apps/chunk_app.py @@ -28,6 +28,7 @@ from api.db.services.llm_service import LLMBundle from common.metadata_utils import apply_meta_data_filter from api.db.services.search_service import SearchService from api.db.services.user_service import UserTenantService +from api.db.joint_services.tenant_model_service import get_model_config_by_id, get_tenant_default_model_by_type, get_model_config_by_type_and_name from api.utils.api_utils import ( get_data_error_result, get_json_result, @@ -165,13 +166,21 @@ async def set(): if not tenant_id: return get_data_error_result(message="Tenant not found!") - embd_id = DocumentService.get_embd_id(req["doc_id"]) - embd_mdl = LLMBundle(tenant_id, LLMType.EMBEDDING, embd_id) - e, doc = DocumentService.get_by_id(req["doc_id"]) if not e: return get_data_error_result(message="Document not found!") + tenant_embd_id = DocumentService.get_tenant_embd_id(req["doc_id"]) + if tenant_embd_id: + embd_model_config = get_model_config_by_id(tenant_embd_id) + else: + embd_id = DocumentService.get_embd_id(req["doc_id"]) + if embd_id: + embd_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.EMBEDDING, embd_id) + else: + embd_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.EMBEDDING) + embd_mdl = LLMBundle(tenant_id, embd_model_config) + _d = d if doc.parser_id == ParserType.QA: arr = [ @@ -324,8 +333,16 @@ async def create(): if kb.pagerank: d[PAGERANK_FLD] = kb.pagerank - embd_id = DocumentService.get_embd_id(req["doc_id"]) - embd_mdl = LLMBundle(tenant_id, LLMType.EMBEDDING.value, embd_id) + tenant_embd_id = DocumentService.get_tenant_embd_id(req["doc_id"]) + if tenant_embd_id: + embd_model_config = get_model_config_by_id(tenant_embd_id) + else: + embd_id = DocumentService.get_embd_id(req["doc_id"]) + if embd_id: + embd_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.EMBEDDING, embd_id) + else: + embd_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.EMBEDDING) + embd_mdl = LLMBundle(tenant_id, embd_model_config) v, c = embd_mdl.encode([doc.name, req["content_with_weight"] if not d["question_kwd"] else "\n".join(d["question_kwd"])]) v = 0.1 * v[0] + 0.9 * v[1] @@ -375,11 +392,17 @@ async def retrieval_test(): search_config = SearchService.get_detail(req.get("search_id", "")).get("search_config", {}) meta_data_filter = search_config.get("meta_data_filter", {}) if meta_data_filter.get("method") in ["auto", "semi_auto"]: - chat_mdl = LLMBundle(user_id, LLMType.CHAT, llm_name=search_config.get("chat_id", "")) + chat_id = search_config.get("chat_id", "") + if chat_id: + chat_model_config = get_model_config_by_type_and_name(user_id, LLMType.CHAT, search_config["chat_id"]) + else: + chat_model_config = get_tenant_default_model_by_type(user_id, LLMType.CHAT) + chat_mdl = LLMBundle(user_id, chat_model_config) else: meta_data_filter = req.get("meta_data_filter") or {} if meta_data_filter.get("method") in ["auto", "semi_auto"]: - chat_mdl = LLMBundle(user_id, LLMType.CHAT) + chat_model_config = get_tenant_default_model_by_type(user_id, LLMType.CHAT) + chat_mdl = LLMBundle(user_id, chat_model_config) if meta_data_filter: metas = DocMetadataService.get_flatted_meta_by_kbs(kb_ids) @@ -404,15 +427,25 @@ async def retrieval_test(): _question = question if langs: _question = await cross_languages(kb.tenant_id, None, _question, langs) - - embd_mdl = LLMBundle(kb.tenant_id, LLMType.EMBEDDING.value, llm_name=kb.embd_id) + if kb.tenant_embd_id: + embd_model_config = get_model_config_by_id(kb.tenant_embd_id) + elif kb.embd_id: + embd_model_config = get_model_config_by_type_and_name(kb.tenant_id, LLMType.EMBEDDING, kb.embd_id) + else: + embd_model_config = get_tenant_default_model_by_type(kb.tenant_id, LLMType.EMBEDDING) + embd_mdl = LLMBundle(kb.tenant_id, embd_model_config) rerank_mdl = None - if req.get("rerank_id"): - rerank_mdl = LLMBundle(kb.tenant_id, LLMType.RERANK.value, llm_name=req["rerank_id"]) + if req.get("tenant_rerank_id"): + rerank_model_config = get_model_config_by_id(req["tenant_rerank_id"]) + rerank_mdl = LLMBundle(kb.tenant_id, rerank_model_config) + elif req.get("rerank_id"): + rerank_model_config = get_model_config_by_type_and_name(kb.tenant_id, LLMType.RERANK.value, req["rerank_id"]) + rerank_mdl = LLMBundle(kb.tenant_id, rerank_model_config) if req.get("keyword", False): - chat_mdl = LLMBundle(kb.tenant_id, LLMType.CHAT) + default_chat_model_config = get_tenant_default_model_by_type(kb.tenant_id, LLMType.CHAT) + chat_mdl = LLMBundle(kb.tenant_id, default_chat_model_config) _question += await keyword_extraction(chat_mdl, _question) labels = label_question(_question, [kb]) @@ -432,11 +465,12 @@ async def retrieval_test(): ) if use_kg: + default_chat_model_config = get_tenant_default_model_by_type(user_id, LLMType.CHAT) ck = await settings.kg_retriever.retrieval(_question, tenant_ids, kb_ids, embd_mdl, - LLMBundle(kb.tenant_id, LLMType.CHAT)) + LLMBundle(kb.tenant_id, default_chat_model_config)) if ck["content_with_weight"]: ranks["chunks"].insert(0, ck) ranks["chunks"] = settings.retriever.retrieval_by_children(ranks["chunks"], tenant_ids) diff --git a/api/apps/conversation_app.py b/api/apps/conversation_app.py index 482d6d0756..f7518d2255 100644 --- a/api/apps/conversation_app.py +++ b/api/apps/conversation_app.py @@ -27,7 +27,8 @@ from api.db.services.dialog_service import DialogService, async_ask, async_chat, from api.db.services.llm_service import LLMBundle from api.db.services.search_service import SearchService from api.db.services.tenant_llm_service import TenantLLMService -from api.db.services.user_service import TenantService, UserTenantService +from api.db.services.user_service import UserTenantService +from api.db.joint_services.tenant_model_service import get_model_config_by_type_and_name, get_tenant_default_model_by_type from api.utils.api_utils import get_data_error_result, get_json_result, get_request_json, server_error_response, validate_request from rag.prompts.template import load_prompt from rag.prompts.generator import chunks_format @@ -280,15 +281,12 @@ async def sequence2txt(): os.close(fd) await uploaded.save(temp_audio_path) - tenants = TenantService.get_info_by(current_user.id) - if not tenants: - return get_data_error_result(message="Tenant not found!") + try: + default_asr_model_config = get_tenant_default_model_by_type(current_user.id, LLMType.SPEECH2TEXT) + except Exception as e: + return get_data_error_result(message=str(e)) - asr_id = tenants[0]["asr_id"] - if not asr_id: - return get_data_error_result(message="No default ASR model is set") - - asr_mdl=LLMBundle(tenants[0]["tenant_id"], LLMType.SPEECH2TEXT, asr_id) + asr_mdl=LLMBundle(current_user.id, default_asr_model_config) if not stream_mode: text = asr_mdl.transcription(temp_audio_path) try: @@ -317,15 +315,12 @@ async def tts(): req = await get_request_json() text = req["text"] - tenants = TenantService.get_info_by(current_user.id) - if not tenants: - return get_data_error_result(message="Tenant not found!") + try: + default_tts_model_config = get_tenant_default_model_by_type(current_user.id, LLMType.TTS) + except Exception as e: + return get_data_error_result(message=str(e)) - tts_id = tenants[0]["tts_id"] - if not tts_id: - return get_data_error_result(message="No default TTS model is set") - - tts_mdl = LLMBundle(tenants[0]["tenant_id"], LLMType.TTS, tts_id) + tts_mdl = LLMBundle(current_user.id, default_tts_model_config) def stream_audio(): try: @@ -458,7 +453,11 @@ async def related_questions(): question = req["question"] chat_id = search_config.get("chat_id", "") - chat_mdl = LLMBundle(current_user.id, LLMType.CHAT, chat_id) + if chat_id: + chat_model_config = get_model_config_by_type_and_name(current_user.id, LLMType.CHAT, chat_id) + else: + chat_model_config = get_tenant_default_model_by_type(current_user.id, LLMType.CHAT) + chat_mdl = LLMBundle(current_user.id, chat_model_config) gen_conf = search_config.get("llm_setting", {"temperature": 0.9}) if "parameter" in gen_conf: diff --git a/api/apps/dialog_app.py b/api/apps/dialog_app.py index 9b7617797d..043fb39de5 100644 --- a/api/apps/dialog_app.py +++ b/api/apps/dialog_app.py @@ -22,6 +22,7 @@ from api.db.services.tenant_llm_service import TenantLLMService from api.db.services.knowledgebase_service import KnowledgebaseService from api.db.services.user_service import TenantService, UserTenantService from api.utils.api_utils import get_data_error_result, get_json_result, get_request_json, server_error_response, validate_request +from api.utils.tenant_utils import ensure_tenant_model_id_for_params from common.misc_utils import get_uuid from common.constants import RetCode from api.apps import login_required, current_user @@ -33,9 +34,10 @@ import logging @login_required async def set_dialog(): req = await get_request_json() - dialog_id = req.get("dialog_id", "") + dialog_info = ensure_tenant_model_id_for_params(current_user.id, req) + dialog_id = dialog_info.get("dialog_id", "") is_create = not dialog_id - name = req.get("name", "New Dialog") + name = dialog_info.get("name", "New Dialog") if not isinstance(name, str): return get_data_error_result(message="Dialog name must be string.") if name.strip() == "": @@ -57,22 +59,22 @@ async def set_dialog(): name = duplicate_name(_name_exists, name=name) - description = req.get("description", "A helpful dialog") - icon = req.get("icon", "") - top_n = req.get("top_n", 6) - top_k = req.get("top_k", 1024) - rerank_id = req.get("rerank_id", "") + description = dialog_info.get("description", "A helpful dialog") + icon = dialog_info.get("icon", "") + top_n = dialog_info.get("top_n", 6) + top_k = dialog_info.get("top_k", 1024) + rerank_id = dialog_info.get("rerank_id", "") if not rerank_id: - req["rerank_id"] = "" - similarity_threshold = req.get("similarity_threshold", 0.1) - vector_similarity_weight = req.get("vector_similarity_weight", 0.3) - llm_setting = req.get("llm_setting", {}) - meta_data_filter = req.get("meta_data_filter", {}) - prompt_config = req["prompt_config"] + dialog_info["rerank_id"] = "" + similarity_threshold = dialog_info.get("similarity_threshold", 0.1) + vector_similarity_weight = dialog_info.get("vector_similarity_weight", 0.3) + llm_setting = dialog_info.get("llm_setting", {}) + meta_data_filter = dialog_info.get("meta_data_filter", {}) + prompt_config = dialog_info["prompt_config"] # Set default parameters for datasets with knowledge retrieval # All datasets with {knowledge} in system prompt need "knowledge" parameter to enable retrieval - kb_ids = req.get("kb_ids", []) + kb_ids = dialog_info.get("kb_ids", []) parameters = prompt_config.get("parameters") logging.debug(f"set_dialog: kb_ids={kb_ids}, parameters={parameters}, is_create={not is_create}") # Check if parameters is missing, None, or empty list @@ -85,7 +87,7 @@ async def set_dialog(): if not is_create: # only for chat updating - if not req.get("kb_ids", []) and not prompt_config.get("tavily_api_key") and "{knowledge}" in prompt_config.get("system", ""): + if not dialog_info.get("kb_ids", []) and not prompt_config.get("tavily_api_key") and "{knowledge}" in prompt_config.get("system", ""): return get_data_error_result(message="Please remove `{knowledge}` in system prompt since no dataset / Tavily used here.") for p in prompt_config.get("parameters", []): @@ -99,27 +101,30 @@ async def set_dialog(): e, tenant = TenantService.get_by_id(current_user.id) if not e: return get_data_error_result(message="Tenant not found!") - kbs = KnowledgebaseService.get_by_ids(req.get("kb_ids", [])) + kbs = KnowledgebaseService.get_by_ids(dialog_info.get("kb_ids", [])) embd_ids = [TenantLLMService.split_model_name_and_factory(kb.embd_id)[0] for kb in kbs] # remove vendor suffix for comparison embd_count = len(set(embd_ids)) if embd_count > 1: return get_data_error_result(message=f'Datasets use different embedding models: {[kb.embd_id for kb in kbs]}"') - llm_id = req.get("llm_id", tenant.llm_id) + llm_id = dialog_info.get("llm_id", tenant.llm_id) + tenant_llm_id = dialog_info.get("tenant_llm_id", tenant.tenant_llm_id) if not dialog_id: dia = { "id": get_uuid(), "tenant_id": current_user.id, "name": name, - "kb_ids": req.get("kb_ids", []), + "kb_ids": dialog_info.get("kb_ids", []), "description": description, "llm_id": llm_id, + "tenant_llm_id": tenant_llm_id, "llm_setting": llm_setting, "prompt_config": prompt_config, "meta_data_filter": meta_data_filter, "top_n": top_n, "top_k": top_k, "rerank_id": rerank_id, + "tenant_rerank_id": dialog_info.get("tenant_rerank_id", 0), "similarity_threshold": similarity_threshold, "vector_similarity_weight": vector_similarity_weight, "icon": icon @@ -128,16 +133,16 @@ async def set_dialog(): return get_data_error_result(message="Fail to new a dialog!") return get_json_result(data=dia) else: - del req["dialog_id"] - if "kb_names" in req: - del req["kb_names"] - if not DialogService.update_by_id(dialog_id, req): + del dialog_info["dialog_id"] + if "kb_names" in dialog_info: + del dialog_info["kb_names"] + if not DialogService.update_by_id(dialog_id, dialog_info): return get_data_error_result(message="Dialog not found!") e, dia = DialogService.get_by_id(dialog_id) if not e: return get_data_error_result(message="Fail to update a dialog!") dia = dia.to_dict() - dia.update(req) + dia.update(dialog_info) dia["kb_ids"], dia["kb_names"] = get_kb_names(dia["kb_ids"]) return get_json_result(data=dia) except Exception as e: diff --git a/api/apps/kb_app.py b/api/apps/kb_app.py index efb028bf15..8a57bcd637 100644 --- a/api/apps/kb_app.py +++ b/api/apps/kb_app.py @@ -31,6 +31,7 @@ from api.db.services.file_service import FileService from api.db.services.pipeline_operation_log_service import PipelineOperationLogService from api.db.services.task_service import TaskService, GRAPH_RAPTOR_FAKE_DOC_ID from api.db.services.user_service import TenantService, UserTenantService +from api.db.joint_services.tenant_model_service import get_model_config_by_type_and_name, get_model_config_by_id from api.utils.api_utils import ( get_error_data_result, server_error_response, @@ -44,6 +45,7 @@ from api.db import VALID_FILE_TYPES from api.db.services.knowledgebase_service import KnowledgebaseService from api.db.db_models import File from api.utils.api_utils import get_json_result +from api.utils.tenant_utils import ensure_tenant_model_id_for_params from rag.nlp import search from api.constants import DATASET_NAME_LIMIT from rag.utils.redis_conn import REDIS_CONN @@ -57,11 +59,12 @@ from api.apps import login_required, current_user @validate_request("name") async def create(): req = await get_request_json() + create_dict = ensure_tenant_model_id_for_params(current_user.id, req) e, res = KnowledgebaseService.create_with_name( - name = req.pop("name", None), + name = create_dict.pop("name", None), tenant_id = current_user.id, - parser_id = req.pop("parser_id", None), - **req + parser_id = create_dict.pop("parser_id", None), + **create_dict ) if not e: @@ -81,30 +84,31 @@ async def create(): @not_allowed_parameters("id", "tenant_id", "created_by", "create_time", "update_time", "create_date", "update_date", "created_by") async def update(): req = await get_request_json() - if not isinstance(req["name"], str): + update_dict = ensure_tenant_model_id_for_params(current_user.id, req) + if not isinstance(update_dict["name"], str): return get_data_error_result(message="Dataset name must be string.") - if req["name"].strip() == "": + if update_dict["name"].strip() == "": return get_data_error_result(message="Dataset name can't be empty.") - if len(req["name"].encode("utf-8")) > DATASET_NAME_LIMIT: + if len(update_dict["name"].encode("utf-8")) > DATASET_NAME_LIMIT: return get_data_error_result( - message=f"Dataset name length is {len(req['name'])} which is large than {DATASET_NAME_LIMIT}") - req["name"] = req["name"].strip() + message=f"Dataset name length is {len(update_dict['name'])} which is large than {DATASET_NAME_LIMIT}") + update_dict["name"] = update_dict["name"].strip() if settings.DOC_ENGINE_INFINITY: - parser_id = req.get("parser_id") + parser_id = update_dict.get("parser_id") if isinstance(parser_id, str) and parser_id.lower() == "tag": return get_json_result( code=RetCode.OPERATING_ERROR, message="The chunking method Tag has not been supported by Infinity yet.", data=False, ) - if "pagerank" in req and req["pagerank"] > 0: + if "pagerank" in update_dict and update_dict["pagerank"] > 0: return get_json_result( code=RetCode.DATA_ERROR, message="'pagerank' can only be set when doc_engine is elasticsearch", data=False, ) - if not KnowledgebaseService.accessible4deletion(req["kb_id"], current_user.id): + if not KnowledgebaseService.accessible4deletion(update_dict["kb_id"], current_user.id): return get_json_result( data=False, message='No authorization.', @@ -112,15 +116,15 @@ async def update(): ) try: if not KnowledgebaseService.query( - created_by=current_user.id, id=req["kb_id"]): + created_by=current_user.id, id=update_dict["kb_id"]): return get_json_result( data=False, message='Only owner of dataset authorized for this operation.', code=RetCode.OPERATING_ERROR) - e, kb = KnowledgebaseService.get_by_id(req["kb_id"]) + e, kb = KnowledgebaseService.get_by_id(update_dict["kb_id"]) # Rename folder in FileService - if e and req["name"].lower() != kb.name.lower(): + if e and update_dict["name"].lower() != kb.name.lower(): FileService.filter_update( [ File.tenant_id == kb.tenant_id, @@ -128,33 +132,33 @@ async def update(): File.type == "folder", File.name == kb.name, ], - {"name": req["name"]}, + {"name": update_dict["name"]}, ) if not e: return get_data_error_result( message="Can't find this dataset!") - if req["name"].lower() != kb.name.lower() \ + if update_dict["name"].lower() != kb.name.lower() \ and len( - KnowledgebaseService.query(name=req["name"], tenant_id=current_user.id, status=StatusEnum.VALID.value)) >= 1: + KnowledgebaseService.query(name=update_dict["name"], tenant_id=current_user.id, status=StatusEnum.VALID.value)) >= 1: return get_data_error_result( message="Duplicated dataset name.") - del req["kb_id"] + del update_dict["kb_id"] connectors = [] - if "connectors" in req: - connectors = req["connectors"] - del req["connectors"] - if not KnowledgebaseService.update_by_id(kb.id, req): + if "connectors" in update_dict: + connectors = update_dict["connectors"] + del update_dict["connectors"] + if not KnowledgebaseService.update_by_id(kb.id, update_dict): return get_data_error_result() - if kb.pagerank != req.get("pagerank", 0): - if req.get("pagerank", 0) > 0: + if kb.pagerank != update_dict.get("pagerank", 0): + if update_dict.get("pagerank", 0) > 0: await thread_pool_exec( settings.docStoreConn.update, {"kb_id": kb.id}, - {PAGERANK_FLD: req["pagerank"]}, + {PAGERANK_FLD: update_dict["pagerank"]}, search.index_name(kb.tenant_id), kb.id, ) @@ -176,7 +180,7 @@ async def update(): if errors: logging.error("Link KB errors: ", errors) kb = kb.to_dict() - kb.update(req) + kb.update(update_dict) kb["connectors"] = connectors return get_json_result(data=kb) @@ -943,12 +947,18 @@ async def check_embedding(): return s if s else "None" req = await get_request_json() kb_id = req.get("kb_id", "") + tenant_embd_id = req.get("tenant_embd_id") embd_id = req.get("embd_id", "") n = int(req.get("check_num", 5)) _, kb = KnowledgebaseService.get_by_id(kb_id) tenant_id = kb.tenant_id - - emb_mdl = LLMBundle(tenant_id, LLMType.EMBEDDING, embd_id) + if tenant_embd_id: + embd_model_config = get_model_config_by_id(tenant_embd_id) + elif embd_id: + embd_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.EMBEDDING, embd_id) + else: + return get_error_data_result("`tenant_embd_id` or `embd_id` is required.") + emb_mdl = LLMBundle(tenant_id, embd_model_config) samples = sample_random_chunks_with_vectors(settings.docStoreConn, tenant_id=tenant_id, kb_id=kb_id, n=n) results, eff_sims = [], [] diff --git a/api/apps/llm_app.py b/api/apps/llm_app.py index b15d820c23..91c20fddfa 100644 --- a/api/apps/llm_app.py +++ b/api/apps/llm_app.py @@ -410,6 +410,7 @@ def my_llms(): res[o_dict["llm_factory"]]["llm"].append( { + "id": o_dict["id"], "type": o_dict["model_type"], "name": o_dict["llm_name"], "used_token": o_dict["used_tokens"], @@ -423,7 +424,7 @@ def my_llms(): for o in TenantLLMService.get_my_llms(current_user.id): if o["llm_factory"] not in res: res[o["llm_factory"]] = {"tags": o["tags"], "llm": []} - res[o["llm_factory"]]["llm"].append({"type": o["model_type"], "name": o["llm_name"], "used_token": o["used_tokens"], "status": o["status"]}) + res[o["llm_factory"]]["llm"].append({"id": o["id"], "type": o["model_type"], "name": o["llm_name"], "used_token": o["used_tokens"], "status": o["status"]}) return get_json_result(data=res) except Exception as e: @@ -441,10 +442,12 @@ async def list_app(): TenantLLMService.ensure_mineru_from_env(tenant_id) objs = TenantLLMService.query(tenant_id=tenant_id) facts = set([o.to_dict()["llm_factory"] for o in objs if o.api_key and o.status == StatusEnum.VALID.value]) + tenant_llm_mapping = {f"{o.llm_name}@{o.llm_factory}": o for o in objs} status = {(o.llm_name + "@" + o.llm_factory) for o in objs if o.status == StatusEnum.VALID.value} llms = LLMService.get_all() llms = [m.to_dict() for m in llms if m.status == StatusEnum.VALID.value and m.fid not in weighted and (m.fid == "Builtin" or (m.llm_name + "@" + m.fid) in status)] for m in llms: + m["id"] = tenant_llm_mapping.get(m["llm_name"] + "@" + m["fid"], TenantLLM(id=None)).id m["available"] = m["fid"] in facts or m["llm_name"].lower() == "flag-embedding" or m["fid"] in self_deployed if "tei-" in os.getenv("COMPOSE_PROFILES", "") and m["model_type"] == LLMType.EMBEDDING and m["fid"] == "Builtin" and m["llm_name"] == os.getenv("TEI_MODEL", ""): m["available"] = True @@ -453,7 +456,7 @@ async def list_app(): for o in objs: if o.llm_name + "@" + o.llm_factory in llm_set: continue - llms.append({"llm_name": o.llm_name, "model_type": o.model_type, "fid": o.llm_factory, "available": True, "status": StatusEnum.VALID.value}) + llms.append({"id": o.id, "llm_name": o.llm_name, "model_type": o.model_type, "fid": o.llm_factory, "available": True, "status": StatusEnum.VALID.value}) res = {} for m in llms: diff --git a/api/apps/restful_apis/memory_api.py b/api/apps/restful_apis/memory_api.py index c1cd9e5a9e..7238e480a2 100644 --- a/api/apps/restful_apis/memory_api.py +++ b/api/apps/restful_apis/memory_api.py @@ -20,9 +20,10 @@ import time from quart import request from common.constants import RetCode from common.exceptions import ArgumentException, NotFoundException -from api.apps import login_required +from api.apps import login_required, current_user from api.utils.api_utils import validate_request, get_request_json, get_error_argument_result, get_json_result from api.apps.services import memory_api_service +from api.utils.tenant_utils import ensure_tenant_model_id_for_params @manager.route("/memories", methods=["POST"]) # noqa: F821 @@ -32,13 +33,16 @@ async def create_memory(): timing_enabled = os.getenv("RAGFLOW_API_TIMING") t_start = time.perf_counter() if timing_enabled else None req = await get_request_json() + req = ensure_tenant_model_id_for_params(current_user.id, req) t_parsed = time.perf_counter() if timing_enabled else None try: memory_info = { "name": req["name"], "memory_type": req["memory_type"], "embd_id": req["embd_id"], - "llm_id": req["llm_id"] + "llm_id": req["llm_id"], + "tenant_embd_id": req["tenant_embd_id"], + "tenant_llm_id": req["tenant_llm_id"], } success, res = await memory_api_service.create_memory(memory_info) if timing_enabled: @@ -85,7 +89,7 @@ async def update_memory(memory_id): req = await get_request_json() new_settings = {k: req[k] for k in [ "name", "permissions", "llm_id", "embd_id", "memory_type", "memory_size", "forgetting_policy", "temperature", - "avatar", "description", "system_prompt", "user_prompt" + "avatar", "description", "system_prompt", "user_prompt", "tenant_llm_id", "tenant_embd_id" ] if k in req} try: success, res = await memory_api_service.update_memory(memory_id, new_settings) diff --git a/api/apps/sdk/dify_retrieval.py b/api/apps/sdk/dify_retrieval.py index d1e0b2d8b3..07ad5d2003 100644 --- a/api/apps/sdk/dify_retrieval.py +++ b/api/apps/sdk/dify_retrieval.py @@ -21,6 +21,7 @@ from api.db.services.document_service import DocumentService from api.db.services.doc_metadata_service import DocMetadataService from api.db.services.knowledgebase_service import KnowledgebaseService 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, get_tenant_default_model_by_type from common.metadata_utils import meta_filter, convert_conditions from api.utils.api_utils import apikey_required, build_error_result, get_request_json, validate_request from rag.app.tag import label_question @@ -130,8 +131,11 @@ async def retrieval(tenant_id): e, kb = KnowledgebaseService.get_by_id(kb_id) if not e: return build_error_result(message="Knowledgebase not found!", code=RetCode.NOT_FOUND) - - embd_mdl = LLMBundle(kb.tenant_id, LLMType.EMBEDDING.value, llm_name=kb.embd_id) + if kb.tenant_embd_id: + model_config = get_model_config_by_id(kb.tenant_embd_id) + else: + model_config = get_model_config_by_type_and_name(kb.tenant_id, LLMType.EMBEDDING, kb.embd_id) + embd_mdl = LLMBundle(kb.tenant_id, model_config) if metadata_condition: doc_ids.extend(meta_filter(metas, convert_conditions(metadata_condition), metadata_condition.get("logic", "and"))) if not doc_ids and metadata_condition: @@ -152,11 +156,12 @@ async def retrieval(tenant_id): ranks["chunks"] = settings.retriever.retrieval_by_children(ranks["chunks"], [tenant_id]) if use_kg: + model_config = get_tenant_default_model_by_type(kb.tenant_id, LLMType.CHAT) ck = await settings.kg_retriever.retrieval(question, [tenant_id], [kb_id], embd_mdl, - LLMBundle(kb.tenant_id, LLMType.CHAT)) + LLMBundle(kb.tenant_id, model_config)) if ck["content_with_weight"]: ranks["chunks"].insert(0, ck) diff --git a/api/apps/sdk/doc.py b/api/apps/sdk/doc.py index 991e292382..7c1b3aa864 100644 --- a/api/apps/sdk/doc.py +++ b/api/apps/sdk/doc.py @@ -36,6 +36,7 @@ from api.db.services.knowledgebase_service import KnowledgebaseService from api.db.services.llm_service import LLMBundle from api.db.services.tenant_llm_service import TenantLLMService from api.db.services.task_service import TaskService, queue_tasks, cancel_all_task_of +from api.db.joint_services.tenant_model_service import get_model_config_by_id, get_tenant_default_model_by_type, get_model_config_by_type_and_name from common.metadata_utils import meta_filter, convert_conditions from api.utils.api_utils import check_duplicate_ids, construct_json_result, get_error_data_result, get_parser_config, get_result, server_error_response, token_required, \ get_request_json @@ -1248,8 +1249,13 @@ async def add_chunk(tenant_id, dataset_id, document_id): d["tag_kwd"] = req["tag_kwd"] if "tag_feas" in req: d["tag_feas"] = req["tag_feas"] - embd_id = DocumentService.get_embd_id(document_id) - embd_mdl = TenantLLMService.model_instance(tenant_id, LLMType.EMBEDDING.value, embd_id) + tenant_embd_id = DocumentService.get_tenant_embd_id(document_id) + if tenant_embd_id: + model_config = get_model_config_by_id(tenant_embd_id) + else: + embd_id = DocumentService.get_embd_id(document_id) + model_config = get_model_config_by_type_and_name(tenant_id, LLMType.EMBEDDING.value, embd_id) + embd_mdl = TenantLLMService.model_instance(model_config) v, c = embd_mdl.encode([doc.name, req["content"] if not d["question_kwd"] else "\n".join(d["question_kwd"])]) v = 0.1 * v[0] + 0.9 * v[1] d["q_%d_vec" % len(v)] = v.tolist() @@ -1446,8 +1452,13 @@ async def update_chunk(tenant_id, dataset_id, document_id, chunk_id): d["tag_kwd"] = req["tag_kwd"] if "tag_feas" in req: d["tag_feas"] = req["tag_feas"] - embd_id = DocumentService.get_embd_id(document_id) - embd_mdl = TenantLLMService.model_instance(tenant_id, LLMType.EMBEDDING.value, embd_id) + tenant_embd_id = DocumentService.get_tenant_embd_id(document_id) + if tenant_embd_id: + model_config = get_model_config_by_id(tenant_embd_id) + else: + embd_id = DocumentService.get_embd_id(document_id) + model_config = get_model_config_by_type_and_name(tenant_id, LLMType.EMBEDDING.value, embd_id) + embd_mdl = TenantLLMService.model_instance(model_config) if doc.parser_id == ParserType.QA: arr = [t for t in re.split(r"[\n\t]", d["content_with_weight"]) if len(t) > 1] if len(arr) != 2: @@ -1616,17 +1627,26 @@ async def retrieval_test(tenant_id): e, kb = KnowledgebaseService.get_by_id(kb_ids[0]) if not e: return get_error_data_result(message="Dataset not found!") - embd_mdl = LLMBundle(kb.tenant_id, LLMType.EMBEDDING, llm_name=kb.embd_id) + if kb.tenant_embd_id: + embd_model_config = get_model_config_by_id(kb.tenant_embd_id) + else: + embd_model_config = get_model_config_by_type_and_name(kb.tenant_id, LLMType.EMBEDDING, kb.embd_id) + embd_mdl = LLMBundle(kb.tenant_id, embd_model_config) rerank_mdl = None - if req.get("rerank_id"): - rerank_mdl = LLMBundle(kb.tenant_id, LLMType.RERANK, llm_name=req["rerank_id"]) + if req.get("tenant_rerank_id"): + rerank_model_config = get_model_config_by_id(req["tenant_rerank_id"]) + rerank_mdl = LLMBundle(kb.tenant_id, rerank_model_config) + elif req.get("rerank_id"): + rerank_model_config = get_model_config_by_type_and_name(kb.tenant_id, LLMType.RERANK, req["rerank_id"]) + rerank_mdl = LLMBundle(kb.tenant_id, rerank_model_config) if langs: question = await cross_languages(kb.tenant_id, None, question, langs) if req.get("keyword", False): - chat_mdl = LLMBundle(kb.tenant_id, LLMType.CHAT) + chat_model_config = get_tenant_default_model_by_type(kb.tenant_id, LLMType.CHAT) + chat_mdl = LLMBundle(kb.tenant_id, chat_model_config) question += await keyword_extraction(chat_mdl, question) ranks = await settings.retriever.retrieval( @@ -1645,13 +1665,15 @@ async def retrieval_test(tenant_id): rank_feature=label_question(question, kbs), ) if toc_enhance: - chat_mdl = LLMBundle(kb.tenant_id, LLMType.CHAT) + chat_model_config = get_tenant_default_model_by_type(kb.tenant_id, LLMType.CHAT) + chat_mdl = LLMBundle(kb.tenant_id, chat_model_config) cks = await settings.retriever.retrieval_by_toc(question, ranks["chunks"], tenant_ids, chat_mdl, size) if cks: ranks["chunks"] = cks ranks["chunks"] = settings.retriever.retrieval_by_children(ranks["chunks"], tenant_ids) if use_kg: - ck = await settings.kg_retriever.retrieval(question, [k.tenant_id for k in kbs], kb_ids, embd_mdl, LLMBundle(kb.tenant_id, LLMType.CHAT)) + chat_model_config = get_tenant_default_model_by_type(kb.tenant_id, LLMType.CHAT) + ck = await settings.kg_retriever.retrieval(question, [k.tenant_id for k in kbs], kb_ids, embd_mdl, LLMBundle(kb.tenant_id, chat_model_config)) if ck["content_with_weight"]: ranks["chunks"].insert(0, ck) diff --git a/api/apps/sdk/session.py b/api/apps/sdk/session.py index 18bdc4ee53..9553baf1a8 100644 --- a/api/apps/sdk/session.py +++ b/api/apps/sdk/session.py @@ -40,7 +40,9 @@ from api.db.services.knowledgebase_service import KnowledgebaseService from api.db.services.llm_service import LLMBundle from common.metadata_utils import apply_meta_data_filter, convert_conditions, meta_filter from api.db.services.search_service import SearchService -from api.db.services.user_service import TenantService,UserTenantService +from api.db.services.user_service import UserTenantService +from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, get_model_config_by_id, \ + get_model_config_by_type_and_name from common.misc_utils import get_uuid from api.utils.api_utils import check_duplicate_ids, get_data_openai, get_error_data_result, get_json_result, \ get_result, get_request_json, server_error_response, token_required, validate_request @@ -881,7 +883,8 @@ async def related_questions(tenant_id): return get_error_data_result("`question` is required.") question = req["question"] industry = req.get("industry", "") - chat_mdl = LLMBundle(tenant_id, LLMType.CHAT) + chat_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.CHAT) + chat_mdl = LLMBundle(tenant_id, chat_model_config) prompt = """ Objective: To generate search terms related to the user's search keywords, helping users find more valuable information. Instructions: @@ -1109,6 +1112,7 @@ async def retrieval_test_embedded(): top = int(req.get("top_k", 1024)) langs = req.get("cross_languages", []) rerank_id = req.get("rerank_id", "") + tenant_rerank_id = req.get("tenant_rerank_id", "") tenant_id = objs[0].tenant_id if not tenant_id: return get_error_data_result(message="permission denined.") @@ -1125,7 +1129,12 @@ async def retrieval_test_embedded(): search_config = SearchService.get_detail(req.get("search_id", "")).get("search_config", {}) meta_data_filter = search_config.get("meta_data_filter", {}) if meta_data_filter.get("method") in ["auto", "semi_auto"]: - chat_mdl = LLMBundle(tenant_id, LLMType.CHAT, llm_name=search_config.get("chat_id", "")) + 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) + else: + chat_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.CHAT) + chat_mdl = LLMBundle(tenant_id, chat_model_config) # Apply search_config settings if not explicitly provided in request if not req.get("similarity_threshold"): similarity_threshold = float(search_config.get("similarity_threshold", similarity_threshold)) @@ -1138,7 +1147,8 @@ async def retrieval_test_embedded(): else: meta_data_filter = req.get("meta_data_filter") or {} if meta_data_filter.get("method") in ["auto", "semi_auto"]: - chat_mdl = LLMBundle(tenant_id, LLMType.CHAT) + chat_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.CHAT) + chat_mdl = LLMBundle(tenant_id, chat_model_config) if meta_data_filter: metas = DocMetadataService.get_flatted_meta_by_kbs(kb_ids) @@ -1160,15 +1170,23 @@ async def retrieval_test_embedded(): if langs: _question = await cross_languages(kb.tenant_id, None, _question, langs) - - embd_mdl = LLMBundle(kb.tenant_id, LLMType.EMBEDDING.value, llm_name=kb.embd_id) + if kb.tenant_embd_id: + embd_model_config = get_model_config_by_id(kb.tenant_embd_id) + else: + embd_model_config = get_model_config_by_type_and_name(kb.tenant_id, LLMType.EMBEDDING, kb.embd_id) + embd_mdl = LLMBundle(kb.tenant_id, embd_model_config) rerank_mdl = None - if rerank_id: - rerank_mdl = LLMBundle(kb.tenant_id, LLMType.RERANK.value, llm_name=rerank_id) + if tenant_rerank_id: + rerank_model_config = get_model_config_by_id(tenant_rerank_id) + rerank_mdl = LLMBundle(kb.tenant_id, rerank_model_config) + elif rerank_id: + rerank_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.RERANK, rerank_id) + rerank_mdl = LLMBundle(kb.tenant_id, rerank_model_config) if req.get("keyword", False): - chat_mdl = LLMBundle(kb.tenant_id, LLMType.CHAT) + default_chat_model = get_tenant_default_model_by_type(kb.tenant_id, LLMType.CHAT) + chat_mdl = LLMBundle(kb.tenant_id, default_chat_model) _question += await keyword_extraction(chat_mdl, _question) labels = label_question(_question, [kb]) @@ -1177,8 +1195,9 @@ async def retrieval_test_embedded(): local_doc_ids, rerank_mdl=rerank_mdl, highlight=req.get("highlight"), rank_feature=labels ) if use_kg: + default_chat_model = get_tenant_default_model_by_type(kb.tenant_id, LLMType.CHAT) ck = await settings.kg_retriever.retrieval(_question, tenant_ids, kb_ids, embd_mdl, - LLMBundle(kb.tenant_id, LLMType.CHAT)) + LLMBundle(kb.tenant_id, default_chat_model)) if ck["content_with_weight"]: ranks["chunks"].insert(0, ck) @@ -1222,7 +1241,11 @@ async def related_questions_embedded(): question = req["question"] chat_id = search_config.get("chat_id", "") - chat_mdl = LLMBundle(tenant_id, LLMType.CHAT, chat_id) + if chat_id: + chat_model_config = get_model_config_by_type_and_name(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) gen_conf = search_config.get("llm_setting", {"temperature": 0.9}) prompt = load_prompt("related_question") @@ -1323,15 +1346,11 @@ async def sequence2txt(tenant_id): os.close(fd) await uploaded.save(temp_audio_path) - tenants = TenantService.get_info_by(tenant_id) - if not tenants: - return get_error_data_result(message="Tenant not found!") - - asr_id = tenants[0]["asr_id"] - if not asr_id: - return get_error_data_result(message="No default ASR model is set") - - asr_mdl=LLMBundle(tenants[0]["tenant_id"], LLMType.SPEECH2TEXT, asr_id) + try: + default_asr_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.SPEECH2TEXT) + except Exception as e: + return get_error_data_result(message=str(e)) + asr_mdl=LLMBundle(tenant_id, default_asr_model_config) if not stream_mode: text = asr_mdl.transcription(temp_audio_path) try: @@ -1360,15 +1379,11 @@ async def tts(tenant_id): req = await get_request_json() text = req["text"] - tenants = TenantService.get_info_by(tenant_id) - if not tenants: - return get_error_data_result(message="Tenant not found!") - - tts_id = tenants[0]["tts_id"] - if not tts_id: - return get_error_data_result(message="No default TTS model is set") - - tts_mdl = LLMBundle(tenants[0]["tenant_id"], LLMType.TTS, tts_id) + try: + default_tts_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.TTS) + except Exception as e: + return get_error_data_result(message=str(e)) + tts_mdl = LLMBundle(tenant_id, default_tts_model_config) def stream_audio(): try: diff --git a/api/apps/services/memory_api_service.py b/api/apps/services/memory_api_service.py index e49fe7ed0f..e6aa0de2a6 100644 --- a/api/apps/services/memory_api_service.py +++ b/api/apps/services/memory_api_service.py @@ -35,7 +35,9 @@ async def create_memory(memory_info: dict): "name": str, "memory_type": list[str], "embd_id": str, - "llm_id": str + "llm_id": str, + "tenant_embd_id": str, + "tenant_llm_id": str } """ # check name length @@ -58,7 +60,9 @@ async def create_memory(memory_info: dict): name=memory_name, memory_type=memory_type, embd_id=memory_info["embd_id"], - llm_id=memory_info["llm_id"] + llm_id=memory_info["llm_id"], + tenant_llm_id=memory_info["tenant_llm_id"], + tenant_embd_id=memory_info["tenant_embd_id"] ) if success: return True, format_ret_data_from_memory(res) @@ -103,6 +107,10 @@ async def update_memory(memory_id: str, new_memory_setting: dict): update_dict["llm_id"] = new_memory_setting["llm_id"] if new_memory_setting.get("embd_id"): update_dict["embd_id"] = new_memory_setting["embd_id"] + if new_memory_setting.get("tenant_llm_id"): + update_dict["tenant_llm_id"] = new_memory_setting["tenant_llm_id"] + if new_memory_setting.get("tenant_embd_id"): + update_dict["tenant_embd_id"] = new_memory_setting["tenant_embd_id"] if new_memory_setting.get("memory_type"): memory_type = set(new_memory_setting["memory_type"]) invalid_type = memory_type - {e.name.lower() for e in MemoryType} @@ -146,7 +154,7 @@ async def update_memory(memory_id: str, new_memory_setting: dict): return True, memory_dict # check memory empty when update embd_id, memory_type memory_size = get_memory_size_cache(memory_id, current_memory.tenant_id) - not_allowed_update = [f for f in ["embd_id", "memory_type"] if f in to_update and memory_size > 0] + not_allowed_update = [f for f in ["tenant_embd_id", "embd_id", "memory_type"] if f in to_update and memory_size > 0] if not_allowed_update: raise ArgumentException(f"Can't update {not_allowed_update} when memory isn't empty.") if "memory_type" in to_update: diff --git a/api/apps/user_app.py b/api/apps/user_app.py index 3eb8e6c3d3..e08e434d49 100644 --- a/api/apps/user_app.py +++ b/api/apps/user_app.py @@ -45,6 +45,7 @@ from api.utils.api_utils import ( validate_request, ) from api.utils.crypt import decrypt +from api.utils.tenant_utils import ensure_tenant_model_id_for_params from rag.utils.redis_conn import REDIS_CONN from api.apps import login_required, current_user, login_user, logout_user from api.utils.web_utils import ( @@ -841,7 +842,8 @@ async def set_tenant_info(): req = await get_request_json() try: tid = req.pop("tenant_id") - TenantService.update_by_id(tid, req) + update_dict = ensure_tenant_model_id_for_params(tid, req) + TenantService.update_by_id(tid, update_dict) return get_json_result(data=True) except Exception as e: return server_error_response(e) diff --git a/api/db/db_models.py b/api/db/db_models.py index 563c079e9d..b735fbce64 100644 --- a/api/db/db_models.py +++ b/api/db/db_models.py @@ -43,6 +43,7 @@ from peewee import ( Metadata, Model, TextField, + PrimaryKeyField, ) from playhouse.migrate import MySQLMigrator, PostgresqlMigrator, migrate from playhouse.pool import PooledMySQLDatabase, PooledPostgresqlDatabase @@ -737,11 +738,17 @@ class Tenant(DataBaseModel): name = CharField(max_length=100, null=True, help_text="Tenant name", index=True) public_key = CharField(max_length=255, null=True, index=True) llm_id = CharField(max_length=128, null=False, help_text="default llm ID", index=True) + tenant_llm_id = IntegerField(null=True, help_text="id in tenant_llm", index=True) embd_id = CharField(max_length=128, null=False, help_text="default embedding model ID", index=True) + tenant_embd_id = IntegerField(null=True, help_text="id in tenant_llm", index=True) asr_id = CharField(max_length=128, null=False, help_text="default ASR model ID", index=True) + tenant_asr_id = IntegerField(null=True, help_text="id in tenant_llm", index=True) img2txt_id = CharField(max_length=128, null=False, help_text="default image to text model ID", index=True) + tenant_img2txt_id = IntegerField(null=True, help_text="id in tenant_llm", index=True) rerank_id = CharField(max_length=128, null=False, help_text="default rerank model ID", index=True) + 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) 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) @@ -808,14 +815,15 @@ class LLM(DataBaseModel): class TenantLLM(DataBaseModel): + id = PrimaryKeyField() tenant_id = CharField(max_length=32, null=False, index=True) llm_factory = CharField(max_length=128, null=False, help_text="LLM factory name", index=True) model_type = CharField(max_length=128, null=True, help_text="LLM, Text Embedding, Image2Text, ASR", index=True) llm_name = CharField(max_length=128, null=True, help_text="LLM name", default="", index=True) api_key = TextField(null=True, help_text="API KEY") api_base = CharField(max_length=255, null=True, help_text="API Base") - max_tokens = IntegerField(default=8192, index=True) - used_tokens = IntegerField(default=0, index=True) + max_tokens = IntegerField(default=8192, help_text="Max context token num", index=True) + used_tokens = IntegerField(default=0, help_text="Used token num", index=True) status = CharField(max_length=1, null=False, help_text="is it validate(0: wasted, 1: validate)", default="1", index=True) def __str__(self): @@ -823,7 +831,9 @@ class TenantLLM(DataBaseModel): class Meta: db_table = "tenant_llm" - primary_key = CompositeKey("tenant_id", "llm_factory", "llm_name") + indexes = ( + (("tenant_id", "llm_factory", "llm_name"), True), + ) class TenantLangfuse(DataBaseModel): @@ -847,6 +857,7 @@ class Knowledgebase(DataBaseModel): language = CharField(max_length=32, null=True, default="Chinese" if "zh_CN" in os.getenv("LANG", "") else "English", help_text="English|Chinese", index=True) description = TextField(null=True, help_text="KB description") embd_id = CharField(max_length=128, null=False, help_text="default embedding model ID", index=True) + tenant_embd_id = IntegerField(null=True, help_text="id in tenant_llm", index=True) permission = CharField(max_length=16, null=False, help_text="me|team", default="me", index=True) created_by = CharField(max_length=32, null=False, index=True) doc_num = IntegerField(default=0, index=True) @@ -954,6 +965,7 @@ class Dialog(DataBaseModel): icon = TextField(null=True, help_text="icon base64 string") language = CharField(max_length=32, null=True, default="Chinese" if "zh_CN" in os.getenv("LANG", "") else "English", help_text="English|Chinese", index=True) llm_id = CharField(max_length=128, null=False, help_text="default llm ID") + tenant_llm_id = IntegerField(null=True, help_text="id in tenant_llm", index=True) llm_setting = JSONField(null=False, default={"temperature": 0.1, "top_p": 0.3, "frequency_penalty": 0.7, "presence_penalty": 0.4, "max_tokens": 512}) prompt_type = CharField(max_length=16, null=False, default="simple", help_text="simple|advanced", index=True) @@ -973,7 +985,7 @@ class Dialog(DataBaseModel): do_refer = CharField(max_length=1, null=False, default="1", help_text="it needs to insert reference index into answer or not") rerank_id = CharField(max_length=128, null=False, help_text="default rerank model ID") - + tenant_rerank_id = IntegerField(null=True, help_text="id in tenant_llm", index=True) kb_ids = JSONField(null=False, default=[]) status = CharField(max_length=1, null=True, help_text="is it validate(0: wasted, 1: validate)", default="1", index=True) @@ -1295,7 +1307,9 @@ class Memory(DataBaseModel): memory_type = IntegerField(null=False, default=1, index=True, help_text="Bit flags (LSB->MSB): 1=raw, 2=semantic, 4=episodic, 8=procedural. E.g., 5 enables raw + episodic.") storage_type = CharField(max_length=32, default='table', null=False, index=True, help_text="table|graph") embd_id = CharField(max_length=128, null=False, index=False, help_text="embedding model ID") + tenant_embd_id = IntegerField(null=True, help_text="id in tenant_llm", index=True) llm_id = CharField(max_length=128, null=False, index=False, help_text="chat model ID") + tenant_llm_id = IntegerField(null=True, help_text="id in tenant_llm", index=True) permissions = CharField(max_length=16, null=False, index=True, help_text="me|team", default="me") description = TextField(null=True, help_text="description") memory_size = IntegerField(default=5242880, null=False, index=False) @@ -1351,6 +1365,23 @@ def alter_db_rename_column(migrator, table_name, old_column_name, new_column_nam def migrate_add_unique_email(migrator): """Deduplicates user emails and add UNIQUE constraint to email column (idempotent)""" + # step 0: check if UNIQUE index on email already exists + try: + cursor = DB.execute_sql(""" + SELECT COUNT(*) + FROM information_schema.statistics + WHERE table_schema = DATABASE() + AND table_name = 'user' + AND index_name = 'user_email' + AND non_unique = 0 + """) + result = cursor.fetchone() + if result and result[0] > 0: + logging.info("UNIQUE index on user.email already exists, skipping migration") + return + except Exception as ex: + logging.warning("Failed to check if UNIQUE index exists on user.email: %s, continuing with migration", ex) + # step 1: rename duplicate rows so the UNIQUE constraint can be applied try: duplicates = User.select(User.email).group_by(User.email).having(fn.COUNT(User.id) > 1).tuples() @@ -1385,6 +1416,62 @@ def migrate_add_unique_email(migrator): logging.critical("Failed to add UNIQUE constraint on user.email: %s", ex) + +def update_tenant_llm_to_id_primary_key(): + """Add ID and set to primary key step by step.""" + try: + with DB.atomic(): + # 0. Check if exist ID + cursor = DB.execute_sql(""" + SELECT COLUMN_NAME + FROM INFORMATION_SCHEMA.COLUMNS + WHERE TABLE_SCHEMA = DATABASE() + AND TABLE_NAME = 'tenant_llm' + AND COLUMN_NAME = 'id' + """) + if cursor.rowcount > 0: + return + + # 1. Add nullable column + DB.execute_sql("ALTER TABLE tenant_llm ADD COLUMN temp_id INT NULL") + + # 2. Set ID + DB.execute_sql("SET @row = 0;") + DB.execute_sql("UPDATE tenant_llm SET temp_id = (@row := @row + 1) ORDER BY tenant_id, llm_factory, llm_name;") + + # 3. Drop old primary key + DB.execute_sql("ALTER TABLE tenant_llm DROP PRIMARY KEY") + + # 4. Update ID column to primary key + DB.execute_sql(""" + ALTER TABLE tenant_llm + MODIFY COLUMN temp_id INT NOT NULL AUTO_INCREMENT PRIMARY KEY + """) + + # 5. Add unique key + DB.execute_sql(""" + ALTER TABLE tenant_llm + ADD CONSTRAINT uk_tenant_llm UNIQUE (tenant_id, llm_factory, llm_name) + """) + + # 6. rename + DB.execute_sql("ALTER TABLE tenant_llm RENAME COLUMN temp_id TO id") + + logging.info("Successfully updated tenant_llm to id primary key.") + + except Exception as e: + logging.error(str(e)) + cursor = DB.execute_sql(""" + SELECT COLUMN_NAME + FROM INFORMATION_SCHEMA.COLUMNS + WHERE TABLE_SCHEMA = DATABASE() + AND TABLE_NAME = 'tenant_llm' + AND COLUMN_NAME = 'temp_id' + """) + if cursor.rowcount > 0: + DB.execute_sql("ALTER TABLE tenant_llm DROP COLUMN temp_id") + + def migrate_db(): logging.disable(logging.ERROR) migrator = DatabaseMigrator[settings.DATABASE_TYPE.upper()].value(DB) @@ -1436,6 +1523,18 @@ def migrate_db(): alter_db_add_column(migrator, "api_4_conversation", "exp_user_id", CharField(max_length=255, null=True, help_text="exp_user_id", index=True)) # Migrate system_settings.value from CharField to TextField for longer sandbox configs alter_db_column_type(migrator, "system_settings", "value", TextField(null=False, help_text="Configuration value (JSON, string, etc.)")) + update_tenant_llm_to_id_primary_key() + alter_db_add_column(migrator, "tenant", "tenant_llm_id", IntegerField(null=True, help_text="id in tenant_llm", index=True)) + alter_db_add_column(migrator, "tenant", "tenant_embd_id", IntegerField(null=True, help_text="id in tenant_llm", index=True)) + alter_db_add_column(migrator, "tenant", "tenant_asr_id", IntegerField(null=True, help_text="id in tenant_llm", index=True)) + alter_db_add_column(migrator, "tenant", "tenant_img2txt_id", IntegerField(null=True, help_text="id in tenant_llm", index=True)) + alter_db_add_column(migrator, "tenant", "tenant_rerank_id", IntegerField(null=True, help_text="id in tenant_llm", index=True)) + alter_db_add_column(migrator, "tenant", "tenant_tts_id", IntegerField(null=True, help_text="id in tenant_llm", index=True)) + alter_db_add_column(migrator, "knowledgebase", "tenant_embd_id", IntegerField(null=True, help_text="id in tenant_llm", index=True)) + alter_db_add_column(migrator, "dialog", "tenant_llm_id", IntegerField(null=True, help_text="id in tenant_llm", index=True)) + alter_db_add_column(migrator, "dialog", "tenant_rerank_id", IntegerField(null=True, help_text="id in tenant_llm", index=True)) + alter_db_add_column(migrator, "memory", "tenant_embd_id", IntegerField(null=True, help_text="id in tenant_llm", index=True)) + alter_db_add_column(migrator, "memory", "tenant_llm_id", IntegerField(null=True, help_text="id in tenant_llm", index=True)) logging.disable(logging.NOTSET) # this is after re-enabling logging to allow logging changed user emails migrate_add_unique_email(migrator) diff --git a/api/db/init_data.py b/api/db/init_data.py index 525ae5bc5c..229d3ffc01 100644 --- a/api/db/init_data.py +++ b/api/db/init_data.py @@ -24,16 +24,19 @@ 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 +from api.db.db_models import init_database_tables as init_web_db, LLMFactories, LLM, TenantLLM, Knowledgebase, Dialog, Memory 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.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 from common.constants import LLMType from common.file_utils import get_project_base_directory from common import settings @@ -90,13 +93,15 @@ def init_superuser(nickname=DEFAULT_SUPERUSER_NICKNAME, email=DEFAULT_SUPERUSER_ f"Super user initialized. email: {email},A default password has been set; changing the password after login is strongly recommended.") if tenant["llm_id"]: - chat_mdl = LLMBundle(tenant["id"], LLMType.CHAT, tenant["llm_id"]) + chat_model_config = get_tenant_default_model_by_type(tenant["id"], LLMType.CHAT) + chat_mdl = LLMBundle(tenant["id"], chat_model_config) msg = asyncio.run(chat_mdl.async_chat(system="", history=[{"role": "user", "content": "Hello!"}], gen_conf={})) if msg.find("ERROR: ") == 0: logging.error("'{}' doesn't work. {}".format( tenant["llm_id"], msg)) if tenant["embd_id"]: - embd_mdl = LLMBundle(tenant["id"], LLMType.EMBEDDING, tenant["embd_id"]) + embd_model_config = get_tenant_default_model_by_type(tenant["id"], LLMType.EMBEDDING) + embd_mdl = LLMBundle(tenant["id"], embd_model_config) v, c = embd_mdl.encode(["Hello!"]) if c == 0: logging.error("'{}' doesn't work!".format(tenant["embd_id"])) @@ -185,6 +190,7 @@ 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(): @@ -213,6 +219,105 @@ 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_dict += 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() diff --git a/api/db/joint_services/memory_message_service.py b/api/db/joint_services/memory_message_service.py index 8f66212472..9c07f5fbbb 100644 --- a/api/db/joint_services/memory_message_service.py +++ b/api/db/joint_services/memory_message_service.py @@ -25,8 +25,8 @@ from api.db.db_utils import bulk_insert_into_db 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.tenant_llm_service import TenantLLMService 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.utils.memory_utils import get_memory_type_human from memory.services.messages import MessageService from memory.services.query import MsgTextQuery, get_vector @@ -53,11 +53,12 @@ async def save_to_memory(memory_id: str, message_dict: dict): tenant_id = memory.tenant_id extracted_content = await extract_by_llm( tenant_id, - memory.llm_id, + memory.tenant_llm_id, {"temperature": memory.temperature}, get_memory_type_human(memory.memory_type), message_dict.get("user_input", ""), - message_dict.get("agent_response", "") + message_dict.get("agent_response", ""), + llm_id=memory.llm_id ) if memory.memory_type != MemoryType.RAW.value else [] # if only RAW, no need to extract raw_message_id = REDIS_CONN.generate_auto_increment_id(namespace="memory") message_list = [{ @@ -107,12 +108,13 @@ async def save_extracted_to_memory_only(memory_id: str, message_dict, source_mes tenant_id = memory.tenant_id extracted_content = await extract_by_llm( tenant_id, - memory.llm_id, + memory.tenant_llm_id, {"temperature": memory.temperature}, get_memory_type_human(memory.memory_type), message_dict.get("user_input", ""), message_dict.get("agent_response", ""), - task_id=task_id + task_id=task_id, + llm_id=memory.llm_id ) message_list = [{ "message_id": REDIS_CONN.generate_auto_increment_id(namespace="memory"), @@ -139,11 +141,8 @@ async def save_extracted_to_memory_only(memory_id: str, message_dict, source_mes return await embed_and_save(memory, message_list, task_id) -async def extract_by_llm(tenant_id: str, llm_id: str, extract_conf: dict, memory_type: List[str], user_input: str, - agent_response: str, system_prompt: str = "", user_prompt: str="", task_id: str=None) -> List[dict]: - llm_type = TenantLLMService.llm_id2llm_type(llm_id) - if not llm_type: - raise RuntimeError(f"Unknown type of LLM '{llm_id}'") +async def extract_by_llm(tenant_id: str, tenant_llm_id: int, extract_conf: dict, memory_type: List[str], user_input: str, + agent_response: str, system_prompt: str = "", user_prompt: str="", task_id: str=None, llm_id: str = "") -> List[dict]: if not system_prompt: system_prompt = PromptAssembler.assemble_system_prompt({"memory_type": memory_type}) conversation_content = f"User Input: {user_input}\nAgent Response: {agent_response}" @@ -154,7 +153,11 @@ async def extract_by_llm(tenant_id: str, llm_id: str, extract_conf: dict, memory 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)}) - llm = LLMBundle(tenant_id, llm_type, llm_id) + if tenant_llm_id: + llm_config = get_model_config_by_id(tenant_llm_id) + else: + llm_config = get_model_config_by_type_and_name(tenant_id, LLMType.CHAT, llm_id) + llm = LLMBundle(tenant_id, llm_config) if task_id: TaskService.update_progress(task_id, {"progress": 0.15, "progress_msg": timestamp_to_date(current_timestamp())+ " " + "Prepared prompts and LLM."}) res = await llm.async_chat(system_prompt, user_prompts, extract_conf) @@ -170,7 +173,11 @@ async def extract_by_llm(tenant_id: str, llm_id: str, extract_conf: dict, memory async def embed_and_save(memory, message_list: list[dict], task_id: str=None): - embedding_model = LLMBundle(memory.tenant_id, llm_type=LLMType.EMBEDDING, llm_name=memory.embd_id) + if memory.tenant_embd_id: + embd_model_config = get_model_config_by_id(memory.tenant_embd_id) + else: + embd_model_config = get_model_config_by_type_and_name(memory.tenant_id, LLMType.EMBEDDING, memory.embd_id) + embedding_model = LLMBundle(memory.tenant_id, embd_model_config) if task_id: TaskService.update_progress(task_id, {"progress": 0.65, "progress_msg": timestamp_to_date(current_timestamp())+ " " + "Prepared embedding model."}) vector_list, _ = embedding_model.encode([msg["content"] for msg in message_list]) @@ -239,7 +246,11 @@ def query_message(filter_dict: dict, params: dict): question = params["query"] question = question.strip() memory = memory_list[0] - embd_model = LLMBundle(memory.tenant_id, llm_type=LLMType.EMBEDDING, llm_name=memory.embd_id) + if memory.tenant_embd_id: + embd_model_config = get_model_config_by_id(memory.tenant_embd_id) + else: + embd_model_config = get_model_config_by_type_and_name(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"]) keywords_similarity_weight = params.get("keywords_similarity_weight", 0.7) diff --git a/api/db/joint_services/tenant_model_service.py b/api/db/joint_services/tenant_model_service.py new file mode 100644 index 0000000000..32607dcf9f --- /dev/null +++ b/api/db/joint_services/tenant_model_service.py @@ -0,0 +1,91 @@ +# +# 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. +# +import os +import enum +from common import settings +from common.constants import LLMType +from api.db.services.llm_service import LLMService +from api.db.services.tenant_llm_service import TenantLLMService, TenantService + + +def get_model_config_by_id(tenant_model_id: int) -> 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") + config_dict = model_config.to_dict() + llm = LLMService.query(llm_name=config_dict["llm_name"]) + if 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_config = TenantLLMService.get_api_key(tenant_id, model_name) + 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) + if model_type == LLMType.EMBEDDING and fid == "Builtin" and "tei-" in os.getenv("COMPOSE_PROFILES", "") and pure_model_name == os.getenv("TEI_MODEL", ""): + # 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, + } + else: + model_config = TenantLLMService.get_api_key(tenant_id, pure_model_name) + if not model_config: + raise LookupError(f"Tenant Model with name {model_name} not found") + config_dict = model_config.to_dict() + else: + # model_name without @factory + config_dict = model_config.to_dict() + llm = LLMService.query(llm_name=config_dict["llm_name"]) + if 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: + raise LookupError("Tenant not found") + model_type_val = model_type if isinstance(model_type, str) else model_type.value + model_name: str = "" + match model_type_val: + case LLMType.EMBEDDING.value: + model_name = tenant.embd_id + case LLMType.SPEECH2TEXT.value: + model_name = tenant.asr_id + case LLMType.IMAGE2TEXT.value: + model_name = tenant.img2txt_id + case LLMType.CHAT.value: + model_name = tenant.llm_id + case LLMType.RERANK.value: + model_name = tenant.rerank_id + case LLMType.TTS.value: + model_name = tenant.tts_id + case LLMType.OCR.value: + raise Exception("OCR model name is required") + case _: + 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) diff --git a/api/db/services/dialog_service.py b/api/db/services/dialog_service.py index 1dcab82a7f..7a549b69d0 100644 --- a/api/db/services/dialog_service.py +++ b/api/db/services/dialog_service.py @@ -34,6 +34,7 @@ from api.db.services.langfuse_service import TenantLangfuseService from api.db.services.llm_service import LLMBundle from common.metadata_utils import apply_meta_data_filter 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 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 @@ -179,6 +180,28 @@ class DialogService(CommonService): offset += limit return res + @classmethod + @DB.connection_context() + def get_null_tenant_llm_id_row(cls): + fields = [ + cls.model.id, + cls.model.tenant_id, + cls.model.llm_id + ] + objs = cls.model.select(*fields).where(cls.model.tenant_llm_id.is_null()) + return list(objs) + + @classmethod + @DB.connection_context() + def get_null_tenant_rerank_id_row(cls): + fields = [ + cls.model.id, + cls.model.tenant_id, + cls.model.rerank_id + ] + objs = cls.model.select(*fields).where(cls.model.tenant_rerank_id.is_null()) + return list(objs) + async def async_chat_solo(dialog, messages, stream=True): llm_type = TenantLLMService.llm_id2llm_type(dialog.llm_id) @@ -191,22 +214,15 @@ async def async_chat_solo(dialog, messages, stream=True): else: text_attachments, image_files = split_file_attachments(messages[-1]["files"], raw=True) attachments = "\n\n".join(text_attachments) - - if llm_type == "image2text": - llm_model_config = TenantLLMService.get_model_config(dialog.tenant_id, LLMType.IMAGE2TEXT, dialog.llm_id) - else: - llm_model_config = TenantLLMService.get_model_config(dialog.tenant_id, LLMType.CHAT, dialog.llm_id) - factory = llm_model_config.get("llm_factory", "") if llm_model_config else "" - - if llm_type == "image2text": - chat_mdl = LLMBundle(dialog.tenant_id, LLMType.IMAGE2TEXT, dialog.llm_id) - else: - chat_mdl = LLMBundle(dialog.tenant_id, LLMType.CHAT, dialog.llm_id) + model_config = get_model_config_by_id(dialog.tenant_llm_id) + chat_mdl = LLMBundle(dialog.tenant_id, model_config) + factory = model_config.get("llm_factory", "") if model_config else "" prompt_config = dialog.prompt_config tts_mdl = None if prompt_config.get("tts"): - tts_mdl = LLMBundle(dialog.tenant_id, LLMType.TTS) + default_tts_model = get_tenant_default_model_by_type(dialog.tenant_id, LLMType.TTS) + tts_mdl = LLMBundle(dialog.tenant_id, default_tts_model) 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 @@ -241,20 +257,27 @@ def get_models(dialog): raise Exception("**ERROR**: Knowledge bases use different embedding models.") if embedding_list: - embd_mdl = LLMBundle(dialog.tenant_id, LLMType.EMBEDDING, embedding_list[0]) + embd_model_config = get_model_config_by_type_and_name(dialog.tenant_id, LLMType.EMBEDDING, embedding_list[0]) + embd_mdl = LLMBundle(dialog.tenant_id, embd_model_config) if not embd_mdl: raise LookupError("Embedding model(%s) not found" % embedding_list[0]) - if TenantLLMService.llm_id2llm_type(dialog.llm_id) == "image2text": - chat_mdl = LLMBundle(dialog.tenant_id, LLMType.IMAGE2TEXT, dialog.llm_id) + if dialog.tenant_llm_id: + chat_model_config = get_model_config_by_id(dialog.tenant_llm_id) + elif dialog.llm_id: + chat_model_config = get_model_config_by_type_and_name(dialog.tenant_id, LLMType.CHAT, dialog.llm_id) else: - chat_mdl = LLMBundle(dialog.tenant_id, LLMType.CHAT, dialog.llm_id) + 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_mdl = LLMBundle(dialog.tenant_id, LLMType.RERANK, dialog.rerank_id) + rerank_model_config = get_model_config_by_type_and_name(dialog.tenant_id, LLMType.RERANK, dialog.rerank_id) + rerank_mdl = LLMBundle(dialog.tenant_id, rerank_model_config) if dialog.prompt_config.get("tts"): - tts_mdl = LLMBundle(dialog.tenant_id, LLMType.TTS) + default_tts_model_config = get_tenant_default_model_by_type(dialog.tenant_id, LLMType.TTS) + tts_mdl = LLMBundle(dialog.tenant_id, default_tts_model_config) return kbs, embd_mdl, rerank_mdl, chat_mdl, tts_mdl @@ -603,8 +626,9 @@ async def async_chat(dialog, messages, stream=True, **kwargs): kbinfos["chunks"].extend(tav_res["chunks"]) kbinfos["doc_aggs"].extend(tav_res["doc_aggs"]) if prompt_config.get("use_kg"): + default_chat_model = get_tenant_default_model_by_type(dialog.tenant_id, LLMType.CHAT) ck = await settings.kg_retriever.retrieval(" ".join(questions), tenant_ids, dialog.kb_ids, embd_mdl, - LLMBundle(dialog.tenant_id, LLMType.CHAT)) + LLMBundle(dialog.tenant_id, default_chat_model)) if ck["content_with_weight"]: kbinfos["chunks"].insert(0, ck) @@ -1341,11 +1365,13 @@ 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_mdl = LLMBundle(tenant_id, LLMType.EMBEDDING, embedding_list[0]) - chat_mdl = LLMBundle(tenant_id, LLMType.CHAT, chat_llm_name) + embd_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.EMBEDDING, embedding_list[0]) + embd_mdl = LLMBundle(tenant_id, embd_model_config) + chat_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.CHAT, chat_llm_name) + chat_mdl = LLMBundle(tenant_id, chat_model_config) if rerank_id: - rerank_mdl = LLMBundle(tenant_id, LLMType.RERANK, rerank_id) + rerank_model_config = get_model_config_by_type_and_name(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])) @@ -1417,13 +1443,22 @@ 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"} - embedding_list = list(set([kb.embd_id for kb in kbs])) + tenant_embedding_list = list(set([kb.tenant_embd_id for kb in kbs])) tenant_ids = list(set([kb.tenant_id for kb in kbs])) - - embd_mdl = LLMBundle(tenant_id, LLMType.EMBEDDING, llm_name=embedding_list[0]) - chat_mdl = LLMBundle(tenant_id, LLMType.CHAT, llm_name=search_config.get("chat_id", "")) + if tenant_embedding_list[0]: + embd_model_config = get_model_config_by_id(tenant_embedding_list[0]) + else: + embd_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.EMBEDDING, kbs[0].embd_id) + embd_mdl = LLMBundle(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) + 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_mdl = LLMBundle(tenant_id, LLMType.RERANK, rerank_id) + rerank_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.RERANK, rerank_id) + rerank_mdl = LLMBundle(tenant_id, rerank_model_config) if meta_data_filter: metas = DocMetadataService.get_flatted_meta_by_kbs(kb_ids) diff --git a/api/db/services/document_service.py b/api/db/services/document_service.py index a2fee62d0b..9390b79415 100644 --- a/api/db/services/document_service.py +++ b/api/db/services/document_service.py @@ -661,6 +661,19 @@ class DocumentService(CommonService): return None return docs[0]["embd_id"] + @classmethod + @DB.connection_context() + def get_tenant_embd_id(cls, doc_id): + docs = cls.model.select( + Knowledgebase.tenant_embd_id).join( + Knowledgebase, on=( + Knowledgebase.id == cls.model.kb_id)).where( + cls.model.id == doc_id, Knowledgebase.status == StatusEnum.VALID.value) + docs = docs.dicts() + if not docs: + return None + return docs[0]["tenant_embd_id"] + @classmethod @DB.connection_context() def get_chunking_config(cls, doc_id): @@ -1007,6 +1020,7 @@ def doc_upload_and_parse(conversation_id, file_objs, user_id): from api.db.services.file_service import FileService from api.db.services.llm_service import LLMBundle from api.db.services.user_service import TenantService + 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 rag.app import audio, email, naive, picture, presentation e, conv = ConversationService.get_by_id(conversation_id) @@ -1022,8 +1036,11 @@ def doc_upload_and_parse(conversation_id, file_objs, user_id): e, kb = KnowledgebaseService.get_by_id(kb_id) if not e: raise LookupError("Can't find this dataset!") - - embd_mdl = LLMBundle(kb.tenant_id, LLMType.EMBEDDING, llm_name=kb.embd_id, lang=kb.language) + if kb.tenant_embd_id: + embd_model_config = get_model_config_by_id(kb.tenant_embd_id) + else: + embd_model_config = get_model_config_by_type_and_name(kb.tenant_id, LLMType.EMBEDDING, kb.embd_id) + embd_mdl = LLMBundle(kb.tenant_id, embd_model_config, lang=kb.language) err, files = FileService.upload_document(kb, file_objs, user_id) assert not err, "\n".join(err) @@ -1101,7 +1118,8 @@ def doc_upload_and_parse(conversation_id, file_objs, user_id): try_create_idx = True _, tenant = TenantService.get_by_id(kb.tenant_id) - llm_bdl = LLMBundle(kb.tenant_id, LLMType.CHAT, tenant.llm_id) + tenant_llm_config = get_tenant_default_model_by_type(kb.tenant_id, LLMType.CHAT) + llm_bdl = LLMBundle(kb.tenant_id, tenant_llm_config) for doc_id in docids: cks = [c for c in docs if c["doc_id"] == doc_id] diff --git a/api/db/services/knowledgebase_service.py b/api/db/services/knowledgebase_service.py index 1f8b096daa..dcd4038874 100644 --- a/api/db/services/knowledgebase_service.py +++ b/api/db/services/knowledgebase_service.py @@ -564,3 +564,14 @@ class KnowledgebaseService(CommonService): 'update_date': datetime_format(datetime.now()) } return cls.model.update(update_dict).where(cls.model.id == kb_id).execute() + + @classmethod + @DB.connection_context() + def get_null_tenant_embd_id_row(cls): + fields = [ + cls.model.id, + cls.model.tenant_id, + cls.model.embd_id + ] + objs = cls.model.select(*fields).where(cls.model.tenant_embd_id.is_null()) + return list(objs) diff --git a/api/db/services/llm_service.py b/api/db/services/llm_service.py index 16d847e7f1..d24cc4b941 100644 --- a/api/db/services/llm_service.py +++ b/api/db/services/llm_service.py @@ -83,18 +83,18 @@ def get_init_tenant_llm(user_id): class LLMBundle(LLM4Tenant): - def __init__(self, tenant_id, llm_type, llm_name=None, lang="Chinese", **kwargs): - super().__init__(tenant_id, llm_type, llm_name, lang, **kwargs) + def __init__(self, tenant_id: str, model_config: dict, lang="Chinese", **kwargs): + super().__init__(tenant_id, model_config, lang, **kwargs) def bind_tools(self, toolcall_session, tools): if not self.is_tools: - logging.warning(f"Model {self.llm_name} does not support tool call, but you have assigned one or more tools to it!") + logging.warning(f"Model {self.model_config['llm_name']} does not support tool call, but you have assigned one or more tools to it!") return self.mdl.bind_tools(toolcall_session, tools) def encode(self, texts: list): if self.langfuse: - generation = self.langfuse.start_generation(trace_context=self.trace_context, name="encode", model=self.llm_name, input={"texts": texts}) + generation = self.langfuse.start_generation(trace_context=self.trace_context, name="encode", model=self.model_config["llm_name"], input={"texts": texts}) safe_texts = [] for text in texts: @@ -106,9 +106,9 @@ class LLMBundle(LLM4Tenant): safe_texts.append(text) embeddings, used_tokens = self.mdl.encode(safe_texts) - - llm_name = getattr(self, "llm_name", None) - if not TenantLLMService.increase_usage(self.tenant_id, self.llm_type, used_tokens, llm_name): + 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(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 /EMBEDDING used_tokens: {}".format(used_tokens)) if self.langfuse: @@ -119,11 +119,12 @@ class LLMBundle(LLM4Tenant): def encode_queries(self, query: str): if self.langfuse: - generation = self.langfuse.start_generation(trace_context=self.trace_context, name="encode_queries", model=self.llm_name, input={"query": query}) + generation = self.langfuse.start_generation(trace_context=self.trace_context, name="encode_queries", model=self.model_config["llm_name"], input={"query": query}) emd, used_tokens = self.mdl.encode_queries(query) - llm_name = getattr(self, "llm_name", None) - if not TenantLLMService.increase_usage(self.tenant_id, self.llm_type, used_tokens, llm_name): + 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 /EMBEDDING used_tokens: {}".format(used_tokens)) if self.langfuse: @@ -134,10 +135,10 @@ class LLMBundle(LLM4Tenant): def similarity(self, query: str, texts: list): if self.langfuse: - generation = self.langfuse.start_generation(trace_context=self.trace_context, name="similarity", model=self.llm_name, input={"query": query, "texts": texts}) + generation = self.langfuse.start_generation(trace_context=self.trace_context, 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(self.tenant_id, self.llm_type, used_tokens): + 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)) if self.langfuse: @@ -148,10 +149,10 @@ class LLMBundle(LLM4Tenant): def describe(self, image, max_tokens=300): if self.langfuse: - generation = self.langfuse.start_generation(trace_context=self.trace_context, name="describe", metadata={"model": self.llm_name}) + generation = self.langfuse.start_generation(trace_context=self.trace_context, name="describe", metadata={"model": self.model_config["llm_name"]}) txt, used_tokens = self.mdl.describe(image) - if not TenantLLMService.increase_usage(self.tenant_id, self.llm_type, used_tokens): + 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)) if self.langfuse: @@ -162,10 +163,10 @@ class LLMBundle(LLM4Tenant): def describe_with_prompt(self, image, prompt): if self.langfuse: - generation = self.langfuse.start_generation(trace_context=self.trace_context, name="describe_with_prompt", metadata={"model": self.llm_name, "prompt": prompt}) + generation = self.langfuse.start_generation(trace_context=self.trace_context, 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(self.tenant_id, self.llm_type, used_tokens): + 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)) if self.langfuse: @@ -176,10 +177,10 @@ class LLMBundle(LLM4Tenant): def transcription(self, audio): if self.langfuse: - generation = self.langfuse.start_generation(trace_context=self.trace_context, name="transcription", metadata={"model": self.llm_name}) + generation = self.langfuse.start_generation(trace_context=self.trace_context, name="transcription", metadata={"model": self.model_config["llm_name"]}) txt, used_tokens = self.mdl.transcription(audio) - if not TenantLLMService.increase_usage(self.tenant_id, self.llm_type, used_tokens): + 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)) if self.langfuse: @@ -196,7 +197,7 @@ class LLMBundle(LLM4Tenant): generation = self.langfuse.start_generation( trace_context=self.trace_context, name="stream_transcription", - metadata={"model": self.llm_name}, + metadata={"model": self.model_config["llm_name"]}, ) final_text = "" used_tokens = 0 @@ -215,7 +216,7 @@ class LLMBundle(LLM4Tenant): finally: if final_text: used_tokens = num_tokens_from_string(final_text) - TenantLLMService.increase_usage(self.tenant_id, self.llm_type, used_tokens) + TenantLLMService.increase_usage_by_id(self.model_config["id"], used_tokens) if self.langfuse: generation.update( @@ -230,11 +231,11 @@ class LLMBundle(LLM4Tenant): generation = self.langfuse.start_generation( trace_context=self.trace_context, name="stream_transcription", - metadata={"model": self.llm_name}, + metadata={"model": self.model_config["llm_name"]}, ) full_text, used_tokens = mdl.transcription(audio) - if not TenantLLMService.increase_usage(self.tenant_id, self.llm_type, used_tokens): + 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}") if self.langfuse: @@ -256,7 +257,7 @@ class LLMBundle(LLM4Tenant): for chunk in self.mdl.tts(text): if isinstance(chunk, int): - if not TenantLLMService.increase_usage(self.tenant_id, self.llm_type, chunk, self.llm_name): + 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)) return yield chunk @@ -375,7 +376,7 @@ class LLMBundle(LLM4Tenant): generation = None if self.langfuse: - generation = self.langfuse.start_generation(trace_context=self.trace_context, name="chat", model=self.llm_name, input={"system": system, "history": history}) + generation = self.langfuse.start_generation(trace_context=self.trace_context, name="chat", model=self.model_config["llm_name"], input={"system": system, "history": history}) chat_partial = partial(base_fn, system, history, gen_conf) use_kwargs = self._clean_param(chat_partial, **kwargs) @@ -392,8 +393,8 @@ class LLMBundle(LLM4Tenant): if not self.verbose_tool_use: txt = re.sub(r".*?", "", txt, flags=re.DOTALL) - if used_tokens and not TenantLLMService.increase_usage(self.tenant_id, self.llm_type, used_tokens, self.llm_name): - logging.error("LLMBundle.async_chat can't update token usage for {}/CHAT llm_name: {}, used_tokens: {}".format(self.tenant_id, self.llm_name, used_tokens)) + 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 generation: generation.update(output={"output": txt}, usage_details={"total_tokens": used_tokens}) @@ -413,7 +414,7 @@ class LLMBundle(LLM4Tenant): generation = None if self.langfuse: - generation = self.langfuse.start_generation(trace_context=self.trace_context, name="chat_streamly", model=self.llm_name, input={"system": system, "history": history}) + generation = self.langfuse.start_generation(trace_context=self.trace_context, name="chat_streamly", model=self.model_config["llm_name"], input={"system": system, "history": history}) if stream_fn: chat_partial = partial(stream_fn, system, history, gen_conf) @@ -437,8 +438,8 @@ class LLMBundle(LLM4Tenant): generation.update(output={"error": str(e)}) generation.end() raise - if total_tokens and not TenantLLMService.increase_usage(self.tenant_id, self.llm_type, total_tokens, self.llm_name): - logging.error("LLMBundle.async_chat_streamly can't update token usage for {}/CHAT llm_name: {}, used_tokens: {}".format(self.tenant_id, self.llm_name, total_tokens)) + 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 generation: generation.update(output={"output": ans}, usage_details={"total_tokens": total_tokens}) generation.end() @@ -456,7 +457,7 @@ class LLMBundle(LLM4Tenant): generation = None if self.langfuse: - generation = self.langfuse.start_generation(trace_context=self.trace_context, name="chat_streamly", model=self.llm_name, input={"system": system, "history": history}) + generation = self.langfuse.start_generation(trace_context=self.trace_context, name="chat_streamly", model=self.model_config["llm_name"], input={"system": system, "history": history}) if stream_fn: chat_partial = partial(stream_fn, system, history, gen_conf) @@ -480,8 +481,8 @@ class LLMBundle(LLM4Tenant): generation.update(output={"error": str(e)}) generation.end() raise - if total_tokens and not TenantLLMService.increase_usage(self.tenant_id, self.llm_type, total_tokens, self.llm_name): - logging.error("LLMBundle.async_chat_streamly can't update token usage for {}/CHAT llm_name: {}, used_tokens: {}".format(self.tenant_id, self.llm_name, total_tokens)) + 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 generation: generation.update(output={"output": ans}, usage_details={"total_tokens": total_tokens}) generation.end() diff --git a/api/db/services/memory_service.py b/api/db/services/memory_service.py index 215a198fe2..d2433d01d0 100644 --- a/api/db/services/memory_service.py +++ b/api/db/services/memory_service.py @@ -107,7 +107,7 @@ class MemoryService(CommonService): @classmethod @DB.connection_context() - def create_memory(cls, tenant_id: str, name: str, memory_type: List[str], embd_id: str, llm_id: str): + 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): # Deduplicate name within tenant memory_name = duplicate_name( cls.query, @@ -126,7 +126,9 @@ 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, @@ -168,3 +170,25 @@ class MemoryService(CommonService): @DB.connection_context() def delete_memory(cls, memory_id: str): return cls.delete_by_id(memory_id) + + @classmethod + @DB.connection_context() + def get_null_tenant_embd_id_row(cls): + fields = [ + cls.model.id, + cls.model.tenant_id, + cls.model.embd_id + ] + objs = cls.model.select(*fields).where(cls.model.tenant_embd_id.is_null()) + return list(objs) + + @classmethod + @DB.connection_context() + def get_null_tenant_llm_id_row(cls): + fields = [ + cls.model.id, + cls.model.tenant_id, + cls.model.llm_id + ] + objs = cls.model.select(*fields).where(cls.model.tenant_llm_id.is_null()) + return list(objs) diff --git a/api/db/services/tenant_llm_service.py b/api/db/services/tenant_llm_service.py index 5bd663734a..17480c1413 100644 --- a/api/db/services/tenant_llm_service.py +++ b/api/db/services/tenant_llm_service.py @@ -60,7 +60,7 @@ class TenantLLMService(CommonService): @classmethod @DB.connection_context() def get_my_llms(cls, tenant_id): - fields = [cls.model.llm_factory, LLMFactories.logo, LLMFactories.tags, cls.model.model_type, cls.model.llm_name, cls.model.used_tokens, cls.model.status] + fields = [cls.model.id, cls.model.llm_factory, LLMFactories.logo, LLMFactories.tags, cls.model.model_type, cls.model.llm_name, cls.model.used_tokens, cls.model.status] objs = cls.model.select(*fields).join(LLMFactories, on=(cls.model.llm_factory == LLMFactories.name)).where(cls.model.tenant_id == tenant_id, ~cls.model.api_key.is_null()).dicts() return list(objs) @@ -133,34 +133,35 @@ class TenantLLMService(CommonService): @classmethod @DB.connection_context() - def model_instance(cls, tenant_id, llm_type, llm_name=None, lang="Chinese", **kwargs): - model_config = TenantLLMService.get_model_config(tenant_id, llm_type, llm_name) + def model_instance(cls, model_config: dict, lang="Chinese", **kwargs): + if not model_config: + raise LookupError("Model config is required") kwargs.update({"provider": model_config["llm_factory"]}) - if llm_type == LLMType.EMBEDDING.value: + if model_config["model_type"] == LLMType.EMBEDDING.value: if model_config["llm_factory"] not in EmbeddingModel: return None return EmbeddingModel[model_config["llm_factory"]](model_config["api_key"], model_config["llm_name"], base_url=model_config["api_base"]) - elif llm_type == LLMType.RERANK: + elif model_config["model_type"] == LLMType.RERANK: if model_config["llm_factory"] not in RerankModel: return None return RerankModel[model_config["llm_factory"]](model_config["api_key"], model_config["llm_name"], base_url=model_config["api_base"]) - elif llm_type == LLMType.IMAGE2TEXT.value: + elif model_config["model_type"] == LLMType.IMAGE2TEXT.value: if model_config["llm_factory"] not in CvModel: return None return CvModel[model_config["llm_factory"]](model_config["api_key"], model_config["llm_name"], lang, base_url=model_config["api_base"], **kwargs) - elif llm_type == LLMType.CHAT.value: + elif model_config["model_type"] == LLMType.CHAT.value: if model_config["llm_factory"] not in ChatModel: return None return ChatModel[model_config["llm_factory"]](model_config["api_key"], model_config["llm_name"], base_url=model_config["api_base"], **kwargs) - elif llm_type == LLMType.SPEECH2TEXT: + elif model_config["model_type"] == LLMType.SPEECH2TEXT: if model_config["llm_factory"] not in Seq2txtModel: return None return Seq2txtModel[model_config["llm_factory"]](key=model_config["api_key"], model_name=model_config["llm_name"], lang=lang, base_url=model_config["api_base"]) - elif llm_type == LLMType.TTS: + elif model_config["model_type"] == LLMType.TTS: if model_config["llm_factory"] not in TTSModel: return None return TTSModel[model_config["llm_factory"]]( @@ -169,7 +170,7 @@ class TenantLLMService(CommonService): base_url=model_config["api_base"], ) - elif llm_type == LLMType.OCR: + elif model_config["model_type"] == LLMType.OCR: if model_config["llm_factory"] not in OcrModel: return None return OcrModel[model_config["llm_factory"]]( @@ -218,6 +219,16 @@ class TenantLLMService(CommonService): return num + @classmethod + @DB.connection_context() + def increase_usage_by_id(cls, tenant_model_id: int, used_tokens: int): + try: + update_cnt = cls.model.update(used_tokens=cls.model.used_tokens + used_tokens).where(cls.model.id == tenant_model_id).execute() + except Exception as e: + logging.exception(f"TenantLLMService.increase_usage got exception {e}, Failed to update used_tokens for tenant_model_id {tenant_model_id}") + return 0 + return update_cnt + @classmethod @DB.connection_context() def get_openai_models(cls): @@ -376,13 +387,12 @@ class TenantLLMService(CommonService): class LLM4Tenant: - def __init__(self, tenant_id, llm_type, llm_name=None, lang="Chinese", **kwargs): + def __init__(self, tenant_id: str, model_config: dict, lang="Chinese", **kwargs): self.tenant_id = tenant_id - self.llm_type = llm_type - self.llm_name = llm_name - self.mdl = TenantLLMService.model_instance(tenant_id, llm_type, llm_name, lang=lang, **kwargs) - assert self.mdl, "Can't find model for {}/{}/{}".format(tenant_id, llm_type, llm_name) - model_config = TenantLLMService.get_model_config(tenant_id, llm_type, llm_name) + self.llm_name = model_config["llm_name"] + self.model_config = model_config + self.mdl = TenantLLMService.model_instance(model_config, lang=lang, **kwargs) + assert self.mdl, "Can't find model for {}/{}/{}".format(tenant_id, model_config["llm_type"], model_config["llm_name"]) self.max_length = model_config.get("max_tokens", 8192) self.is_tools = model_config.get("is_tools", False) diff --git a/api/db/services/user_service.py b/api/db/services/user_service.py index 20d8c3230f..6804dbd445 100644 --- a/api/db/services/user_service.py +++ b/api/db/services/user_service.py @@ -226,6 +226,12 @@ class TenantService(CommonService): hash_obj = hashlib.sha256(tenant_id.encode("utf-8")) return int(hash_obj.hexdigest(), 16)%len(settings.MINIO) + @classmethod + @DB.connection_context() + def get_null_tenant_model_id_rows(cls): + objs = cls.model.select().orwhere(cls.model.tenant_llm_id.is_null(), cls.model.tenant_embd_id.is_null(), cls.model.tenant_asr_id.is_null(), cls.model.tenant_tts_id.is_null(), cls.model.tenant_rerank_id.is_null(), cls.model.tenant_img2txt_id.is_null()) + return list(objs) + class UserTenantService(CommonService): """Service class for managing user-tenant relationship operations. diff --git a/api/utils/tenant_utils.py b/api/utils/tenant_utils.py new file mode 100644 index 0000000000..cf362e3c23 --- /dev/null +++ b/api/utils/tenant_utils.py @@ -0,0 +1,26 @@ +# +# 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.services.tenant_llm_service import TenantLLMService + +def ensure_tenant_model_id_for_params(tenant_id: str, param_dict: dict) -> dict: + for key in ["llm_id", "embd_id", "asr_id", "img2txt_id", "rerank_id", "tts_id"]: + if param_dict.get(key) and not param_dict.get(f"tenant_{key}"): + tenant_model = TenantLLMService.get_api_key(tenant_id, param_dict[key]) + if tenant_model: + param_dict.update({f"tenant_{key}": tenant_model.id}) + else: + param_dict.update({f"tenant_{key}": 0}) + return param_dict diff --git a/deepdoc/parser/figure_parser.py b/deepdoc/parser/figure_parser.py index d10795a431..3b85e5648f 100644 --- a/deepdoc/parser/figure_parser.py +++ b/deepdoc/parser/figure_parser.py @@ -20,6 +20,7 @@ from PIL import Image from common.constants import LLMType from api.db.services.llm_service import LLMBundle +from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type from common.connection_utils import timeout from rag.app.picture import vision_llm_chunk as picture_vision_llm_chunk from rag.prompts.generator import vision_llm_figure_describe_prompt, vision_llm_figure_describe_prompt_with_context @@ -47,7 +48,8 @@ def vision_figure_parser_docx_wrapper(sections, tbls, callback=None,**kwargs): if not sections: return tbls try: - vision_model = LLMBundle(kwargs["tenant_id"], LLMType.IMAGE2TEXT) + vision_model_config = get_tenant_default_model_by_type(kwargs["tenant_id"], LLMType.IMAGE2TEXT) + vision_model = LLMBundle(kwargs["tenant_id"], vision_model_config) callback(0.7, "Visual model detected. Attempting to enhance figure extraction...") except Exception: vision_model = None @@ -66,7 +68,8 @@ def vision_figure_parser_figure_xlsx_wrapper(images,callback=None, **kwargs): if not images: return [] try: - vision_model = LLMBundle(kwargs["tenant_id"], LLMType.IMAGE2TEXT) + vision_model_config = get_tenant_default_model_by_type(kwargs["tenant_id"], LLMType.IMAGE2TEXT) + vision_model = LLMBundle(kwargs["tenant_id"], vision_model_config) callback(0.2, "Visual model detected. Attempting to enhance Excel image extraction...") except Exception: vision_model = None @@ -94,7 +97,8 @@ def vision_figure_parser_pdf_wrapper(tbls, callback=None, **kwargs): parser_config = kwargs.get("parser_config", {}) context_size = max(0, int(parser_config.get("image_context_size", 0) or 0)) try: - vision_model = LLMBundle(kwargs["tenant_id"], LLMType.IMAGE2TEXT) + vision_model_config = get_tenant_default_model_by_type(kwargs["tenant_id"], LLMType.IMAGE2TEXT) + vision_model = LLMBundle(kwargs["tenant_id"], vision_model_config) callback(0.7, "Visual model detected. Attempting to enhance figure extraction...") except Exception: vision_model = None @@ -132,7 +136,8 @@ def vision_figure_parser_docx_wrapper_naive(chunks, idx_lst, callback=None, **kw if not chunks: return [] try: - vision_model = LLMBundle(kwargs["tenant_id"], LLMType.IMAGE2TEXT) + vision_model_config = get_tenant_default_model_by_type(kwargs["tenant_id"], LLMType.IMAGE2TEXT) + vision_model = LLMBundle(kwargs["tenant_id"], vision_model_config) callback(0.7, "Visual model detected. Attempting to enhance figure extraction...") except Exception: vision_model = None diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh index 5dd300a78d..f6557b3e4e 100755 --- a/docker/entrypoint.sh +++ b/docker/entrypoint.sh @@ -212,10 +212,16 @@ function ensure_docling() { || uv pip install -i https://pypi.tuna.tsinghua.edu.cn/simple --extra-index-url https://pypi.org/simple --no-cache-dir "docling${DOCLING_PIN}" } +function ensure_db_init() { + echo "Initializing database tables..." + "$PY" -c "from api.db.db_models import init_database_tables as init_web_db; init_web_db()" +} + # ----------------------------------------------------------------------------- # Start components based on flags # ----------------------------------------------------------------------------- ensure_docling +ensure_db_init if [[ "${ENABLE_WEBSERVER}" -eq 1 ]]; then echo "Starting nginx..." diff --git a/rag/app/audio.py b/rag/app/audio.py index 5bcb3d2573..29ef625fad 100644 --- a/rag/app/audio.py +++ b/rag/app/audio.py @@ -20,6 +20,7 @@ import tempfile from common.constants import LLMType from api.db.services.llm_service import LLMBundle +from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type from rag.nlp import rag_tokenizer, tokenize @@ -45,7 +46,8 @@ def chunk(filename, binary, tenant_id, lang, callback=None, **kwargs): tmp_path = os.path.abspath(tmpf.name) callback(0.1, "USE Sequence2Txt LLM to transcription the audio") - seq2txt_mdl = LLMBundle(tenant_id, LLMType.SPEECH2TEXT, lang=lang) + seq2txt_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.SPEECH2TEXT) + seq2txt_mdl = LLMBundle(tenant_id, seq2txt_model_config, lang=lang) ans = seq2txt_mdl.transcription(tmp_path) callback(0.8, "Sequence2Txt LLM respond: %s ..." % ans[:32]) diff --git a/rag/app/naive.py b/rag/app/naive.py index e7f6730d80..0620190744 100644 --- a/rag/app/naive.py +++ b/rag/app/naive.py @@ -32,6 +32,7 @@ from common.token_utils import num_tokens_from_string from common.constants import LLMType from api.db.services.llm_service import LLMBundle +from api.db.joint_services.tenant_model_service import get_model_config_by_type_and_name, get_tenant_default_model_by_type from rag.utils.file_utils import extract_embed_file, extract_links_from_pdf, extract_links_from_docx, extract_html from rag.utils.lazy_image import LazyDocxImage from deepdoc.parser import DocxParser, ExcelParser, HtmlParser, JsonParser, MarkdownElementExtractor, MarkdownParser, PdfParser, TxtParser @@ -129,7 +130,8 @@ def by_mineru( if mineru_llm_name: try: - ocr_model = LLMBundle(tenant_id=tenant_id, llm_type=LLMType.OCR, llm_name=mineru_llm_name, lang=lang) + ocr_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.OCR, mineru_llm_name) + ocr_model = LLMBundle(tenant_id=tenant_id, model_config=ocr_model_config, lang=lang) pdf_parser = ocr_model.mdl sections, tables = pdf_parser.parse_pdf( filepath=filename, @@ -208,7 +210,8 @@ def by_paddleocr( if paddleocr_llm_name: try: - ocr_model = LLMBundle(tenant_id=tenant_id, llm_type=LLMType.OCR, llm_name=paddleocr_llm_name, lang=lang) + ocr_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.OCR, paddleocr_llm_name) + ocr_model = LLMBundle(tenant_id=tenant_id, model_config=ocr_model_config, lang=lang) pdf_parser = ocr_model.mdl sections, tables = pdf_parser.parse_pdf( filepath=filename, @@ -236,10 +239,10 @@ def by_plaintext(filename, binary=None, from_page=0, to_page=100000, callback=No tenant_id = kwargs.get("tenant_id") if not tenant_id: raise ValueError("tenant_id is required when using vision layout recognizer") + vision_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.IMAGE2TEXT, layout_recognizer) vision_model = LLMBundle( tenant_id, - LLMType.IMAGE2TEXT, - llm_name=layout_recognizer, + model_config=vision_model_config, lang=kwargs.get("lang", "Chinese"), ) pdf_parser = VisionParser(vision_model=vision_model, **kwargs) @@ -939,7 +942,8 @@ def chunk(filename, binary=None, from_page=0, to_page=100000, lang="Chinese", ca is_markdown = True try: - vision_model = LLMBundle(kwargs["tenant_id"], LLMType.IMAGE2TEXT) + vision_model_config = get_tenant_default_model_by_type(kwargs["tenant_id"], LLMType.IMAGE2TEXT) + vision_model = LLMBundle(kwargs["tenant_id"], vision_model_config) callback(0.2, "Visual model detected. Attempting to enhance figure extraction...") except Exception as e: logging.warning(f"Failed to detect figure extraction: {e}") diff --git a/rag/app/picture.py b/rag/app/picture.py index 67540772ab..d58f923eb8 100644 --- a/rag/app/picture.py +++ b/rag/app/picture.py @@ -22,6 +22,7 @@ import numpy as np from PIL import Image from api.db.services.llm_service import LLMBundle +from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type from common.constants import LLMType from common.string_utils import clean_markdown_block from deepdoc.vision import OCR @@ -50,7 +51,8 @@ def chunk(filename, binary, tenant_id, lang, callback=None, **kwargs): "doc_type_kwd": "video", } ) - cv_mdl = LLMBundle(tenant_id, llm_type=LLMType.IMAGE2TEXT, lang=lang) + cv_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.IMAGE2TEXT) + cv_mdl = LLMBundle(tenant_id, model_config=cv_model_config, lang=lang) video_prompt = str(parser_config.get("video_prompt", "") or "") ans = asyncio.run( cv_mdl.async_chat(system="", history=[], gen_conf={}, video_bytes=binary, filename=filename, video_prompt=video_prompt)) @@ -78,7 +80,8 @@ def chunk(filename, binary, tenant_id, lang, callback=None, **kwargs): try: callback(0.4, "Use CV LLM to describe the picture.") - cv_mdl = LLMBundle(tenant_id, LLMType.IMAGE2TEXT, lang=lang) + cv_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.IMAGE2TEXT) + cv_mdl = LLMBundle(tenant_id, model_config=cv_model_config, lang=lang) with io.BytesIO() as img_binary: img.save(img_binary, format="JPEG") img_binary.seek(0) diff --git a/rag/benchmark.py b/rag/benchmark.py index 93b93adcf3..0e3b256f87 100644 --- a/rag/benchmark.py +++ b/rag/benchmark.py @@ -25,6 +25,7 @@ from common import settings from common.constants import LLMType from api.db.services.llm_service import LLMBundle from api.db.services.knowledgebase_service import KnowledgebaseService +from api.db.joint_services.tenant_model_service import get_model_config_by_id, get_model_config_by_type_and_name from common.misc_utils import get_uuid from rag.nlp import tokenize, search from ranx import evaluate @@ -42,7 +43,11 @@ class Benchmark: e, self.kb = KnowledgebaseService.get_by_id(kb_id) self.similarity_threshold = self.kb.similarity_threshold self.vector_similarity_weight = self.kb.vector_similarity_weight - self.embd_mdl = LLMBundle(self.kb.tenant_id, LLMType.EMBEDDING, llm_name=self.kb.embd_id, lang=self.kb.language) + if self.kb.tenant_embd_id: + embd_model_config = get_model_config_by_id(self.kb.tenant_embd_id) + else: + embd_model_config = get_model_config_by_type_and_name(self.kb.tenant_id, LLMType.EMBEDDING, self.kb.embd_id) + self.embd_mdl = LLMBundle(self.kb.tenant_id, embd_model_config, lang=self.kb.language) self.tenant_id = '' self.index_name = '' self.initialized_index = False diff --git a/rag/flow/parser/parser.py b/rag/flow/parser/parser.py index e42d007d08..9ee6994841 100644 --- a/rag/flow/parser/parser.py +++ b/rag/flow/parser/parser.py @@ -27,6 +27,7 @@ from PIL import Image from api.db.services.file2document_service import File2DocumentService from api.db.services.file_service import FileService from api.db.services.llm_service import LLMBundle +from api.db.joint_services.tenant_model_service import get_model_config_by_type_and_name, get_tenant_default_model_by_type from common import settings from common.constants import LLMType from common.misc_utils import get_uuid @@ -351,7 +352,8 @@ class Parser(ProcessBase): raise RuntimeError("MinerU model not configured. Please add MinerU in Model Providers or set MINERU_* env.") tenant_id = self._canvas._tenant_id - ocr_model = LLMBundle(tenant_id, LLMType.OCR, llm_name=parser_model_name, lang=conf.get("lang", "Chinese")) + ocr_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.OCR, parser_model_name) + ocr_model = LLMBundle(tenant_id, ocr_model_config, lang=conf.get("lang", "Chinese")) pdf_parser = ocr_model.mdl lines, _ = pdf_parser.parse_pdf( @@ -432,7 +434,8 @@ class Parser(ProcessBase): raise RuntimeError("PaddleOCR model not configured. Please add PaddleOCR in Model Providers or set PADDLEOCR_* env.") tenant_id = self._canvas._tenant_id - ocr_model = LLMBundle(tenant_id, LLMType.OCR, llm_name=parser_model_name) + ocr_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.OCR, parser_model_name) + ocr_model = LLMBundle(tenant_id, ocr_model_config) pdf_parser = ocr_model.mdl lines, _ = pdf_parser.parse_pdf( @@ -453,7 +456,11 @@ class Parser(ProcessBase): } bboxes.append(box) else: - vision_model = LLMBundle(self._canvas._tenant_id, LLMType.IMAGE2TEXT, llm_name=conf.get("parse_method"), lang=self._param.setups["pdf"].get("lang")) + if conf.get("parse_method"): + vision_model_config = get_model_config_by_type_and_name(self._canvas._tenant_id, LLMType.IMAGE2TEXT, conf["parse_method"]) + else: + vision_model_config = get_tenant_default_model_by_type(self._canvas._tenant_id, LLMType.IMAGE2TEXT) + vision_model = LLMBundle(self._canvas._tenant_id, vision_model_config, lang=self._param.setups["pdf"].get("lang")) lines, _ = VisionParser(vision_model=vision_model)(blob, callback=self.callback) bboxes = [] for t, poss in lines: @@ -796,7 +803,8 @@ class Parser(ProcessBase): else: lang = conf["lang"] # use VLM to describe the picture - cv_model = LLMBundle(self._canvas.get_tenant_id(), LLMType.IMAGE2TEXT, llm_name=conf["parse_method"], lang=lang) + cv_model_config = get_model_config_by_type_and_name(self._canvas.get_tenant_id(), LLMType.IMAGE2TEXT, conf["parse_method"]) + cv_model = LLMBundle(self._canvas.get_tenant_id(), cv_model_config, lang=lang) img_binary = io.BytesIO() img.save(img_binary, format="JPEG") img_binary.seek(0) @@ -827,8 +835,8 @@ class Parser(ProcessBase): tmpf.write(blob) tmpf.flush() tmp_path = os.path.abspath(tmpf.name) - - seq2txt_mdl = LLMBundle(self._canvas.get_tenant_id(), LLMType.SPEECH2TEXT, llm_name=conf["llm_id"]) + seq2txt_model_config = get_model_config_by_type_and_name(self._canvas.get_tenant_id(), LLMType.SPEECH2TEXT, conf["llm_id"]) + seq2txt_mdl = LLMBundle(self._canvas.get_tenant_id(), seq2txt_model_config) txt = seq2txt_mdl.transcription(tmp_path) self.set_output("text", txt) @@ -838,8 +846,8 @@ class Parser(ProcessBase): conf = self._param.setups["video"] self.set_output("output_format", conf["output_format"]) - - cv_mdl = LLMBundle(self._canvas.get_tenant_id(), LLMType.IMAGE2TEXT, llm_name=conf["llm_id"]) + cv_model_config = get_model_config_by_type_and_name(self._canvas.get_tenant_id(), LLMType.IMAGE2TEXT, conf["llm_id"]) + cv_mdl = LLMBundle(self._canvas.get_tenant_id(), cv_model_config) video_prompt = str(conf.get("prompt", "") or "") txt = asyncio.run(cv_mdl.async_chat(system="", history=[], gen_conf={}, video_bytes=blob, filename=name, video_prompt=video_prompt)) diff --git a/rag/flow/tokenizer/tokenizer.py b/rag/flow/tokenizer/tokenizer.py index 617c3e62a0..dcf4751064 100644 --- a/rag/flow/tokenizer/tokenizer.py +++ b/rag/flow/tokenizer/tokenizer.py @@ -21,7 +21,7 @@ import numpy as np from common.constants import LLMType from api.db.services.knowledgebase_service import KnowledgebaseService from api.db.services.llm_service import LLMBundle -from api.db.services.user_service import TenantService +from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, get_model_config_by_id, get_model_config_by_type_and_name from common.connection_utils import timeout from rag.flow.base import ProcessBase, ProcessParamBase from rag.flow.tokenizer.schema import TokenizerFromUpstream @@ -55,11 +55,13 @@ class Tokenizer(ProcessBase): token_count = 0 if self._canvas._kb_id: e, kb = KnowledgebaseService.get_by_id(self._canvas._kb_id) - embedding_id = kb.embd_id + if kb.tenant_embd_id: + embd_model_config = get_model_config_by_id(kb.tenant_embd_id) + else: + embd_model_config = get_model_config_by_type_and_name(self._canvas._tenant_id, LLMType.EMBEDDING, kb.embd_id) else: - e, ten = TenantService.get_by_id(self._canvas._tenant_id) - embedding_id = ten.embd_id - embedding_model = LLMBundle(self._canvas._tenant_id, LLMType.EMBEDDING, llm_name=embedding_id) + embd_model_config = get_tenant_default_model_by_type(self._canvas._tenant_id, LLMType.EMBEDDING) + embedding_model = LLMBundle(self._canvas._tenant_id, embd_model_config) texts = [] for c in chunks: txt = "" diff --git a/rag/graphrag/general/smoke.py b/rag/graphrag/general/smoke.py index 0070270379..02c1ab5cf6 100644 --- a/rag/graphrag/general/smoke.py +++ b/rag/graphrag/general/smoke.py @@ -24,7 +24,7 @@ from common.constants import LLMType from api.db.services.document_service import DocumentService from api.db.services.knowledgebase_service import KnowledgebaseService from api.db.services.llm_service import LLMBundle -from api.db.services.user_service import TenantService +from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, get_model_config_by_id, get_model_config_by_type_and_name from rag.graphrag.general.graph_extractor import GraphExtractor from rag.graphrag.general.index import update_graph, with_resolution, with_community from common import settings @@ -71,10 +71,14 @@ async def main(): ) ] - _, tenant = TenantService.get_by_id(args.tenant_id) - llm_bdl = LLMBundle(args.tenant_id, LLMType.CHAT, tenant.llm_id) + llm_config = get_tenant_default_model_by_type(args.tenant_id, LLMType.CHAT) + llm_bdl = LLMBundle(args.tenant_id, llm_config) _, kb = KnowledgebaseService.get_by_id(kb_id) - embed_bdl = LLMBundle(args.tenant_id, LLMType.EMBEDDING, kb.embd_id) + if kb.tenant_embd_id: + embd_model_config = get_model_config_by_id(kb.tenant_embd_id) + else: + embd_model_config = get_model_config_by_type_and_name(args.tenant_id, LLMType.EMBEDDING, kb.embd_id) + embed_bdl = LLMBundle(args.tenant_id, embd_model_config) graph, doc_ids = await update_graph( GraphExtractor, diff --git a/rag/graphrag/light/smoke.py b/rag/graphrag/light/smoke.py index 2688e0bb60..18af251518 100644 --- a/rag/graphrag/light/smoke.py +++ b/rag/graphrag/light/smoke.py @@ -24,7 +24,7 @@ from common.constants import LLMType from api.db.services.document_service import DocumentService from api.db.services.knowledgebase_service import KnowledgebaseService from api.db.services.llm_service import LLMBundle -from api.db.services.user_service import TenantService +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 rag.graphrag.general.index import update_graph from rag.graphrag.light.graph_extractor import GraphExtractor from common import settings @@ -72,10 +72,14 @@ async def main(): ) ] - _, tenant = TenantService.get_by_id(args.tenant_id) - llm_bdl = LLMBundle(args.tenant_id, LLMType.CHAT, tenant.llm_id) + llm_config = get_tenant_default_model_by_type(args.tenant_id, LLMType.CHAT) + llm_bdl = LLMBundle(args.tenant_id, llm_config) _, kb = KnowledgebaseService.get_by_id(kb_id) - embed_bdl = LLMBundle(args.tenant_id, LLMType.EMBEDDING, kb.embd_id) + if kb.tenant_embd_id: + embd_model_config = get_model_config_by_id(kb.tenant_embd_id) + else: + embd_model_config = get_model_config_by_type_and_name(args.tenant_id, LLMType.EMBEDDING, kb.embd_id) + embed_bdl = LLMBundle(args.tenant_id, embd_model_config) graph, doc_ids = await update_graph( GraphExtractor, diff --git a/rag/graphrag/search.py b/rag/graphrag/search.py index 6b6ebb82a3..552b64ec99 100644 --- a/rag/graphrag/search.py +++ b/rag/graphrag/search.py @@ -318,7 +318,7 @@ if __name__ == "__main__": from common.constants import LLMType from api.db.services.knowledgebase_service import KnowledgebaseService from api.db.services.llm_service import LLMBundle - from api.db.services.user_service import TenantService + from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, get_model_config_by_id, get_model_config_by_type_and_name from rag.nlp import search settings.init_settings() @@ -329,10 +329,14 @@ if __name__ == "__main__": args = parser.parse_args() kb_id = args.kb_id - _, tenant = TenantService.get_by_id(args.tenant_id) - llm_bdl = LLMBundle(args.tenant_id, LLMType.CHAT, tenant.llm_id) + llm_config = get_tenant_default_model_by_type(args.tenant_id, LLMType.CHAT) + llm_bdl = LLMBundle(args.tenant_id, llm_config) _, kb = KnowledgebaseService.get_by_id(kb_id) - embed_bdl = LLMBundle(args.tenant_id, LLMType.EMBEDDING, kb.embd_id) + if kb.tenant_embd_id: + embd_model_config = get_model_config_by_id(kb.tenant_embd_id) + else: + embd_model_config = get_model_config_by_type_and_name(args.tenant_id, LLMType.EMBEDDING, kb.embd_id) + embed_bdl = LLMBundle(args.tenant_id, embd_model_config) kg = KGSearch(settings.docStoreConn) print(asyncio.run(kg.retrieval({"question": args.question, "kb_ids": [kb_id]}, diff --git a/rag/prompts/generator.py b/rag/prompts/generator.py index 609f2a6bcc..b9c2113d8f 100644 --- a/rag/prompts/generator.py +++ b/rag/prompts/generator.py @@ -225,12 +225,14 @@ async def full_question(tenant_id=None, llm_id=None, messages=[], language=None, from common.constants import LLMType from api.db.services.llm_service import LLMBundle from api.db.services.tenant_llm_service import TenantLLMService + from api.db.joint_services.tenant_model_service import get_model_config_by_type_and_name if not chat_mdl: if TenantLLMService.llm_id2llm_type(llm_id) == "image2text": - chat_mdl = LLMBundle(tenant_id, LLMType.IMAGE2TEXT, llm_id) + chat_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.IMAGE2TEXT, llm_id) else: - chat_mdl = LLMBundle(tenant_id, LLMType.CHAT, llm_id) + chat_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.CHAT, llm_id) + chat_mdl = LLMBundle(tenant_id, chat_model_config) conv = [] for m in messages: if m["role"] not in ["user", "assistant"]: @@ -259,12 +261,13 @@ async def cross_languages(tenant_id, llm_id, query, languages=[]): from common.constants import LLMType from api.db.services.llm_service import LLMBundle from api.db.services.tenant_llm_service import TenantLLMService + from api.db.joint_services.tenant_model_service import get_model_config_by_type_and_name if llm_id and TenantLLMService.llm_id2llm_type(llm_id) == "image2text": - chat_mdl = LLMBundle(tenant_id, LLMType.IMAGE2TEXT, llm_id) + chat_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.IMAGE2TEXT, llm_id) else: - chat_mdl = LLMBundle(tenant_id, LLMType.CHAT, llm_id) - + chat_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.CHAT, llm_id) + chat_mdl = LLMBundle(tenant_id, chat_model_config) rendered_sys_prompt = PROMPT_JINJA_ENV.from_string(CROSS_LANGUAGES_SYS_PROMPT_TEMPLATE).render() rendered_user_prompt = PROMPT_JINJA_ENV.from_string(CROSS_LANGUAGES_USER_PROMPT_TEMPLATE).render(query=query, languages=languages) diff --git a/rag/svr/task_executor.py b/rag/svr/task_executor.py index 4b8a272891..2b3ba46151 100644 --- a/rag/svr/task_executor.py +++ b/rag/svr/task_executor.py @@ -65,6 +65,7 @@ from api.db.services.doc_metadata_service import DocMetadataService from api.db.services.llm_service import LLMBundle from api.db.services.task_service import TaskService, has_canceled, CANVAS_DEBUG_DOC_ID, GRAPH_RAPTOR_FAKE_DOC_ID from api.db.services.file2document_service import File2DocumentService +from api.db.joint_services.tenant_model_service import get_model_config_by_type_and_name, get_tenant_default_model_by_type from common.versions import get_ragflow_version from api.db.db_models import close_connection from rag.app import laws, paper, presentation, manual, qa, table, book, resume, picture, naive, one, audio, \ @@ -342,7 +343,8 @@ async def build_chunks(task, progress_callback): if task["parser_config"].get("auto_keywords", 0): st = timer() progress_callback(msg="Start to generate keywords for every chunk ...") - chat_mdl = LLMBundle(task["tenant_id"], LLMType.CHAT, llm_name=task["llm_id"], lang=task["language"]) + chat_model_config = get_model_config_by_type_and_name(task["tenant_id"], LLMType.CHAT, task["llm_id"]) + chat_mdl = LLMBundle(task["tenant_id"], chat_model_config, lang=task["language"]) async def doc_keyword_extraction(chat_mdl, d, topn): cached = get_llm_cache(chat_mdl.llm_name, d["content_with_weight"], "keywords", {"topn": topn}) @@ -375,7 +377,8 @@ async def build_chunks(task, progress_callback): if task["parser_config"].get("auto_questions", 0): st = timer() progress_callback(msg="Start to generate questions for every chunk ...") - chat_mdl = LLMBundle(task["tenant_id"], LLMType.CHAT, llm_name=task["llm_id"], lang=task["language"]) + chat_model_config = get_model_config_by_type_and_name(task["tenant_id"], LLMType.CHAT, task["llm_id"]) + chat_mdl = LLMBundle(task["tenant_id"], chat_model_config, lang=task["language"]) async def doc_question_proposal(chat_mdl, d, topn): cached = get_llm_cache(chat_mdl.llm_name, d["content_with_weight"], "question", {"topn": topn}) @@ -407,7 +410,8 @@ async def build_chunks(task, progress_callback): if task["parser_config"].get("enable_metadata", False) and task["parser_config"].get("metadata"): st = timer() progress_callback(msg="Start to generate meta-data for every chunk ...") - chat_mdl = LLMBundle(task["tenant_id"], LLMType.CHAT, llm_name=task["llm_id"], lang=task["language"]) + chat_model_config = get_model_config_by_type_and_name(task["tenant_id"], LLMType.CHAT, task["llm_id"]) + chat_mdl = LLMBundle(task["tenant_id"], chat_model_config, lang=task["language"]) async def gen_metadata_task(chat_mdl, d): cached = get_llm_cache(chat_mdl.llm_name, d["content_with_weight"], "metadata", @@ -461,8 +465,8 @@ async def build_chunks(task, progress_callback): set_tags_to_cache(kb_ids, all_tags) else: all_tags = json.loads(all_tags) - - chat_mdl = LLMBundle(task["tenant_id"], LLMType.CHAT, llm_name=task["llm_id"], lang=task["language"]) + chat_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.CHAT, task["llm_id"]) + chat_mdl = LLMBundle(task["tenant_id"], chat_model_config, lang=task["language"]) docs_to_tag = [] for d in docs: @@ -517,7 +521,8 @@ async def build_chunks(task, progress_callback): def build_TOC(task, docs, progress_callback): progress_callback(msg="Start to generate table of content ...") - chat_mdl = LLMBundle(task["tenant_id"], LLMType.CHAT, llm_name=task["llm_id"], lang=task["language"]) + chat_model_config = get_model_config_by_type_and_name(task["tenant_id"], LLMType.CHAT, task["llm_id"]) + chat_mdl = LLMBundle(task["tenant_id"], chat_model_config, lang=task["language"]) docs = sorted(docs, key=lambda d: ( d.get("page_num_int", 0)[0] if isinstance(d.get("page_num_int", 0), list) else d.get("page_num_int", 0), d.get("top_int", 0)[0] if isinstance(d.get("top_int", 0), list) else d.get("top_int", 0) @@ -668,7 +673,8 @@ async def run_dataflow(task: dict): set_progress(task_id, prog=0.82, msg="\n-------------------------------------\nStart to embedding...") e, kb = KnowledgebaseService.get_by_id(task["kb_id"]) embedding_id = kb.embd_id - embedding_model = LLMBundle(task["tenant_id"], LLMType.EMBEDDING, llm_name=embedding_id) + embd_model_config = get_model_config_by_type_and_name(task["tenant_id"], LLMType.EMBEDDING, embedding_id) + embedding_model = LLMBundle(task["tenant_id"], embd_model_config) @timeout(60) def batch_encode(txts): @@ -985,7 +991,11 @@ async def do_handle_task(task): try: # bind embedding model - embedding_model = LLMBundle(task_tenant_id, LLMType.EMBEDDING, llm_name=task_embedding_id, lang=task_language) + if task_embedding_id: + embd_model_config = get_model_config_by_type_and_name(task_tenant_id, LLMType.EMBEDDING, task_embedding_id) + else: + embd_model_config = get_tenant_default_model_by_type(task_tenant_id, LLMType.EMBEDDING) + embedding_model = LLMBundle(task_tenant_id, embd_model_config, lang=task_language) vts, _ = embedding_model.encode(["ok"]) vector_size = len(vts[0]) except Exception as e: @@ -1037,7 +1047,8 @@ async def do_handle_task(task): return # bind LLM for raptor - chat_model = LLMBundle(task_tenant_id, LLMType.CHAT, llm_name=kb_task_llm_id, lang=task_language) + chat_model_config = get_model_config_by_type_and_name(task_dataset_id, LLMType.CHAT, kb_task_llm_id) + chat_model = LLMBundle(task_tenant_id, chat_model_config, lang=task_language) # run RAPTOR async with kg_limiter: chunks, token_count = await run_raptor_for_kb( @@ -1081,7 +1092,8 @@ async def do_handle_task(task): graphrag_conf = kb_parser_config.get("graphrag", {}) start_ts = timer() - chat_model = LLMBundle(task_tenant_id, LLMType.CHAT, llm_name=kb_task_llm_id, lang=task_language) + chat_model_config = get_model_config_by_type_and_name(task_tenant_id, LLMType.CHAT, kb_task_llm_id) + chat_model = LLMBundle(task_tenant_id, chat_model_config, lang=task_language) with_resolution = graphrag_conf.get("resolution", False) with_community = graphrag_conf.get("community", False) async with kg_limiter: diff --git a/test/testcases/test_http_api/test_chat_assistant_management/test_chat_sdk_routes_unit.py b/test/testcases/test_http_api/test_chat_assistant_management/test_chat_sdk_routes_unit.py index 63fef2de89..cb3ca0ae82 100644 --- a/test/testcases/test_http_api/test_chat_assistant_management/test_chat_sdk_routes_unit.py +++ b/test/testcases/test_http_api/test_chat_assistant_management/test_chat_sdk_routes_unit.py @@ -44,9 +44,10 @@ class _AwaitableValue: class _DummyKB: - def __init__(self, embd_id="embd@factory", chunk_num=1): + def __init__(self, embd_id="embd@factory", chunk_num=1, tenant_embd_id=1): self.embd_id = embd_id self.chunk_num = chunk_num + self.tenant_embd_id = tenant_embd_id def to_json(self): return {"id": "kb-1"} diff --git a/test/testcases/test_http_api/test_dataset_management/test_dify_retrieval_routes_unit.py b/test/testcases/test_http_api/test_dataset_management/test_dify_retrieval_routes_unit.py index 38ec8d9fa0..b7c625cbe7 100644 --- a/test/testcases/test_http_api/test_dataset_management/test_dify_retrieval_routes_unit.py +++ b/test/testcases/test_http_api/test_dataset_management/test_dify_retrieval_routes_unit.py @@ -45,9 +45,10 @@ class _AwaitableValue: class _DummyKB: - def __init__(self, tenant_id="tenant-1", embd_id="embd-1"): + def __init__(self, tenant_id="tenant-1", embd_id="embd-1", tenant_embd_id=1): self.tenant_id = tenant_id self.embd_id = embd_id + self.tenant_embd_id = tenant_embd_id class _DummyRetriever: @@ -102,6 +103,138 @@ def _load_dify_retrieval_module(monkeypatch): monkeypatch.setitem(sys.modules, "deepdoc.parser.utils", deepdoc_parser_utils) monkeypatch.setitem(sys.modules, "xgboost", ModuleType("xgboost")) + # Mock tenant_llm_service for TenantLLMService and TenantService + tenant_llm_service_mod = ModuleType("api.db.services.tenant_llm_service") + + class _MockModelConfig: + def __init__(self, tenant_id, model_name): + self.tenant_id = tenant_id + self.llm_name = model_name + self.llm_factory = "Builtin" + self.api_key = "fake-api-key" + self.api_base = "https://api.example.com" + self.model_type = "chat" + self.max_tokens = 8192 + self.used_tokens = 0 + self.status = 1 + self.id = 1 + + def to_dict(self): + return { + "tenant_id": self.tenant_id, + "llm_name": self.llm_name, + "llm_factory": self.llm_factory, + "api_key": self.api_key, + "api_base": self.api_base, + "model_type": self.model_type, + "max_tokens": self.max_tokens, + "used_tokens": self.used_tokens, + "status": self.status, + "id": self.id + } + + class _StubTenantService: + @staticmethod + def get_by_id(tenant_id): + # Return a mock tenant with default model configurations + return True, SimpleNamespace( + id=tenant_id, + llm_id="chat-model", + embd_id="embd-model", + asr_id="asr-model", + img2txt_id="img2txt-model", + rerank_id="rerank-model", + tts_id="tts-model" + ) + + class _StubTenantLLMService: + @staticmethod + def get_api_key(tenant_id, model_name): + return _MockModelConfig(tenant_id, model_name) + + @staticmethod + def split_model_name_and_factory(model_name): + if "@" in model_name: + parts = model_name.split("@") + return parts[0], parts[1] + return model_name, None + + tenant_llm_service_mod.TenantService = _StubTenantService + tenant_llm_service_mod.TenantLLMService = _StubTenantLLMService + monkeypatch.setitem(sys.modules, "api.db.services.tenant_llm_service", tenant_llm_service_mod) + + # Mock llm_service for LLMService + llm_service_mod = ModuleType("api.db.services.llm_service") + + class _StubLLM: + def __init__(self, llm_name): + self.llm_name = llm_name + self.is_tools = False + + class _StubLLMBundle: + def __init__(self, tenant_id: str, model_config: dict, lang="Chinese", **kwargs): + self.tenant_id = tenant_id + self.model_config = model_config + self.lang = lang + + def encode(self, texts: list): + import numpy as np + # Return mock embeddings and token usage + return [np.array([0.1, 0.2, 0.3]) for _ in texts], len(texts) * 10 + + llm_service_mod.LLMService = SimpleNamespace( + query=lambda llm_name: [_StubLLM(llm_name)] if llm_name else [] + ) + llm_service_mod.LLMBundle = _StubLLMBundle + monkeypatch.setitem(sys.modules, "api.db.services.llm_service", llm_service_mod) + + # Mock tenant_model_service to ensure it uses mocked services + tenant_model_service_mod = ModuleType("api.db.joint_services.tenant_model_service") + + class _MockModelConfig2: + def __init__(self, tenant_id, model_name): + self.tenant_id = tenant_id + self.llm_name = model_name + self.llm_factory = "Builtin" + self.api_key = "fake-api-key" + self.api_base = "https://api.example.com" + self.model_type = "chat" + self.max_tokens = 8192 + self.used_tokens = 0 + self.status = 1 + self.id = 1 + + def to_dict(self): + return { + "tenant_id": self.tenant_id, + "llm_name": self.llm_name, + "llm_factory": self.llm_factory, + "api_key": self.api_key, + "api_base": self.api_base, + "model_type": self.model_type, + "max_tokens": self.max_tokens, + "used_tokens": self.used_tokens, + "status": self.status, + "id": self.id + } + + def _get_model_config_by_id(tenant_model_id: int) -> dict: + return _MockModelConfig2("tenant-1", "model-1").to_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") + return _MockModelConfig2(tenant_id, model_name).to_dict() + + def _get_tenant_default_model_by_type(tenant_id: str, model_type): + # Return mock tenant with default model configurations + return _MockModelConfig2(tenant_id, "chat-model").to_dict() + + tenant_model_service_mod.get_model_config_by_id = _get_model_config_by_id + tenant_model_service_mod.get_model_config_by_type_and_name = _get_model_config_by_type_and_name + tenant_model_service_mod.get_tenant_default_model_by_type = _get_tenant_default_model_by_type + monkeypatch.setitem(sys.modules, "api.db.joint_services.tenant_model_service", tenant_model_service_mod) + module_name = "test_dify_retrieval_routes_unit_module" module_path = repo_root / "api" / "apps" / "sdk" / "dify_retrieval.py" spec = importlib.util.spec_from_file_location(module_name, module_path) @@ -134,7 +267,6 @@ def test_retrieval_success_with_metadata_and_kg(monkeypatch): monkeypatch.setattr(module, "jsonify", lambda payload: payload) monkeypatch.setattr(module.DocMetadataService, "get_meta_by_kbs", lambda _kb_ids: [{"doc_id": "doc-1"}]) monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, _DummyKB())) - monkeypatch.setattr(module, "LLMBundle", lambda *_args, **_kwargs: object()) monkeypatch.setattr(module, "convert_conditions", lambda cond: cond.get("conditions", [])) monkeypatch.setattr(module, "meta_filter", lambda *_args, **_kwargs: []) @@ -185,7 +317,6 @@ def test_retrieval_not_found_exception_mapping(monkeypatch): _set_request_json(monkeypatch, module, {"knowledge_id": "kb-1", "query": "hello"}) monkeypatch.setattr(module.DocMetadataService, "get_meta_by_kbs", lambda _kb_ids: []) monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, _DummyKB())) - monkeypatch.setattr(module, "LLMBundle", lambda *_args, **_kwargs: object()) monkeypatch.setattr(module, "label_question", lambda *_args, **_kwargs: []) class _BrokenRetriever: @@ -205,7 +336,6 @@ def test_retrieval_generic_exception_mapping(monkeypatch): _set_request_json(monkeypatch, module, {"knowledge_id": "kb-1", "query": "hello"}) monkeypatch.setattr(module.DocMetadataService, "get_meta_by_kbs", lambda _kb_ids: []) monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, _DummyKB())) - monkeypatch.setattr(module, "LLMBundle", lambda *_args, **_kwargs: object()) monkeypatch.setattr(module, "label_question", lambda *_args, **_kwargs: []) class _BrokenRetriever: diff --git a/test/testcases/test_http_api/test_file_management_within_dataset/test_doc_sdk_routes_unit.py b/test/testcases/test_http_api/test_file_management_within_dataset/test_doc_sdk_routes_unit.py index 635d09b555..23ac8fcf67 100644 --- a/test/testcases/test_http_api/test_file_management_within_dataset/test_doc_sdk_routes_unit.py +++ b/test/testcases/test_http_api/test_file_management_within_dataset/test_doc_sdk_routes_unit.py @@ -151,6 +151,149 @@ def _load_doc_module(monkeypatch): monkeypatch.setitem(sys.modules, "deepdoc.parser.utils", deepdoc_parser_utils) monkeypatch.setitem(sys.modules, "xgboost", ModuleType("xgboost")) + # Mock tenant_llm_service for TenantLLMService and TenantService + tenant_llm_service_mod = ModuleType("api.db.services.tenant_llm_service") + + class _MockModelConfig: + def __init__(self, tenant_id, model_name): + self.tenant_id = tenant_id + self.llm_name = model_name + self.llm_factory = "Builtin" + self.api_key = "fake-api-key" + self.api_base = "https://api.example.com" + self.model_type = "embedding" + self.max_tokens = 8192 + self.used_tokens = 0 + self.status = 1 + self.id = 1 + + def to_dict(self): + return { + "tenant_id": self.tenant_id, + "llm_name": self.llm_name, + "llm_factory": self.llm_factory, + "api_key": self.api_key, + "api_base": self.api_base, + "model_type": self.model_type, + "max_tokens": self.max_tokens, + "used_tokens": self.used_tokens, + "status": self.status, + "id": self.id + } + + class _StubTenantService: + @staticmethod + def get_by_id(tenant_id): + return True, SimpleNamespace( + id=tenant_id, + llm_id="chat-model", + embd_id="embd-model", + asr_id="asr-model", + img2txt_id="img2txt-model", + rerank_id="rerank-model", + tts_id="tts-model" + ) + + class _StubTenantLLMService: + @staticmethod + def get_api_key(tenant_id, model_name): + return _MockModelConfig(tenant_id, model_name) + + @staticmethod + def split_model_name_and_factory(model_name): + if "@" in model_name: + parts = model_name.split("@") + return parts[0], parts[1] + return model_name, None + + @staticmethod + def get_by_id(tenant_model_id): + return True, _MockModelConfig("tenant-1", "model-1") + + @staticmethod + def model_instance(model_config): + class _EmbedModel: + def encode(self, texts): + import numpy as np + return [np.array([0.2, 0.8]), np.array([0.3, 0.7])], 1 + return _EmbedModel() + + tenant_llm_service_mod.TenantService = _StubTenantService + tenant_llm_service_mod.TenantLLMService = _StubTenantLLMService + monkeypatch.setitem(sys.modules, "api.db.services.tenant_llm_service", tenant_llm_service_mod) + + # Mock LLMService + llm_service_mod = ModuleType("api.db.services.llm_service") + + class _StubLLM: + def __init__(self, llm_name): + self.llm_name = llm_name + self.is_tools = False + + class _StubLLMBundle: + def __init__(self, tenant_id: str, model_config: dict, lang="Chinese", **kwargs): + self.tenant_id = tenant_id + self.model_config = model_config + self.lang = lang + + def encode(self, texts: list): + import numpy as np + # Return mock embeddings and token usage + return [np.array([0.2, 0.8]), np.array([0.3, 0.7])], len(texts) * 10 + + llm_service_mod.LLMService = SimpleNamespace( + query=lambda llm_name: [_StubLLM(llm_name)] if llm_name else [] + ) + llm_service_mod.LLMBundle = _StubLLMBundle + monkeypatch.setitem(sys.modules, "api.db.services.llm_service", llm_service_mod) + + # Mock tenant_model_service to ensure it uses mocked services + tenant_model_service_mod = ModuleType("api.db.joint_services.tenant_model_service") + + class _MockModelConfig2: + def __init__(self, tenant_id, model_name): + self.tenant_id = tenant_id + self.llm_name = model_name + self.llm_factory = "Builtin" + self.api_key = "fake-api-key" + self.api_base = "https://api.example.com" + self.model_type = "embedding" + self.max_tokens = 8192 + self.used_tokens = 0 + self.status = 1 + self.id = 1 + + def to_dict(self): + return { + "tenant_id": self.tenant_id, + "llm_name": self.llm_name, + "llm_factory": self.llm_factory, + "api_key": self.api_key, + "api_base": self.api_base, + "model_type": self.model_type, + "max_tokens": self.max_tokens, + "used_tokens": self.used_tokens, + "status": self.status, + "id": self.id + } + + def _get_model_config_by_id(tenant_model_id: int) -> dict: + return _MockModelConfig2("tenant-1", "model-1").to_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") + return _MockModelConfig2(tenant_id, model_name).to_dict() + + def _get_tenant_default_model_by_type(tenant_id: str, model_type): + # Return mock tenant with default model configurations + return _MockModelConfig2(tenant_id, "chat-model").to_dict() + + tenant_model_service_mod.get_model_config_by_id = _get_model_config_by_id + tenant_model_service_mod.get_model_config_by_type_and_name = _get_model_config_by_type_and_name + tenant_model_service_mod.get_tenant_default_model_by_type = _get_tenant_default_model_by_type + monkeypatch.setitem(sys.modules, "api.db.joint_services.tenant_model_service", tenant_model_service_mod) + module_path = repo_root / "api" / "apps" / "sdk" / "doc.py" spec = importlib.util.spec_from_file_location("test_doc_sdk_routes_unit", module_path) module = importlib.util.module_from_spec(spec) @@ -762,6 +905,7 @@ class TestDocRoutesUnit: monkeypatch.setattr(module.rag_tokenizer, "fine_grained_tokenize", lambda text: text or "") monkeypatch.setattr(module.rag_tokenizer, "is_chinese", lambda _text: False) monkeypatch.setattr(module.DocumentService, "get_embd_id", lambda _doc_id: "embd") + monkeypatch.setattr(module.DocumentService, "get_tenant_embd_id", lambda _doc_id: 1) class _EmbedModel: def encode(self, _texts): @@ -851,8 +995,8 @@ class TestDocRoutesUnit: "get_request_json", lambda: _AwaitableValue({"dataset_ids": ["ds-1"], "question": "q", "highlight": "True"}), ) - monkeypatch.setattr(module.KnowledgebaseService, "get_by_ids", lambda _ids: [SimpleNamespace(embd_id="m1", tenant_id="tenant-1")]) - monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _id: (True, SimpleNamespace(tenant_id="tenant-1", embd_id="m1"))) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_ids", lambda _ids: [SimpleNamespace(embd_id="m1", tenant_id="tenant-1", tenant_embd_id=1)]) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _id: (True, SimpleNamespace(tenant_id="tenant-1", embd_id="m1", tenant_embd_id=1))) class _Retriever: async def retrieval(self, *_args, **_kwargs): @@ -865,7 +1009,7 @@ class TestDocRoutesUnit: monkeypatch.setattr(module, "label_question", lambda *_args, **_kwargs: {}) monkeypatch.setattr(module.settings, "retriever", _Retriever()) res = _run(module.retrieval_test.__wrapped__("tenant-1")) - assert res["code"] == 0 + assert res["code"] == 0, res["message"] assert res["data"]["chunks"] == [] monkeypatch.setattr( @@ -974,7 +1118,7 @@ class TestDocRoutesUnit: } ), ) - monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _id: (True, SimpleNamespace(tenant_id="tenant-1", embd_id="m1"))) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _id: (True, SimpleNamespace(tenant_id="tenant-1", embd_id="m1", tenant_embd_id=1))) monkeypatch.setattr(module, "cross_languages", _cross_languages) monkeypatch.setattr(module, "keyword_extraction", _keyword_extraction) monkeypatch.setattr(module.settings, "retriever", _FeatureRetriever()) @@ -982,7 +1126,7 @@ class TestDocRoutesUnit: monkeypatch.setattr(module, "label_question", lambda *_args, **_kwargs: {}) monkeypatch.setattr(module, "LLMBundle", lambda *_args, **_kwargs: SimpleNamespace()) res = _run(module.retrieval_test.__wrapped__("tenant-1")) - assert res["code"] == 0 + assert res["code"] == 0, res["message"] assert feature_calls["cross"] == ("fr",) assert feature_calls["keyword"] == "q-xl" assert feature_calls["retrieval_question"] == "q-xl-kw" diff --git a/test/testcases/test_http_api/test_session_management/test_session_sdk_routes_unit.py b/test/testcases/test_http_api/test_session_management/test_session_sdk_routes_unit.py index 9d0e6c2ddb..c1bbd10347 100644 --- a/test/testcases/test_http_api/test_session_management/test_session_sdk_routes_unit.py +++ b/test/testcases/test_http_api/test_session_management/test_session_sdk_routes_unit.py @@ -121,6 +121,132 @@ def _load_session_module(monkeypatch): common_pkg.__path__ = [str(repo_root / "common")] monkeypatch.setitem(sys.modules, "common", common_pkg) + # Mock common.constants module + from enum import Enum + from strenum import StrEnum + + class _StubLLMType(StrEnum): + CHAT = "chat" + EMBEDDING = "embedding" + SPEECH2TEXT = "speech2text" + IMAGE2TEXT = "image2text" + RERANK = "rerank" + TTS = "tts" + OCR = "ocr" + + class _StubParserType(StrEnum): + PRESENTATION = "presentation" + LAWS = "laws" + MANUAL = "manual" + PAPER = "paper" + RESUME = "resume" + BOOK = "book" + QA = "qa" + TABLE = "table" + NAIVE = "naive" + PICTURE = "picture" + ONE = "one" + AUDIO = "audio" + EMAIL = "email" + KG = "knowledge_graph" + TAG = "tag" + + class _StubRetCode(int, Enum): + SUCCESS = 0 + NOT_EFFECTIVE = 10 + EXCEPTION_ERROR = 100 + ARGUMENT_ERROR = 101 + DATA_ERROR = 102 + OPERATING_ERROR = 103 + CONNECTION_ERROR = 105 + RUNNING = 106 + PERMISSION_ERROR = 108 + AUTHENTICATION_ERROR = 109 + BAD_REQUEST = 400 + UNAUTHORIZED = 401 + SERVER_ERROR = 500 + FORBIDDEN = 403 + NOT_FOUND = 404 + CONFLICT = 409 + + class _StubStatusEnum(str, Enum): + VALID = "1" + INVALID = "0" + + class _StubActiveEnum(Enum): + ACTIVE = "1" + INACTIVE = "0" + + class _StubStorage(Enum): + MINIO = 1 + AZURE_SPN = 2 + AZURE_SAS = 3 + AWS_S3 = 4 + OSS = 5 + OPENDAL = 6 + GCS = 7 + + class _StubMCPServerType(StrEnum): + SSE = "sse" + STREAMABLE_HTTP = "streamable-http" + + class _StubTaskStatus(StrEnum): + UNSTART = "0" + RUNNING = "1" + CANCEL = "2" + DONE = "3" + FAIL = "4" + SCHEDULE = "5" + + class _StubFileSource(StrEnum): + LOCAL = "" + KNOWLEDGEBASE = "knowledgebase" + S3 = "s3" + NOTION = "notion" + DISCORD = "discord" + CONFLUENCE = "confluence" + GMAIL = "gmail" + GOOGLE_DRIVE = "google_drive" + JIRA = "jira" + SHAREPOINT = "sharepoint" + SLACK = "slack" + TEAMS = "teams" + WEBDAV = "webdav" + MOODLE = "moodle" + DROPBOX = "dropbox" + BOX = "box" + R2 = "r2" + OCI_STORAGE = "oci_storage" + GOOGLE_CLOUD_STORAGE = "google_cloud_storage" + AIRTABLE = "airtable" + ASANA = "asana" + GITHUB = "github" + GITLAB = "gitlab" + IMAP = "imap" + BITBUCKET = "bitbucket" + ZENDESK = "zendesk" + SEAFILE = "seafile" + MYSQL = "mysql" + POSTGRESQL = "postgresql" + + common_constants_mod = ModuleType("common.constants") + common_constants_mod.LLMType = _StubLLMType + common_constants_mod.ParserType = _StubParserType + common_constants_mod.RetCode = _StubRetCode + common_constants_mod.StatusEnum = _StubStatusEnum + common_constants_mod.ActiveEnum = _StubActiveEnum + common_constants_mod.Storage = _StubStorage + common_constants_mod.MCPServerType = _StubMCPServerType + common_constants_mod.TaskStatus = _StubTaskStatus + common_constants_mod.FileSource = _StubFileSource + common_constants_mod.SERVICE_CONF = "service_conf.yaml" + common_constants_mod.RAG_FLOW_SERVICE_NAME = "ragflow" + common_constants_mod.SVR_QUEUE_NAME = "rag_flow_svr_queue" + common_constants_mod.SVR_CONSUMER_GROUP_NAME = "rag_flow_svr_task_broker" + common_constants_mod.PAGERANK_FLD = "pagerank_fea" + common_constants_mod.TAG_FLD = "tag_feas" + monkeypatch.setitem(sys.modules, "common.constants", common_constants_mod) + deepdoc_pkg = ModuleType("deepdoc") deepdoc_parser_pkg = ModuleType("deepdoc.parser") deepdoc_parser_pkg.__path__ = [] @@ -166,6 +292,180 @@ def _load_session_module(monkeypatch): monkeypatch.setitem(sys.modules, "deepdoc.parser.utils", deepdoc_parser_utils) monkeypatch.setitem(sys.modules, "xgboost", ModuleType("xgboost")) + # Mock tenant_llm_service for TenantLLMService and TenantService + tenant_llm_service_mod = ModuleType("api.db.services.tenant_llm_service") + + class _MockModelConfig: + def __init__(self, tenant_id, model_name): + self.tenant_id = tenant_id + self.llm_name = model_name + self.llm_factory = "Builtin" + self.api_key = "fake-api-key" + self.api_base = "https://api.example.com" + self.model_type = "chat" + self.max_tokens = 8192 + self.used_tokens = 0 + self.status = 1 + self.id = 1 + + def to_dict(self): + return { + "tenant_id": self.tenant_id, + "llm_name": self.llm_name, + "llm_factory": self.llm_factory, + "api_key": self.api_key, + "api_base": self.api_base, + "model_type": self.model_type, + "max_tokens": self.max_tokens, + "used_tokens": self.used_tokens, + "status": self.status, + "id": self.id + } + + class _StubTenantService: + @staticmethod + def get_by_id(tenant_id): + # Return a mock tenant with default model configurations + return True, SimpleNamespace( + id=tenant_id, + llm_id="chat-model", + embd_id="embd-model", + asr_id="asr-model", + img2txt_id="img2txt-model", + rerank_id="rerank-model", + tts_id="tts-model" + ) + + class _StubTenantLLMService: + @staticmethod + def get_api_key(tenant_id, model_name): + return _MockModelConfig(tenant_id, model_name) + + @staticmethod + def split_model_name_and_factory(model_name): + if "@" in model_name: + parts = model_name.split("@") + return parts[0], parts[1] + return model_name, None + + class _StubLLMFactoriesService: + @staticmethod + def query(**_kwargs): + return [] + + tenant_llm_service_mod.TenantService = _StubTenantService + tenant_llm_service_mod.TenantLLMService = _StubTenantLLMService + tenant_llm_service_mod.LLMFactoriesService = _StubLLMFactoriesService + monkeypatch.setitem(sys.modules, "api.db.services.tenant_llm_service", tenant_llm_service_mod) + + # Mock LLMService + llm_service_mod = ModuleType("api.db.services.llm_service") + + class _StubLLM: + def __init__(self, llm_name): + self.llm_name = llm_name + self.is_tools = False + + llm_service_mod.LLMService = SimpleNamespace( + query=lambda llm_name: [_StubLLM(llm_name)] if llm_name else [] + ) + + class _StubLLMBundle: + def __init__(self, tenant_id: str, model_config: dict, lang="Chinese", **kwargs): + self.tenant_id = tenant_id + self.model_config = model_config + self.lang = lang + + async def async_chat(self, prompt, messages, options): + return "mock response" + + def transcription(self, audio_path): + return "mock transcription" + + llm_service_mod.LLMBundle = _StubLLMBundle + monkeypatch.setitem(sys.modules, "api.db.services.llm_service", llm_service_mod) + + # Mock tenant_model_service to ensure it uses mocked services + tenant_model_service_mod = ModuleType("api.db.joint_services.tenant_model_service") + + class _MockModelConfig2: + def __init__(self, tenant_id, model_name, model_type="chat"): + self.tenant_id = tenant_id + self.llm_name = model_name + self.llm_factory = "Builtin" + self.api_key = "fake-api-key" + self.api_base = "https://api.example.com" + self.model_type = model_type + self.max_tokens = 8192 + self.used_tokens = 0 + self.status = 1 + self.id = 1 + + def to_dict(self): + return { + "tenant_id": self.tenant_id, + "llm_name": self.llm_name, + "llm_factory": self.llm_factory, + "api_key": self.api_key, + "api_base": self.api_base, + "model_type": self.model_type, + "max_tokens": self.max_tokens, + "used_tokens": self.used_tokens, + "status": self.status, + "id": self.id + } + + def _get_model_config_by_id(tenant_model_id: int) -> dict: + return _MockModelConfig2("tenant-1", "model-1").to_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") + return _MockModelConfig2(tenant_id, model_name, model_type).to_dict() + + def _get_tenant_default_model_by_type(tenant_id: str, model_type): + # Check if tenant exists + from api.db.services.tenant_llm_service import TenantService + exist, tenant = TenantService.get_by_id(tenant_id) + if not exist: + raise LookupError("Tenant not found!") + # Return mock tenant with default model configurations + model_type_val = model_type if isinstance(model_type, str) else model_type.value + model_name = "" + if model_type_val == "embedding": + model_name = tenant.embd_id + elif model_type_val == "speech2text": + model_name = tenant.asr_id + elif model_type_val == "image2text": + model_name = tenant.img2txt_id + elif model_type_val == "chat": + model_name = tenant.llm_id + elif model_type_val == "rerank": + model_name = tenant.rerank_id + elif model_type_val == "tts": + model_name = tenant.tts_id + elif model_type_val == "ocr": + raise Exception("OCR model name is required") + if not model_name: + # Use friendly model type names + friendly_names = { + "embedding": "Embedding", + "speech2text": "ASR", + "image2text": "Image2Text", + "chat": "Chat", + "rerank": "Rerank", + "tts": "TTS", + "ocr": "OCR" + } + friendly_name = friendly_names.get(model_type_val, model_type_val) + raise Exception(f"No default {friendly_name} model is set") + return _MockModelConfig2(tenant_id, model_name, model_type_val).to_dict() + + tenant_model_service_mod.get_model_config_by_id = _get_model_config_by_id + tenant_model_service_mod.get_model_config_by_type_and_name = _get_model_config_by_type_and_name + tenant_model_service_mod.get_tenant_default_model_by_type = _get_tenant_default_model_by_type + monkeypatch.setitem(sys.modules, "api.db.joint_services.tenant_model_service", tenant_model_service_mod) + agent_pkg = ModuleType("agent") agent_pkg.__path__ = [] agent_canvas_mod = ModuleType("agent.canvas") @@ -200,6 +500,29 @@ def _load_session_module(monkeypatch): module.manager = _DummyManager() monkeypatch.setitem(sys.modules, "test_session_sdk_routes_unit_module", module) spec.loader.exec_module(module) + + # Add TenantService to module for test compatibility + class _StubTenantServiceForTest: + @staticmethod + def get_info_by(tenant_id): + # Return mock tenant info for tests + return [] + + @staticmethod + def get_by_id(tenant_id): + # Return mock tenant by id + return True, SimpleNamespace( + id=tenant_id, + llm_id="chat-model", + embd_id="embd-model", + asr_id="asr-model", + img2txt_id="img2txt-model", + rerank_id="rerank-model", + tts_id="tts-model" + ) + + module.TenantService = _StubTenantServiceForTest + return module @@ -1028,9 +1351,12 @@ def test_searchbots_retrieval_test_embedded_matrix_unit(monkeypatch): llm_calls = [] - def _fake_llm_bundle(tenant_id, llm_type, *args, **kwargs): - llm_calls.append((tenant_id, llm_type, args, kwargs)) - return SimpleNamespace(tenant_id=tenant_id, llm_type=llm_type, args=args, kwargs=kwargs) + def _fake_llm_bundle(tenant_id, model_config, *args, **kwargs): + # Extract llm_type from model_config for comparison + llm_type = model_config.get("model_type") if isinstance(model_config, dict) else model_config + llm_name = model_config.get("llm_name") if isinstance(model_config, dict) else None + llm_calls.append((tenant_id, llm_type, llm_name, args, kwargs)) + return SimpleNamespace(tenant_id=tenant_id, llm_type=llm_type, llm_name=llm_name, args=args, kwargs=kwargs) monkeypatch.setattr(module, "LLMBundle", _fake_llm_bundle) monkeypatch.setattr( @@ -1128,7 +1454,7 @@ def test_searchbots_retrieval_test_embedded_matrix_unit(monkeypatch): monkeypatch.setattr( module.KnowledgebaseService, "get_by_id", - lambda _kb_id: (True, SimpleNamespace(tenant_id="tenant-kb", embd_id="embd-model")), + lambda _kb_id: (True, SimpleNamespace(tenant_id="tenant-kb", embd_id="embd-model", tenant_embd_id=None)), ) res = _run(handler()) assert res["code"] == 0 @@ -1143,7 +1469,7 @@ def test_searchbots_retrieval_test_embedded_matrix_unit(monkeypatch): assert retrieval_capture["local_doc_ids"] == ["doc-filtered"] assert retrieval_capture["rank_feature"] == ["label-1"] assert retrieval_capture["rerank_mdl"] is not None - assert any(call[1] == module.LLMType.EMBEDDING.value and call[3].get("llm_name") == "embd-model" for call in llm_calls) + assert any(call[1] == module.LLMType.EMBEDDING.value and call[2] == "embd-model" for call in llm_calls) llm_calls.clear() @@ -1178,7 +1504,7 @@ def test_searchbots_retrieval_test_embedded_matrix_unit(monkeypatch): monkeypatch.setattr( module.KnowledgebaseService, "get_by_id", - lambda _kb_id: (True, SimpleNamespace(tenant_id="tenant-kb", embd_id="embd-model")), + lambda _kb_id: (True, SimpleNamespace(tenant_id="tenant-kb", embd_id="embd-model", tenant_embd_id=None)), ) res = _run(handler()) assert res["code"] == 0 @@ -1247,7 +1573,11 @@ def test_searchbots_related_questions_embedded_matrix_unit(monkeypatch): res = _run(handler()) assert res["code"] == 0 assert res["data"] == ["Alpha", "Beta"] - assert captured["bundle_args"] == ("tenant-1", module.LLMType.CHAT, "chat-x") + # LLMBundle is called with (tenant_id, model_config) + # model_config is a dict with model_type, llm_name, etc. + assert captured["bundle_args"][0] == "tenant-1" + assert captured["bundle_args"][1].get("model_type") == module.LLMType.CHAT + assert captured["bundle_args"][1].get("llm_name") == "chat-x" assert captured["options"] == {"temperature": 0.2} assert "Keywords: solar" in captured["messages"][0]["content"] @@ -1361,12 +1691,14 @@ def test_sequence2txt_embedded_validation_and_stream_matrix_unit(monkeypatch): assert "Unsupported audio format: .txt" in res["message"] _set_request({"stream": "false"}, {"file": _DummyUploadFile("audio.wav")}) - monkeypatch.setattr(module.TenantService, "get_info_by", lambda _tid: []) + tenant_llm_service = sys.modules["api.db.services.tenant_llm_service"] + monkeypatch.setattr(tenant_llm_service.TenantService, "get_by_id", lambda _tid: (False, None)) res = _run(handler("tenant-1")) assert res["message"] == "Tenant not found!" _set_request({"stream": "false"}, {"file": _DummyUploadFile("audio.wav")}) - monkeypatch.setattr(module.TenantService, "get_info_by", lambda _tid: [{"tenant_id": "tenant-1", "asr_id": ""}]) + tenant_llm_service = sys.modules["api.db.services.tenant_llm_service"] + monkeypatch.setattr(tenant_llm_service.TenantService, "get_by_id", lambda _tid: (True, SimpleNamespace(asr_id="", tts_id="", llm_id="", embd_id="", img2txt_id="", rerank_id=""))) res = _run(handler("tenant-1")) assert res["message"] == "No default ASR model is set" @@ -1378,7 +1710,7 @@ def test_sequence2txt_embedded_validation_and_stream_matrix_unit(monkeypatch): return [] _set_request({"stream": "false"}, {"file": _DummyUploadFile("audio.wav")}) - monkeypatch.setattr(module.TenantService, "get_info_by", lambda _tid: [{"tenant_id": "tenant-1", "asr_id": "asr-x"}]) + monkeypatch.setattr(tenant_llm_service.TenantService, "get_by_id", lambda _tid: (True, SimpleNamespace(asr_id="asr-x", tts_id="", llm_id="", embd_id="", img2txt_id="", rerank_id=""))) monkeypatch.setattr(module, "LLMBundle", lambda *_args, **_kwargs: _SyncASR()) monkeypatch.setattr(module.os, "remove", lambda _path: (_ for _ in ()).throw(RuntimeError("cleanup fail"))) res = _run(handler("tenant-1")) @@ -1420,14 +1752,15 @@ def test_sequence2txt_embedded_validation_and_stream_matrix_unit(monkeypatch): def test_tts_embedded_stream_and_error_matrix_unit(monkeypatch): module = _load_session_module(monkeypatch) handler = inspect.unwrap(module.tts) - monkeypatch.setattr(module, "Response", _StubResponse) monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"text": "A。B"})) + monkeypatch.setattr(module, "Response", _StubResponse) - monkeypatch.setattr(module.TenantService, "get_info_by", lambda _tid: []) + tenant_llm_service = sys.modules["api.db.services.tenant_llm_service"] + monkeypatch.setattr(tenant_llm_service.TenantService, "get_by_id", lambda _tid: (False, None)) res = _run(handler("tenant-1")) assert res["message"] == "Tenant not found!" - monkeypatch.setattr(module.TenantService, "get_info_by", lambda _tid: [{"tenant_id": "tenant-1", "tts_id": ""}]) + monkeypatch.setattr(tenant_llm_service.TenantService, "get_by_id", lambda _tid: (True, SimpleNamespace(asr_id="", tts_id="", llm_id="", embd_id="", img2txt_id="", rerank_id=""))) res = _run(handler("tenant-1")) assert res["message"] == "No default TTS model is set" @@ -1437,7 +1770,7 @@ def test_tts_embedded_stream_and_error_matrix_unit(monkeypatch): return [] yield f"chunk-{txt}".encode("utf-8") - monkeypatch.setattr(module.TenantService, "get_info_by", lambda _tid: [{"tenant_id": "tenant-1", "tts_id": "tts-x"}]) + monkeypatch.setattr(tenant_llm_service.TenantService, "get_by_id", lambda _tid: (True, SimpleNamespace(asr_id="", tts_id="tts-x", llm_id="", embd_id="", img2txt_id="", rerank_id=""))) monkeypatch.setattr(module, "LLMBundle", lambda *_args, **_kwargs: _TTSOk()) resp = _run(handler("tenant-1")) assert resp.mimetype == "audio/mpeg" diff --git a/test/testcases/test_sdk_api/test_memory_management/conftest.py b/test/testcases/test_sdk_api/test_memory_management/conftest.py index 516b408967..7027d541e6 100644 --- a/test/testcases/test_sdk_api/test_memory_management/conftest.py +++ b/test/testcases/test_sdk_api/test_memory_management/conftest.py @@ -31,7 +31,7 @@ def add_memory_func(client, request): payload = { "name": f"test_memory_{i}", "memory_type": ["raw"] + random.choices(["semantic", "episodic", "procedural"], k=random.randint(0, 3)), - "embd_id": "BAAI/bge-large-zh-v1.5@SILICONFLOW", + "embd_id": "BAAI/bge-small-en-v1.5@Builtin", "llm_id": "glm-4-flash@ZHIPU-AI" } res = client.create_memory(**payload) diff --git a/test/testcases/test_sdk_api/test_memory_management/test_create_memory.py b/test/testcases/test_sdk_api/test_memory_management/test_create_memory.py index 2c9a3e7c7d..0e90b1fb9d 100644 --- a/test/testcases/test_sdk_api/test_memory_management/test_create_memory.py +++ b/test/testcases/test_sdk_api/test_memory_management/test_create_memory.py @@ -36,7 +36,7 @@ class TestAuthorization: def test_auth_invalid(self, invalid_auth, expected_message): client = RAGFlow(invalid_auth, HOST_ADDRESS) with pytest.raises(Exception) as exception_info: - client.create_memory(**{"name": "test_memory", "memory_type": ["raw"], "embd_id": "BAAI/bge-large-zh-v1.5@SILICONFLOW", "llm_id": "glm-4-flash@ZHIPU-AI"}) + client.create_memory(**{"name": "test_memory", "memory_type": ["raw"], "embd_id": "BAAI/bge-small-en-v1.5@Builtin", "llm_id": "glm-4-flash@ZHIPU-AI"}) assert str(exception_info.value) == expected_message, str(exception_info.value) @@ -50,7 +50,7 @@ class TestMemoryCreate: payload = { "name": name, "memory_type": ["raw"] + random.choices(["semantic", "episodic", "procedural"], k=random.randint(0, 3)), - "embd_id": "BAAI/bge-large-zh-v1.5@SILICONFLOW", + "embd_id": "BAAI/bge-small-en-v1.5@Builtin", "llm_id": "glm-4-flash@ZHIPU-AI" } memory = client.create_memory(**payload) @@ -72,7 +72,7 @@ class TestMemoryCreate: payload = { "name": name, "memory_type": ["raw"] + random.choices(["semantic", "episodic", "procedural"], k=random.randint(0, 3)), - "embd_id": "BAAI/bge-large-zh-v1.5@SILICONFLOW", + "embd_id": "BAAI/bge-small-en-v1.5@Builtin", "llm_id": "glm-4-flash@ZHIPU-AI" } with pytest.raises(Exception) as exception_info: @@ -86,7 +86,7 @@ class TestMemoryCreate: payload = { "name": name, "memory_type": ["something"], - "embd_id": "BAAI/bge-large-zh-v1.5@SILICONFLOW", + "embd_id": "BAAI/bge-small-en-v1.5@Builtin", "llm_id": "glm-4-flash@ZHIPU-AI" } with pytest.raises(Exception) as exception_info: @@ -99,7 +99,7 @@ class TestMemoryCreate: payload = { "name": name, "memory_type": ["raw"] + random.choices(["semantic", "episodic", "procedural"], k=random.randint(0, 3)), - "embd_id": "BAAI/bge-large-zh-v1.5@SILICONFLOW", + "embd_id": "BAAI/bge-small-en-v1.5@Builtin", "llm_id": "glm-4-flash@ZHIPU-AI" } res1 = client.create_memory(**payload) diff --git a/test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py b/test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py index fbc447db42..5182500841 100644 --- a/test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py +++ b/test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py @@ -207,6 +207,10 @@ def _load_chunk_module(monkeypatch): EMBEDDING = SimpleNamespace(value="embedding") CHAT = SimpleNamespace(value="chat") RERANK = SimpleNamespace(value="rerank") + SPEECH2TEXT = SimpleNamespace(value="speech2text") + IMAGE2TEXT = SimpleNamespace(value="image2text") + TTS = SimpleNamespace(value="tts") + OCR = SimpleNamespace(value="ocr") constants_mod.RetCode = _DummyRetCode constants_mod.LLMType = _DummyLLMType @@ -301,6 +305,10 @@ def _load_chunk_module(monkeypatch): def get_embd_id(_doc_id): return "embed-1" + @staticmethod + def get_tenant_embd_id(_doc_id): + return 1 + @staticmethod def decrement_chunk_num(*args): _DocumentService.decrement_calls.append(args) @@ -327,13 +335,24 @@ def _load_chunk_module(monkeypatch): @staticmethod def get_by_id(_kb_id): - return True, SimpleNamespace(pagerank=0.6) + return True, SimpleNamespace(pagerank=0.6, tenant_embd_id=2, tenant_llm_id=1) kb_service_mod.KnowledgebaseService = _KnowledgebaseService monkeypatch.setitem(sys.modules, "api.db.services.knowledgebase_service", kb_service_mod) services_pkg.knowledgebase_service = kb_service_mod + class _DummyLLMService: + @staticmethod + def query(**_kwargs): + return [SimpleNamespace( + llm_name="gpt-3.5-turbo", + model_type="chat", + max_tokens=8192, + is_tools=True + )] + llm_service_mod = ModuleType("api.db.services.llm_service") + llm_service_mod.LLMService = _DummyLLMService llm_service_mod.LLMBundle = _DummyLLMBundle monkeypatch.setitem(sys.modules, "api.db.services.llm_service", llm_service_mod) services_pkg.llm_service = llm_service_mod @@ -343,6 +362,77 @@ def _load_chunk_module(monkeypatch): monkeypatch.setitem(sys.modules, "api.db.services.search_service", search_service_mod) services_pkg.search_service = search_service_mod + tenant_llm_service_mod = ModuleType("api.db.services.tenant_llm_service") + + class _MockTableObject: + def __init__(self, **kwargs): + for key, value in kwargs.items(): + setattr(self, key, value) + + def to_dict(self): + return {k: v for k, v in self.__dict__.items()} + + class _TenantLLMService: + @staticmethod + def get_by_id(tenant_model_id): + return True, _MockTableObject( + id=tenant_model_id, + tenant_id="tenant-1", + llm_factory="", + model_type="chat", + llm_name="gpt-3.5-turbo", + api_key="fake-api-key", + api_base="https://api.example.com", + max_tokens=8192, + used_tokens=0, + status=1 + ) + + @staticmethod + def get_api_key(tenant_id, model_name): + return _MockTableObject( + id=1, + tenant_id=tenant_id, + llm_factory="", + model_type="chat", + llm_name=model_name, + api_key="fake-api-key", + api_base="https://api.example.com", + max_tokens=8192, + used_tokens=0, + status=1 + ) + + @staticmethod + def split_model_name_and_factory(model_name): + if "@" in model_name: + parts = model_name.rsplit("@", 1) + return parts[0], parts[1] + return model_name, None + + @staticmethod + def increase_usage_by_id(model_id, used_tokens): + return True + + class _TenantService: + @staticmethod + def get_by_id(tenant_id): + return True, SimpleNamespace( + llm_id="gpt-3.5-turbo", + tenant_llm_id=1, + embd_id="text-embedding-ada-002", + tenant_embd_id=2, + asr_id="whisper-1", + img2txt_id="gpt-4-vision-preview", + rerank_id="bge-reranker", + tts_id="tts-1" + ) + + tenant_llm_service_mod.TenantLLMService = _TenantLLMService + tenant_llm_service_mod.TenantService = _TenantService + monkeypatch.setitem(sys.modules, "api.db.services.tenant_llm_service", tenant_llm_service_mod) + services_pkg.tenant_llm_service = tenant_llm_service_mod + user_service_mod = ModuleType("api.db.services.user_service") class _UserTenantService: @@ -775,7 +865,7 @@ def test_retrieval_test_branch_matrix_unit(monkeypatch): assert "Knowledgebase not found!" in res["message"], res retriever = _Retriever(mode="ok") - monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, SimpleNamespace(tenant_id="tenant-kb", embd_id="embd-1")), raising=False) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, SimpleNamespace(tenant_id="tenant-kb", embd_id="embd-1", tenant_embd_id=2)), raising=False) monkeypatch.setattr(module.settings, "retriever", retriever) monkeypatch.setattr(module.settings, "kg_retriever", _KgRetriever(), raising=False) _set_request_json( diff --git a/test/testcases/test_web_api/test_conversation_app/test_conversation_routes_unit.py b/test/testcases/test_web_api/test_conversation_app/test_conversation_routes_unit.py index 06caa19a4a..e53257a097 100644 --- a/test/testcases/test_web_api/test_conversation_app/test_conversation_routes_unit.py +++ b/test/testcases/test_web_api/test_conversation_app/test_conversation_routes_unit.py @@ -143,6 +143,19 @@ def _load_conversation_module(monkeypatch): apps_mod.login_required = lambda func: func monkeypatch.setitem(sys.modules, "api.apps", apps_mod) + # Create user_service module with TenantService stub if not already exists + if "api.db.services.user_service" not in sys.modules: + user_service_mod = ModuleType("api.db.services.user_service") + user_service_mod.UserService = SimpleNamespace() # Dummy UserService class + user_service_mod.TenantService = SimpleNamespace( + get_info_by=lambda _uid: [], + get_by_id=lambda _uid: (False, None) + ) + user_service_mod.UserTenantService = SimpleNamespace( + query=lambda **_kwargs: [] + ) + monkeypatch.setitem(sys.modules, "api.db.services.user_service", user_service_mod) + module_name = "test_conversation_routes_unit_module" module_path = repo_root / "api" / "apps" / "conversation_app.py" spec = importlib.util.spec_from_file_location(module_name, module_path) @@ -519,15 +532,15 @@ def test_sequence2txt_validation_and_transcription_paths(monkeypatch): wav_file = _DummyUploadedFile("audio.wav") monkeypatch.setattr(module, "request", _DummyRequest(form={"stream": "false"}, files={"file": wav_file})) - monkeypatch.setattr(module.TenantService, "get_info_by", lambda _uid: []) + monkeypatch.setattr(sys.modules["api.db.joint_services.tenant_model_service"].TenantService, "get_by_id", lambda _uid: (False, None)) res = _run(module.sequence2txt()) - assert res["message"] == "Tenant not found!" + assert res["message"] == "Tenant not found" wav_file = _DummyUploadedFile("audio.wav") monkeypatch.setattr(module, "request", _DummyRequest(form={"stream": "false"}, files={"file": wav_file})) - monkeypatch.setattr(module.TenantService, "get_info_by", lambda _uid: [{"tenant_id": "tenant-1", "asr_id": ""}]) + monkeypatch.setattr(sys.modules["api.db.joint_services.tenant_model_service"].TenantService, "get_by_id", lambda _uid: (True, SimpleNamespace(tenant_id="tenant-1", asr_id=""))) res = _run(module.sequence2txt()) - assert res["message"] == "No default ASR model is set" + assert res["message"] == "No default speech2text model is set." class _SyncAsr: def transcription(self, _path): @@ -538,7 +551,8 @@ def test_sequence2txt_validation_and_transcription_paths(monkeypatch): wav_file = _DummyUploadedFile("audio.wav") monkeypatch.setattr(module, "request", _DummyRequest(form={"stream": "false"}, files={"file": wav_file})) - monkeypatch.setattr(module.TenantService, "get_info_by", lambda _uid: [{"tenant_id": "tenant-1", "asr_id": "asr-model"}]) + monkeypatch.setattr(sys.modules["api.db.joint_services.tenant_model_service"].TenantService, "get_by_id", lambda _uid: (True, SimpleNamespace(tenant_id="tenant-1", asr_id="asr-model"))) + monkeypatch.setattr(module.TenantLLMService, "get_api_key", lambda tenant_id, model_name: SimpleNamespace(to_dict=lambda: {"llm_factory": "test", "llm_name": "asr-model"})) monkeypatch.setattr(module, "LLMBundle", lambda *_args, **_kwargs: _SyncAsr()) monkeypatch.setattr(module.os, "remove", lambda _path: (_ for _ in ()).throw(RuntimeError("remove failed"))) res = _run(module.sequence2txt()) @@ -579,13 +593,13 @@ def test_sequence2txt_validation_and_transcription_paths(monkeypatch): def test_tts_request_parse_entry(monkeypatch): module = _load_conversation_module(monkeypatch) _set_request_json(monkeypatch, module, {"text": "A。B"}) - monkeypatch.setattr(module.TenantService, "get_info_by", lambda _uid: []) + monkeypatch.setattr(sys.modules["api.db.joint_services.tenant_model_service"].TenantService, "get_by_id", lambda _uid: (False, None)) res = _run(module.tts()) - assert res["message"] == "Tenant not found!" + assert res["message"] == "Tenant not found" - monkeypatch.setattr(module.TenantService, "get_info_by", lambda _uid: [{"tenant_id": "tenant-1", "tts_id": ""}]) + monkeypatch.setattr(sys.modules["api.db.joint_services.tenant_model_service"].TenantService, "get_by_id", lambda _uid: (True, SimpleNamespace(tenant_id="tenant-1", tts_id=""))) res = _run(module.tts()) - assert res["message"] == "No default TTS model is set" + assert res["message"] == "No default tts model is set." class _TTSOk: def tts(self, txt): @@ -593,7 +607,8 @@ def test_tts_request_parse_entry(monkeypatch): return [] yield f"chunk-{txt}".encode("utf-8") - monkeypatch.setattr(module.TenantService, "get_info_by", lambda _uid: [{"tenant_id": "tenant-1", "tts_id": "tts-x"}]) + monkeypatch.setattr(sys.modules["api.db.joint_services.tenant_model_service"].TenantService, "get_by_id", lambda _uid: (True, SimpleNamespace(tenant_id="tenant-1", tts_id="tts-x"))) + monkeypatch.setattr(module.TenantLLMService, "get_api_key", lambda tenant_id, model_name: SimpleNamespace(to_dict=lambda: {"llm_factory": "test", "llm_name": model_name})) monkeypatch.setattr(module, "LLMBundle", lambda *_args, **_kwargs: _TTSOk()) resp = _run(module.tts()) assert resp.mimetype == "audio/mpeg" @@ -749,18 +764,18 @@ def test_mindmap_and_related_questions_matrix_unit(monkeypatch): llm_calls["options"] = options return "1. Alpha\n2. Beta\nignored" - def _fake_bundle(tenant_id, llm_type, chat_id): - llm_calls["bundle"] = (tenant_id, llm_type, chat_id) + def _fake_bundle(tenant_id, model_config, lang="Chinese", **kwargs): + llm_calls["bundle"] = (tenant_id, model_config) return _FakeChat() monkeypatch.setattr(module, "LLMBundle", _fake_bundle) monkeypatch.setattr(module, "load_prompt", lambda name: f"prompt-{name}") + monkeypatch.setattr(module.TenantLLMService, "get_api_key", lambda tenant_id, model_name: SimpleNamespace(to_dict=lambda: {"llm_factory": "test", "llm_name": model_name})) _set_request_json(monkeypatch, module, {"question": "solar", "search_id": "search-1"}) res = _run(module.related_questions.__wrapped__()) assert res["code"] == 0 assert res["data"] == ["Alpha", "Beta"] assert llm_calls["bundle"][0] == "user-1" - assert llm_calls["bundle"][2] == "chat-x" assert llm_calls["options"] == {"temperature": 0.2} assert llm_calls["prompt"] == "prompt-related_question" assert "Keywords: solar" in llm_calls["messages"][0]["content"] diff --git a/test/testcases/test_web_api/test_dialog_app/test_dialog_routes_unit.py b/test/testcases/test_web_api/test_dialog_app/test_dialog_routes_unit.py index 565c516d9e..b5b9e1ab73 100644 --- a/test/testcases/test_web_api/test_dialog_app/test_dialog_routes_unit.py +++ b/test/testcases/test_web_api/test_dialog_app/test_dialog_routes_unit.py @@ -137,11 +137,34 @@ def _load_dialog_module(monkeypatch): tenant_llm_service_mod = ModuleType("api.db.services.tenant_llm_service") + class _MockTableObject: + def __init__(self, **kwargs): + for key, value in kwargs.items(): + setattr(self, key, value) + + def to_dict(self): + return {k: v for k, v in self.__dict__.items()} + class _TenantLLMService: @staticmethod def split_model_name_and_factory(embd_id): return embd_id.split("@") + @staticmethod + def get_api_key(tenant_id, model_name): + return _MockTableObject( + id=1, + tenant_id=tenant_id, + llm_factory="", + model_type="chat", + llm_name=model_name, + api_key="fake-api-key", + api_base="https://api.example.com", + max_tokens=8192, + used_tokens=0, + status=1 + ) + tenant_llm_service_mod.TenantLLMService = _TenantLLMService monkeypatch.setitem(sys.modules, "api.db.services.tenant_llm_service", tenant_llm_service_mod) @@ -253,8 +276,8 @@ def test_set_dialog_branch_matrix_unit(monkeypatch): monkeypatch.setattr(module, "duplicate_name", _dup_name) monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: [SimpleNamespace(name="new dialog")]) - monkeypatch.setattr(module.TenantService, "get_by_id", lambda _id: (True, SimpleNamespace(llm_id="llm-x"))) - monkeypatch.setattr(module.KnowledgebaseService, "get_by_ids", lambda _ids: [SimpleNamespace(embd_id="embd-a@builtin")]) + monkeypatch.setattr(module.TenantService, "get_by_id", lambda _id: (True, SimpleNamespace(llm_id="llm-x", tenant_llm_id=1))) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_ids", lambda _ids: [SimpleNamespace(embd_id="embd-a@builtin", tenant_embd_id=2)]) monkeypatch.setattr(module.TenantLLMService, "split_model_name_and_factory", lambda embd_id: embd_id.split("@")) monkeypatch.setattr(module.DialogService, "save", lambda **kwargs: captured.update(kwargs) or False) _set_request_json( @@ -301,7 +324,7 @@ def test_set_dialog_branch_matrix_unit(monkeypatch): res = _run(handler()) assert res["message"] == "Tenant not found!" - monkeypatch.setattr(module.TenantService, "get_by_id", lambda _id: (True, SimpleNamespace(llm_id="llm-x"))) + monkeypatch.setattr(module.TenantService, "get_by_id", lambda _id: (True, SimpleNamespace(llm_id="llm-x", tenant_llm_id=1))) monkeypatch.setattr( module, "get_request_json", @@ -316,7 +339,7 @@ def test_set_dialog_branch_matrix_unit(monkeypatch): monkeypatch.setattr( module.KnowledgebaseService, "get_by_ids", - lambda _ids: [SimpleNamespace(embd_id="embd-a@f1"), SimpleNamespace(embd_id="embd-b@f2")], + lambda _ids: [SimpleNamespace(embd_id="embd-a@f1", tenant_embd_id=2), SimpleNamespace(embd_id="embd-b@f2", tenant_embd_id=2)], ) monkeypatch.setattr(module.TenantLLMService, "split_model_name_and_factory", lambda embd_id: embd_id.split("@")) res = _run(handler()) diff --git a/test/testcases/test_web_api/test_llm_app/test_llm_list_unit.py b/test/testcases/test_web_api/test_llm_app/test_llm_list_unit.py index 6253a5ade0..dea30e68e8 100644 --- a/test/testcases/test_web_api/test_llm_app/test_llm_list_unit.py +++ b/test/testcases/test_web_api/test_llm_app/test_llm_list_unit.py @@ -51,11 +51,19 @@ class _DummyTenantLLMModel: llm_factory = _ExprField("llm_factory") llm_name = _ExprField("llm_name") + def __init__(self, id=None, **kwargs): + self.id = id + self.api_key = None + self.status = None + for key, value in kwargs.items(): + setattr(self, key, value) + class _TenantLLMRow: def __init__( self, *, + id, llm_name, llm_factory, model_type, @@ -65,6 +73,7 @@ class _TenantLLMRow: api_base="", max_tokens=8192, ): + self.id = id self.llm_name = llm_name self.llm_factory = llm_factory self.model_type = model_type @@ -76,6 +85,7 @@ class _TenantLLMRow: def to_dict(self): return { + "id": self.id, "llm_name": self.llm_name, "llm_factory": self.llm_factory, "model_type": self.model_type, @@ -246,8 +256,8 @@ def test_list_app_grouping_availability_and_merge(monkeypatch): monkeypatch.setattr(module.TenantLLMService, "ensure_mineru_from_env", lambda tenant_id: ensure_calls.append(tenant_id)) tenant_rows = [ - _TenantLLMRow(llm_name="fast-emb", llm_factory="FastEmbed", model_type="embedding", api_key="k1", status="1"), - _TenantLLMRow(llm_name="tenant-only", llm_factory="CustomFactory", model_type="chat", api_key="k2", status="1"), + _TenantLLMRow(id=1, llm_name="fast-emb", llm_factory="FastEmbed", model_type="embedding", api_key="k1", status="1"), + _TenantLLMRow(id=2, llm_name="tenant-only", llm_factory="CustomFactory", model_type="chat", api_key="k2", status="1"), ] monkeypatch.setattr(module.TenantLLMService, "query", lambda **_kwargs: tenant_rows) @@ -263,7 +273,7 @@ def test_list_app_grouping_availability_and_merge(monkeypatch): monkeypatch.setenv("TEI_MODEL", "tei-embed") res = _run(module.list_app()) - assert res["code"] == 0 + assert res["code"] == 0, res["message"] assert ensure_calls == ["tenant-1"] data = res["data"] @@ -291,8 +301,8 @@ def test_list_app_model_type_filter(monkeypatch): module.TenantLLMService, "query", lambda **_kwargs: [ - _TenantLLMRow(llm_name="fast-emb", llm_factory="FastEmbed", model_type="embedding", api_key="k1", status="1"), - _TenantLLMRow(llm_name="tenant-only", llm_factory="CustomFactory", model_type="chat", api_key="k2", status="1"), + _TenantLLMRow(id=1, llm_name="fast-emb", llm_factory="FastEmbed", model_type="embedding", api_key="k1", status="1"), + _TenantLLMRow(id=2, llm_name="tenant-only", llm_factory="CustomFactory", model_type="chat", api_key="k2", status="1"), ], ) monkeypatch.setattr( @@ -306,7 +316,7 @@ def test_list_app_model_type_filter(monkeypatch): monkeypatch.setattr(module, "request", SimpleNamespace(args={"model_type": "chat"})) res = _run(module.list_app()) - assert res["code"] == 0 + assert res["code"] == 0, res["message"] assert list(res["data"].keys()) == ["CustomFactory"] assert res["data"]["CustomFactory"][0]["model_type"] == "chat" @@ -799,7 +809,7 @@ def test_add_llm_model_type_probe_and_persistence_matrix_unit(monkeypatch): monkeypatch.setattr(module.TenantLLMService, "filter_update", lambda _filters, _payload: False) monkeypatch.setattr(module.TenantLLMService, "save", lambda **kwargs: saved.append(kwargs) or True) res = _call({"llm_factory": "FChatPass", "llm_name": "m", "model_type": module.LLMType.CHAT.value, "api_key": "k"}) - assert res["code"] == 0 + assert res["code"] == 0, res["message"] assert res["data"] is True assert saved assert saved[0]["llm_factory"] == "FChatPass" @@ -841,6 +851,7 @@ def test_my_llms_include_details_and_exception_unit(monkeypatch): "query", lambda **_kwargs: [ _TenantLLMRow( + id=1, llm_name="chat-model", llm_factory="FactoryX", model_type="chat", diff --git a/test/testcases/test_web_api/test_memory_app/conftest.py b/test/testcases/test_web_api/test_memory_app/conftest.py index 7fdd78f53f..3c7fa3bc96 100644 --- a/test/testcases/test_web_api/test_memory_app/conftest.py +++ b/test/testcases/test_web_api/test_memory_app/conftest.py @@ -32,7 +32,7 @@ def add_memory_func(request, WebApiAuth): payload = { "name": f"test_memory_{i}", "memory_type": ["raw"] + random.choices(["semantic", "episodic", "procedural"], k=random.randint(0, 3)), - "embd_id": "BAAI/bge-large-zh-v1.5@SILICONFLOW", + "embd_id": "BAAI/bge-small-en-v1.5@Builtin", "llm_id": "glm-4-flash@ZHIPU-AI" } res = create_memory(WebApiAuth, payload) diff --git a/test/testcases/test_web_api/test_memory_app/test_create_memory.py b/test/testcases/test_web_api/test_memory_app/test_create_memory.py index d2b0f967bc..ad168911c7 100644 --- a/test/testcases/test_web_api/test_memory_app/test_create_memory.py +++ b/test/testcases/test_web_api/test_memory_app/test_create_memory.py @@ -45,7 +45,7 @@ class TestMemoryCreate: payload = { "name": name, "memory_type": ["raw"] + random.choices(["semantic", "episodic", "procedural"], k=random.randint(0, 3)), - "embd_id": "BAAI/bge-large-zh-v1.5@SILICONFLOW", + "embd_id": "BAAI/bge-small-en-v1.5@Builtin", "llm_id": "glm-4-flash@ZHIPU-AI" } res = create_memory(WebApiAuth, payload) @@ -68,7 +68,7 @@ class TestMemoryCreate: payload = { "name": name, "memory_type": ["raw"] + random.choices(["semantic", "episodic", "procedural"], k=random.randint(0, 3)), - "embd_id": "BAAI/bge-large-zh-v1.5@SILICONFLOW", + "embd_id": "BAAI/bge-small-en-v1.5@Builtin", "llm_id": "glm-4-flash@ZHIPU-AI" } res = create_memory(WebApiAuth, payload) @@ -80,7 +80,7 @@ class TestMemoryCreate: payload = { "name": name, "memory_type": ["something"], - "embd_id": "BAAI/bge-large-zh-v1.5@SILICONFLOW", + "embd_id": "BAAI/bge-small-en-v1.5@Builtin", "llm_id": "glm-4-flash@ZHIPU-AI" } res = create_memory(WebApiAuth, payload) @@ -92,7 +92,7 @@ class TestMemoryCreate: payload = { "name": name, "memory_type": ["raw"] + random.choices(["semantic", "episodic", "procedural"], k=random.randint(0, 3)), - "embd_id": "BAAI/bge-large-zh-v1.5@SILICONFLOW", + "embd_id": "BAAI/bge-small-en-v1.5@Builtin", "llm_id": "glm-4-flash@ZHIPU-AI" } res1 = create_memory(WebApiAuth, payload) diff --git a/test/testcases/test_web_api/test_user_app/test_user_app_unit.py b/test/testcases/test_web_api/test_user_app/test_user_app_unit.py index 9ad71c7acb..5cecafe235 100644 --- a/test/testcases/test_web_api/test_user_app/test_user_app_unit.py +++ b/test/testcases/test_web_api/test_user_app/test_user_app_unit.py @@ -216,11 +216,34 @@ def _load_user_app(monkeypatch): tenant_llm_service_mod = ModuleType("api.db.services.tenant_llm_service") + class _MockTableObject: + def __init__(self, **kwargs): + for key, value in kwargs.items(): + setattr(self, key, value) + + def to_dict(self): + return {k: v for k, v in self.__dict__.items()} + class _StubTenantLLMService: @staticmethod def insert_many(_payload): return True + @staticmethod + def get_api_key(tenant_id, model_name): + return _MockTableObject( + id=1, + tenant_id=tenant_id, + llm_factory="", + model_type="chat", + llm_name=model_name, + api_key="fake-api-key", + api_base="https://api.example.com", + max_tokens=8192, + used_tokens=0, + status=1 + ) + tenant_llm_service_mod.TenantLLMService = _StubTenantLLMService monkeypatch.setitem(sys.modules, "api.db.services.tenant_llm_service", tenant_llm_service_mod)