rig_memory/
saved.rs

1//! A saved agent's values read without the host that saved them: the keys
2//! rig saves an agent's model and title under, typed reads of any saved
3//! value, and which agent of a session is its main one.
4
5use std::collections::BTreeMap;
6
7use serde::de::DeserializeOwned;
8use serde_json::Value;
9
10use crate::session::skipped;
11use crate::{SavedAgent, Session};
12
13/// The key an agent's model is saved under: a string, `vendor/model`.
14pub const MODEL: &str = "rig::model";
15
16/// The key an agent's title is saved under, such as a subagent's task: a
17/// string.
18pub const TITLE: &str = "rig::title";
19
20impl SavedAgent {
21    /// Its saved value under `key`, such as a plugin's `goals.goal`, read
22    /// as a `T`: `None` when it has none, an error when the value is not a
23    /// `T`.
24    pub fn value<T: DeserializeOwned>(&self, key: &str) -> Option<Result<T, serde_json::Error>> {
25        value(&self.saved, key)
26    }
27
28    /// Its model, `vendor/model`, when one was saved ([`MODEL`]).
29    pub fn model(&self) -> Option<String> {
30        self.value(MODEL)?.ok()
31    }
32
33    /// Its title, when one was saved ([`TITLE`]).
34    pub fn title(&self) -> Option<String> {
35        self.value(TITLE)?.ok()
36    }
37}
38
39/// The value of `key` in `state`, parsed.
40fn value<T: DeserializeOwned>(
41    state: &BTreeMap<String, Value>,
42    key: &str,
43) -> Option<Result<T, serde_json::Error>> {
44    state.get(key).map(T::deserialize)
45}
46
47/// Where an agent sorts among the agents of a session: those nothing
48/// spawned first, each group by id. The first is the main agent, the one
49/// the user talks to when none is named.
50pub fn main_order(spawned: bool, id: &str) -> (bool, &str) {
51    (spawned, id)
52}
53
54impl Session {
55    /// The session's own saved value under `key` (the [`SESSION`](crate::SESSION)
56    /// log's), read as a `T`: `None` when it has none, an error when the
57    /// value is not a `T`.
58    pub fn value<T: DeserializeOwned>(&self, key: &str) -> Option<Result<T, serde_json::Error>> {
59        value(&self.saved, key)
60    }
61
62    /// The main agent: the first by [`main_order`] of the agents a restore
63    /// brings back (not retired, nor spawned by one retired).
64    pub fn main_agent(&self) -> Option<&SavedAgent> {
65        let logs = self.agents.iter().map(|agent| {
66            let parent = agent.parent.as_deref();
67            (agent.id.as_str(), parent, agent.retired)
68        });
69        let skipped = skipped(logs);
70        self.agents
71            .iter()
72            .filter(|agent| !skipped.contains(agent.id.as_str()))
73            .min_by_key(|agent| main_order(agent.parent.is_some(), &agent.id))
74    }
75}
76
77#[cfg(test)]
78mod tests;
rig_memory/saved.rs

rig_memory/
saved.rs
