Skip to main content

work_buddy/tools/
worklog.rs

1//! Tools to manage the worklog
2
3use anyhow::{Result, anyhow};
4use chrono::format::SecondsFormat;
5use chrono::{DateTime, Datelike, Days, FixedOffset, Local, NaiveDate};
6use helpers::TruncatedDisplay;
7use llimorse::{Agent, CallableTool};
8use schemars::JsonSchema;
9use serde::{Deserialize, Serialize};
10use std::path::PathBuf;
11use std::sync::Arc;
12use std::{fmt, fs, io};
13use tokio::sync::Mutex;
14
15/// Directory in which the worklogs are stored
16#[derive(Debug)]
17pub struct WorklogDirectory {
18    /// Base path of the directory containing the worklogs
19    path: PathBuf,
20}
21
22/// A worklog entry
23#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema)]
24struct WorklogEntry {
25    /// LLM-settable part of the entry
26    #[serde(flatten)]
27    settable: WorklogSettableEntry,
28
29    /// When this entry was created
30    timestamp: String,
31}
32
33/// The LLM-settable part of a worklog entry
34#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema)]
35struct WorklogSettableEntry {
36    /// Summary of the work that was actually done in this chunk
37    summary: String,
38
39    /// How many minutes were spent, roughly, on this work item
40    effort_minutes: usize,
41
42    /// What components were touched, i.e. which projects were affected
43    #[serde(default, skip_serializing_if = "Vec::is_empty")]
44    components: Vec<String>,
45
46    /// What is the underlying reason for working on this, what has caused the need to work on it;
47    /// in case of bugs, what has the investigation unveiled to be the cause of the bug?
48    #[serde(default, skip_serializing_if = "Option::is_none")]
49    root_cause: Option<String>,
50
51    /// What is the overarching approach being taken here, regarding the whole underlying problem?
52    #[serde(default, skip_serializing_if = "Option::is_none")]
53    fix_approach: Option<String>,
54
55    /// Persistent and globally reachable references of the work: Issue tickets, merge requests,
56    /// commits.
57    #[serde(default, skip_serializing_if = "Vec::is_empty")]
58    references: Vec<WorklogReference>,
59
60    /// Arbitrary tags you would like to give to this work item to better find it later
61    #[serde(default, skip_serializing_if = "Vec::is_empty")]
62    tags: Vec<WorklogTag>,
63
64    /// A detailed explanation of what was done for this work item, capturing everything you have
65    /// to offer
66    narrative: String,
67}
68
69/// Persistent and globally reachable references of work
70#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema)]
71#[serde(rename_all = "snake_case", tag = "type")]
72#[schemars(inline)]
73enum WorklogReference {
74    /// Gitlab work items, github issues, Jira tickets, ...
75    Issue {
76        /// URL linking to the issue
77        url: String,
78    },
79
80    /// Github pull requests, gitlab merge requests, ...
81    MergeRequest {
82        /// URL linking to the MR/PR
83        url: String,
84    },
85
86    /// Commit in a repository noted in `components`
87    Commit {
88        /// The commit hash
89        hash: String,
90    },
91}
92
93/// Tags for worklog items
94#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize, JsonSchema)]
95#[serde(rename_all = "snake_case")]
96#[schemars(inline)]
97enum WorklogTag {
98    /// A bug to be fixed
99    Bug,
100
101    /// Code review
102    Review,
103
104    /// A new feature to be implemented
105    Feature,
106
107    /// Investigation for either the root cause of a bug or how to implement a new feature
108    Investigation,
109
110    /// Some general chore like refactoring
111    Chore,
112
113    /// A meeting with one or more other people
114    Meeting,
115
116    /// Testing work
117    Tests,
118}
119
120impl WorklogDirectory {
121    /// Create an instance storing logs in the given directory.
122    pub fn new(path: PathBuf) -> Self {
123        WorklogDirectory { path }
124    }
125
126    /// Add relevant tools for the worklog to `agent`.
127    pub fn add_tools(self, agent: &mut Agent) {
128        let this = Arc::new(Mutex::new(self));
129
130        agent.add_tool(WorklogAdd::new(Arc::clone(&this)));
131        agent.add_tool(WorklogQuery::new(this));
132    }
133
134    /// Append `entry` to the worklog.
135    fn push(&mut self, entry: WorklogSettableEntry) -> Result<()> {
136        let now = Local::now();
137        let entry = WorklogEntry {
138            settable: entry,
139            timestamp: now.to_rfc3339_opts(SecondsFormat::Secs, false),
140        };
141
142        let mut week = self.load_week(now.date_naive())?;
143        week.push(entry);
144        self.store_week(now.date_naive(), week)?;
145
146        Ok(())
147    }
148
149    /// Find the worklog filename corresponding to `date`.
150    fn date_to_fname<D: Datelike>(date: D) -> String {
151        let week = date.iso_week();
152        format!("{}-w{}.json", week.year(), week.week())
153    }
154
155    /// Load the worklog entries for the week in which `date` is.
156    fn load_week(&self, date: NaiveDate) -> Result<Vec<WorklogEntry>> {
157        let path = self.path.join(Self::date_to_fname(date));
158        let content = match fs::read_to_string(&path) {
159            Ok(content) => content,
160            Err(err) if err.kind() == io::ErrorKind::NotFound => String::from("[]"),
161            Err(err) => return Err(err.into()),
162        };
163
164        serde_json::from_str(&content).map_err(|err| {
165            anyhow!(
166                "{}: Failed to load week log from JSON: {err}",
167                path.display()
168            )
169        })
170    }
171
172    /// Store the worklog entries for the week in which `date` is.
173    fn store_week(&self, date: NaiveDate, week: Vec<WorklogEntry>) -> Result<()> {
174        let path = self.path.join(Self::date_to_fname(date));
175
176        let json = serde_json::to_string(&week)
177            .map_err(|err| anyhow!("Failed to convert week log to JSON: {err}"))?;
178
179        fs::write(&path, json)
180            .map_err(|err| anyhow!("{}: Failed to write week log file: {err}", path.display()))?;
181
182        Ok(())
183    }
184}
185
186impl fmt::Display for WorklogEntry {
187    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
188        write!(f, "{} timestamp={}", self.settable, self.timestamp)
189    }
190}
191
192impl fmt::Display for WorklogSettableEntry {
193    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
194        write!(
195            f,
196            "summary={:?} effort_mins={} ",
197            self.summary, self.effort_minutes
198        )?;
199        if !self.components.is_empty() {
200            write!(f, "components={:?} ", self.components)?;
201        }
202        if let Some(root_cause) = &self.root_cause {
203            write!(f, "root_cause={root_cause:?} ")?;
204        }
205        if let Some(fix_approach) = &self.fix_approach {
206            write!(f, "fix_approach={fix_approach:?} ")?;
207        }
208        if !self.references.is_empty() {
209            write!(f, "refs={:?} ", self.references)?;
210        }
211        if !self.tags.is_empty() {
212            write!(f, "tags={:?} ", self.tags)?;
213        }
214        write!(f, "narrative={:?}", self.narrative.truncated_display(100))?;
215
216        Ok(())
217    }
218}
219
220llimorse::tool! {
221    'name: "worklog_add";
222
223    /// Log work the user has done in the worklog.
224    #[derive(Debug)]
225    'params: pub struct WorklogAddParams {
226        /// Worklog entry to add
227        #[serde(flatten)]
228        entry: WorklogSettableEntry,
229    }
230
231    /// Result of adding a worklog entry.
232    #[derive(Debug)]
233    'result: pub struct WorklogAddResult {
234        /// Summary of the entry just added
235        summary: String,
236    }
237
238    /// Log work the user has done in the worklog.
239    #[derive(Debug)]
240    'state: pub struct WorklogAdd {
241        /// Underlying worklog “database”
242        storage: Arc<Mutex<WorklogDirectory>>,
243    }
244}
245
246impl std::fmt::Display for WorklogAddParams {
247    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
248        write!(f, "{}", self.entry)
249    }
250}
251
252impl std::fmt::Display for WorklogAddResult {
253    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
254        write!(f, "summary={:?}", self.summary)
255    }
256}
257
258impl WorklogAdd {
259    /// Create a worklog_add tool for the given worklog directory.
260    pub fn new(dir: Arc<Mutex<WorklogDirectory>>) -> Self {
261        WorklogAdd { storage: dir }
262    }
263}
264
265impl CallableTool for WorklogAdd {
266    async fn execute(&self, _agent: &Agent, params: WorklogAddParams) -> Result<WorklogAddResult> {
267        let summary = params.entry.summary.clone();
268        let mut storage = self.storage.lock().await;
269        storage.push(params.entry)?;
270        Ok(WorklogAddResult { summary })
271    }
272}
273
274llimorse::tool! {
275    'name: "worklog_query";
276
277    /// Query the worklog for entries.
278    #[derive(Debug)]
279    'params: pub struct WorklogQueryParams {
280        /// Date to start the search from (inclusive, YYYY-mm-dd)
281        start_date: String,
282
283        /// Date to end the search at (inclusive, YYYY-mm-dd)
284        end_date: String,
285
286        /// Return only entries that include any of these tags
287        #[serde(default, skip_serializing_if = "Option::is_none")]
288        tags: Option<Vec<WorklogTag>>,
289
290        /// Return only entries that touched any of these components
291        #[serde(default, skip_serializing_if = "Option::is_none")]
292        components: Option<Vec<String>>,
293
294        /// Return only entries whose `summary` or `narrative` fields contain any of these strings
295        #[serde(default, skip_serializing_if = "Option::is_none")]
296        match_strings: Option<Vec<String>>,
297    }
298
299    /// Result of a worklog query.
300    #[derive(Debug)]
301    'result: pub struct WorklogQueryResult {
302        /// Matching worklog entries
303        matching: Vec<WorklogEntry>,
304    }
305
306    /// Query the worklog for entries.
307    #[derive(Debug)]
308    'state: pub struct WorklogQuery {
309        /// Underlying worklog “database”
310        storage: Arc<Mutex<WorklogDirectory>>,
311    }
312}
313
314impl std::fmt::Display for WorklogQueryParams {
315    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
316        write!(f, "dates=[{}..{}]", self.start_date, self.end_date)?;
317
318        if let Some(tags) = &self.tags {
319            write!(f, " tags={tags:?}")?;
320        }
321        if let Some(components) = &self.components {
322            write!(f, " tags={components:?}")?;
323        }
324        if let Some(match_strings) = &self.match_strings {
325            write!(f, " tags={match_strings:?}")?;
326        }
327
328        Ok(())
329    }
330}
331
332impl std::fmt::Display for WorklogQueryResult {
333    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
334        write!(f, "[")?;
335        let result_count = self.matching.len();
336        for (i, entry) in self.matching.iter().enumerate() {
337            if i == result_count - 1 {
338                write!(f, "{entry}")?;
339            } else {
340                write!(f, "{entry}; ")?;
341            }
342        }
343        write!(f, "]")
344    }
345}
346
347impl WorklogQuery {
348    /// Create a worklog_query tool for the given worklog directory.
349    pub fn new(dir: Arc<Mutex<WorklogDirectory>>) -> Self {
350        WorklogQuery { storage: dir }
351    }
352}
353
354impl CallableTool for WorklogQuery {
355    async fn execute(
356        &self,
357        _agent: &Agent,
358        params: WorklogQueryParams,
359    ) -> Result<WorklogQueryResult> {
360        let start = NaiveDate::parse_from_str(&params.start_date, "%Y-%m-%d")
361            .map_err(|err| anyhow!("Failed to parse start_date parameter: {err}"))?;
362        let end = NaiveDate::parse_from_str(&params.end_date, "%Y-%m-%d")
363            .map_err(|err| anyhow!("Failed to parse end_date parameter: {err}"))?;
364
365        let storage = self.storage.lock().await;
366        let mut everything = Vec::new();
367        let mut current = start;
368        while current <= end {
369            everything.append(&mut storage.load_week(current)?);
370            current = current + Days::new(7);
371        }
372
373        let mut filtered = everything
374            .into_iter()
375            .filter(|item| {
376                let Ok(timestamp) = DateTime::<FixedOffset>::parse_from_rfc3339(&item.timestamp)
377                else {
378                    // Ignore database corruption...
379                    return false;
380                };
381
382                let date = timestamp.date_naive();
383                if date < start || date > end {
384                    return false;
385                }
386                if let Some(tags) = &params.tags
387                    && !item.settable.tags.iter().any(|tag| tags.contains(tag))
388                {
389                    return false;
390                }
391                if let Some(components) = &params.components
392                    && !item
393                        .settable
394                        .components
395                        .iter()
396                        .any(|component| components.contains(component))
397                {
398                    return false;
399                }
400                if let Some(match_strings) = &params.match_strings
401                    && !match_strings.iter().any(|string| {
402                        item.settable.summary.contains(string)
403                            || item.settable.narrative.contains(string)
404                    })
405                {
406                    return false;
407                }
408
409                true
410            })
411            .collect::<Vec<_>>();
412
413        filtered.sort_by(|item1, item2| {
414            // Yes, would be nice if we could bring the timestamp from above here, but, well,
415            // whatever.  Unwrap is safe, we only have items with valid timestamps.
416            let ts1 = DateTime::<FixedOffset>::parse_from_rfc3339(&item1.timestamp).unwrap();
417            let ts2 = DateTime::<FixedOffset>::parse_from_rfc3339(&item2.timestamp).unwrap();
418
419            ts1.cmp(&ts2)
420        });
421
422        Ok(WorklogQueryResult { matching: filtered })
423    }
424}