AgentCoord/backend/server.py

import nest_asyncio
nest_asyncio.apply()
import os, sys, functools
print = functools.partial(print, flush=True)   # 全局 flush
sys.stdout.reconfigure(line_buffering=True)    # 3.7+ 有效
import asyncio

from flask import Flask, request, jsonify
import json
from DataProcess import Add_Collaboration_Brief_FrontEnd
from AgentCoord.RehearsalEngine_V2.ExecutePlan import executePlan
from AgentCoord.PlanEngine.basePlan_Generator import generate_basePlan
from AgentCoord.PlanEngine.fill_stepTask import fill_stepTask
from AgentCoord.PlanEngine.fill_stepTask_TaskProcess import (
    fill_stepTask_TaskProcess,
)
from AgentCoord.PlanEngine.branch_PlanOutline import branch_PlanOutline
from AgentCoord.PlanEngine.branch_TaskProcess import branch_TaskProcess
from AgentCoord.PlanEngine.AgentSelectModify import (
    AgentSelectModify_init,
    AgentSelectModify_addAspect,
)
import os
import yaml
import argparse

# initialize global variables
yaml_file = os.path.join(os.getcwd(), "config", "config.yaml")
yaml_data = {}
try:
    with open(yaml_file, "r", encoding="utf-8") as file:
        yaml_data = yaml.safe_load(file)
except Exception:
    yaml_data = {}
USE_CACHE: bool = os.getenv("USE_CACHE")
if USE_CACHE is None:
    USE_CACHE = yaml_data.get("USE_CACHE", False)
else:
    USE_CACHE = USE_CACHE.lower() in ["true", "1", "yes"]
AgentBoard = None
AgentProfile_Dict = {}
Request_Cache: dict[str, str] = {}
app = Flask(__name__)


from jsonschema import validate, ValidationError

AGENT_SELECTION_SCHEMA = {
    "type": "object",
    "properties": {
        "AgentSelectionPlan": {
            "type": "array",
            "items": {
                "type": "string",
                "minLength": 1,      # 不允许空字符串
                "pattern": r"^\S+$"  # 不允许仅空白
            },
            "minItems": 1          # 至少选一个
        }
    },
    "required": ["AgentSelectionPlan"],
    "additionalProperties": False
}


BASE_PLAN_SCHEMA = {
    "type": "object",
    "properties": {
        "Plan_Outline": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "StepName":        {"type": "string"},
                    "TaskContent":     {"type": "string"},
                    "InputObject_List":{"type": "array", "items": {"type": "string"}},
                    "OutputObject":    {"type": "string"},
                },
                "required": ["StepName", "TaskContent", "InputObject_List", "OutputObject"],
                "additionalProperties": False,
            },
        }
    },
    "required": ["Plan_Outline"],
    "additionalProperties": False,
}


def safe_join(iterable, sep=""):
    """保证 join 前全是 str，None 变空串"""
    return sep.join("" if x is None else str(x) for x in iterable)

def clean_agent_board(board):
    """把 AgentBoard 洗成只含 str 的字典列表"""
    if not board:
        return []
    return [
        {"Name": (a.get("Name") or "").strip(),
         "Profile": (a.get("Profile") or "").strip()}
        for a in board
        if a and a.get("Name")
    ]

def clean_plan_outline(outline):
    """清洗 Plan_Outline 里的 None"""
    if not isinstance(outline, list):
        return []
    for step in outline:
        if not isinstance(step, dict):
            continue
        step["InputObject_List"] = [
            str(i) for i in step.get("InputObject_List", []) if i is not None
        ]
        step["OutputObject"] = str(step.get("OutputObject") or "")
        step["StepName"] = str(step.get("StepName") or "")
        step["TaskContent"] = str(step.get("TaskContent") or "")
    return outline

@app.route("/fill_stepTask_TaskProcess", methods=["post"])
def Handle_fill_stepTask_TaskProcess():
    incoming_data = request.get_json()
    # print(f"[DEBUG] fill_stepTask_TaskProcess received data: {incoming_data}", flush=True)
    
    # 验证必需参数
    General_Goal = incoming_data.get("General Goal", "").strip()
    stepTask_lackTaskProcess = incoming_data.get("stepTask_lackTaskProcess")
    
    if not General_Goal:
        return jsonify({"error": "General Goal is required and cannot be empty"}), 400
    if not stepTask_lackTaskProcess:
        return jsonify({"error": "stepTask_lackTaskProcess is required"}), 400
    
    requestIdentifier = str(
        (
            "/fill_stepTask_TaskProcess",
            General_Goal,
            stepTask_lackTaskProcess,
        )
    )

    if USE_CACHE:
        if requestIdentifier in Request_Cache:
            return jsonify(Request_Cache[requestIdentifier])

    try:
        filled_stepTask = fill_stepTask_TaskProcess(
            General_Goal=General_Goal,
            stepTask=stepTask_lackTaskProcess,
            AgentProfile_Dict=AgentProfile_Dict,
        )
    except Exception as e:
        print(f"[ERROR] fill_stepTask_TaskProcess failed: {e}", flush=True)
        return jsonify({"error": str(e)}), 500
    filled_stepTask = Add_Collaboration_Brief_FrontEnd(filled_stepTask)
    Request_Cache[requestIdentifier] = filled_stepTask
    response = jsonify(filled_stepTask)
    return response


