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
83
84
85
86
87
88
89
90
91
92

use bevy_ecs::{
    hierarchy::{ChildOf, Children},
    prelude::*,
};
use serde::{Deserialize, Serialize};

#[derive(Component, Clone, Debug, PartialEq, Eq)]
pub struct Session {
    pub title: String,
}

#[derive(Component, Clone, Debug, Default, PartialEq, Eq)]
pub struct ChatMessage;

#[derive(Component, Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum ChatMessageRole {
    User,
    Assistant,
    System,
}

#[derive(Component, Clone, Debug, PartialEq, Eq)]
pub struct ChatMessageText(pub String);

#[derive(Bundle)]
pub struct SessionBundle {
    pub session: Session,
}

impl SessionBundle {
    pub fn new(title: impl Into<String>) -> Self {
        Self {
            session: Session {
                title: title.into(),
            },
        }
    }
}

#[derive(Bundle)]
pub struct ChatMessageBundle {
    pub message: ChatMessage,
    pub role: ChatMessageRole,
    pub text: ChatMessageText,
    pub child_of: ChildOf,
}

impl ChatMessageBundle {
    pub fn new(session: Entity, role: ChatMessageRole, text: impl Into<String>) -> Self {
        Self {
            message: ChatMessage,
            role,
            text: ChatMessageText(text.into()),
            child_of: ChildOf(session),
        }
    }
}

pub fn spawn_session(world: &mut World, title: impl Into<String>) -> Entity {
    world.spawn(SessionBundle::new(title)).id()
}

pub fn spawn_chat_message(
    world: &mut World,
    session: Entity,
    role: ChatMessageRole,
    text: impl Into<String>,
) -> Entity {
    world
        .spawn(ChatMessageBundle::new(session, role, text))
        .id()
}

pub fn collect_transcript(world: &World, session: Entity) -> Vec<(ChatMessageRole, String)> {
    let mut transcript = Vec::new();

    let Some(children) = world.get::<Children>(session) else {
        return transcript;
    };

    for child in children.iter() {
        let Some(role) = world.get::<ChatMessageRole>(child) else {
            continue;
        };
        let Some(text) = world.get::<ChatMessageText>(child) else {
            continue;
        };
        transcript.push((role.clone(), text.0.clone()));
    }

    transcript
}