Something went wrong. Try again.
A 5e storytelling engine with an LLM DM
Something went wrong. Try again.
12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079use std::collections::VecDeque;use std::fmt;use std::io;use std::time::Duration;
use crate::chat::sse::SseReader;use crate::chat::wire::{ ChatChunk, ChatRequest, Message, ToolCall, ToolCallDelta, ToolCallFunction,};
/// The maximum length of a status error's body, in characters.const STATUS_ERROR_BODY_LIMIT: usize = 2_000;
/// How many times the client sends one request before it gives up.const MAX_ATTEMPTS: u32 = 3;
/// The time limits on one chat completion request.////// `stream` uses the defaults. The tests build shorter ones so a stalled/// server or a backoff does not hold the suite.#[derive(Clone, Copy)]struct Limits { /// How long the socket and the TLS handshake may take. connect: Duration, /// How long the client waits for the response headers after it sends /// the request. response: Duration, /// How long the whole streamed body may take. body: Duration, /// How long to wait before the second attempt. Each further attempt /// doubles it. backoff: Duration,}
impl Default for Limits { fn default() -> Self { Self { connect: Duration::from_secs(10), // This waits for the first bytes of the reply, and a // reasoning model spends its whole deliberation before any // byte reaches the wire: the provider holds the stream // silent while the model thinks, so a high reasoning effort // makes minutes of silence normal. Like `body`, this is a // bound on a dead provider, not a budget for the model. response: Duration::from_secs(300), // ureq measures this over the whole body, and the body of a // streamed turn is the whole turn. So this is a bound on a // dead connection, not a budget for the model: a healthy turn // finishes well inside 10 minutes, and a provider that stops // writing mid-stream frees the worker thread at 10 minutes // rather than never. body: Duration::from_secs(600), backoff: Duration::from_millis(500), } }}
/// Reports whether a status is worth sending the request again.////// 429 means the provider is rate limiting, and a 5xx is a failure on its/// side that the next attempt may not hit. Every other status describes/// the request itself, so a second attempt gives the same answer.fn worth_retrying(status: u16) -> bool { status == 429 || (500..600).contains(&status)}
/// A client for one chat completion endpoint.pub struct Client { pub api_base: String, pub api_key: String, pub model: String, /// How much the model deliberates before it answers, sent to the /// provider verbatim on every request. `None` omits the field, so the /// provider's own default applies. pub reasoning_effort: Option<String>,}
impl Client { /// Starts a streaming chat completion for `messages`, offering `tools`. /// /// Sends `messages` with `self.model` to `{api_base}/chat/completions`. /// A trailing slash on `api_base` is tolerated. An empty `tools` omits /// the field from the request rather than sending an empty array. /// /// A rate limit or a server error costs one of `MAX_ATTEMPTS` tries, /// with a doubling wait between them. Retrying here is safe because no /// byte has reached the caller yet: the caller sees one request, and a /// failure leaves it with nothing to undo. Once the body starts, the /// client never sends the request again. pub fn stream( &self, messages: &[Message], tools: &[serde_json::Value], ) -> Result<ChatStream, ChatError> { self.stream_within(messages, tools, Limits::default()) }
/// Runs `stream` under `limits`. fn stream_within( &self, messages: &[Message], tools: &[serde_json::Value], limits: Limits, ) -> Result<ChatStream, ChatError> { let request = ChatRequest { model: self.model.clone(), messages: messages.to_vec(), stream: true, tools: tools.to_vec(), reasoning_effort: self.reasoning_effort.clone(), }; let body = serde_json::to_vec(&request).expect("ChatRequest always serializes to JSON");
let config = ureq::Agent::config_builder() .http_status_as_error(false) .timeout_connect(Some(limits.connect)) .timeout_recv_response(Some(limits.response)) .timeout_recv_body(Some(limits.body)) .build(); let agent: ureq::Agent = config.into();
let base = self.api_base.trim_end_matches('/'); let url = format!("{base}/chat/completions");
let mut attempts_left = MAX_ATTEMPTS - 1; let mut backoff = limits.backoff; loop { let mut response = agent .post(&url) .header("Authorization", format!("Bearer {}", self.api_key)) .header("Content-Type", "application/json") .send(body.as_slice()) .map_err(ChatError::Transport)?;
let status = response.status(); if status.is_success() { return Ok(ChatStream::new(response.into_body().into_reader())); }
let body_text = response.body_mut().read_to_string().unwrap_or_default(); let truncated: String = body_text.chars().take(STATUS_ERROR_BODY_LIMIT).collect(); if attempts_left == 0 || !worth_retrying(status.as_u16()) { return Err(ChatError::Status { status: status.as_u16(), body: truncated, }); } attempts_left -= 1; std::thread::sleep(backoff); backoff *= 2; } }}
/// One piece of a streamed chat completion response.#[derive(Debug, Clone, PartialEq)]pub enum StreamItem { /// A piece of the reply's narration, as it arrives. Text(String), /// A tool call, once every fragment for it has arrived. Call(ToolCall),}
/// A streamed chat completion response, yielding text as it arrives and/// each tool call as soon as its fragments are complete.////// The stream ends when the server sends `[DONE]` or closes the connection.////// OpenAI-compatible providers number tool calls in ways that disagree./// Some give every call the index 0, and some send the fragments of two/// calls one after the other, 0, 1, 0, 1. So the index is a hint, not a/// key: a fragment that carries an id joins the call with that id, and a/// fragment without one joins the newest open call at its index. An id no/// open call has starts a new call, even at an index already in use.////// A call leaves the stream when a later call starts and its arguments/// parse as one whole JSON value, so the caller can run a tool while the/// model still writes the next call. Every call still open when the stream/// ends leaves then.pub struct ChatStream { events: SseReader<ureq::BodyReader<'static>>, /// The calls still being assembled, oldest first. open: VecDeque<OpenToolCall>, /// Items ready to yield before the next chunk is read: one chunk can /// both complete a call and carry text. queue: VecDeque<StreamItem>,}
/// A tool call being assembled from fragments, with the index the provider/// gave it.struct OpenToolCall { index: usize, partial: PartialToolCall,}
/// A tool call being assembled from fragments.#[derive(Default)]struct PartialToolCall { id: String, name: String, arguments: String,}
impl PartialToolCall { /// Reports whether the arguments seen so far are one whole JSON value. /// /// A provider that splits arguments sends a prefix of a JSON value, /// and a prefix never parses. So a parse that succeeds means every /// argument fragment has arrived. fn arguments_are_whole(&self) -> bool { serde_json::from_str::<serde_json::Value>(&self.arguments).is_ok() }
fn into_tool_call(self) -> ToolCall { ToolCall { id: self.id, kind: "function".to_string(), function: ToolCallFunction { name: self.name, arguments: self.arguments, }, } }}
impl ChatStream { fn new(reader: ureq::BodyReader<'static>) -> Self { Self { events: SseReader::new(reader), open: VecDeque::new(), queue: VecDeque::new(), } }
/// Merges one chunk's tool-call fragments into the calls being /// assembled, queuing finished calls whenever a new call starts. fn absorb_tool_call_fragments(&mut self, fragments: Vec<ToolCallDelta>) { for fragment in fragments { match self.position_of(&fragment) { Some(position) => self.merge_into(position, fragment), None => { self.queue_finished_calls(); self.open.push_back(OpenToolCall { index: fragment.index, partial: PartialToolCall::default(), }); let newest = self.open.len() - 1; self.merge_into(newest, fragment); } } } }
/// Finds the open call `fragment` belongs to, or `None` when it starts /// a new one. fn position_of(&self, fragment: &ToolCallDelta) -> Option<usize> { match &fragment.id { Some(id) => self.open.iter().position(|open| open.partial.id == *id), None => self .open .iter() .rposition(|open| open.index == fragment.index), } }
/// Copies whatever `fragment` carries into the open call at `position`. fn merge_into(&mut self, position: usize, fragment: ToolCallDelta) { let partial = &mut self.open[position].partial; if let Some(id) = fragment.id { partial.id = id; } let Some(function) = fragment.function else { return; }; if let Some(name) = function.name { partial.name = name; } if let Some(arguments) = function.arguments { partial.arguments.push_str(&arguments); } }
/// Queues each finished call from the front of `open`, stopping at the /// first one still waiting on fragments. Stopping there keeps the /// calls in the order the provider started them. fn queue_finished_calls(&mut self) { while self .open .front() .is_some_and(|open| open.partial.arguments_are_whole()) { let finished = self.open.pop_front().expect("front is Some, just checked"); self.queue .push_back(StreamItem::Call(finished.partial.into_tool_call())); } }}
impl Iterator for ChatStream { type Item = Result<StreamItem, ChatError>;
fn next(&mut self) -> Option<Self::Item> { loop { if let Some(item) = self.queue.pop_front() { return Some(Ok(item)); }
let payload = match self.events.next() { None => { return self .open .pop_front() .map(|open| Ok(StreamItem::Call(open.partial.into_tool_call()))); } Some(Ok(payload)) => payload, Some(Err(error)) => return Some(Err(ChatError::Io(error))), }; let chunk: ChatChunk = match serde_json::from_str(&payload) { Ok(chunk) => chunk, Err(source) => return Some(Err(ChatError::Parse { source, payload })), }; let Some(choice) = chunk.choices.into_iter().next() else { continue; }; if let Some(fragments) = choice.delta.tool_calls { self.absorb_tool_call_fragments(fragments); } if let Some(content) = choice.delta.content && !content.is_empty() { self.queue.push_back(StreamItem::Text(content)); } } }}
/// Everything that can go wrong sending a chat completion request or reading/// its streamed response.#[derive(Debug)]pub enum ChatError { /// The server answered with a non-success HTTP status. Status { status: u16, body: String }, /// Reading the streamed response body failed. Io(io::Error), /// A streamed chunk was not valid JSON. Parse { source: serde_json::Error, payload: String, }, /// The request failed before a response arrived. Transport(ureq::Error),}
impl fmt::Display for ChatError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { ChatError::Status { status, body } => { write!(f, "the chat server answered with status {status}: {body}") } ChatError::Io(error) => write!(f, "reading the chat response failed: {error}"), ChatError::Parse { source, payload } => { write!( f, "could not parse chat chunk: {source} (payload: {payload})" ) } ChatError::Transport(error) => write!(f, "could not reach the chat server: {error}"), } }}
impl std::error::Error for ChatError {}
#[cfg(test)]mod tests { use super::*; use crate::chat::testing::{ CapturedRequest, fake_server as fake_server_serving, http_response, read_request, sse_response, }; use std::io::Write; use std::net::TcpListener; use std::sync::mpsc::Receiver; use std::thread::JoinHandle;
/// Serves `response` once to the first connection it accepts, then sends /// the request it received back through the returned channel. /// /// `response` must be a complete HTTP/1.1 response, including status /// line and headers. fn fake_server(response: Vec<u8>) -> (String, Receiver<CapturedRequest>, JoinHandle<()>) { fake_server_serving(vec![response]) }
/// Accepts one connection, reads the request, writes `sent`, and then /// holds the connection open for `hold` without writing anything more. fn stalling_server(sent: Vec<u8>, hold: Duration) -> (String, JoinHandle<()>) { let listener = TcpListener::bind("127.0.0.1:0").unwrap(); let addr = listener.local_addr().unwrap(); let handle = std::thread::spawn(move || { let (mut stream, _) = listener.accept().unwrap(); read_request(&stream); stream.write_all(&sent).unwrap(); std::thread::sleep(hold); }); (format!("http://{addr}"), handle) }
/// Limits short enough that a stalled server or a backoff costs the /// suite a fraction of a second. fn quick_limits() -> Limits { Limits { connect: Duration::from_secs(1), response: Duration::from_millis(200), body: Duration::from_millis(200), backoff: Duration::from_millis(1), } }
fn client_for(api_base: String) -> Client { Client { api_base, api_key: "sk-test".to_string(), model: "gpt-4o-mini".to_string(), reasoning_effort: None, } }
/// An address that always refuses a connection. /// /// Port 0 is never assigned to a live listener, so connecting to it /// fails the same way a closed port does, without binding and /// dropping a listener whose port another test's `fake_server` could /// then reuse before this test connects to it. fn dead_address() -> String { "http://127.0.0.1:0".to_string() }
/// The `StreamItem::Call` a finished tool call arrives as. fn call_item(id: &str, name: &str, arguments: &str) -> StreamItem { StreamItem::Call(ToolCall { id: id.to_string(), kind: "function".to_string(), function: ToolCallFunction { name: name.to_string(), arguments: arguments.to_string(), }, }) }
/// One SSE event carrying `fragments` as a chunk's `tool_calls`. fn tool_call_event(fragments: serde_json::Value) -> String { let chunk = serde_json::json!({ "choices": [{"delta": {"tool_calls": fragments}, "finish_reason": null}], }); format!("data: {chunk}\n\n") }
fn a_message() -> Vec<Message> { vec![Message { role: crate::chat::wire::Role::User, content: "I open the door.".to_string(), tool_calls: None, tool_call_id: None, }] }
#[test] fn streams_two_deltas_in_order_and_stops_at_done() { let body = "data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"},\"finish_reason\":null}]}\n\n\ data: {\"choices\":[{\"delta\":{\"content\":\" world\"},\"finish_reason\":null}]}\n\n\ data: [DONE]\n\n"; let (url, requests, server) = fake_server(sse_response(body)); let client = client_for(url);
let items: Vec<StreamItem> = client .stream(&a_message(), &[]) .unwrap() .map(|item| item.unwrap()) .collect();
assert_eq!( items, vec![ StreamItem::Text("Hello".to_string()), StreamItem::Text(" world".to_string()), ] ); server.join().unwrap(); let request = requests.recv().unwrap(); assert_eq!(request.method, "POST"); assert_eq!(request.path, "/chat/completions"); assert_eq!( request.headers.get("authorization"), Some(&"Bearer sk-test".to_string()) ); let sent: serde_json::Value = serde_json::from_str(&request.body).unwrap(); assert_eq!(sent["model"], "gpt-4o-mini"); assert_eq!(sent["stream"], true); assert_eq!(sent["messages"][0]["content"], "I open the door."); assert!(sent.get("tools").is_none()); }
#[test] fn stream_puts_the_given_tools_on_the_request() { let (url, requests, server) = fake_server(sse_response("data: [DONE]\n\n")); let client = client_for(url); let tools = vec![serde_json::json!({"type": "function", "function": {"name": "roll"}})];
client.stream(&a_message(), &tools).unwrap().for_each(drop);
server.join().unwrap(); let request = requests.recv().unwrap(); let sent: serde_json::Value = serde_json::from_str(&request.body).unwrap(); assert_eq!(sent["tools"], serde_json::json!(tools)); }
#[test] fn stream_puts_the_reasoning_effort_on_the_request_when_set() { let (url, requests, server) = fake_server(sse_response("data: [DONE]\n\n")); let client = Client { reasoning_effort: Some("high".to_string()), ..client_for(url) };
client.stream(&a_message(), &[]).unwrap().for_each(drop);
server.join().unwrap(); let request = requests.recv().unwrap(); let sent: serde_json::Value = serde_json::from_str(&request.body).unwrap(); assert_eq!(sent["reasoning_effort"], "high"); }
#[test] fn stream_omits_reasoning_effort_from_the_request_when_unset() { let (url, requests, server) = fake_server(sse_response("data: [DONE]\n\n")); let client = client_for(url);
client.stream(&a_message(), &[]).unwrap().for_each(drop);
server.join().unwrap(); let request = requests.recv().unwrap(); let sent: serde_json::Value = serde_json::from_str(&request.body).unwrap(); assert!(sent.get("reasoning_effort").is_none()); }
#[test] fn a_trailing_slash_on_api_base_does_not_double_the_path() { let (url, requests, server) = fake_server(sse_response("data: [DONE]\n\n")); let client = client_for(format!("{url}/"));
client.stream(&a_message(), &[]).unwrap().for_each(drop);
server.join().unwrap(); let request = requests.recv().unwrap(); assert_eq!(request.path, "/chat/completions"); }
#[test] fn a_non_success_status_becomes_a_status_error_with_the_body() { let (url, _requests, server) = fake_server(http_response( "400 Bad Request", "text/plain", "unknown model", )); let client = client_for(url);
let error = client.stream(&a_message(), &[]).err().unwrap();
assert!(matches!(error, ChatError::Status { status: 400, .. })); assert_eq!( error.to_string(), "the chat server answered with status 400: unknown model" ); server.join().unwrap(); }
#[test] fn a_rate_limited_request_goes_out_again_and_the_next_answer_streams() { let rate_limited = || http_response("429 Too Many Requests", "text/plain", "slow down"); let body = "data: {\"choices\":[{\"delta\":{\"content\":\"Hi\"},\"finish_reason\":null}]}\n\n\ data: [DONE]\n\n"; let (url, requests, server) = fake_server_serving(vec![rate_limited(), rate_limited(), sse_response(body)]); let client = client_for(url);
let items: Vec<StreamItem> = client .stream_within(&a_message(), &[], quick_limits()) .unwrap() .map(|item| item.unwrap()) .collect();
assert_eq!(items, vec![StreamItem::Text("Hi".to_string())]); server.join().unwrap(); assert_eq!(requests.iter().count(), 3); }
#[test] fn a_server_error_on_every_attempt_gives_up_with_the_status() { let failing = || http_response("503 Service Unavailable", "text/plain", "no capacity"); let (url, requests, server) = fake_server_serving(vec![failing(), failing(), failing()]); let client = client_for(url);
let error = client .stream_within(&a_message(), &[], quick_limits()) .err() .unwrap();
assert!(matches!(error, ChatError::Status { status: 503, .. })); server.join().unwrap(); assert_eq!(requests.iter().count(), MAX_ATTEMPTS as usize); }
#[test] fn a_server_that_takes_the_request_and_never_answers_fails_the_request() { let (url, server) = stalling_server(Vec::new(), Duration::from_millis(400)); let client = client_for(url);
let error = client .stream_within(&a_message(), &[], quick_limits()) .err() .unwrap();
assert!(matches!(error, ChatError::Transport(_))); server.join().unwrap(); }
#[test] fn a_body_that_stops_arriving_ends_the_stream_with_an_io_error() { let sse_body = "data: {\"choices\":[{\"delta\":{\"content\":\"Hi\"},\"finish_reason\":null}]}\n\n"; let head = format!( "HTTP/1.1 200 OK\r\n\ Content-Type: text/event-stream\r\n\ Content-Length: {}\r\n\ Connection: close\r\n\ \r\n\ {sse_body}", sse_body.len() + 100, ) .into_bytes(); let (url, server) = stalling_server(head, Duration::from_millis(400)); let client = client_for(url);
let mut stream = client .stream_within(&a_message(), &[], quick_limits()) .unwrap();
assert_eq!( stream.next().unwrap().unwrap(), StreamItem::Text("Hi".to_string()) ); assert!(matches!(stream.next(), Some(Err(ChatError::Io(_))))); server.join().unwrap(); }
#[test] fn a_chunk_that_is_not_valid_json_yields_a_parse_error() { let (url, _requests, server) = fake_server(sse_response("data: not json\n\n")); let client = client_for(url);
let error = client.stream(&a_message(), &[]).unwrap().next().unwrap();
assert!(matches!(error, Err(ChatError::Parse { .. }))); server.join().unwrap(); }
#[test] fn chunks_with_empty_choices_or_no_content_are_skipped() { let body = "data: {\"choices\":[]}\n\n\ data: {\"choices\":[{\"delta\":{},\"finish_reason\":null}]}\n\n\ data: {\"choices\":[{\"delta\":{\"content\":\"\"},\"finish_reason\":null}]}\n\n\ data: {\"choices\":[{\"delta\":{\"content\":\"Hi\"},\"finish_reason\":null}]}\n\n\ data: [DONE]\n\n"; let (url, _requests, server) = fake_server(sse_response(body)); let client = client_for(url);
let items: Vec<StreamItem> = client .stream(&a_message(), &[]) .unwrap() .map(|item| item.unwrap()) .collect();
assert_eq!(items, vec![StreamItem::Text("Hi".to_string())]); server.join().unwrap(); }
#[test] fn a_plain_text_response_yields_only_text_items() { let body = "data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"},\"finish_reason\":null}]}\n\n\ data: [DONE]\n\n"; let (url, _requests, server) = fake_server(sse_response(body)); let client = client_for(url);
let items: Vec<StreamItem> = client .stream(&a_message(), &[]) .unwrap() .map(|item| item.unwrap()) .collect();
assert_eq!(items, vec![StreamItem::Text("Hello".to_string())]); server.join().unwrap(); }
#[test] fn tool_call_argument_fragments_assemble_across_chunks() { let body = "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_1\",\"function\":{\"name\":\"roll\",\"arguments\":\"\"}}]},\"finish_reason\":null}]}\n\n\ data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"{\\\"notation\\\":\"}}]},\"finish_reason\":null}]}\n\n\ data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"\\\"1d20\\\"}\"}}]},\"finish_reason\":null}]}\n\n\ data: [DONE]\n\n"; let (url, _requests, server) = fake_server(sse_response(body)); let client = client_for(url);
let items: Vec<StreamItem> = client .stream(&a_message(), &[]) .unwrap() .map(|item| item.unwrap()) .collect();
assert_eq!( items, vec![StreamItem::Call(ToolCall { id: "call_1".to_string(), kind: "function".to_string(), function: ToolCallFunction { name: "roll".to_string(), arguments: "{\"notation\":\"1d20\"}".to_string(), }, })] ); server.join().unwrap(); }
#[test] fn a_fragment_with_an_id_but_no_function_field_still_lets_later_fragments_fill_it_in() { let body = "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_1\"}]},\"finish_reason\":null}]}\n\n\ data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"name\":\"roll\",\"arguments\":\"{}\"}}]},\"finish_reason\":null}]}\n\n\ data: [DONE]\n\n"; let (url, _requests, server) = fake_server(sse_response(body)); let client = client_for(url);
let items: Vec<StreamItem> = client .stream(&a_message(), &[]) .unwrap() .map(|item| item.unwrap()) .collect();
assert_eq!( items, vec![StreamItem::Call(ToolCall { id: "call_1".to_string(), kind: "function".to_string(), function: ToolCallFunction { name: "roll".to_string(), arguments: "{}".to_string(), }, })] ); server.join().unwrap(); }
#[test] fn tool_calls_at_distinct_indexes_assemble_into_separate_calls_in_index_order() { let body = "data: {\"choices\":[{\"delta\":{\"tool_calls\":[\ {\"index\":0,\"id\":\"call_1\",\"function\":{\"name\":\"roll\",\"arguments\":\"{}\"}},\ {\"index\":1,\"id\":\"call_2\",\"function\":{\"name\":\"roll\",\"arguments\":\"{}\"}}\ ]},\"finish_reason\":null}]}\n\n\ data: [DONE]\n\n"; let (url, _requests, server) = fake_server(sse_response(body)); let client = client_for(url);
let items: Vec<StreamItem> = client .stream(&a_message(), &[]) .unwrap() .map(|item| item.unwrap()) .collect();
assert_eq!( items, vec![ StreamItem::Call(ToolCall { id: "call_1".to_string(), kind: "function".to_string(), function: ToolCallFunction { name: "roll".to_string(), arguments: "{}".to_string(), }, }), StreamItem::Call(ToolCall { id: "call_2".to_string(), kind: "function".to_string(), function: ToolCallFunction { name: "roll".to_string(), arguments: "{}".to_string(), }, }), ] ); server.join().unwrap(); }
#[test] fn two_calls_reported_at_the_same_index_stay_separate() { let body = format!( "{}{}data: [DONE]\n\n", tool_call_event(serde_json::json!([{ "index": 0, "id": "call_1", "function": {"name": "roll", "arguments": "{\"notation\":\"1d20\"}"}, }])), tool_call_event(serde_json::json!([{ "index": 0, "id": "call_2", "function": {"name": "roll", "arguments": "{\"notation\":\"1d6\"}"}, }])), ); let (url, _requests, server) = fake_server(sse_response(&body)); let client = client_for(url);
let items: Vec<StreamItem> = client .stream(&a_message(), &[]) .unwrap() .map(|item| item.unwrap()) .collect();
assert_eq!( items, vec![ call_item("call_1", "roll", "{\"notation\":\"1d20\"}"), call_item("call_2", "roll", "{\"notation\":\"1d6\"}"), ] ); server.join().unwrap(); }
#[test] fn interleaved_fragments_for_two_indexes_assemble_into_two_whole_calls() { let body = format!( "{}{}{}{}data: [DONE]\n\n", tool_call_event(serde_json::json!([{ "index": 0, "id": "call_1", "function": {"name": "roll", "arguments": "{\"notation\":"}, }])), tool_call_event(serde_json::json!([{ "index": 1, "id": "call_2", "function": {"name": "roll", "arguments": "{\"notation\":"}, }])), tool_call_event(serde_json::json!([{ "index": 0, "function": {"arguments": "\"1d20\"}"}, }])), tool_call_event(serde_json::json!([{ "index": 1, "function": {"arguments": "\"1d6\"}"}, }])), ); let (url, _requests, server) = fake_server(sse_response(&body)); let client = client_for(url);
let items: Vec<StreamItem> = client .stream(&a_message(), &[]) .unwrap() .map(|item| item.unwrap()) .collect();
assert_eq!( items, vec![ call_item("call_1", "roll", "{\"notation\":\"1d20\"}"), call_item("call_2", "roll", "{\"notation\":\"1d6\"}"), ] ); server.join().unwrap(); }
#[test] fn a_whole_call_arriving_in_one_fragment_yields_one_call() { let body = format!( "{}data: [DONE]\n\n", tool_call_event(serde_json::json!([{ "index": 0, "id": "call_1", "function": {"name": "roll", "arguments": "{\"notation\":\"1d20\"}"}, }])), ); let (url, _requests, server) = fake_server(sse_response(&body)); let client = client_for(url);
let items: Vec<StreamItem> = client .stream(&a_message(), &[]) .unwrap() .map(|item| item.unwrap()) .collect();
assert_eq!( items, vec![call_item("call_1", "roll", "{\"notation\":\"1d20\"}")] ); server.join().unwrap(); }
#[test] fn a_fragment_that_repeats_the_id_continues_the_same_call() { let body = format!( "{}{}data: [DONE]\n\n", tool_call_event(serde_json::json!([{ "index": 0, "id": "call_1", "function": {"name": "roll", "arguments": "{\"notation\":"}, }])), tool_call_event(serde_json::json!([{ "index": 0, "id": "call_1", "function": {"arguments": "\"1d20\"}"}, }])), ); let (url, _requests, server) = fake_server(sse_response(&body)); let client = client_for(url);
let items: Vec<StreamItem> = client .stream(&a_message(), &[]) .unwrap() .map(|item| item.unwrap()) .collect();
assert_eq!( items, vec![call_item("call_1", "roll", "{\"notation\":\"1d20\"}")] ); server.join().unwrap(); }
#[test] fn text_and_completed_calls_interleave_in_stream_order() { let body = "data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"},\"finish_reason\":null}]}\n\n\ data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_1\",\"function\":{\"name\":\"roll\",\"arguments\":\"{}\"}}]},\"finish_reason\":null}]}\n\n\ data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":1,\"id\":\"call_2\",\"function\":{\"name\":\"roll\",\"arguments\":\"{}\"}}]},\"finish_reason\":null}]}\n\n\ data: {\"choices\":[{\"delta\":{\"content\":\"world\"},\"finish_reason\":null}]}\n\n\ data: [DONE]\n\n"; let (url, _requests, server) = fake_server(sse_response(body)); let client = client_for(url);
let items: Vec<StreamItem> = client .stream(&a_message(), &[]) .unwrap() .map(|item| item.unwrap()) .collect();
assert_eq!( items, vec![ StreamItem::Text("Hello".to_string()), StreamItem::Call(ToolCall { id: "call_1".to_string(), kind: "function".to_string(), function: ToolCallFunction { name: "roll".to_string(), arguments: "{}".to_string(), }, }), StreamItem::Text("world".to_string()), StreamItem::Call(ToolCall { id: "call_2".to_string(), kind: "function".to_string(), function: ToolCallFunction { name: "roll".to_string(), arguments: "{}".to_string(), }, }), ] ); server.join().unwrap(); }
#[test] fn eof_without_done_ends_the_stream_cleanly() { let body = "data: {\"choices\":[{\"delta\":{\"content\":\"Hi\"},\"finish_reason\":null}]}\n\n"; let (url, _requests, server) = fake_server(sse_response(body)); let client = client_for(url);
let mut stream = client.stream(&a_message(), &[]).unwrap();
assert_eq!( stream.next().unwrap().unwrap(), StreamItem::Text("Hi".to_string()) ); assert!(stream.next().is_none()); server.join().unwrap(); }
#[test] fn a_connection_closed_before_the_promised_content_length_yields_an_io_error() { let sse_body = "data: {\"choices\":[{\"delta\":{\"content\":\"Hi\"},\"finish_reason\":null}]}\n\n"; let response = format!( "HTTP/1.1 200 OK\r\n\ Content-Type: text/event-stream\r\n\ Content-Length: {}\r\n\ Connection: close\r\n\ \r\n\ {sse_body}", sse_body.len() + 100, ) .into_bytes(); let (url, _requests, server) = fake_server(response); let client = client_for(url);
let mut stream = client.stream(&a_message(), &[]).unwrap();
assert_eq!( stream.next().unwrap().unwrap(), StreamItem::Text("Hi".to_string()) ); assert!(matches!(stream.next(), Some(Err(ChatError::Io(_))))); server.join().unwrap(); }
#[test] fn a_connection_to_a_dead_port_returns_a_transport_error() { let client = client_for(dead_address());
let error = client.stream(&a_message(), &[]).err().unwrap();
assert!(matches!(error, ChatError::Transport(_))); }
#[test] fn status_error_displays_the_status_and_body() { let error = ChatError::Status { status: 500, body: "boom".to_string(), };
assert_eq!( error.to_string(), "the chat server answered with status 500: boom" ); }
#[test] fn io_error_displays_the_underlying_message() { let error = ChatError::Io(io::Error::other("connection reset"));
assert_eq!( error.to_string(), "reading the chat response failed: connection reset" ); }
#[test] fn parse_error_displays_the_source_and_payload() { let source = serde_json::from_str::<ChatChunk>("not json").unwrap_err(); let error = ChatError::Parse { source, payload: "not json".to_string(), };
let message = error.to_string(); assert!(message.starts_with("could not parse chat chunk:")); assert!(message.contains("payload: not json")); }
#[test] fn transport_error_displays_the_underlying_message() { let client = client_for(dead_address());
let error = client.stream(&a_message(), &[]).err().unwrap();
assert!( error .to_string() .starts_with("could not reach the chat server:") ); }}