crates/hub/src/assistants.rs at e71846aa07b9 · andymccall/oxygenhub

# e71846aa07b9 oxygenhub / crates / hub / src / assistants.rs

19.1 KiB · Rust · last changed AI assistants: settings per site, organisation and repository 1 month ago source blame history raw

//! AI assistants (hub spec §13.4, #97): chat-completion endpoints the hub
//! may ask to review a proposal or draft its description. Assistants are
//! kept at three levels, site, organisation and repository, and combined
//! downwards: a repository sees its organisation's and the site's beside
//! its own, the innermost of a name wins, and one disabled at an outer
//! level is off for everything inside it. Each level adds a prompt; the
//! prompt in force for a repository is the site's, the organisation's and
//! the repository's own, each under a heading. Tokens are sealed with
//! ChaCha20-Poly1305 under `keys/assistants.key`, shown masked, and never
//! leave the server.

use crate::db::{now_ms, Pool};
use crate::Shared;
use rand::RngCore;
use ring::aead;
use sqlx::FromRow;
use std::path::Path;
use std::time::Duration;

/// The request shapes an assistant may speak.
pub const APIS: [&str; 2] = ["chat-completions", "anthropic-messages"];
/// Characters of context (prompt, description, diffs) a call may carry
/// unless the assistant says otherwise.
pub const DEFAULT_BUDGET: i64 = 120_000;
/// How long one call may take.
pub const CALL_TIMEOUT: Duration = Duration::from_secs(60);
/// The most an answer is kept.
const MAX_ANSWER: usize = 64 * 1024;

/// Which level an assistant or a prompt belongs to.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Scope {
    Site,
    Org(i64),
    Repo(i64),
}

impl Scope {
    pub fn kind(self) -> &'static str {
        match self {
            Scope::Site => "site",
            Scope::Org(_) => "org",
            Scope::Repo(_) => "repo",
        }
    }

    pub fn id(self) -> i64 {
        match self {
            Scope::Site => 0,
            Scope::Org(id) | Scope::Repo(id) => id,
        }
    }
}

#[derive(Clone, Debug, FromRow)]
pub struct Assistant {
    pub id: i64,
    pub scope: String,
    pub scope_id: i64,
    pub name: String,
    pub endpoint: String,
    /// One of [`APIS`].
    pub api: String,
    pub model: String,
    /// Extra request headers, one `Name: value` per line.
    pub headers: String,
    /// The sealed token; empty when the endpoint needs none.
    pub token: Vec<u8>,
    pub enabled: bool,
    pub budget: i64,
    pub created_ms: i64,
    pub updated_ms: i64,
}

impl Assistant {
    pub fn scope(&self) -> Scope {
        match self.scope.as_str() {
            "org" => Scope::Org(self.scope_id),
            "repo" => Scope::Repo(self.scope_id),
            _ => Scope::Site,
        }
    }

    pub fn has_token(&self) -> bool {
        !self.token.is_empty()
    }

    /// The extra headers as pairs; malformed lines are skipped.
    pub fn header_list(&self) -> Vec<(String, String)> {
        parse_headers(&self.headers).unwrap_or_default()
    }
}

/// What a settings form supplies for an assistant.
#[derive(Clone, Debug, Default)]
pub struct Fields {
    pub name: String,
    pub endpoint: String,
    pub api: String,
    pub model: String,
    pub headers: String,
    pub enabled: bool,
    pub budget: i64,
}

/// `Name: value` lines to pairs: `Err(line)` names the first bad one.
pub fn parse_headers(text: &str) -> Result<Vec<(String, String)>, String> {
    let mut out = vec![];
    for line in text.lines().map(str::trim).filter(|l| !l.is_empty()) {
        let Some((name, value)) = line.split_once(':') else { return Err(line.to_string()) };
        let (name, value) = (name.trim(), value.trim());
        let name_ok = !name.is_empty() && name.chars().all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_');
        if !name_ok || value.chars().any(|c| c.is_control()) {
            return Err(line.to_string());
        }
        out.push((name.to_string(), value.to_string()));
    }
    Ok(out)
}

// -------------------------------------------------------------- sealing

/// The key that seals tokens, kept in `<data>/keys/assistants.key`.
#[derive(Clone)]
pub struct SecretKey([u8; 32]);

