mirror of
https://github.com/langgenius/dify.git
synced 2026-07-15 01:17:04 +08:00
69 lines
2.5 KiB
Python
69 lines
2.5 KiB
Python
"""Celery tasks that execute Agent backend lifecycle-only session cleanup."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
|
|
from celery import shared_task
|
|
|
|
from clients.agent_backend.factory import create_agent_backend_run_client
|
|
from clients.agent_backend.request_builder import AgentBackendRunRequestBuilder
|
|
from clients.agent_backend.session_cleanup import (
|
|
AgentBackendSessionCleanupPayload,
|
|
cleanup_agent_backend_session,
|
|
)
|
|
from configs import dify_config
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _create_agent_backend_client():
|
|
if not (dify_config.AGENT_BACKEND_USE_FAKE or dify_config.AGENT_BACKEND_BASE_URL):
|
|
return None
|
|
return create_agent_backend_run_client(
|
|
base_url=dify_config.AGENT_BACKEND_BASE_URL,
|
|
use_fake=dify_config.AGENT_BACKEND_USE_FAKE,
|
|
fake_scenario=dify_config.AGENT_BACKEND_FAKE_SCENARIO,
|
|
)
|
|
|
|
|
|
def _run_cleanup_task(payload_dict: dict[str, object]) -> None:
|
|
payload = AgentBackendSessionCleanupPayload.model_validate(payload_dict)
|
|
result = cleanup_agent_backend_session(
|
|
payload=payload,
|
|
client=_create_agent_backend_client(),
|
|
request_builder=AgentBackendRunRequestBuilder(),
|
|
)
|
|
if result.status == "succeeded":
|
|
return
|
|
|
|
log_fields = {
|
|
"tenant_id": payload.metadata.get("tenant_id"),
|
|
"app_id": payload.metadata.get("app_id"),
|
|
"workflow_run_id": payload.metadata.get("workflow_run_id"),
|
|
"node_id": payload.metadata.get("node_id"),
|
|
"conversation_id": payload.metadata.get("conversation_id"),
|
|
"agent_id": payload.metadata.get("agent_id"),
|
|
"previous_agent_backend_run_id": payload.metadata.get("previous_agent_backend_run_id"),
|
|
"failed_agent_backend_run_id": payload.metadata.get("failed_agent_backend_run_id"),
|
|
"cleanup_run_id": result.cleanup_run_id,
|
|
"reason": result.reason,
|
|
}
|
|
if result.status == "skipped":
|
|
logger.info("Agent backend session cleanup skipped: %s", log_fields)
|
|
return
|
|
|
|
logger.warning("Agent backend session cleanup failed: %s", log_fields)
|
|
|
|
|
|
@shared_task(queue="workflow_storage")
|
|
def cleanup_workflow_agent_runtime_session(payload_dict: dict[str, object]) -> None:
|
|
"""Run one workflow-owned Agent backend cleanup payload."""
|
|
_run_cleanup_task(payload_dict)
|
|
|
|
|
|
@shared_task(queue="conversation")
|
|
def cleanup_conversation_agent_runtime_session(payload_dict: dict[str, object]) -> None:
|
|
"""Run one conversation-owned Agent backend cleanup payload."""
|
|
_run_cleanup_task(payload_dict)
|