| 1 | use crate::*; |
| 2 | use axum::Json; |
| 3 | |
| 4 | pub async fn route( |
| 5 | app: Arc<App>, |
| 6 | method: &Method, |
| 7 | parts: &[&str], |
| 8 | me: &Value, |
| 9 | headers: &HeaderMap, |
| 10 | ) -> Result<Response> { |
| 11 | need(me, "ai")?; |
| 12 | if method == Method::POST && me["viewing"] == true { |
| 13 | return Err(Error::new( |
| 14 | 403, |
| 15 | "Switch back to your account to access Local AI.", |
| 16 | )); |
| 17 | } |
| 18 | if !matches!( |
| 19 | (method, parts), |
| 20 | (&Method::GET, []) |
| 21 | | (&Method::POST, ["key"]) |
| 22 | | (&Method::GET, ["snow-codex" | "snow-claude"]) |
| 23 | ) { |
| 24 | return Err(Error::new(404, "No Local AI setting here.")); |
| 25 | } |
| 26 | if method == Method::POST |
| 27 | && headers.get("origin").and_then(|v| v.to_str().ok()) |
| 28 | != Some(app.mcp.origin.origin().ascii_serialization().as_str()) |
| 29 | { |
| 30 | return Err(Error::new( |
| 31 | 403, |
| 32 | "Open Local AI in Snowglobe to copy its key.", |
| 33 | )); |
| 34 | } |
| 35 | let jobs = core::nomad(&app, "/v1/jobs").await?; |
| 36 | let id = array(&jobs) |
| 37 | .iter() |
| 38 | .filter(|job| job["Stop"] != true) |
| 39 | .map(|job| string(&job["ID"])) |
| 40 | .find(|id| *id == "local-ai" || id.starts_with("local-ai-preview-")) |
| 41 | .ok_or_else(|| { |
| 42 | Error::new( |
| 43 | 503, |
| 44 | "Local AI is offline. Retry when the service is running.", |
| 45 | ) |
| 46 | })?; |
| 47 | let job = core::nomad(&app, &format!("/v1/job/{}", encoded(id))).await?; |
| 48 | let hostname = string(&job["Meta"]["studio_hostname"]); |
| 49 | if hostname.is_empty() |
| 50 | || !hostname |
| 51 | .bytes() |
| 52 | .all(|c| c.is_ascii_alphanumeric() || matches!(c, b'.' | b'-')) |
| 53 | { |
| 54 | return Err(Error::new(502, "Local AI has no valid endpoint.")); |
| 55 | } |
| 56 | let endpoint = format!("https://{hostname}"); |
| 57 | let catalog: Value = serde_json::from_slice( |
| 58 | &tokio::fs::read(app.repo.join("service/local-ai/catalog.json")).await?, |
| 59 | )?; |
| 60 | let model = &catalog["models"][0]; |
| 61 | let response = match parts { |
| 62 | [] => Json(json!({"endpoint":endpoint,"model":model["display_name"],"context":model["context_window"]})).into_response(), |
| 63 | ["key"] => { |
| 64 | let secret = core::nomad(&app, &format!("/v1/var/nomad/jobs/{}", encoded(id))).await?; |
| 65 | let key = string(&secret["Items"]["api_key"]); |
| 66 | if key.is_empty() { return Err(Error::new(503, "Local AI has no API key configured.")); } |
| 67 | Json(json!({"key":key})).into_response() |
| 68 | } |
| 69 | [name] => { |
| 70 | let client = name.strip_prefix("snow-").unwrap(); |
| 71 | let quote = |value: &str| format!("'{}'", value.replace('\'', "'\\''")); |
| 72 | let script = include_str!("../ai/launch.sh") |
| 73 | .replace("@CLIENT@", client) |
| 74 | .replace("@ENDPOINT@", "e(&endpoint)) |
| 75 | .replace("@MODEL@", "e(string(&model["slug"]))) |
| 76 | .replace("@CATALOG@", &serde_json::to_string_pretty(&catalog)?) |
| 77 | .replace("@CONTEXT@", &model["context_window"].to_string()) |
| 78 | .replace("@COMPACT@", &model["auto_compact_token_limit"].to_string()) |
| 79 | .replace("@REASONING@", string(&model["default_reasoning_level"])); |
| 80 | ([("content-type", "text/x-shellscript; charset=utf-8"), ("content-disposition", &format!("attachment; filename=\"{name}\""))], script).into_response() |
| 81 | } |
| 82 | _ => unreachable!(), |
| 83 | }; |
| 84 | let mut response = response; |
| 85 | response |
| 86 | .headers_mut() |
| 87 | .insert("cache-control", "no-store".parse().unwrap()); |
| 88 | Ok(response) |
| 89 | } |