work_buddy/tools/
knowledge.rs

1//! Tools to manage contextual knowledge
2
3use anyhow::{Result, anyhow};
4use helpers::TruncatedDisplay;
5use llimorse::{Agent, CallableTool};
6use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8use std::io::{self, Write};
9use std::path::PathBuf;
10use std::sync::Arc;
11use std::{fmt, fs};
12use tokio::sync::Mutex;
13
14/// Knowledge file
15#[derive(Debug)]
16pub struct KnowledgeFile {
17    /// Path where to load/store the content
18    path: PathBuf,
19
20    /// Knowledge content
21    content: HashMap<String, KnowledgeOrAlias>,
22}
23
24/// Knowledge entry, or an alias
25#[derive(Debug, Deserialize, Serialize)]
26#[serde(rename_all = "snake_case")]
27pub enum KnowledgeOrAlias {
28    /// Knowledge content, describing the keyword
29    Content(String),
30
31    /// Link to a different keyword, both of which are described by the same content
32    Alias(String),
33}
34
35impl KnowledgeFile {
36    /// Open the knowledge file under the given `path`.
37    pub fn open(path: PathBuf) -> Result<Self> {
38        let content = match fs::read_to_string(&path) {
39            Ok(content) => content,
40            Err(err) if err.kind() == io::ErrorKind::NotFound => {
41                let mut file = fs::File::create_new(&path).map_err(|c_err| {
42                    anyhow!("Opening file failed: {err}; and creating failed, too: {c_err}")
43                })?;
44
45                let content = String::from("{}");
46                file.write_all(content.as_bytes())?;
47                content
48            }
49            Err(err) => return Err(err.into()),
50        };
51
52        let content = serde_json::from_str(&content)?;
53
54        Ok(KnowledgeFile { path, content })
55    }
56
57    /// Add relevant tools for this file to `agent`.
58    pub fn add_tools(self, agent: &mut Agent) {
59        let this = Arc::new(Mutex::new(self));
60
61        agent.add_tool(KnowledgeUpsert::new(Arc::clone(&this)));
62        agent.add_tool(KnowledgeQuery::new(this));
63    }
64
65    /// Write the contents into the file.
66    fn write(&self) -> Result<()> {
67        let json = serde_json::to_string(&self.content)
68            .map_err(|err| anyhow!("Failed to convert knowledge database to JSON: {err}"))?;
69
70        fs::write(&self.path, json)
71            .map_err(|err| anyhow!("Failed to write knowledge database file: {err}"))?;
72
73        Ok(())
74    }
75}
76
77llimorse::tool! {
78    'name: "knowledge_upsert";
79
80    /// Add new content to the knowledge database (to explain keywords, e.g. projects, components,
81    /// etc.), or modify an existing entry.
82    #[derive(Debug)]
83    'params: pub struct KnowledgeUpsertParams {
84        /// The keyword which this information describes
85        keyword: String,
86
87        /// The content to store under this keyword; by default, leave the existing content
88        /// unchanged.
89        #[serde(default, skip_serializing_if = "Option::is_none")]
90        content: Option<String>,
91
92        /// Alias keywords to add to this one, i.e., keywords that should resolve to the exact same
93        /// content when queried.
94        #[serde(default, skip_serializing_if = "Option::is_none")]
95        add_aliases: Option<Vec<String>>,
96    }
97
98    /// Result of modifying the knowledge database
99    #[derive(Debug)]
100    'result: pub struct KnowledgeUpsertResult {
101        /// The keyword whose content was updated
102        keyword: String,
103
104        /// List of aliases that resolve to the same content
105        aliases: Vec<String>,
106    }
107
108    /// Add new content to the knowledge database (to explain keywords, e.g. projects, components,
109    /// etc.), or modify an existing entry.
110    #[derive(Debug)]
111    'state: pub struct KnowledgeUpsert {
112        /// Knowledge database
113        db: Arc<Mutex<KnowledgeFile>>,
114    }
115}
116
117impl fmt::Display for KnowledgeUpsertParams {
118    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
119        write!(f, "key={}", self.keyword)?;
120        if let Some(content) = &self.content {
121            write!(f, " content={:?}", content.truncated_display(100))?;
122        }
123        if let Some(aliases) = &self.add_aliases {
124            write!(f, " add_aliases={aliases:?}")?;
125        }
126        Ok(())
127    }
128}
129
130impl fmt::Display for KnowledgeUpsertResult {
131    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
132        write!(f, "key={} aliases={:?}", self.keyword, self.aliases)
133    }
134}
135
136impl KnowledgeUpsert {
137    /// Create a knowledge_upsert tool for the given knowledge storage.
138    pub fn new(storage: Arc<Mutex<KnowledgeFile>>) -> Self {
139        KnowledgeUpsert { db: storage }
140    }
141}
142
143impl CallableTool for KnowledgeUpsert {
144    async fn execute(
145        &self,
146        _agent: &Agent,
147        params: KnowledgeUpsertParams,
148    ) -> Result<KnowledgeUpsertResult> {
149        if params.content.is_none() && params.add_aliases.is_none() {
150            return Err(anyhow!("Must update content or add aliases"));
151        }
152
153        let mut db = self.db.lock().await;
154
155        // Check first, when we can still generate errors without having modified anything
156        if let Some(aliases) = &params.add_aliases {
157            for alias in aliases {
158                if let Some(entry) = db.content.get(alias) {
159                    let err = match entry {
160                        KnowledgeOrAlias::Content(_) => {
161                            anyhow!("Cannot make '{alias}' an alias: Already contains content")
162                        }
163                        KnowledgeOrAlias::Alias(link) => {
164                            anyhow!("Cannot make '{alias}' an alias: Already an alias of '{link}'")
165                        }
166                    };
167                    return Err(err);
168                }
169            }
170        }
171
172        let resolved_keyword = if let Some(mut entry) = db.content.get_mut(&params.keyword) {
173            let mut resolved = params.keyword.clone();
174            while let KnowledgeOrAlias::Alias(link) = entry {
175                let link = link.clone(); // Cannot `get_mut()` while having this borrowed
176                entry = db.content.get_mut(&link).ok_or_else(|| {
177                    anyhow!("Dead link: '{resolved}' points to '{link}', which does not exist")
178                })?;
179                resolved = link.clone();
180            }
181            if let Some(content) = params.content {
182                *entry = KnowledgeOrAlias::Content(content);
183            }
184            resolved
185        } else if let Some(content) = params.content {
186            db.content
187                .insert(params.keyword.clone(), KnowledgeOrAlias::Content(content));
188            params.keyword.clone()
189        } else {
190            return Err(anyhow!(
191                "{} does not have an entry yet, and you cannot add aliases to an empty entry",
192                params.keyword
193            ));
194        };
195
196        if let Some(aliases) = params.add_aliases {
197            for alias in aliases {
198                let old = db
199                    .content
200                    .insert(alias, KnowledgeOrAlias::Alias(resolved_keyword.clone()));
201                // Loop above must have verified this
202                assert!(old.is_none());
203            }
204        }
205
206        db.write()?;
207
208        let all_aliases = db
209            .content
210            .iter()
211            .filter_map(|(kw, entry)| match entry {
212                KnowledgeOrAlias::Content(_) => None,
213                KnowledgeOrAlias::Alias(link) if link == &resolved_keyword => Some(kw.clone()),
214                KnowledgeOrAlias::Alias(_) => None,
215            })
216            .collect();
217
218        Ok(KnowledgeUpsertResult {
219            keyword: params.keyword,
220            aliases: all_aliases,
221        })
222    }
223}
224
225llimorse::tool! {
226    'name: "knowledge_query";
227
228    /// Query content from the knowledge database, by keyword.
229    #[derive(Debug)]
230    'params: pub struct KnowledgeQueryParams {
231        /// The keyword whose content to look up. If not specified, return all keywords in the DB.
232        #[serde(default, skip_serializing_if = "Option::is_none")]
233        keyword: Option<String>,
234    }
235
236    /// Result of querying the knowledge database
237    #[derive(Debug)]
238    'result: pub struct KnowledgeQueryResult {
239        /// The keyword that was looked up
240        #[serde(default, skip_serializing_if = "Option::is_none")]
241        keyword: Option<String>,
242
243        /// Content that describes the keyword
244        #[serde(default, skip_serializing_if = "Option::is_none")]
245        content: Option<String>,
246
247        /// All keywords in the database
248        #[serde(default, skip_serializing_if = "Vec::is_empty")]
249        keywords: Vec<String>,
250    }
251
252    /// Query content from the knowledge database, by keyword.
253    #[derive(Debug)]
254    'state: pub struct KnowledgeQuery {
255        /// Knowledge database
256        db: Arc<Mutex<KnowledgeFile>>,
257    }
258}
259
260impl fmt::Display for KnowledgeQueryParams {
261    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
262        if let Some(key) = &self.keyword {
263            write!(f, "key={key}")?;
264        }
265        Ok(())
266    }
267}
268
269impl fmt::Display for KnowledgeQueryResult {
270    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
271        if !self.keywords.is_empty() {
272            write!(f, "keywords={:?}", self.keywords)?;
273            assert!(self.keyword.is_none());
274            assert!(self.content.is_none());
275        }
276        if let Some(key) = &self.keyword {
277            write!(f, "key={key}")?;
278        }
279        if let Some(content) = &self.content {
280            write!(f, " content={:?}", content.truncated_display(100))?;
281        }
282        Ok(())
283    }
284}
285
286impl KnowledgeQuery {
287    /// Create a knowledge_query tool for the given knowledge storage.
288    pub fn new(storage: Arc<Mutex<KnowledgeFile>>) -> Self {
289        KnowledgeQuery { db: storage }
290    }
291}
292
293impl CallableTool for KnowledgeQuery {
294    async fn execute(
295        &self,
296        _agent: &Agent,
297        params: KnowledgeQueryParams,
298    ) -> Result<KnowledgeQueryResult> {
299        let db = self.db.lock().await;
300
301        let Some(key) = params.keyword else {
302            return Ok(KnowledgeQueryResult {
303                keyword: None,
304                content: None,
305                keywords: db.content.keys().cloned().collect(),
306            });
307        };
308
309        let mut entry = db
310            .content
311            .get(&key)
312            .ok_or_else(|| anyhow!("Keyword {key} not present in the database"))?;
313        let mut looked_up = &key;
314        while let KnowledgeOrAlias::Alias(link) = entry {
315            entry = db.content.get(link).ok_or_else(|| {
316                anyhow!("Dead link: '{looked_up}' points to '{link}', which does not exist")
317            })?;
318            looked_up = link;
319        }
320
321        let KnowledgeOrAlias::Content(content) = entry else {
322            panic!(
323                "Exhaustively went through all alias links, for some reason this is still no content entry: {entry:?}"
324            );
325        };
326
327        Ok(KnowledgeQueryResult {
328            keyword: Some(key),
329            content: Some(content.clone()),
330            keywords: Vec::new(),
331        })
332    }
333}
work_buddy/tools/knowledge.rs

work_buddy/tools/
knowledge.rs
