141 lines
4.6 KiB
Rust
141 lines
4.6 KiB
Rust
//! A server-sent-events reader: turns a streamed body into `data:` payloads.
|
|
//!
|
|
//! `llama-server` streams a completion as SSE — one JSON object per `data:` line, then
|
|
//! `data: [DONE]`. Lines end in `\n` or `\r\n`. Blank lines, comment lines (`:`) and any field
|
|
//! other than `data` carry nothing and are skipped. After `data:` one leading space is dropped, and
|
|
//! only one. Every byte the peer sends is read in bounded chunks; a line is only UTF-8 decoded once
|
|
//! its whole length has arrived.
|
|
|
|
use std::io::{ErrorKind, Read};
|
|
|
|
use crate::http::is_timeout;
|
|
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
pub enum SseItem {
|
|
/// The text after `data:`, without the line ending.
|
|
Data(String),
|
|
Done,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub enum SseError {
|
|
Io(std::io::Error),
|
|
Timeout,
|
|
Truncated,
|
|
LineTooLong,
|
|
NotUtf8,
|
|
}
|
|
|
|
impl std::fmt::Display for SseError {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
match self {
|
|
SseError::Io(e) => write!(f, "io error: {e}"),
|
|
SseError::Timeout => write!(f, "read timed out before the line arrived"),
|
|
SseError::Truncated => write!(f, "the body ended in the middle of a line"),
|
|
SseError::LineTooLong => write!(f, "a line exceeded the limit"),
|
|
SseError::NotUtf8 => write!(f, "a line was not valid UTF-8"),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl std::error::Error for SseError {
|
|
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
|
|
match self {
|
|
SseError::Io(e) => Some(e),
|
|
_ => None,
|
|
}
|
|
}
|
|
}
|
|
|
|
pub struct Events<R> {
|
|
reader: R,
|
|
buf: Vec<u8>,
|
|
max_line: usize,
|
|
ended: bool,
|
|
}
|
|
|
|
impl<R: Read> Events<R> {
|
|
pub fn new(reader: R, max_line: usize) -> Self {
|
|
Events {
|
|
reader,
|
|
buf: Vec::new(),
|
|
max_line,
|
|
ended: false,
|
|
}
|
|
}
|
|
|
|
/// The next item, or `None` when the stream has ended cleanly between lines.
|
|
pub fn next_item(&mut self) -> Result<Option<SseItem>, SseError> {
|
|
loop {
|
|
if self.ended {
|
|
return Ok(None);
|
|
}
|
|
let line = match self.read_line() {
|
|
Ok(None) => {
|
|
self.ended = true;
|
|
return Ok(None);
|
|
}
|
|
Ok(Some(line)) => line,
|
|
Err(e) => return Err(e),
|
|
};
|
|
match self.process_line(line)? {
|
|
Some(item) => return Ok(Some(item)),
|
|
None => continue,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Reads one complete line, without its ending. Returns `None` on a clean end with nothing
|
|
/// buffered, and `Truncated` on a clean end with a partial line left over.
|
|
fn read_line(&mut self) -> Result<Option<Vec<u8>>, SseError> {
|
|
loop {
|
|
if let Some(pos) = self.buf.iter().position(|&b| b == b'\n') {
|
|
let mut line: Vec<u8> = self.buf.drain(..=pos).collect();
|
|
line.pop(); // drop the '\n'
|
|
if line.last() == Some(&b'\r') {
|
|
line.pop();
|
|
}
|
|
if line.len() > self.max_line {
|
|
return Err(SseError::LineTooLong);
|
|
}
|
|
return Ok(Some(line));
|
|
}
|
|
if self.buf.len() > self.max_line {
|
|
return Err(SseError::LineTooLong);
|
|
}
|
|
let mut tmp = [0u8; 4096];
|
|
let n = match self.reader.read(&mut tmp) {
|
|
Ok(0) => {
|
|
if self.buf.is_empty() {
|
|
return Ok(None);
|
|
}
|
|
return Err(SseError::Truncated);
|
|
}
|
|
Ok(n) => n,
|
|
Err(e) if e.kind() == ErrorKind::Interrupted => continue,
|
|
Err(e) if is_timeout(&e) => return Err(SseError::Timeout),
|
|
Err(e) if e.kind() == ErrorKind::UnexpectedEof => return Err(SseError::Truncated),
|
|
Err(e) => return Err(SseError::Io(e)),
|
|
};
|
|
self.buf.extend_from_slice(&tmp[..n]);
|
|
}
|
|
}
|
|
|
|
fn process_line(&self, line: Vec<u8>) -> Result<Option<SseItem>, SseError> {
|
|
let rest = match line.strip_prefix(b"data:") {
|
|
Some(rest) => rest,
|
|
None => return Ok(None),
|
|
};
|
|
let rest = match rest.split_first() {
|
|
Some((&b' ', rest)) => rest,
|
|
_ => rest,
|
|
};
|
|
let text = std::str::from_utf8(rest).map_err(|_| SseError::NotUtf8)?;
|
|
if text == "[DONE]" {
|
|
Ok(Some(SseItem::Done))
|
|
} else {
|
|
Ok(Some(SseItem::Data(text.to_string())))
|
|
}
|
|
}
|
|
}
|