//! AI assistants (hub spec §13.4, #97): chat-completion endpoints the hub //! may ask to review a proposal or draft its description. Assistants are //! kept at three levels, site, organisation and repository, and combined //! downwards: a repository sees its organisation's and the site's beside //! its own, the innermost of a name wins, and one disabled at an outer //! level is off for everything inside it. Each level adds a prompt; the //! prompt in force for a repository is the site's, the organisation's and //! the repository's own, each under a heading. Tokens are sealed with //! ChaCha20-Poly1305 under `keys/assistants.key`, shown masked, and never //! leave the server. use crate::db::{now_ms, Pool}; use crate::Shared; use rand::RngCore; use ring::aead; use sqlx::FromRow; use std::path::Path; use std::time::Duration; /// The request shapes an assistant may speak. pub const APIS: [&str; 2] = ["chat-completions", "anthropic-messages"]; /// Characters of context (prompt, description, diffs) a call may carry /// unless the assistant says otherwise. pub const DEFAULT_BUDGET: i64 = 120_000; /// How long a call may take while a person waits for it (a draft, a /// connection test): inside Cloudflare's 100 seconds for a proxied request. pub const CALL_TIMEOUT: Duration = Duration::from_secs(90); /// How long a call answered in the background may take (a review, a /// summary): a model that thinks first needs minutes on a large proposal /// (oxygenhub#201). pub const BACKGROUND_TIMEOUT: Duration = Duration::from_secs(300); /// The longest answer asked for: room for a model that thinks before it /// writes (oxygenhub#201: 4096 was spent on thinking alone). pub const MAX_ANSWER_TOKENS: u32 = 16_000; /// The most an answer is kept. const MAX_ANSWER: usize = 64 * 1024; pub use crate::levels::Scope; #[derive(Clone, Debug, FromRow)] pub struct Assistant { pub id: i64, pub scope: String, pub scope_id: i64, pub name: String, pub endpoint: String, /// One of [`APIS`]. pub api: String, pub model: String, /// Extra request headers, one `Name: value` per line. pub headers: String, /// The sealed token; empty when the endpoint needs none. pub token: Vec, pub enabled: bool, pub budget: i64, pub created_ms: i64, pub updated_ms: i64, } impl Assistant { pub fn scope(&self) -> Scope { match self.scope.as_str() { "org" => Scope::Org(self.scope_id), "repo" => Scope::Repo(self.scope_id), _ => Scope::Site, } } pub fn has_token(&self) -> bool { !self.token.is_empty() } /// The extra headers as pairs; malformed lines are skipped. pub fn header_list(&self) -> Vec<(String, String)> { parse_headers(&self.headers).unwrap_or_default() } } /// What a settings form supplies for an assistant. #[derive(Clone, Debug, Default)] pub struct Fields { pub name: String, pub endpoint: String, pub api: String, pub model: String, pub headers: String, pub enabled: bool, pub budget: i64, } /// `Name: value` lines to pairs: `Err(line)` names the first bad one. pub fn parse_headers(text: &str) -> Result, String> { let mut out = vec![]; for line in text.lines().map(str::trim).filter(|l| !l.is_empty()) { let Some((name, value)) = line.split_once(':') else { return Err(line.to_string()) }; let (name, value) = (name.trim(), value.trim()); let name_ok = !name.is_empty() && name.chars().all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_'); if !name_ok || value.chars().any(|c| c.is_control()) { return Err(line.to_string()); } out.push((name.to_string(), value.to_string())); } Ok(out) } // -------------------------------------------------------------- sealing /// The key that seals tokens, kept in `/keys/assistants.key`. #[derive(Clone)] pub struct SecretKey([u8; 32]); impl SecretKey { pub fn load_or_create(keys: &Path) -> anyhow::Result { SecretKey::load_or_create_named(keys, "assistants.key") } /// The same, for another key file beside it (`secrets.key` seals Flow /// secrets, #107). pub fn load_or_create_named(keys: &Path, file: &str) -> anyhow::Result { std::fs::create_dir_all(keys)?; let path = keys.join(file); if let Ok(bytes) = std::fs::read(&path) { if bytes.len() == 32 { return Ok(SecretKey(bytes.try_into().unwrap())); } } let mut k = [0u8; 32]; rand::rngs::OsRng.fill_bytes(&mut k); std::fs::write(&path, k)?; #[cfg(unix)] { use std::os::unix::fs::PermissionsExt; std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600))?; } Ok(SecretKey(k)) } fn key(&self) -> aead::LessSafeKey { aead::LessSafeKey::new(aead::UnboundKey::new(&aead::CHACHA20_POLY1305, &self.0).expect("key length")) } /// Seal a token: a random nonce, then the ciphertext and tag. An /// empty token seals to nothing, meaning "no token". pub fn seal(&self, plain: &str) -> Vec { if plain.is_empty() { return vec![]; } let mut nonce = [0u8; 12]; rand::rngs::OsRng.fill_bytes(&mut nonce); let mut out = plain.as_bytes().to_vec(); self.key().seal_in_place_append_tag(aead::Nonce::assume_unique_for_key(nonce), aead::Aad::empty(), &mut out).expect("seal"); let mut sealed = nonce.to_vec(); sealed.extend(out); sealed } /// The token back, or `None` when the bytes were sealed under another /// key or damaged. Nothing opens to an empty token. pub fn open(&self, sealed: &[u8]) -> Option { if sealed.is_empty() { return Some(String::new()); } let (nonce, rest) = sealed.split_at_checked(12)?; let mut buf = rest.to_vec(); let nonce = aead::Nonce::try_assume_unique_for_key(nonce).ok()?; let plain = self.key().open_in_place(nonce, aead::Aad::empty(), &mut buf).ok()?; String::from_utf8(plain.to_vec()).ok() } } // ------------------------------------------------------------- storage pub async fn list(pool: &Pool, scope: Scope) -> sqlx::Result> { sqlx::query_as("SELECT * FROM assistants WHERE scope = $1 AND scope_id = $2 ORDER BY name") .bind(scope.kind()) .bind(scope.id()) .fetch_all(pool) .await } pub async fn get(pool: &Pool, id: i64) -> sqlx::Result> { sqlx::query_as("SELECT * FROM assistants WHERE id = $1").bind(id).fetch_optional(pool).await } /// Add an assistant; `None` when the level already has one of that name. pub async fn create(pool: &Pool, scope: Scope, f: &Fields, token: &[u8]) -> sqlx::Result> { let now = now_ms(); let r = sqlx::query_as( "INSERT INTO assistants (scope, scope_id, name, endpoint, api, model, headers, token, enabled, budget, created_ms, updated_ms) \ VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) ON CONFLICT DO NOTHING RETURNING *", ) .bind(scope.kind()) .bind(scope.id()) .bind(&f.name) .bind(&f.endpoint) .bind(&f.api) .bind(&f.model) .bind(&f.headers) .bind(token) .bind(f.enabled) .bind(f.budget) .bind(now) .bind(now) .fetch_optional(pool) .await?; Ok(r) } /// Change an assistant's details; `token` is `None` to keep the stored /// one. False when the name is taken at that level. pub async fn update(pool: &Pool, id: i64, f: &Fields, token: Option<&[u8]>) -> sqlx::Result { let r = match token { Some(t) => sqlx::query( "UPDATE assistants SET name = $1, endpoint = $2, api = $3, model = $4, headers = $5, enabled = $6, budget = $7, token = $8, updated_ms = $9 \ WHERE id = $10 AND NOT EXISTS (SELECT 1 FROM assistants o WHERE o.scope = assistants.scope AND o.scope_id = assistants.scope_id AND o.name = $11 AND o.id <> $12)", ) .bind(&f.name) .bind(&f.endpoint) .bind(&f.api) .bind(&f.model) .bind(&f.headers) .bind(f.enabled) .bind(f.budget) .bind(t) .bind(now_ms()) .bind(id) .bind(&f.name) .bind(id) .execute(pool) .await?, None => sqlx::query( "UPDATE assistants SET name = $1, endpoint = $2, api = $3, model = $4, headers = $5, enabled = $6, budget = $7, updated_ms = $8 \ WHERE id = $9 AND NOT EXISTS (SELECT 1 FROM assistants o WHERE o.scope = assistants.scope AND o.scope_id = assistants.scope_id AND o.name = $10 AND o.id <> $11)", ) .bind(&f.name) .bind(&f.endpoint) .bind(&f.api) .bind(&f.model) .bind(&f.headers) .bind(f.enabled) .bind(f.budget) .bind(now_ms()) .bind(id) .bind(&f.name) .bind(id) .execute(pool) .await?, }; Ok(r.rows_affected() == 1) } /// A renamed assistant keeps its name on the reviews it answered and the /// events that mention it, for the repositories its level covers: the /// name is what people read, not a key. pub async fn rename_in_history(pool: &Pool, scope: Scope, owner_hint: &str, old: &str, new: &str) -> sqlx::Result<()> { if old == new || old.is_empty() || new.is_empty() { return Ok(()); } let like = format!("%{old}%"); match scope { Scope::Site => { for table in ["proposal_reviews", "proposal_comments"] { sqlx::query(&format!("UPDATE {table} SET via = $1 WHERE via = $2")).bind(new).bind(old).execute(pool).await?; } sqlx::query("UPDATE proposal_events SET event = replace(event, $1, $2) WHERE event LIKE $3") .bind(old) .bind(new) .bind(&like) .execute(pool) .await?; } Scope::Org(_) => { let in_org = "proposal_id IN (SELECT p.id FROM proposals p JOIN repositories r ON r.id = p.repo_id WHERE r.owner = $1)"; for table in ["proposal_reviews", "proposal_comments"] { sqlx::query(&format!("UPDATE {table} SET via = $1 WHERE via = $2 AND {in_org}")) .bind(new) .bind(old) .bind(owner_hint) .execute(pool) .await?; } sqlx::query(&format!("UPDATE proposal_events SET event = replace(event, $1, $2) WHERE event LIKE $3 AND {in_org}")) .bind(old) .bind(new) .bind(&like) .bind(owner_hint) .execute(pool) .await?; } Scope::Repo(id) => { let in_repo = "proposal_id IN (SELECT id FROM proposals WHERE repo_id = $1)"; for table in ["proposal_reviews", "proposal_comments"] { sqlx::query(&format!("UPDATE {table} SET via = $1 WHERE via = $2 AND {in_repo}")) .bind(new) .bind(old) .bind(id) .execute(pool) .await?; } sqlx::query(&format!("UPDATE proposal_events SET event = replace(event, $1, $2) WHERE event LIKE $3 AND {in_repo}")) .bind(old) .bind(new) .bind(&like) .bind(id) .execute(pool) .await?; } } Ok(()) } pub async fn set_enabled(pool: &Pool, id: i64, enabled: bool) -> sqlx::Result<()> { sqlx::query("UPDATE assistants SET enabled = $1, updated_ms = $2 WHERE id = $3") .bind(enabled) .bind(now_ms()) .bind(id) .execute(pool) .await?; Ok(()) } pub async fn delete(pool: &Pool, id: i64) -> sqlx::Result<()> { sqlx::query("DELETE FROM assistants WHERE id = $1").bind(id).execute(pool).await?; Ok(()) } /// Forget a level's assistants and prompt, when the organisation or /// repository goes. pub async fn delete_scope(pool: &Pool, scope: Scope) -> sqlx::Result<()> { sqlx::query("DELETE FROM assistants WHERE scope = $1 AND scope_id = $2").bind(scope.kind()).bind(scope.id()).execute(pool).await?; sqlx::query("DELETE FROM assistant_prompts WHERE scope = $1 AND scope_id = $2") .bind(scope.kind()) .bind(scope.id()) .execute(pool) .await?; Ok(()) } /// The prompt a level adds, empty when none. pub async fn prompt(pool: &Pool, scope: Scope) -> sqlx::Result { Ok(sqlx::query_scalar("SELECT prompt FROM assistant_prompts WHERE scope = $1 AND scope_id = $2") .bind(scope.kind()) .bind(scope.id()) .fetch_optional(pool) .await? .unwrap_or_default()) } pub async fn set_prompt(pool: &Pool, scope: Scope, text: &str) -> sqlx::Result<()> { if text.trim().is_empty() { sqlx::query("DELETE FROM assistant_prompts WHERE scope = $1 AND scope_id = $2") .bind(scope.kind()) .bind(scope.id()) .execute(pool) .await?; } else { sqlx::query( "INSERT INTO assistant_prompts (scope, scope_id, prompt, updated_ms) VALUES ($1, $2, $3, $4) \ ON CONFLICT (scope, scope_id) DO UPDATE SET prompt = excluded.prompt, updated_ms = excluded.updated_ms", ) .bind(scope.kind()) .bind(scope.id()) .bind(text.trim()) .bind(now_ms()) .execute(pool) .await?; } Ok(()) } // ----------------------------------------------------- what applies where /// An assistant seen from a level: where it comes from, and whether it is /// on once outer levels have had their say. #[derive(Clone, Debug)] pub struct Available { pub assistant: Assistant, /// `the site`, `the acme organisation`, or empty for the level's own. pub from: String, /// Off because a level above disabled an assistant of this name. pub off_above: bool, } impl Available { pub fn on(&self) -> bool { self.assistant.enabled && !self.off_above } } /// Every assistant a level sees, outermost first, one per name: the /// innermost definition of a name wins, and a name disabled at an outer /// level is off inside it whatever the inner level says. pub async fn available(pool: &Pool, scope: Scope, owner: &str) -> sqlx::Result> { let mut levels = crate::levels::above(pool, scope, owner).await?; levels.push((scope, String::new())); let mut out: Vec = vec![]; let mut off: Vec = vec![]; for (level, from) in levels { for a in list(pool, level).await? { let off_above = off.iter().any(|n| n.eq_ignore_ascii_case(&a.name)); if !a.enabled && !off.iter().any(|n| n.eq_ignore_ascii_case(&a.name)) { off.push(a.name.clone()); } out.retain(|o| !o.assistant.name.eq_ignore_ascii_case(&a.name)); out.push(Available { assistant: a, from: from.clone(), off_above }); } } Ok(out) } /// The assistants a repository may ask, by name. pub async fn enabled_for_repo(pool: &Pool, repo: &crate::db::repos::RepoRow) -> sqlx::Result> { Ok(available(pool, Scope::Repo(repo.id), &repo.owner).await?.into_iter().filter(Available::on).map(|a| a.assistant).collect()) } /// The prompts a level inherits, outermost first: (where from, text), /// only the levels that set one. pub async fn inherited_prompts(pool: &Pool, scope: Scope, owner: &str) -> sqlx::Result> { let mut out = vec![]; for (level, from) in crate::levels::above(pool, scope, owner).await? { let text = prompt(pool, level).await?; if !text.is_empty() { out.push((from, text)); } } Ok(out) } /// The prompt in force for a repository: each level's text under a /// heading, outermost first, so a project adds to what the site says. pub async fn prompt_for_repo(pool: &Pool, repo: &crate::db::repos::RepoRow) -> sqlx::Result { let scope = Scope::Repo(repo.id); let mut parts = inherited_prompts(pool, scope, &repo.owner).await?; let own = prompt(pool, scope).await?; if !own.is_empty() { parts.push(("this repository".to_string(), own)); } Ok(parts.into_iter().map(|(from, text)| format!("## Instructions from {from}\n\n{text}\n")).collect::>().join("\n")) } // -------------------------------------------------------------- calling /// Ask an assistant: the system prompt and the request, the answer's /// text. Errors are for the person who asked: a status and the first /// line of what the endpoint said, or why no answer came. pub async fn ask(app: &Shared, a: &Assistant, system: &str, user: &str, wait: Duration) -> Result { let Some(token) = app.assistant_key.open(&a.token) else { return Err("its token cannot be read: set it again".into()); }; let client = crate::net::client(&a.endpoint, &app.config.outbound.allow, wait).await?; let mut req = client.post(&a.endpoint).header("content-type", "application/json"); let body = match a.api.as_str() { "anthropic-messages" => { if !token.is_empty() { req = req.header("x-api-key", &token); } req = req.header("anthropic-version", "2023-06-01"); serde_json::json!({ "model": a.model, "max_tokens": MAX_ANSWER_TOKENS, "system": system, "messages": [{ "role": "user", "content": user }], }) } _ => { if !token.is_empty() { req = req.header("authorization", format!("Bearer {token}")); } // No length here: providers differ (`max_tokens`, or // `max_completion_tokens` for newer OpenAI models, with // different ceilings), and their defaults are generous. serde_json::json!({ "model": a.model, "messages": [ { "role": "system", "content": system }, { "role": "user", "content": user }, ], }) } }; for (name, value) in a.header_list() { req = req.header(name, value); } let r = req.json(&body).send().await.map_err(|e| { if e.is_timeout() { format!("no answer within {} seconds", wait.as_secs()) } else { scrub(&e.to_string(), &token) } })?; let status = r.status(); let text: String = r.text().await.unwrap_or_default().chars().take(MAX_ANSWER).collect(); if !status.is_success() { let why = serde_json::from_str::(&text) .ok() .and_then(|v| v["error"]["message"].as_str().or(v["error"].as_str()).or(v["message"].as_str()).map(str::to_string)) .unwrap_or_else(|| text.lines().next().unwrap_or("").chars().take(160).collect()); return Err(scrub(&format!("{} {}", status.as_u16(), why.trim()), &token)); } let v: serde_json::Value = serde_json::from_str(&text).map_err(|_| "the answer was not JSON".to_string())?; let (answer, stop) = answer_of(&a.api, &v); finish(answer, stop) } /// Why an answer ended, as the person who asked needs to know it. #[derive(Debug, PartialEq, Eq)] enum Stop { /// It finished. Done, /// It reached the length it was allowed (thinking counts against it). Length, /// The provider declined to answer, with its category if it gave one. Declined(Option), } /// The text of a provider's answer (text parts only: thinking is not /// part of the answer), and why it ended. fn answer_of(api: &str, v: &serde_json::Value) -> (String, Stop) { match api { "anthropic-messages" => { let text = v["content"] .as_array() .map(|parts| { parts .iter() .filter(|p| p["type"].as_str().is_none_or(|t| t == "text")) .filter_map(|p| p["text"].as_str()) .collect::>() .join("") }) .unwrap_or_default(); let stop = match v["stop_reason"].as_str() { Some("max_tokens") => Stop::Length, Some("refusal") => Stop::Declined(v["stop_details"]["category"].as_str().map(str::to_string)), _ => Stop::Done, }; (text, stop) } _ => { let choice = &v["choices"][0]; let text = match &choice["message"]["content"] { serde_json::Value::String(s) => s.clone(), serde_json::Value::Array(parts) => parts.iter().filter_map(|p| p["text"].as_str()).collect::>().join(""), _ => String::new(), }; let stop = match choice["finish_reason"].as_str() { Some("length") => Stop::Length, Some("content_filter") => Stop::Declined(None), _ => Stop::Done, }; (text, stop) } } } /// The answer, or why there is none. An answer cut short is kept, with a /// note saying so. fn finish(answer: String, stop: Stop) -> Result { match stop { Stop::Declined(category) => Err(match category { Some(c) => format!("it declined to answer ({c})"), None => "it declined to answer".into(), }), Stop::Length if answer.trim().is_empty() => { Err("it ran out of room before writing an answer (it spent its length limit thinking): try a smaller proposal".into()) } Stop::Length => Ok(format!("{}\n\n*(Cut short: the answer reached its length limit.)*", answer.trim_end())), Stop::Done if answer.trim().is_empty() => Err("the answer carried no text".into()), Stop::Done => Ok(answer), } } /// Providers the add form offers, with what they need (#109): /// (id, label, name, endpoint, api). pub const PRESETS: &[(&str, &str, &str, &str, &str)] = &[ ("anthropic", "Anthropic (Claude)", "Claude", "https://api.anthropic.com/v1/messages", "anthropic-messages"), ("openai", "OpenAI (ChatGPT)", "ChatGPT", "https://api.openai.com/v1/chat/completions", "chat-completions"), ("gemini", "Google Gemini", "Gemini", "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions", "chat-completions"), ("deepseek", "DeepSeek", "DeepSeek", "https://api.deepseek.com/chat/completions", "chat-completions"), ("mistral", "Mistral", "Mistral", "https://api.mistral.ai/v1/chat/completions", "chat-completions"), ("perplexity", "Perplexity", "Perplexity", "https://api.perplexity.ai/chat/completions", "chat-completions"), ("ollama", "Ollama (on your network)", "Ollama", "http://localhost:11434/v1/chat/completions", "chat-completions"), ]; /// Where an endpoint lists its models: `…/chat/completions` → `…/models`, /// `…/messages` → `…/models`. pub fn models_url(endpoint: &str) -> Option { let e = endpoint.trim_end_matches('/'); e.strip_suffix("/chat/completions").or_else(|| e.strip_suffix("/messages")).map(|base| format!("{base}/models")) } /// The chat models an endpoint offers for this token, sorted; `Err` when /// it cannot say (no list, refused, unreachable). Embedding, speech and /// image models are left out. pub async fn list_models(endpoint: &str, api: &str, token: &str, allow: &[String]) -> Result, String> { let url = models_url(endpoint).ok_or("this endpoint has no model list next to it")?; let client = crate::net::client(&url, allow, std::time::Duration::from_secs(15)).await?; let mut req = client.get(&url); if api == "anthropic-messages" { req = req.header("anthropic-version", "2023-06-01"); if !token.is_empty() { req = req.header("x-api-key", token); } } else if !token.is_empty() { req = req.header("authorization", format!("Bearer {token}")); } let r = req.send().await.map_err(|e| scrub(&e.to_string(), token))?; let status = r.status(); if !status.is_success() { return Err(format!("the model list answered {}", status.as_u16())); } let v: serde_json::Value = r.json().await.map_err(|_| "the model list was not JSON".to_string())?; let skip = ["embed", "tts", "whisper", "dall-e", "moderation", "audio", "image", "realtime", "transcribe", "search-preview"]; let mut out: Vec = v["data"] .as_array() .or(v["models"].as_array()) .ok_or("the model list had no data")? .iter() .filter_map(|m| m["id"].as_str().or(m["name"].as_str())) .map(|id| id.strip_prefix("models/").unwrap_or(id).to_string()) .filter(|id| !skip.iter().any(|s| id.to_lowercase().contains(s))) .collect(); out.sort(); out.dedup(); if out.is_empty() { return Err("the model list was empty".into()); } Ok(out) } /// A tiny exchange, for the settings page's TEST: the start of the answer. pub async fn test(app: &Shared, a: &Assistant) -> Result { let answer = ask(app, a, "You are being connected to an OxygenHub site. Answer briefly.", "Reply with the single word OK.", CALL_TIMEOUT) .await?; Ok(answer.trim().chars().take(80).collect()) } /// Never let the token through in an error message. fn scrub(text: &str, token: &str) -> String { if token.is_empty() { text.to_string() } else { text.replace(token, "***") } } #[cfg(test)] mod tests { use super::*; #[test] fn tokens_seal_and_open_only_under_their_key() { let dir = tempfile::tempdir().unwrap(); let k = SecretKey::load_or_create(dir.path()).unwrap(); let sealed = k.seal("sk-very-secret"); assert_ne!(sealed, b"sk-very-secret"); assert!(!String::from_utf8_lossy(&sealed).contains("secret")); assert_eq!(k.open(&sealed).as_deref(), Some("sk-very-secret")); assert_eq!(SecretKey::load_or_create(dir.path()).unwrap().open(&sealed).as_deref(), Some("sk-very-secret"), "the key is kept"); let other = SecretKey::load_or_create(tempfile::tempdir().unwrap().path()).unwrap(); assert_eq!(other.open(&sealed), None); assert_eq!(k.open(&[1, 2, 3]), None); assert_eq!(k.seal(""), Vec::::new()); assert_eq!(k.open(&[]).as_deref(), Some("")); } #[test] fn answers_say_why_they_stopped() { let a = |v: serde_json::Value| { let (text, stop) = answer_of("anthropic-messages", &v); finish(text, stop) }; // #201: all of the length spent thinking. let thought = serde_json::json!({ "stop_reason": "max_tokens", "content": [{ "type": "thinking", "thinking": "" }] }); assert!(a(thought).unwrap_err().contains("ran out of room")); let cut = serde_json::json!({ "stop_reason": "max_tokens", "content": [{ "type": "thinking", "thinking": "" }, { "type": "text", "text": "Half" }] }); assert!(a(cut).unwrap().starts_with("Half\n\n*(Cut short")); let declined = serde_json::json!({ "stop_reason": "refusal", "stop_details": { "category": "cyber" }, "content": [] }); assert_eq!(a(declined).unwrap_err(), "it declined to answer (cyber)"); let ok = serde_json::json!({ "stop_reason": "end_turn", "content": [{ "type": "text", "text": "Fine." }] }); assert_eq!(a(ok).unwrap(), "Fine."); let chat = |v: serde_json::Value| { let (text, stop) = answer_of("chat-completions", &v); finish(text, stop) }; assert!(chat(serde_json::json!({ "choices": [{ "finish_reason": "length", "message": { "content": "" } }] })) .unwrap_err() .contains("ran out of room")); assert_eq!( chat(serde_json::json!({ "choices": [{ "finish_reason": "content_filter", "message": { "content": "" } }] })).unwrap_err(), "it declined to answer" ); assert_eq!( chat(serde_json::json!({ "choices": [{ "finish_reason": "stop", "message": { "content": "Yes." } }] })).unwrap(), "Yes." ); } #[test] fn a_refusal_for_length_is_told_from_others() { for e in [ "400 This model's maximum context length is 8192 tokens. Please reduce the length of the messages.", "400 prompt is too long: 210000 tokens > 200000 maximum", "413 Request too large", "400 input length exceeds the context window", ] { assert!(super::too_long(e), "{e}"); } for e in ["401 invalid api key", "500 too long a queue", "no answer within 90 seconds", "400 model not found"] { assert!(!super::too_long(e), "{e}"); } } #[test] fn headers_parse_or_name_the_bad_line() { assert_eq!( parse_headers("X-Org: acme\n\n OpenAI-Beta : v1 ").unwrap(), vec![("X-Org".into(), "acme".into()), ("OpenAI-Beta".into(), "v1".into())] ); assert_eq!(parse_headers("no colon").unwrap_err(), "no colon"); assert_eq!(parse_headers("bad name: x").unwrap_err(), "bad name: x"); } } // ------------------------------------------------------- reviewing /// At most this many assistant calls per repository per hour. pub const CALLS_PER_HOUR: usize = 30; /// Who is answering what, and recent calls per repository. #[derive(Default)] pub struct Gate { busy: std::collections::HashSet<(i64, String)>, recent: std::collections::HashMap>, } impl Gate { /// Take the slot for (proposal, assistant), or say why not. fn take(&mut self, repo_id: i64, proposal_id: i64, name: &str) -> Result<(), String> { let key = (proposal_id, name.to_lowercase()); if self.busy.contains(&key) { return Err(format!("{name} is already answering this proposal")); } let now = now_ms(); let calls = self.recent.entry(repo_id).or_default(); calls.retain(|t| now - *t < 3_600_000); if calls.len() >= CALLS_PER_HOUR { return Err(format!("this repository has asked assistants {CALLS_PER_HOUR} times in the last hour; try later")); } calls.push(now); self.busy.insert(key); Ok(()) } fn release(&mut self, proposal_id: i64, name: &str) { self.busy.remove(&(proposal_id, name.to_lowercase())); } /// The assistants answering on a proposal now (lower-case names, sorted). pub fn answering(&self, proposal_id: i64) -> Vec { let mut v: Vec = self.busy.iter().filter(|(p, _)| *p == proposal_id).map(|(_, n)| n.clone()).collect(); v.sort(); v } } /// What a reviewer is told about a proposal: its title, description, /// changes and diff, cut to `budget` characters. pub fn review_context( row: &crate::db::repos::RepoRow, p: &crate::db::proposals::ProposalRow, loaded: &crate::web::proposals::Loaded, budget: usize, ) -> String { let mut head = format!("# Proposal #{}: {}\n\nRepository {}/{} · into {} · by {}\n\n", p.number, p.title, row.owner, row.name, p.target, p.author); if !p.body.trim().is_empty() { head.push_str("## Description\n\n"); head.push_str(p.body.trim()); head.push_str("\n\n"); } context_text(head, &loaded.changes, &loaded.diffs, budget) } /// What a drafter is told: the title so far, the changes and the diff. pub fn draft_context( row: &crate::db::repos::RepoRow, title: &str, target: &str, changes: &[crate::repos::view::ChangeView], diffs: &[oxy_core::repo::FileDiff], budget: usize, ) -> String { let head = format!( "# Proposal: {}\n\nRepository {}/{} · into {}\n\n", if title.trim().is_empty() { "(untitled)" } else { title.trim() }, row.owner, row.name, target ); context_text(head, changes, diffs, budget) } /// `head`, the list of changes, then the diff as `-`/`+` lines per hunk, /// binary files noted, the diff cut so the whole fits `budget`. fn context_text(mut text: String, changes: &[crate::repos::view::ChangeView], diffs: &[oxy_core::repo::FileDiff], budget: usize) -> String { text.push_str(&format!("## Changes ({})\n\n", changes.len())); for c in changes { text.push_str(&format!("- {} {}\n", c.short, c.title)); let body = c.body.trim(); if !body.is_empty() { for line in body.lines().take(20) { text.push_str(" "); text.push_str(line); text.push('\n'); } } } // Every file with its size first, so a model told only part of the // diff still knows the whole (oxygenhub#248). text.push_str(&format!("\n## Files ({})\n\n", diffs.len())); let mut parts: Vec<(String, String)> = Vec::with_capacity(diffs.len()); for f in diffs { let (added, removed) = f.hunks.iter().fold((0, 0), |(a, r), h| (a + h.added.len(), r + h.removed.len())); text.push_str(&format!("- {} ({}, +{added} -{removed})\n", f.path, f.kind)); let mut diff = format!("### {} ({})\n", f.path, f.kind); if let Some(from) = &f.from { diff.push_str(&format!("renamed from {from}\n")); } let binary = f.hunks.iter().any(|h| h.added.iter().chain(h.removed.iter().map(|(_, l)| l)).any(|l| l.contains(&0))); if binary { diff.push_str("(binary)\n\n"); parts.push((f.path.clone(), diff)); continue; } diff.push_str("```\n"); for h in &f.hunks { diff.push_str(&format!("@@ after line {} @@\n", h.base_line)); for (_, l) in &h.removed { diff.push('-'); diff.push_str(String::from_utf8_lossy(l).trim_end_matches('\n')); diff.push('\n'); } for l in &h.added { diff.push('+'); diff.push_str(String::from_utf8_lossy(l).trim_end_matches('\n')); diff.push('\n'); } } diff.push_str("```\n\n"); parts.push((f.path.clone(), diff)); } text.push_str("\n## Diff\n\n"); // Whole files, in order, while they fit. One that doesn't is cut where // the room ends if there is a fair amount left, else left out (a // smaller one after it may still go in); either way it is named. let head = text.chars().count(); let mut room = budget.saturating_sub(head + 200); let mut left_out = Vec::new(); for (path, diff) in parts { let n = diff.chars().count(); if n <= room { text.push_str(&diff); room -= n; } else if room >= 1_500 { let keep = room - 80; text.extend(diff.chars().take(keep)); text.push_str("\n(the rest of this file was cut to fit)\n```\n\n"); left_out.push(format!("{path} (cut)")); room = 0; } else { left_out.push(path); } } if !left_out.is_empty() { let names: String = left_out.join(", ").chars().take(room.max(400)).collect(); text.push_str(&format!("(left out to fit the context budget, {} file(s): {names})\n", left_out.len())); } text } /// The system prompt for a review: who the assistant is, what a review /// here is, then the prompts the levels set. pub fn review_system_prompt(a: &Assistant, site: &str, prompts: &str) -> String { let mut s = format!( "You are {name}, an AI assistant reviewing a proposal on {site}, an OxygenHub site. \ A person asked for your review; it is recorded as a review comment made via you, under their name. \ You cannot approve: only people approve. Answer in Markdown: start with a one-paragraph summary, \ then findings (correctness first, then clarity), each naming the file and quoting the lines it is about; \ keep praise out unless something is notably well done; say so plainly when you find nothing wrong.", name = a.name ); if !prompts.trim().is_empty() { s.push_str("\n\n"); s.push_str(prompts.trim()); } s } /// Ask `a` to review proposal `p` for `by`: records the request, calls the /// endpoint in the background, and stores the answer as a `comment` review /// via the assistant (or an event saying why none came). `Err` is for the /// person who asked, when nothing was started. pub async fn request_review( app: &Shared, row: &crate::db::repos::RepoRow, p: &crate::db::proposals::ProposalRow, by: &crate::db::accounts::Account, a: Assistant, ) -> Result<(), String> { app.assistant_gate.lock().unwrap_or_else(|e| e.into_inner()).take(row.id, p.id, &a.name)?; let pool = &app.db; crate::db::proposals::add_event(pool, p.id, &by.username, &format!("asked {} to review", a.name), "").await.ok(); app.bus.live(&crate::events::topic::proposal(row.id, p.number), "assistant"); crate::db::accounts::audit( pool, Some(by), "assistant.asked", &format!("{}/{}", row.owner, row.name), &format!("{} #{}", a.name, p.number), ) .await .ok(); let (app, row, p, by) = (app.clone(), row.clone(), p.clone(), by.clone()); tokio::spawn(async move { let outcome = answer_review(&app, &row, &p, &a).await; app.assistant_gate.lock().unwrap_or_else(|e| e.into_inner()).release(p.id, &a.name); match outcome { Ok(body) => { if crate::db::proposals::add_review_via(&app.db, p.id, by.id, &by.username, "comment", &body, &a.name).await.is_ok() { crate::db::proposals::add_event(&app.db, p.id, &by.username, &format!("review.comment via {}", a.name), "").await.ok(); crate::web::proposals::reviewed(&app, &row, &p, &by, "comment").await; crate::hooks::emit( &app, &row, "review", &by.username, serde_json::json!({ "number": p.number, "verdict": "comment", "via": a.name }), ) .await; } } Err(why) => { crate::db::proposals::add_event(&app.db, p.id, &by.username, &format!("{} did not answer", a.name), &why).await.ok(); app.bus.live(&crate::events::topic::proposal(row.id, p.number), "assistant"); } } }); Ok(()) } async fn answer_review( app: &Shared, row: &crate::db::repos::RepoRow, p: &crate::db::proposals::ProposalRow, a: &Assistant, ) -> Result { let repo = crate::repos::names::RepoName::new(&row.owner, &row.name).ok_or("bad repository name")?; let loaded = crate::web::proposals::load(app, &repo, p.clone()).await.map_err(|e| format!("the proposal could not be loaded: {e}"))?; let budget = usize::try_from(a.budget).unwrap_or(DEFAULT_BUDGET as usize).max(1_000); let prompts = prompt_for_repo(&app.db, row).await.map_err(|e| e.to_string())?; let system = review_system_prompt(a, &app.config.site.name, &prompts); ask_within(app, a, &system, budget, BACKGROUND_TIMEOUT, |room| review_context(row, p, &loaded, room), |_| {}).await } /// Ask with the context cut to fit `budget` characters, system prompt /// included (oxygenhub#248). A provider that says the prompt is too long /// for its model is asked once more with a third of the room; whatever /// fails is said plainly, with what to do about it. async fn ask_within( app: &Shared, a: &Assistant, system: &str, budget: usize, wait: Duration, context: impl Fn(usize) -> String, mut part: impl FnMut(&str), ) -> Result { let room = |b: usize| b.saturating_sub(system.chars().count()); // A refusal for length comes before any part, so asking again starts // the answer afresh. match ask_streaming(app, a, system, &context(room(budget)), wait, &mut part).await { Err(e) if too_long(&e) => { let smaller = (budget / 3).max(1_000); tracing::info!("assistant {}: too long at {budget} characters ({e}); asking again at {smaller}", a.name); ask_streaming(app, a, system, &context(room(smaller)), wait, &mut part).await.map_err(|e2| { format!( "the proposal was too large for {}'s model at {budget} characters and again at {smaller} ({e2}): lower this assistant's context budget", a.name ) }) } Err(e) if e.starts_with("no answer within") => Err(format!( "{e}: a large proposal on a slow or local model can take longer; lower this assistant's context budget, or use a faster model" )), other => other, } } /// Whether a provider's refusal says the prompt is too long for its model. fn too_long(e: &str) -> bool { let l = e.to_lowercase(); (l.starts_with("400") || l.starts_with("413") || l.starts_with("422")) && [ "too long", "too large", "context length", "context window", "maximum context", "too many tokens", "exceeds", "reduce the length", "input length", ] .iter() .any(|k| l.contains(k)) } // -------------------------------------------------------- summaries /// The system prompt for a summary posted as a comment (oxygenhub#190). pub fn summary_system_prompt(a: &Assistant, site: &str, prompts: &str) -> String { let mut s = format!( "You are {name}, an AI assistant working on {site}, an OxygenHub site, with a next-generation version control system \ called Oxygen (its command is oxy). A proposal is Oxygen's equivalent of a merge request or pull request. \ Please summarise the proposal and comment on it. A person asked for this; it is posted as a comment made via you, \ under their name. Answer in Markdown: a short summary of what the proposal changes and why, then anything a reviewer \ should look at closely; keep to what the description and the diff show.", name = a.name ); if !prompts.trim().is_empty() { s.push_str("\n\n"); s.push_str(prompts.trim()); } s } /// Ask `a` to summarise proposal `p` for `by`: like [`request_review`], /// but the answer is a conversation comment via the assistant. pub async fn request_summary( app: &Shared, row: &crate::db::repos::RepoRow, p: &crate::db::proposals::ProposalRow, by: &crate::db::accounts::Account, a: Assistant, ) -> Result<(), String> { app.assistant_gate.lock().unwrap_or_else(|e| e.into_inner()).take(row.id, p.id, &a.name)?; let pool = &app.db; crate::db::proposals::add_event(pool, p.id, &by.username, &format!("asked {} for a summary", a.name), "").await.ok(); app.bus.live(&crate::events::topic::proposal(row.id, p.number), "assistant"); crate::db::accounts::audit( pool, Some(by), "assistant.asked", &format!("{}/{}", row.owner, row.name), &format!("{} #{} summary", a.name, p.number), ) .await .ok(); let (app, row, p, by) = (app.clone(), row.clone(), p.clone(), by.clone()); tokio::spawn(async move { let outcome = answer_summary(&app, &row, &p, &a).await; app.assistant_gate.lock().unwrap_or_else(|e| e.into_inner()).release(p.id, &a.name); match outcome { Ok(body) => { crate::db::proposals::add_comment_via(&app.db, p.id, by.id, &by.username, &body, &a.name).await.ok(); app.bus.live(&crate::events::topic::proposal(row.id, p.number), "comment"); } Err(why) => { crate::db::proposals::add_event(&app.db, p.id, &by.username, &format!("{} did not answer", a.name), &why).await.ok(); app.bus.live(&crate::events::topic::proposal(row.id, p.number), "assistant"); } } }); Ok(()) } async fn answer_summary( app: &Shared, row: &crate::db::repos::RepoRow, p: &crate::db::proposals::ProposalRow, a: &Assistant, ) -> Result { let repo = crate::repos::names::RepoName::new(&row.owner, &row.name).ok_or("bad repository name")?; let loaded = crate::web::proposals::load(app, &repo, p.clone()).await.map_err(|e| format!("the proposal could not be loaded: {e}"))?; let budget = usize::try_from(a.budget).unwrap_or(DEFAULT_BUDGET as usize).max(1_000); let prompts = prompt_for_repo(&app.db, row).await.map_err(|e| e.to_string())?; let system = summary_system_prompt(a, &app.config.site.name, &prompts); ask_within(app, a, &system, budget, BACKGROUND_TIMEOUT, |room| review_context(row, p, &loaded, room), |_| {}).await } // -------------------------------------------------------- drafting /// The system prompt for drafting a proposal's description. pub fn draft_system_prompt(a: &Assistant, site: &str, prompts: &str) -> String { let mut s = format!( "You are {name}, an AI assistant drafting the description of a proposal on {site}, an OxygenHub site, \ for the person who will open or edit it; they read and change your draft before saving, and the draft is theirs. \ Write Markdown for the description field only: no title line, no greeting, no sign-off. \ Say what the changes do and why, what a reviewer should look at, and how to verify them; \ keep to what the diff shows, and keep it short.", name = a.name ); if !prompts.trim().is_empty() { s.push_str("\n\n"); s.push_str(prompts.trim()); } s } /// What a description is drafted from. pub struct Draft<'a> { /// The proposal's id, or minus the repository's for one not yet opened: /// one draft per assistant at a time for it. pub slot: i64, pub title: &'a str, pub target: &'a str, pub changes: &'a [crate::repos::view::ChangeView], pub diffs: &'a [oxy_core::repo::FileDiff], } /// Ask `a` for a description, now; the answer goes into the editor, /// nothing is saved. `Err` is for the person. pub async fn draft_description( app: &Shared, row: &crate::db::repos::RepoRow, by: &crate::db::accounts::Account, a: &Assistant, d: Draft<'_>, ) -> Result { draft_description_streaming(app, row, by, a, d, |_| {}).await } /// [`draft_description`], handing each part of the draft to `part` as it /// comes, for a page that shows it being written (oxygenhub#248). pub async fn draft_description_streaming( app: &Shared, row: &crate::db::repos::RepoRow, by: &crate::db::accounts::Account, a: &Assistant, d: Draft<'_>, part: impl FnMut(&str), ) -> Result { let Draft { slot, title, target, changes, diffs } = d; app.assistant_gate.lock().unwrap_or_else(|e| e.into_inner()).take(row.id, slot, &a.name)?; crate::db::accounts::audit(&app.db, Some(by), "assistant.drafted", &format!("{}/{}", row.owner, row.name), &a.name).await.ok(); let budget = usize::try_from(a.budget).unwrap_or(DEFAULT_BUDGET as usize).max(1_000); let prompts = prompt_for_repo(&app.db, row).await.map_err(|e| e.to_string()); let outcome = match prompts { Ok(prompts) => { let system = draft_system_prompt(a, &app.config.site.name, &prompts); ask_within(app, a, &system, budget, CALL_TIMEOUT, |room| draft_context(row, title, target, changes, diffs, room), part).await } Err(e) => Err(e), }; app.assistant_gate.lock().unwrap_or_else(|e| e.into_inner()).release(slot, &a.name); outcome.map(|t| t.trim().to_string()) }