♻️ Refactor move some business logic from agent_app.py to agent_service.py …

Phinease · web-flow · commit da32c0cf7477 · 2025-08-21T09:15:00.000+08:00
diff --git a/backend/apps/agent_app.py b/backend/apps/agent_app.py
@@ -2,22 +2,13 @@
 from typing import Optional
 
 from fastapi import HTTPException, APIRouter, Header, Request, Body
-from fastapi.responses import StreamingResponse, JSONResponse
-from nexent.core.agents.run_agent import agent_run
-
-from database.agent_db import delete_related_agent
-from utils.auth_utils import get_current_user_info, get_current_user_id
-from agents.create_agent_info import create_agent_run_info
+from fastapi.responses import JSONResponse
 from consts.model import AgentRequest, AgentInfoRequest, AgentIDRequest, ConversationResponse, AgentImportRequest
 from services.agent_service import get_agent_info_impl, \
     get_creating_sub_agent_info_impl, update_agent_info_impl, delete_agent_impl, export_agent_impl, import_agent_impl, \
-    list_all_agent_info_impl, insert_related_agent_impl
-from services.conversation_management_service import save_conversation_user, save_conversation_assistant
-from services.memory_config_service import build_memory_context
-from utils.config_utils import config_manager
-from utils.thread_utils import submit
-from agents.agent_run_manager import agent_run_manager
-from agents.preprocess_manager import preprocess_manager
+    list_all_agent_info_impl, insert_related_agent_impl, run_agent_stream, stop_agent_tasks
+from database.agent_db import delete_related_agent
+from utils.auth_utils import get_current_user_info, get_current_user_id
 
 
 router = APIRouter(prefix="/agent")
