rig_ecs/
fork.rs

1//! Forking an agent: [`ForkAgent`] makes a new agent in the same session
2//! whose log is a copy of another's up to a turn boundary, so it goes on
3//! from there by itself, a compaction made before the cut can be undone in
4//! it, and its requests share the source's prompt cache.
5//!
6//! ```
7//! use rig_ecs::prelude::*;
8//!
9//! // Shows each fork as it is made.
10//! fn show(forked: On<AgentForked>, mut shown: Local<Option<Entity>>) {
11//!     *shown = Some(forked.fork);
12//! }
13//!
14//! App::new().add_observer(show);
15//! ```
16
17use bevy_ecs::prelude::*;
18use bevy_ecs::reflect::{AppTypeRegistry, ReflectComponent};
19use bevy_reflect::PartialReflect;
20use rig_core::conversation::MessageId;
21
22use super::agent::{Agent, AgentId, CacheKey, Conversation, Notice};
23use super::model::Connection;
24use super::restore::SavedKeys;
25use super::session_log::{ReflectSaved, SessionLog, save_conversation};
26
27/// Fork the agent `agent`: a new agent in the same session whose log is a
28/// copy of `agent`'s up to a turn boundary at or before `through` (its
29/// end with `None`; rig-memory's `Writer::fork_log`), its conversation that
30/// copy. A user message that asked for a turn is at no boundary, so a fork
31/// through it ends just before it. The new agent has every saved
32/// component `agent` has now (its model, reasoning setting, tools, system
33/// prompt and working directory among them) and its connection, and shares
34/// its prompt cache ([`CacheKey`]). It has no parent. [`AgentForked`]
35/// names it; a fork that cannot be made, such as without a session log or
36/// through a message the agent never had, is a notice. It is made once the
37/// commands run, after the source's unsaved changes are saved.
38#[derive(EntityEvent, Clone, Debug)]
39pub struct ForkAgent {
40    /// The agent forked.
41    #[event_target]
42    pub agent: Entity,
43    /// The last message the fork may keep; `None` for the end.
44    pub through: Option<MessageId>,
45}
46
47/// An agent made by a [`ForkAgent`].
48#[derive(Event, Clone, Debug)]
49pub struct AgentForked {
50    /// The agent forked.
51    pub from: Entity,
52    /// The new agent.
53    pub fork: Entity,
54    /// The message the fork was asked through, as [`ForkAgent`] said.
55    pub through: Option<MessageId>,
56}
57
58/// Forks an agent, as [`ForkAgent`] says.
59pub(crate) fn fork_agent(fork: On<ForkAgent>, mut commands: Commands) {
60    let (from, through) = (fork.agent, fork.through);
61    commands.queue(
62        move |world: &mut World| match fork_now(world, from, through) {
63            Ok(fork) => world.trigger(AgentForked {
64                from,
65                fork,
66                through,
67            }),
68            Err(why) => {
69                let why = format!("Could not fork the agent: {why}.");
70                world.write_message(Notice::error(from, why));
71            }
72        },
73    );
74}
75
76/// Forks `from` through `through`; the new agent.
77fn fork_now(world: &mut World, from: Entity, through: Option<MessageId>) -> Result<Entity, String> {
78    let id = world
79        .get::<AgentId>(from)
80        .cloned()
81        .ok_or("it is no agent")?;
82    let log = world
83        .get_resource::<SessionLog>()
84        .filter(|log| log.0.is_live())
85        .cloned()
86        .ok_or("the session is not logged")?;
87    let registry = world
88        .get_resource::<AppTypeRegistry>()
89        .cloned()
90        .ok_or("the app has no type registry")?;
91    if let Some(mut conversation) = world.get_mut::<Conversation>(from) {
92        let saved = save_conversation(&log, from, &id, conversation.bypass_change_detection());
93        world.write_message_batch(saved);
94    }
95    let to = AgentId::default();
96    let copy = log
97        .0
98        .fork_log(&id.0, &to.0, through)
99        .map_err(|failure| failure.to_string())?;
100    let registry = registry.read();
101    let saved: Vec<(&ReflectComponent, Box<dyn PartialReflect>)> = registry
102        .iter_with_data::<ReflectSaved>()
103        .filter_map(|(registration, _)| {
104            let component = registration.data::<ReflectComponent>()?;
105            let value = component.reflect(world.get_entity(from).ok()?)?;
106            let value = match value.reflect_clone() {
107                Ok(value) => value.into_partial_reflect(),
108                Err(_) => value.to_dynamic().ok()?,
109            };
110            Some((component, value))
111        })
112        .collect();
113    let key = CacheKey(CacheKey::of(world.get::<CacheKey>(from), &id).to_owned());
114    let connection = world.get::<Connection>(from).cloned();
115    let conversation = Conversation(copy.conversation);
116    let mut spawned = world.spawn((Name::new("agent"), Agent, to.clone(), conversation));
117    for (component, value) in &saved {
118        component.insert(&mut spawned, value.as_ref(), &registry);
119    }
120    spawned.insert(key);
121    if let Some(connection) = connection {
122        spawned.insert(connection);
123    }
124    let fork = spawned.id();
125    // A value the copied log holds that the source has no more is gone in
126    // the fork too.
127    let keys = world.get_resource::<SavedKeys>();
128    for key in copy.saved.keys() {
129        let path = keys.and_then(|keys| keys.path(key)).unwrap_or(key);
130        let registered = registry.get_with_type_path(path);
131        let component = registered.filter(|it| it.data::<ReflectSaved>().is_some());
132        let component = component.and_then(|it| it.data::<ReflectComponent>());
133        let held = |component: &ReflectComponent| {
134            let entity = world.get_entity(fork).ok();
135            entity.is_some_and(|entity| component.contains(entity))
136        };
137        if component.is_some_and(|component| !held(component)) {
138            log.0.save_value(&to.0, key, None, false);
139        }
140    }
141    Ok(fork)
142}
143
144#[cfg(test)]
145mod tests;
rig_ecs/fork.rs

rig_ecs/
fork.rs