impl SecretKey {
    pub fn load_or_create(keys: &Path) -> anyhow::Result<SecretKey> {
        std::fs::create_dir_all(keys)?;
        let path = keys.join("assistants.key");
        if let Ok(bytes) = std::fs::read(&path) {
            if bytes.len() == 32 {
                return Ok(SecretKey(bytes.try_into().unwrap()));
            }
        }
        let mut k = [0u8; 32];
        rand::rngs::OsRng.fill_bytes(&mut k);
        std::fs::write(&path, k)?;
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600))?;
        }
        Ok(SecretKey(k))
    }

    fn key(&self) -> aead::LessSafeKey {
        aead::LessSafeKey::new(aead::UnboundKey::new(&aead::CHACHA20_POLY1305, &self.0).expect("key length"))
    }

    /// Seal a token: a random nonce, then the ciphertext and tag. An
    /// empty token seals to nothing, meaning "no token".
    pub fn seal(&self, plain: &str) -> Vec<u8> {
        if plain.is_empty() {
            return vec![];
        }
        let mut nonce = [0u8; 12];
        rand::rngs::OsRng.fill_bytes(&mut nonce);
        let mut out = plain.as_bytes().to_vec();
        self.key().seal_in_place_append_tag(aead::Nonce::assume_unique_for_key(nonce), aead::Aad::empty(), &mut out).expect("seal");
        let mut sealed = nonce.to_vec();
        sealed.extend(out);
        sealed
    }

    /// The token back, or `None` when the bytes were sealed under another
    /// key or damaged. Nothing opens to an empty token.
    pub fn open(&self, sealed: &[u8]) -> Option<String> {
        if sealed.is_empty() {
            return Some(String::new());
        }
        let (nonce, rest) = sealed.split_at_checked(12)?;
        let mut buf = rest.to_vec();
        let nonce = aead::Nonce::try_assume_unique_for_key(nonce).ok()?;
        let plain = self.key().open_in_place(nonce, aead::Aad::empty(), &mut buf).ok()?;
        String::from_utf8(plain.to_vec()).ok()
    }
}

// ------------------------------------------------------------- storage

pub async fn list(pool: &Pool, scope: Scope) -> sqlx::Result<Vec<Assistant>> {
    sqlx::query_as("SELECT * FROM assistants WHERE scope = ? AND scope_id = ? ORDER BY name")
        .bind(scope.kind())
        .bind(scope.id())
        .fetch_all(pool)
        .await
}

pub async fn get(pool: &Pool, id: i64) -> sqlx::Result<Option<Assistant>> {
    sqlx::query_as("SELECT * FROM assistants WHERE id = ?").bind(id).fetch_optional(pool).await
}

/// Add an assistant; `None` when the level already has one of that name.
pub async fn create(pool: &Pool, scope: Scope, f: &Fields, token: &[u8]) -> sqlx::Result<Option<Assistant>> {
    let now = now_ms();
    let r = sqlx::query_as(
        "INSERT INTO assistants (scope, scope_id, name, endpoint, api, model, headers, token, enabled, budget, created_ms, updated_ms) \
         VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT DO NOTHING RETURNING *",
    )
    .bind(scope.kind())
    .bind(scope.id())
    .bind(&f.name)
    .bind(&f.endpoint)
    .bind(&f.api)
    .bind(&f.model)
    .bind(&f.headers)
    .bind(token)
    .bind(f.enabled)
    .bind(f.budget)
    .bind(now)
    .bind(now)
    .fetch_optional(pool)
    .await?;
    Ok(r)
}

/// Change an assistant's details; `token` is `None` to keep the stored
/// one. False when the name is taken at that level.
pub async fn update(pool: &Pool, id: i64, f: &Fields, token: Option<&[u8]>) -> sqlx::Result<bool> {
    let r = match token {
        Some(t) => sqlx::query(
            "UPDATE assistants SET name = ?, endpoint = ?, api = ?, model = ?, headers = ?, enabled = ?, budget = ?, token = ?, updated_ms = ? \
             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 <> ?)",
        )
        .bind(&f.name)
        .bind(&f.endpoint)
        .bind(&f.api)
        .bind(&f.model)
        .bind(&f.headers)
        .bind(f.enabled)
        .bind(f.budget)
        .bind(t)
        .bind(now_ms())
        .bind(id)
        .bind(&f.name)
        .bind(id)
        .execute(pool)
        .await?,
        None => sqlx::query(
            "UPDATE assistants SET name = ?, endpoint = ?, api = ?, model = ?, headers = ?, enabled = ?, budget = ?, updated_ms = ? \
             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 <> ?)",
        )
        .bind(&f.name)
        .bind(&f.endpoint)
        .bind(&f.api)
        .bind(&f.model)
        .bind(&f.headers)
        .bind(f.enabled)
        .bind(f.budget)
        .bind(now_ms())
        .bind(id)
        .bind(&f.name)
        .bind(id)
        .execute(pool)
        .await?,
    };
    Ok(r.rows_affected() == 1)
}

pub async fn set_enabled(pool: &Pool, id: i64, enabled: bool) -> sqlx::Result<()> {
    sqlx::query("UPDATE assistants SET enabled = ?, updated_ms = ? WHERE id = ?")
        .bind(enabled)
        .bind(now_ms())
        .bind(id)
        .execute(pool)
        .await?;
    Ok(())
}

pub async fn delete(pool: &Pool, id: i64) -> sqlx::Result<()> {
    sqlx::query("DELETE FROM assistants WHERE id = ?").bind(id).execute(pool).await?;
    Ok(())
}

/// Forget a level's assistants and prompt, when the organisation or
/// repository goes.
pub async fn delete_scope(pool: &Pool, scope: Scope) -> sqlx::Result<()> {
    sqlx::query("DELETE FROM assistants WHERE scope = ? AND scope_id = ?").bind(scope.kind()).bind(scope.id()).execute(pool).await?;
    sqlx::query("DELETE FROM assistant_prompts WHERE scope = ? AND scope_id = ?").bind(scope.kind()).bind(scope.id()).execute(pool).await?;
    Ok(())
}

