//! Minimal OpenAI-compatible chat client (llama-server / LiteLLM proxy).
use anyhow::{Context, Result};
use base64::Engine as _;
use futures_util::StreamExt as _;
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::path::Path;
use std::time::Duration;
use crate::config;
use crate::tools::{self, Policy, ToolCall};
#[derive(Serialize)]
struct ChatRequest<'a> {
model: &'a str,
messages: &'a [ChatMessage],
#[serde(skip_serializing_if = "Option::is_none")]
tools: Option<serde_json::Value>,
max_tokens: u32,
temperature: f32,
stream: bool,
}
#[derive(Serialize, Deserialize, Clone)]
pub struct ChatMessage {
pub role: String,
/// Null for assistant turns that only carry tool calls.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub content: Option<Content>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub tool_calls: Vec<ToolCall>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub tool_call_id: Option<String>,
}
impl ChatMessage {
fn system(text: &str) -> Self {
Self {
role: "system".into(),
content: Some(Content::Text(text.to_string())),
tool_calls: Vec::new(),
tool_call_id: None,
}
}
fn user(content: Content) -> Self {
Self {
role: "user".into(),
content: Some(content),
tool_calls: Vec::new(),
tool_call_id: None,
}
}
fn assistant(content: Option<String>, tool_calls: Vec<ToolCall>) -> Self {
Self {
role: "assistant".into(),
content: content.map(Content::Text),
tool_calls,
tool_call_id: None,
}
}
fn tool_result(id: String, text: String) -> Self {
Self {
role: "tool".into(),
content: Some(Content::Text(text)),
tool_calls: Vec::new(),
tool_call_id: Some(id),
}
}
}
/// Message content: plain text, or a list of text/image parts.
#[derive(Serialize, Deserialize, Clone)]
#[serde(untagged)]
pub enum Content {
Text(String),
Parts(Vec<Part>),
}
#[derive(Serialize, Deserialize, Clone)]
#[serde(tag = "type")]
pub enum Part {
#[serde(rename = "text")]
Text { text: String },
#[serde(rename = "image_url")]
ImageUrl { image_url: ImageUrl },
}
#[derive(Serialize, Deserialize, Clone)]
pub struct ImageUrl {
pub url: String,
}
/// A local image to attach to the prompt.
pub struct Image {
pub mime: &'static str,
pub data: Vec<u8>,
}
impl Image {
/// Guess the MIME type from the path's extension. Returns None when the
/// extension isn't a supported image type, so callers can decide whether
/// that's an error (one-shot argv) or just not an image (REPL input).
pub fn mime_for(path: &Path) -> Option<&'static str> {
match path
.extension()
.and_then(|e| e.to_str())
.unwrap_or_default()
.to_ascii_lowercase()
.as_str()
{
"png" => Some("image/png"),
"jpg" | "jpeg" => Some("image/jpeg"),
"gif" => Some("image/gif"),
"webp" => Some("image/webp"),
_ => None,
}
}
/// Read an image from disk, guessing the MIME type from the extension.
pub fn load(path: &Path) -> Result<Self> {
let mime = Self::mime_for(path).ok_or_else(|| {
let ext = path
.extension()
.and_then(|e| e.to_str())
.unwrap_or_default();
anyhow::anyhow!(
"unsupported image type '.{ext}' (supported: png, jpg, jpeg, gif, webp)"
)
})?;
let data =
std::fs::read(path).with_context(|| format!("reading image {}", path.display()))?;
if data.is_empty() {
anyhow::bail!("image {} is empty", path.display());
}
if data.len() > config::MAX_IMAGE_BYTES {
anyhow::bail!(
"image {} is {} MB, over the {} MB limit",
path.display(),
data.len() / 1_000_000,
config::MAX_IMAGE_BYTES / 1_000_000
);
}
Ok(Self { mime, data })
}
fn data_url(&self) -> String {
let b64 = base64::engine::general_purpose::STANDARD.encode(&self.data);
format!("data:{};base64,{b64}", self.mime)
}
}
/// Streamed chunk of a completion.
#[derive(Deserialize)]
struct StreamChunk {
#[serde(default)]
choices: Vec<StreamChoice>,
}
#[derive(Deserialize)]
struct StreamChoice {
#[serde(default)]
delta: Option<StreamDelta>,
}
#[derive(Deserialize)]
struct StreamDelta {
#[serde(default)]
content: Option<Content>,
#[serde(default)]
tool_calls: Vec<ToolCallDelta>,
}
/// Tool calls arrive as fragments spread across many frames, keyed by `index`.
/// The name and id show up once; `arguments` is a JSON string built up a few
/// characters at a time.
#[derive(Deserialize)]
struct ToolCallDelta {
#[serde(default)]
index: usize,
#[serde(default)]
id: Option<String>,
#[serde(default)]
function: Option<FunctionDelta>,
}
#[derive(Deserialize)]
struct FunctionDelta {
#[serde(default)]
name: Option<String>,
#[serde(default)]
arguments: Option<String>,
}
/// Accumulates fragmented tool calls, keyed by their wire index.
#[derive(Default)]
struct ToolCallAccumulator {
calls: BTreeMap<usize, ToolCall>,
}
impl ToolCallAccumulator {
fn push(&mut self, deltas: Vec<ToolCallDelta>) {
for delta in deltas {
let call = self.calls.entry(delta.index).or_insert_with(|| ToolCall {
id: String::new(),
kind: "function".into(),
function: crate::tools::FunctionCall {
name: String::new(),
arguments: String::new(),
},
});
if let Some(id) = delta.id {
call.id = id;
}
if let Some(function) = delta.function {
if let Some(name) = function.name {
call.function.name = name;
}
if let Some(arguments) = function.arguments {
call.function.arguments.push_str(&arguments);
}
}
}
}
fn finish(self) -> Vec<ToolCall> {
let mut calls: Vec<ToolCall> = self.calls.into_values().collect();
for (n, call) in calls.iter_mut().enumerate() {
// Some servers omit ids; synthesize one so the tool reply still matches.
if call.id.is_empty() {
call.id = format!("call_{n}");
}
}
calls
}
}
/// An ongoing conversation: holds the message history and streams each reply.
///
/// One `send` may span several requests: if the model asks to call tools, the
/// calls run, their results go back as `tool` messages, and generation
/// continues until it produces a plain answer.
pub struct Chat {
client: reqwest::Client,
messages: Vec<ChatMessage>,
policy: Policy,
}
impl Chat {
pub fn new(policy: Policy) -> Result<Self> {
let client = reqwest::Client::builder()
.connect_timeout(Duration::from_secs(config::CONNECT_TIMEOUT_SECS))
// Per-chunk idle timeout, not a total cap: long generations are
// fine, but a stalled connection fails instead of hanging forever.
.read_timeout(Duration::from_secs(config::READ_TIMEOUT_SECS))
.build()
.context("building HTTP client")?;
let mut chat = Self {
client,
messages: Vec::new(),
policy,
};
chat.reset();
Ok(chat)
}
/// Drop the history and start over from the system prompt.
pub fn reset(&mut self) {
self.messages.clear();
self.messages
.push(ChatMessage::system(config::SYSTEM_PROMPT));
}
/// Number of turns so far, excluding the system prompt.
pub fn turn_count(&self) -> usize {
self.messages.len().saturating_sub(1)
}
/// Append a user message and stream the reply into `on_item`, which
/// receives text deltas and tool events as one ordered stream.
///
/// If a request fails, every message from this turn is rolled back so the
/// history stays consistent with what the model actually saw.
pub async fn send(
&mut self,
user_message: &str,
images: &[Image],
mut on_item: impl FnMut(Item),
) -> Result<()> {
let mark = self.messages.len();
self.messages
.push(ChatMessage::user(user_content(user_message, images)));
match self.run_turn(&mut on_item).await {
Ok(()) => Ok(()),
Err(err) => {
self.messages.truncate(mark);
Err(err)
}
}
}
async fn run_turn(&mut self, on_item: &mut impl FnMut(Item)) -> Result<()> {
for round in 0..config::MAX_TOOL_ROUNDS {
let (reply, tool_calls) = self.stream_reply(on_item).await?;
if tool_calls.is_empty() {
self.messages.push(ChatMessage::assistant(
(!reply.is_empty()).then_some(reply),
Vec::new(),
));
return Ok(());
}
// Out of rounds: stop calling tools but still ask for an answer,
// so everything gathered so far gets summarized rather than lost.
if round + 1 == config::MAX_TOOL_ROUNDS {
self.messages
.push(ChatMessage::assistant(Some(reply), tool_calls));
self.messages.push(ChatMessage::user(Content::Text(
"Tool budget exhausted. Answer using only what you have gathered so far."
.into(),
)));
let (final_reply, _) = self.stream_reply(on_item).await?;
self.messages.push(ChatMessage::assistant(
(!final_reply.is_empty()).then_some(final_reply),
Vec::new(),
));
return Ok(());
}
// Record the assistant turn, tool calls included: the server needs
// them to match the results that follow.
self.messages
.push(ChatMessage::assistant(Some(reply), tool_calls.clone()));
for call in &tool_calls {
on_item(Item::ToolStart {
name: call.function.name.clone(),
summary: call.summary(),
});
let result = tools::execute(call, self.policy, &mut blocking_confirm).await;
on_item(Item::ToolDone {
outcome: outcome_summary(&result),
});
self.messages
.push(ChatMessage::tool_result(call.id.clone(), result));
}
}
Ok(())
}
async fn stream_reply(
&self,
on_item: &mut impl FnMut(Item),
) -> Result<(String, Vec<ToolCall>)> {
let body = ChatRequest {
model: config::MODEL,
messages: &self.messages,
tools: Some(tools::specs()),
max_tokens: config::MAX_TOKENS,
temperature: config::TEMPERATURE,
stream: true,
};
let url = format!("{}/v1/chat/completions", config::BASE_URL);
let resp = self
.client
.post(&url)
.bearer_auth(config::api_key())
.json(&body)
.send()
.await
.with_context(|| format!("request to {url}"))?;
let status = resp.status();
if !status.is_success() {
let text = resp.text().await.context("reading error body")?;
anyhow::bail!(
"HTTP {status}: {}",
text.chars().take(4000).collect::<String>()
);
}
let mut stream = resp.bytes_stream();
let mut reply = String::new();
let mut buf = String::new();
let mut accumulator = ToolCallAccumulator::default();
'outer: while let Some(item) = stream.next().await {
let bytes = item.context("reading stream chunk")?;
buf.push_str(&String::from_utf8_lossy(&bytes));
// SSE frames are line-delimited; the last one is incomplete until more arrives.
while let Some(pos) = buf.find('\n') {
let line: String = buf.drain(..=pos).collect();
let line = line.trim();
let Some(data) = line.strip_prefix("data:").map(str::trim) else {
continue; // comments, "event:" lines, blank lines
};
if data == "[DONE]" {
break 'outer;
}
let Ok(chunk) = serde_json::from_str::<StreamChunk>(data) else {
continue; // tolerate frames we don't understand
};
for choice in chunk.choices {
let Some(delta) = choice.delta else { continue };
if let Some(text) = delta_content(delta.content).filter(|t| !t.is_empty()) {
reply.push_str(&text);
on_item(Item::Text(text));
}
accumulator.push(delta.tool_calls);
}
}
}
Ok((reply, accumulator.finish()))
}
}
/// Something surfaced to the UI mid-turn: a chunk of the model's answer, or a
/// note that a tool is running and how it turned out. One type so callers can
/// hand a single callback to `Chat::send`.
pub enum Item {
Text(String),
ToolStart { name: String, summary: String },
ToolDone { outcome: String },
}
/// One line describing how a tool call went, for the status display.
fn outcome_summary(result: &str) -> String {
let trimmed = result.trim_start();
if let Some(rest) = trimmed.strip_prefix("Error:") {
return format!(
"error: {}",
rest.trim().chars().take(70).collect::<String>()
);
}
if trimmed.starts_with("Skipped:") {
return "skipped".into();
}
let lines = result.lines().count();
let first = result.trim().lines().next().unwrap_or("");
let shown: String = first.chars().take(48).collect();
format!("{lines} line(s) — {shown}")
}
/// Confirmation has to block on terminal input. That's safe here: it only
/// happens while a tool is running, with no other request in flight, so
/// parking the task costs nothing.
///
/// Defaults to yes: a bare Enter proceeds, `n` declines.
fn blocking_confirm(_call: &ToolCall) -> bool {
use std::io::Write as _;
eprint!(" run it? [Y/n] ");
let _ = std::io::stderr().flush();
let mut answer = String::new();
if std::io::stdin().read_line(&mut answer).unwrap_or(0) == 0 {
return true; // EOF, nothing left to ask
}
!matches!(answer.trim(), "n" | "N" | "no" | "No")
}
/// Build the content for a user turn: plain text, or text plus image parts.
/// Empty text with images still yields image-only parts.
fn user_content(user_message: &str, images: &[Image]) -> Content {
if images.is_empty() {
return Content::Text(user_message.to_string());
}
let mut parts = Vec::with_capacity(images.len() + 1);
if !user_message.is_empty() {
parts.push(Part::Text {
text: user_message.to_string(),
});
}
for image in images {
parts.push(Part::ImageUrl {
image_url: ImageUrl {
url: image.data_url(),
},
});
}
Content::Parts(parts)
}
/// Delta content is a string on every server we target, but some emit content
/// parts; flatten either shape.
fn delta_content(content: Option<Content>) -> Option<String> {
match content? {
Content::Text(text) => Some(text),
Content::Parts(parts) => {
let mut out = String::new();
for part in parts {
if let Part::Text { text } = part {
out.push_str(&text);
}
}
Some(out)
}
}
}