AgentCoord/backend/AgentCoord/PlanEngine/branch_TaskProcess.py

from AgentCoord.util.converter import read_LLM_Completion
import AgentCoord.util as util
from typing import List
from pydantic import BaseModel, Field
import json


ACT_SET = """
- Propose (Propose something that contribute to the current task):
  - Description Format Example:
    Propose some suggestion for the development of the story from emotional aspects.

- Critique (Provide feedback to the action result of other agents):
  - Description Format Example:
    Critique the logic soundness of the story outline written by Alice.

- Improve (Improve the result of a previous action):
  - Description Format Example:
    Improve the story outline written by Bob based on the feedback by Jordan"

- Finalize (Deliver the final result based on previous actions):
  - Description Format Example:
    Summarize the discussion and submit the final version of the Story Outline.

"""

PROMPT_TASK_PROCESS_BRANCHING = """
## Instruction
Based on "Existing Steps", your task is to comeplete the "Remaining Steps" for the "Task for Current Step".
Note: "Modification Requirement" specifies how to modify the "Baseline Completion" for a better/alternative solution.

**IMPORTANT LANGUAGE REQUIREMENT: You must respond in Chinese (中文) for the Description field and all explanations, while keeping ID, ActionType, and AgentName in their original format.**

## General Goal (The general goal for the collaboration plan, you just design the plan for one of its step (i.e. "Task for Current Step"))
{General_Goal}

## Task for Current Step
{Current_Task_Description}

## Action Set (The list of Action Type Available)
{Act_Set}

## Existing Steps
{Existing_Steps}

## Baseline Completion
{Baseline_Completion}

## Modification Requirement
{Modification_Requirement}

## Output Format Example (Specify the output format)
```json
{{
    "Remaining Steps": [
    {{
      "ID": "Action4",
      "ActionType": "Propose",
      "AgentName": "Mia",
      "Description": "提议关于人工智能情感发展的心理学理论，重点关注爱与依恋的概念。",
      "ImportantInput": [
        "InputObject:Story Outline"
      ]
    }},
    {{
      "ID": "Action5",
      "ActionType": "Critique",
      "AgentName": "Noah",
      "Description": "对Mia提出的心理学理论进行批判性评估，分析其在AI情感发展场景中的适用性和局限性。",
      "ImportantInput": [
        "ActionResult:Action4"
      ]
    }},
    {{
      "ID": "Action6",
      "ActionType": "Improve",
      "AgentName": "Liam",
      "Description": "基于Noah的批判性反馈，改进和完善心理学理论框架，使其更贴合AI情感发展的实际需求。",
      "ImportantInput": [
        "ActionResult:Action4",
        "ActionResult:Action5"
      ]
    }},
    {{
      "ID": "Action7",
      "ActionType": "Finalize",
      "AgentName": "Mia",
      "Description": "综合所有提议、批判和改进意见，整合并提交最终的AI情感发展心理学理论框架。",
      "ImportantInput": [
        "ActionResult:Action4",
        "ActionResult:Action5",
        "ActionResult:Action6"
      ]
    }}
  ]
}}

## Format Explaination (Explain the Output Format):
ImportantInput: Specify if there is any previous result that should be taken special consideration during the execution the action. Should be of format "InputObject:xx" or "ActionResult:xx".
InputObject_List: List existing objects that should be utilized in current step.
AgentName: Specify the agent who will perform the action, You CAN ONLY USE THE NAME APPEARS IN "AgentInvolved".
ActionType: Specify the type of action. **CRITICAL REQUIREMENTS:**
  1. The "Remaining Steps" MUST include ALL FOUR action types in the following order: Propose -> Critique -> Improve -> Finalize
  2. Each action type (Propose, Critique, Improve, Finalize) MUST appear at least once
  3. The actions must follow the sequence: Propose actions first, then Critique actions, then Improve actions, and Finalize must be the last action
  4. Even if only one agent is involved in a phase, that phase must still have its corresponding action type
  5. The last action must ALWAYS be of type "Finalize"

"""


class JSON_Step(BaseModel):
    ID: str
    ActionType: str
    AgentName: str
    Description: str
    ImportantInput: List[str]


class JSON_TASK_PROCESS_BRANCHING(BaseModel):
    Remaining_Steps: List[JSON_Step] = Field(..., alias="Remaining Steps")


def branch_TaskProcess(
    branch_Number,
    Modification_Requirement,
    Existing_Steps,
    Baseline_Completion,
    stepTaskExisting,
    General_Goal,
    AgentProfile_Dict,
):
    Current_Task_Description = {
        "TaskName": stepTaskExisting["StepName"],
        "AgentInvolved": [
            {"Name": name, "Profile": AgentProfile_Dict[name]}
            for name in stepTaskExisting["AgentSelection"]
        ],
        "InputObject_List": stepTaskExisting["InputObject_List"],
        "OutputObject": stepTaskExisting["OutputObject"],
        "CurrentTaskDescription": util.generate_template_sentence_for_CollaborationBrief(
            stepTaskExisting["InputObject_List"],
            stepTaskExisting["OutputObject"],
            stepTaskExisting["AgentSelection"],
            stepTaskExisting["TaskContent"],
        ),
    }
    prompt = PROMPT_TASK_PROCESS_BRANCHING.format(
        Modification_Requirement=Modification_Requirement,
        Current_Task_Description=json.dumps(
            Current_Task_Description, indent=4
        ),
        Existing_Steps=json.dumps(Existing_Steps, indent=4),
        Baseline_Completion=json.dumps(Baseline_Completion, indent=4),
        General_Goal=General_Goal,
        Act_Set=ACT_SET,
    )
    print(prompt)
    branch_List = []
    for i in range(branch_Number):
        messages = [
            {
                "role": "system",
                "content": f" The JSON object must use the schema: {json.dumps(JSON_TASK_PROCESS_BRANCHING.model_json_schema(), indent=2)}",
            },
            {"role": "system", "content": prompt},
        ]
        Remaining_Steps = read_LLM_Completion(messages, useGroq=False)[
            "Remaining Steps"
        ]

        branch_List.append(Remaining_Steps)
    return branch_List