zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-20 06:31:02 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refactor: Task Executor (#15154)

Browse Source 
### What problem does this PR solve?

1. Break huge function into smaller pieces
2. Add unit test for the smaller pieces function
3. Layer-ed design
a. infra layer - task_context.py, recording_context.py,
write_operation_interceptor.py, ...
    b. service layer - *_service.py
    c. business layer - task_handler.py
4. Default behavior: use "refactor-ed version" - can switch to original
version by change env variable

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
- [x] Refactoring
- [x] Performance Improvement

---------

Co-authored-by: Liu An <asiro@qq.com>
Co-authored-by: Zhichang Yu <yuzhichang@gmail.com>
This commit is contained in:
Jack
2026-05-27 21:54:17 +08:00
committed by GitHub
parent 0071e98c11
commit f0cb7a544b
55 changed files with 12707 additions and 465 deletions
Download Patch File Download Diff File Expand all files Collapse all files

106
api/db/joint_services/memory_message_service.py
View File

@@ -161,19 +161,19 @@ async def extract_by_llm(tenant_id: str, tenant_llm_id: int, extract_conf: dict,
)
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)
res_json = get_json_result_from_llm_response(res)
if task_id:
TaskService.update_progress(task_id, {"progress": 0.35, "progress_msg": timestamp_to_date(current_timestamp())+ " " + "Get extracted result from LLM."})
return [{
"content": extracted_content["content"],
"valid_at": format_iso_8601_to_ymd_hms(extracted_content["valid_at"]),
"invalid_at": format_iso_8601_to_ymd_hms(extracted_content["invalid_at"]) if extracted_content.get("invalid_at") else "",
"message_type": message_type
} for message_type, extracted_content_list in res_json.items() for extracted_content in extracted_content_list]
with LLMBundle(tenant_id, llm_config) as llm:
if task_id:
TaskService.update_progress(task_id, {"progress": 0.15, "progress_msg": timestamp_to_date(current_timestamp())+ " " + "Prepared prompts and LLM."})
res = await llm.async_chat(system_prompt, user_prompts, extract_conf)
res_json = get_json_result_from_llm_response(res)
if task_id:
TaskService.update_progress(task_id, {"progress": 0.35, "progress_msg": timestamp_to_date(current_timestamp())+ " " + "Get extracted result from LLM."})
return [{
"content": extracted_content["content"],
"valid_at": format_iso_8601_to_ymd_hms(extracted_content["valid_at"]),
"invalid_at": format_iso_8601_to_ymd_hms(extracted_content["invalid_at"]) if extracted_content.get("invalid_at") else "",
"message_type": message_type
} for message_type, extracted_content_list in res_json.items() for extracted_content in extracted_content_list]
async def embed_and_save(memory, message_list: list[dict], task_id: str=None):
@@ -185,48 +185,48 @@ async def embed_and_save(memory, message_list: list[dict], task_id: str=None):
)
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])
for idx, msg in enumerate(message_list):
msg["content_embed"] = vector_list[idx]
if task_id:
TaskService.update_progress(task_id, {"progress": 0.85, "progress_msg": timestamp_to_date(current_timestamp())+ " " + "Embedded extracted content."})
vector_dimension = len(vector_list[0])
if not MessageService.has_index(memory.tenant_id, memory.id):
created = MessageService.create_index(memory.tenant_id, memory.id, vector_size=vector_dimension)
if not created:
error_msg = "Failed to create message index."
if task_id:
TaskService.update_progress(task_id, {"progress": -1, "progress_msg": timestamp_to_date(current_timestamp())+ " " + error_msg})
return False, error_msg
new_msg_size = sum([MessageService.calculate_message_size(m) for m in message_list])
current_memory_size = get_memory_size_cache(memory.tenant_id, memory.id)
if new_msg_size + current_memory_size > memory.memory_size:
size_to_delete = current_memory_size + new_msg_size - memory.memory_size
if memory.forgetting_policy == "FIFO":
message_ids_to_delete, delete_size = MessageService.pick_messages_to_delete_by_fifo(memory.id, memory.tenant_id,
size_to_delete)
MessageService.delete_message({"message_id": message_ids_to_delete}, memory.tenant_id, memory.id)
decrease_memory_size_cache(memory.id, delete_size)
else:
error_msg = "Failed to insert message into memory. Memory size reached limit and cannot decide which to delete."
if task_id:
TaskService.update_progress(task_id, {"progress": -1, "progress_msg": timestamp_to_date(current_timestamp())+ " " + error_msg})
return False, error_msg
fail_cases = MessageService.insert_message(message_list, memory.tenant_id, memory.id)
if fail_cases:
error_msg = "Failed to insert message into memory. Details: " + "; ".join(fail_cases)
with LLMBundle(memory.tenant_id, embd_model_config) as embedding_model:
if task_id:
TaskService.update_progress(task_id, {"progress": -1, "progress_msg": timestamp_to_date(current_timestamp())+ " " + error_msg})
return False, error_msg
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])
for idx, msg in enumerate(message_list):
msg["content_embed"] = vector_list[idx]
if task_id:
TaskService.update_progress(task_id, {"progress": 0.85, "progress_msg": timestamp_to_date(current_timestamp())+ " " + "Embedded extracted content."})
vector_dimension = len(vector_list[0])
if not MessageService.has_index(memory.tenant_id, memory.id):
created = MessageService.create_index(memory.tenant_id, memory.id, vector_size=vector_dimension)
if not created:
error_msg = "Failed to create message index."
if task_id:
TaskService.update_progress(task_id, {"progress": -1, "progress_msg": timestamp_to_date(current_timestamp())+ " " + error_msg})
return False, error_msg
if task_id:
TaskService.update_progress(task_id, {"progress": 0.95, "progress_msg": timestamp_to_date(current_timestamp())+ " " + "Saved messages to storage."})
increase_memory_size_cache(memory.id, new_msg_size)
return True, "Message saved successfully."
new_msg_size = sum([MessageService.calculate_message_size(m) for m in message_list])
current_memory_size = get_memory_size_cache(memory.tenant_id, memory.id)
if new_msg_size + current_memory_size > memory.memory_size:
size_to_delete = current_memory_size + new_msg_size - memory.memory_size
if memory.forgetting_policy == "FIFO":
message_ids_to_delete, delete_size = MessageService.pick_messages_to_delete_by_fifo(memory.id, memory.tenant_id,
size_to_delete)
MessageService.delete_message({"message_id": message_ids_to_delete}, memory.tenant_id, memory.id)
decrease_memory_size_cache(memory.id, delete_size)
else:
error_msg = "Failed to insert message into memory. Memory size reached limit and cannot decide which to delete."
if task_id:
TaskService.update_progress(task_id, {"progress": -1, "progress_msg": timestamp_to_date(current_timestamp())+ " " + error_msg})
return False, error_msg
fail_cases = MessageService.insert_message(message_list, memory.tenant_id, memory.id)
if fail_cases:
error_msg = "Failed to insert message into memory. Details: " + "; ".join(fail_cases)
if task_id:
TaskService.update_progress(task_id, {"progress": -1, "progress_msg": timestamp_to_date(current_timestamp())+ " " + error_msg})
return False, error_msg
if task_id:
TaskService.update_progress(task_id, {"progress": 0.95, "progress_msg": timestamp_to_date(current_timestamp())+ " " + "Saved messages to storage."})
increase_memory_size_cache(memory.id, new_msg_size)
return True, "Message saved successfully."
def query_message(filter_dict: dict, params: dict):

