OXYGENHUB
# c4027e038d25 oxygenhub / crates / hub / src / assistants.rs
42.1 KiB · Rust · last changed Outbound requests only to public addresses 20 days ago source blame history raw
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;
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;
35
/// The most an answer is kept.
36
const MAX_ANSWER: usize = 64 * 1024;
37
38
pub use crate::levels::Scope;
39
40
#[derive(Clone, Debug, FromRow)]
41
pub struct Assistant {
42
    pub id: i64,
43
    pub scope: String,
44
    pub scope_id: i64,
45
    pub name: String,
46
    pub endpoint: String,
47
    /// One of [`APIS`].
48
    pub api: String,
49
    pub model: String,
50
    /// Extra request headers, one `Name: value` per line.
51
    pub headers: String,
52
    /// The sealed token; empty when the endpoint needs none.
53
    pub token: Vec<u8>,
54
    pub enabled: bool,
55
    pub budget: i64,
56
    pub created_ms: i64,
57
    pub updated_ms: i64,
58
}
59
60
impl Assistant {
61
    pub fn scope(&self) -> Scope {
62
        match self.scope.as_str() {
63
            "org" => Scope::Org(self.scope_id),
64
            "repo" => Scope::Repo(self.scope_id),
65
            _ => Scope::Site,
66
        }
67
    }
68
69
    pub fn has_token(&self) -> bool {
70
        !self.token.is_empty()
71
    }
72
73
    /// The extra headers as pairs; malformed lines are skipped.
74
    pub fn header_list(&self) -> Vec<(String, String)> {
75
        parse_headers(&self.headers).unwrap_or_default()
76
    }
77
}
78
79
/// What a settings form supplies for an assistant.
80
#[derive(Clone, Debug, Default)]
81
pub struct Fields {
82
    pub name: String,
83
    pub endpoint: String,
84
    pub api: String,
85
    pub model: String,
86
    pub headers: String,
87
    pub enabled: bool,
88
    pub budget: i64,
89
}
90
91
/// `Name: value` lines to pairs: `Err(line)` names the first bad one.
92
pub fn parse_headers(text: &str) -> Result<Vec<(String, String)>, String> {
93
    let mut out = vec![];
94
    for line in text.lines().map(str::trim).filter(|l| !l.is_empty()) {
95
        let Some((name, value)) = line.split_once(':') else { return Err(line.to_string()) };
96
        let (name, value) = (name.trim(), value.trim());
97
        let name_ok = !name.is_empty() && name.chars().all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_');
98
        if !name_ok || value.chars().any(|c| c.is_control()) {
99
            return Err(line.to_string());
100
        }
101
        out.push((name.to_string(), value.to_string()));
102
    }
103
    Ok(out)
104
}
105
106
// -------------------------------------------------------------- sealing
107
108
/// The key that seals tokens, kept in `<data>/keys/assistants.key`.
109
#[derive(Clone)]
110
pub struct SecretKey([u8; 32]);
111
112
impl SecretKey {
113
    pub fn load_or_create(keys: &Path) -> anyhow::Result<SecretKey> {
114
        SecretKey::load_or_create_named(keys, "assistants.key")
115
    }
116
117
    /// The same, for another key file beside it (`secrets.key` seals Flow
118
    /// secrets, #107).
119
    pub fn load_or_create_named(keys: &Path, file: &str) -> anyhow::Result<SecretKey> {
120
        std::fs::create_dir_all(keys)?;
121
        let path = keys.join(file);
122
        if let Ok(bytes) = std::fs::read(&path) {
123
            if bytes.len() == 32 {
124
                return Ok(SecretKey(bytes.try_into().unwrap()));
125
            }
126
        }
127
        let mut k = [0u8; 32];
128
        rand::rngs::OsRng.fill_bytes(&mut k);
129
        std::fs::write(&path, k)?;
130
        #[cfg(unix)]
131
        {
132
            use std::os::unix::fs::PermissionsExt;
133
            std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600))?;
134
        }
135
        Ok(SecretKey(k))
136
    }
137
138
    fn key(&self) -> aead::LessSafeKey {
139
        aead::LessSafeKey::new(aead::UnboundKey::new(&aead::CHACHA20_POLY1305, &self.0).expect("key length"))
140
    }
141
142
    /// Seal a token: a random nonce, then the ciphertext and tag. An
143
    /// empty token seals to nothing, meaning "no token".
144
    pub fn seal(&self, plain: &str) -> Vec<u8> {
145
        if plain.is_empty() {
146
            return vec![];
147
        }
148
        let mut nonce = [0u8; 12];
149
        rand::rngs::OsRng.fill_bytes(&mut nonce);
150
        let mut out = plain.as_bytes().to_vec();
151
        self.key().seal_in_place_append_tag(aead::Nonce::assume_unique_for_key(nonce), aead::Aad::empty(), &mut out).expect("seal");
152
        let mut sealed = nonce.to_vec();
153
        sealed.extend(out);
154
        sealed
155
    }
156
157
    /// The token back, or `None` when the bytes were sealed under another
158
    /// key or damaged. Nothing opens to an empty token.
159
    pub fn open(&self, sealed: &[u8]) -> Option<String> {
160
        if sealed.is_empty() {
161
            return Some(String::new());
162
        }
163
        let (nonce, rest) = sealed.split_at_checked(12)?;
164
        let mut buf = rest.to_vec();
165
        let nonce = aead::Nonce::try_assume_unique_for_key(nonce).ok()?;
166
        let plain = self.key().open_in_place(nonce, aead::Aad::empty(), &mut buf).ok()?;
167
        String::from_utf8(plain.to_vec()).ok()
168
    }
