From 56d73d0c2c833215e716b3b43523a675dc21f93d Mon Sep 17 00:00:00 2001 From: Wang Qi Date: Mon, 18 May 2026 15:55:59 +0800 Subject: [PATCH] Refactor: speed up ragflow server, save startup memory (#14973) ### What problem does this PR solve? Refactor: speed up ragflow server, save startup memory, saved 200MiB, and 5-9 seconds start time. ##### Before 1241292 | | \_ python3 api/ragflow_server.py RAGFlow server is ready after 25.61845850944519s initialization. ##### After 1019968 | | \_ python3 api/ragflow_server.py RAGFlow server is ready after 16.205134391784668s initialization. ### Type of change - [x] Refactoring --- api/apps/llm_app.py | 8 +++++-- api/apps/restful_apis/agent_api.py | 33 ++++++++++++++++++++++----- api/apps/restful_apis/chunk_api.py | 15 ++++++++++-- api/db/services/tenant_llm_service.py | 3 ++- 4 files changed, 48 insertions(+), 11 deletions(-) diff --git a/api/apps/llm_app.py b/api/apps/llm_app.py index e1d4954df4..58997fb4df 100644 --- a/api/apps/llm_app.py +++ b/api/apps/llm_app.py @@ -25,8 +25,6 @@ from api.db.services.llm_service import LLMService from api.utils.api_utils import get_allowed_llm_factories, get_data_error_result, get_json_result, get_request_json, server_error_response, validate_request from common.constants import StatusEnum, LLMType from api.db.db_models import TenantLLM -from rag.utils.base64_image import test_image -from rag.llm import EmbeddingModel, ChatModel, RerankModel, CvModel, TTSModel, OcrModel, Seq2txtModel def _resolve_my_llm_is_tools(o_dict: dict) -> bool: @@ -78,6 +76,8 @@ def factories(): @validate_request("llm_factory", "api_key") async def set_api_key(): req = await get_request_json() + from rag.llm import ChatModel, EmbeddingModel, RerankModel + # test if api key works chat_passed, embd_passed, rerank_passed = False, False, False factory = req["llm_factory"] @@ -178,6 +178,8 @@ async def set_api_key(): @validate_request("llm_factory") async def add_llm(): req = await get_request_json() + from rag.llm import ChatModel, CvModel, EmbeddingModel, OcrModel, RerankModel, Seq2txtModel, TTSModel + factory = req["llm_factory"] api_key = req.get("api_key", "x") llm_name = req.get("llm_name") @@ -318,6 +320,8 @@ async def add_llm(): msg += f"\nFail to access model({factory}/{mdl_nm})." + str(e) case LLMType.IMAGE2TEXT.value: + from rag.utils.base64_image import test_image + assert factory in CvModel, f"Image to text model from {factory} is not supported yet." mdl = CvModel[factory](key=model_api_key, model_name=mdl_nm, base_url=model_base_url) try: diff --git a/api/apps/restful_apis/agent_api.py b/api/apps/restful_apis/agent_api.py index f88ce90b3f..8478571ac1 100644 --- a/api/apps/restful_apis/agent_api.py +++ b/api/apps/restful_apis/agent_api.py @@ -29,9 +29,6 @@ from functools import partial, wraps import jwt from quart import Response, jsonify, request -from agent.canvas import Canvas -from agent.component import LLM -from agent.dsl_migration import normalize_chunker_dsl from api.apps import current_user, login_required from api.apps.services.canvas_replica_service import CanvasReplicaService from api.db import CanvasCategory @@ -64,9 +61,6 @@ from common.ssrf_guard import assert_host_is_safe from common.constants import RetCode from common.misc_utils import get_uuid, thread_pool_exec from peewee import MySQLDatabase, PostgresqlDatabase -from rag.flow.pipeline import Pipeline -from rag.nlp import search -from rag.utils.redis_conn import REDIS_CONN def _require_canvas_access_sync(func): @@ -195,6 +189,8 @@ def list_agent_sessions(agent_id, tenant_id): @add_tenant_id_to_kwargs @_require_canvas_access_async async def create_agent_session(agent_id, tenant_id): + from agent.canvas import Canvas + req = await get_request_json() user_id = req.get("user_id") or request.args.get("user_id", tenant_id) release_mode = bool(req.get("release", request.args.get("release", False))) @@ -522,6 +518,8 @@ async def upload_agent_file(agent_id, tenant_id): @_require_canvas_access_sync def get_agent_component_input_form(agent_id, component_id, tenant_id): try: + from agent.canvas import Canvas + exists, user_canvas = UserCanvasService.get_by_id(agent_id) if not exists: return get_data_error_result(message="canvas not found.") @@ -539,6 +537,9 @@ def get_agent_component_input_form(agent_id, component_id, tenant_id): async def debug_agent_component(agent_id, component_id, tenant_id): req = await get_request_json() try: + from agent.canvas import Canvas + from agent.component import LLM + _, user_canvas = UserCanvasService.get_by_id(agent_id) canvas = Canvas(json.dumps(user_canvas.dsl), tenant_id, canvas_id=user_canvas.id) canvas.reset() @@ -597,6 +598,8 @@ def get_agent(agent_id, tenant_id): released_versions.sort(key=lambda version: version.update_time, reverse=True) last_publish_time = released_versions[0].update_time + from agent.dsl_migration import normalize_chunker_dsl + canvas["dsl"] = normalize_chunker_dsl(canvas.get("dsl", {})) canvas["last_publish_time"] = last_publish_time @@ -642,6 +645,8 @@ def get_agent_version(agent_id, version_id, tenant_id): @_require_canvas_access_async async def get_agent_logs(agent_id, message_id, tenant_id): try: + from rag.utils.redis_conn import REDIS_CONN + binary = await thread_pool_exec(REDIS_CONN.get, f"{agent_id}-{message_id}-logs") if not binary: return get_json_result(data={}) @@ -719,6 +724,8 @@ async def update_agent(agent_id, tenant_id): @_require_canvas_access_async async def reset_agent(agent_id, tenant_id): try: + from agent.canvas import Canvas + exists, user_canvas = UserCanvasService.get_by_id(agent_id) if not exists: return get_data_error_result(message="canvas not found.") @@ -747,6 +754,8 @@ async def reset_agent(agent_id, tenant_id): @login_required @add_tenant_id_to_kwargs async def rerun_agent(tenant_id): + from rag.nlp import search + req = await get_request_json() doc = PipelineOperationLogService.get_documents_info(req["id"]) if not doc: @@ -1042,6 +1051,8 @@ async def agent_chat_completion(tenant_id, agent_id=None): dsl_str = json.dumps(replica_dsl, ensure_ascii=False) if cvs.canvas_category == CanvasCategory.DataFlow: + from rag.flow.pipeline import Pipeline + task_id = get_uuid() Pipeline( dsl_str, @@ -1064,6 +1075,8 @@ async def agent_chat_completion(tenant_id, agent_id=None): return get_json_result(data={"message_id": task_id}) try: + from agent.canvas import Canvas + canvas = Canvas(dsl_str, str(tenant_id), canvas_id=agent_id, custom_header=custom_header) except Exception as exc: return server_error_response(exc) @@ -1349,6 +1362,8 @@ async def webhook(agent_id: str): now = time.time() try: + from rag.utils.redis_conn import REDIS_CONN + res = REDIS_CONN.lua_token_bucket( keys=[key], args=[capacity, rate, now, cost], @@ -1456,6 +1471,8 @@ async def webhook(agent_id: str): if not isinstance(cvs.dsl, str): dsl = json.dumps(cvs.dsl, ensure_ascii=False) try: + from agent.canvas import Canvas + canvas = Canvas(dsl, cvs.user_id, agent_id, canvas_id=agent_id) except Exception as e: resp=get_data_error_result(code=RetCode.BAD_REQUEST,message=str(e)) @@ -1709,6 +1726,8 @@ async def webhook(agent_id: str): response_cfg = webhook_cfg.get("response", {}) def append_webhook_trace(agent_id: str, start_ts: float,event: dict, ttl=600): + from rag.utils.redis_conn import REDIS_CONN + key = f"webhook-trace-{agent_id}-logs" raw = REDIS_CONN.get(key) @@ -1908,6 +1927,8 @@ async def webhook_trace(agent_id: str): webhook_id = request.args.get("webhook_id") key = f"webhook-trace-{agent_id}-logs" + from rag.utils.redis_conn import REDIS_CONN + raw = REDIS_CONN.get(key) if since_ts is None: diff --git a/api/apps/restful_apis/chunk_api.py b/api/apps/restful_apis/chunk_api.py index d3a30710e8..fe45209dd0 100644 --- a/api/apps/restful_apis/chunk_api.py +++ b/api/apps/restful_apis/chunk_api.py @@ -43,8 +43,6 @@ from common.constants import LLMType, ParserType, RetCode from common.misc_utils import thread_pool_exec from common.string_utils import is_content_empty, remove_redundant_spaces from common.tag_feature_utils import validate_tag_features -from rag.app.qa import beAdoc, rmPrefix -from rag.nlp import rag_tokenizer, search class Chunk(BaseModel): @@ -107,6 +105,8 @@ def _get_dataset_tenant_id(dataset_id): @login_required @add_tenant_id_to_kwargs async def list_chunks(tenant_id, dataset_id, document_id): + from rag.nlp import search + if not KnowledgebaseService.accessible(kb_id=dataset_id, user_id=tenant_id): return get_error_data_result(message=f"You don't own the dataset {dataset_id}.") dataset_tenant_id = _get_dataset_tenant_id(dataset_id) @@ -191,6 +191,8 @@ async def list_chunks(tenant_id, dataset_id, document_id): @login_required @add_tenant_id_to_kwargs async def get_chunk(tenant_id, dataset_id, document_id, chunk_id): + from rag.nlp import search + if not KnowledgebaseService.accessible(kb_id=dataset_id, user_id=tenant_id): return get_error_data_result(message=f"You don't own the dataset {dataset_id}.") dataset_tenant_id = _get_dataset_tenant_id(dataset_id) @@ -214,6 +216,8 @@ async def get_chunk(tenant_id, dataset_id, document_id, chunk_id): @login_required @add_tenant_id_to_kwargs async def add_chunk(tenant_id, dataset_id, document_id): + from rag.nlp import rag_tokenizer, search + if not KnowledgebaseService.accessible(kb_id=dataset_id, user_id=tenant_id): return get_error_data_result(message=f"You don't own the dataset {dataset_id}.") dataset_tenant_id = _get_dataset_tenant_id(dataset_id) @@ -303,6 +307,8 @@ async def add_chunk(tenant_id, dataset_id, document_id): @login_required @add_tenant_id_to_kwargs async def rm_chunk(tenant_id, dataset_id, document_id): + from rag.nlp import search + if not KnowledgebaseService.accessible(kb_id=dataset_id, user_id=tenant_id): return get_error_data_result(message=f"You don't own the dataset {dataset_id}.") dataset_tenant_id = _get_dataset_tenant_id(dataset_id) @@ -350,6 +356,9 @@ async def rm_chunk(tenant_id, dataset_id, document_id): @login_required @add_tenant_id_to_kwargs async def update_chunk(tenant_id, dataset_id, document_id, chunk_id): + from rag.app.qa import beAdoc, rmPrefix + from rag.nlp import rag_tokenizer, search + if not KnowledgebaseService.accessible(kb_id=dataset_id, user_id=tenant_id): return get_error_data_result(message=f"You don't own the dataset {dataset_id}.") dataset_tenant_id = _get_dataset_tenant_id(dataset_id) @@ -436,6 +445,8 @@ async def update_chunk(tenant_id, dataset_id, document_id, chunk_id): @login_required @add_tenant_id_to_kwargs async def switch_chunks(tenant_id, dataset_id, document_id): + from rag.nlp import search + if not KnowledgebaseService.accessible(kb_id=dataset_id, user_id=tenant_id): return get_error_data_result(message=f"You don't own the dataset {dataset_id}.") dataset_tenant_id = _get_dataset_tenant_id(dataset_id) diff --git a/api/db/services/tenant_llm_service.py b/api/db/services/tenant_llm_service.py index ee2eab6648..f14f97fcef 100644 --- a/api/db/services/tenant_llm_service.py +++ b/api/db/services/tenant_llm_service.py @@ -24,7 +24,6 @@ from api.db.db_models import DB, LLMFactories, TenantLLM from api.db.services.common_service import CommonService from api.db.services.langfuse_service import TenantLangfuseService from api.db.services.user_service import TenantService -from rag.llm import ChatModel, CvModel, EmbeddingModel, OcrModel, RerankModel, Seq2txtModel, TTSModel class LLMFactoriesService(CommonService): @@ -183,6 +182,8 @@ class TenantLLMService(CommonService): def model_instance(cls, model_config: dict, lang="Chinese", **kwargs): if not model_config: raise LookupError("Model config is required") + from rag.llm import ChatModel, CvModel, EmbeddingModel, OcrModel, RerankModel, Seq2txtModel, TTSModel + kwargs.update({"provider": model_config["llm_factory"]}) api_key = model_config.get("api_key_payload", model_config["api_key"]) if model_config["model_type"] == LLMType.EMBEDDING.value: