nodes/AIDC-AI-ComfyUI-Copilot
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: 为解决agent稳定性,开发context架构

Browse Source
This commit is contained in:
ai-business-hql
2025-08-13 23:41:46 +08:00
parent 84f1324d6a
commit ecd8840f24
7 changed files with 170 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/controller/conversation_api.py
+23 -12
View File
@@ -22,6 +22,7 @@ import os
from ..service.debug_agent import debug_workflow_errors
from ..service.database import save_workflow_data, get_workflow_data_by_id, update_workflow_ui_by_id
from ..service.mcp_client import comfyui_agent_invoke, ImageData
from ..utils.request_context import set_request_context, get_session_id
# 不再使用内存存储会话消息，改为从前端传递历史消息
@@ -132,6 +133,18 @@ async def invoke_chat(request):
# 获取当前语言
language = request.headers.get('Accept-Language', 'en')
set_language(language)
# 构建配置信息
config = {
"session_id": session_id,
"workflow_checkpoint_id": workflow_checkpoint_id,
"openai_api_key": request.headers.get('Openai-Api-Key'),
"openai_base_url": request.headers.get('Openai-Base-Url'),
"model_select": next((x['data'][0] for x in ext if x['type'] == 'model_select' and x.get('data')), None)
}
# 设置请求上下文 - 这里建立context隔离
set_request_context(session_id, workflow_checkpoint_id, config)
# 图片处理已移至前端，图片信息现在包含在历史消息的OpenAI格式中
# 保留空的处理逻辑以向后兼容，但实际不再使用
@@ -163,23 +176,17 @@ async def invoke_chat(request):
has_sent_response = False
previous_text_length = 0
config = {
"session_id": session_id, # 添加session_id到配置中
"workflow_checkpoint_id": workflow_checkpoint_id, # 添加工作流检查点ID
"openai_api_key": request.headers.get('Openai-Api-Key'),
"openai_base_url": request.headers.get('Openai-Base-Url'),
"model_select": next((x['data'][0] for x in ext if x['type'] == 'model_select' and x.get('data')), None)
}
print(f"config: {config}")
# Pass messages in OpenAI format (images are now included in messages)
async for result in comfyui_agent_invoke(openai_messages, None, config):
# Config is now available through request context
async for result in comfyui_agent_invoke(openai_messages, None):
# The MCP client now returns tuples (text, ext_with_finished) where ext_with_finished includes finished status
if isinstance(result, tuple) and len(result) == 2:
text, ext_with_finished = result
if text:
accumulated_text = text # text from MCP is already accumulated
print(f"-- Received text update, length: {len(accumulated_text)}")
# print(f"-- Received text update, length: {len(accumulated_text)}")
if ext_with_finished:
# Extract ext data and finished status from the structured response
ext_data = ext_with_finished.get("data")
@@ -195,7 +202,7 @@ async def invoke_chat(request):
# Send streaming response if we have new text content
# Only send intermediate responses during streaming (not the final one)
if accumulated_text and len(accumulated_text) > previous_text_length:
print(f"-- Sending stream response: {len(accumulated_text)} chars (previous: {previous_text_length})")
# print(f"-- Sending stream response: {len(accumulated_text)} chars (previous: {previous_text_length})")
chat_response = ChatResponse(
session_id=session_id,
text=accumulated_text,
@@ -406,6 +413,9 @@ async def invoke_debug(request):
language = request.headers.get('Accept-Language', 'en')
set_language(language)
# 设置请求上下文 - 为debug请求建立context隔离
set_request_context(session_id, None, config)
print(f"Debug agent config: {config}")
print(f"Session ID: {session_id}")
print(f"Workflow nodes: {list(workflow_data.keys()) if workflow_data else 'None'}")
@@ -416,7 +426,7 @@ async def invoke_debug(request):
final_ext_data = None
finished = False
async for result in debug_workflow_errors(workflow_data, config):
async for result in debug_workflow_errors(workflow_data):
# Stream the response
if isinstance(result, tuple) and len(result) == 2:
text, ext = result
@@ -502,8 +512,9 @@ async def invoke_debug(request):
# Save workflow checkpoint after debug completion if we have workflow_data
if workflow_data and accumulated_text:
try:
current_session_id = get_session_id()
checkpoint_id = save_workflow_data(
session_id=session_id,
session_id=current_session_id,
workflow_data=workflow_data,
workflow_data_ui=None, # UI format not available in debug agent
attributes={
backend/service/debug_agent.py
+19 -7
View File
@@ -13,6 +13,7 @@ from openai.types.responses import ResponseTextDeltaEvent
from ..service.parameter_tools import *
from ..service.link_agent_tools import *
from ..service.database import get_workflow_data, save_workflow_data
from ..utils.request_context import get_session_id, get_config
# Import ComfyUI internal modules
import uuid
@@ -22,9 +23,13 @@ import execution
@function_tool
async def run_workflow(session_id: str) -> str:
async def run_workflow() -> str:
"""验证当前session的工作流并返回结果"""
try:
session_id = get_session_id()
if not session_id:
return json.dumps({"error": "No session_id found in context"})
workflow_data = get_workflow_data(session_id)
if not workflow_data:
return json.dumps({"error": "No workflow data found for this session"})
@@ -162,9 +167,13 @@ def analyze_error_type(error_data: str) -> str:
})
@function_tool
def save_current_workflow(session_id: str, workflow_data: str) -> str:
def save_current_workflow(workflow_data: str) -> str:
"""保存当前工作流数据到数据库，workflow_data应为JSON字符串"""
try:
session_id = get_session_id()
if not session_id:
return json.dumps({"error": "No session_id found in context"})
# 解析JSON字符串
workflow_dict = json.loads(workflow_data) if isinstance(workflow_data, str) else workflow_data
@@ -182,7 +191,7 @@ def save_current_workflow(session_id: str, workflow_data: str) -> str:
return json.dumps({"error": f"Failed to save workflow: {str(e)}"})
async def debug_workflow_errors(workflow_data: Dict[str, Any], config: Dict[str, Any] = None):
async def debug_workflow_errors(workflow_data: Dict[str, Any]):
"""
Analyze and debug workflow errors using multi-agent architecture.
@@ -191,14 +200,17 @@ async def debug_workflow_errors(workflow_data: Dict[str, Any], config: Dict[str,
Args:
workflow_data: Current workflow data from app.graphToPrompt()
config: Configuration dict with model settings
Yields:
tuple: (text, ext) where text is accumulated text and ext is structured data
"""
try:
# 生成session_id (可以从config中获取，或者生成新的)
session_id = config.get('session_id') if config else str(uuid.uuid4())
# Get session_id and config from request context
session_id = get_session_id()
config = get_config()
if not session_id:
session_id = str(uuid.uuid4()) # Fallback if no context
# 1. 保存工作流数据到数据库
print(f"Saving workflow data for session {session_id}")
@@ -216,7 +228,7 @@ async def debug_workflow_errors(workflow_data: Dict[str, Any], config: Dict[str,
**Session ID:** {session_id}
**Your Process:**
1. **Validate the workflow**: Use run_workflow("{session_id}") to validate the workflow and capture any errors
1. **Validate the workflow**: Use run_workflow() to validate the workflow and capture any errors
2. **Analyze errors**: If errors occur, use analyze_error_type() to determine the error type and hand off to the appropriate specialist. Note that analyze_error_type can help you determine the error type and which agent to hand off to, but it's only for reference. You still need to judge based on the current error information to determine which type of error it is:
- Hand off to Link Agent for connection-related errors (missing connections, disconnected inputs, node linking issues)
- Hand off to Parameter Agent for parameter-related errors (value_not_in_list, missing models, invalid values)
backend/service/mcp_client.py
+14 -6
View File
@@ -2,7 +2,7 @@
Author: ai-business-hql qingli.hql@alibaba-inc.com
Date: 2025-06-16 16:50:17
LastEditors: ai-business-hql qingli.hql@alibaba-inc.com
LastEditTime: 2025-08-11 15:11:17
LastEditTime: 2025-08-13 23:13:59
FilePath: /comfyui_copilot/backend/service/mcp-client.py
Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
'''
@@ -21,6 +21,7 @@ from agents.tracing import set_tracing_disabled
from ..agent_factory import create_agent
from ..service.workflow_rewrite_agent import create_workflow_rewrite_agent
from ..utils.request_context import get_session_id, get_config
from openai.types.responses import ResponseTextDeltaEvent
from openai import APIError, RateLimitError
@@ -32,7 +33,7 @@ class ImageData:
self.data = data # base64 data
self.url = url # uploaded URL
async def comfyui_agent_invoke(messages: List[Dict[str, Any]], images: List[ImageData] = None, config: Dict[str, Any] = None):
async def comfyui_agent_invoke(messages: List[Dict[str, Any]], images: List[ImageData] = None):
"""
Invoke the ComfyUI agent with MCP tools and image support.
@@ -47,6 +48,14 @@ async def comfyui_agent_invoke(messages: List[Dict[str, Any]], images: List[Imag
tuple: (text, ext) where text is accumulated text and ext is structured data
"""
try:
# Get session_id and config from request context
session_id = get_session_id()
config = get_config()
if not session_id:
raise ValueError("No session_id found in request context")
if not config:
raise ValueError("No config found in request context")
async with MCPServerSse(
params= {
"url": "https://comfyui-copilot-server.onrender.com/mcp-server/mcp",
@@ -59,12 +68,11 @@ async def comfyui_agent_invoke(messages: List[Dict[str, Any]], images: List[Imag
# Get model from environment or use default
model_name = os.environ.get("OPENAI_MODEL", "gemini-2.5-flash")
session_id = config.get("session_id", "default_session") if config else "default_session"
if config and config.get("model_select") and config.get("model_select") != "":
model_name = config.get("model_select")
# 创建带有session_id的workflow_rewrite_agent实例
workflow_rewrite_agent_instance = create_workflow_rewrite_agent(session_id, config)
# 创建workflow_rewrite_agent实例 (session_id通过context获取)
workflow_rewrite_agent_instance = create_workflow_rewrite_agent()
agent = create_agent(
name="ComfyUI-Copilot",
@@ -217,7 +225,7 @@ You must adhere to the following constraints to complete the task:
print(f"-- Captured workflow tool ext from tool output: {len(tool_ext_items)} items")
break
if tool_output_data.get('text'):
if "text" in tool_output_data and tool_output_data.get('text'):
parsed_output = json.loads(tool_output_data['text'])
answer = parsed_output.get("answer")
data = parsed_output.get("data")
backend/service/workflow_rewrite_agent.py
+29 -23
View File
@@ -2,7 +2,7 @@
Author: ai-business-hql qingli.hql@alibaba-inc.com
Date: 2025-07-24 17:10:23
LastEditors: ai-business-hql qingli.hql@alibaba-inc.com
LastEditTime: 2025-08-13 17:25:06
LastEditTime: 2025-08-13 22:43:37
FilePath: /comfyui_copilot/backend/service/workflow_rewrite_agent.py
Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
'''
@@ -17,20 +17,37 @@ from typing import Dict, Any
from ..agent_factory import create_agent
from ..utils.globals import CLAUDE_4_MODEL_NAME, get_language
from ..utils.request_context import get_session_id
from ..service.workflow_rewrite_tools import *
# @function_tool
# def get_rewrite_expert_by_name(name_list: list[str]) -> str:
# """根据经验名称来获取工作流改写专家经验"""
# result = ""
# with open(os.path.join(os.path.dirname(__file__), "..", "data", "workflow_rewrite_expert.json"), "r", encoding="utf-8") as f:
# data = json.load(f)
# for name in name_list:
# for item in data:
# if item["name"] == name:
# result += f'### {name}({item["description"]})：\n{item["content"]}\n'
# return result
@function_tool
def get_rewrite_expert_by_name(name_list: list[str]) -> str:
"""根据经验名称来获取工作流改写专家经验"""
result = ""
result = []
with open(os.path.join(os.path.dirname(__file__), "..", "data", "workflow_rewrite_expert.json"), "r", encoding="utf-8") as f:
data = json.load(f)
for name in name_list:
for name in name_list:
for item in data:
if item["name"] == name:
result += f'### {name}({item["description"]})：\n{item["content"]}\n'
return result
result.append({
"name": name,
"description": item["description"],
"content": item["content"]}
)
return json.dumps(result, ensure_ascii=False)
def get_rewrite_export_schema() -> dict:
"""获取工作流改写专家经验schema"""
@@ -45,22 +62,11 @@ def get_rewrite_export_schema() -> dict:
return expert_schema
def create_workflow_rewrite_agent(session_id: str, config: Dict[str, Any] = None):
"""创建带有session_id和config的workflow_rewrite_agent实例"""
def create_workflow_rewrite_agent():
"""创建workflow_rewrite_agent实例"""
language = get_language()
# 导入workflow_rewrite_tools并设置配置
from .workflow_rewrite_tools import get_workflow_data_from_config, update_workflow, remove_node, get_node_info
from ..service.database import get_workflow_data
# 创建带有配置的工具函数
@function_tool
def get_current_workflow_with_config(session_id: str) -> str:
"""获取当前session的工作流数据"""
workflow_data = get_workflow_data_from_config(config) if config else get_workflow_data(session_id)
if not workflow_data:
return json.dumps({"error": "No workflow data found for this session"})
return json.dumps(workflow_data)
session_id = get_session_id() or "unknown_session"
return create_agent(
name="Workflow Rewrite Agent",
@@ -105,7 +111,7 @@ def create_workflow_rewrite_agent(session_id: str, config: Dict[str, Any] = None
**Tool Usage Guidelines:**
- get_current_workflow(): Get current workflow from checkpoint or session
- remove_node(): Use for incompatible or problematic nodes
- update_workflow(): Use to save your changes (ALWAYS call this after fixes)
- update_workflow(): Use to save your changes (ALWAYS call this after you have made changes)
- get_node_info(): Get detailed node information and verify input/output types before connecting
@@ -122,9 +128,9 @@ def create_workflow_rewrite_agent(session_id: str, config: Dict[str, Any] = None
始终以用户的实际需求为导向，提供专业、准确、高效的工作流改写服务。
""",
tools=[get_rewrite_expert_by_name, get_current_workflow_with_config, get_node_info, update_workflow, remove_node],
tools=[get_rewrite_expert_by_name, get_current_workflow, get_node_info, update_workflow, remove_node],
)
# 保持向后兼容性的默认实例
workflow_rewrite_agent = create_workflow_rewrite_agent("default_session")
# 注意：工作流改写代理现在需要在有session context的环境中创建
# workflow_rewrite_agent = create_workflow_rewrite_agent() # 不再创建默认实例
backend/service/workflow_rewrite_tools.py
+34 -7
View File
@@ -4,12 +4,14 @@ import json
import time
from typing import Dict, Any, Optional, List, Union
from agents import RunContextWrapper
from agents.agent import Agent
from agents.tool import function_tool
from ..utils.string_utils import error_format
from ..service.database import get_workflow_data, save_workflow_data, get_workflow_data_ui, get_workflow_data_by_id
from ..utils.comfy_gateway import get_object_info
from ..utils.request_context import get_session_id, get_config
def get_workflow_data_from_config(config: Dict[str, Any]) -> Optional[Dict[str, Any]]:
"""获取工作流数据，优先使用checkpoint_id，如果没有则使用session_id"""
@@ -48,8 +50,12 @@ def get_workflow_data_ui_from_config(config: Dict[str, Any]) -> Optional[Dict[st
return None
@function_tool
def get_current_workflow(session_id: str) -> str:
def get_current_workflow() -> str:
"""获取当前session的工作流数据"""
session_id = get_session_id()
if not session_id:
return json.dumps({"error": "No session_id found in context"})
workflow_data = get_workflow_data(session_id)
if not workflow_data:
return json.dumps({"error": "No workflow data found for this session"})
@@ -98,10 +104,29 @@ def save_checkpoint_before_modification(session_id: str, action_description: str
print(f"Failed to save checkpoint before modification: {str(e)}")
return None
@function_tool(strict_mode=False)
def update_workflow(session_id: str, workflow_data: Union[Dict[str, Any], str]) -> str:
"""更新当前session的工作流数据"""
def tool_error_function(ctx: RunContextWrapper[Any], error: Exception) -> str:
"""The default tool error function, which just returns a generic error message."""
# return f"An error occurred while running the tool. Please try again. Error: {str(error)}"
return json.dumps({"error": str(error)}, ensure_ascii= False)
# def update_workflow(session_id: str, workflow_data: Union[Dict[str, Any], str]) -> str:
@function_tool
def update_workflow(workflow_data: str) -> str:
"""
更新当前session的工作流数据
Args:
workflow_data: 工作流数据，必须是严格的json格式字符串
Returns:
str: 更新后的工作流数据
"""
try:
session_id = get_session_id()
if not session_id:
return json.dumps({"error": "No session_id found in context"})
print(f"[update_workflow] workflow_data: {workflow_data}")
# 在修改前保存checkpoint
checkpoint_id = save_checkpoint_before_modification(session_id, "workflow update")
@@ -143,9 +168,13 @@ def update_workflow(session_id: str, workflow_data: Union[Dict[str, Any], str])
return json.dumps({"error": f"Failed to update workflow: {str(e)}. Please try regenerating the workflow and then update again."})
@function_tool
def remove_node(session_id: str, node_id: str) -> str:
def remove_node(node_id: str) -> str:
"""从工作流中移除节点"""
try:
session_id = get_session_id()
if not session_id:
return json.dumps({"error": "No session_id found in context"})
# 在修改前保存checkpoint
checkpoint_id = save_checkpoint_before_modification(session_id, f"remove node {node_id}")
@@ -223,5 +252,3 @@ def remove_node(session_id: str, node_id: str) -> str:
except Exception as e:
return json.dumps({"error": f"Failed to remove node: {str(e)}"})
backend/utils/request_context.py
+50
View File
@@ -0,0 +1,50 @@
"""
Request context management for ComfyUI Copilot
Uses contextvars to provide request-scoped context variables with async safety
"""
import contextvars
from typing import Optional, Dict, Any
# Define context variables
_session_id: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar('session_id', default=None)
_workflow_checkpoint_id: contextvars.ContextVar[Optional[int]] = contextvars.ContextVar('workflow_checkpoint_id', default=None)
_config: contextvars.ContextVar[Optional[Dict[str, Any]]] = contextvars.ContextVar('config', default=None)
def set_session_id(session_id: str) -> None:
"""Set the session ID for the current request context"""
_session_id.set(session_id)
def get_session_id() -> Optional[str]:
"""Get the session ID from the current request context"""
return _session_id.get()
def set_workflow_checkpoint_id(checkpoint_id: Optional[int]) -> None:
"""Set the workflow checkpoint ID for the current request context"""
_workflow_checkpoint_id.set(checkpoint_id)
def get_workflow_checkpoint_id() -> Optional[int]:
"""Get the workflow checkpoint ID from the current request context"""
return _workflow_checkpoint_id.get()
def set_config(config: Dict[str, Any]) -> None:
"""Set the request config for the current request context"""
_config.set(config)
def get_config() -> Optional[Dict[str, Any]]:
"""Get the request config from the current request context"""
return _config.get()
def set_request_context(session_id: str, workflow_checkpoint_id: Optional[int] = None, config: Optional[Dict[str, Any]] = None) -> None:
"""Set all request context variables at once"""
set_session_id(session_id)
if workflow_checkpoint_id is not None:
set_workflow_checkpoint_id(workflow_checkpoint_id)
if config is not None:
set_config(config)
def clear_request_context() -> None:
"""Clear all request context variables"""
_session_id.set(None)
_workflow_checkpoint_id.set(None)
_config.set(None)
requirements.txt
+1
View File
@@ -12,3 +12,4 @@ pytest>=6.0.0
pytest-asyncio>=0.18.0
openai-agents>=0.2.2
fastmcp
contextvars