| 1 | |
|
| 2 | |
|
| 3 | |
|
| 4 | |
|
| 5 | |
|
| 6 | |
|
| 7 | |
|
| 8 | |
|
| 9 | |
|
| 10 | |
|
| 11 | |
|
| 12 | | use crate::db::{now_ms, Pool}; |
| 13 | | use crate::Shared; |
| 14 | | use rand::RngCore; |
| 15 | | use ring::aead; |
| 16 | | use sqlx::FromRow; |
| 17 | | use std::path::Path; |
| 18 | | use std::time::Duration; |
| 19 | |
|
| 20 | |
|
| 21 | | pub const APIS: [&str; 2] = ["chat-completions", "anthropic-messages"]; |
| 22 | |
|
| 23 | |
|
| 24 | | pub const DEFAULT_BUDGET: i64 = 120_000; |
| 25 | |
|
| 26 | |
|
| 27 | | pub const CALL_TIMEOUT: Duration = Duration::from_secs(90); |
| 28 | |
|
| 29 | |
|
| 30 | |
|
| 31 | | pub const BACKGROUND_TIMEOUT: Duration = Duration::from_secs(300); |
| 32 | |
|
| 33 | |
|
| 34 | | pub const MAX_ANSWER_TOKENS: u32 = 16_000; |
| 35 | |
|
| 36 | | const MAX_ANSWER: usize = 64 * 1024; |
| 37 | |
|
| 38 | |
|
| 39 | | const STREAM_LIMIT: Duration = Duration::from_secs(15 * 60); |
| 40 | |
|
| 41 | | pub use crate::levels::Scope; |
| 42 | |
|
| 43 | | #[derive(Clone, Debug, FromRow)] |
| 44 | | pub struct Assistant { |
| 45 | | pub id: i64, |
| 46 | | pub scope: String, |
| 47 | | pub scope_id: i64, |
| 48 | | pub name: String, |
| 49 | | pub endpoint: String, |
| 50 | | |
| 51 | | pub api: String, |
| 52 | | pub model: String, |
| 53 | | |
| 54 | | pub headers: String, |
| 55 | | |
| 56 | | pub token: Vec<u8>, |
| 57 | | pub enabled: bool, |
| 58 | | pub budget: i64, |
| 59 | | pub created_ms: i64, |
| 60 | | pub updated_ms: i64, |
| 61 | | } |
| 62 | |
|
| 63 | | impl Assistant { |
| 64 | 6× | pub fn scope(&self) -> Scope { |
| 65 | 6× | match self.scope.as_str() { |
| 66 | 6× | "org" => Scope::Org(self.scope_id), |
| 67 | 6× | "repo" => Scope::Repo(self.scope_id), |
| 68 | 6× | _ => Scope::Site, |
| 69 | | } |
| 70 | 6× | } |
| 71 | |
|
| 72 | 36× | pub fn has_token(&self) -> bool { |
| 73 | 36× | !self.token.is_empty() |
| 74 | 36× | } |
| 75 | |
|
| 76 | | |
| 77 | 16× | pub fn header_list(&self) -> Vec<(String, String)> { |
| 78 | 16× | parse_headers(&self.headers).unwrap_or_default() |
| 79 | 16× | } |
| 80 | | } |
| 81 | |
|
| 82 | |
|
| 83 | | #[derive(Clone, Debug, Default)] |
| 84 | | pub struct Fields { |
| 85 | | pub name: String, |
| 86 | | pub endpoint: String, |
| 87 | | pub api: String, |
| 88 | | pub model: String, |
| 89 | | pub headers: String, |
| 90 | | pub enabled: bool, |
| 91 | | pub budget: i64, |
| 92 | | } |
| 93 | |
|
| 94 | |
|
| 95 | 33× | pub fn parse_headers(text: &str) -> Result<Vec<(String, String)>, String> { |
| 96 | 33× | let mut out = vec![]; |
| 97 | 33× | for line in text.lines().map(str::trim).filter(|l| !l.is_empty()) { |
| 98 | 4× | let Some((name, value)) = line.split_once(':') else { return Err(line.to_string()) }; |
| 99 | 3× | let (name, value) = (name.trim(), value.trim()); |
| 100 | 20× | let name_ok = !name.is_empty() && name.chars().all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_'); |
| 101 | 6× | if !name_ok || value.chars().any(|c| c.is_control()) { |
| 102 | 1× | return Err(line.to_string()); |
| 103 | 2× | } |
| 104 | 2× | out.push((name.to_string(), value.to_string())); |
| 105 | | } |
| 106 | 31× | Ok(out) |
| 107 | 33× | } |
| 108 | |
|
| 109 | |
|
| 110 | |
|
| 111 | |
|
| 112 | | #[derive(Clone)] |
| 113 | | pub struct SecretKey([u8; 32]); |
| 114 | |
|
| 115 | | impl SecretKey { |
| 116 | 211× | pub fn load_or_create(keys: &Path) -> anyhow::Result<SecretKey> { |
| 117 | 211× | SecretKey::load_or_create_named(keys, "assistants.key") |
| 118 | 211× | } |
| 119 | |
|
| 120 | | |
| 121 | | |
| 122 | 419× | pub fn load_or_create_named(keys: &Path, file: &str) -> anyhow::Result<SecretKey> { |
| 123 | 419× | std::fs::create_dir_all(keys)?; |
| 124 | 419× | let path = keys.join(file); |
| 125 | 419× | if let Ok(bytes) = std::fs::read(&path) { |
| 126 | 1× | if bytes.len() == 32 { |
| 127 | 1× | return Ok(SecretKey(bytes.try_into().unwrap())); |
| 128 | 0× | } |
| 129 | 418× | } |
| 130 | 418× | let mut k = [0u8; 32]; |
| 131 | 418× | rand::rngs::OsRng.fill_bytes(&mut k); |
| 132 | 418× | std::fs::write(&path, k)?; |
| 133 | | #[cfg(unix)] |
| 134 | | { |
| 135 | | use std::os::unix::fs::PermissionsExt; |
| 136 | 418× | std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600))?; |
| 137 | | } |
| 138 | 418× | Ok(SecretKey(k)) |
| 139 | 419× | } |
| 140 | |
|
| 141 | 99× | fn key(&self) -> aead::LessSafeKey { |
| 142 | 99× | aead::LessSafeKey::new(aead::UnboundKey::new(&aead::CHACHA20_POLY1305, &self.0).expect("key length")) |
| 143 | 99× | } |
| 144 | |
|
| 145 | | |
| 146 | | |
| 147 | 23× | pub fn seal(&self, plain: &str) -> Vec<u8> { |
| 148 | 23× | if plain.is_empty() { |
| 149 | 2× | return vec![]; |
| 150 | 21× | } |
| 151 | 21× | let mut nonce = [0u8; 12]; |
| 152 | 21× | rand::rngs::OsRng.fill_bytes(&mut nonce); |
| 153 | 21× | let mut out = plain.as_bytes().to_vec(); |
| 154 | 21× | self.key().seal_in_place_append_tag(aead::Nonce::assume_unique_for_key(nonce), aead::Aad::empty(), &mut out).expect("seal"); |
| 155 | 21× | let mut sealed = nonce.to_vec(); |
| 156 | 21× | sealed.extend(out); |
| 157 | 21× | sealed |
| 158 | 23× | } |
| 159 | |
|
| 160 | | |
| 161 | | |
| 162 | 408× | pub fn open(&self, sealed: &[u8]) -> Option<String> { |
| 163 | 408× | if sealed.is_empty() { |
| 164 | 329× | return Some(String::new()); |
| 165 | 79× | } |
| 166 | 79× | let (nonce, rest) = sealed.split_at_checked(12)?; |
| 167 | 78× | let mut buf = rest.to_vec(); |
| 168 | 78× | let nonce = aead::Nonce::try_assume_unique_for_key(nonce).ok()?; |
| 169 | 78× | let plain = self.key().open_in_place(nonce, aead::Aad::empty(), &mut buf).ok()?; |
| 170 | 77× | String::from_utf8(plain.to_vec()).ok() |
| 171 | 408× | } |
| 172 | | } |
| 173 | |
|
| 174 | |
|
| 175 | |
|
| 176 | 646× | pub async fn list(pool: &Pool, scope: Scope) -> sqlx::Result<Vec<Assistant>> { |
| 177 | 646× | sqlx::query_as("SELECT * FROM assistants WHERE scope = $1 AND scope_id = $2 ORDER BY name") |
| 178 | 646× | .bind(scope.kind()) |
| 179 | 646× | .bind(scope.id()) |
| 180 | 646× | .fetch_all(pool) |
| 181 | 646× | .await |
| 182 | 646× | } |
| 183 | |
|
| 184 | 7× | pub async fn get(pool: &Pool, id: i64) -> sqlx::Result<Option<Assistant>> { |
| 185 | 7× | sqlx::query_as("SELECT * FROM assistants WHERE id = $1").bind(id).fetch_optional(pool).await |
| 186 | 7× | } |
| 187 | |
|
| 188 | |
|
| 189 | 10× | pub async fn create(pool: &Pool, scope: Scope, f: &Fields, token: &[u8]) -> sqlx::Result<Option<Assistant>> { |
| 190 | 10× | let now = now_ms(); |
| 191 | 10× | let r = sqlx::query_as( |
| 192 | 10× | "INSERT INTO assistants (scope, scope_id, name, endpoint, api, model, headers, token, enabled, budget, created_ms, updated_ms) \ |
| 193 | 10× | VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) ON CONFLICT DO NOTHING RETURNING *", |
| 194 | 10× | ) |
| 195 | 10× | .bind(scope.kind()) |
| 196 | 10× | .bind(scope.id()) |
| 197 | 10× | .bind(&f.name) |
| 198 | 10× | .bind(&f.endpoint) |
| 199 | 10× | .bind(&f.api) |
| 200 | 10× | .bind(&f.model) |
| 201 | 10× | .bind(&f.headers) |
| 202 | 10× | .bind(token) |
| 203 | 10× | .bind(f.enabled) |
| 204 | 10× | .bind(f.budget) |
| 205 | 10× | .bind(now) |
| 206 | 10× | .bind(now) |
| 207 | 10× | .fetch_optional(pool) |
| 208 | 10× | .await?; |
| 209 | 10× | Ok(r) |
| 210 | 10× | } |
| 211 | |
|
| 212 | |
|
| 213 | |
|
| 214 | 3× | pub async fn update(pool: &Pool, id: i64, f: &Fields, token: Option<&[u8]>) -> sqlx::Result<bool> { |
| 215 | 3× | let r = match token { |
| 216 | 0× | Some(t) => sqlx::query( |
| 217 | 0× | "UPDATE assistants SET name = $1, endpoint = $2, api = $3, model = $4, headers = $5, enabled = $6, budget = $7, token = $8, updated_ms = $9 \ |
| 218 | 0× | 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)", |
| 219 | 0× | ) |
| 220 | 0× | .bind(&f.name) |
| 221 | 0× | .bind(&f.endpoint) |
| 222 | 0× | .bind(&f.api) |
| 223 | 0× | .bind(&f.model) |
| 224 | 0× | .bind(&f.headers) |
| 225 | 0× | .bind(f.enabled) |
| 226 | 0× | .bind(f.budget) |
| 227 | 0× | .bind(t) |
| 228 | 0× | .bind(now_ms()) |
| 229 | 0× | .bind(id) |
| 230 | 0× | .bind(&f.name) |
| 231 | 0× | .bind(id) |
| 232 | 0× | .execute(pool) |
| 233 | 0× | .await?, |
| 234 | 3× | None => sqlx::query( |
| 235 | 3× | "UPDATE assistants SET name = $1, endpoint = $2, api = $3, model = $4, headers = $5, enabled = $6, budget = $7, updated_ms = $8 \ |
| 236 | 3× | 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)", |
| 237 | 3× | ) |
| 238 | 3× | .bind(&f.name) |
| 239 | 3× | .bind(&f.endpoint) |
| 240 | 3× | .bind(&f.api) |
| 241 | 3× | .bind(&f.model) |
| 242 | 3× | .bind(&f.headers) |
| 243 | 3× | .bind(f.enabled) |
| 244 | 3× | .bind(f.budget) |
| 245 | 3× | .bind(now_ms()) |
| 246 | 3× | .bind(id) |
| 247 | 3× | .bind(&f.name) |
| 248 | 3× | .bind(id) |
| 249 | 3× | .execute(pool) |
| 250 | 3× | .await?, |
| 251 | | }; |
| 252 | 3× | Ok(r.rows_affected() == 1) |
| 253 | 3× | } |
| 254 | |
|
| 255 | |
|
| 256 | |
|
| 257 | |
|
| 258 | 1× | pub async fn rename_in_history(pool: &Pool, scope: Scope, owner_hint: &str, old: &str, new: &str) -> sqlx::Result<()> { |
| 259 | 1× | if old == new || old.is_empty() || new.is_empty() { |
| 260 | 0× | return Ok(()); |
| 261 | 1× | } |
| 262 | 1× | let like = format!("%{old}%"); |
| 263 | 1× | match scope { |
| 264 | | Scope::Site => { |
| 265 | 2× | for table in ["proposal_reviews", "proposal_comments"] { |
| 266 | 2× | sqlx::query(&format!("UPDATE {table} SET via = $1 WHERE via = $2")).bind(new).bind(old).execute(pool).await?; |
| 267 | | } |
| 268 | 1× | sqlx::query("UPDATE proposal_events SET event = replace(event, $1, $2) WHERE event LIKE $3") |
| 269 | 1× | .bind(old) |
| 270 | 1× | .bind(new) |
| 271 | 1× | .bind(&like) |
| 272 | 1× | .execute(pool) |
| 273 | 1× | .await?; |
| 274 | | } |
| 275 | | Scope::Org(_) => { |
| 276 | 0× | 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)"; |
| 277 | 0× | for table in ["proposal_reviews", "proposal_comments"] { |
| 278 | 0× | sqlx::query(&format!("UPDATE {table} SET via = $1 WHERE via = $2 AND {in_org}")) |
| 279 | 0× | .bind(new) |
| 280 | 0× | .bind(old) |
| 281 | 0× | .bind(owner_hint) |
| 282 | 0× | .execute(pool) |
| 283 | 0× | .await?; |
| 284 | | } |
| 285 | 0× | sqlx::query(&format!("UPDATE proposal_events SET event = replace(event, $1, $2) WHERE event LIKE $3 AND {in_org}")) |
| 286 | 0× | .bind(old) |
| 287 | 0× | .bind(new) |
| 288 | 0× | .bind(&like) |
| 289 | 0× | .bind(owner_hint) |
| 290 | 0× | .execute(pool) |
| 291 | 0× | .await?; |
| 292 | | } |
| 293 | 0× | Scope::Repo(id) => { |
| 294 | 0× | let in_repo = "proposal_id IN (SELECT id FROM proposals WHERE repo_id = $1)"; |
| 295 | 0× | for table in ["proposal_reviews", "proposal_comments"] { |
| 296 | 0× | sqlx::query(&format!("UPDATE {table} SET via = $1 WHERE via = $2 AND {in_repo}")) |
| 297 | 0× | .bind(new) |
| 298 | 0× | .bind(old) |
| 299 | 0× | .bind(id) |
| 300 | 0× | .execute(pool) |
| 301 | 0× | .await?; |
| 302 | | } |
| 303 | 0× | sqlx::query(&format!("UPDATE proposal_events SET event = replace(event, $1, $2) WHERE event LIKE $3 AND {in_repo}")) |
| 304 | 0× | .bind(old) |
| 305 | 0× | .bind(new) |
| 306 | 0× | .bind(&like) |
| 307 | 0× | .bind(id) |
| 308 | 0× | .execute(pool) |
| 309 | 0× | .await?; |
| 310 | | } |
| 311 | | } |
| 312 | 1× | Ok(()) |
| 313 | 1× | } |
| 314 | |
|
| 315 | 2× | pub async fn set_enabled(pool: &Pool, id: i64, enabled: bool) -> sqlx::Result<()> { |
| 316 | 2× | sqlx::query("UPDATE assistants SET enabled = $1, updated_ms = $2 WHERE id = $3") |
| 317 | 2× | .bind(enabled) |
| 318 | 2× | .bind(now_ms()) |
| 319 | 2× | .bind(id) |
| 320 | 2× | .execute(pool) |
| 321 | 2× | .await?; |
| 322 | 2× | Ok(()) |
| 323 | 2× | } |
| 324 | |
|
| 325 | 1× | pub async fn delete(pool: &Pool, id: i64) -> sqlx::Result<()> { |
| 326 | 1× | sqlx::query("DELETE FROM assistants WHERE id = $1").bind(id).execute(pool).await?; |
| 327 | 1× | Ok(()) |
| 328 | 1× | } |
| 329 | |
|
| 330 | |
|
| 331 | |
|
| 332 | 1× | pub async fn delete_scope(pool: &Pool, scope: Scope) -> sqlx::Result<()> { |
| 333 | 1× | sqlx::query("DELETE FROM assistants WHERE scope = $1 AND scope_id = $2").bind(scope.kind()).bind(scope.id()).execute(pool).await?; |
| 334 | 1× | sqlx::query("DELETE FROM assistant_prompts WHERE scope = $1 AND scope_id = $2") |
| 335 | 1× | .bind(scope.kind()) |
| 336 | 1× | .bind(scope.id()) |
| 337 | 1× | .execute(pool) |
| 338 | 1× | .await?; |
| 339 | 1× | Ok(()) |
| 340 | 1× | } |
| 341 | |
|
| 342 | |
|
| 343 | 454× | pub async fn prompt(pool: &Pool, scope: Scope) -> sqlx::Result<String> { |
| 344 | 454× | Ok(sqlx::query_scalar("SELECT prompt FROM assistant_prompts WHERE scope = $1 AND scope_id = $2") |
| 345 | 454× | .bind(scope.kind()) |
| 346 | 454× | .bind(scope.id()) |
| 347 | 454× | .fetch_optional(pool) |
| 348 | 454× | .await? |
| 349 | 454× | .unwrap_or_default()) |
| 350 | 454× | } |
| 351 | |
|
| 352 | 5× | pub async fn set_prompt(pool: &Pool, scope: Scope, text: &str) -> sqlx::Result<()> { |
| 353 | 5× | if text.trim().is_empty() { |
| 354 | 0× | sqlx::query("DELETE FROM assistant_prompts WHERE scope = $1 AND scope_id = $2") |
| 355 | 0× | .bind(scope.kind()) |
| 356 | 0× | .bind(scope.id()) |
| 357 | 0× | .execute(pool) |
| 358 | 0× | .await?; |
| 359 | | } else { |
| 360 | 5× | sqlx::query( |
| 361 | 5× | "INSERT INTO assistant_prompts (scope, scope_id, prompt, updated_ms) VALUES ($1, $2, $3, $4) \ |
| 362 | 5× | ON CONFLICT (scope, scope_id) DO UPDATE SET prompt = excluded.prompt, updated_ms = excluded.updated_ms", |
| 363 | 5× | ) |
| 364 | 5× | .bind(scope.kind()) |
| 365 | 5× | .bind(scope.id()) |
| 366 | 5× | .bind(text.trim()) |
| 367 | 5× | .bind(now_ms()) |
| 368 | 5× | .execute(pool) |
| 369 | 5× | .await?; |
| 370 | | } |
| 371 | 5× | Ok(()) |
| 372 | 5× | } |
| 373 | |
|
| 374 | |
|
| 375 | |
|
| 376 | |
|
| 377 | |
|
| 378 | | #[derive(Clone, Debug)] |
| 379 | | pub struct Available { |
| 380 | | pub assistant: Assistant, |
| 381 | | |
| 382 | | pub from: String, |
| 383 | | |
| 384 | | pub off_above: bool, |
| 385 | | } |
| 386 | |
|
| 387 | | impl Available { |
| 388 | 59× | pub fn on(&self) -> bool { |
| 389 | 59× | self.assistant.enabled && !self.off_above |
| 390 | 59× | } |
| 391 | | } |
| 392 | |
|
| 393 | |
|
| 394 | |
|
| 395 | |
|
| 396 | 362× | pub async fn available(pool: &Pool, scope: Scope, owner: &str) -> sqlx::Result<Vec<Available>> { |
| 397 | 362× | let mut levels = crate::levels::above(pool, scope, owner).await?; |
| 398 | 362× | levels.push((scope, String::new())); |
| 399 | 362× | let mut out: Vec<Available> = vec![]; |
| 400 | 362× | let mut off: Vec<String> = vec![]; |
| 401 | 636× | for (level, from) in levels { |
| 402 | 636× | for a in list(pool, level).await? { |
| 403 | 69× | let off_above = off.iter().any(|n| n.eq_ignore_ascii_case(&a.name)); |
| 404 | 69× | if !a.enabled && !off.iter().any(|n| n.eq_ignore_ascii_case(&a.name)) { |
| 405 | 4× | off.push(a.name.clone()); |
| 406 | 65× | } |
| 407 | 69× | out.retain(|o| !o.assistant.name.eq_ignore_ascii_case(&a.name)); |
| 408 | 69× | out.push(Available { assistant: a, from: from.clone(), off_above }); |
| 409 | | } |
| 410 | | } |
| 411 | 362× | Ok(out) |
| 412 | 362× | } |
| 413 | |
|
| 414 | |
|
| 415 | 94× | pub async fn enabled_for_repo(pool: &Pool, repo: &crate::db::repos::RepoRow) -> sqlx::Result<Vec<Assistant>> { |
| 416 | 94× | Ok(available(pool, Scope::Repo(repo.id), &repo.owner).await?.into_iter().filter(Available::on).map(|a| a.assistant).collect()) |
| 417 | 94× | } |
| 418 | |
|
| 419 | |
|
| 420 | |
|
| 421 | 271× | pub async fn inherited_prompts(pool: &Pool, scope: Scope, owner: &str) -> sqlx::Result<Vec<(String, String)>> { |
| 422 | 271× | let mut out = vec![]; |
| 423 | 271× | for (level, from) in crate::levels::above(pool, scope, owner).await? { |
| 424 | 183× | let text = prompt(pool, level).await?; |
| 425 | 183× | if !text.is_empty() { |
| 426 | 17× | out.push((from, text)); |
| 427 | 166× | } |
| 428 | | } |
| 429 | 271× | Ok(out) |
| 430 | 271× | } |
| 431 | |
|
| 432 | |
|
| 433 | |
|
| 434 | 7× | pub async fn prompt_for_repo(pool: &Pool, repo: &crate::db::repos::RepoRow) -> sqlx::Result<String> { |
| 435 | 7× | let scope = Scope::Repo(repo.id); |
| 436 | 7× | let mut parts = inherited_prompts(pool, scope, &repo.owner).await?; |
| 437 | 7× | let own = prompt(pool, scope).await?; |
| 438 | 7× | if !own.is_empty() { |
| 439 | 0× | parts.push(("this repository".to_string(), own)); |
| 440 | 7× | } |
| 441 | 7× | Ok(parts.into_iter().map(|(from, text)| format!("## Instructions from {from}\n\n{text}\n")).collect::<Vec<_>>().join("\n")) |
| 442 | 7× | } |
| 443 | |
|
| 444 | |
|
| 445 | |
|
| 446 | |
|
| 447 | |
|
| 448 | |
|
| 449 | |
|
| 450 | |
|
| 451 | |
|
| 452 | |
|
| 453 | 9× | pub async fn ask(app: &Shared, a: &Assistant, system: &str, user: &str, wait: Duration) -> Result<String, String> { |
| 454 | 9× | ask_streaming(app, a, system, user, wait, |_| {}).await |
| 455 | 9× | } |
| 456 | |
|
| 457 | |
|
| 458 | 16× | pub async fn ask_streaming( |
| 459 | 16× | app: &Shared, |
| 460 | 16× | a: &Assistant, |
| 461 | 16× | system: &str, |
| 462 | 16× | user: &str, |
| 463 | 16× | wait: Duration, |
| 464 | 16× | mut part: impl FnMut(&str), |
| 465 | 16× | ) -> Result<String, String> { |
| 466 | 16× | let Some(token) = app.assistant_key.open(&a.token) else { |
| 467 | 0× | return Err("its token cannot be read: set it again".into()); |
| 468 | | }; |
| 469 | 16× | let client = crate::net::client(&a.endpoint, &app.config.outbound.allow, STREAM_LIMIT.max(wait)).await?; |
| 470 | 16× | let mut req = client.post(&a.endpoint).header("content-type", "application/json"); |
| 471 | 16× | let body = match a.api.as_str() { |
| 472 | 16× | "anthropic-messages" => { |
| 473 | 0× | if !token.is_empty() { |
| 474 | 0× | req = req.header("x-api-key", &token); |
| 475 | 0× | } |
| 476 | 0× | req = req.header("anthropic-version", "2023-06-01"); |
| 477 | 0× | serde_json::json!({ |
| 478 | 0× | "model": a.model, |
| 479 | 0× | "max_tokens": MAX_ANSWER_TOKENS, |
| 480 | 0× | "system": system, |
| 481 | 0× | "messages": [{ "role": "user", "content": user }], |
| 482 | 0× | "stream": true, |
| 483 | | }) |
| 484 | | } |
| 485 | | _ => { |
| 486 | 16× | if !token.is_empty() { |
| 487 | 14× | req = req.header("authorization", format!("Bearer {token}")); |
| 488 | 14× | } |
| 489 | | |
| 490 | | |
| 491 | | |
| 492 | 16× | serde_json::json!({ |
| 493 | 16× | "model": a.model, |
| 494 | 16× | "messages": [ |
| 495 | 16× | { "role": "system", "content": system }, |
| 496 | 16× | { "role": "user", "content": user }, |
| 497 | | ], |
| 498 | 16× | "stream": true, |
| 499 | | }) |
| 500 | | } |
| 501 | | }; |
| 502 | 16× | for (name, value) in a.header_list() { |
| 503 | 0× | req = req.header(name, value); |
| 504 | 0× | } |
| 505 | 16× | let quiet = || format!("no answer within {} seconds", wait.as_secs()); |
| 506 | 16× | let r = tokio::time::timeout(wait, req.json(&body).send()).await.map_err(|_| quiet())?.map_err(|e| { |
| 507 | 0× | if e.is_timeout() { |
| 508 | 0× | quiet() |
| 509 | | } else { |
| 510 | 0× | scrub(&e.to_string(), &token) |
| 511 | | } |
| 512 | 0× | })?; |
| 513 | 16× | let status = r.status(); |
| 514 | 16× | let streamed = |
| 515 | 16× | r.headers().get(reqwest::header::CONTENT_TYPE).and_then(|v| v.to_str().ok()).is_some_and(|t| t.starts_with("text/event-stream")); |
| 516 | 16× | if status.is_success() && streamed { |
| 517 | 0× | let (answer, stop) = read_stream(&a.api, r, wait, &mut part).await.map_err(|e| scrub(&e, &token))?; |
| 518 | 0× | return finish(answer, stop); |
| 519 | 16× | } |
| 520 | 16× | let text: String = |
| 521 | 16× | tokio::time::timeout(wait, r.text()).await.map_err(|_| quiet())?.unwrap_or_default().chars().take(MAX_ANSWER).collect(); |
| 522 | 16× | if !status.is_success() { |
| 523 | 1× | let why = serde_json::from_str::<serde_json::Value>(&text) |
| 524 | 1× | .ok() |
| 525 | 1× | .and_then(|v| v["error"]["message"].as_str().or(v["error"].as_str()).or(v["message"].as_str()).map(str::to_string)) |
| 526 | 1× | .unwrap_or_else(|| text.lines().next().unwrap_or("").chars().take(160).collect()); |
| 527 | 1× | return Err(scrub(&format!("{} {}", status.as_u16(), why.trim()), &token)); |
| 528 | 15× | } |
| 529 | 15× | let v: serde_json::Value = serde_json::from_str(&text).map_err(|_| "the answer was not JSON".to_string())?; |
| 530 | 15× | let (answer, stop) = answer_of(&a.api, &v); |
| 531 | 15× | if !answer.is_empty() { |
| 532 | 15× | part(&answer); |
| 533 | 15× | } |
| 534 | 15× | finish(answer, stop) |
| 535 | 16× | } |
| 536 | |
|
| 537 | |
|
| 538 | |
|
| 539 | 2× | async fn read_stream(api: &str, mut r: reqwest::Response, quiet: Duration, part: &mut impl FnMut(&str)) -> Result<(String, Stop), String> { |
| 540 | 2× | let mut events = Events::new(api); |
| 541 | | loop { |
| 542 | 10× | let chunk = match tokio::time::timeout(quiet, r.chunk()).await { |
| 543 | 1× | Err(_) if events.text.is_empty() => return Err(format!("no answer within {} seconds", quiet.as_secs())), |
| 544 | | |
| 545 | 1× | Err(_) => return Ok((events.text, Stop::Quiet)), |
| 546 | 0× | Ok(Err(e)) if events.text.is_empty() => return Err(e.to_string()), |
| 547 | 0× | Ok(Err(_)) => return Ok((events.text, Stop::Quiet)), |
| 548 | 0× | Ok(Ok(None)) => return Ok((events.text, events.stop.unwrap_or(Stop::Done))), |
| 549 | 9× | Ok(Ok(Some(c))) => c, |
| 550 | | }; |
| 551 | 9× | let before = events.text.len(); |
| 552 | 9× | events.feed(&chunk)?; |
| 553 | 9× | if events.text.len() > before { |
| 554 | 8× | part(&events.text[before..]); |
| 555 | 8× | } |
| 556 | 9× | if events.done || events.text.len() >= MAX_ANSWER { |
| 557 | 1× | return Ok((events.text, events.stop.unwrap_or(Stop::Done))); |
| 558 | 8× | } |
| 559 | | } |
| 560 | 2× | } |
| 561 | |
|
| 562 | |
|
| 563 | |
|
| 564 | |
|
| 565 | | struct Events { |
| 566 | | anthropic: bool, |
| 567 | | buf: Vec<u8>, |
| 568 | | text: String, |
| 569 | | stop: Option<Stop>, |
| 570 | | done: bool, |
| 571 | | } |
| 572 | |
|
| 573 | | impl Events { |
| 574 | 10× | fn new(api: &str) -> Events { |
| 575 | 10× | Events { anthropic: api == "anthropic-messages", buf: Vec::new(), text: String::new(), stop: None, done: false } |
| 576 | 10× | } |
| 577 | |
|
| 578 | 676× | fn feed(&mut self, bytes: &[u8]) -> Result<(), String> { |
| 579 | 676× | self.buf.extend_from_slice(bytes); |
| 580 | 26322× | while let Some(end) = self.buf.iter().position(|b| *b == b'\n') { |
| 581 | 93× | let line: Vec<u8> = self.buf.drain(..=end).collect(); |
| 582 | 93× | let line = String::from_utf8_lossy(&line); |
| 583 | 93× | let Some(data) = line.trim_end_matches(['\r', '\n']).strip_prefix("data:") else { |
| 584 | 48× | continue; |
| 585 | | }; |
| 586 | 45× | let data = data.trim_start(); |
| 587 | 45× | if data == "[DONE]" { |
| 588 | 1× | self.done = true; |
| 589 | 1× | continue; |
| 590 | 44× | } |
| 591 | 44× | let Ok(v) = serde_json::from_str::<serde_json::Value>(data) else { |
| 592 | 0× | continue; |
| 593 | | }; |
| 594 | 44× | self.event(&v)?; |
| 595 | | } |
| 596 | 675× | Ok(()) |
| 597 | 676× | } |
| 598 | |
|
| 599 | 44× | fn event(&mut self, v: &serde_json::Value) -> Result<(), String> { |
| 600 | 44× | if self.anthropic { |
| 601 | 41× | match v["type"].as_str() { |
| 602 | 41× | Some("content_block_delta") if v["delta"]["type"] == "text_delta" => { |
| 603 | 18× | self.text.push_str(v["delta"]["text"].as_str().unwrap_or("")); |
| 604 | 18× | } |
| 605 | 18× | Some("message_delta") => { |
| 606 | 6× | self.stop = match v["delta"]["stop_reason"].as_str() { |
| 607 | 6× | Some("max_tokens") => Some(Stop::Length), |
| 608 | 5× | Some("refusal") => Some(Stop::Declined(v["delta"]["stop_details"]["category"].as_str().map(str::to_string))), |
| 609 | 5× | Some(_) => Some(Stop::Done), |
| 610 | 0× | None => self.stop.take(), |
| 611 | | }; |
| 612 | | } |
| 613 | 12× | Some("message_stop") => self.done = true, |
| 614 | 6× | Some("error") => { |
| 615 | 1× | return Err(v["error"]["message"].as_str().unwrap_or("the provider reported an error").to_string()); |
| 616 | | } |
| 617 | 10× | _ => {} |
| 618 | | } |
| 619 | | } else { |
| 620 | 3× | if let Some(e) = v["error"]["message"].as_str() { |
| 621 | 0× | return Err(e.to_string()); |
| 622 | 3× | } |
| 623 | 3× | let choice = &v["choices"][0]; |
| 624 | 3× | if let Some(t) = choice["delta"]["content"].as_str() { |
| 625 | 2× | self.text.push_str(t); |
| 626 | 2× | } |
| 627 | 3× | match choice["finish_reason"].as_str() { |
| 628 | 1× | Some("length") => self.stop = Some(Stop::Length), |
| 629 | 0× | Some("content_filter") => self.stop = Some(Stop::Declined(None)), |
| 630 | 0× | Some(_) => self.stop = Some(Stop::Done), |
| 631 | 2× | None => {} |
| 632 | | } |
| 633 | | } |
| 634 | 43× | Ok(()) |
| 635 | 44× | } |
| 636 | | } |
| 637 | |
|
| 638 | |
|
| 639 | | #[derive(Debug, PartialEq, Eq)] |
| 640 | | enum Stop { |
| 641 | | |
| 642 | | Done, |
| 643 | | |
| 644 | | Length, |
| 645 | | |
| 646 | | Declined(Option<String>), |
| 647 | | |
| 648 | | Quiet, |
| 649 | | } |
| 650 | |
|
| 651 | |
|
| 652 | |
|
| 653 | 22× | fn answer_of(api: &str, v: &serde_json::Value) -> (String, Stop) { |
| 654 | 22× | match api { |
| 655 | 22× | "anthropic-messages" => { |
| 656 | 4× | let text = v["content"] |
| 657 | 4× | .as_array() |
| 658 | 4× | .map(|parts| { |
| 659 | 4× | parts |
| 660 | 4× | .iter() |
| 661 | 4× | .filter(|p| p["type"].as_str().is_none_or(|t| t == "text")) |
| 662 | 4× | .filter_map(|p| p["text"].as_str()) |
| 663 | 4× | .collect::<Vec<_>>() |
| 664 | 4× | .join("") |
| 665 | 4× | }) |
| 666 | 4× | .unwrap_or_default(); |
| 667 | 4× | let stop = match v["stop_reason"].as_str() { |
| 668 | 4× | Some("max_tokens") => Stop::Length, |
| 669 | 2× | Some("refusal") => Stop::Declined(v["stop_details"]["category"].as_str().map(str::to_string)), |
| 670 | 1× | _ => Stop::Done, |
| 671 | | }; |
| 672 | 4× | (text, stop) |
| 673 | | } |
| 674 | | _ => { |
| 675 | 18× | let choice = &v["choices"][0]; |
| 676 | 18× | let text = match &choice["message"]["content"] { |
| 677 | 18× | serde_json::Value::String(s) => s.clone(), |
| 678 | 0× | serde_json::Value::Array(parts) => parts.iter().filter_map(|p| p["text"].as_str()).collect::<Vec<_>>().join(""), |
| 679 | 0× | _ => String::new(), |
| 680 | | }; |
| 681 | 18× | let stop = match choice["finish_reason"].as_str() { |
| 682 | 3× | Some("length") => Stop::Length, |
| 683 | 2× | Some("content_filter") => Stop::Declined(None), |
| 684 | 16× | _ => Stop::Done, |
| 685 | | }; |
| 686 | 18× | (text, stop) |
| 687 | | } |
| 688 | | } |
| 689 | 22× | } |
| 690 | |
|
| 691 | |
|
| 692 | |
|
| 693 | 23× | fn finish(answer: String, stop: Stop) -> Result<String, String> { |
| 694 | 17× | match stop { |
| 695 | 2× | Stop::Declined(category) => Err(match category { |
| 696 | 1× | Some(c) => format!("it declined to answer ({c})"), |
| 697 | 1× | None => "it declined to answer".into(), |
| 698 | | }), |
| 699 | 3× | Stop::Length if answer.trim().is_empty() => { |
| 700 | 2× | Err("it ran out of room before writing an answer (it spent its length limit thinking): try a smaller proposal".into()) |
| 701 | | } |
| 702 | 1× | Stop::Length => Ok(format!("{}\n\n*(Cut short: the answer reached its length limit.)*", answer.trim_end())), |
| 703 | 1× | Stop::Quiet if answer.trim().is_empty() => Err("it went quiet before writing an answer".into()), |
| 704 | 1× | Stop::Quiet => Ok(format!("{}\n\n*(Cut short: the assistant went quiet part of the way through.)*", answer.trim_end())), |
| 705 | 17× | Stop::Done if answer.trim().is_empty() => Err("the answer carried no text".into()), |
| 706 | 17× | Stop::Done => Ok(answer), |
| 707 | | } |
| 708 | 23× | } |
| 709 | |
|
| 710 | |
|
| 711 | |
|
| 712 | | pub const PRESETS: &[(&str, &str, &str, &str, &str)] = &[ |
| 713 | | ("anthropic", "Anthropic (Claude)", "Claude", "https://api.anthropic.com/v1/messages", "anthropic-messages"), |
| 714 | | ("openai", "OpenAI (ChatGPT)", "ChatGPT", "https://api.openai.com/v1/chat/completions", "chat-completions"), |
| 715 | | ("gemini", "Google Gemini", "Gemini", "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions", "chat-completions"), |
| 716 | | ("deepseek", "DeepSeek", "DeepSeek", "https://api.deepseek.com/chat/completions", "chat-completions"), |
| 717 | | ("mistral", "Mistral", "Mistral", "https://api.mistral.ai/v1/chat/completions", "chat-completions"), |
| 718 | | ("perplexity", "Perplexity", "Perplexity", "https://api.perplexity.ai/chat/completions", "chat-completions"), |
| 719 | | ("ollama", "Ollama (on your network)", "Ollama", "http://localhost:11434/v1/chat/completions", "chat-completions"), |
| 720 | | ]; |
| 721 | |
|
| 722 | |
|
| 723 | |
|
| 724 | 16× | pub fn models_url(endpoint: &str) -> Option<String> { |
| 725 | 16× | let e = endpoint.trim_end_matches('/'); |
| 726 | 16× | e.strip_suffix("/chat/completions").or_else(|| e.strip_suffix("/messages")).map(|base| format!("{base}/models")) |
| 727 | 16× | } |
| 728 | |
|
| 729 | |
|
| 730 | |
|
| 731 | |
|
| 732 | 16× | pub async fn list_models(endpoint: &str, api: &str, token: &str, allow: &[String]) -> Result<Vec<String>, String> { |
| 733 | 16× | let url = models_url(endpoint).ok_or("this endpoint has no model list next to it")?; |
| 734 | 16× | let client = crate::net::client(&url, allow, std::time::Duration::from_secs(15)).await?; |
| 735 | 16× | let mut req = client.get(&url); |
| 736 | 16× | if api == "anthropic-messages" { |
| 737 | 0× | req = req.header("anthropic-version", "2023-06-01"); |
| 738 | 0× | if !token.is_empty() { |
| 739 | 0× | req = req.header("x-api-key", token); |
| 740 | 0× | } |
| 741 | 16× | } else if !token.is_empty() { |
| 742 | 15× | req = req.header("authorization", format!("Bearer {token}")); |
| 743 | 15× | } |
| 744 | 16× | let r = req.send().await.map_err(|e| scrub(&e.to_string(), token))?; |
| 745 | 16× | let status = r.status(); |
| 746 | 16× | if !status.is_success() { |
| 747 | 12× | return Err(format!("the model list answered {}", status.as_u16())); |
| 748 | 4× | } |
| 749 | 4× | let v: serde_json::Value = r.json().await.map_err(|_| "the model list was not JSON".to_string())?; |
| 750 | 4× | let skip = ["embed", "tts", "whisper", "dall-e", "moderation", "audio", "image", "realtime", "transcribe", "search-preview"]; |
| 751 | 4× | let mut out: Vec<String> = v["data"] |
| 752 | 4× | .as_array() |
| 753 | 4× | .or(v["models"].as_array()) |
| 754 | 4× | .ok_or("the model list had no data")? |
| 755 | 4× | .iter() |
| 756 | 12× | .filter_map(|m| m["id"].as_str().or(m["name"].as_str())) |
| 757 | 12× | .map(|id| id.strip_prefix("models/").unwrap_or(id).to_string()) |
| 758 | 84× | .filter(|id| !skip.iter().any(|s| id.to_lowercase().contains(s))) |
| 759 | 4× | .collect(); |
| 760 | 4× | out.sort(); |
| 761 | 4× | out.dedup(); |
| 762 | 4× | if out.is_empty() { |
| 763 | 0× | return Err("the model list was empty".into()); |
| 764 | 4× | } |
| 765 | 4× | Ok(out) |
| 766 | 16× | } |
| 767 | |
|
| 768 | |
|
| 769 | 7× | pub async fn test(app: &Shared, a: &Assistant) -> Result<String, String> { |
| 770 | 7× | let answer = |
| 771 | 7× | ask(app, a, "You are being connected to an OxygenHub site. Answer briefly.", "Reply with the single word OK.", CALL_TIMEOUT) |
| 772 | 7× | .await?; |
| 773 | 7× | Ok(answer.trim().chars().take(80).collect()) |
| 774 | 7× | } |
| 775 | |
|
| 776 | |
|
| 777 | 1× | fn scrub(text: &str, token: &str) -> String { |
| 778 | 1× | if token.is_empty() { |
| 779 | 0× | text.to_string() |
| 780 | | } else { |
| 781 | 1× | text.replace(token, "***") |
| 782 | | } |
| 783 | 1× | } |
| 784 | |
|
| 785 | | #[cfg(test)] |
| 786 | | mod tests { |
| 787 | | use super::*; |
| 788 | |
|
| 789 | | |
| 790 | | |
| 791 | | #[test] |
| 792 | 1× | fn a_draft_s_title_line_is_split_off() { |
| 793 | 1× | assert_eq!( |
| 794 | 1× | split_title("Title: feat(flow): hand jobs back\n\n**Summary.** Runners leave."), |
| 795 | 1× | (Some("feat(flow): hand jobs back".to_string()), "**Summary.** Runners leave.".to_string()) |
| 796 | | ); |
| 797 | 1× | assert_eq!(split_title(" title: \"Quoted\"\nBody"), (Some("Quoted".to_string()), "Body".to_string())); |
| 798 | 1× | assert_eq!(split_title("**Summary.** No title here."), (None, "**Summary.** No title here.".to_string())); |
| 799 | 1× | assert_eq!(split_title("Title:\nBody"), (None, "Body".to_string()), "an empty title is none"); |
| 800 | 1× | } |
| 801 | |
|
| 802 | | |
| 803 | | |
| 804 | | #[test] |
| 805 | 1× | fn anthropic_events_give_the_text_whole() { |
| 806 | 1× | let sse = "event: message_start\ndata: {\"type\":\"message_start\"}\n\n\ |
| 807 | 1× | data: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"thinking_delta\",\"thinking\":\"hmm\"}}\n\n\ |
| 808 | 1× | data: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"Adds a \"}}\n\n\ |
| 809 | 1× | data: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"line — café.\"}}\n\n\ |
| 810 | 1× | data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"}}\n\n\ |
| 811 | 1× | data: {\"type\":\"message_stop\"}\n\n"; |
| 812 | 5× | for size in [1, 3, 7, 64, sse.len()] { |
| 813 | 5× | let mut e = Events::new("anthropic-messages"); |
| 814 | 626× | for chunk in sse.as_bytes().chunks(size) { |
| 815 | 626× | e.feed(chunk).unwrap(); |
| 816 | 626× | } |
| 817 | 5× | assert_eq!((e.text.as_str(), e.stop.as_ref(), e.done), ("Adds a line — café.", Some(&Stop::Done), true), "chunks of {size}"); |
| 818 | | } |
| 819 | 1× | let mut cut = Events::new("anthropic-messages"); |
| 820 | 1× | cut.feed(b"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"max_tokens\"}}\n").unwrap(); |
| 821 | 1× | assert_eq!(cut.stop, Some(Stop::Length)); |
| 822 | 1× | let mut bad = Events::new("anthropic-messages"); |
| 823 | 1× | let e = bad.feed(b"data: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"Overloaded\"}}\n").unwrap_err(); |
| 824 | 1× | assert_eq!(e, "Overloaded"); |
| 825 | 1× | } |
| 826 | |
|
| 827 | | #[test] |
| 828 | 1× | fn chat_completion_events_give_the_text() { |
| 829 | 1× | let sse = "data: {\"choices\":[{\"delta\":{\"role\":\"assistant\"}}]}\n\n\ |
| 830 | 1× | data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n\ |
| 831 | 1× | data: {\"choices\":[{\"delta\":{\"content\":\", world\"},\"finish_reason\":\"length\"}]}\n\n\ |
| 832 | 1× | data: [DONE]\n\n"; |
| 833 | 1× | let mut e = Events::new("chat-completions"); |
| 834 | 39× | for chunk in sse.as_bytes().chunks(5) { |
| 835 | 39× | e.feed(chunk).unwrap(); |
| 836 | 39× | } |
| 837 | 1× | assert_eq!((e.text.as_str(), e.stop.as_ref(), e.done), ("Hello, world", Some(&Stop::Length), true)); |
| 838 | 1× | } |
| 839 | |
|
| 840 | | |
| 841 | | |
| 842 | | |
| 843 | | #[tokio::test] |
| 844 | 1× | async fn a_slow_steady_stream_is_read_and_a_silent_one_cut() { |
| 845 | 2× | async fn serve(parts: usize, gap_ms: u64, then_silent: bool) -> String { |
| 846 | 2× | let router = axum::Router::new().route( |
| 847 | 2× | "/v1/messages", |
| 848 | 2× | axum::routing::post(move || async move { |
| 849 | 2× | let (tx, rx) = tokio::sync::mpsc::channel::<Result<String, std::io::Error>>(4); |
| 850 | 2× | tokio::spawn(async move { |
| 851 | 8× | for i in 0..parts { |
| 852 | 8× | tokio::time::sleep(Duration::from_millis(gap_ms)).await; |
| 853 | 8× | let ev = format!( |
| 854 | | "data: {{\"type\":\"content_block_delta\",\"delta\":{{\"type\":\"text_delta\",\"text\":\"{i} \"}}}}\n\n" |
| 855 | | ); |
| 856 | 8× | if tx.send(Ok(ev)).await.is_err() { |
| 857 | 0× | return; |
| 858 | 8× | } |
| 859 | | } |
| 860 | 2× | if then_silent { |
| 861 | 1× | tokio::time::sleep(Duration::from_secs(5)).await; |
| 862 | 1× | } |
| 863 | 1× | let _ = tx.send(Ok("data: {\"type\":\"message_stop\"}\n\n".into())).await; |
| 864 | 1× | }); |
| 865 | 2× | let body = axum::body::Body::from_stream(tokio_stream::wrappers::ReceiverStream::new(rx)); |
| 866 | 2× | ([(axum::http::header::CONTENT_TYPE, "text/event-stream")], body) |
| 867 | 4× | }), |
| 868 | | ); |
| 869 | 2× | let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); |
| 870 | 2× | let port = listener.local_addr().unwrap().port(); |
| 871 | 2× | tokio::spawn(async move { axum::serve(listener, router).await.unwrap() }); |
| 872 | 2× | format!("http://127.0.0.1:{port}/v1/messages") |
| 873 | 2× | } |
| 874 | 1× | let quiet = Duration::from_millis(400); |
| 875 | | |
| 876 | 1× | let url = serve(6, 200, false).await; |
| 877 | 1× | let r = reqwest::Client::new().post(&url).send().await.unwrap(); |
| 878 | 1× | let mut parts = Vec::new(); |
| 879 | 6× | let (text, stop) = read_stream("anthropic-messages", r, quiet, &mut |p: &str| parts.push(p.to_string())).await.unwrap(); |
| 880 | 1× | assert_eq!((text.as_str(), stop), ("0 1 2 3 4 5 ", Stop::Done)); |
| 881 | 1× | assert_eq!(parts.len(), 6, "each part handed out as it came"); |
| 882 | | |
| 883 | 1× | let url = serve(2, 50, true).await; |
| 884 | 1× | let r = reqwest::Client::new().post(&url).send().await.unwrap(); |
| 885 | 2× | let (text, stop) = read_stream("anthropic-messages", r, quiet, &mut |_: &str| {}).await.unwrap(); |
| 886 | 1× | assert_eq!((text.as_str(), stop), ("0 1 ", Stop::Quiet)); |
| 887 | 1× | assert!(finish(text, Stop::Quiet).unwrap().contains("went quiet")); |
| 888 | 1× | } |
| 889 | |
|
| 890 | | #[test] |
| 891 | 1× | fn tokens_seal_and_open_only_under_their_key() { |
| 892 | 1× | let dir = tempfile::tempdir().unwrap(); |
| 893 | 1× | let k = SecretKey::load_or_create(dir.path()).unwrap(); |
| 894 | 1× | let sealed = k.seal("sk-very-secret"); |
| 895 | 1× | assert_ne!(sealed, b"sk-very-secret"); |
| 896 | 1× | assert!(!String::from_utf8_lossy(&sealed).contains("secret")); |
| 897 | 1× | assert_eq!(k.open(&sealed).as_deref(), Some("sk-very-secret")); |
| 898 | 1× | assert_eq!(SecretKey::load_or_create(dir.path()).unwrap().open(&sealed).as_deref(), Some("sk-very-secret"), "the key is kept"); |
| 899 | 1× | let other = SecretKey::load_or_create(tempfile::tempdir().unwrap().path()).unwrap(); |
| 900 | 1× | assert_eq!(other.open(&sealed), None); |
| 901 | 1× | assert_eq!(k.open(&[1, 2, 3]), None); |
| 902 | 1× | assert_eq!(k.seal(""), Vec::<u8>::new()); |
| 903 | 1× | assert_eq!(k.open(&[]).as_deref(), Some("")); |
| 904 | 1× | } |
| 905 | |
|
| 906 | | #[test] |
| 907 | 1× | fn answers_say_why_they_stopped() { |
| 908 | 4× | let a = |v: serde_json::Value| { |
| 909 | 4× | let (text, stop) = answer_of("anthropic-messages", &v); |
| 910 | 4× | finish(text, stop) |
| 911 | 4× | }; |
| 912 | | |
| 913 | 1× | let thought = serde_json::json!({ "stop_reason": "max_tokens", "content": [{ "type": "thinking", "thinking": "" }] }); |
| 914 | 1× | assert!(a(thought).unwrap_err().contains("ran out of room")); |
| 915 | 1× | let cut = serde_json::json!({ "stop_reason": "max_tokens", "content": [{ "type": "thinking", "thinking": "" }, { "type": "text", "text": "Half" }] }); |
| 916 | 1× | assert!(a(cut).unwrap().starts_with("Half\n\n*(Cut short")); |
| 917 | 1× | let declined = serde_json::json!({ "stop_reason": "refusal", "stop_details": { "category": "cyber" }, "content": [] }); |
| 918 | 1× | assert_eq!(a(declined).unwrap_err(), "it declined to answer (cyber)"); |
| 919 | 1× | let ok = serde_json::json!({ "stop_reason": "end_turn", "content": [{ "type": "text", "text": "Fine." }] }); |
| 920 | 1× | assert_eq!(a(ok).unwrap(), "Fine."); |
| 921 | 3× | let chat = |v: serde_json::Value| { |
| 922 | 3× | let (text, stop) = answer_of("chat-completions", &v); |
| 923 | 3× | finish(text, stop) |
| 924 | 3× | }; |
| 925 | 1× | assert!(chat(serde_json::json!({ "choices": [{ "finish_reason": "length", "message": { "content": "" } }] })) |
| 926 | 1× | .unwrap_err() |
| 927 | 1× | .contains("ran out of room")); |
| 928 | 1× | assert_eq!( |
| 929 | 1× | chat(serde_json::json!({ "choices": [{ "finish_reason": "content_filter", "message": { "content": "" } }] })).unwrap_err(), |
| 930 | | "it declined to answer" |
| 931 | | ); |
| 932 | 1× | assert_eq!( |
| 933 | 1× | chat(serde_json::json!({ "choices": [{ "finish_reason": "stop", "message": { "content": "Yes." } }] })).unwrap(), |
| 934 | | "Yes." |
| 935 | | ); |
| 936 | 1× | } |
| 937 | |
|
| 938 | | #[test] |
| 939 | 1× | fn a_refusal_for_length_is_told_from_others() { |
| 940 | 4× | for e in [ |
| 941 | 1× | "400 This model's maximum context length is 8192 tokens. Please reduce the length of the messages.", |
| 942 | 1× | "400 prompt is too long: 210000 tokens > 200000 maximum", |
| 943 | 1× | "413 Request too large", |
| 944 | 1× | "400 input length exceeds the context window", |
| 945 | 1× | ] { |
| 946 | 4× | assert!(super::too_long(e), "{e}"); |
| 947 | | } |
| 948 | 4× | for e in ["401 invalid api key", "500 too long a queue", "no answer within 90 seconds", "400 model not found"] { |
| 949 | 4× | assert!(!super::too_long(e), "{e}"); |
| 950 | | } |
| 951 | 1× | } |
| 952 | |
|
| 953 | | #[test] |
| 954 | 1× | fn headers_parse_or_name_the_bad_line() { |
| 955 | 1× | assert_eq!( |
| 956 | 1× | parse_headers("X-Org: acme\n\n OpenAI-Beta : v1 ").unwrap(), |
| 957 | 1× | vec![("X-Org".into(), "acme".into()), ("OpenAI-Beta".into(), "v1".into())] |
| 958 | | ); |
| 959 | 1× | assert_eq!(parse_headers("no colon").unwrap_err(), "no colon"); |
| 960 | 1× | assert_eq!(parse_headers("bad name: x").unwrap_err(), "bad name: x"); |
| 961 | 1× | } |
| 962 | | } |
| 963 | |
|
| 964 | |
|
| 965 | |
|
| 966 | |
|
| 967 | | pub const CALLS_PER_HOUR: usize = 30; |
| 968 | |
|
| 969 | |
|
| 970 | | #[derive(Default)] |
| 971 | | pub struct Gate { |
| 972 | | busy: std::collections::HashSet<(i64, String)>, |
| 973 | | recent: std::collections::HashMap<i64, Vec<i64>>, |
| 974 | | } |
| 975 | |
|
| 976 | | impl Gate { |
| 977 | | |
| 978 | 6× | fn take(&mut self, repo_id: i64, proposal_id: i64, name: &str) -> Result<(), String> { |
| 979 | 6× | let key = (proposal_id, name.to_lowercase()); |
| 980 | 6× | if self.busy.contains(&key) { |
| 981 | 0× | return Err(format!("{name} is already answering this proposal")); |
| 982 | 6× | } |
| 983 | 6× | let now = now_ms(); |
| 984 | 6× | let calls = self.recent.entry(repo_id).or_default(); |
| 985 | 10× | calls.retain(|t| now - *t < 3_600_000); |
| 986 | 6× | if calls.len() >= CALLS_PER_HOUR { |
| 987 | 0× | return Err(format!("this repository has asked assistants {CALLS_PER_HOUR} times in the last hour; try later")); |
| 988 | 6× | } |
| 989 | 6× | calls.push(now); |
| 990 | 6× | self.busy.insert(key); |
| 991 | 6× | Ok(()) |
| 992 | 6× | } |
| 993 | |
|
| 994 | 6× | fn release(&mut self, proposal_id: i64, name: &str) { |
| 995 | 6× | self.busy.remove(&(proposal_id, name.to_lowercase())); |
| 996 | 6× | } |
| 997 | |
|
| 998 | | |
| 999 | 77× | pub fn answering(&self, proposal_id: i64) -> Vec<String> { |
| 1000 | 77× | let mut v: Vec<String> = self.busy.iter().filter(|(p, _)| *p == proposal_id).map(|(_, n)| n.clone()).collect(); |
| 1001 | 77× | v.sort(); |
| 1002 | 77× | v |
| 1003 | 77× | } |
| 1004 | | } |
| 1005 | |
|
| 1006 | |
|
| 1007 | |
|
| 1008 | 2× | pub fn review_context( |
| 1009 | 2× | row: &crate::db::repos::RepoRow, |
| 1010 | 2× | p: &crate::db::proposals::ProposalRow, |
| 1011 | 2× | loaded: &crate::web::proposals::Loaded, |
| 1012 | 2× | budget: usize, |
| 1013 | 2× | ) -> String { |
| 1014 | 2× | let mut head = |
| 1015 | 2× | format!("# Proposal #{}: {}\n\nRepository {}/{} · into {} · by {}\n\n", p.number, p.title, row.owner, row.name, p.target, p.author); |
| 1016 | 2× | if !p.body.trim().is_empty() { |
| 1017 | 2× | head.push_str("## Description\n\n"); |
| 1018 | 2× | head.push_str(p.body.trim()); |
| 1019 | 2× | head.push_str("\n\n"); |
| 1020 | 2× | } |
| 1021 | 2× | context_text(head, &loaded.changes, &loaded.diffs, budget) |
| 1022 | 2× | } |
| 1023 | |
|
| 1024 | |
|
| 1025 | 5× | pub fn draft_context( |
| 1026 | 5× | row: &crate::db::repos::RepoRow, |
| 1027 | 5× | title: &str, |
| 1028 | 5× | target: &str, |
| 1029 | 5× | changes: &[crate::repos::view::ChangeView], |
| 1030 | 5× | diffs: &[oxy_core::repo::FileDiff], |
| 1031 | 5× | budget: usize, |
| 1032 | 5× | ) -> String { |
| 1033 | 5× | let head = format!( |
| 1034 | | "# Proposal: {}\n\nRepository {}/{} · into {}\n\n", |
| 1035 | 5× | if title.trim().is_empty() { "(untitled)" } else { title.trim() }, |
| 1036 | | row.owner, |
| 1037 | | row.name, |
| 1038 | | target |
| 1039 | | ); |
| 1040 | 5× | context_text(head, changes, diffs, budget) |
| 1041 | 5× | } |
| 1042 | |
|
| 1043 | |
|
| 1044 | |
|
| 1045 | 7× | fn context_text(mut text: String, changes: &[crate::repos::view::ChangeView], diffs: &[oxy_core::repo::FileDiff], budget: usize) -> String { |
| 1046 | 7× | text.push_str(&format!("## Changes ({})\n\n", changes.len())); |
| 1047 | 7× | for c in changes { |
| 1048 | 7× | text.push_str(&format!("- {} {}\n", c.short, c.title)); |
| 1049 | 7× | let body = c.body.trim(); |
| 1050 | 7× | if !body.is_empty() { |
| 1051 | 0× | for line in body.lines().take(20) { |
| 1052 | 0× | text.push_str(" "); |
| 1053 | 0× | text.push_str(line); |
| 1054 | 0× | text.push('\n'); |
| 1055 | 0× | } |
| 1056 | 7× | } |
| 1057 | | } |
| 1058 | | |
| 1059 | | |
| 1060 | 7× | text.push_str(&format!("\n## Files ({})\n\n", diffs.len())); |
| 1061 | 7× | let mut parts: Vec<(String, String)> = Vec::with_capacity(diffs.len()); |
| 1062 | 7× | for f in diffs { |
| 1063 | 7× | let (added, removed) = f.hunks.iter().fold((0, 0), |(a, r), h| (a + h.added.len(), r + h.removed.len())); |
| 1064 | 7× | text.push_str(&format!("- {} ({}, +{added} -{removed})\n", f.path, f.kind)); |
| 1065 | 7× | let mut diff = format!("### {} ({})\n", f.path, f.kind); |
| 1066 | 7× | if let Some(from) = &f.from { |
| 1067 | 0× | diff.push_str(&format!("renamed from {from}\n")); |
| 1068 | 7× | } |
| 1069 | 405× | let binary = f.hunks.iter().any(|h| h.added.iter().chain(h.removed.iter().map(|(_, l)| l)).any(|l| l.contains(&0))); |
| 1070 | 7× | if binary { |
| 1071 | 0× | diff.push_str("(binary)\n\n"); |
| 1072 | 0× | parts.push((f.path.clone(), diff)); |
| 1073 | 0× | continue; |
| 1074 | 7× | } |
| 1075 | 7× | diff.push_str("```\n"); |
| 1076 | 7× | for h in &f.hunks { |
| 1077 | 7× | diff.push_str(&format!("@@ after line {} @@\n", h.base_line)); |
| 1078 | 7× | for (_, l) in &h.removed { |
| 1079 | 0× | diff.push('-'); |
| 1080 | 0× | diff.push_str(String::from_utf8_lossy(l).trim_end_matches('\n')); |
| 1081 | 0× | diff.push('\n'); |
| 1082 | 0× | } |
| 1083 | 405× | for l in &h.added { |
| 1084 | 405× | diff.push('+'); |
| 1085 | 405× | diff.push_str(String::from_utf8_lossy(l).trim_end_matches('\n')); |
| 1086 | 405× | diff.push('\n'); |
| 1087 | 405× | } |
| 1088 | | } |
| 1089 | 7× | diff.push_str("```\n\n"); |
| 1090 | 7× | parts.push((f.path.clone(), diff)); |
| 1091 | | } |
| 1092 | 7× | text.push_str("\n## Diff\n\n"); |
| 1093 | | |
| 1094 | | |
| 1095 | | |
| 1096 | 7× | let head = text.chars().count(); |
| 1097 | 7× | let mut room = budget.saturating_sub(head + 200); |
| 1098 | 7× | let mut left_out = Vec::new(); |
| 1099 | 7× | for (path, diff) in parts { |
| 1100 | 7× | let n = diff.chars().count(); |
| 1101 | 7× | if n <= room { |
| 1102 | 5× | text.push_str(&diff); |
| 1103 | 5× | room -= n; |
| 1104 | 5× | } else if room >= 1_500 { |
| 1105 | 2× | let keep = room - 80; |
| 1106 | 2× | text.extend(diff.chars().take(keep)); |
| 1107 | 2× | text.push_str("\n(the rest of this file was cut to fit)\n```\n\n"); |
| 1108 | 2× | left_out.push(format!("{path} (cut)")); |
| 1109 | 2× | room = 0; |
| 1110 | 2× | } else { |
| 1111 | 0× | left_out.push(path); |
| 1112 | 0× | } |
| 1113 | | } |
| 1114 | 7× | if !left_out.is_empty() { |
| 1115 | 2× | let names: String = left_out.join(", ").chars().take(room.max(400)).collect(); |
| 1116 | 2× | text.push_str(&format!("(left out to fit the context budget, {} file(s): {names})\n", left_out.len())); |
| 1117 | 5× | } |
| 1118 | 7× | text |
| 1119 | 7× | } |
| 1120 | |
|
| 1121 | |
|
| 1122 | |
|
| 1123 | 1× | pub fn review_system_prompt(a: &Assistant, site: &str, prompts: &str) -> String { |
| 1124 | 1× | let mut s = format!( |
| 1125 | | "You are {name}, an AI assistant reviewing a proposal on {site}, an OxygenHub site. \ |
| 1126 | | A person asked for your review; it is recorded as a review comment made via you, under their name. \ |
| 1127 | | You cannot approve: only people approve. Answer in Markdown: start with a one-paragraph summary, \ |
| 1128 | | then findings (correctness first, then clarity), each naming the file and quoting the lines it is about; \ |
| 1129 | | keep praise out unless something is notably well done; say so plainly when you find nothing wrong.", |
| 1130 | | name = a.name |
| 1131 | | ); |
| 1132 | 1× | if !prompts.trim().is_empty() { |
| 1133 | 1× | s.push_str("\n\n"); |
| 1134 | 1× | s.push_str(prompts.trim()); |
| 1135 | 1× | } |
| 1136 | 1× | s |
| 1137 | 1× | } |
| 1138 | |
|
| 1139 | |
|
| 1140 | |
|
| 1141 | |
|
| 1142 | |
|
| 1143 | 1× | pub async fn request_review( |
| 1144 | 1× | app: &Shared, |
| 1145 | 1× | row: &crate::db::repos::RepoRow, |
| 1146 | 1× | p: &crate::db::proposals::ProposalRow, |
| 1147 | 1× | by: &crate::db::accounts::Account, |
| 1148 | 1× | a: Assistant, |
| 1149 | 1× | ) -> Result<(), String> { |
| 1150 | 1× | app.assistant_gate.lock().unwrap_or_else(|e| e.into_inner()).take(row.id, p.id, &a.name)?; |
| 1151 | 1× | let pool = &app.db; |
| 1152 | 1× | crate::db::proposals::add_event(pool, p.id, &by.username, &format!("asked {} to review", a.name), "").await.ok(); |
| 1153 | 1× | app.bus.live(&crate::events::topic::proposal(row.id, p.number), "assistant"); |
| 1154 | 1× | crate::db::accounts::audit( |
| 1155 | 1× | pool, |
| 1156 | 1× | Some(by), |
| 1157 | 1× | "assistant.asked", |
| 1158 | 1× | &format!("{}/{}", row.owner, row.name), |
| 1159 | 1× | &format!("{} #{}", a.name, p.number), |
| 1160 | 1× | ) |
| 1161 | 1× | .await |
| 1162 | 1× | .ok(); |
| 1163 | 1× | let (app, row, p, by) = (app.clone(), row.clone(), p.clone(), by.clone()); |
| 1164 | 1× | tokio::spawn(async move { |
| 1165 | 1× | let outcome = answer_review(&app, &row, &p, &a).await; |
| 1166 | 1× | app.assistant_gate.lock().unwrap_or_else(|e| e.into_inner()).release(p.id, &a.name); |
| 1167 | 1× | match outcome { |
| 1168 | 1× | Ok(body) => { |
| 1169 | 1× | if crate::db::proposals::add_review_via(&app.db, p.id, by.id, &by.username, "comment", &body, &a.name).await.is_ok() { |
| 1170 | 1× | crate::db::proposals::add_event(&app.db, p.id, &by.username, &format!("review.comment via {}", a.name), "").await.ok(); |
| 1171 | 1× | crate::web::proposals::reviewed(&app, &row, &p, &by, "comment").await; |
| 1172 | 1× | crate::hooks::emit( |
| 1173 | 1× | &app, |
| 1174 | 1× | &row, |
| 1175 | 1× | "review", |
| 1176 | 1× | &by.username, |
| 1177 | 1× | serde_json::json!({ "number": p.number, "verdict": "comment", "via": a.name }), |
| 1178 | 1× | ) |
| 1179 | 1× | .await; |
| 1180 | 0× | } |
| 1181 | | } |
| 1182 | 0× | Err(why) => { |
| 1183 | 0× | crate::db::proposals::add_event(&app.db, p.id, &by.username, &format!("{} did not answer", a.name), &why).await.ok(); |
| 1184 | 0× | app.bus.live(&crate::events::topic::proposal(row.id, p.number), "assistant"); |
| 1185 | | } |
| 1186 | | } |
| 1187 | 1× | }); |
| 1188 | 1× | Ok(()) |
| 1189 | 1× | } |
| 1190 | |
|
| 1191 | 1× | async fn answer_review( |
| 1192 | 1× | app: &Shared, |
| 1193 | 1× | row: &crate::db::repos::RepoRow, |
| 1194 | 1× | p: &crate::db::proposals::ProposalRow, |
| 1195 | 1× | a: &Assistant, |
| 1196 | 1× | ) -> Result<String, String> { |
| 1197 | 1× | let repo = crate::repos::names::RepoName::new(&row.owner, &row.name).ok_or("bad repository name")?; |
| 1198 | 1× | let loaded = crate::web::proposals::load(app, &repo, p.clone()).await.map_err(|e| format!("the proposal could not be loaded: {e}"))?; |
| 1199 | 1× | let budget = usize::try_from(a.budget).unwrap_or(DEFAULT_BUDGET as usize).max(1_000); |
| 1200 | 1× | let prompts = prompt_for_repo(&app.db, row).await.map_err(|e| e.to_string())?; |
| 1201 | 1× | let system = review_system_prompt(a, &app.config.site.name, &prompts); |
| 1202 | 1× | ask_within(app, a, &system, budget, BACKGROUND_TIMEOUT, |room| review_context(row, p, &loaded, room), |_| {}).await |
| 1203 | 1× | } |
| 1204 | |
|
| 1205 | |
|
| 1206 | |
|
| 1207 | |
|
| 1208 | |
|
| 1209 | 6× | async fn ask_within( |
| 1210 | 6× | app: &Shared, |
| 1211 | 6× | a: &Assistant, |
| 1212 | 6× | system: &str, |
| 1213 | 6× | budget: usize, |
| 1214 | 6× | wait: Duration, |
| 1215 | 6× | context: impl Fn(usize) -> String, |
| 1216 | 6× | mut part: impl FnMut(&str), |
| 1217 | 6× | ) -> Result<String, String> { |
| 1218 | 7× | let room = |b: usize| b.saturating_sub(system.chars().count()); |
| 1219 | | |
| 1220 | | |
| 1221 | 6× | match ask_streaming(app, a, system, &context(room(budget)), wait, &mut part).await { |
| 1222 | 1× | Err(e) if too_long(&e) => { |
| 1223 | 1× | let smaller = (budget / 3).max(1_000); |
| 1224 | 1× | tracing::info!("assistant {}: too long at {budget} characters ({e}); asking again at {smaller}", a.name); |
| 1225 | 1× | ask_streaming(app, a, system, &context(room(smaller)), wait, &mut part).await.map_err(|e2| { |
| 1226 | 0× | format!( |
| 1227 | | "the proposal was too large for {}'s model at {budget} characters and again at {smaller} ({e2}): lower this assistant's context budget", |
| 1228 | | a.name |
| 1229 | | ) |
| 1230 | 0× | }) |
| 1231 | | } |
| 1232 | 0× | Err(e) if e.starts_with("no answer within") => Err(format!( |
| 1233 | 0× | "{e}: a large proposal on a slow or local model can take longer; lower this assistant's context budget, or use a faster model" |
| 1234 | 0× | )), |
| 1235 | 5× | other => other, |
| 1236 | | } |
| 1237 | 6× | } |
| 1238 | |
|
| 1239 | |
|
| 1240 | 9× | fn too_long(e: &str) -> bool { |
| 1241 | 9× | let l = e.to_lowercase(); |
| 1242 | 9× | (l.starts_with("400") || l.starts_with("413") || l.starts_with("422")) |
| 1243 | 6× | && [ |
| 1244 | 6× | "too long", |
| 1245 | 6× | "too large", |
| 1246 | 6× | "context length", |
| 1247 | 6× | "context window", |
| 1248 | 6× | "maximum context", |
| 1249 | 6× | "too many tokens", |
| 1250 | 6× | "exceeds", |
| 1251 | 6× | "reduce the length", |
| 1252 | 6× | "input length", |
| 1253 | 6× | ] |
| 1254 | 6× | .iter() |
| 1255 | 22× | .any(|k| l.contains(k)) |
| 1256 | 9× | } |
| 1257 | |
|
| 1258 | |
|
| 1259 | |
|
| 1260 | |
|
| 1261 | 1× | pub fn summary_system_prompt(a: &Assistant, site: &str, prompts: &str) -> String { |
| 1262 | 1× | let mut s = format!( |
| 1263 | | "You are {name}, an AI assistant working on {site}, an OxygenHub site, with a next-generation version control system \ |
| 1264 | | called Oxygen (its command is oxy). A proposal is Oxygen's equivalent of a merge request or pull request. \ |
| 1265 | | Please summarise the proposal and comment on it. A person asked for this; it is posted as a comment made via you, \ |
| 1266 | | under their name. Answer in Markdown: a short summary of what the proposal changes and why, then anything a reviewer \ |
| 1267 | | should look at closely; keep to what the description and the diff show.", |
| 1268 | | name = a.name |
| 1269 | | ); |
| 1270 | 1× | if !prompts.trim().is_empty() { |
| 1271 | 1× | s.push_str("\n\n"); |
| 1272 | 1× | s.push_str(prompts.trim()); |
| 1273 | 1× | } |
| 1274 | 1× | s |
| 1275 | 1× | } |
| 1276 | |
|
| 1277 | |
|
| 1278 | |
|
| 1279 | 1× | pub async fn request_summary( |
| 1280 | 1× | app: &Shared, |
| 1281 | 1× | row: &crate::db::repos::RepoRow, |
| 1282 | 1× | p: &crate::db::proposals::ProposalRow, |
| 1283 | 1× | by: &crate::db::accounts::Account, |
| 1284 | 1× | a: Assistant, |
| 1285 | 1× | ) -> Result<(), String> { |
| 1286 | 1× | app.assistant_gate.lock().unwrap_or_else(|e| e.into_inner()).take(row.id, p.id, &a.name)?; |
| 1287 | 1× | let pool = &app.db; |
| 1288 | 1× | crate::db::proposals::add_event(pool, p.id, &by.username, &format!("asked {} for a summary", a.name), "").await.ok(); |
| 1289 | 1× | app.bus.live(&crate::events::topic::proposal(row.id, p.number), "assistant"); |
| 1290 | 1× | crate::db::accounts::audit( |
| 1291 | 1× | pool, |
| 1292 | 1× | Some(by), |
| 1293 | 1× | "assistant.asked", |
| 1294 | 1× | &format!("{}/{}", row.owner, row.name), |
| 1295 | 1× | &format!("{} #{} summary", a.name, p.number), |
| 1296 | 1× | ) |
| 1297 | 1× | .await |
| 1298 | 1× | .ok(); |
| 1299 | 1× | let (app, row, p, by) = (app.clone(), row.clone(), p.clone(), by.clone()); |
| 1300 | 1× | tokio::spawn(async move { |
| 1301 | 1× | let outcome = answer_summary(&app, &row, &p, &a).await; |
| 1302 | 1× | app.assistant_gate.lock().unwrap_or_else(|e| e.into_inner()).release(p.id, &a.name); |
| 1303 | 1× | match outcome { |
| 1304 | 1× | Ok(body) => { |
| 1305 | 1× | crate::db::proposals::add_comment_via(&app.db, p.id, by.id, &by.username, &body, &a.name).await.ok(); |
| 1306 | 1× | app.bus.live(&crate::events::topic::proposal(row.id, p.number), "comment"); |
| 1307 | | } |
| 1308 | 0× | Err(why) => { |
| 1309 | 0× | crate::db::proposals::add_event(&app.db, p.id, &by.username, &format!("{} did not answer", a.name), &why).await.ok(); |
| 1310 | 0× | app.bus.live(&crate::events::topic::proposal(row.id, p.number), "assistant"); |
| 1311 | | } |
| 1312 | | } |
| 1313 | 1× | }); |
| 1314 | 1× | Ok(()) |
| 1315 | 1× | } |
| 1316 | |
|
| 1317 | 1× | async fn answer_summary( |
| 1318 | 1× | app: &Shared, |
| 1319 | 1× | row: &crate::db::repos::RepoRow, |
| 1320 | 1× | p: &crate::db::proposals::ProposalRow, |
| 1321 | 1× | a: &Assistant, |
| 1322 | 1× | ) -> Result<String, String> { |
| 1323 | 1× | let repo = crate::repos::names::RepoName::new(&row.owner, &row.name).ok_or("bad repository name")?; |
| 1324 | 1× | let loaded = crate::web::proposals::load(app, &repo, p.clone()).await.map_err(|e| format!("the proposal could not be loaded: {e}"))?; |
| 1325 | 1× | let budget = usize::try_from(a.budget).unwrap_or(DEFAULT_BUDGET as usize).max(1_000); |
| 1326 | 1× | let prompts = prompt_for_repo(&app.db, row).await.map_err(|e| e.to_string())?; |
| 1327 | 1× | let system = summary_system_prompt(a, &app.config.site.name, &prompts); |
| 1328 | 1× | ask_within(app, a, &system, budget, BACKGROUND_TIMEOUT, |room| review_context(row, p, &loaded, room), |_| {}).await |
| 1329 | 1× | } |
| 1330 | |
|
| 1331 | |
|
| 1332 | |
|
| 1333 | |
|
| 1334 | |
|
| 1335 | 4× | pub fn draft_system_prompt(a: &Assistant, site: &str, prompts: &str, with_title: bool) -> String { |
| 1336 | 4× | let shape = if with_title { |
| 1337 | 2× | "Begin with one line `Title: ` and a title of at most 72 characters, in the style of the changes' own titles \ |
| 1338 | 2× | (Conventional Commits if they use them), then a blank line, then the description in Markdown: \ |
| 1339 | 2× | no greeting, no sign-off." |
| 1340 | | } else { |
| 1341 | 2× | "Write Markdown for the description field only: no title line, no greeting, no sign-off." |
| 1342 | | }; |
| 1343 | 4× | let mut s = format!( |
| 1344 | | "You are {name}, an AI assistant drafting the description of a proposal on {site}, an OxygenHub site, \ |
| 1345 | | for the person who will open or edit it; they read and change your draft before saving, and the draft is theirs. \ |
| 1346 | | {shape} \ |
| 1347 | | Say what the changes do and why, what a reviewer should look at, and how to verify them; \ |
| 1348 | | keep to what the diff shows, and keep it short.", |
| 1349 | | name = a.name |
| 1350 | | ); |
| 1351 | 4× | if !prompts.trim().is_empty() { |
| 1352 | 3× | s.push_str("\n\n"); |
| 1353 | 3× | s.push_str(prompts.trim()); |
| 1354 | 3× | } |
| 1355 | 4× | s |
| 1356 | 4× | } |
| 1357 | |
|
| 1358 | |
|
| 1359 | | pub struct Draft<'a> { |
| 1360 | | |
| 1361 | | |
| 1362 | | pub slot: i64, |
| 1363 | | pub title: &'a str, |
| 1364 | | pub target: &'a str, |
| 1365 | | pub changes: &'a [crate::repos::view::ChangeView], |
| 1366 | | pub diffs: &'a [oxy_core::repo::FileDiff], |
| 1367 | | |
| 1368 | | pub with_title: bool, |
| 1369 | | } |
| 1370 | |
|
| 1371 | |
|
| 1372 | |
|
| 1373 | 6× | pub fn split_title(text: &str) -> (Option<String>, String) { |
| 1374 | 6× | let text = text.trim_start(); |
| 1375 | 6× | let (first, rest) = text.split_once('\n').unwrap_or((text, "")); |
| 1376 | 6× | let Some(t) = first.trim().strip_prefix("Title:").or_else(|| first.trim().strip_prefix("title:")) else { |
| 1377 | 3× | return (None, text.trim().to_string()); |
| 1378 | | }; |
| 1379 | 6× | let t: String = t.trim().trim_matches(|c| c == '"' || c == '`' || c == '*').trim().chars().take(120).collect(); |
| 1380 | 3× | ((!t.is_empty()).then_some(t), rest.trim().to_string()) |
| 1381 | 6× | } |
| 1382 | |
|
| 1383 | |
|
| 1384 | |
|
| 1385 | 3× | pub async fn draft_description( |
| 1386 | 3× | app: &Shared, |
| 1387 | 3× | row: &crate::db::repos::RepoRow, |
| 1388 | 3× | by: &crate::db::accounts::Account, |
| 1389 | 3× | a: &Assistant, |
| 1390 | 3× | d: Draft<'_>, |
| 1391 | 3× | ) -> Result<String, String> { |
| 1392 | 3× | draft_description_streaming(app, row, by, a, d, |_| {}).await |
| 1393 | 3× | } |
| 1394 | |
|
| 1395 | |
|
| 1396 | |
|
| 1397 | 4× | pub async fn draft_description_streaming( |
| 1398 | 4× | app: &Shared, |
| 1399 | 4× | row: &crate::db::repos::RepoRow, |
| 1400 | 4× | by: &crate::db::accounts::Account, |
| 1401 | 4× | a: &Assistant, |
| 1402 | 4× | d: Draft<'_>, |
| 1403 | 4× | part: impl FnMut(&str), |
| 1404 | 4× | ) -> Result<String, String> { |
| 1405 | 4× | let Draft { slot, title, target, changes, diffs, with_title } = d; |
| 1406 | 4× | app.assistant_gate.lock().unwrap_or_else(|e| e.into_inner()).take(row.id, slot, &a.name)?; |
| 1407 | 4× | crate::db::accounts::audit(&app.db, Some(by), "assistant.drafted", &format!("{}/{}", row.owner, row.name), &a.name).await.ok(); |
| 1408 | 4× | let budget = usize::try_from(a.budget).unwrap_or(DEFAULT_BUDGET as usize).max(1_000); |
| 1409 | 4× | let prompts = prompt_for_repo(&app.db, row).await.map_err(|e| e.to_string()); |
| 1410 | 4× | let outcome = match prompts { |
| 1411 | 4× | Ok(prompts) => { |
| 1412 | 4× | let system = draft_system_prompt(a, &app.config.site.name, &prompts, with_title); |
| 1413 | 5× | ask_within(app, a, &system, budget, CALL_TIMEOUT, |room| draft_context(row, title, target, changes, diffs, room), part).await |
| 1414 | | } |
| 1415 | 0× | Err(e) => Err(e), |
| 1416 | | }; |
| 1417 | 4× | app.assistant_gate.lock().unwrap_or_else(|e| e.into_inner()).release(slot, &a.name); |
| 1418 | 4× | outcome.map(|t| t.trim().to_string()) |
| 1419 | 4× | } |