| 1 | //! [`Provider::Local`](crate::llm::Provider::Local): the turn loop on an on-device model, with |
| 2 | //! onde. |
| 3 | //! |
| 4 | //! The rest of `ed-acp` keeps the conversation as OpenAI chat-completions JSON and calls |
| 5 | //! [`LlmClient::complete`](crate::llm::LlmClient::complete) once per model turn with all of it. |
| 6 | //! This module converts that JSON to onde's [`ChatTurn`]s and runs |
| 7 | //! [`ChatEngine::complete`], which takes the whole conversation each call and leaves the |
| 8 | //! engine's own history alone. One engine serves the process; it loads the requested model on |
| 9 | //! first use (downloading it the first time) and swaps it when the session picks another. |
| 10 | //! |
| 11 | //! onde doesn't stream tool-aware turns yet, so a turn's reasoning and text are reported |
| 12 | //! through `on_delta` once the turn ends. |
| 13 | //! |
| 14 | //! Without the `local` feature the provider is still selectable but every turn fails with an |
| 15 | //! error that says how to enable it. |
| 16 | |
| 17 | #[cfg(feature = "local")] |
| 18 | pub use enabled::{catalog, complete}; |
| 19 | |
| 20 | #[cfg(not(feature = "local"))] |
| 21 | pub use disabled::{catalog, complete}; |
| 22 | |
| 23 | #[cfg(not(feature = "local"))] |
| 24 | mod disabled { |
| 25 | use serde_json::Value; |
| 26 | |
| 27 | use crate::llm::{Completion, Delta, ModelInfo}; |
| 28 | |
| 29 | pub fn catalog() -> Vec<ModelInfo> { |
| 30 | Vec::new() |
| 31 | } |
| 32 | |
| 33 | pub async fn complete( |
| 34 | _model: &str, |
| 35 | _messages: &[Value], |
| 36 | _tools: &Value, |
| 37 | _on_delta: impl FnMut(Delta<'_>), |
| 38 | ) -> anyhow::Result<Completion> { |
| 39 | anyhow::bail!( |
| 40 | "this agent was built without on-device inference; rebuild it with ed-acp's `local` \ |
| 41 | feature, or choose another provider" |
| 42 | ) |
| 43 | } |
| 44 | } |
| 45 | |
| 46 | #[cfg(feature = "local")] |
| 47 | mod enabled { |
| 48 | use std::sync::OnceLock; |
| 49 | |
| 50 | use anyhow::{Context, Result, anyhow, bail}; |
| 51 | use onde::inference::models::{SUPPORTED_MODEL_INFO, TOOL_CALLING_MODELS}; |
| 52 | use onde::inference::{ |
| 53 | ChatEngine, ChatTurn, GgufModelConfig, ToolCallRequest as OndeCall, ToolCallingSupport, |
| 54 | ToolDefinition, |
| 55 | }; |
| 56 | use serde_json::Value; |
| 57 | use tokio::sync::Mutex; |
| 58 | |
| 59 | use crate::llm::{Completion, Delta, ModelInfo, ToolCallRequest, Usage}; |
| 60 | |
| 61 | /// The process's engine and the model loaded in it. The mutex is held across a turn, so |
| 62 | /// sessions take turns on the one model instead of swapping it under each other. |
| 63 | struct Local { |
| 64 | engine: ChatEngine, |
| 65 | loaded: Mutex<Option<String>>, |
| 66 | } |
| 67 | |
| 68 | fn local() -> &'static Local { |
| 69 | static LOCAL: OnceLock<Local> = OnceLock::new(); |
| 70 | LOCAL.get_or_init(|| Local { |
| 71 | engine: ChatEngine::new(), |
| 72 | loaded: Mutex::new(None), |
| 73 | }) |
| 74 | } |
| 75 | |
| 76 | /// onde's tool-calling models, the ones an agent can drive. |
| 77 | pub fn catalog() -> Vec<ModelInfo> { |
| 78 | TOOL_CALLING_MODELS |
| 79 | .iter() |
| 80 | .map(|id| ModelInfo { |
| 81 | id: (*id).to_string(), |
| 82 | owned_by: SUPPORTED_MODEL_INFO |
| 83 | .iter() |
| 84 | .find(|info| info.id == *id) |
| 85 | .map(|info| format!("{} (on device)", info.name)), |
| 86 | }) |
| 87 | .collect() |
| 88 | } |
| 89 | |
| 90 | pub async fn complete( |
| 91 | model: &str, |
| 92 | messages: &[Value], |
| 93 | tools: &Value, |
| 94 | mut on_delta: impl FnMut(Delta<'_>), |
| 95 | ) -> Result<Completion> { |
| 96 | let turns = to_turns(messages)?; |
| 97 | let tools = to_tools(tools); |
| 98 | let local = local(); |
| 99 | let mut loaded = local.loaded.lock().await; |
| 100 | if loaded.as_deref() != Some(model) { |
| 101 | let config = GgufModelConfig::from_supported_model_id(model) |
| 102 | .ok_or_else(|| anyhow!("{model} is not an on-device model onde supports"))?; |
| 103 | if onde::inference::models::tool_calling_support(model) != ToolCallingSupport::Supported |
| 104 | { |
| 105 | bail!("{model} can't call tools reliably, so it can't run an agent"); |
| 106 | } |
| 107 | on_delta(Delta::Reasoning(&format!( |
| 108 | "Loading {} on this device. The first use downloads it.\n", |
| 109 | config.display_name |
| 110 | ))); |
| 111 | tracing::info!("loading on-device model {model}"); |
| 112 | *loaded = None; |
| 113 | local |
| 114 | .engine |
| 115 | .load_gguf_model(config, None, None) |
| 116 | .await |
| 117 | .with_context(|| format!("loading {model}"))?; |
| 118 | *loaded = Some(model.to_string()); |
| 119 | } |
| 120 | let result = local |
| 121 | .engine |
| 122 | .complete(turns, &tools, None) |
| 123 | .await |
| 124 | .context("on-device inference")?; |
| 125 | drop(loaded); |
| 126 | |
| 127 | if let Some(reasoning) = &result.reasoning { |
| 128 | on_delta(Delta::Reasoning(reasoning)); |
| 129 | } |
| 130 | if !result.text.is_empty() { |
| 131 | on_delta(Delta::Text(&result.text)); |
| 132 | } |
| 133 | Ok(Completion { |
| 134 | content: result.text, |
| 135 | tool_calls: result |
| 136 | .tool_calls |
| 137 | .into_iter() |
| 138 | .map(|call| ToolCallRequest { |
| 139 | id: call.id, |
| 140 | name: call.function_name, |
| 141 | arguments: call.arguments, |
| 142 | extra_content: None, |
| 143 | }) |
| 144 | .collect(), |
| 145 | finish_reason: Some(result.finish_reason), |
| 146 | usage: Some(Usage { |
| 147 | prompt_tokens: result.prompt_tokens as u64, |
| 148 | completion_tokens: result.completion_tokens as u64, |
| 149 | total_tokens: (result.prompt_tokens + result.completion_tokens) as u64, |
| 150 | }), |
| 151 | }) |
| 152 | } |
| 153 | |
| 154 | /// OpenAI chat messages, as the turn loop keeps them, to onde turns. |
| 155 | pub(super) fn to_turns(messages: &[Value]) -> Result<Vec<ChatTurn>> { |
| 156 | messages |
| 157 | .iter() |
| 158 | .map(|message| { |
| 159 | let role = message.get("role").and_then(Value::as_str).unwrap_or(""); |
| 160 | let content = text_of(message.get("content")); |
| 161 | Ok(match role { |
| 162 | "system" | "developer" => ChatTurn::System(content), |
| 163 | "user" => ChatTurn::User(content), |
| 164 | "assistant" => ChatTurn::Assistant { |
| 165 | content, |
| 166 | tool_calls: message |
| 167 | .get("tool_calls") |
| 168 | .and_then(Value::as_array) |
| 169 | .map(|calls| calls.iter().map(to_call).collect()) |
| 170 | .unwrap_or_default(), |
| 171 | }, |
| 172 | "tool" => ChatTurn::Tool { |
| 173 | tool_call_id: message |
| 174 | .get("tool_call_id") |
| 175 | .and_then(Value::as_str) |
| 176 | .unwrap_or_default() |
| 177 | .to_string(), |
| 178 | content, |
| 179 | }, |
| 180 | other => bail!("unexpected message role {other:?}"), |
| 181 | }) |
| 182 | }) |
| 183 | .collect() |
| 184 | } |
| 185 | |
| 186 | fn to_call(call: &Value) -> OndeCall { |
| 187 | let field = |pointer: &str| { |
| 188 | call.pointer(pointer) |
| 189 | .and_then(Value::as_str) |
| 190 | .unwrap_or_default() |
| 191 | .to_string() |
| 192 | }; |
| 193 | OndeCall { |
| 194 | id: field("/id"), |
| 195 | function_name: field("/function/name"), |
| 196 | arguments: field("/function/arguments"), |
| 197 | } |
| 198 | } |
| 199 | |
| 200 | /// A message's text: a string, or the text parts of a content array. On-device models here |
| 201 | /// are text-only, so other parts become a short placeholder. |
| 202 | fn text_of(content: Option<&Value>) -> String { |
| 203 | match content { |
| 204 | Some(Value::String(text)) => text.clone(), |
| 205 | Some(Value::Array(parts)) => parts |
| 206 | .iter() |
| 207 | .map(|part| match part.get("type").and_then(Value::as_str) { |
| 208 | Some("text") => part |
| 209 | .get("text") |
| 210 | .and_then(Value::as_str) |
| 211 | .unwrap_or_default() |
| 212 | .to_string(), |
| 213 | Some(kind) => format!("[{kind} omitted: the on-device model reads text only]"), |
| 214 | None => String::new(), |
| 215 | }) |
| 216 | .filter(|text| !text.is_empty()) |
| 217 | .collect::<Vec<_>>() |
| 218 | .join("\n"), |
| 219 | _ => String::new(), |
| 220 | } |
| 221 | } |
| 222 | |
| 223 | /// OpenAI `tools` entries to onde tool definitions. |
| 224 | pub(super) fn to_tools(tools: &Value) -> Vec<ToolDefinition> { |
| 225 | tools |
| 226 | .as_array() |
| 227 | .map(|tools| { |
| 228 | tools |
| 229 | .iter() |
| 230 | .filter_map(|tool| { |
| 231 | let function = tool.get("function")?; |
| 232 | Some(ToolDefinition { |
| 233 | name: function.get("name")?.as_str()?.to_string(), |
| 234 | description: function |
| 235 | .get("description") |
| 236 | .and_then(Value::as_str) |
| 237 | .unwrap_or_default() |
| 238 | .to_string(), |
| 239 | parameters_schema: function |
| 240 | .get("parameters") |
| 241 | .cloned() |
| 242 | .unwrap_or_else(|| serde_json::json!({"type": "object"})) |
| 243 | .to_string(), |
| 244 | }) |
| 245 | }) |
| 246 | .collect() |
| 247 | }) |
| 248 | .unwrap_or_default() |
| 249 | } |
| 250 | } |
| 251 | |
| 252 | #[cfg(all(test, feature = "local"))] |
| 253 | mod tests { |
| 254 | use onde::inference::{ChatTurn, ToolCallRequest}; |
| 255 | use serde_json::json; |
| 256 | |
| 257 | use super::enabled::{catalog, to_tools, to_turns}; |
| 258 | use crate::llm::LOCAL_DEFAULT_MODEL; |
| 259 | |
| 260 | #[test] |
| 261 | fn the_default_model_is_in_the_catalog() { |
| 262 | assert!(catalog().iter().any(|m| m.id == LOCAL_DEFAULT_MODEL)); |
| 263 | assert!(catalog().iter().all(|m| { |
| 264 | m.owned_by |
| 265 | .as_deref() |
| 266 | .is_some_and(|o| o.ends_with("(on device)")) |
| 267 | })); |
| 268 | } |
| 269 | |
| 270 | #[test] |
| 271 | fn messages_become_turns() { |
| 272 | let messages = vec![ |
| 273 | json!({"role": "system", "content": "You are SplitFire."}), |
| 274 | json!({"role": "user", "content": [ |
| 275 | {"type": "text", "text": "What key is this?"}, |
| 276 | {"type": "image_url", "image_url": {"url": "data:..."}} |
| 277 | ]}), |
| 278 | json!({"role": "assistant", "content": "", "tool_calls": [ |
| 279 | {"id": "call_1", "type": "function", |
| 280 | "function": {"name": "analyze_audio", "arguments": "{\"path\":\"a.wav\"}"}} |
| 281 | ]}), |
| 282 | json!({"role": "tool", "tool_call_id": "call_1", "content": "Key: A minor"}), |
| 283 | json!({"role": "assistant", "content": "A minor."}), |
| 284 | ]; |
| 285 | let turns = to_turns(&messages).unwrap(); |
| 286 | assert_eq!( |
| 287 | turns, |
| 288 | vec![ |
| 289 | ChatTurn::System("You are SplitFire.".into()), |
| 290 | ChatTurn::User( |
| 291 | "What key is this?\n[image_url omitted: the on-device model reads text only]" |
| 292 | .into() |
| 293 | ), |
| 294 | ChatTurn::Assistant { |
| 295 | content: String::new(), |
| 296 | tool_calls: vec![ToolCallRequest { |
| 297 | id: "call_1".into(), |
| 298 | function_name: "analyze_audio".into(), |
| 299 | arguments: "{\"path\":\"a.wav\"}".into(), |
| 300 | }], |
| 301 | }, |
| 302 | ChatTurn::Tool { |
| 303 | tool_call_id: "call_1".into(), |
| 304 | content: "Key: A minor".into(), |
| 305 | }, |
| 306 | ChatTurn::Assistant { |
| 307 | content: "A minor.".into(), |
| 308 | tool_calls: vec![], |
| 309 | }, |
| 310 | ] |
| 311 | ); |
| 312 | assert!(to_turns(&[json!({"role": "narrator", "content": "x"})]).is_err()); |
| 313 | } |
| 314 | |
| 315 | #[test] |
| 316 | fn tools_become_definitions() { |
| 317 | let tools = json!([ |
| 318 | {"type": "function", "function": { |
| 319 | "name": "theory_chord", "description": "Spell a chord.", |
| 320 | "parameters": {"type": "object", "properties": {"symbol": {"type": "string"}}} |
| 321 | }}, |
| 322 | {"type": "function", "function": {"name": "no_params"}}, |
| 323 | {"type": "not-a-function"} |
| 324 | ]); |
| 325 | let defs = to_tools(&tools); |
| 326 | assert_eq!(defs.len(), 2); |
| 327 | assert_eq!(defs[0].name, "theory_chord"); |
| 328 | assert_eq!(defs[0].description, "Spell a chord."); |
| 329 | let schema: serde_json::Value = serde_json::from_str(&defs[0].parameters_schema).unwrap(); |
| 330 | assert_eq!(schema["properties"]["symbol"]["type"], "string"); |
| 331 | assert_eq!(defs[1].parameters_schema, r#"{"type":"object"}"#); |
| 332 | } |
| 333 | } |