@app.route("/agentSelectModify_init", methods=["POST"])
def Handle_agentSelectModify_init():
    incoming = request.get_json(silent=True) or {}
    general_goal = (incoming.get("General Goal") or "").strip()
    step_task   = incoming.get("stepTask")
    if not general_goal or not step_task:
        return jsonify({"error": "Missing field"}), 400

    if not AgentBoard:                      # 空 Board 直接返回
        return jsonify({"AgentSelectionPlan": []})

    req_id = str(("/agentSelectModify_init", general_goal, step_task))
    if USE_CACHE and req_id in Request_Cache:
        return jsonify(Request_Cache[req_id])

    try:
        clean_board = clean_agent_board(AgentBoard)
        raw = AgentSelectModify_init(stepTask=step_task,
                                   General_Goal=general_goal,
                                   Agent_Board=clean_board)
        if not isinstance(raw, dict):
            raise ValueError("model returned non-dict")
        plan = raw.get("AgentSelectionPlan") or []
        cleaned = [str(x).strip() for x in plan if x is not None and str(x).strip()]
        raw["AgentSelectionPlan"] = cleaned
        validate(instance=raw, schema=AGENT_SELECTION_SCHEMA)
    except Exception as exc:
        print(f"[ERROR] AgentSelectModify_init: {exc}")
        return jsonify({"error": str(exc)}), 500

    if USE_CACHE:
        Request_Cache[req_id] = raw
    return jsonify(raw)


@app.route("/agentSelectModify_addAspect", methods=["post"])
def Handle_agentSelectModify_addAspect():
    incoming_data = request.get_json()
    requestIdentifier = str(
        ("/agentSelectModify_addAspect", incoming_data["aspectList"])
    )

    if USE_CACHE:
        if requestIdentifier in Request_Cache:
            return jsonify(Request_Cache[requestIdentifier])

    scoreTable = AgentSelectModify_addAspect(
        aspectList=incoming_data["aspectList"], Agent_Board=AgentBoard or []
    )
    Request_Cache[requestIdentifier] = scoreTable
    response = jsonify(scoreTable)
    return response


@app.route("/fill_stepTask", methods=["post"])
def Handle_fill_stepTask():
    incoming_data = request.get_json()
    # print(f"[DEBUG] fill_stepTask received data: {incoming_data}", flush=True)
    
    # 验证必需参数
    General_Goal = incoming_data.get("General Goal", "").strip()
    stepTask = incoming_data.get("stepTask")
    
    if not General_Goal:
        return jsonify({"error": "General Goal is required and cannot be empty"}), 400
    if not stepTask:
        return jsonify({"error": "stepTask is required"}), 400
    
    requestIdentifier = str(
        (
            "/fill_stepTask",
            General_Goal,
            stepTask,
        )
    )

    if USE_CACHE:
        if requestIdentifier in Request_Cache:
            return jsonify(Request_Cache[requestIdentifier])

    try:
        filled_stepTask = fill_stepTask(
            General_Goal=General_Goal,
            stepTask=stepTask,
            Agent_Board=AgentBoard,
            AgentProfile_Dict=AgentProfile_Dict,
        )
    except Exception as e:
        print(f"[ERROR] fill_stepTask failed: {e}", flush=True)
        return jsonify({"error": str(e)}), 500
    filled_stepTask = Add_Collaboration_Brief_FrontEnd(filled_stepTask)
    Request_Cache[requestIdentifier] = filled_stepTask
    response = jsonify(filled_stepTask)
    return response


@app.route("/branch_PlanOutline", methods=["post"])
def Handle_branch_PlanOutline():
    incoming_data = request.get_json()
    requestIdentifier = str(
        (
            "/branch_PlanOutline",
            incoming_data["branch_Number"],
            incoming_data["Modification_Requirement"],
            incoming_data["Existing_Steps"],
            incoming_data["Baseline_Completion"],
            incoming_data["Initial Input Object"],
            incoming_data["General Goal"],
        )
    )

    if USE_CACHE:
        if requestIdentifier in Request_Cache:
            return jsonify(Request_Cache[requestIdentifier])

    branchList = branch_PlanOutline(
        branch_Number=incoming_data["branch_Number"],
        Modification_Requirement=incoming_data["Modification_Requirement"],
        Existing_Steps=incoming_data["Existing_Steps"],
        Baseline_Completion=incoming_data["Baseline_Completion"],
        InitialObject_List=incoming_data["Initial Input Object"],
        General_Goal=incoming_data["General Goal"],
    )
    branchList = Add_Collaboration_Brief_FrontEnd(branchList)
    Request_Cache[requestIdentifier] = branchList
    response = jsonify(branchList)
    return response


