rig-ecs 0.44.0

rig inside a Bevy World: effects as entities, handlers as entities, the driver as a system.
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
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

//! Hello, model: the bus in a `World` in under thirty lines of user code.
//! Add the plugin, register a handler, spawn an effect, observe the answer.
//! The model is a scripted mock, so no provider feature and no key: the
//! shape is the same with a real `ModelAdapter` registered under the
//! key.

#![allow(
    clippy::expect_used,
    clippy::unwrap_used,
    clippy::indexing_slicing,
    clippy::panic,
    reason = "an example: user code, thirty lines, a mock behind it"
)]

use bevy_app::{App, AppExit, Startup};
use bevy_ecs::prelude::*;
use rig_core::serve::Dispatch;
use rig_core::{
    completion::{
        CompletionRequest, CompletionResponse, Message, ModelRef, ProviderCapabilities, Usage,
    },
    effect::{EffectKind, FamilyDescriptor, HandlerDescriptor, HandlerKey, Outcome},
    message::AssistantContent,
    serve::{Serve, ServingPolicy},
};
use rig_ecs::bus::{EffectOutcome, Handlers, PendingEffect};

// ---- the user's program: the next thirty lines ----

fn main() {
    let mut app = App::new();
    app.add_plugins(rig_ecs::bus::BusPlugin::with_policy(
        ServingPolicy::default(),
    ))
    .add_systems(Startup, (register_the_model, ask).chain())
    .add_observer(print_the_answer)
    .run();
}

fn register_the_model(mut handlers: Handlers) {
    if let Err(report) = handlers.register("model", Mock) {
        eprintln!("could not register the model: {report}");
    }
}

fn ask(mut commands: Commands) {
    let request = CompletionRequest::from(vec![Message::user("hello?")]);
    commands.spawn(PendingEffect::new(
        "model",
        EffectKind::Completion {
            request,
            stream: false,
        },
    ));
}

fn print_the_answer(
    answered: On<Add, EffectOutcome>,
    outcomes: Query<&EffectOutcome>,
    mut exit: MessageWriter<AppExit>,
) {
    if let Ok(outcome) = outcomes.get(answered.event().entity) {
        println!("the model said: {}", text(&outcome.0));
        exit.write(AppExit::Success);
    }
}

// ---- the mock, in place of a provider ----

struct Mock;

impl Serve for Mock {
    type Family = rig_core::effect::family::Completion;

    fn descriptor(&self) -> HandlerDescriptor {
        HandlerDescriptor {
            key: HandlerKey::from("model"),
            family: FamilyDescriptor::Completion {
                model: ModelRef::new("mock"),
                capabilities: ProviderCapabilities::default(),
            },
            layers: Vec::new(),
        }
    }

    async fn serve(&self, _kind: EffectKind, _dispatch: Dispatch) -> rig_core::serve::Reply {
        let response = CompletionResponse::new(
            vec![AssistantContent::text("hello from the world")],
            Usage::default(),
            rig_core::message::Origin::new("example.api", "mock", ""),
            serde_json::json!({}),
        );
        rig_core::serve::Reply::Outcome(Ok(Outcome::Completion(response)))
    }
}

fn text(outcome: &Result<Outcome, rig_core::error::ErrorReport>) -> String {
    match outcome {
        Ok(Outcome::Completion(response)) => response
            .choice
            .iter()
            .filter_map(|content| match content {
                AssistantContent::Text(text) => Some(text.text.clone()),
                _ => None,
            })
            .collect(),
        Ok(other) => format!("a {} answer", other.family()),
        Err(report) => format!("failed: {report}"),
    }
}