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