@app.route("/branch_TaskProcess", methods=["post"])
def Handle_branch_TaskProcess():
    incoming_data = request.get_json()
    requestIdentifier = str(
        (
            "/branch_TaskProcess",
            incoming_data["branch_Number"],
            incoming_data["Modification_Requirement"],
            incoming_data["Existing_Steps"],
            incoming_data["Baseline_Completion"],
            incoming_data["stepTaskExisting"],
            incoming_data["General Goal"],
        )
    )

    if USE_CACHE:
        if requestIdentifier in Request_Cache:
            return jsonify(Request_Cache[requestIdentifier])

    branchList = branch_TaskProcess(
        branch_Number=incoming_data["branch_Number"],
        Modification_Requirement=incoming_data["Modification_Requirement"],
        Existing_Steps=incoming_data["Existing_Steps"],
        Baseline_Completion=incoming_data["Baseline_Completion"],
        stepTaskExisting=incoming_data["stepTaskExisting"],
        General_Goal=incoming_data["General Goal"],
        AgentProfile_Dict=AgentProfile_Dict,
    )
    Request_Cache[requestIdentifier] = branchList
    response = jsonify(branchList)
    return response


@app.route("/generate_basePlan", methods=["POST"])
def Handle_generate_basePlan():
    incoming = request.get_json(silent=True) or {}
    general_goal = (incoming.get("General Goal") or "").strip()
    initial_objs = incoming.get("Initial Input Object") or []

    if not general_goal:
        return jsonify({"error": "General Goal is required"}), 400

    # 1. 空 Board 直接短路
    if not AgentBoard:
        print("[SKIP] AgentBoard empty")
        out = Add_Collaboration_Brief_FrontEnd({"Plan_Outline": []})
        return jsonify(out)

    req_id = str(("/generate_basePlan", general_goal, initial_objs))
    if USE_CACHE and req_id in Request_Cache:
        return jsonify(Request_Cache[req_id])

    try:
        # 2. 洗 Board → 调模型 → 洗返回
        clean_board = clean_agent_board(AgentBoard)
        raw_plan = asyncio.run(
            generate_basePlan(
                General_Goal=general_goal,
                Agent_Board=clean_board,
                AgentProfile_Dict=AgentProfile_Dict,
                InitialObject_List=initial_objs,
            )
        )
        raw_plan["Plan_Outline"] = clean_plan_outline(raw_plan.get("Plan_Outline"))
        validate(instance=raw_plan, schema=BASE_PLAN_SCHEMA)   # 可选，二次校验
    except Exception as exc:
        print(f"[ERROR] generate_basePlan failed: {exc}")
        return jsonify({"error": "model call failed", "detail": str(exc)}), 500

    out = Add_Collaboration_Brief_FrontEnd(raw_plan)
    if USE_CACHE:
        Request_Cache[req_id] = out
    return jsonify(out)


@app.route("/executePlan", methods=["post"])
def Handle_executePlan():
    incoming_data = request.get_json()
    requestIdentifier = str(
        (
            "/executePlan",
            incoming_data["num_StepToRun"],
            incoming_data["RehearsalLog"],
            incoming_data["plan"],
        )
    )

    if USE_CACHE:
        if requestIdentifier in Request_Cache:
            return jsonify(Request_Cache[requestIdentifier])

    RehearsalLog = executePlan(
        incoming_data["plan"],
        incoming_data["num_StepToRun"],
        incoming_data["RehearsalLog"],
        AgentProfile_Dict,
    )
    Request_Cache[requestIdentifier] = RehearsalLog
    response = jsonify(RehearsalLog)
    return response


@app.route("/_saveRequestCashe", methods=["post"])
def Handle_saveRequestCashe():
    with open(
        os.path.join(os.getcwd(), "RequestCache", "Request_Cache.json"), "w"
    ) as json_file:
        json.dump(Request_Cache, json_file, indent=4)
    response = jsonify(
        {"code": 200, "content": "request cashe sucessfully saved"}
    )
    return response


@app.route("/setAgents", methods=["POST"])
def set_agents():
    global AgentBoard, AgentProfile_Dict
    board_in = request.json or []
    # 先清洗再赋值
    AgentBoard = clean_agent_board(board_in)
    AgentProfile_Dict = {a["Name"]: a["Profile"] for a in AgentBoard}
    return jsonify({"code": 200, "content": "AgentBoard set successfully"})


def init():
    global AgentBoard, AgentProfile_Dict, Request_Cache
    with open(
        os.path.join(os.getcwd(), "RequestCache", "Request_Cache.json"), "r"
    ) as json_file:
        Request_Cache = json.load(json_file)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="start the backend for AgentCoord"
    )
    parser.add_argument(
        "--port",
        type=int,
        default=8017,
        help="set the port number, 8017 by defaul.",
    )
    args = parser.parse_args()
    init()
    app.run(host="0.0.0.0", port=args.port, debug=True)