bevy_rig 0.1.0

Bevy ECS primitives and systems for modeling providers, agents, tools, sessions, runs, and workflows on top of Rig.
Documentation 
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

use bevy_app::App;
use bevy_rig::prelude::*;
use serde_json::json;

fn main() {
    let mut app = App::new();
    app.add_plugins(BevyRigPlugin);

    let workflow = {
        let world = app.world_mut();

        let provider = spawn_provider(
            world,
            ProviderSpec::new(ProviderKind::OpenAi, "openai"),
            ProviderCapabilities::text_tooling(),
        );
        let model = spawn_model(
            world,
            provider,
            ModelSpec::new("gpt-4o-mini"),
            ModelCapabilities::chat_with_tools(),
            128_000,
        )
        .expect("model should register");

        let planner = spawn_agent_from_model(world, "planner", model)
            .expect("planner agent should spawn")
            .agent;
        let summarizer = spawn_agent_from_model(world, "summarizer", model)
            .expect("summarizer agent should spawn")
            .agent;
        let lookup_tool = world
            .spawn(ToolBundle::new(ToolSpec::new(
                "lookup",
                "Dummy lookup tool",
                json!({"type":"object","properties":{}}),
            )))
            .id();

        let workflow = spawn_workflow(
            world,
            WorkflowSpec::new("research", "Plan, call a tool, then summarize"),
        );
        let plan_node = spawn_workflow_node(world, workflow, WorkflowNodeKind::Agent, "plan")
            .expect("plan node");
        let tool_node = spawn_workflow_node(world, workflow, WorkflowNodeKind::Tool, "lookup")
            .expect("tool node");
        let summarize_node =
            spawn_workflow_node(world, workflow, WorkflowNodeKind::Agent, "summarize")
                .expect("summarize node");

        bind_workflow_node(world, plan_node, planner).expect("planner binding");
        bind_workflow_node(world, tool_node, lookup_tool).expect("tool binding");
        bind_workflow_node(world, summarize_node, summarizer).expect("summarizer binding");
        set_workflow_entry(world, workflow, plan_node).expect("entry node");
        connect_workflow_nodes(world, plan_node, tool_node, Some("needs_lookup"))
            .expect("plan -> tool");
        connect_workflow_nodes(world, tool_node, summarize_node, None::<String>)
            .expect("tool -> summarize");

        workflow
    };

    for node in reachable_workflow_nodes(app.world(), workflow).expect("workflow traversal") {
        let name = &app
            .world()
            .get::<WorkflowNodeName>(node)
            .expect("workflow node name")
            .0;
        let kind = app
            .world()
            .get::<WorkflowNodeKind>(node)
            .expect("workflow node kind");
        let binding = app
            .world()
            .get::<WorkflowBinding>(node)
            .expect("workflow binding")
            .0;

        println!("{name}: {kind:?} -> {binding:?}");
    }
}