169
}
170
171
// ------------------------------------------------------------- storage
172
173
pub async fn list(pool: &Pool, scope: Scope) -> sqlx::Result<Vec<Assistant>> {
174
    sqlx::query_as("SELECT * FROM assistants WHERE scope = ? AND scope_id = ? ORDER BY name")
175
        .bind(scope.kind())
176
        .bind(scope.id())
177
        .fetch_all(pool)
178
        .await
179
}
180
181
pub async fn get(pool: &Pool, id: i64) -> sqlx::Result<Option<Assistant>> {
182
    sqlx::query_as("SELECT * FROM assistants WHERE id = ?").bind(id).fetch_optional(pool).await
183
}
184
185
/// Add an assistant; `None` when the level already has one of that name.
186
pub async fn create(pool: &Pool, scope: Scope, f: &Fields, token: &[u8]) -> sqlx::Result<Option<Assistant>> {
187
    let now = now_ms();
188
    let r = sqlx::query_as(
189
        "INSERT INTO assistants (scope, scope_id, name, endpoint, api, model, headers, token, enabled, budget, created_ms, updated_ms) \
190
         VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT DO NOTHING RETURNING *",
191
    )
192
    .bind(scope.kind())
193
    .bind(scope.id())
194
    .bind(&f.name)
195
    .bind(&f.endpoint)
196
    .bind(&f.api)
197
    .bind(&f.model)
198
    .bind(&f.headers)
199
    .bind(token)
200
    .bind(f.enabled)
201
    .bind(f.budget)
202
    .bind(now)
203
    .bind(now)
204
    .fetch_optional(pool)
205
    .await?;
206
    Ok(r)
207
}
208
209
/// Change an assistant's details; `token` is `None` to keep the stored
210
/// one. False when the name is taken at that level.
211
pub async fn update(pool: &Pool, id: i64, f: &Fields, token: Option<&[u8]>) -> sqlx::Result<bool> {
212
    let r = match token {
213
        Some(t) => sqlx::query(
214
            "UPDATE assistants SET name = ?, endpoint = ?, api = ?, model = ?, headers = ?, enabled = ?, budget = ?, token = ?, updated_ms = ? \
215
             WHERE id = ? AND NOT EXISTS (SELECT 1 FROM assistants o WHERE o.scope = assistants.scope AND o.scope_id = assistants.scope_id AND o.name = ? AND o.id <> ?)",
216
        )
217
        .bind(&f.name)
218
        .bind(&f.endpoint)
219
        .bind(&f.api)
220
        .bind(&f.model)
221
        .bind(&f.headers)
222
        .bind(f.enabled)
223
        .bind(f.budget)
224
        .bind(t)
225
        .bind(now_ms())
226
        .bind(id)
227
        .bind(&f.name)
228
        .bind(id)
229
        .execute(pool)
230
        .await?,
231
        None => sqlx::query(
232
            "UPDATE assistants SET name = ?, endpoint = ?, api = ?, model = ?, headers = ?, enabled = ?, budget = ?, updated_ms = ? \
233
             WHERE id = ? AND NOT EXISTS (SELECT 1 FROM assistants o WHERE o.scope = assistants.scope AND o.scope_id = assistants.scope_id AND o.name = ? AND o.id <> ?)",
234
        )
235
        .bind(&f.name)
236
        .bind(&f.endpoint)
237
        .bind(&f.api)
238
        .bind(&f.model)
239
        .bind(&f.headers)
240
        .bind(f.enabled)
241
        .bind(f.budget)
242
        .bind(now_ms())
243
        .bind(id)
244
        .bind(&f.name)
245
        .bind(id)
246
        .execute(pool)
247
        .await?,
248
    };
249
    Ok(r.rows_affected() == 1)
250
}
251
252
/// A renamed assistant keeps its name on the reviews it answered and the
253
/// events that mention it, for the repositories its level covers: the
254
/// name is what people read, not a key.
255
pub async fn rename_in_history(pool: &Pool, scope: Scope, owner_hint: &str, old: &str, new: &str) -> sqlx::Result<()> {
256
    if old == new || old.is_empty() || new.is_empty() {
257
        return Ok(());
258
    }
259
    let like = format!("%{old}%");
260
    match scope {
261
        Scope::Site => {
262
            for table in ["proposal_reviews", "proposal_comments"] {
263
                sqlx::query(&format!("UPDATE {table} SET via = ? WHERE via = ?")).bind(new).bind(old).execute(pool).await?;
264
            }
265
            sqlx::query("UPDATE proposal_events SET event = replace(event, ?, ?) WHERE event LIKE ?")
266
                .bind(old)
267
                .bind(new)
268
                .bind(&like)
269
                .execute(pool)
270
                .await?;
271
        }
272
        Scope::Org(_) => {
273
            let in_org = "proposal_id IN (SELECT p.id FROM proposals p JOIN repositories r ON r.id = p.repo_id WHERE r.owner = ?)";
274
            for table in ["proposal_reviews", "proposal_comments"] {
275
                sqlx::query(&format!("UPDATE {table} SET via = ? WHERE via = ? AND {in_org}"))
276
                    .bind(new)
277
                    .bind(old)
278
                    .bind(owner_hint)
279
                    .execute(pool)
280
                    .await?;
281
            }
282
            sqlx::query(&format!("UPDATE proposal_events SET event = replace(event, ?, ?) WHERE event LIKE ? AND {in_org}"))
283
                .bind(old)
284
                .bind(new)
285
                .bind(&like)
286
                .bind(owner_hint)
287
                .execute(pool)
288
                .await?;
289
        }
290
        Scope::Repo(id) => {
291
            let in_repo = "proposal_id IN (SELECT id FROM proposals WHERE repo_id = ?)";
292
            for table in ["proposal_reviews", "proposal_comments"] {
293
                sqlx::query(&format!("UPDATE {table} SET via = ? WHERE via = ? AND {in_repo}"))
294
                    .bind(new)
295
                    .bind(old)
296
                    .bind(id)
297
                    .execute(pool)
298
                    .await?;
299
            }
300
            sqlx::query(&format!("UPDATE proposal_events SET event = replace(event, ?, ?) WHERE event LIKE ? AND {in_repo}"))
301
                .bind(old)
302
                .bind(new)
303
                .bind(&like)
304
                .bind(id)
305
                .execute(pool)
306
                .await?;
307
        }
308
    }
309
    Ok(())
310
}
311
312
pub async fn set_enabled(pool: &Pool, id: i64, enabled: bool) -> sqlx::Result<()> {
313
    sqlx::query("UPDATE assistants SET enabled = ?, updated_ms = ? WHERE id = ?")
314
        .bind(enabled)
315
        .bind(now_ms())
316
        .bind(id)
317
        .execute(pool)
318
        .await?;
319
    Ok(())
320
}
321
322
pub async fn delete(pool: &Pool, id: i64) -> sqlx::Result<()> {
323
    sqlx::query("DELETE FROM assistants WHERE id = ?").bind(id).execute(pool).await?;
324
    Ok(())
325
}
326
327
/// Forget a level's assistants and prompt, when the organisation or
328
/// repository goes.
329
pub async fn delete_scope(pool: &Pool, scope: Scope) -> sqlx::Result<()> {
330
    sqlx::query("DELETE FROM assistants WHERE scope = ? AND scope_id = ?").bind(scope.kind()).bind(scope.id()).execute(pool).await?;
331
    sqlx::query("DELETE FROM assistant_prompts WHERE scope = ? AND scope_id = ?").bind(scope.kind()).bind(scope.id()).execute(pool).await?;
332
    Ok(())
333
}
334
335
/// The prompt a level adds, empty when none.
336
pub async fn prompt(pool: &Pool, scope: Scope) -> sqlx::Result<String> {
337
    Ok(sqlx::query_scalar("SELECT prompt FROM assistant_prompts WHERE scope = ? AND scope_id = ?")
338
        .bind(scope.kind())
339
        .bind(scope.id())
340
        .fetch_optional(pool)
341
        .await?
342
        .unwrap_or_default())
343
}
344
345
pub async fn set_prompt(pool: &Pool, scope: Scope, text: &str) -> sqlx::Result<()> {
346
    if text.trim().is_empty() {
347
        sqlx::query("DELETE FROM assistant_prompts WHERE scope = ? AND scope_id = ?")
348
            .bind(scope.kind())
349
            .bind(scope.id())
350
            .execute(pool)
351
            .await?;
352
    } else {
353
        sqlx::query(
354
            "INSERT INTO assistant_prompts (scope, scope_id, prompt, updated_ms) VALUES (?, ?, ?, ?) \
355
             ON CONFLICT (scope, scope_id) DO UPDATE SET prompt = excluded.prompt, updated_ms = excluded.updated_ms",
356
        )
357
        .bind(scope.kind())
358
        .bind(scope.id())
359
        .bind(text.trim())
360
        .bind(now_ms())
361
        .execute(pool)
362
        .await?;
363
    }
364
    Ok(())
365
}
366
367
// ----------------------------------------------------- what applies where
368
369
/// An assistant seen from a level: where it comes from, and whether it is
370
/// on once outer levels have had their say.
371
#[derive(Clone, Debug)]
372
pub struct Available {
373
    pub assistant: Assistant,
374
    /// `the site`, `the acme organisation`, or empty for the level's own.
375
    pub from: String,
376
    /// Off because a level above disabled an assistant of this name.
377
    pub off_above: bool,
378
}
379
380
impl Available {
381
    pub fn on(&self) -> bool {
382
        self.assistant.enabled && !self.off_above
383
    }
384
}
385
386
/// Every assistant a level sees, outermost first, one per name: the
387
/// innermost definition of a name wins, and a name disabled at an outer
388
/// level is off inside it whatever the inner level says.
389
pub async fn available(pool: &Pool, scope: Scope, owner: &str) -> sqlx::Result<Vec<Available>> {
390
    let mut levels = crate::levels::above(pool, scope, owner).await?;
391
    levels.push((scope, String::new()));
392
    let mut out: Vec<Available> = vec![];
393
    let mut off: Vec<String> = vec![];
394
    for (level, from) in levels {
395
        for a in list(pool, level).await? {
396
            let off_above = off.iter().any(|n| n.eq_ignore_ascii_case(&a.name));
397
            if !a.enabled && !off.iter().any(|n| n.eq_ignore_ascii_case(&a.name)) {
398
                off.push(a.name.clone());
399
            }
400
            out.retain(|o| !o.assistant.name.eq_ignore_ascii_case(&a.name));
401
            out.push(Available { assistant: a, from: from.clone(), off_above });
402
        }
403
    }
404
    Ok(out)
405
}
406
407
/// The assistants a repository may ask, by name.
408
pub async fn enabled_for_repo(pool: &Pool, repo: &crate::db::repos::RepoRow) -> sqlx::Result<Vec<Assistant>> {
409
    Ok(available(pool, Scope::Repo(repo.id), &repo.owner).await?.into_iter().filter(Available::on).map(|a| a.assistant).collect())
410
}
411
412
/// The prompts a level inherits, outermost first: (where from, text),
413
/// only the levels that set one.
414
pub async fn inherited_prompts(pool: &Pool, scope: Scope, owner: &str) -> sqlx::Result<Vec<(String, String)>> {
415
    let mut out = vec![];
416
    for (level, from) in crate::levels::above(pool, scope, owner).await? {
417
        let text = prompt(pool, level).await?;
418
        if !text.is_empty() {
419
            out.push((from, text));
420
        }
421
    }
422
    Ok(out)
423
}
424
425
/// The prompt in force for a repository: each level's text under a
426
/// heading, outermost first, so a project adds to what the site says.
427
pub async fn prompt_for_repo(pool: &Pool, repo: &crate::db::repos::RepoRow) -> sqlx::Result<String> {
428
    let scope = Scope::Repo(repo.id);
429
    let mut parts = inherited_prompts(pool, scope, &repo.owner).await?;
430
    let own = prompt(pool, scope).await?;
431
    if !own.is_empty() {
432
        parts.push(("this repository".to_string(), own));
433
    }
434
    Ok(parts.into_iter().map(|(from, text)| format!("## Instructions from {from}\n\n{text}\n")).collect::<Vec<_>>().join("\n"))
435
}
436
437
// -------------------------------------------------------------- calling
438
439
/// Ask an assistant: the system prompt and the request, the answer's
440
/// text. Errors are for the person who asked: a status and the first
441
/// line of what the endpoint said, or why no answer came.
442
pub async fn ask(app: &Shared, a: &Assistant, system: &str, user: &str, wait: Duration) -> Result<String, String> {
443
    let Some(token) = app.assistant_key.open(&a.token) else {
444
        return Err("its token cannot be read: set it again".into());
445
    };
446
    let client = crate::net::client(&a.endpoint, &app.config.outbound.allow, wait).await?;
447
    let mut req = client.post(&a.endpoint).header("content-type", "application/json");
448
    let body = match a.api.as_str() {
449
        "anthropic-messages" => {
450
            if !token.is_empty() {
451
                req = req.header("x-api-key", &token);
452
            }
453
            req = req.header("anthropic-version", "2023-06-01");
454
            serde_json::json!({
455
                "model": a.model,
456
                "max_tokens": MAX_ANSWER_TOKENS,
457
                "system": system,
458
                "messages": [{ "role": "user", "content": user }],
459
            })
460
        }
461
        _ => {
462
            if !token.is_empty() {
463
                req = req.header("authorization", format!("Bearer {token}"));
464
            }
465
            // No length here: providers differ (`max_tokens`, or
466
            // `max_completion_tokens` for newer OpenAI models, with
467
            // different ceilings), and their defaults are generous.
468
            serde_json::json!({
469
                "model": a.model,
470
                "messages": [
471
                    { "role": "system", "content": system },
472
                    { "role": "user", "content": user },
473
                ],
474
            })
475
        }
476
    };
477
    for (name, value) in a.header_list() {
478
        req = req.header(name, value);
479
    }
480
    let r = req.json(&body).send().await.map_err(|e| {
481
        if e.is_timeout() {
482
            format!("no answer within {} seconds", wait.as_secs())
483
        } else {
484
            scrub(&e.to_string(), &token)
485
        }
486
    })?;
487
    let status = r.status();
488
    let text: String = r.text().await.unwrap_or_default().chars().take(MAX_ANSWER).collect();
489
    if !status.is_success() {
490
        let why = serde_json::from_str::<serde_json::Value>(&text)
491
            .ok()
492
            .and_then(|v| v["error"]["message"].as_str().or(v["error"].as_str()).or(v["message"].as_str()).map(str::to_string))
493
            .unwrap_or_else(|| text.lines().next().unwrap_or("").chars().take(160).collect());
494
        return Err(scrub(&format!("{} {}", status.as_u16(), why.trim()), &token));
495
    }
496
    let v: serde_json::Value = serde_json::from_str(&text).map_err(|_| "the answer was not JSON".to_string())?;
497
    let (answer, stop) = answer_of(&a.api, &v);
498
    finish(answer, stop)
499
}
500
501
/// Why an answer ended, as the person who asked needs to know it.
502
#[derive(Debug, PartialEq, Eq)]
503
enum Stop {
504
    /// It finished.
505
    Done,
506
    /// It reached the length it was allowed (thinking counts against it).
507
    Length,
508
    /// The provider declined to answer, with its category if it gave one.
509
    Declined(Option<String>),
510
}
511
512
/// The text of a provider's answer (text parts only: thinking is not
513
/// part of the answer), and why it ended.
514
fn answer_of(api: &str, v: &serde_json::Value) -> (String, Stop) {
515
    match api {
516
        "anthropic-messages" => {
517
            let text = v["content"]
518
                .as_array()
519
                .map(|parts| {
520
                    parts
521
                        .iter()
522
                        .filter(|p| p["type"].as_str().is_none_or(|t| t == "text"))
523
                        .filter_map(|p| p["text"].as_str())
524
                        .collect::<Vec<_>>()
525
                        .join("")
526
                })
527
                .unwrap_or_default();
528
            let stop = match v["stop_reason"].as_str() {
529
                Some("max_tokens") => Stop::Length,
530
                Some("refusal") => Stop::Declined(v["stop_details"]["category"].as_str().map(str::to_string)),
531
                _ => Stop::Done,
532
            };
533
            (text, stop)
534
        }
535
        _ => {
536
            let choice = &v["choices"][0];
537
            let text = match &choice["message"]["content"] {
538
                serde_json::Value::String(s) => s.clone(),
539
                serde_json::Value::Array(parts) => parts.iter().filter_map(|p| p["text"].as_str()).collect::<Vec<_>>().join(""),
540
                _ => String::new(),
541
            };
542
            let stop = match choice["finish_reason"].as_str() {
543
                Some("length") => Stop::Length,
544
                Some("content_filter") => Stop::Declined(None),
545
                _ => Stop::Done,
546
            };
547
            (text, stop)
548
        }
549
    }
550
}
551
552
/// The answer, or why there is none. An answer cut short is kept, with a
553
/// note saying so.
554
fn finish(answer: String, stop: Stop) -> Result<String, String> {
555
    match stop {
556
        Stop::Declined(category) => Err(match category {
557
            Some(c) => format!("it declined to answer ({c})"),
558
            None => "it declined to answer".into(),
559
        }),
560
        Stop::Length if answer.trim().is_empty() => {
561
            Err("it ran out of room before writing an answer (it spent its length limit thinking): try a smaller proposal".into())
562
        }
563
        Stop::Length => Ok(format!("{}\n\n*(Cut short: the answer reached its length limit.)*", answer.trim_end())),
564
        Stop::Done if answer.trim().is_empty() => Err("the answer carried no text".into()),
565
        Stop::Done => Ok(answer),
566
    }
567
}
568
569
/// Providers the add form offers, with what they need (#109):
570
/// (id, label, name, endpoint, api).
571
pub const PRESETS: &[(&str, &str, &str, &str, &str)] = &[
572
    ("anthropic", "Anthropic (Claude)", "Claude", "https://api.anthropic.com/v1/messages", "anthropic-messages"),
573
    ("openai", "OpenAI (ChatGPT)", "ChatGPT", "https://api.openai.com/v1/chat/completions", "chat-completions"),
574
    ("gemini", "Google Gemini", "Gemini", "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions", "chat-completions"),
575
    ("deepseek", "DeepSeek", "DeepSeek", "https://api.deepseek.com/chat/completions", "chat-completions"),
576
    ("mistral", "Mistral", "Mistral", "https://api.mistral.ai/v1/chat/completions", "chat-completions"),
577
    ("perplexity", "Perplexity", "Perplexity", "https://api.perplexity.ai/chat/completions", "chat-completions"),
578
    ("ollama", "Ollama (on your network)", "Ollama", "http://localhost:11434/v1/chat/completions", "chat-completions"),
579
];
580
581
/// Where an endpoint lists its models: `…/chat/completions` → `…/models`,
582
/// `…/messages` → `…/models`.
583
pub fn models_url(endpoint: &str) -> Option<String> {
584
    let e = endpoint.trim_end_matches('/');
585
    e.strip_suffix("/chat/completions").or_else(|| e.strip_suffix("/messages")).map(|base| format!("{base}/models"))
586
}
587
588
/// The chat models an endpoint offers for this token, sorted; `Err` when
589
/// it cannot say (no list, refused, unreachable). Embedding, speech and
590
/// image models are left out.
591
pub async fn list_models(endpoint: &str, api: &str, token: &str, allow: &[String]) -> Result<Vec<String>, String> {
592
    let url = models_url(endpoint).ok_or("this endpoint has no model list next to it")?;
593
    let client = crate::net::client(&url, allow, std::time::Duration::from_secs(15)).await?;
594
    let mut req = client.get(&url);
595
    if api == "anthropic-messages" {
596
        req = req.header("anthropic-version", "2023-06-01");
597
        if !token.is_empty() {
598
            req = req.header("x-api-key", token);
599
        }
600
    } else if !token.is_empty() {
601
        req = req.header("authorization", format!("Bearer {token}"));
602
    }
603
    let r = req.send().await.map_err(|e| scrub(&e.to_string(), token))?;
604
    let status = r.status();
605
    if !status.is_success() {
606
        return Err(format!("the model list answered {}", status.as_u16()));
607
    }
608
    let v: serde_json::Value = r.json().await.map_err(|_| "the model list was not JSON".to_string())?;
609
    let skip = ["embed", "tts", "whisper", "dall-e", "moderation", "audio", "image", "realtime", "transcribe", "search-preview"];
610
    let mut out: Vec<String> = v["data"]
611
        .as_array()
612
        .or(v["models"].as_array())
613
        .ok_or("the model list had no data")?
614
        .iter()
615
        .filter_map(|m| m["id"].as_str().or(m["name"].as_str()))
616
        .map(|id| id.strip_prefix("models/").unwrap_or(id).to_string())
617
        .filter(|id| !skip.iter().any(|s| id.to_lowercase().contains(s)))
618
        .collect();
619
    out.sort();
620
    out.dedup();
621
    if out.is_empty() {
622
        return Err("the model list was empty".into());
623
    }
624
    Ok(out)
625
}
626
627
/// A tiny exchange, for the settings page's TEST: the start of the answer.
628
pub async fn test(app: &Shared, a: &Assistant) -> Result<String, String> {
629
    let answer =
630
        ask(app, a, "You are being connected to an OxygenHub site. Answer briefly.", "Reply with the single word OK.", CALL_TIMEOUT)
631
            .await?;
632
    Ok(answer.trim().chars().take(80).collect())
633
}
634
635
/// Never let the token through in an error message.
636
fn scrub(text: &str, token: &str) -> String {
637
    if token.is_empty() {
638
        text.to_string()
639
    } else {
640
        text.replace(token, "***")
641
    }
642
}
643
644
#[cfg(test)]
645
mod tests {
646
    use super::*;
647
648
    #[test]
649
    fn tokens_seal_and_open_only_under_their_key() {
650
        let dir = tempfile::tempdir().unwrap();
651
        let k = SecretKey::load_or_create(dir.path()).unwrap();
652
        let sealed = k.seal("sk-very-secret");
653
        assert_ne!(sealed, b"sk-very-secret");
654
        assert!(!String::from_utf8_lossy(&sealed).contains("secret"));
655
        assert_eq!(k.open(&sealed).as_deref(), Some("sk-very-secret"));
656
        assert_eq!(SecretKey::load_or_create(dir.path()).unwrap().open(&sealed).as_deref(), Some("sk-very-secret"), "the key is kept");
657
        let other = SecretKey::load_or_create(tempfile::tempdir().unwrap().path()).unwrap();
658
        assert_eq!(other.open(&sealed), None);
659
        assert_eq!(k.open(&[1, 2, 3]), None);
660
        assert_eq!(k.seal(""), Vec::<u8>::new());
661
        assert_eq!(k.open(&[]).as_deref(), Some(""));
662
    }
663
664
    #[test]
665
    fn answers_say_why_they_stopped() {
666
        let a = |v: serde_json::Value| {
667
            let (text, stop) = answer_of("anthropic-messages", &v);
668
            finish(text, stop)
669
        };
670
        // #201: all of the length spent thinking.
671
        let thought = serde_json::json!({ "stop_reason": "max_tokens", "content": [{ "type": "thinking", "thinking": "" }] });
672
        assert!(a(thought).unwrap_err().contains("ran out of room"));
673
        let cut = serde_json::json!({ "stop_reason": "max_tokens", "content": [{ "type": "thinking", "thinking": "" }, { "type": "text", "text": "Half" }] });
674
        assert!(a(cut).unwrap().starts_with("Half\n\n*(Cut short"));
675
        let declined = serde_json::json!({ "stop_reason": "refusal", "stop_details": { "category": "cyber" }, "content": [] });
676
        assert_eq!(a(declined).unwrap_err(), "it declined to answer (cyber)");
677
        let ok = serde_json::json!({ "stop_reason": "end_turn", "content": [{ "type": "text", "text": "Fine." }] });
678
        assert_eq!(a(ok).unwrap(), "Fine.");
679
        let chat = |v: serde_json::Value| {
680
            let (text, stop) = answer_of("chat-completions", &v);
681
            finish(text, stop)
682
        };
683
        assert!(chat(serde_json::json!({ "choices": [{ "finish_reason": "length", "message": { "content": "" } }] }))
684
            .unwrap_err()
685
            .contains("ran out of room"));
686
        assert_eq!(
687
            chat(serde_json::json!({ "choices": [{ "finish_reason": "content_filter", "message": { "content": "" } }] })).unwrap_err(),
688
            "it declined to answer"
689
        );
690
        assert_eq!(
691
            chat(serde_json::json!({ "choices": [{ "finish_reason": "stop", "message": { "content": "Yes." } }] })).unwrap(),
692
            "Yes."
693
        );
694
    }
695
696
    #[test]
697
    fn headers_parse_or_name_the_bad_line() {
698
        assert_eq!(
699
            parse_headers("X-Org: acme\n\n OpenAI-Beta : v1 ").unwrap(),
700
            vec![("X-Org".into(), "acme".into()), ("OpenAI-Beta".into(), "v1".into())]
701
        );
702
        assert_eq!(parse_headers("no colon").unwrap_err(), "no colon");
703
        assert_eq!(parse_headers("bad name: x").unwrap_err(), "bad name: x");
704
    }
705
}
706
707
// ------------------------------------------------------- reviewing
708
709
/// At most this many assistant calls per repository per hour.
710
pub const CALLS_PER_HOUR: usize = 30;
711
712
/// Who is answering what, and recent calls per repository.
713
#[derive(Default)]
714
pub struct Gate {
715
    busy: std::collections::HashSet<(i64, String)>,
716
    recent: std::collections::HashMap<i64, Vec<i64>>,
717
}
718
719
impl Gate {
720
    /// Take the slot for (proposal, assistant), or say why not.
721
    fn take(&mut self, repo_id: i64, proposal_id: i64, name: &str) -> Result<(), String> {
722
        let key = (proposal_id, name.to_lowercase());
723
        if self.busy.contains(&key) {
724
            return Err(format!("{name} is already answering this proposal"));
725
        }
726
        let now = now_ms();
727
        let calls = self.recent.entry(repo_id).or_default();
728
        calls.retain(|t| now - *t < 3_600_000);
729
        if calls.len() >= CALLS_PER_HOUR {
730
            return Err(format!("this repository has asked assistants {CALLS_PER_HOUR} times in the last hour; try later"));
731
        }
732
        calls.push(now);
733
        self.busy.insert(key);
734
        Ok(())
735
    }
736
737
    fn release(&mut self, proposal_id: i64, name: &str) {
738
        self.busy.remove(&(proposal_id, name.to_lowercase()));
739
    }
740
}
741
742
/// What a reviewer is told about a proposal: its title, description,
743
/// changes and diff, cut to `budget` characters.
744
pub fn review_context(
745
    row: &crate::db::repos::RepoRow,
746
    p: &crate::db::proposals::ProposalRow,
747
    loaded: &crate::web::proposals::Loaded,
748
    budget: usize,
749
) -> String {
750
    let mut head =
751
        format!("# Proposal #{}: {}\n\nRepository {}/{} · into {} · by {}\n\n", p.number, p.title, row.owner, row.name, p.target, p.author);
752
    if !p.body.trim().is_empty() {
753
        head.push_str("## Description\n\n");
754
        head.push_str(p.body.trim());
755
        head.push_str("\n\n");
756
    }
757
    context_text(head, &loaded.changes, &loaded.diffs, budget)
758
}
759
760
/// What a drafter is told: the title so far, the changes and the diff.
761
pub fn draft_context(
762
    row: &crate::db::repos::RepoRow,
763
    title: &str,
764
    target: &str,
765
    changes: &[crate::repos::view::ChangeView],
766
    diffs: &[oxy_core::repo::FileDiff],
767
    budget: usize,
768
) -> String {
769
    let head = format!(
770
        "# Proposal: {}\n\nRepository {}/{} · into {}\n\n",
771
        if title.trim().is_empty() { "(untitled)" } else { title.trim() },
772
        row.owner,
773
        row.name,
774
        target
775
    );
776
    context_text(head, changes, diffs, budget)
777
}
778
779
/// `head`, the list of changes, then the diff as `-`/`+` lines per hunk,
780
/// binary files noted, the diff cut so the whole fits `budget`.
781
fn context_text(mut text: String, changes: &[crate::repos::view::ChangeView], diffs: &[oxy_core::repo::FileDiff], budget: usize) -> String {
782
    text.push_str(&format!("## Changes ({})\n\n", changes.len()));
783
    for c in changes {
784
        text.push_str(&format!("- {} {}\n", c.short, c.title));
785
        let body = c.body.trim();
786
        if !body.is_empty() {
787
            for line in body.lines().take(20) {
788
                text.push_str("  ");
789
                text.push_str(line);
790
                text.push('\n');
791
            }
792
        }
793
    }
794
    text.push_str("\n## Diff\n\n");
795
    let head = text.chars().count();
796
    let mut diff = String::new();
797
    for f in diffs {
798
        diff.push_str(&format!("### {} ({})\n", f.path, f.kind));
799
        if let Some(from) = &f.from {
800
            diff.push_str(&format!("renamed from {from}\n"));
801
        }
802
        let binary = f.hunks.iter().any(|h| h.added.iter().chain(h.removed.iter().map(|(_, l)| l)).any(|l| l.contains(&0)));
803
        if binary {
804
            diff.push_str("(binary)\n\n");
805
            continue;
806
        }
807
        diff.push_str("```\n");
808
        for h in &f.hunks {
809
            diff.push_str(&format!("@@ after line {} @@\n", h.base_line));
810
            for (_, l) in &h.removed {
811
                diff.push('-');
812
                diff.push_str(String::from_utf8_lossy(l).trim_end_matches('\n'));
813
                diff.push('\n');
814
            }
815
            for l in &h.added {
816
                diff.push('+');
817
                diff.push_str(String::from_utf8_lossy(l).trim_end_matches('\n'));
818
                diff.push('\n');
819
            }
820
        }
821
        diff.push_str("```\n\n");
822
    }
823
    let room = budget.saturating_sub(head + 80);
824
    if diff.chars().count() > room {
825
        let cut: String = diff.chars().take(room).collect();
826
        text.push_str(&cut);
827
        text.push_str("\n\n(the diff was cut here to fit the context budget)\n");
828
    } else {
829
        text.push_str(&diff);
830
    }
831
    text
832
}
833
834
/// The system prompt for a review: who the assistant is, what a review
835
/// here is, then the prompts the levels set.
836
pub fn review_system_prompt(a: &Assistant, site: &str, prompts: &str) -> String {
837
    let mut s = format!(
838
        "You are {name}, an AI assistant reviewing a proposal on {site}, an OxygenHub site. \
839
         A person asked for your review; it is recorded as a review comment made via you, under their name. \
840
         You cannot approve: only people approve. Answer in Markdown: start with a one-paragraph summary, \
841
         then findings (correctness first, then clarity), each naming the file and quoting the lines it is about; \
842
         keep praise out unless something is notably well done; say so plainly when you find nothing wrong.",
843
        name = a.name
844
    );
845
    if !prompts.trim().is_empty() {
846
        s.push_str("\n\n");
847
        s.push_str(prompts.trim());
848
    }
849
    s
850
}
851
852
/// Ask `a` to review proposal `p` for `by`: records the request, calls the
853
/// endpoint in the background, and stores the answer as a `comment` review
854
/// via the assistant (or an event saying why none came). `Err` is for the
855
/// person who asked, when nothing was started.
856
pub async fn request_review(
857
    app: &Shared,
858
    row: &crate::db::repos::RepoRow,
859
    p: &crate::db::proposals::ProposalRow,
860
    by: &crate::db::accounts::Account,
861
    a: Assistant,
862
) -> Result<(), String> {
863
    app.assistant_gate.lock().unwrap_or_else(|e| e.into_inner()).take(row.id, p.id, &a.name)?;
864
    let pool = &app.db;
865
    crate::db::proposals::add_event(pool, p.id, &by.username, &format!("asked {} to review", a.name), "").await.ok();
866
    crate::db::accounts::audit(
867
        pool,
868
        Some(by),
869
        "assistant.asked",
870
        &format!("{}/{}", row.owner, row.name),
871
        &format!("{} #{}", a.name, p.number),
872
    )
873
    .await
874
    .ok();
875
    let (app, row, p, by) = (app.clone(), row.clone(), p.clone(), by.clone());
876
    tokio::spawn(async move {
877
        let outcome = answer_review(&app, &row, &p, &a).await;
878
        app.assistant_gate.lock().unwrap_or_else(|e| e.into_inner()).release(p.id, &a.name);
879
        match outcome {
880
            Ok(body) => {
881
                if crate::db::proposals::add_review_via(&app.db, p.id, by.id, &by.username, "comment", &body, &a.name).await.is_ok() {
882
                    crate::db::proposals::add_event(&app.db, p.id, &by.username, &format!("review.comment via {}", a.name), "").await.ok();
883
                    crate::web::proposals::reviewed(&app, &row, &p, &by, "comment").await;
884
                    crate::hooks::emit(
885
                        &app,
886
                        &row,
887
                        "review",
888
                        &by.username,
889
                        serde_json::json!({ "number": p.number, "verdict": "comment", "via": a.name }),
890
                    )
891
                    .await;
892
                }
893
            }
894
            Err(why) => {
895
                crate::db::proposals::add_event(&app.db, p.id, &by.username, &format!("{} did not answer", a.name), &why).await.ok();
896
            }
897
        }
898
    });
899
    Ok(())
900
}
901
902
async fn answer_review(
903
    app: &Shared,
904
    row: &crate::db::repos::RepoRow,
905
    p: &crate::db::proposals::ProposalRow,
906
    a: &Assistant,
907
) -> Result<String, String> {
908
    let repo = crate::repos::names::RepoName::new(&row.owner, &row.name).ok_or("bad repository name")?;
909
    let loaded = crate::web::proposals::load(app, &repo, p.clone()).await.map_err(|e| format!("the proposal could not be loaded: {e}"))?;
910
    let budget = usize::try_from(a.budget).unwrap_or(DEFAULT_BUDGET as usize).max(1_000);
911
    let prompts = prompt_for_repo(&app.db, row).await.map_err(|e| e.to_string())?;
912
    let system = review_system_prompt(a, &app.config.site.name, &prompts);
913
    let user = review_context(row, p, &loaded, budget.saturating_sub(system.chars().count()));
914
    ask(app, a, &system, &user, BACKGROUND_TIMEOUT).await
915
}
916
917
// -------------------------------------------------------- summaries
918
919
/// The system prompt for a summary posted as a comment (oxygenhub#190).
920
pub fn summary_system_prompt(a: &Assistant, site: &str, prompts: &str) -> String {
921
    let mut s = format!(
922
        "You are {name}, an AI assistant working on {site}, an OxygenHub site, with a next-generation version control system \
923
         called Oxygen (its command is oxy). A proposal is Oxygen's equivalent of a merge request or pull request. \
924
         Please summarise the proposal and comment on it. A person asked for this; it is posted as a comment made via you, \
925
         under their name. Answer in Markdown: a short summary of what the proposal changes and why, then anything a reviewer \
926
         should look at closely; keep to what the description and the diff show.",
927
        name = a.name
928
    );
929
    if !prompts.trim().is_empty() {
930
        s.push_str("\n\n");
931
        s.push_str(prompts.trim());
932
    }
933
    s
934
}
935
936
/// Ask `a` to summarise proposal `p` for `by`: like [`request_review`],
937
/// but the answer is a conversation comment via the assistant.
938
pub async fn request_summary(
939
    app: &Shared,
940
    row: &crate::db::repos::RepoRow,
941
    p: &crate::db::proposals::ProposalRow,
942
    by: &crate::db::accounts::Account,
943
    a: Assistant,
944
) -> Result<(), String> {
945
    app.assistant_gate.lock().unwrap_or_else(|e| e.into_inner()).take(row.id, p.id, &a.name)?;
946
    let pool = &app.db;
947
    crate::db::proposals::add_event(pool, p.id, &by.username, &format!("asked {} for a summary", a.name), "").await.ok();
948
    crate::db::accounts::audit(
949
        pool,
950
        Some(by),
951
        "assistant.asked",
952
        &format!("{}/{}", row.owner, row.name),
953
        &format!("{} #{} summary", a.name, p.number),
954
    )
955
    .await
956
    .ok();
957
    let (app, row, p, by) = (app.clone(), row.clone(), p.clone(), by.clone());
958
    tokio::spawn(async move {
959
        let outcome = answer_summary(&app, &row, &p, &a).await;
960
        app.assistant_gate.lock().unwrap_or_else(|e| e.into_inner()).release(p.id, &a.name);
961
        match outcome {
962
            Ok(body) => {
963
                crate::db::proposals::add_comment_via(&app.db, p.id, by.id, &by.username, &body, &a.name).await.ok();
964
            }
965
            Err(why) => {
966
                crate::db::proposals::add_event(&app.db, p.id, &by.username, &format!("{} did not answer", a.name), &why).await.ok();
967
            }
968
        }
969
    });
970
    Ok(())
971
}
972
973
async fn answer_summary(
974
    app: &Shared,
975
    row: &crate::db::repos::RepoRow,
976
    p: &crate::db::proposals::ProposalRow,
977
    a: &Assistant,
978
) -> Result<String, String> {
979
    let repo = crate::repos::names::RepoName::new(&row.owner, &row.name).ok_or("bad repository name")?;
980
    let loaded = crate::web::proposals::load(app, &repo, p.clone()).await.map_err(|e| format!("the proposal could not be loaded: {e}"))?;
981
    let budget = usize::try_from(a.budget).unwrap_or(DEFAULT_BUDGET as usize).max(1_000);
982
    let prompts = prompt_for_repo(&app.db, row).await.map_err(|e| e.to_string())?;
983
    let system = summary_system_prompt(a, &app.config.site.name, &prompts);
984
    let user = review_context(row, p, &loaded, budget.saturating_sub(system.chars().count()));
985
    ask(app, a, &system, &user, BACKGROUND_TIMEOUT).await
986
}
987
988
// -------------------------------------------------------- drafting
989
990
/// The system prompt for drafting a proposal's description.
991
pub fn draft_system_prompt(a: &Assistant, site: &str, prompts: &str) -> String {
992
    let mut s = format!(
993
        "You are {name}, an AI assistant drafting the description of a proposal on {site}, an OxygenHub site, \
994
         for the person who will open or edit it; they read and change your draft before saving, and the draft is theirs. \
995
         Write Markdown for the description field only: no title line, no greeting, no sign-off. \
996
         Say what the changes do and why, what a reviewer should look at, and how to verify them; \
997
         keep to what the diff shows, and keep it short.",
998
        name = a.name
999
    );
1000
    if !prompts.trim().is_empty() {
1001
        s.push_str("\n\n");
1002
        s.push_str(prompts.trim());
1003
    }
1004
    s
1005
}
1006
1007
/// What a description is drafted from.
1008
pub struct Draft<'a> {
1009
    /// The proposal's id, or minus the repository's for one not yet opened:
1010
    /// one draft per assistant at a time for it.
1011
    pub slot: i64,
1012
    pub title: &'a str,
1013
    pub target: &'a str,
1014
    pub changes: &'a [crate::repos::view::ChangeView],
1015
    pub diffs: &'a [oxy_core::repo::FileDiff],
1016
}
1017
1018
/// Ask `a` for a description, now; the answer goes into the editor,
1019
/// nothing is saved. `Err` is for the person.
1020
pub async fn draft_description(
1021
    app: &Shared,
1022
    row: &crate::db::repos::RepoRow,
1023
    by: &crate::db::accounts::Account,
1024
    a: &Assistant,
1025
    d: Draft<'_>,
1026
) -> Result<String, String> {
1027
    let Draft { slot, title, target, changes, diffs } = d;
1028
    app.assistant_gate.lock().unwrap_or_else(|e| e.into_inner()).take(row.id, slot, &a.name)?;
1029
    crate::db::accounts::audit(&app.db, Some(by), "assistant.drafted", &format!("{}/{}", row.owner, row.name), &a.name).await.ok();
1030
    let budget = usize::try_from(a.budget).unwrap_or(DEFAULT_BUDGET as usize).max(1_000);
1031
    let prompts = prompt_for_repo(&app.db, row).await.map_err(|e| e.to_string());
1032
    let outcome = match prompts {
1033
        Ok(prompts) => {
1034
            let system = draft_system_prompt(a, &app.config.site.name, &prompts);
1035
            let user = draft_context(row, title, target, changes, diffs, budget.saturating_sub(system.chars().count()));
1036
            ask(app, a, &system, &user, CALL_TIMEOUT).await
1037
        }
1038
        Err(e) => Err(e),
1039
    };
1040
    app.assistant_gate.lock().unwrap_or_else(|e| e.into_inner()).release(slot, &a.name);
1041
    outcome.map(|t| t.trim().to_string())
1042
}