6
api/db/services/document_service.py
View File

@@ -1098,12 +1098,12 @@ def queue_raptor_o_graphrag_tasks(sample_doc, ty, priority, fake_doc_id="", doc_
task["doc_ids"] = doc_ids
DocumentService.begin2parse(task["doc_id"], keep_progress=True)
assert REDIS_CONN.queue_product(settings.get_svr_queue_name(priority), message=task), "Can't access Redis. Please check the Redis' status."
assert REDIS_CONN.queue_product(settings.get_svr_queue_name(priority, ty), message=task), "Can't access Redis. Please check the Redis' status."
return task["id"]
def get_queue_length(priority):
group_info = REDIS_CONN.queue_info(settings.get_svr_queue_name(priority), SVR_CONSUMER_GROUP_NAME)
def get_queue_length(priority, suffix="common"):
group_info = REDIS_CONN.queue_info(settings.get_svr_queue_name(priority, suffix), SVR_CONSUMER_GROUP_NAME)
if not group_info:
return 0
return int(group_info.get("lag", 0) or 0)

15
api/db/services/llm_service.py
View File

@@ -86,6 +86,19 @@ class LLMBundle(LLM4Tenant):
def __init__(self, tenant_id: str, model_config: dict, lang="Chinese", **kwargs):
super().__init__(tenant_id, model_config, lang, **kwargs)
def close(self):
"""Release resources held by this LLMBundle instance."""
super().close()
def __enter__(self):
"""Enter context manager."""
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""Exit context manager and release resources."""
self.close()
return False
def bind_tools(self, toolcall_session, tools):
if not self.is_tools:
logging.warning(f"Model {self.model_config['llm_name']} does not support tool call, but you have assigned one or more tools to it!")
@@ -124,7 +137,7 @@ class LLMBundle(LLM4Tenant):
embeddings, used_tokens = self.mdl.encode(safe_texts)
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))
logging.debug("LLMBundle.encode_queries query: {}, emd len: {}, used_tokens: {}. Builtin model don't need to update token usage".format(texts, len(embeddings), used_tokens))
elif not TenantLLMService.increase_usage_by_id(self.model_config["id"], used_tokens):
logging.error("LLMBundle.encode can't update token usage for <tenant redacted>/EMBEDDING used_tokens: {}".format(used_tokens))

7
api/db/services/task_service.py
View File

@@ -419,6 +419,9 @@ def queue_tasks(doc: dict, bucket: str, name: str, priority: int):
else:
parse_task_array.append(new_task())
# Determine suffix based on parser_id (consistent with SAAS version line 444)
suffix = "common" if doc["parser_id"] != "resume" else "resume"
chunking_config = DocumentService.get_chunking_config(doc["id"])
for task in parse_task_array:
hasher = xxhash.xxh64()
@@ -456,7 +459,7 @@ def queue_tasks(doc: dict, bucket: str, name: str, priority: int):
unfinished_task_array = [task for task in parse_task_array if task["progress"] < 1.0]
for unfinished_task in unfinished_task_array:
assert REDIS_CONN.queue_product(
settings.get_svr_queue_name(priority), message=unfinished_task
settings.get_svr_queue_name(priority, suffix), message=unfinished_task
), "Can't access Redis. Please check the Redis' status."
@@ -547,7 +550,7 @@ def queue_dataflow(tenant_id:str, flow_id:str, task_id:str, doc_id:str=CANVAS_DE
task["file"] = file
if not REDIS_CONN.queue_product(
settings.get_svr_queue_name(priority), message=task
settings.get_svr_queue_name(priority, "common"), message=task
):
return False, "Can't access Redis. Please check the Redis' status."

28
api/db/services/tenant_llm_service.py
View File

@@ -520,3 +520,31 @@ class LLM4Tenant:
except Exception:
# Skip langfuse tracing if connection fails
pass
def close(self):
"""Release resources held by this LLM4Tenant instance.
This method should be called when the instance is no longer needed
to properly release resources such as:
- Langfuse tracing client (flush and shutdown)
- Underlying model instance resources (HTTP sessions, etc.)
"""
# Flush and shutdown Langfuse client if it was initialized
if self.langfuse:
try:
self.langfuse.flush()
if hasattr(self.langfuse, 'shutdown'):
self.langfuse.shutdown()
except Exception:
# Ignore errors during cleanup
pass
finally:
self.langfuse = None
# Release underlying model instance if it has a close method
if self.mdl and hasattr(self.mdl, 'close') and callable(getattr(self.mdl, 'close')):
try:
self.mdl.close()
except Exception:
# Ignore errors during cleanup
pass