@@ -30,43 +21,10 @@ async def agent_run_api(agent_request: AgentRequest, http_request: Request, auth
     """
     Agent execution API endpoint
     """
-    user_id, tenant_id, language = get_current_user_info(authorization, http_request)
-    memory_context = build_memory_context(user_id, tenant_id, agent_request.agent_id)
-
-    agent_run_info = await create_agent_run_info(agent_id=agent_request.agent_id,
-                                                 minio_files=agent_request.minio_files,
-                                                 query=agent_request.query,
-                                                 history=agent_request.history,
-                                                 authorization=authorization,
-                                                 language=language)
-
-    agent_run_manager.register_agent_run(agent_request.conversation_id, agent_run_info)
-    # Save user message only if not in debug mode
-    if not agent_request.is_debug:
-        submit(save_conversation_user, agent_request, authorization)
-
-    async def generate():
-        messages = []
-        try:
-            async for chunk in agent_run(agent_run_info, memory_context):
-                messages.append(chunk)
-                yield f"data: {chunk}\n\n"
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Agent run error: {str(e)}")
-        finally:
-            # Save assistant message only if not in debug mode
-            if not agent_request.is_debug:
-                submit(save_conversation_assistant, agent_request, messages, authorization)
-            # Unregister agent run instance for both debug and non-debug modes
-            agent_run_manager.unregister_agent_run(agent_request.conversation_id)
-
-    return StreamingResponse(
-        generate(),
-        media_type="text/event-stream",
-        headers={
-            "Cache-Control": "no-cache",
-            "Connection": "keep-alive"
-        }
+    return await run_agent_stream(
+        agent_request=agent_request,
+        http_request=http_request,
+        authorization=authorization
     )
 
 
@@ -75,21 +33,8 @@ async def agent_stop_api(conversation_id: int):
     """
     stop agent run and preprocess tasks for specified conversation_id
     """
-    # Stop agent run
-    agent_stopped = agent_run_manager.stop_agent_run(conversation_id)
-    
-    # Stop preprocess tasks
-    preprocess_stopped = preprocess_manager.stop_preprocess_tasks(conversation_id)
-    
-    if agent_stopped or preprocess_stopped:
-        message_parts = []
-        if agent_stopped:
-            message_parts.append("agent run")
-        if preprocess_stopped:
-            message_parts.append("preprocess tasks")
-        
-        message = f"successfully stopped {' and '.join(message_parts)} for conversation_id {conversation_id}"
-        return {"status": "success", "message": message}
+    if stop_agent_tasks(conversation_id).get("status") == "success":
+        return {"status": "success", "message": "agent run and preprocess tasks stopped successfully"}
     else:
         raise HTTPException(status_code=404, detail=f"no running agent or preprocess tasks found for conversation_id {conversation_id}")
 
diff --git a/backend/apps/conversation_management_app.py b/backend/apps/conversation_management_app.py
@@ -2,8 +2,6 @@
 from typing import Dict, Any, Optional
 
 from fastapi import HTTPException, APIRouter, Header, Request
-from fastapi.encoders import jsonable_encoder
-from pydantic import BaseModel
 
 from consts.model import ConversationResponse, ConversationRequest, RenameRequest, GenerateTitleRequest, OpinionRequest, MessageIdRequest
 from services.conversation_management_service import (
diff --git a/backend/database/agent_db.py b/backend/database/agent_db.py
@@ -25,6 +25,20 @@ def search_agent_info_by_agent_id(agent_id: int, tenant_id: str):
 
         return agent_dict
 
+def search_agent_id_by_agent_name(agent_name: str, tenant_id: str):
+    """
+    Search agent id by agent name
+    """
+    with get_db_session() as session:
+        agent = session.query(AgentInfo).filter(
+            AgentInfo.name == agent_name,
+            AgentInfo.tenant_id == tenant_id,
+            AgentInfo.delete_flag != 'Y').first()
+        if not agent:
+            raise ValueError("agent not found")
+        return agent.agent_id
+
+
 def search_blank_sub_agent_by_main_agent_id(tenant_id: str):
     """
     Search blank sub agent by main agent id
diff --git a/backend/services/agent_service.py b/backend/services/agent_service.py
@@ -3,22 +3,30 @@
 import logging
 from collections import deque
 
-from fastapi import Header
-from fastapi.responses import JSONResponse
+from fastapi import Header, Request, HTTPException
+from fastapi.responses import JSONResponse, StreamingResponse
+from consts.model import AgentRequest
 from agents.create_agent_info import create_tool_config_list
 from consts.model import AgentInfoRequest, ExportAndImportAgentInfo, ExportAndImportDataFormat, ToolInstanceInfoRequest, MCPInfo
 from database.agent_db import create_agent, query_all_enabled_tool_instances, \
      search_blank_sub_agent_by_main_agent_id, \
     search_tools_for_sub_agent, search_agent_info_by_agent_id, update_agent, delete_agent_by_id, query_all_tools, \
     create_or_update_tool_by_tool_info, check_tool_is_available, query_all_agent_info_by_tenant_id, \
-    query_sub_agents_id_list, insert_related_agent, delete_all_related_agent
+    query_sub_agents_id_list, insert_related_agent, delete_all_related_agent, search_agent_id_by_agent_name
 from database.remote_mcp_db import get_mcp_server_by_name_and_tenant, check_mcp_name_exists
 from services.remote_mcp_service import add_remote_mcp_server_list
 from services.tool_configuration_service import update_tool_list
+from services.conversation_management_service import save_conversation_user, save_conversation_assistant
 
-from utils.auth_utils import get_current_user_id
+from utils.auth_utils import get_current_user_info
 from utils.memory_utils import build_memory_config
+from utils.thread_utils import submit
 from nexent.memory.memory_service import clear_memory
+from nexent.core.agents.run_agent import agent_run
+from services.memory_config_service import build_memory_context
+from agents.create_agent_info import create_agent_run_info
+from agents.agent_run_manager import agent_run_manager
+from agents.preprocess_manager import preprocess_manager
 
 
 logger = logging.getLogger("agent_service")
@@ -69,7 +77,7 @@ def get_agent_info_impl(agent_id: int, tenant_id: str):
 
 
 def get_creating_sub_agent_info_impl(authorization: str = Header(None)):
-    user_id, tenant_id = get_current_user_id(authorization)
+    user_id, tenant_id, _ = get_current_user_info(authorization)
     
     try:
         sub_agent_id = get_creating_sub_agent_id_service(tenant_id, user_id)
@@ -100,7 +108,7 @@ def get_creating_sub_agent_info_impl(authorization: str = Header(None)):
             "sub_agent_id_list": query_sub_agents_id_list(main_agent_id=sub_agent_id, tenant_id=tenant_id)}
 
 def update_agent_info_impl(request: AgentInfoRequest, authorization: str = Header(None)):
-    user_id, tenant_id = get_current_user_id(authorization)
+    user_id, tenant_id, _ = get_current_user_info(authorization)
     
     try:
         update_agent(request.agent_id, request, tenant_id, user_id)
@@ -109,7 +117,7 @@ def update_agent_info_impl(request: AgentInfoRequest, authorization: str = Heade
         raise ValueError(f"Failed to update agent info: {str(e)}")
 
 async def delete_agent_impl(agent_id: int, authorization: str = Header(None)):
-    user_id, tenant_id = get_current_user_id(authorization)
+    user_id, tenant_id, _ = get_current_user_info(authorization)
 
     try:
         delete_agent_by_id(agent_id, tenant_id, user_id)
@@ -183,7 +191,7 @@ async def export_agent_impl(agent_id: int, authorization: str = Header(None)) ->
         This function recursively finds all managed sub-agents and exports the detailed configuration of each agent (including tools, prompts, etc.) as a dictionary, and finally returns it as a formatted JSON string for frontend download and backup.
     """
 
-    user_id, tenant_id = get_current_user_id(authorization)
+    user_id, tenant_id, _ = get_current_user_info(authorization)
 
     export_agent_dict = {}
     search_list = deque([agent_id])
@@ -251,7 +259,7 @@ async def import_agent_impl(agent_info: ExportAndImportDataFormat, authorization
     """
     Import agent using DFS
     """
-    user_id, tenant_id = get_current_user_id(authorization)
+    user_id, tenant_id, _ = get_current_user_info(authorization)
     agent_id = agent_info.agent_id
 
     # First, add MCP servers if any
@@ -459,4 +467,121 @@ def insert_related_agent_impl(parent_agent_id, child_agent_id, tenant_id):
         return JSONResponse(
             status_code=400,
             content={"message":"Failed to insert relation", "status": "error"}
-        )
+        )
+
+
+# Helper function for run_agent_stream, used to prepare context for an agent run
+async def prepare_agent_run(agent_request: AgentRequest, http_request: Request, authorization: str):
+    """
+    Prepare for an agent run by creating context and run info, and registering the run.
+    """
+    user_id, tenant_id, language = get_current_user_info(authorization, http_request)
+
+    memory_context = build_memory_context(user_id, tenant_id, agent_request.agent_id)
+    agent_run_info = await create_agent_run_info(agent_id=agent_request.agent_id,
+                                                 minio_files=agent_request.minio_files,
+                                                 query=agent_request.query,
+                                                 history=agent_request.history,
+                                                 authorization=authorization,
+                                                 language=language)
+    agent_run_manager.register_agent_run(agent_request.conversation_id, agent_run_info)
+    return agent_run_info, memory_context
+
+
+# Helper function for run_agent_stream, used to save messages for either user or assistant
+def save_messages(agent_request, target:str, messages=None, authorization=None):
+    if target == "user":
+        if messages is not None:
+            raise ValueError("Messages should be None when saving for user.")
+        submit(save_conversation_user, agent_request, authorization)
+    elif target == "assistant":
+        if messages is None:
+            raise ValueError("Messages cannot be None when saving for assistant.")
+        submit(save_conversation_assistant, agent_request, messages, authorization)
+
+
+# Helper function for run_agent_stream, used to generate stream response
+async def generate_stream(agent_run_info, memory_context, agent_request: AgentRequest, authorization: str):
+    messages = []
+    try:
+        async for chunk in agent_run(agent_run_info, memory_context):
+            messages.append(chunk)
+            yield f"data: {chunk}\n\n"
+    except Exception as e:
+        logger.error(f"Agent run error: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Agent run error: {str(e)}")
+    finally:
+        # Save assistant message only if not in debug mode
+        if not agent_request.is_debug:
+            save_messages(agent_request, target="assistant", messages=messages, authorization=authorization)
+        # Unregister agent run instance for both debug and non-debug modes
+        agent_run_manager.unregister_agent_run(agent_request.conversation_id)
+
+
+async def run_agent_stream(agent_request: AgentRequest, http_request: Request, authorization: str):
+    """
+    Start an agent run and stream responses, using explicit user/tenant context.
+    Mirrors the logic of agent_app.agent_run_api but reusable by services.
+    """
+    agent_run_info, memory_context = await prepare_agent_run(
+        agent_request=agent_request,
+        http_request=http_request,
+        authorization=authorization
+    )
+
+    # Save user message only if not in debug mode
+    if not agent_request.is_debug:
+        save_messages(
+            agent_request, 
+            target="user", 
+            authorization=authorization
+        )
+
+    return StreamingResponse(
+        generate_stream(agent_run_info, memory_context, agent_request, authorization),
+        media_type="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive"
+        }
+    )
+
+
+def stop_agent_tasks(conversation_id: int):
+    """
+    Stop agent run and preprocess tasks for the specified conversation_id.
+    Matches the behavior of agent_app.agent_stop_api.
+    """
+    # Stop agent run
+    agent_stopped = agent_run_manager.stop_agent_run(conversation_id)
+
+    # Stop preprocess tasks
+    preprocess_stopped = preprocess_manager.stop_preprocess_tasks(conversation_id)
+
+    if agent_stopped or preprocess_stopped:
+        message_parts = []
+        if agent_stopped:
+            message_parts.append("agent run")
+        if preprocess_stopped:
+            message_parts.append("preprocess tasks")
+
+        message = f"successfully stopped {' and '.join(message_parts)} for conversation_id {conversation_id}"
+        logging.info(message)
+        return {"status": "success", "message": message}
+    else:
+        message = f"no running agent or preprocess tasks found for conversation_id {conversation_id}"
+        logging.error(message)
+        return {"status": "error", "message": message}
+
+
+def get_agent_id_by_name(agent_name: str, tenant_id: str) -> int:
+    """
+    Resolve unique agent id by its unique name under the same tenant.
+    """
+    if not agent_name:
+        raise HTTPException(status_code=400, detail="agent_name required")
+    try:
+        return search_agent_id_by_agent_name(agent_name, tenant_id)
+    except Exception as _:
+        logger.error(f"Failed to find agent id with '{agent_name}' in tenant {tenant_id}")
+        raise HTTPException(status_code=404, detail="agent not found")
diff --git a/sdk/nexent/vector_database/elasticsearch_core.py b/sdk/nexent/vector_database/elasticsearch_core.py
@@ -59,7 +59,7 @@ def __init__(
             api_key=self.api_key,
             verify_certs=verify_certs,
             ssl_show_warn=ssl_show_warn,
-            timeout=20,
+            request_timeout=20,
             max_retries=3,  # Reduce retries for faster failure detection
             retry_on_timeout=True,
             retry_on_status=[502, 503, 504],  # Retry on these status codes,
diff --git a/test/backend/app/test_agent_app.py b/test/backend/app/test_agent_app.py
diff --git a/test/backend/services/test_agent_service.py b/test/backend/services/test_agent_service.py