/// The prompt a level adds, empty when none.
pub async fn prompt(pool: &Pool, scope: Scope) -> sqlx::Result<String> {
    Ok(sqlx::query_scalar("SELECT prompt FROM assistant_prompts WHERE scope = ? AND scope_id = ?")
        .bind(scope.kind())
        .bind(scope.id())
        .fetch_optional(pool)
        .await?
        .unwrap_or_default())
}

pub async fn set_prompt(pool: &Pool, scope: Scope, text: &str) -> sqlx::Result<()> {
    if text.trim().is_empty() {
        sqlx::query("DELETE FROM assistant_prompts WHERE scope = ? AND scope_id = ?")
            .bind(scope.kind())
            .bind(scope.id())
            .execute(pool)
            .await?;
    } else {
        sqlx::query(
            "INSERT INTO assistant_prompts (scope, scope_id, prompt, updated_ms) VALUES (?, ?, ?, ?) \
             ON CONFLICT (scope, scope_id) DO UPDATE SET prompt = excluded.prompt, updated_ms = excluded.updated_ms",
        )
        .bind(scope.kind())
        .bind(scope.id())
        .bind(text.trim())
        .bind(now_ms())
        .execute(pool)
        .await?;
    }
    Ok(())
}

// ----------------------------------------------------- what applies where

/// The levels above a scope, outermost first: the site, then the
/// organisation when the owner is one.
pub async fn levels_above(pool: &Pool, scope: Scope, owner: &str) -> sqlx::Result<Vec<(Scope, String)>> {
    let mut out = vec![];
    match scope {
        Scope::Site => {}
        Scope::Org(_) => out.push((Scope::Site, "the site".to_string())),
        Scope::Repo(_) => {
            out.push((Scope::Site, "the site".to_string()));
            if let Some(org) = crate::db::accounts::org_by_name(pool, owner).await? {
                out.push((Scope::Org(org.id), format!("the {} organisation", org.name)));
            }
        }
    }
    Ok(out)
}

/// An assistant seen from a level: where it comes from, and whether it is
/// on once outer levels have had their say.
#[derive(Clone, Debug)]
pub struct Available {
    pub assistant: Assistant,
    /// `the site`, `the acme organisation`, or empty for the level's own.
    pub from: String,
    /// Off because a level above disabled an assistant of this name.
    pub off_above: bool,
}

impl Available {
    pub fn on(&self) -> bool {
        self.assistant.enabled && !self.off_above
    }
}

/// Every assistant a level sees, outermost first, one per name: the
/// innermost definition of a name wins, and a name disabled at an outer
/// level is off inside it whatever the inner level says.
pub async fn available(pool: &Pool, scope: Scope, owner: &str) -> sqlx::Result<Vec<Available>> {
    let mut levels = levels_above(pool, scope, owner).await?;
    levels.push((scope, String::new()));
    let mut out: Vec<Available> = vec![];
    let mut off: Vec<String> = vec![];
    for (level, from) in levels {
        for a in list(pool, level).await? {
            let off_above = off.iter().any(|n| n.eq_ignore_ascii_case(&a.name));
            if !a.enabled && !off.iter().any(|n| n.eq_ignore_ascii_case(&a.name)) {
                off.push(a.name.clone());
            }
            out.retain(|o| !o.assistant.name.eq_ignore_ascii_case(&a.name));
            out.push(Available { assistant: a, from: from.clone(), off_above });
        }
    }
    Ok(out)
}

/// The assistants a repository may ask, by name.
pub async fn enabled_for_repo(pool: &Pool, repo: &crate::db::repos::RepoRow) -> sqlx::Result<Vec<Assistant>> {
    Ok(available(pool, Scope::Repo(repo.id), &repo.owner).await?.into_iter().filter(Available::on).map(|a| a.assistant).collect())
}

/// The prompts a level inherits, outermost first: (where from, text),
/// only the levels that set one.
pub async fn inherited_prompts(pool: &Pool, scope: Scope, owner: &str) -> sqlx::Result<Vec<(String, String)>> {
    let mut out = vec![];
    for (level, from) in levels_above(pool, scope, owner).await? {
        let text = prompt(pool, level).await?;
        if !text.is_empty() {
            out.push((from, text));
        }
    }
    Ok(out)
}

/// The prompt in force for a repository: each level's text under a
/// heading, outermost first, so a project adds to what the site says.
pub async fn prompt_for_repo(pool: &Pool, repo: &crate::db::repos::RepoRow) -> sqlx::Result<String> {
    let scope = Scope::Repo(repo.id);
    let mut parts = inherited_prompts(pool, scope, &repo.owner).await?;
    let own = prompt(pool, scope).await?;
    if !own.is_empty() {
        parts.push(("this repository".to_string(), own));
    }
    Ok(parts.into_iter().map(|(from, text)| format!("## Instructions from {from}\n\n{text}\n")).collect::<Vec<_>>().join("\n"))
}

// -------------------------------------------------------------- calling

