1 //! Inference backend abstraction.
2 //!
3 //! The agent loop only needs to send a turn (optionally with tools) and return
4 //! tool results. This module defines that seam as the `InferenceBackend` trait
5 //! plus a few neutral types, with two implementations:
6 //!
7 //! - `LocalBackend` runs on-device through the `onde` crate (`ChatEngine`).
8 //! - `OpenAiBackend` talks to any OpenAI-compatible HTTP endpoint, configured by
9 //! `base_url`, `api_key`, and `model`.
10 //!
11 //! The trait exposes neither `onde` nor OpenAI types, so the loop does not depend
12 //! on a specific backend.
13 //!
14 //! The whole backend seam is wired up only through the interactive client, which
15 //! is `#[cfg(unix)]` (see `run_interactive` in `main.rs` and `mod tui` in
16 //! `chat.rs`). On non-Unix targets the binary runs ACP-only and drives `onde`
17 //! directly, so every item here is legitimately unused there. Suppress the
18 //! dead-code lint on those targets only — Unix builds still get full coverage.
19 #![cfg_attr(not(unix), allow(dead_code))]
20
21 use std::sync::Arc;
22
23 use async_trait::async_trait;
24 use onde::inference::{ChatEngine, ToolDefinition};
25 use serde::Deserialize;
26 use tokio::sync::Mutex;
27
28 // ── Neutral types ───────────────────────────────────────────────────────────────
29
30 /// A tool the model may call, in a provider-neutral form. `parameters_schema` is
31 /// a JSON Schema encoded as a string (matching how siGit already declares tools).
32 #[derive(Debug, Clone)]
33 pub struct ToolSpec {
34 pub name: String,
35 pub description: String,
36 pub parameters_schema: String,
37 }
38
39 /// A tool call requested by the model.
40 #[derive(Debug, Clone)]
41 pub struct ToolCall {
42 pub id: String,
43 pub name: String,
44 /// Arguments as a JSON-encoded string.
45 pub arguments: String,
46 }
47
48 /// The output of executing one tool call, fed back to the model.
49 #[derive(Debug, Clone)]
50 pub struct ToolResult {
51 pub tool_call_id: String,
52 pub content: String,
53 }
54
55 /// The result of one assistant turn: free text and/or tool calls.
56 #[derive(Debug, Clone, Default)]
57 pub struct TurnResult {
58 pub text: String,
59 pub tool_calls: Vec<ToolCall>,
60 }
61
62 /// Backend errors are plain strings. Callers map them to ACP errors.
63 pub type BackendError = String;
64
65 /// A sink for streaming assistant text deltas to the UI as they are produced.
66 ///
67 /// When a caller passes `Some(sink)`, a streaming-capable backend forwards each
68 /// text fragment through it as the model emits it; the returned [`TurnResult`]
69 /// still carries the fully assembled text (and any tool calls). When the sink is
70 /// `None`, the backend runs in non-streaming mode. Unbounded so the inference
71 /// task never blocks on a slow consumer.
72 pub type TokenSink = tokio::sync::mpsc::UnboundedSender<String>;
73
74 // ── The trait ───────────────────────────────────────────────────────────────────
75
76 /// A swappable inference backend driving siGit Code's agent loop.
77 #[async_trait]
78 pub trait InferenceBackend: Send + Sync {
79 /// Start an assistant turn from a new user message, offering `tools`.
80 ///
81 /// If `sink` is `Some`, text is streamed through it as it is generated. A
82 /// backend may decline to stream a given round (for example, on-device
83 /// inference cannot stream while it is still deciding whether to call a
84 /// tool); in that case the text is delivered only via the returned result.
85 async fn send_message_with_tools(
86 &self,
87 text: &str,
88 tools: &[ToolSpec],
89 sink: Option<&TokenSink>,
90 ) -> Result<TurnResult, BackendError>;
91
92 /// Continue the turn by returning tool results. `tools` may be `None` on the
93 /// final round to force a text answer. `sink` streams that text when set.
94 async fn send_tool_results(
95 &self,
96 results: Vec<ToolResult>,
97 tools: Option<&[ToolSpec]>,
98 sink: Option<&TokenSink>,
99 ) -> Result<TurnResult, BackendError>;
100
101 /// Whether inference runs over the network (a configured provider) rather
102 /// than on-device. Drives UI labelling so the displayed model can't claim a
103 /// local model while requests actually go to the cloud.
104 fn is_remote(&self) -> bool;
105 }
106
107 // ── Local backend (onde ChatEngine) ──────────────────────────────────────────────
108
109 /// On-device inference. A thin adapter over `onde::ChatEngine`.
110 pub struct LocalBackend {
111 engine: Arc<ChatEngine>,
112 }
113
114 impl LocalBackend {
115 pub fn new(engine: Arc<ChatEngine>) -> Self {
116 Self { engine }
117 }
118 }
119
120 fn to_onde_tools(tools: &[ToolSpec]) -> Vec<ToolDefinition> {
121 tools
122 .iter()
123 .map(|tool| ToolDefinition {
124 name: tool.name.clone(),
125 description: tool.description.clone(),
126 parameters_schema: tool.parameters_schema.clone(),
127 })
128 .collect()
129 }
130
131 #[async_trait]
132 impl InferenceBackend for LocalBackend {
133 async fn send_message_with_tools(
134 &self,
135 text: &str,
136 tools: &[ToolSpec],
137 sink: Option<&TokenSink>,
138 ) -> Result<TurnResult, BackendError> {
139 // onde's tool-aware path is non-streaming: it has to buffer the whole
140 // reply to detect tool calls. We can only stream when no tools are on
141 // offer (a plain answer), which is exactly the tools-disabled case.
142 if let Some(sink) = sink
143 && tools.is_empty()
144 {
145 let rx = self
146 .engine
147 .stream_message(text)
148 .await
149 .map_err(|error| error.to_string())?;
150 return drain_onde_stream(rx, sink).await;
151 }
152
153 let onde_tools = to_onde_tools(tools);
154 let result = self
155 .engine
156 .send_message_with_tools(text, &onde_tools)
157 .await
158 .map_err(|error| error.to_string())?;
159 Ok(onde_result_to_turn(result))
160 }
161
162 async fn send_tool_results(
163 &self,
164 results: Vec<ToolResult>,
165 tools: Option<&[ToolSpec]>,
166 sink: Option<&TokenSink>,
167 ) -> Result<TurnResult, BackendError> {
168 let onde_results: Vec<onde::inference::ToolResult> = results
169 .into_iter()
170 .map(|result| onde::inference::ToolResult {
171 tool_call_id: result.tool_call_id,
172 content: result.content,
173 })
174 .collect();
175
176 // The final round passes `tools = None` to force a text answer; that's
177 // the only round onde can stream, since no further tool calls are parsed.
178 if let Some(sink) = sink
179 && tools.is_none()
180 {
181 let rx = self
182 .engine
183 .stream_tool_results(onde_results, None)
184 .await
185 .map_err(|error| error.to_string())?;
186 return drain_onde_stream(rx, sink).await;
187 }
188
189 let onde_tools = tools.map(to_onde_tools);
190 let result = self
191 .engine
192 .send_tool_results(onde_results, onde_tools.as_deref())
193 .await
194 .map_err(|error| error.to_string())?;
195 Ok(onde_result_to_turn(result))
196 }
197
198 fn is_remote(&self) -> bool {
199 false
200 }
201 }
202
203 /// Drain an onde streaming receiver, forwarding each token to `sink` and
204 /// assembling the full text. onde reports stream failures as a final chunk whose
205 /// `finish_reason` is `"error: …"`; surface those as a backend error.
206 async fn drain_onde_stream(
207 mut rx: tokio::sync::mpsc::Receiver<onde::inference::StreamChunk>,
208 sink: &TokenSink,
209 ) -> Result<TurnResult, BackendError> {
210 let mut text = String::new();
211 while let Some(chunk) = rx.recv().await {
212 if !chunk.delta.is_empty() {
213 text.push_str(&chunk.delta);
214 // The receiver is the UI; if it's gone the turn is being cancelled,
215 // so stop assembling rather than spinning the model to completion.
216 if sink.send(chunk.delta).is_err() {
217 break;
218 }
219 }
220 if chunk.done {
221 if let Some(reason) = chunk.finish_reason
222 && let Some(message) = reason.strip_prefix("error: ")
223 {
224 return Err(message.to_string());
225 }
226 break;
227 }
228 }
229 Ok(TurnResult {
230 text,
231 tool_calls: Vec::new(),
232 })
233 }
234
235 /// Convert an `onde` tool-aware result into the neutral [`TurnResult`].
236 fn onde_result_to_turn(result: onde::inference::ToolAwareResult) -> TurnResult {
237 TurnResult {
238 text: result.text,
239 tool_calls: result
240 .tool_calls
241 .into_iter()
242 .map(|call| ToolCall {
243 id: call.id,
244 name: call.function_name,
245 arguments: call.arguments,
246 })
247 .collect(),
248 }
249 }
250
251 // ── OpenAI-compatible backend ─────────────────────────────────────────────────────
252
253 /// Inference against any OpenAI-compatible Chat Completions endpoint.
254 ///
255 /// Conversation state is held client-side and replayed on every request, so the
256 /// endpoint can be stateless. Standard OpenAI function-calling is used end to
257 /// end (`tools`, `choices[].message.tool_calls`, `role: "tool"` follow-ups).
258 pub struct OpenAiBackend {
259 base_url: String,
260 api_key: String,
261 model: String,
262 http: reqwest::Client,
263 /// The full message list sent on each request (system + turns + tool results).
264 history: Mutex<Vec<serde_json::Value>>,
265 }
266
267 impl OpenAiBackend {
268 /// Build a backend for `{base_url, api_key, model}`, seeding the optional
269 /// system prompt. `base_url` should include the API root (e.g. ending in
270 /// `/v1`); the chat path is appended.
271 pub fn new(
272 base_url: impl Into<String>,
273 api_key: impl Into<String>,
274 model: impl Into<String>,
275 system_prompt: Option<String>,
276 ) -> Self {
277 let mut history = Vec::new();
278 if let Some(prompt) = system_prompt {
279 history.push(serde_json::json!({ "role": "system", "content": prompt }));
280 }
281 Self {
282 base_url: base_url.into(),
283 api_key: api_key.into(),
284 model: model.into(),
285 http: reqwest::Client::new(),
286 history: Mutex::new(history),
287 }
288 }
289
290 fn tools_json(tools: &[ToolSpec]) -> Vec<serde_json::Value> {
291 tools
292 .iter()
293 .map(|tool| {
294 // parameters_schema is a JSON string; parse it, defaulting to an
295 // empty object schema if malformed.
296 let parameters: serde_json::Value = serde_json::from_str(&tool.parameters_schema)
297 .unwrap_or_else(|_| serde_json::json!({ "type": "object", "properties": {} }));
298 serde_json::json!({
299 "type": "function",
300 "function": {
301 "name": tool.name,
302 "description": tool.description,
303 "parameters": parameters,
304 }
305 })
306 })
307 .collect()
308 }
309
310 /// POST the current history (plus `tools`) and apply the assistant reply to
311 /// history, returning the neutral turn result. Streams via SSE when `sink`
312 /// is set; otherwise reads a single JSON response.
313 async fn complete(
314 &self,
315 tools: Option<&[ToolSpec]>,
316 sink: Option<&TokenSink>,
317 ) -> Result<TurnResult, BackendError> {
318 let url = format!("{}/chat/completions", self.base_url.trim_end_matches('/'));
319 let streaming = sink.is_some();
320
321 let mut body = serde_json::json!({
322 "model": self.model,
323 "messages": *self.history.lock().await,
324 "stream": streaming,
325 });
326 if let Some(tools) = tools
327 && !tools.is_empty()
328 {
329 body["tools"] = serde_json::Value::Array(Self::tools_json(tools));
330 }
331
332 let response = self
333 .http
334 .post(&url)
335 .bearer_auth(&self.api_key)
336 .json(&body)
337 .send()
338 .await
339 .map_err(|error| format!("request to {url} failed: {error}"))?;
340
341 if !response.status().is_success() {
342 let status = response.status();
343 let detail = response.text().await.unwrap_or_default();
344 return Err(format!("endpoint returned {status}: {detail}"));
345 }
346
347 if let Some(sink) = sink {
348 self.consume_stream(response, sink).await
349 } else {
350 self.consume_json(response).await
351 }
352 }
353
354 /// Parse a single non-streaming chat-completion response.
355 async fn consume_json(&self, response: reqwest::Response) -> Result<TurnResult, BackendError> {
356 let parsed: ChatCompletion = response
357 .json()
358 .await
359 .map_err(|error| format!("response parse error: {error}"))?;
360
361 let message = parsed
362 .choices
363 .into_iter()
364 .next()
365 .map(|choice| choice.message)
366 .ok_or_else(|| "endpoint returned no choices".to_string())?;
367
368 let text = message.content.clone().unwrap_or_default();
369 let tool_calls: Vec<ToolCall> = message
370 .tool_calls
371 .iter()
372 .flatten()
373 .map(|call| ToolCall {
374 id: call.id.clone(),
375 name: call.function.name.clone(),
376 arguments: call.function.arguments.clone(),
377 })
378 .collect();
379
380 // Record the assistant turn so later tool results have context.
381 self.history.lock().await.push(message.into_history_value());
382
383 Ok(TurnResult { text, tool_calls })
384 }
385
386 /// Consume an OpenAI Server-Sent Events stream, forwarding content deltas to
387 /// `sink` and reassembling any tool calls (which arrive fragmented across
388 /// chunks, keyed by `index`).
389 async fn consume_stream(
390 &self,
391 response: reqwest::Response,
392 sink: &TokenSink,
393 ) -> Result<TurnResult, BackendError> {
394 use futures::StreamExt;
395
396 let mut stream = response.bytes_stream();
397 // Newlines are ASCII, so splitting raw bytes on `\n` never bisects a
398 // multibyte UTF-8 sequence; we only lossily decode whole lines.
399 let mut buffer: Vec<u8> = Vec::new();
400 let mut text = String::new();
401 let mut tool_accum: Vec<StreamingToolCall> = Vec::new();
402 let mut done = false;
403
404 while let Some(item) = stream.next().await {
405 let bytes = item.map_err(|error| format!("stream read error: {error}"))?;
406 buffer.extend_from_slice(&bytes);
407
408 while let Some(pos) = buffer.iter().position(|&b| b == b'\n') {
409 let line: Vec<u8> = buffer.drain(..=pos).collect();
410 let line = String::from_utf8_lossy(&line);
411 let line = line.trim();
412
413 let Some(data) = line.strip_prefix("data:") else {
414 continue;
415 };
416 let data = data.trim();
417 if data == "[DONE]" {
418 done = true;
419 break;
420 }
421 if data.is_empty() {
422 continue;
423 }
424
425 let chunk: StreamCompletion = match serde_json::from_str(data) {
426 Ok(chunk) => chunk,
427 // Skip keep-alive comments and anything we can't parse rather
428 // than aborting a turn over one malformed frame.
429 Err(_) => continue,
430 };
431
432 let Some(choice) = chunk.choices.into_iter().next() else {
433 continue;
434 };
435 if let Some(content) = choice.delta.content
436 && !content.is_empty()
437 {
438 text.push_str(&content);
439 if sink.send(content).is_err() {
440 // Consumer dropped (turn cancelled) — stop reading.
441 done = true;
442 break;
443 }
444 }
445 for delta in choice.delta.tool_calls.into_iter().flatten() {
446 let index = delta.index.unwrap_or(0) as usize;
447 if tool_accum.len() <= index {
448 tool_accum.resize_with(index + 1, StreamingToolCall::default);
449 }
450 let slot = &mut tool_accum[index];
451 if let Some(id) = delta.id {
452 slot.id = id;
453 }
454 if let Some(function) = delta.function {
455 if let Some(name) = function.name {
456 slot.name = name;
457 }
458 if let Some(arguments) = function.arguments {
459 slot.arguments.push_str(&arguments);
460 }
461 }
462 }
463 }
464
465 if done {
466 break;
467 }
468 }
469
470 let tool_calls: Vec<ToolCall> = tool_accum
471 .iter()
472 .filter(|call| !call.name.is_empty())
473 .enumerate()
474 .map(|(index, call)| ToolCall {
475 id: if call.id.is_empty() {
476 format!("call_{index}")
477 } else {
478 call.id.clone()
479 },
480 name: call.name.clone(),
481 arguments: call.arguments.clone(),
482 })
483 .collect();
484
485 // Record the assistant turn so later tool results have context.
486 self.history
487 .lock()
488 .await
489 .push(streamed_assistant_history(&text, &tool_calls));
490
491 Ok(TurnResult { text, tool_calls })
492 }
493 }
494
495 /// One tool call being reassembled from streamed deltas.
496 #[derive(Default)]
497 struct StreamingToolCall {
498 id: String,
499 name: String,
500 arguments: String,
501 }
502
503 /// Rebuild the assistant message for replay in history after a streamed turn,
504 /// preserving any tool calls so the follow-up request is well-formed. Mirrors
505 /// [`ResponseMessage::into_history_value`] for the non-streaming path.
506 fn streamed_assistant_history(text: &str, tool_calls: &[ToolCall]) -> serde_json::Value {
507 let mut message = serde_json::json!({ "role": "assistant" });
508 message["content"] = if text.is_empty() {
509 serde_json::Value::Null
510 } else {
511 serde_json::Value::String(text.to_string())
512 };
513 if !tool_calls.is_empty() {
514 message["tool_calls"] = serde_json::json!(
515 tool_calls
516 .iter()
517 .map(|call| serde_json::json!({
518 "id": call.id,
519 "type": "function",
520 "function": {
521 "name": call.name,
522 "arguments": call.arguments,
523 }
524 }))
525 .collect::<Vec<_>>()
526 );
527 }
528 message
529 }
530
531 #[async_trait]
532 impl InferenceBackend for OpenAiBackend {
533 async fn send_message_with_tools(
534 &self,
535 text: &str,
536 tools: &[ToolSpec],
537 sink: Option<&TokenSink>,
538 ) -> Result<TurnResult, BackendError> {
539 self.history
540 .lock()
541 .await
542 .push(serde_json::json!({ "role": "user", "content": text }));
543 self.complete(Some(tools), sink).await
544 }
545
546 async fn send_tool_results(
547 &self,
548 results: Vec<ToolResult>,
549 tools: Option<&[ToolSpec]>,
550 sink: Option<&TokenSink>,
551 ) -> Result<TurnResult, BackendError> {
552 {
553 let mut history = self.history.lock().await;
554 for result in results {
555 history.push(serde_json::json!({
556 "role": "tool",
557 "tool_call_id": result.tool_call_id,
558 "content": result.content,
559 }));
560 }
561 }
562 self.complete(tools, sink).await
563 }
564
565 fn is_remote(&self) -> bool {
566 true
567 }
568 }
569
570 // ── OpenAI response shapes ────────────────────────────────────────────────────────
571
572 #[derive(Debug, Deserialize)]
573 struct ChatCompletion {
574 #[serde(default)]
575 choices: Vec<CompletionChoice>,
576 }
577
578 #[derive(Debug, Deserialize)]
579 struct CompletionChoice {
580 message: ResponseMessage,
581 }
582
583 #[derive(Debug, Deserialize)]
584 struct ResponseMessage {
585 #[serde(default)]
586 content: Option<String>,
587 #[serde(default)]
588 tool_calls: Option<Vec<ResponseToolCall>>,
589 }
590
591 impl ResponseMessage {
592 /// Reconstruct the assistant message for replay in history, preserving any
593 /// tool calls so the follow-up request is well-formed.
594 fn into_history_value(self) -> serde_json::Value {
595 let mut message = serde_json::json!({ "role": "assistant" });
596 message["content"] = match self.content {
597 Some(text) => serde_json::Value::String(text),
598 None => serde_json::Value::Null,
599 };
600 if let Some(tool_calls) = self.tool_calls {
601 message["tool_calls"] = serde_json::json!(
602 tool_calls
603 .into_iter()
604 .map(|call| serde_json::json!({
605 "id": call.id,
606 "type": "function",
607 "function": {
608 "name": call.function.name,
609 "arguments": call.function.arguments,
610 }
611 }))
612 .collect::<Vec<_>>()
613 );
614 }
615 message
616 }
617 }
618
619 #[derive(Debug, Deserialize)]
620 struct ResponseToolCall {
621 id: String,
622 function: ResponseFunction,
623 }
624
625 #[derive(Debug, Deserialize)]
626 struct ResponseFunction {
627 name: String,
628 #[serde(default)]
629 arguments: String,
630 }
631
632 // ── OpenAI streaming (SSE) chunk shapes ─────────────────────────────────────────
633
634 #[derive(Debug, Deserialize)]
635 struct StreamCompletion {
636 #[serde(default)]
637 choices: Vec<StreamChoice>,
638 }
639
640 #[derive(Debug, Deserialize)]
641 struct StreamChoice {
642 #[serde(default)]
643 delta: StreamDelta,
644 }
645
646 #[derive(Debug, Default, Deserialize)]
647 struct StreamDelta {
648 #[serde(default)]
649 content: Option<String>,
650 #[serde(default)]
651 tool_calls: Option<Vec<StreamToolCallDelta>>,
652 }
653
654 #[derive(Debug, Deserialize)]
655 struct StreamToolCallDelta {
656 #[serde(default)]
657 index: Option<u32>,
658 #[serde(default)]
659 id: Option<String>,
660 #[serde(default)]
661 function: Option<StreamFunctionDelta>,
662 }
663
664 #[derive(Debug, Deserialize)]
665 struct StreamFunctionDelta {
666 #[serde(default)]
667 name: Option<String>,
668 #[serde(default)]
669 arguments: Option<String>,
670 }
671
672 #[cfg(test)]
673 mod tests {
674 use super::*;
675
676 #[test]
677 fn tools_json_wraps_function_schema() {
678 let tools = vec![ToolSpec {
679 name: "read_file".to_string(),
680 description: "Read a file".to_string(),
681 parameters_schema: r#"{"type":"object","properties":{"path":{"type":"string"}}}"#
682 .to_string(),
683 }];
684 let json = OpenAiBackend::tools_json(&tools);
685 assert_eq!(json[0]["type"], "function");
686 assert_eq!(json[0]["function"]["name"], "read_file");
687 assert_eq!(
688 json[0]["function"]["parameters"]["properties"]["path"]["type"],
689 "string"
690 );
691 }
692
693 #[test]
694 fn malformed_schema_falls_back_to_empty_object() {
695 let tools = vec![ToolSpec {
696 name: "x".to_string(),
697 description: String::new(),
698 parameters_schema: "not json".to_string(),
699 }];
700 let json = OpenAiBackend::tools_json(&tools);
701 assert_eq!(json[0]["function"]["parameters"]["type"], "object");
702 }
703
704 #[test]
705 fn streamed_assistant_history_omits_empty_tool_calls() {
706 let value = streamed_assistant_history("hello", &[]);
707 assert_eq!(value["role"], "assistant");
708 assert_eq!(value["content"], "hello");
709 assert!(value.get("tool_calls").is_none());
710 }
711
712 #[test]
713 fn streamed_assistant_history_preserves_tool_calls() {
714 let calls = vec![ToolCall {
715 id: "call_0".to_string(),
716 name: "read_file".to_string(),
717 arguments: r#"{"path":"a.rs"}"#.to_string(),
718 }];
719 let value = streamed_assistant_history("", &calls);
720 assert!(value["content"].is_null());
721 assert_eq!(value["tool_calls"][0]["id"], "call_0");
722 assert_eq!(value["tool_calls"][0]["type"], "function");
723 assert_eq!(value["tool_calls"][0]["function"]["name"], "read_file");
724 assert_eq!(
725 value["tool_calls"][0]["function"]["arguments"],
726 r#"{"path":"a.rs"}"#
727 );
728 }
729
730 #[test]
731 fn assistant_message_with_tool_calls_round_trips() {
732 let message = ResponseMessage {
733 content: None,
734 tool_calls: Some(vec![ResponseToolCall {
735 id: "call_1".to_string(),
736 function: ResponseFunction {
737 name: "read_file".to_string(),
738 arguments: r#"{"path":"a.rs"}"#.to_string(),
739 },
740 }]),
741 };
742 let value = message.into_history_value();
743 assert_eq!(value["role"], "assistant");
744 assert!(value["content"].is_null());
745 assert_eq!(value["tool_calls"][0]["id"], "call_1");
746 assert_eq!(value["tool_calls"][0]["type"], "function");
747 assert_eq!(value["tool_calls"][0]["function"]["name"], "read_file");
748 }
749 }