AgentWorld/examples/planner_coder_reviewer.py at main · ScienceIntelligence/AgentWorld · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
from __future__ import annotations

import json

from agentworld import AgentGraph, DefaultOperator, StaticController
from agentworld.controller.base import ControllerEvent
from agentworld.graph.reducers import append_list


def planner_script(_request):
    return [
        ControllerEvent(
            kind="message_completed",
            payload={
                "kind": "plan",
                "text": "Break the work into implementation and review steps.",
            },
        ),
        ControllerEvent(
            kind="completed",
            payload={
                "state_patch": {
                    "steps": ["implement core package", "review graph result"],
                    "status_log": ["planner completed"],
                }
            },
        ),
    ]


def coder_script(request):
    payload = json.loads(request.instruction)
    return [
        ControllerEvent(
            kind="message_completed",
            payload={
                "kind": "observation",
                "text": (
                    f"Coder received {len(payload['skills'])} skills "
                    f"and {len(request.metadata)} metadata keys and produced a draft."
                ),
            },
        ),
        ControllerEvent(
            kind="artifact_created",
            payload={
                "kind": "note",
                "path": "artifacts/draft.txt",
                "description": "Mock draft produced by coder.",
            },
        ),
        ControllerEvent(
            kind="completed",
            payload={
                "state_patch": {
                    "draft_ready": True,
                    "status_log": ["coder completed"],
                }
            },
        ),
    ]


def reviewer_script(_request):
    return [
        ControllerEvent(
            kind="message_completed",
            payload={
                "kind": "review",
                "text": "Draft looks consistent with the plan.",
            },
        ),
        ControllerEvent(
            kind="completed",
            payload={
                "state_patch": {
                    "review_passed": True,
                    "status_log": ["reviewer completed"],
                }
            },
        ),
    ]


def main() -> None:
    graph = AgentGraph(reducers={"steps": append_list, "status_log": append_list}, name="planner-coder-reviewer")
    graph.add_operator("planner_op", DefaultOperator("planner_op", StaticController(planner_script), role="planner"))
    graph.add_operator("coder_op", DefaultOperator("coder_op", StaticController(coder_script), role="coder"))
    graph.add_operator("reviewer_op", DefaultOperator("reviewer_op", StaticController(reviewer_script), role="reviewer"))

    graph.add_node(
        "plan",
        operator="planner_op",
        objective="Plan the work",
        skills=["research-paper-search", "literature-synthesis"],
    )
    graph.add_node(
        "implement",
        operator="coder_op",
        objective="Implement the work",
        skills=["experiment-planning", "dataset-triage"],
    )
    graph.add_node(
        "review",
        operator="reviewer_op",
        objective="Review the work",
        skills=["citation-audit", "result-audit"],
    )
    graph.add_edge("plan", "implement")
    graph.add_edge("implement", "review")

    result = graph.compile().invoke({"task": "Start the MAS implementation"})
    print(json.dumps(result.state, indent=2, ensure_ascii=True))
    print(f"messages={len(result.messages)} artifacts={len(result.artifacts)} status={result.status}")


if __name__ == "__main__":
    main()