/// Ask an assistant: the system prompt and the request, the answer's
/// text. Errors are for the person who asked: a status and the first
/// line of what the endpoint said, or why no answer came.
pub async fn ask(app: &Shared, a: &Assistant, system: &str, user: &str) -> Result<String, String> {
    let Some(token) = app.assistant_key.open(&a.token) else {
        return Err("its token cannot be read: set it again".into());
    };
    let client = reqwest::Client::builder()
        .timeout(CALL_TIMEOUT)
        .user_agent(format!("oxygenhub/{}", env!("CARGO_PKG_VERSION")))
        .build()
        .map_err(|e| e.to_string())?;
    let mut req = client.post(&a.endpoint).header("content-type", "application/json");
    let body = match a.api.as_str() {
        "anthropic-messages" => {
            if !token.is_empty() {
                req = req.header("x-api-key", &token);
            }
            req = req.header("anthropic-version", "2023-06-01");
            serde_json::json!({
                "model": a.model,
                "max_tokens": 4096,
                "system": system,
                "messages": [{ "role": "user", "content": user }],
            })
        }
        _ => {
            if !token.is_empty() {
                req = req.header("authorization", format!("Bearer {token}"));
            }
            serde_json::json!({
                "model": a.model,
                "messages": [
                    { "role": "system", "content": system },
                    { "role": "user", "content": user },
                ],
            })
        }
    };
    for (name, value) in a.header_list() {
        req = req.header(name, value);
    }
    let r = req.json(&body).send().await.map_err(|e| {
        if e.is_timeout() {
            "no answer within 60 seconds".to_string()
        } else {
            scrub(&e.to_string(), &token)
        }
    })?;
    let status = r.status();
    let text: String = r.text().await.unwrap_or_default().chars().take(MAX_ANSWER).collect();
    if !status.is_success() {
        let why = serde_json::from_str::<serde_json::Value>(&text)
            .ok()
            .and_then(|v| v["error"]["message"].as_str().or(v["error"].as_str()).or(v["message"].as_str()).map(str::to_string))
            .unwrap_or_else(|| text.lines().next().unwrap_or("").chars().take(160).collect());
        return Err(scrub(&format!("{} {}", status.as_u16(), why.trim()), &token));
    }
    let v: serde_json::Value = serde_json::from_str(&text).map_err(|_| "the answer was not JSON".to_string())?;
    let answer = match a.api.as_str() {
        "anthropic-messages" => v["content"]
            .as_array()
            .map(|parts| parts.iter().filter_map(|p| p["text"].as_str()).collect::<Vec<_>>().join(""))
            .unwrap_or_default(),
        _ => match &v["choices"][0]["message"]["content"] {
            serde_json::Value::String(s) => s.clone(),
            serde_json::Value::Array(parts) => parts.iter().filter_map(|p| p["text"].as_str()).collect::<Vec<_>>().join(""),
            _ => String::new(),
        },
    };
    if answer.trim().is_empty() {
        return Err("the answer carried no text".into());
    }
    Ok(answer)
}

/// A tiny exchange, for the settings page's TEST: the start of the answer.
pub async fn test(app: &Shared, a: &Assistant) -> Result<String, String> {
    let answer = ask(app, a, "You are being connected to an OxygenHub site. Answer briefly.", "Reply with the single word OK.").await?;
    Ok(answer.trim().chars().take(80).collect())
}

