547 lines
21 KiB
Python
547 lines
21 KiB
Python
import argparse
|
||
import asyncio
|
||
import json
|
||
import threading
|
||
import traceback
|
||
import logging
|
||
from typing import Any, Dict, Iterable, AsyncIterable, AsyncGenerator, Optional
|
||
import cozeloop
|
||
import uvicorn
|
||
import time
|
||
from fastapi import FastAPI, HTTPException, Request
|
||
from fastapi.responses import StreamingResponse, JSONResponse
|
||
from langchain_core.runnables import RunnableConfig
|
||
from langgraph.graph import StateGraph, END
|
||
from langgraph.graph.state import CompiledStateGraph
|
||
from coze_coding_utils.runtime_ctx.context import new_context, Context
|
||
from coze_coding_utils.helper import graph_helper
|
||
from coze_coding_utils.log.node_log import LOG_FILE
|
||
from coze_coding_utils.log.write_log import setup_logging, request_context
|
||
from coze_coding_utils.log.config import LOG_LEVEL
|
||
from coze_coding_utils.error.classifier import ErrorClassifier, classify_error
|
||
from coze_coding_utils.helper.stream_runner import AgentStreamRunner, WorkflowStreamRunner,agent_stream_handler,workflow_stream_handler, RunOpt
|
||
|
||
setup_logging(
|
||
log_file=LOG_FILE,
|
||
max_bytes=100 * 1024 * 1024, # 100MB
|
||
backup_count=5,
|
||
log_level=LOG_LEVEL,
|
||
use_json_format=True,
|
||
console_output=True
|
||
)
|
||
|
||
logger = logging.getLogger(__name__)
|
||
from coze_coding_utils.helper.agent_helper import to_stream_input
|
||
from coze_coding_utils.openai.handler import OpenAIChatHandler
|
||
from coze_coding_utils.log.parser import LangGraphParser
|
||
from coze_coding_utils.log.err_trace import extract_core_stack
|
||
from coze_coding_utils.log.loop_trace import init_run_config, init_agent_config
|
||
|
||
|
||
# 超时配置常量
|
||
TIMEOUT_SECONDS = 900 # 15分钟
|
||
|
||
class GraphService:
|
||
def __init__(self):
|
||
# 用于跟踪正在运行的任务（使用asyncio.Task）
|
||
self.running_tasks: Dict[str, asyncio.Task] = {}
|
||
# 错误分类器
|
||
self.error_classifier = ErrorClassifier()
|
||
# stream runner
|
||
self._agent_stream_runner = AgentStreamRunner()
|
||
self._workflow_stream_runner = WorkflowStreamRunner()
|
||
self._graph = None
|
||
self._graph_lock = threading.Lock()
|
||
|
||
def _get_graph(self, ctx=Context):
|
||
if graph_helper.is_agent_proj():
|
||
return graph_helper.get_agent_instance("agents.agent", ctx)
|
||
|
||
if self._graph is not None:
|
||
return self._graph
|
||
with self._graph_lock:
|
||
if self._graph is not None:
|
||
return self._graph
|
||
self._graph = graph_helper.get_graph_instance("graphs.graph")
|
||
return self._graph
|
||
|
||
@staticmethod
|
||
def _sse_event(data: Any, event_id: Any = None) -> str:
|
||
id_line = f"id: {event_id}\n" if event_id else ""
|
||
return f"{id_line}event: message\ndata: {json.dumps(data, ensure_ascii=False, default=str)}\n\n"
|
||
|
||
def _get_stream_runner(self):
|
||
if graph_helper.is_agent_proj():
|
||
return self._agent_stream_runner
|
||
else:
|
||
return self._workflow_stream_runner
|
||
|
||
# 流式运行（原始迭代器）：本地调用使用
|
||
def stream(self, payload: Dict[str, Any], run_config: RunnableConfig, ctx=Context) -> Iterable[Any]:
|
||
graph = self._get_graph(ctx)
|
||
stream_runner = self._get_stream_runner()
|
||
for chunk in stream_runner.stream(payload, graph, run_config, ctx):
|
||
yield chunk
|
||
|
||
# 同步运行：本地/HTTP 通用
|
||
async def run(self, payload: Dict[str, Any], ctx=None) -> Dict[str, Any]:
|
||
if ctx is None:
|
||
ctx = new_context("run")
|
||
|
||
run_id = ctx.run_id
|
||
logger.info(f"Starting run with run_id: {run_id}")
|
||
|
||
try:
|
||
graph = self._get_graph(ctx)
|
||
# custom tracer
|
||
run_config = init_run_config(graph, ctx)
|
||
run_config["configurable"] = {"thread_id": ctx.run_id}
|
||
|
||
# 直接调用，LangGraph会在当前任务上下文中执行
|
||
# 如果当前任务被取消，LangGraph的执行也会被取消
|
||
return await graph.ainvoke(payload, config=run_config, context=ctx)
|
||
|
||
except asyncio.CancelledError:
|
||
logger.info(f"Run {run_id} was cancelled")
|
||
return {"status": "cancelled", "run_id": run_id, "message": "Execution was cancelled"}
|
||
except Exception as e:
|
||
# 使用错误分类器分类错误
|
||
err = self.error_classifier.classify(e, {"node_name": "run", "run_id": run_id})
|
||
# 记录详细的错误信息和堆栈跟踪
|
||
logger.error(
|
||
f"Error in GraphService.run: [{err.code}] {err.message}\n"
|
||
f"Category: {err.category.name}\n"
|
||
f"Traceback:\n{extract_core_stack()}"
|
||
)
|
||
# 保留原始异常堆栈，便于上层返回真正的报错位置
|
||
raise
|
||
finally:
|
||
# 清理任务记录
|
||
self.running_tasks.pop(run_id, None)
|
||
|
||
# 流式运行（SSE 格式化）：HTTP 路由使用
|
||
async def stream_sse(self, payload: Dict[str, Any], ctx=None, run_opt: Optional[RunOpt] = None) -> AsyncGenerator[str, None]:
|
||
if ctx is None:
|
||
ctx = new_context(method="stream_sse")
|
||
if run_opt is None:
|
||
run_opt = RunOpt()
|
||
|
||
run_id = ctx.run_id
|
||
logger.info(f"Starting stream with run_id: {run_id}")
|
||
graph = self._get_graph(ctx)
|
||
if graph_helper.is_agent_proj():
|
||
run_config = init_agent_config(graph, ctx)
|
||
else:
|
||
run_config = init_run_config(graph, ctx) # vibeflow
|
||
|
||
is_workflow = not graph_helper.is_agent_proj()
|
||
|
||
try:
|
||
async for chunk in self.astream(payload, graph, run_config=run_config, ctx=ctx, run_opt=run_opt):
|
||
if is_workflow and isinstance(chunk, tuple):
|
||
event_id, data = chunk
|
||
yield self._sse_event(data, event_id)
|
||
else:
|
||
yield self._sse_event(chunk)
|
||
finally:
|
||
# 清理任务记录
|
||
self.running_tasks.pop(run_id, None)
|
||
cozeloop.flush()
|
||
|
||
# 取消执行 - 使用asyncio的标准方式
|
||
def cancel_run(self, run_id: str, ctx: Optional[Context] = None) -> Dict[str, Any]:
|
||
"""
|
||
取消指定run_id的执行
|
||
|
||
使用asyncio.Task.cancel()来取消任务,这是标准的Python异步取消机制。
|
||
LangGraph会在节点之间检查CancelledError,实现优雅的取消。
|
||
"""
|
||
logger.info(f"Attempting to cancel run_id: {run_id}")
|
||
|
||
# 查找对应的任务
|
||
if run_id in self.running_tasks:
|
||
task = self.running_tasks[run_id]
|
||
if not task.done():
|
||
# 使用asyncio的标准取消机制
|
||
# 这会在下一个await点抛出CancelledError
|
||
task.cancel()
|
||
logger.info(f"Cancellation requested for run_id: {run_id}")
|
||
return {
|
||
"status": "success",
|
||
"run_id": run_id,
|
||
"message": "Cancellation signal sent, task will be cancelled at next await point"
|
||
}
|
||
else:
|
||
logger.info(f"Task already completed for run_id: {run_id}")
|
||
return {
|
||
"status": "already_completed",
|
||
"run_id": run_id,
|
||
"message": "Task has already completed"
|
||
}
|
||
else:
|
||
logger.warning(f"No active task found for run_id: {run_id}")
|
||
return {
|
||
"status": "not_found",
|
||
"run_id": run_id,
|
||
"message": "No active task found with this run_id. Task may have already completed or run_id is invalid."
|
||
}
|
||
|
||
# 运行指定节点：本地/HTTP 通用
|
||
async def run_node(self, node_id: str, payload: Dict[str, Any], ctx=None) -> Any:
|
||
if ctx is None or Context.run_id == "":
|
||
ctx = new_context(method="node_run")
|
||
|
||
_graph = self._get_graph()
|
||
node_func, input_cls, output_cls = graph_helper.get_graph_node_func_with_inout(_graph.get_graph(), node_id)
|
||
if node_func is None or input_cls is None:
|
||
raise KeyError(f"node_id '{node_id}' not found")
|
||
|
||
parser = LangGraphParser(_graph)
|
||
metadata = parser.get_node_metadata(node_id) or {}
|
||
|
||
_g = StateGraph(input_cls, input_schema=input_cls, output_schema=output_cls)
|
||
_g.add_node("sn", node_func, metadata=metadata)
|
||
_g.set_entry_point("sn")
|
||
_g.add_edge("sn", END)
|
||
_graph = _g.compile()
|
||
|
||
run_config = init_run_config(_graph, ctx)
|
||
return await _graph.ainvoke(payload, config=run_config)
|
||
|
||
def graph_inout_schema(self) -> Any:
|
||
if graph_helper.is_agent_proj():
|
||
return {"input_schema": {}, "output_schema": {}}
|
||
builder = getattr(self._get_graph(), 'builder', None)
|
||
if builder is not None:
|
||
input_cls = getattr(builder, 'input_schema', None) or self.graph.get_input_schema()
|
||
output_cls = getattr(builder, 'output_schema', None) or self.graph.get_output_schema()
|
||
else:
|
||
logger.warning(f"No builder input schema found for graph_inout_schema, using graph input schema instead")
|
||
input_cls = self.graph.get_input_schema()
|
||
output_cls = self.graph.get_output_schema()
|
||
|
||
return {
|
||
"input_schema": input_cls.model_json_schema(),
|
||
"output_schema": output_cls.model_json_schema(),
|
||
"code":0,
|
||
"msg":""
|
||
}
|
||
|
||
async def astream(self, payload: Dict[str, Any], graph: CompiledStateGraph, run_config: RunnableConfig, ctx=Context, run_opt: Optional[RunOpt] = None) -> AsyncIterable[Any]:
|
||
stream_runner = self._get_stream_runner()
|
||
async for chunk in stream_runner.astream(payload, graph, run_config, ctx, run_opt):
|
||
yield chunk
|
||
|
||
|
||
service = GraphService()
|
||
app = FastAPI()
|
||
|
||
# OpenAI 兼容接口处理器
|
||
openai_handler = OpenAIChatHandler(service)
|
||
|
||
|
||
HEADER_X_RUN_ID = "x-run-id"
|
||
@app.post("/run")
|
||
async def http_run(request: Request) -> Dict[str, Any]:
|
||
global result
|
||
raw_body = await request.body()
|
||
try:
|
||
body_text = raw_body.decode("utf-8")
|
||
except Exception as e:
|
||
body_text = str(raw_body)
|
||
raise HTTPException(status_code=400,
|
||
detail=f"Invalid JSON format: {body_text}, traceback: {traceback.format_exc()}, error: {e}")
|
||
|
||
ctx = new_context(method="run", headers=request.headers)
|
||
# 优先使用上游指定的 run_id，保证 cancel 能精确匹配
|
||
upstream_run_id = request.headers.get(HEADER_X_RUN_ID)
|
||
if upstream_run_id:
|
||
ctx.run_id = upstream_run_id
|
||
run_id = ctx.run_id
|
||
request_context.set(ctx)
|
||
|
||
logger.info(
|
||
f"Received request for /run: "
|
||
f"run_id={run_id}, "
|
||
f"query={dict(request.query_params)}, "
|
||
f"body={body_text}"
|
||
)
|
||
|
||
try:
|
||
payload = await request.json()
|
||
|
||
# 创建任务并记录 - 这是关键，让我们可以通过run_id取消任务
|
||
task = asyncio.create_task(service.run(payload, ctx))
|
||
service.running_tasks[run_id] = task
|
||
|
||
try:
|
||
result = await asyncio.wait_for(task, timeout=float(TIMEOUT_SECONDS))
|
||
except asyncio.TimeoutError:
|
||
logger.error(f"Run execution timeout after {TIMEOUT_SECONDS}s for run_id: {run_id}")
|
||
task.cancel()
|
||
try:
|
||
result = await task
|
||
except asyncio.CancelledError:
|
||
return {
|
||
"status": "timeout",
|
||
"run_id": run_id,
|
||
"message": f"Execution timeout: exceeded {TIMEOUT_SECONDS} seconds"
|
||
}
|
||
|
||
if not result:
|
||
result = {}
|
||
if isinstance(result, dict):
|
||
result["run_id"] = run_id
|
||
return result
|
||
|
||
except json.JSONDecodeError as e:
|
||
logger.error(f"JSON decode error in http_run: {e}, traceback: {traceback.format_exc()}")
|
||
raise HTTPException(status_code=400, detail=f"Invalid JSON format, {extract_core_stack()}")
|
||
|
||
except asyncio.CancelledError:
|
||
logger.info(f"Request cancelled for run_id: {run_id}")
|
||
result = {"status": "cancelled", "run_id": run_id, "message": "Execution was cancelled"}
|
||
return result
|
||
|
||
except Exception as e:
|
||
# 使用错误分类器获取错误信息
|
||
error_response = service.error_classifier.get_error_response(e, {"node_name": "http_run", "run_id": run_id})
|
||
logger.error(
|
||
f"Unexpected error in http_run: [{error_response['error_code']}] {error_response['error_message']}, "
|
||
f"traceback: {traceback.format_exc()}", exc_info=True
|
||
)
|
||
raise HTTPException(
|
||
status_code=500,
|
||
detail={
|
||
"error_code": error_response["error_code"],
|
||
"error_message": error_response["error_message"],
|
||
"stack_trace": extract_core_stack(),
|
||
}
|
||
)
|
||
finally:
|
||
cozeloop.flush()
|
||
|
||
|
||
HEADER_X_WORKFLOW_STREAM_MODE = "x-workflow-stream-mode"
|
||
|
||
|
||
def _register_task(run_id: str, task: asyncio.Task):
|
||
service.running_tasks[run_id] = task
|
||
|
||
|
||
@app.post("/stream_run")
|
||
async def http_stream_run(request: Request):
|
||
ctx = new_context(method="stream_run", headers=request.headers)
|
||
# 优先使用上游指定的 run_id，保证 cancel 能精确匹配
|
||
upstream_run_id = request.headers.get(HEADER_X_RUN_ID)
|
||
if upstream_run_id:
|
||
ctx.run_id = upstream_run_id
|
||
workflow_stream_mode = request.headers.get(HEADER_X_WORKFLOW_STREAM_MODE, "").lower()
|
||
workflow_debug = workflow_stream_mode == "debug"
|
||
request_context.set(ctx)
|
||
raw_body = await request.body()
|
||
try:
|
||
body_text = raw_body.decode("utf-8")
|
||
except Exception as e:
|
||
body_text = str(raw_body)
|
||
raise HTTPException(status_code=400,
|
||
detail=f"Invalid JSON format: {body_text}, traceback: {extract_core_stack()}, error: {e}")
|
||
run_id = ctx.run_id
|
||
is_agent = graph_helper.is_agent_proj()
|
||
logger.info(
|
||
f"Received request for /stream_run: "
|
||
f"run_id={run_id}, "
|
||
f"is_agent_project={is_agent}, "
|
||
f"query={dict(request.query_params)}, "
|
||
f"body={body_text}"
|
||
)
|
||
try:
|
||
payload = await request.json()
|
||
except json.JSONDecodeError as e:
|
||
logger.error(f"JSON decode error in http_stream_run: {e}, traceback: {traceback.format_exc()}")
|
||
raise HTTPException(status_code=400, detail=f"Invalid JSON format:{extract_core_stack()}")
|
||
|
||
if is_agent:
|
||
stream_generator = agent_stream_handler(
|
||
payload=payload,
|
||
ctx=ctx,
|
||
run_id=run_id,
|
||
stream_sse_func=service.stream_sse,
|
||
sse_event_func=service._sse_event,
|
||
error_classifier=service.error_classifier,
|
||
register_task_func=_register_task,
|
||
)
|
||
else:
|
||
stream_generator = workflow_stream_handler(
|
||
payload=payload,
|
||
ctx=ctx,
|
||
run_id=run_id,
|
||
stream_sse_func=service.stream_sse,
|
||
sse_event_func=service._sse_event,
|
||
error_classifier=service.error_classifier,
|
||
register_task_func=_register_task,
|
||
run_opt=RunOpt(workflow_debug=workflow_debug),
|
||
)
|
||
|
||
response = StreamingResponse(stream_generator, media_type="text/event-stream")
|
||
return response
|
||
|
||
@app.post("/cancel/{run_id}")
|
||
async def http_cancel(run_id: str, request: Request):
|
||
"""
|
||
取消指定run_id的执行
|
||
|
||
使用asyncio.Task.cancel()实现取消,这是Python标准的异步任务取消机制。
|
||
LangGraph会在节点之间的await点检查CancelledError,实现优雅取消。
|
||
"""
|
||
ctx = new_context(method="cancel", headers=request.headers)
|
||
request_context.set(ctx)
|
||
logger.info(f"Received cancel request for run_id: {run_id}")
|
||
result = service.cancel_run(run_id, ctx)
|
||
return result
|
||
|
||
|
||
@app.post(path="/node_run/{node_id}")
|
||
async def http_node_run(node_id: str, request: Request):
|
||
raw_body = await request.body()
|
||
try:
|
||
body_text = raw_body.decode("utf-8")
|
||
except UnicodeDecodeError:
|
||
body_text = str(raw_body)
|
||
raise HTTPException(status_code=400, detail=f"Invalid JSON format: {body_text}")
|
||
ctx = new_context(method="node_run", headers=request.headers)
|
||
request_context.set(ctx)
|
||
logger.info(
|
||
f"Received request for /node_run/{node_id}: "
|
||
f"query={dict(request.query_params)}, "
|
||
f"body={body_text}",
|
||
)
|
||
|
||
try:
|
||
payload = await request.json()
|
||
except json.JSONDecodeError as e:
|
||
logger.error(f"JSON decode error in http_node_run: {e}, traceback: {traceback.format_exc()}")
|
||
raise HTTPException(status_code=400, detail=f"Invalid JSON format:{extract_core_stack()}")
|
||
try:
|
||
return await service.run_node(node_id, payload, ctx)
|
||
except KeyError:
|
||
raise HTTPException(status_code=404,
|
||
detail=f"node_id '{node_id}' not found or input miss required fields, traceback: {extract_core_stack()}")
|
||
except Exception as e:
|
||
# 使用错误分类器获取错误信息
|
||
error_response = service.error_classifier.get_error_response(e, {"node_name": node_id})
|
||
logger.error(
|
||
f"Unexpected error in http_node_run: [{error_response['error_code']}] {error_response['error_message']}, "
|
||
f"traceback: {traceback.format_exc()}", exc_info=True
|
||
)
|
||
raise HTTPException(
|
||
status_code=500,
|
||
detail={
|
||
"error_code": error_response["error_code"],
|
||
"error_message": error_response["error_message"],
|
||
"stack_trace": extract_core_stack(),
|
||
}
|
||
)
|
||
finally:
|
||
cozeloop.flush()
|
||
|
||
|
||
@app.post("/v1/chat/completions")
|
||
async def openai_chat_completions(request: Request):
|
||
"""OpenAI Chat Completions API 兼容接口"""
|
||
ctx = new_context(method="openai_chat", headers=request.headers)
|
||
request_context.set(ctx)
|
||
|
||
logger.info(f"Received request for /v1/chat/completions: run_id={ctx.run_id}")
|
||
|
||
try:
|
||
payload = await request.json()
|
||
return await openai_handler.handle(payload, ctx)
|
||
except json.JSONDecodeError as e:
|
||
logger.error(f"JSON decode error in openai_chat_completions: {e}")
|
||
raise HTTPException(status_code=400, detail="Invalid JSON format")
|
||
finally:
|
||
cozeloop.flush()
|
||
|
||
|
||
@app.get("/health")
|
||
async def health_check():
|
||
try:
|
||
# 这里可以添加更多的健康检查逻辑
|
||
return {
|
||
"status": "ok",
|
||
"message": "Service is running",
|
||
}
|
||
except Exception as e:
|
||
raise HTTPException(status_code=503, detail=str(e))
|
||
|
||
|
||
@app.get(path="/graph_parameter")
|
||
async def http_graph_inout_parameter(request: Request):
|
||
return service.graph_inout_schema()
|
||
|
||
def parse_args():
|
||
parser = argparse.ArgumentParser(description="Start FastAPI server")
|
||
parser.add_argument("-m", type=str, default="http", help="Run mode, support http,flow,node")
|
||
parser.add_argument("-n", type=str, default="", help="Node ID for single node run")
|
||
parser.add_argument("-p", type=int, default=5000, help="HTTP server port")
|
||
parser.add_argument("-i", type=str, default="", help="Input JSON string for flow/node mode")
|
||
return parser.parse_args()
|
||
|
||
|
||
def parse_input(input_str: str) -> Dict[str, Any]:
|
||
"""Parse input string, support both JSON string and plain text"""
|
||
if not input_str:
|
||
return {"text": "你好"}
|
||
|
||
# Try to parse as JSON first
|
||
try:
|
||
return json.loads(input_str)
|
||
except json.JSONDecodeError:
|
||
# If not valid JSON, treat as plain text
|
||
return {"text": input_str}
|
||
|
||
def start_http_server(port):
|
||
workers = 1
|
||
reload = False
|
||
if graph_helper.is_dev_env():
|
||
reload = True
|
||
|
||
logger.info(f"Start HTTP Server, Port: {port}, Workers: {workers}")
|
||
uvicorn.run("main:app", host="0.0.0.0", port=port, reload=reload, workers=workers)
|
||
|
||
if __name__ == "__main__":
|
||
args = parse_args()
|
||
if args.m == "http":
|
||
start_http_server(args.p)
|
||
elif args.m == "flow":
|
||
payload = parse_input(args.i)
|
||
result = asyncio.run(service.run(payload))
|
||
print(json.dumps(result, ensure_ascii=False, indent=2))
|
||
elif args.m == "node" and args.n:
|
||
payload = parse_input(args.i)
|
||
result = asyncio.run(service.run_node(args.n, payload))
|
||
print(json.dumps(result, ensure_ascii=False, indent=2))
|
||
elif args.m == "agent":
|
||
agent_ctx = new_context(method="agent")
|
||
for chunk in service.stream(
|
||
{
|
||
"type": "query",
|
||
"session_id": "1",
|
||
"message": "你好",
|
||
"content": {
|
||
"query": {
|
||
"prompt": [
|
||
{
|
||
"type": "text",
|
||
"content": {"text": "现在几点了？请调用工具获取当前时间"},
|
||
}
|
||
]
|
||
}
|
||
},
|
||
},
|
||
run_config={"configurable": {"session_id": "1"}},
|
||
ctx=agent_ctx,
|
||
):
|
||
print(chunk)
|