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