/// Never let the token through in an error message.
fn scrub(text: &str, token: &str) -> String {
    if token.is_empty() {
        text.to_string()
    } else {
        text.replace(token, "***")
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn tokens_seal_and_open_only_under_their_key() {
        let dir = tempfile::tempdir().unwrap();
        let k = SecretKey::load_or_create(dir.path()).unwrap();
        let sealed = k.seal("sk-very-secret");
        assert_ne!(sealed, b"sk-very-secret");
        assert!(!String::from_utf8_lossy(&sealed).contains("secret"));
        assert_eq!(k.open(&sealed).as_deref(), Some("sk-very-secret"));
        assert_eq!(SecretKey::load_or_create(dir.path()).unwrap().open(&sealed).as_deref(), Some("sk-very-secret"), "the key is kept");
        let other = SecretKey::load_or_create(tempfile::tempdir().unwrap().path()).unwrap();
        assert_eq!(other.open(&sealed), None);
        assert_eq!(k.open(&[1, 2, 3]), None);
        assert_eq!(k.seal(""), Vec::<u8>::new());
        assert_eq!(k.open(&[]).as_deref(), Some(""));
    }

    #[test]
    fn headers_parse_or_name_the_bad_line() {
        assert_eq!(
            parse_headers("X-Org: acme\n\n OpenAI-Beta : v1 ").unwrap(),
            vec![("X-Org".into(), "acme".into()), ("OpenAI-Beta".into(), "v1".into())]
        );
        assert_eq!(parse_headers("no colon").unwrap_err(), "no colon");
        assert_eq!(parse_headers("bad name: x").unwrap_err(), "bad name: x");
    }
}

`/`	Search
`g` `c`	Code
`g` `h`	History
`g` `b`	Bookmarks
`g` `t`	Trust
`y`	Permalink to this state (files and trees)
`?`	This sheet

e71846aa Andy McCall 1 month ago ↶	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	/// Which level an assistant or a prompt belongs to.
	31	#[derive(Clone, Copy, Debug, PartialEq, Eq)]
	32	pub enum Scope {
	33	Site,
	34	Org(i64),
	35	Repo(i64),
	36	}
	37
	38	impl Scope {
	39	pub fn kind(self) -> &'static str {
	40	match self {
	41	Scope::Site => "site",
	42	Scope::Org(_) => "org",
	43	Scope::Repo(_) => "repo",
	44	}
	45	}
	46
	47	pub fn id(self) -> i64 {
	48	match self {
	49	Scope::Site => 0,
	50	Scope::Org(id) \| Scope::Repo(id) => id,
	51	}
	52	}
	53	}
	54
	55	#[derive(Clone, Debug, FromRow)]
	56	pub struct Assistant {
	57	pub id: i64,
	58	pub scope: String,
	59	pub scope_id: i64,
	60	pub name: String,
	61	pub endpoint: String,
	62	/// One of [`APIS`].
	63	pub api: String,
	64	pub model: String,
65	/// Extra request headers, one `Name: value` per line.
66	pub headers: String,
67	/// The sealed token; empty when the endpoint needs none.
68	pub token: Vec<u8>,
69	pub enabled: bool,
70	pub budget: i64,
71	pub created_ms: i64,
72	pub updated_ms: i64,
73	}
74
75	impl Assistant {
76	pub fn scope(&self) -> Scope {
77	match self.scope.as_str() {
78	"org" => Scope::Org(self.scope_id),
79	"repo" => Scope::Repo(self.scope_id),
80	_ => Scope::Site,
81	}
82	}
83
84	pub fn has_token(&self) -> bool {
85	!self.token.is_empty()
86	}
87
88	/// The extra headers as pairs; malformed lines are skipped.
89	pub fn header_list(&self) -> Vec<(String, String)> {
90	parse_headers(&self.headers).unwrap_or_default()
91	}
92	}
93
94	/// What a settings form supplies for an assistant.
95	#[derive(Clone, Debug, Default)]
96	pub struct Fields {
97	pub name: String,
98	pub endpoint: String,
99	pub api: String,
100	pub model: String,
101	pub headers: String,
102	pub enabled: bool,
103	pub budget: i64,
104	}
105
106	/// `Name: value` lines to pairs: `Err(line)` names the first bad one.
107	pub fn parse_headers(text: &str) -> Result<Vec<(String, String)>, String> {
108	let mut out = vec![];
109	for line in text.lines().map(str::trim).filter(\|l\| !l.is_empty()) {
110	let Some((name, value)) = line.split_once(':') else { return Err(line.to_string()) };
111	let (name, value) = (name.trim(), value.trim());
112	let name_ok = !name.is_empty() && name.chars().all(\|c\| c.is_ascii_alphanumeric() \|\| c == '-' \|\| c == '_');
113	if !name_ok \|\| value.chars().any(\|c\| c.is_control()) {
114	return Err(line.to_string());
115	}
116	out.push((name.to_string(), value.to_string()));
117	}
118	Ok(out)
119	}
120
121	// -------------------------------------------------------------- sealing
122
123	/// The key that seals tokens, kept in `<data>/keys/assistants.key`.
124	#[derive(Clone)]
125	pub struct SecretKey([u8; 32]);
126
127	impl SecretKey {
128	pub fn load_or_create(keys: &Path) -> anyhow::Result<SecretKey> {
129	std::fs::create_dir_all(keys)?;
130	let path = keys.join("assistants.key");
131	if let Ok(bytes) = std::fs::read(&path) {
132	if bytes.len() == 32 {
133	return Ok(SecretKey(bytes.try_into().unwrap()));
134	}
135	}
136	let mut k = [0u8; 32];
137	rand::rngs::OsRng.fill_bytes(&mut k);
138	std::fs::write(&path, k)?;
139	#[cfg(unix)]
140	{
141	use std::os::unix::fs::PermissionsExt;
142	std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600))?;
143	}
144	Ok(SecretKey(k))
145	}
146
147	fn key(&self) -> aead::LessSafeKey {
148	aead::LessSafeKey::new(aead::UnboundKey::new(&aead::CHACHA20_POLY1305, &self.0).expect("key length"))
149	}
150
151	/// Seal a token: a random nonce, then the ciphertext and tag. An
152	/// empty token seals to nothing, meaning "no token".
153	pub fn seal(&self, plain: &str) -> Vec<u8> {
154	if plain.is_empty() {
155	return vec![];
156	}
157	let mut nonce = [0u8; 12];
158	rand::rngs::OsRng.fill_bytes(&mut nonce);
159	let mut out = plain.as_bytes().to_vec();
160	self.key().seal_in_place_append_tag(aead::Nonce::assume_unique_for_key(nonce), aead::Aad::empty(), &mut out).expect("seal");
161	let mut sealed = nonce.to_vec();
162	sealed.extend(out);
163	sealed
164	}
165
166	/// The token back, or `None` when the bytes were sealed under another
167	/// key or damaged. Nothing opens to an empty token.
168	pub fn open(&self, sealed: &[u8]) -> Option<String> {
169	if sealed.is_empty() {
170	return Some(String::new());
171	}
172	let (nonce, rest) = sealed.split_at_checked(12)?;
173	let mut buf = rest.to_vec();
174	let nonce = aead::Nonce::try_assume_unique_for_key(nonce).ok()?;
175	let plain = self.key().open_in_place(nonce, aead::Aad::empty(), &mut buf).ok()?;
176	String::from_utf8(plain.to_vec()).ok()
177	}
178	}
179
180	// ------------------------------------------------------------- storage
181
182	pub async fn list(pool: &Pool, scope: Scope) -> sqlx::Result<Vec<Assistant>> {
183	sqlx::query_as("SELECT * FROM assistants WHERE scope = ? AND scope_id = ? ORDER BY name")
184	.bind(scope.kind())
185	.bind(scope.id())
186	.fetch_all(pool)
187	.await
188	}
189
190	pub async fn get(pool: &Pool, id: i64) -> sqlx::Result<Option<Assistant>> {
191	sqlx::query_as("SELECT * FROM assistants WHERE id = ?").bind(id).fetch_optional(pool).await
192	}
193
194	/// Add an assistant; `None` when the level already has one of that name.
195	pub async fn create(pool: &Pool, scope: Scope, f: &Fields, token: &[u8]) -> sqlx::Result<Option<Assistant>> {
196	let now = now_ms();
197	let r = sqlx::query_as(
198	"INSERT INTO assistants (scope, scope_id, name, endpoint, api, model, headers, token, enabled, budget, created_ms, updated_ms) \
199	VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT DO NOTHING RETURNING *",
200	)
201	.bind(scope.kind())
202	.bind(scope.id())
203	.bind(&f.name)
204	.bind(&f.endpoint)
205	.bind(&f.api)
206	.bind(&f.model)
207	.bind(&f.headers)
208	.bind(token)
209	.bind(f.enabled)
210	.bind(f.budget)
211	.bind(now)
212	.bind(now)
213	.fetch_optional(pool)
214	.await?;
215	Ok(r)
216	}
217
218	/// Change an assistant's details; `token` is `None` to keep the stored
219	/// one. False when the name is taken at that level.
220	pub async fn update(pool: &Pool, id: i64, f: &Fields, token: Option<&[u8]>) -> sqlx::Result<bool> {
221	let r = match token {
222	Some(t) => sqlx::query(
223	"UPDATE assistants SET name = ?, endpoint = ?, api = ?, model = ?, headers = ?, enabled = ?, budget = ?, token = ?, updated_ms = ? \
224	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 <> ?)",
225	)
226	.bind(&f.name)
227	.bind(&f.endpoint)
228	.bind(&f.api)
229	.bind(&f.model)
230	.bind(&f.headers)
231	.bind(f.enabled)
232	.bind(f.budget)
233	.bind(t)
234	.bind(now_ms())
235	.bind(id)
236	.bind(&f.name)
237	.bind(id)
238	.execute(pool)
239	.await?,
240	None => sqlx::query(
241	"UPDATE assistants SET name = ?, endpoint = ?, api = ?, model = ?, headers = ?, enabled = ?, budget = ?, updated_ms = ? \
242	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 <> ?)",
243	)
244	.bind(&f.name)
245	.bind(&f.endpoint)
246	.bind(&f.api)
247	.bind(&f.model)
248	.bind(&f.headers)
249	.bind(f.enabled)
250	.bind(f.budget)
251	.bind(now_ms())
252	.bind(id)
253	.bind(&f.name)
254	.bind(id)
255	.execute(pool)
256	.await?,
257	};
258	Ok(r.rows_affected() == 1)
259	}
260
261	pub async fn set_enabled(pool: &Pool, id: i64, enabled: bool) -> sqlx::Result<()> {
262	sqlx::query("UPDATE assistants SET enabled = ?, updated_ms = ? WHERE id = ?")
263	.bind(enabled)
264	.bind(now_ms())
265	.bind(id)
266	.execute(pool)
267	.await?;
268	Ok(())
269	}
270
271	pub async fn delete(pool: &Pool, id: i64) -> sqlx::Result<()> {
272	sqlx::query("DELETE FROM assistants WHERE id = ?").bind(id).execute(pool).await?;
273	Ok(())
274	}
275
276	/// Forget a level's assistants and prompt, when the organisation or
277	/// repository goes.
278	pub async fn delete_scope(pool: &Pool, scope: Scope) -> sqlx::Result<()> {
279	sqlx::query("DELETE FROM assistants WHERE scope = ? AND scope_id = ?").bind(scope.kind()).bind(scope.id()).execute(pool).await?;
280	sqlx::query("DELETE FROM assistant_prompts WHERE scope = ? AND scope_id = ?").bind(scope.kind()).bind(scope.id()).execute(pool).await?;
281	Ok(())
282	}
283
284	/// The prompt a level adds, empty when none.
285	pub async fn prompt(pool: &Pool, scope: Scope) -> sqlx::Result<String> {
286	Ok(sqlx::query_scalar("SELECT prompt FROM assistant_prompts WHERE scope = ? AND scope_id = ?")
287	.bind(scope.kind())
288	.bind(scope.id())
289	.fetch_optional(pool)
290	.await?
291	.unwrap_or_default())
292	}
293
294	pub async fn set_prompt(pool: &Pool, scope: Scope, text: &str) -> sqlx::Result<()> {
295	if text.trim().is_empty() {
296	sqlx::query("DELETE FROM assistant_prompts WHERE scope = ? AND scope_id = ?")
297	.bind(scope.kind())
298	.bind(scope.id())
299	.execute(pool)
300	.await?;
301	} else {
302	sqlx::query(
303	"INSERT INTO assistant_prompts (scope, scope_id, prompt, updated_ms) VALUES (?, ?, ?, ?) \
304	ON CONFLICT (scope, scope_id) DO UPDATE SET prompt = excluded.prompt, updated_ms = excluded.updated_ms",
305	)
306	.bind(scope.kind())
307	.bind(scope.id())
308	.bind(text.trim())
309	.bind(now_ms())
310	.execute(pool)
311	.await?;
312	}
313	Ok(())
314	}
315
316	// ----------------------------------------------------- what applies where
317
318	/// The levels above a scope, outermost first: the site, then the
319	/// organisation when the owner is one.
320	pub async fn levels_above(pool: &Pool, scope: Scope, owner: &str) -> sqlx::Result<Vec<(Scope, String)>> {
321	let mut out = vec![];
322	match scope {
323	Scope::Site => {}
324	Scope::Org(_) => out.push((Scope::Site, "the site".to_string())),
325	Scope::Repo(_) => {
326	out.push((Scope::Site, "the site".to_string()));
327	if let Some(org) = crate::db::accounts::org_by_name(pool, owner).await? {
328	out.push((Scope::Org(org.id), format!("the {} organisation", org.name)));
329	}
330	}
331	}
332	Ok(out)
333	}
334
335	/// An assistant seen from a level: where it comes from, and whether it is
336	/// on once outer levels have had their say.
337	#[derive(Clone, Debug)]
338	pub struct Available {
339	pub assistant: Assistant,
340	/// `the site`, `the acme organisation`, or empty for the level's own.
341	pub from: String,
342	/// Off because a level above disabled an assistant of this name.
343	pub off_above: bool,
344	}
345
346	impl Available {
347	pub fn on(&self) -> bool {
348	self.assistant.enabled && !self.off_above
349	}
350	}
351
352	/// Every assistant a level sees, outermost first, one per name: the
353	/// innermost definition of a name wins, and a name disabled at an outer
354	/// level is off inside it whatever the inner level says.
355	pub async fn available(pool: &Pool, scope: Scope, owner: &str) -> sqlx::Result<Vec<Available>> {
356	let mut levels = levels_above(pool, scope, owner).await?;
357	levels.push((scope, String::new()));
358	let mut out: Vec<Available> = vec![];
359	let mut off: Vec<String> = vec![];
360	for (level, from) in levels {
361	for a in list(pool, level).await? {
362	let off_above = off.iter().any(\|n\| n.eq_ignore_ascii_case(&a.name));
363	if !a.enabled && !off.iter().any(\|n\| n.eq_ignore_ascii_case(&a.name)) {
364	off.push(a.name.clone());
365	}
366	out.retain(\|o\| !o.assistant.name.eq_ignore_ascii_case(&a.name));
367	out.push(Available { assistant: a, from: from.clone(), off_above });
368	}
369	}
370	Ok(out)
371	}
372
373	/// The assistants a repository may ask, by name.
374	pub async fn enabled_for_repo(pool: &Pool, repo: &crate::db::repos::RepoRow) -> sqlx::Result<Vec<Assistant>> {
375	Ok(available(pool, Scope::Repo(repo.id), &repo.owner).await?.into_iter().filter(Available::on).map(\|a\| a.assistant).collect())
376	}
377
378	/// The prompts a level inherits, outermost first: (where from, text),
379	/// only the levels that set one.
380	pub async fn inherited_prompts(pool: &Pool, scope: Scope, owner: &str) -> sqlx::Result<Vec<(String, String)>> {
381	let mut out = vec![];
382	for (level, from) in levels_above(pool, scope, owner).await? {
383	let text = prompt(pool, level).await?;
384	if !text.is_empty() {
385	out.push((from, text));
386	}
387	}
388	Ok(out)
389	}
390
391	/// The prompt in force for a repository: each level's text under a
392	/// heading, outermost first, so a project adds to what the site says.
393	pub async fn prompt_for_repo(pool: &Pool, repo: &crate::db::repos::RepoRow) -> sqlx::Result<String> {
394	let scope = Scope::Repo(repo.id);
395	let mut parts = inherited_prompts(pool, scope, &repo.owner).await?;
396	let own = prompt(pool, scope).await?;
397	if !own.is_empty() {
398	parts.push(("this repository".to_string(), own));
399	}
400	Ok(parts.into_iter().map(\|(from, text)\| format!("## Instructions from {from}\n\n{text}\n")).collect::<Vec<_>>().join("\n"))
401	}
402
403	// -------------------------------------------------------------- calling
404
405	/// Ask an assistant: the system prompt and the request, the answer's
406	/// text. Errors are for the person who asked: a status and the first
407	/// line of what the endpoint said, or why no answer came.
408	pub async fn ask(app: &Shared, a: &Assistant, system: &str, user: &str) -> Result<String, String> {
409	let Some(token) = app.assistant_key.open(&a.token) else {
410	return Err("its token cannot be read: set it again".into());
411	};
412	let client = reqwest::Client::builder()
413	.timeout(CALL_TIMEOUT)
414	.user_agent(format!("oxygenhub/{}", env!("CARGO_PKG_VERSION")))
415	.build()
416	.map_err(\|e\| e.to_string())?;
417	let mut req = client.post(&a.endpoint).header("content-type", "application/json");
418	let body = match a.api.as_str() {
419	"anthropic-messages" => {
420	if !token.is_empty() {
421	req = req.header("x-api-key", &token);
422	}
423	req = req.header("anthropic-version", "2023-06-01");
424	serde_json::json!({
425	"model": a.model,
426	"max_tokens": 4096,
427	"system": system,
428	"messages": [{ "role": "user", "content": user }],
429	})
430	}
431	_ => {
432	if !token.is_empty() {
433	req = req.header("authorization", format!("Bearer {token}"));
434	}
435	serde_json::json!({
436	"model": a.model,
437	"messages": [
438	{ "role": "system", "content": system },
439	{ "role": "user", "content": user },
440	],
441	})
442	}
443	};
444	for (name, value) in a.header_list() {
445	req = req.header(name, value);
446	}
447	let r = req.json(&body).send().await.map_err(\|e\| {
448	if e.is_timeout() {
449	"no answer within 60 seconds".to_string()
450	} else {
451	scrub(&e.to_string(), &token)
452	}
453	})?;
454	let status = r.status();
455	let text: String = r.text().await.unwrap_or_default().chars().take(MAX_ANSWER).collect();
456	if !status.is_success() {
457	let why = serde_json::from_str::<serde_json::Value>(&text)
458	.ok()
459	.and_then(\|v\| v["error"]["message"].as_str().or(v["error"].as_str()).or(v["message"].as_str()).map(str::to_string))
460	.unwrap_or_else(\|\| text.lines().next().unwrap_or("").chars().take(160).collect());
461	return Err(scrub(&format!("{} {}", status.as_u16(), why.trim()), &token));
462	}
463	let v: serde_json::Value = serde_json::from_str(&text).map_err(\|_\| "the answer was not JSON".to_string())?;
464	let answer = match a.api.as_str() {
465	"anthropic-messages" => v["content"]
466	.as_array()
467	.map(\|parts\| parts.iter().filter_map(\|p\| p["text"].as_str()).collect::<Vec<_>>().join(""))
468	.unwrap_or_default(),
469	_ => match &v["choices"][0]["message"]["content"] {
470	serde_json::Value::String(s) => s.clone(),
471	serde_json::Value::Array(parts) => parts.iter().filter_map(\|p\| p["text"].as_str()).collect::<Vec<_>>().join(""),
472	_ => String::new(),
473	},
474	};
475	if answer.trim().is_empty() {
476	return Err("the answer carried no text".into());
477	}
478	Ok(answer)
479	}
480
481	/// A tiny exchange, for the settings page's TEST: the start of the answer.
482	pub async fn test(app: &Shared, a: &Assistant) -> Result<String, String> {
483	let answer = ask(app, a, "You are being connected to an OxygenHub site. Answer briefly.", "Reply with the single word OK.").await?;
484	Ok(answer.trim().chars().take(80).collect())
485	}
486
487	/// Never let the token through in an error message.
488	fn scrub(text: &str, token: &str) -> String {
489	if token.is_empty() {
490	text.to_string()
491	} else {
492	text.replace(token, "***")
493	}
494	}
495
496	#[cfg(test)]
497	mod tests {
498	use super::*;
499
500	#[test]
501	fn tokens_seal_and_open_only_under_their_key() {
502	let dir = tempfile::tempdir().unwrap();
503	let k = SecretKey::load_or_create(dir.path()).unwrap();
504	let sealed = k.seal("sk-very-secret");
505	assert_ne!(sealed, b"sk-very-secret");
506	assert!(!String::from_utf8_lossy(&sealed).contains("secret"));
507	assert_eq!(k.open(&sealed).as_deref(), Some("sk-very-secret"));
508	assert_eq!(SecretKey::load_or_create(dir.path()).unwrap().open(&sealed).as_deref(), Some("sk-very-secret"), "the key is kept");
509	let other = SecretKey::load_or_create(tempfile::tempdir().unwrap().path()).unwrap();
510	assert_eq!(other.open(&sealed), None);
511	assert_eq!(k.open(&[1, 2, 3]), None);
512	assert_eq!(k.seal(""), Vec::<u8>::new());
513	assert_eq!(k.open(&[]).as_deref(), Some(""));
514	}
515
516	#[test]
517	fn headers_parse_or_name_the_bad_line() {
518	assert_eq!(
519	parse_headers("X-Org: acme\n\n OpenAI-Beta : v1 ").unwrap(),
520	vec![("X-Org".into(), "acme".into()), ("OpenAI-Beta".into(), "v1".into())]
521	);
522	assert_eq!(parse_headers("no colon").unwrap_err(), "no colon");
523	assert_eq!(parse_headers("bad name: x").unwrap_err(), "bad name: x");
524	}
525	}

Keyboard