ai-app: a phone interface to Claude Code and llama.cpp sessions
A Rust backend that owns the sessions and an Android app that reads them. The server spawns and adopts CLI processes, normalises everything they emit into one event model, keeps the transcript, and serves it over pinned TLS on a WireGuard interface; the phone streams that, replies, sends images, and imports conversations the machine already has. `AGENTS.md` is the working guide -- what runs where, what has been measured, and the faults that were expensive to find. `PLAN.md` is the design record. History before this point was squashed away. It was a personal project's running commentary and carried a name and a couple of machine paths that have no business in a public repository; the tree is what mattered and the tree is here.
This commit is contained in:
iris committed
commit
b172c464ea
100 files changed
+31795
No files matched your search
@@ -0,0 +1,811 @@
|
||||
//! The llama.cpp driver: a `llama-server` process per session, spoken to
|
||||
//! over its OpenAI-compatible HTTP API and translated into the common
|
||||
//! event model.
|
||||
//!
|
||||
//! Two things make this shaped differently from the Claude driver, and
|
||||
//! both are worth knowing before changing anything here.
|
||||
//!
|
||||
//! **It is spawned but not spoken to over stdio.** The process is started
|
||||
//! through the same [`Transport`] as any other, and then reached over
|
||||
//! HTTP on a loopback port. That is the case the transport's doc comment
|
||||
//! flags: a remote llama-server would need its port forwarded as well as
|
||||
//! its command wrapped, which is not built, so a session on an ssh host
|
||||
//! is refused rather than silently talking to the wrong machine.
|
||||
//!
|
||||
//! **The server is stateless between requests**, so the whole
|
||||
//! conversation goes with every one. It is rebuilt from the session's
|
||||
//! transcript rather than kept in this struct, which is not tidiness: a
|
||||
//! copy in driver memory is invisible to a second device and gone when
|
||||
//! this process restarts, and the app is meant to work across devices.
|
||||
//! The transcript is already the source of truth for everything else, and
|
||||
//! this makes it the source of truth for the prompt too.
|
||||
//!
|
||||
//! That leaves the Claude driver as the odd one out rather than this one:
|
||||
//! the CLI's own memory of a conversation is a cache in front of the same
|
||||
//! transcript, not a second truth. Anyone tempted to "fix" the
|
||||
//! inconsistency should resolve it in this direction.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
|
||||
use anyhow::{Context, Result, bail};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::json;
|
||||
|
||||
use super::driver::{Driver, Event, EventSink, ImageRef, SessionStatus};
|
||||
use super::process;
|
||||
use super::transport::{Launch, Streams, Transport};
|
||||
use crate::config::{ProviderConfig, SessionConfig};
|
||||
|
||||
/// How long to wait for a model to load before giving up on it. Loading
|
||||
/// is mostly disk, and a large quantised model on a cold cache is
|
||||
/// genuinely slow, so this is generous -- the failure it exists for is a
|
||||
/// server that will never answer, not one that is taking its time.
|
||||
const READY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300);
|
||||
|
||||
/// One turn in the conversation this driver keeps on the server's behalf.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
struct Message {
|
||||
role: String,
|
||||
content: String,
|
||||
}
|
||||
|
||||
pub struct LlamaDriver {
|
||||
sink: EventSink,
|
||||
/// Where this session's own llama-server answers.
|
||||
endpoint: String,
|
||||
/// Where the conversation is read back from, one line per event.
|
||||
transcript: PathBuf,
|
||||
/// Sampling settings chosen at spawn, sent with every request.
|
||||
sampling: serde_json::Map<String, serde_json::Value>,
|
||||
/// Set by [`Driver::interrupt`]; the streaming loop checks it between
|
||||
/// chunks and stops, leaving what was generated in the transcript.
|
||||
cancel: Arc<AtomicBool>,
|
||||
/// Where this session's process record lives, so [`Driver::stop`] can
|
||||
/// find the server it has to end.
|
||||
session_dir: PathBuf,
|
||||
}
|
||||
|
||||
impl LlamaDriver {
|
||||
/// Takes charge of this session's `llama-server`: the one already
|
||||
/// loaded if there is one, otherwise a new one.
|
||||
///
|
||||
/// One entry point, for the reason `ClaudeDriver::launch` gives -- the
|
||||
/// choice is not the caller's and a second process is the expensive
|
||||
/// mistake. Here it is expensive in a different currency: two servers
|
||||
/// holding the same model is twice the memory, and the second would
|
||||
/// bind a different port while the phone kept talking to the first.
|
||||
pub fn launch(
|
||||
meta: &SessionConfig,
|
||||
provider: &ProviderConfig,
|
||||
transport: &Transport,
|
||||
models_dir: &Path,
|
||||
transcript: &Path,
|
||||
session_dir: &Path,
|
||||
sink: EventSink,
|
||||
) -> Result<Self> {
|
||||
if !matches!(transport, Transport::Here) {
|
||||
bail!(
|
||||
"llama.cpp sessions can only run on this machine for now: the model is served \
|
||||
over HTTP, and forwarding that port to another host isn't built yet."
|
||||
);
|
||||
}
|
||||
let model = meta.model.as_deref().context(
|
||||
"a llama.cpp session needs a model -- one of the downloaded ones, by its key",
|
||||
)?;
|
||||
let path = model_path(models_dir, model)?;
|
||||
|
||||
// Already loaded and still running: keep talking to it. The
|
||||
// health poll below is what confirms it is really answering, so
|
||||
// adopting a pid whose server has wedged still reports as a
|
||||
// failure rather than as a session that silently never replies.
|
||||
if let Some(process::Record {
|
||||
detail: process::Detail::Http { port },
|
||||
pid,
|
||||
..
|
||||
}) = process::live(session_dir)
|
||||
{
|
||||
tracing::info!(
|
||||
"session {} reattaching to the llama-server it left loaded (pid {pid}, port {port})",
|
||||
meta.id
|
||||
);
|
||||
return Ok(Self::attached(
|
||||
format!("http://127.0.0.1:{port}"),
|
||||
meta,
|
||||
model,
|
||||
transcript,
|
||||
session_dir,
|
||||
sink,
|
||||
));
|
||||
}
|
||||
|
||||
let port = free_port().context("finding a port for llama-server")?;
|
||||
let mut args: Vec<String> = vec![
|
||||
"-m".into(),
|
||||
path.to_string_lossy().into_owned(),
|
||||
"--host".into(),
|
||||
"127.0.0.1".into(),
|
||||
"--port".into(),
|
||||
port.to_string(),
|
||||
];
|
||||
// Settings that belong to the server because they decide how the
|
||||
// model is loaded; the sampling ones ride on each request instead,
|
||||
// so changing them later needn't reload anything.
|
||||
for (key, flag) in [
|
||||
("contextSize", "-c"),
|
||||
("gpuLayers", "-ngl"),
|
||||
("threads", "-t"),
|
||||
] {
|
||||
if let Some(value) = meta.params.get(key) {
|
||||
args.push(flag.to_string());
|
||||
args.push(value.clone());
|
||||
}
|
||||
}
|
||||
|
||||
let program = provider.command.as_deref().unwrap_or("llama-server");
|
||||
let launch = Launch::new(program, args, meta.cwd.as_deref());
|
||||
// Its output goes to files, not pipes. Not only so the process can
|
||||
// outlive this server: nothing ever read those pipes, so a chatty
|
||||
// llama-server filled the 64 KB buffer and blocked mid-load with
|
||||
// no sign of why.
|
||||
let child = transport.spawn(
|
||||
&launch,
|
||||
Streams::Detached {
|
||||
stdin: std::process::Stdio::null(),
|
||||
stdout: log_file(&session_dir.join(SERVER_LOG))?.into(),
|
||||
stderr: log_file(&session_dir.join(SERVER_LOG))?.into(),
|
||||
},
|
||||
)?;
|
||||
let pid = child
|
||||
.id()
|
||||
.context("llama-server exited before it could be recorded")?;
|
||||
tracing::info!(
|
||||
"session {} running {program} for {model} on 127.0.0.1:{port} as pid {pid}",
|
||||
meta.id
|
||||
);
|
||||
// Reaped so it does not become a zombie while this server is still
|
||||
// its parent; the health poll and the record are what actually say
|
||||
// whether the session is alive, because after a restart there is no
|
||||
// `Child` here to ask.
|
||||
tokio::spawn(async move {
|
||||
let mut child = child;
|
||||
let _ = child.wait().await;
|
||||
});
|
||||
|
||||
let record = process::Record::of(pid, process::Detail::Http { port })
|
||||
.context("llama-server was gone before its start time could be read")?;
|
||||
process::write(session_dir, &record);
|
||||
|
||||
Ok(Self::attached(
|
||||
format!("http://127.0.0.1:{port}"),
|
||||
meta,
|
||||
model,
|
||||
transcript,
|
||||
session_dir,
|
||||
sink,
|
||||
))
|
||||
}
|
||||
|
||||
/// The driver for a `llama-server` at `endpoint`, however it got there.
|
||||
///
|
||||
/// Shared by starting one and adopting one, because everything after
|
||||
/// "there is a server at this address" is identical -- including
|
||||
/// waiting for it to answer, which an adopted one still owes: a
|
||||
/// recorded pid says a process exists, not that its model is loaded.
|
||||
fn attached(
|
||||
endpoint: String,
|
||||
meta: &SessionConfig,
|
||||
model: &str,
|
||||
transcript: &Path,
|
||||
session_dir: &Path,
|
||||
sink: EventSink,
|
||||
) -> Self {
|
||||
// Loading is slow enough to be worth saying so: the session shows
|
||||
// as running until the model is in memory, then goes idle, rather
|
||||
// than looking ready and refusing the first message.
|
||||
let _ = sink.send(Event::Status {
|
||||
state: SessionStatus::Running,
|
||||
});
|
||||
{
|
||||
let sink = sink.clone();
|
||||
let endpoint = endpoint.clone();
|
||||
let model = model.to_string();
|
||||
let session_dir = session_dir.to_path_buf();
|
||||
std::thread::spawn(move || match wait_until_ready(&endpoint) {
|
||||
Ok(()) => {
|
||||
tracing::info!("{model} loaded and answering at {endpoint}");
|
||||
let _ = sink.send(Event::Status {
|
||||
state: SessionStatus::Idle,
|
||||
});
|
||||
watch(session_dir, sink);
|
||||
}
|
||||
Err(err) => {
|
||||
let _ = sink.send(Event::Error {
|
||||
message: format!("{model} never became ready: {err:#}"),
|
||||
});
|
||||
let _ = sink.send(Event::Status {
|
||||
state: SessionStatus::Exited,
|
||||
});
|
||||
process::clear(&session_dir);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
let mut sampling = serde_json::Map::new();
|
||||
for (key, field) in [
|
||||
("temperature", "temperature"),
|
||||
("topP", "top_p"),
|
||||
("topK", "top_k"),
|
||||
("maxTokens", "max_tokens"),
|
||||
] {
|
||||
if let Some(raw) = meta.params.get(key)
|
||||
&& let Ok(number) = raw.parse::<f64>()
|
||||
{
|
||||
sampling.insert(field.to_string(), json!(number));
|
||||
}
|
||||
}
|
||||
|
||||
Self {
|
||||
sink,
|
||||
endpoint,
|
||||
transcript: transcript.to_path_buf(),
|
||||
sampling,
|
||||
cancel: Arc::new(AtomicBool::new(false)),
|
||||
session_dir: session_dir.to_path_buf(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Where llama-server's own output goes. One file for both streams: it is
|
||||
/// diagnostics nobody parses, and interleaving them is how it reads in a
|
||||
/// terminal anyway.
|
||||
const SERVER_LOG: &str = "llama-server.log";
|
||||
|
||||
/// How often a loaded server is checked for still being there.
|
||||
///
|
||||
/// Slower than the Claude driver's stdout poll because nothing is waiting
|
||||
/// on it: this only has to notice a server that has gone, and a few
|
||||
/// seconds late costs nothing.
|
||||
const WATCH_INTERVAL: std::time::Duration = std::time::Duration::from_secs(2);
|
||||
|
||||
/// An owner-only log opened for appending, so the two streams pointed at
|
||||
/// it do not overwrite each other and a reattach keeps what came before.
|
||||
fn log_file(path: &Path) -> Result<std::fs::File> {
|
||||
use std::os::unix::fs::OpenOptionsExt;
|
||||
std::fs::OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.mode(0o600)
|
||||
.open(path)
|
||||
.with_context(|| format!("opening {}", path.display()))
|
||||
}
|
||||
|
||||
/// Reports the server going away, for as long as the session is there to
|
||||
/// report it to.
|
||||
///
|
||||
/// Polled rather than waited on, for the reason the Claude driver gives:
|
||||
/// after a restart this server is not the process's parent and has nothing
|
||||
/// to wait on, so liveness has to be a question asked of the record -- and
|
||||
/// asking it two different ways is how the two answers come to disagree.
|
||||
fn watch(session_dir: PathBuf, sink: EventSink) {
|
||||
std::thread::spawn(move || {
|
||||
loop {
|
||||
std::thread::sleep(WATCH_INTERVAL);
|
||||
match process::recorded(&session_dir) {
|
||||
Some((_, process::Liveness::Alive)) => {}
|
||||
// Nothing recorded means the session was stopped or
|
||||
// deleted deliberately, and whoever did that has already
|
||||
// said so.
|
||||
None => return,
|
||||
Some((_, process::Liveness::Dead)) => {
|
||||
let _ = sink.send(Event::Error {
|
||||
message: "llama-server exited".to_string(),
|
||||
});
|
||||
let _ = sink.send(Event::Status {
|
||||
state: SessionStatus::Exited,
|
||||
});
|
||||
process::clear(&session_dir);
|
||||
return;
|
||||
}
|
||||
Some((_, process::Liveness::Unknown)) => {
|
||||
let _ = sink.send(Event::Status {
|
||||
state: SessionStatus::Unknown,
|
||||
});
|
||||
}
|
||||
}
|
||||
if sink.is_closed() {
|
||||
return;
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
impl Driver for LlamaDriver {
|
||||
fn send_user_message(&self, text: String, images: Vec<ImageRef>) {
|
||||
if !images.is_empty() {
|
||||
let _ = self.sink.send(Event::Error {
|
||||
message: "this model can't be sent images".to_string(),
|
||||
});
|
||||
}
|
||||
let sink = self.sink.clone();
|
||||
let endpoint = self.endpoint.clone();
|
||||
let transcript = self.transcript.clone();
|
||||
let sampling = self.sampling.clone();
|
||||
let cancel = Arc::clone(&self.cancel);
|
||||
cancel.store(false, Ordering::Relaxed);
|
||||
|
||||
// Its own thread: the request blocks for as long as the model
|
||||
// takes to generate, which is the whole point of streaming it.
|
||||
std::thread::spawn(move || {
|
||||
// Nothing is ever held back here -- there is no queue to wait
|
||||
// in -- so the message is taken the moment it arrives. Said
|
||||
// anyway, because this is what records it: see `MessageTaken`.
|
||||
let _ = sink.send(Event::MessageTaken {
|
||||
id: None,
|
||||
text: text.clone(),
|
||||
// Never any: this driver refuses images above, and saying
|
||||
// so is what the refusal above is for.
|
||||
images: Vec::new(),
|
||||
});
|
||||
let _ = sink.send(Event::Status {
|
||||
state: SessionStatus::Running,
|
||||
});
|
||||
// Everything before this message, plus this message. Read
|
||||
// rather than remembered, and `text` is appended here rather
|
||||
// than waited for, because the message's own transcript entry
|
||||
// is still on its way when this runs.
|
||||
let mut messages = conversation(&transcript);
|
||||
messages.push(Message {
|
||||
role: "user".into(),
|
||||
content: text,
|
||||
});
|
||||
// The reply is not stored: the deltas below are the durable
|
||||
// record, so the next turn reads back exactly what the phone
|
||||
// was shown -- including a partial one that was interrupted.
|
||||
if let Err(err) = generate(&endpoint, &messages, &sampling, &cancel, &sink) {
|
||||
let _ = sink.send(Event::Error {
|
||||
message: format!("{err:#}"),
|
||||
});
|
||||
}
|
||||
let _ = sink.send(Event::Status {
|
||||
state: SessionStatus::Idle,
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
fn answer_question(&self, _id: &str, _answers: &[String]) {
|
||||
// Nothing here asks questions: this driver has no tools, so no
|
||||
// permission prompts and no AskUserQuestion.
|
||||
}
|
||||
|
||||
fn interrupt(&self) {
|
||||
self.cancel.store(true, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
// Nothing to forward: this process has no notion of what the
|
||||
// conversation is called, and the rename it belongs to has already
|
||||
// happened where the name lives. See `Driver::set_title`.
|
||||
fn set_title(&self, _title: &str) {}
|
||||
|
||||
fn set_permission_mode(&self, _mode: &str) {
|
||||
let _ = self.sink.send(Event::Error {
|
||||
message: "a llama.cpp session runs no tools, so there is nothing for a permission \
|
||||
mode to govern."
|
||||
.to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
fn set_model(&self, _model: &str) {
|
||||
let _ = self.sink.send(Event::Error {
|
||||
message: "a llama.cpp session's model is fixed when it starts, because the server \
|
||||
loads one model into memory. Spawn another session to use a different one."
|
||||
.to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
fn run_command(&self, text: &str) {
|
||||
let _ = self.sink.send(Event::Error {
|
||||
message: format!(
|
||||
"a llama.cpp session has no commands of its own, so {text} means nothing to it."
|
||||
),
|
||||
});
|
||||
}
|
||||
|
||||
fn compact(&self) {
|
||||
let _ = self.sink.send(Event::Error {
|
||||
message: "llama.cpp has no compaction. Clear the session instead, which costs nothing."
|
||||
.to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
fn clear(&self) {
|
||||
// All of it. `conversation` folds from the last of these, so
|
||||
// recording the marker *is* the reset -- there is no driver state
|
||||
// to keep in step with it, which is the same property that makes
|
||||
// a second device see the same conversation this one does.
|
||||
let _ = self.sink.send(Event::Cleared);
|
||||
}
|
||||
|
||||
/// Stops generating and leaves the server loaded.
|
||||
///
|
||||
/// Worth being deliberate about, because the cost is asymmetric and
|
||||
/// points the other way from the Claude driver's: a `llama-server`
|
||||
/// holds its whole model in memory, so a leaked one is gigabytes
|
||||
/// nobody is using. It is left anyway, because the alternative is
|
||||
/// unloading and reloading that model on every backend restart --
|
||||
/// minutes of disk, for a session somebody is in the middle of. The
|
||||
/// record is what keeps it from being *nobody's*: the next run of this
|
||||
/// server adopts it rather than starting a second one.
|
||||
fn detach(&self) {
|
||||
self.cancel.store(true, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
fn stop(&self) {
|
||||
self.cancel.store(true, Ordering::Relaxed);
|
||||
if let Some(record) = process::live(&self.session_dir) {
|
||||
process::stop(&record, process::STOP_GRACE);
|
||||
}
|
||||
process::clear(&self.session_dir);
|
||||
}
|
||||
}
|
||||
|
||||
/// The conversation so far, folded out of the transcript.
|
||||
///
|
||||
/// Consecutive `AssistantText` deltas are one assistant turn, closed by
|
||||
/// the next user message -- which is also what makes an interrupted reply
|
||||
/// come back as the partial text the phone actually saw, rather than
|
||||
/// vanishing or being invented.
|
||||
///
|
||||
/// This must stay a pure function of the transcript and must never
|
||||
/// re-render earlier turns. llama.cpp caches the prompt prefix, so a
|
||||
/// growing conversation reprocesses almost nothing -- but only while
|
||||
/// every turn is byte-identical to last time. Changing how an old turn is
|
||||
/// rendered silently reprocesses the whole history on every message.
|
||||
fn conversation(path: &Path) -> Vec<Message> {
|
||||
let Ok(events) = crate::session::transcript::read_after(path, 0) else {
|
||||
return Vec::new();
|
||||
};
|
||||
let mut messages: Vec<Message> = Vec::new();
|
||||
let mut pending = String::new();
|
||||
// Everything before the last clear is still in the transcript and is
|
||||
// deliberately not in the conversation. Folding from zero here would
|
||||
// put it back, which is the whole of what clearing had to undo.
|
||||
let events = match events.iter().rposition(|e| e.event == Event::Cleared) {
|
||||
Some(at) => &events[at + 1..],
|
||||
None => &events[..],
|
||||
};
|
||||
for event in events.iter().cloned() {
|
||||
match event.event {
|
||||
Event::UserMessage { text, .. } => {
|
||||
if !pending.is_empty() {
|
||||
messages.push(Message {
|
||||
role: "assistant".into(),
|
||||
content: std::mem::take(&mut pending),
|
||||
});
|
||||
}
|
||||
messages.push(Message {
|
||||
role: "user".into(),
|
||||
content: text,
|
||||
});
|
||||
}
|
||||
Event::AssistantText { delta } => pending.push_str(&delta),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
if !pending.is_empty() {
|
||||
messages.push(Message {
|
||||
role: "assistant".into(),
|
||||
content: pending,
|
||||
});
|
||||
}
|
||||
messages
|
||||
}
|
||||
|
||||
/// Where a model key resolves to on disk, refusing anything that climbs
|
||||
/// out of the models directory -- the key arrives from a phone.
|
||||
fn model_path(models_dir: &Path, key: &str) -> Result<PathBuf> {
|
||||
let mut path = models_dir.to_path_buf();
|
||||
for part in key.split('/') {
|
||||
if part.is_empty() || part == "." || part == ".." {
|
||||
bail!("\"{key}\" is not a model key this can resolve");
|
||||
}
|
||||
path.push(part);
|
||||
}
|
||||
if !path.is_file() {
|
||||
bail!("no downloaded model called \"{key}\" -- download it first");
|
||||
}
|
||||
Ok(path)
|
||||
}
|
||||
|
||||
/// An unused loopback port, by asking the OS for one and letting it go.
|
||||
///
|
||||
/// Racy in principle: something else could take it between here and
|
||||
/// llama-server binding. In practice nothing on this machine is hunting
|
||||
/// for ports, and the alternative -- parsing the port back out of the
|
||||
/// server's log -- couples us to its output format for no real gain.
|
||||
fn free_port() -> Result<u16> {
|
||||
let listener = std::net::TcpListener::bind("127.0.0.1:0")?;
|
||||
Ok(listener.local_addr()?.port())
|
||||
}
|
||||
|
||||
/// Polls until the server says it is ready, or gives up.
|
||||
fn wait_until_ready(endpoint: &str) -> Result<()> {
|
||||
let deadline = std::time::Instant::now() + READY_TIMEOUT;
|
||||
let url = format!("{endpoint}/health");
|
||||
loop {
|
||||
if let Ok(response) = ureq::get(&url).call()
|
||||
&& response.status() == 200
|
||||
{
|
||||
return Ok(());
|
||||
}
|
||||
if std::time::Instant::now() > deadline {
|
||||
bail!("gave up after {}s", READY_TIMEOUT.as_secs());
|
||||
}
|
||||
std::thread::sleep(std::time::Duration::from_millis(250));
|
||||
}
|
||||
}
|
||||
|
||||
/// One streamed completion: posts the conversation, emits each delta as it
|
||||
/// arrives. Emits rather than returns: the transcript those events land
|
||||
/// in is what the next turn reads back, so there is nothing to hand up.
|
||||
fn generate(
|
||||
endpoint: &str,
|
||||
messages: &[Message],
|
||||
sampling: &serde_json::Map<String, serde_json::Value>,
|
||||
cancel: &AtomicBool,
|
||||
sink: &EventSink,
|
||||
) -> Result<()> {
|
||||
let mut body = json!({
|
||||
"messages": messages,
|
||||
"stream": true,
|
||||
"stream_options": {"include_usage": true},
|
||||
});
|
||||
let map = body.as_object_mut().expect("built as an object");
|
||||
for (key, value) in sampling {
|
||||
map.insert(key.clone(), value.clone());
|
||||
}
|
||||
|
||||
let mut response = ureq::post(format!("{endpoint}/v1/chat/completions"))
|
||||
.header("Content-Type", "application/json")
|
||||
.send_json(&body)
|
||||
.context("asking llama-server to generate")?;
|
||||
|
||||
let reader = std::io::BufReader::new(response.body_mut().as_reader());
|
||||
let mut tokens = 0u64;
|
||||
// The prompt side only, which is what the model is holding -- the same
|
||||
// definition the other dialects report, so one word on the phone means
|
||||
// one thing whichever kind of session it is.
|
||||
let mut context = None;
|
||||
for line in std::io::BufRead::lines(reader) {
|
||||
if cancel.load(Ordering::Relaxed) {
|
||||
break;
|
||||
}
|
||||
let line = line.context("reading the generation stream")?;
|
||||
// Server-sent events: the payload lines are the ones that matter,
|
||||
// and blank lines separate events.
|
||||
let Some(payload) = line.strip_prefix("data: ") else {
|
||||
continue;
|
||||
};
|
||||
if payload.trim() == "[DONE]" {
|
||||
break;
|
||||
}
|
||||
let Ok(chunk) = serde_json::from_str::<serde_json::Value>(payload) else {
|
||||
continue;
|
||||
};
|
||||
if let Some(usage) = chunk.get("usage") {
|
||||
if let Some(total) = usage
|
||||
.get("total_tokens")
|
||||
.and_then(serde_json::Value::as_u64)
|
||||
{
|
||||
tokens = total;
|
||||
}
|
||||
if let Some(prompt) = usage
|
||||
.get("prompt_tokens")
|
||||
.and_then(serde_json::Value::as_u64)
|
||||
{
|
||||
context = Some(prompt);
|
||||
}
|
||||
}
|
||||
let delta = chunk
|
||||
.get("choices")
|
||||
.and_then(|c| c.get(0))
|
||||
.and_then(|c| c.get("delta"))
|
||||
.and_then(|d| d.get("content"))
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.unwrap_or_default();
|
||||
if !delta.is_empty() {
|
||||
let _ = sink.send(Event::AssistantText {
|
||||
delta: delta.to_string(),
|
||||
});
|
||||
}
|
||||
}
|
||||
if tokens > 0 {
|
||||
let _ = sink.send(Event::UsageDelta { tokens, context });
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::session::transcript::Transcript;
|
||||
|
||||
/// Writes a transcript the way the pump does, so the fold is tested
|
||||
/// against the real file format rather than a hand-built vector.
|
||||
fn transcript_with(events: &[Event]) -> (tempfile::TempDir, PathBuf) {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let path = dir.path().join("transcript.jsonl");
|
||||
let mut transcript = Transcript::open(&path).expect("open");
|
||||
for event in events {
|
||||
transcript.append(event.clone(), 0.0).expect("append");
|
||||
}
|
||||
(dir, path)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deltas_between_user_messages_are_one_assistant_turn() {
|
||||
let (_dir, path) = transcript_with(&[
|
||||
Event::UserMessage {
|
||||
id: None,
|
||||
text: "hello".into(),
|
||||
images: Vec::new(),
|
||||
},
|
||||
Event::AssistantText {
|
||||
delta: "hi ".into(),
|
||||
},
|
||||
Event::AssistantText {
|
||||
delta: "there".into(),
|
||||
},
|
||||
Event::Status {
|
||||
state: SessionStatus::Idle,
|
||||
},
|
||||
Event::UserMessage {
|
||||
id: None,
|
||||
text: "again".into(),
|
||||
images: Vec::new(),
|
||||
},
|
||||
Event::AssistantText {
|
||||
delta: "yes".into(),
|
||||
},
|
||||
]);
|
||||
let messages = conversation(&path);
|
||||
assert_eq!(
|
||||
messages
|
||||
.iter()
|
||||
.map(|m| (m.role.as_str(), m.content.as_str()))
|
||||
.collect::<Vec<_>>(),
|
||||
[
|
||||
("user", "hello"),
|
||||
("assistant", "hi there"),
|
||||
("user", "again"),
|
||||
("assistant", "yes")
|
||||
],
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
/// The interrupted case, which decides what a resumed conversation is
|
||||
/// built from: whatever the phone was shown. The deltas that arrived
|
||||
/// before the stop are in the transcript, so they are in the prompt --
|
||||
/// the model is never told it said something the user did not see, and
|
||||
/// never has a turn silently dropped from under it.
|
||||
fn an_interrupted_reply_stays_in_the_conversation() {
|
||||
let (_dir, path) = transcript_with(&[
|
||||
Event::UserMessage {
|
||||
id: None,
|
||||
text: "count".into(),
|
||||
images: Vec::new(),
|
||||
},
|
||||
Event::AssistantText {
|
||||
delta: "one two".into(),
|
||||
},
|
||||
Event::Status {
|
||||
state: SessionStatus::Idle,
|
||||
},
|
||||
]);
|
||||
let messages = conversation(&path);
|
||||
assert_eq!(messages.len(), 2);
|
||||
assert_eq!(messages[1].content, "one two");
|
||||
}
|
||||
|
||||
#[test]
|
||||
/// Events this driver does not produce must not disturb the fold: a
|
||||
/// transcript can carry errors and status changes from a session that
|
||||
/// was, say, relaunched.
|
||||
fn other_events_are_not_part_of_the_conversation() {
|
||||
let (_dir, path) = transcript_with(&[
|
||||
Event::Status {
|
||||
state: SessionStatus::Running,
|
||||
},
|
||||
Event::UserMessage {
|
||||
id: None,
|
||||
text: "hello".into(),
|
||||
images: Vec::new(),
|
||||
},
|
||||
Event::Error {
|
||||
message: "something went wrong".into(),
|
||||
},
|
||||
Event::AssistantText {
|
||||
delta: "still here".into(),
|
||||
},
|
||||
Event::UsageDelta {
|
||||
tokens: 12,
|
||||
context: Some(12),
|
||||
},
|
||||
]);
|
||||
let messages = conversation(&path);
|
||||
assert_eq!(messages.len(), 2);
|
||||
assert_eq!(messages[0].content, "hello");
|
||||
assert_eq!(messages[1].content, "still here");
|
||||
}
|
||||
|
||||
#[test]
|
||||
/// Clearing decides what the *model* is given, not just what the
|
||||
/// phone draws. Everything above the marker stays in the transcript
|
||||
/// -- a person can still scroll back to it -- and none of it is sent.
|
||||
fn the_conversation_starts_after_the_last_clear() {
|
||||
let (_dir, path) = transcript_with(&[
|
||||
Event::UserMessage {
|
||||
id: None,
|
||||
text: "the long expensive conversation".into(),
|
||||
images: Vec::new(),
|
||||
},
|
||||
Event::AssistantText {
|
||||
delta: "at length".into(),
|
||||
},
|
||||
Event::Cleared,
|
||||
Event::UserMessage {
|
||||
id: None,
|
||||
text: "a fresh start".into(),
|
||||
images: Vec::new(),
|
||||
},
|
||||
Event::AssistantText {
|
||||
delta: "cheaply".into(),
|
||||
},
|
||||
]);
|
||||
let messages = conversation(&path);
|
||||
assert_eq!(messages.len(), 2);
|
||||
assert_eq!(messages[0].content, "a fresh start");
|
||||
assert_eq!(messages[1].content, "cheaply");
|
||||
}
|
||||
|
||||
#[test]
|
||||
/// The *last* one, so clearing twice does not resurrect what the
|
||||
/// first clear dropped.
|
||||
fn only_the_newest_clear_counts() {
|
||||
let (_dir, path) = transcript_with(&[
|
||||
Event::UserMessage {
|
||||
id: None,
|
||||
text: "one".into(),
|
||||
images: Vec::new(),
|
||||
},
|
||||
Event::Cleared,
|
||||
Event::UserMessage {
|
||||
id: None,
|
||||
text: "two".into(),
|
||||
images: Vec::new(),
|
||||
},
|
||||
Event::Cleared,
|
||||
Event::UserMessage {
|
||||
id: None,
|
||||
text: "three".into(),
|
||||
images: Vec::new(),
|
||||
},
|
||||
]);
|
||||
let messages = conversation(&path);
|
||||
assert_eq!(messages.len(), 1);
|
||||
assert_eq!(messages[0].content, "three");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_model_key_cannot_climb_out_of_the_models_directory() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
for attempt in ["../../etc/passwd", "unsloth/../../escape.gguf", ""] {
|
||||
assert!(
|
||||
model_path(dir.path(), attempt).is_err(),
|
||||
"{attempt:?} should have been refused",
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user