PlotWorks/agent.py at main · RootedInData/PlotWorks · 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
from __future__ import annotations

from google.adk import Agent
from google.adk.tools.agent_tool import AgentTool
from google.adk.tools.function_tool import FunctionTool

from .agent_cards import AGENT_CARDS
from .agents import (
    build_animation_developer_agent,
    build_code_planning_agent,
    build_column_decoder_agent,
    build_data_intake_agent,
    build_data_transformation_agent,
    build_eda_agent,
    build_method_research_agent,
    build_plot_review_agent,
    build_publication_plot_agent,
    build_r_plot_developer_agent,
    build_report_agent,
    build_visualization_agent,
    build_visualization_planner_agent,
)
from .llm_factory import build_model
from .prompts import SUPERVISOR_PROMPT
from .tools import (
    execute_generated_python_transform,
    execute_generated_r_animation,
    execute_generated_r_plot,
    save_data_transformations,
    set_ggplot2_case_palette_default,
)

# Specialist agents exposed to the supervisor as callable tools.
data_intake_agent = build_data_intake_agent()
eda_agent = build_eda_agent()
data_transformation_agent = build_data_transformation_agent()
visualization_planner_agent = build_visualization_planner_agent()
visualization_agent = build_visualization_agent()
column_decoder_agent = build_column_decoder_agent()
publication_plot_agent = build_publication_plot_agent()
r_plot_developer_agent = build_r_plot_developer_agent()
animation_developer_agent = build_animation_developer_agent()
plot_review_agent = build_plot_review_agent()
code_planning_agent = build_code_planning_agent()
report_agent = build_report_agent()
method_research_agent = build_method_research_agent()

root_agent = Agent(
    name="PlotWorksSupervisor",
    model=build_model(),
    description=(
        "Supervisor for PlotWorks with deterministic EDA, approved data transformation, "
        "polished Python plotting, approved R recipes, guarded custom R plotting, "
        "and animated visualization."
    ),
    instruction=SUPERVISOR_PROMPT + "\n\n" + AGENT_CARDS,
    tools=[
        AgentTool(agent=data_intake_agent),
        AgentTool(agent=eda_agent),
        AgentTool(agent=data_transformation_agent),
        # Confirmation-required actions live on the user-facing root agent.
        # This lets ADK Web surface and resume the structured confirmation event
        # instead of burying it inside a nested AgentTool invocation.
        FunctionTool(save_data_transformations, require_confirmation=True),
        FunctionTool(execute_generated_python_transform, require_confirmation=True),
        FunctionTool(set_ggplot2_case_palette_default, require_confirmation=True),
        FunctionTool(execute_generated_r_plot, require_confirmation=True),
        FunctionTool(execute_generated_r_animation, require_confirmation=True),
        AgentTool(agent=visualization_planner_agent),
        AgentTool(agent=visualization_agent),
        AgentTool(agent=column_decoder_agent),
        AgentTool(agent=publication_plot_agent),
        AgentTool(agent=r_plot_developer_agent),
        AgentTool(agent=animation_developer_agent),
        AgentTool(agent=plot_review_agent),
        AgentTool(agent=code_planning_agent),
        AgentTool(agent=report_agent),
        AgentTool(agent=method_research_agent),
    ],
)