| 3d086bc3 Andy McCall ↶ |
1 |
|
| 2 |
|
| 3 |
|
| 4 |
|
| 5 |
|
| 6 |
|
| 7 |
|
| 8 |
|
| 9 |
|
| 10 |
use crate::assistants::{self, Assistant, Available, Scope}; |
| 11 |
use crate::db::{self, Pool}; |
| 12 |
use crate::repos::RepoName; |
| 13 |
use crate::Shared; |
| 14 |
use oxy_core::repo::FileEdit; |
| 15 |
use std::collections::HashMap; |
| 16 |
use std::path::PathBuf; |
| 17 |
use std::sync::{LazyLock, Mutex}; |
| 18 |
|
| 19 |
|
| 20 |
pub const MAX_DESCRIPTION: usize = 4000; |
| 21 |
|
| 22 |
pub const MAX_FILES: usize = 40; |
| 23 |
pub const MAX_FILE: usize = 64 * 1024; |
| 24 |
pub const MAX_TOTAL: usize = 256 * 1024; |
| 25 |
|
| 26 |
|
| 27 |
|
| 28 |
pub async fn assistants_for(pool: &Pool, owner: &str) -> sqlx::Result<Vec<Available>> { |
| 29 |
let scope = match db::accounts::org_by_name(pool, owner).await? { |
| 30 |
Some(org) => Scope::Org(org.id), |
| 31 |
None => Scope::Site, |
| 32 |
}; |
| 33 |
Ok(assistants::available(pool, scope, owner).await?.into_iter().filter(Available::on).collect()) |
| 34 |
} |
| 35 |
|
| 36 |
|
| 37 |
#[derive(Clone, Debug, PartialEq, Eq)] |
| 38 |
pub enum State { |
| 39 |
Running { assistant: String }, |
| 40 |
Done { assistant: String, change: String, files: Vec<String> }, |
| 41 |
Failed { assistant: String, error: String }, |
| 42 |
} |
| 43 |
|
| 44 |
static STATES: LazyLock<Mutex<HashMap<PathBuf, State>>> = LazyLock::new(|| Mutex::new(HashMap::new())); |
| 45 |
|
| 46 |
fn key(app: &Shared, name: &RepoName) -> PathBuf { |
| 47 |
app.repos.dir(name) |
| 48 |
} |
| 49 |
|
| 50 |
|
| 51 |
pub fn state(app: &Shared, name: &RepoName) -> Option<State> { |
| 52 |
STATES.lock().unwrap_or_else(|p| p.into_inner()).get(&key(app, name)).cloned() |
| 53 |
} |
| 54 |
|
| 55 |
fn set(app: &Shared, name: &RepoName, s: State) { |
| 56 |
STATES.lock().unwrap_or_else(|p| p.into_inner()).insert(key(app, name), s); |
| 57 |
} |
| 58 |
|
| 59 |
|
| 60 |
pub fn system_prompt(a: &Assistant, site: &str, owner_prompts: &str) -> String { |
| 61 |
let mut s = format!( |
| 62 |
"You are {name}, an AI assistant writing the first files of a new repository on {site}, an OxygenHub site, \ |
| 63 |
for the person who just created it; they read, change and own everything you write. \ |
| 64 |
Make a small, working starting point for what they describe: a README.md saying what it is and how to build, \ |
| 65 |
run and test it; the build files the language's own tools expect; a minimal source layout; and a first test \ |
| 66 |
where the ecosystem has a convention. No secrets, no placeholder credentials, no binary files, nothing \ |
| 67 |
generated (lock files, build output). At most {MAX_FILES} files.\n\n\ |
| 68 |
Answer with files only, each written as\n\n\ |
| 69 |
<<<FILE path/from/the/repository/root>>>\n\ |
| 70 |
the file's whole content\n\ |
| 71 |
<<<END>>>\n\n\ |
| 72 |
with nothing before the first file and nothing after the last.", |
| 73 |
name = a.name |
| 74 |
); |
| 75 |
if !owner_prompts.trim().is_empty() { |
| 76 |
s.push_str("\n\n"); |
| 77 |
s.push_str(owner_prompts.trim()); |
| 78 |
} |
| 79 |
s |
| 80 |
} |
| 81 |
|
| 82 |
|
| 83 |
|
| 84 |
pub fn parse_files(answer: &str) -> Result<Vec<(String, String)>, String> { |
| 85 |
let mut files: Vec<(String, String)> = vec![]; |
| 86 |
let mut current: Option<(String, Vec<&str>)> = None; |
| 87 |
for line in answer.lines() { |
| 88 |
let t = line.trim_end(); |
| 89 |
if let Some((path, body)) = current.as_mut() { |
| 90 |
if t.trim() == "<<<END>>>" { |
| 91 |
let path = std::mem::take(path); |
| 92 |
let content = body.join("\n") + "\n"; |
| 93 |
current = None; |
| 94 |
files.push((path, content)); |
| 95 |
} else { |
| 96 |
body.push(line); |
| 97 |
} |
| 98 |
continue; |
| 99 |
} |
| 100 |
if let Some(rest) = t.trim().strip_prefix("<<<FILE ") { |
| 101 |
let path = rest.strip_suffix(">>>").ok_or_else(|| format!("a file line is not closed with >>>: {t:?}"))?.trim(); |
| 102 |
current = Some((path.to_string(), vec![])); |
| 103 |
} |
| 104 |
} |
| 105 |
if let Some((path, _)) = current { |
| 106 |
return Err(format!("the answer stopped inside {path:?} (no <<<END>>>)")); |
| 107 |
} |
| 108 |
if files.is_empty() { |
| 109 |
return Err("the assistant answered without any files".into()); |
| 110 |
} |
| 111 |
if files.len() > MAX_FILES { |
| 112 |
return Err(format!("the assistant wrote {} files; a starter has at most {MAX_FILES}", files.len())); |
| 113 |
} |
| 114 |
let mut total = 0; |
| 115 |
let mut seen = std::collections::BTreeSet::new(); |
| 116 |
for (path, content) in &files { |
| 117 |
check_path(path)?; |
| 118 |
if !seen.insert(path.clone()) { |
| 119 |
return Err(format!("the assistant wrote {path} twice")); |
| 120 |
} |
| 121 |
if content.len() > MAX_FILE { |
| 122 |
return Err(format!("{path} is {} bytes; a starter's files are at most {MAX_FILE}", content.len())); |
| 123 |
} |
| 124 |
if content.contains('\0') { |
| 125 |
return Err(format!("{path} is not text")); |
| 126 |
} |
| 127 |
total += content.len(); |
| 128 |
} |
| 129 |
if total > MAX_TOTAL { |
| 130 |
return Err(format!("the starter is {total} bytes; at most {MAX_TOTAL}")); |
| 131 |
} |
| 132 |
Ok(files) |
| 133 |
} |
| 134 |
|
| 135 |
fn check_path(path: &str) -> Result<(), String> { |
| 136 |
let bad = |why: &str| Err(format!("the path {path:?} {why}")); |
| 137 |
if path.is_empty() || path.len() > 200 { |
| 138 |
return bad("is empty or longer than 200 characters"); |
| 139 |
} |
| 140 |
if path.starts_with('/') || path.contains('\\') || path.contains(':') { |
| 141 |
return bad("is not relative to the repository's root"); |
| 142 |
} |
| 143 |
for part in path.split('/') { |
| 144 |
if part.is_empty() || part == "." || part == ".." { |
| 145 |
return bad("has an empty, . or .. part"); |
| 146 |
} |
| 147 |
if part.eq_ignore_ascii_case(".oxy") { |
| 148 |
return bad("is inside the repository's own metadata"); |
| 149 |
} |
| 150 |
if part.chars().any(char::is_control) { |
| 151 |
return bad("has a control character"); |
| 152 |
} |
| 153 |
} |
| 154 |
Ok(()) |
| 155 |
} |
| 156 |
|
| 157 |
|
| 158 |
pub async fn is_empty(app: &Shared, name: &RepoName, bookmark: &str) -> Result<bool, String> { |
| 159 |
let b = bookmark.to_string(); |
| 160 |
app.repos |
| 161 |
.read(name, move |r, _| { |
| 162 |
let view = r.view()?; |
| 163 |
Ok(match view.bookmarks.get(&b) { |
| 164 |
None => true, |
| 165 |
Some(st) => r.store().get_state(st)?.heads.is_empty(), |
| 166 |
}) |
| 167 |
}) |
| 168 |
.await |
| 169 |
.map_err(|e| e.to_string()) |
| 170 |
} |
| 171 |
|
| 172 |
|
| 173 |
|
| 174 |
pub fn start(app: Shared, row: db::repos::RepoRow, by: db::accounts::Account, a: Assistant, description: String) { |
| 175 |
let Some(name) = RepoName::new(&row.owner, &row.name) else { return }; |
| 176 |
set(&app, &name, State::Running { assistant: a.name.clone() }); |
| 177 |
tokio::spawn(async move { |
| 178 |
let outcome = run(&app, &row, &by, &a, &description).await; |
| 179 |
let s = match outcome { |
| 180 |
Ok((change, files)) => State::Done { assistant: a.name.clone(), change, files }, |
| 181 |
Err(error) => { |
| 182 |
tracing::warn!("starter for {}/{}: {error}", row.owner, row.name); |
| 183 |
State::Failed { assistant: a.name.clone(), error } |
| 184 |
} |
| 185 |
}; |
| 186 |
set(&app, &name, s); |
| 187 |
}); |
| 188 |
} |
| 189 |
|
| 190 |
async fn run( |
| 191 |
app: &Shared, |
| 192 |
row: &db::repos::RepoRow, |
| 193 |
by: &db::accounts::Account, |
| 194 |
a: &Assistant, |
| 195 |
description: &str, |
| 196 |
) -> Result<(String, Vec<String>), String> { |
| 197 |
let name = RepoName::new(&row.owner, &row.name).ok_or("bad repository name")?; |
| 198 |
let signer = crate::web::edit::signer_for(app, by.id).await?; |
| 199 |
let scope = match db::accounts::org_by_name(&app.db, &row.owner).await.map_err(|e| e.to_string())? { |
| 200 |
Some(org) => Scope::Org(org.id), |
| 201 |
None => Scope::Site, |
| 202 |
}; |
| 203 |
let mut prompts = assistants::inherited_prompts(&app.db, scope, &row.owner).await.map_err(|e| e.to_string())?; |
| 204 |
let own = assistants::prompt(&app.db, scope).await.map_err(|e| e.to_string())?; |
| 205 |
if !own.is_empty() && scope != Scope::Site { |
| 206 |
prompts.push((format!("the {} organisation", row.owner), own)); |
| 207 |
} else if !own.is_empty() { |
| 208 |
prompts.push(("the site".to_string(), own)); |
| 209 |
} |
| 210 |
let prompts: String = |
| 211 |
prompts.into_iter().map(|(from, text)| format!("## Instructions from {from}\n\n{text}\n")).collect::<Vec<_>>().join("\n"); |
| 212 |
let site = app.config.base_url(); |
| 213 |
let user = format!("The repository is {}/{}.\n\nWhat they want:\n\n{}", row.owner, row.name, description.trim()); |
| 214 |
let answer = assistants::ask(app, a, &system_prompt(a, site, &prompts), &user, assistants::BACKGROUND_TIMEOUT).await?; |
| 215 |
let files = parse_files(&answer)?; |
| 216 |
if !is_empty(app, &name, &row.default_bookmark).await? { |
| 217 |
return Err("the repository has files now, so the starter was not added".into()); |
| 218 |
} |
| 219 |
let paths: Vec<String> = files.iter().map(|(p, _)| p.clone()).collect(); |
| 220 |
let edits: Vec<FileEdit> = files.into_iter().map(|(path, content)| FileEdit { path, content: Some(content.into_bytes()) }).collect(); |
| 221 |
let short: String = description.trim().chars().take(400).collect(); |
| 222 |
let message = format!("Starter files drafted by {} ({})\n\nAsked for by {}: {short}", a.name, a.model, by.username); |
| 223 |
let b = row.default_bookmark.clone(); |
| 224 |
let done = crate::repos::edits::commit(&app.repos, &name, b.clone(), b, edits, message, signer).await.map_err(|e| e.to_string())?; |
| 225 |
db::accounts::audit(&app.db, Some(by), "repo.starter", &name.to_string(), &format!("{} files drafted by {}", paths.len(), a.name)) |
| 226 |
.await |
| 227 |
.ok(); |
| 228 |
Ok((done.change.to_hex(), paths)) |
| 229 |
} |
| 230 |
|
| 231 |
#[cfg(test)] |
| 232 |
mod tests { |
| 233 |
use super::*; |
| 234 |
|
| 235 |
#[test] |
| 236 |
fn an_answer_is_read_as_files_and_checked() { |
| 237 |
let answer = "<<<FILE README.md>>>\n# Hello\n<<<END>>>\n<<<FILE src/main.rs>>>\nfn main() {}\n<<<END>>>\n"; |
| 238 |
let files = parse_files(answer).unwrap(); |
| 239 |
assert_eq!( |
| 240 |
files, |
| 241 |
vec![("README.md".to_string(), "# Hello\n".to_string()), ("src/main.rs".to_string(), "fn main() {}\n".to_string())] |
| 242 |
); |
| 243 |
assert!(parse_files("Sure! Here you go.").unwrap_err().contains("without any files")); |
| 244 |
assert!(parse_files("<<<FILE a>>>\nx").unwrap_err().contains("no <<<END>>>")); |
| 245 |
for bad in ["/etc/passwd", "../up", "a//b", "C:\\x", ".oxy/config"] { |
| 246 |
assert!(parse_files(&format!("<<<FILE {bad}>>>\nx\n<<<END>>>")).is_err(), "{bad}"); |
| 247 |
} |
| 248 |
assert!(parse_files("<<<FILE a>>>\n1\n<<<END>>>\n<<<FILE a>>>\n2\n<<<END>>>").unwrap_err().contains("twice")); |
| 249 |
let many: String = (0..=MAX_FILES).map(|i| format!("<<<FILE f{i}>>>\nx\n<<<END>>>\n")).collect(); |
| 250 |
assert!(parse_files(&many).unwrap_err().contains("at most")); |
| 251 |
let big = format!("<<<FILE big>>>\n{}\n<<<END>>>", "x".repeat(MAX_FILE + 1)); |
| 252 |
assert!(parse_files(&big).unwrap_err().contains("at most")); |
| 253 |
} |
| 254 |
} |