rig_tools/
write.rs

1//! The `write` tool.
2
3use std::path::PathBuf;
4use std::sync::Arc;
5
6use rig_core::completion::ToolAnnotations;
7use rig_core::tool::{
8    ChangeKind, FileChange, FilesChanged, Scope, Tool, ToolContext, ToolExecutionError, Touches,
9    args_schema,
10};
11use schemars::JsonSchema;
12use serde::Deserialize;
13
14use crate::fs::{Previous, io_error};
15use crate::workspace::{Local, Workspace};
16use crate::{MAX_FILE_BYTES, Seen, blocking};
17
18/// Writes a whole file, creating its parent directories. Publishes the file
19/// it changed ([`FilesChanged`]): its [resolved](Workspace) path,
20/// and the diff. The text it wrote counts as [`Seen`] whole, so an edit of
21/// the file starts from it; a write that failed forgets the file.
22#[derive(Clone, Debug)]
23pub struct Write {
24    /// Where the file is written.
25    pub workspace: Arc<dyn Workspace>,
26    /// The directory a relative path is written in; `None` is the
27    /// workspace's working directory.
28    pub base: Option<PathBuf>,
29    /// Whether a path that resolves outside `base` is refused, before
30    /// anything is written.
31    pub confine: bool,
32    /// The files seen, shared with [`Read`](crate::Read) and
33    /// [`Edit`](crate::Edit).
34    pub seen: Seen,
35}
36
37impl Default for Write {
38    /// Writes this machine's files, from this process's working directory.
39    fn default() -> Self {
40        Self {
41            workspace: Arc::new(Local),
42            base: None,
43            confine: false,
44            seen: Seen::default(),
45        }
46    }
47}
48
49/// Arguments of [`Write`].
50#[derive(Clone, Debug, PartialEq, Eq, Deserialize, JsonSchema)]
51#[serde(deny_unknown_fields)]
52pub struct WriteArgs {
53    /// The file to write.
54    pub path: String,
55    /// The whole new content.
56    pub content: String,
57}
58
59impl WriteArgs {
60    /// Write `content` to `path`.
61    pub fn new(path: impl Into<String>, content: impl Into<String>) -> Self {
62        Self {
63            path: path.into(),
64            content: content.into(),
65        }
66    }
67}
68
69impl Tool for Write {
70    const NAME: &'static str = "write";
71    type Args = WriteArgs;
72    type Output = String;
73    type Error = ToolExecutionError;
74
75    fn description(&self) -> String {
76        "Write a file with the given content, replacing it if it exists and creating parent \
77         directories. Use it for new files and complete rewrites only; change part of a file \
78         with `edit`."
79            .to_owned()
80    }
81
82    fn parameters(&self) -> serde_json::Value {
83        args_schema::<WriteArgs>()
84    }
85
86    fn annotations(&self) -> ToolAnnotations {
87        ToolAnnotations::default()
88            .with_destructive(true)
89            .with_idempotent(true)
90            .with_open_world(false)
91    }
92
93    fn footprint(&self, args: &WriteArgs) -> Touches {
94        Touches::Writes(Scope::Paths(vec![
95            self.workspace.resolve(self.base.as_deref(), &args.path),
96        ]))
97    }
98
99    async fn call(
100        &self,
101        context: &mut ToolContext,
102        args: WriteArgs,
103    ) -> Result<String, ToolExecutionError> {
104        let resolved = self
105            .workspace
106            .place(self.base.as_deref(), self.confine, &args.path)?;
107        let workspace = self.workspace.clone();
108        let seen = self.seen.clone();
109        let (output, change) = blocking(move || {
110            let path = resolved.to_string_lossy();
111            let before = workspace
112                .replace(&resolved, args.content.as_bytes(), MAX_FILE_BYTES)
113                .map_err(|error| {
114                    seen.forget(&resolved);
115                    io_error(&path, error)
116                })?;
117            seen.saw_whole(&resolved, &args.content);
118            let before = Previous::read(before);
119            let change = Previous::change(before, &path, Some(args.content.as_bytes()));
120            Ok::<_, ToolExecutionError>((wrote(&args.path, &change), change))
121        })
122        .await?;
123        context.insert_result(FilesChanged(vec![change]))?;
124        Ok(output)
125    }
126}
127
128/// The reply to a write of `path` that made `change`: that it created the
129/// file, or replaced it and, when it was read, how many lines it had.
130fn wrote(path: &str, change: &FileChange) -> String {
131    let had = change.before.as_ref().and_then(|before| before.lines);
132    match (change.kind, had) {
133        (ChangeKind::Created, _) => format!("Created {path}."),
134        (_, Some(1)) => format!("Replaced {path}, which had 1 line."),
135        (_, Some(lines)) => format!("Replaced {path}, which had {lines} lines."),
136        (_, None) => format!("Replaced {path}."),
137    }
138}
139
140#[cfg(test)]
141mod tests;
rig_tools/write.rs

rig_tools/
write.rs
