AkurAI-Framework / crates/llm/src/stream.rs

Latest change a81af83654d182e3becb8e2beaed6d61a3bd8e6b - feat(llm): streaming OpenAI-compatible client + /api/assistant SSE endpoint by Ólafur Búi Ólafsson

//! The plain-HTTP/1.1 socket round-trip that streams the SSE response.
//!
//! [`stream_lines`] connects, writes the request, reads the response head
//! (handling `Transfer-Encoding: chunked`, which streaming OpenAI servers use),
//! and invokes a callback for every complete line of the body as it arrives.
//!
//! The socket loop itself is not unit-tested (tests never touch the network).
//! The line-reassembly [`LineBuffer`] — the one tricky pure part, since chunk
//! boundaries don't align with lines — is unit-tested directly.

use std::io::{BufRead, BufReader, Read, Write};
use std::net::TcpStream;
use std::time::Duration;

use crate::error::LlmError;
use crate::request::Endpoint;

/// Per-read timeout. Generous enough for a cold local model's first token and
/// for inter-token gaps; it resets on every read, so total generation time is
/// unbounded as long as tokens keep flowing.
const TIMEOUT: Duration = Duration::from_secs(120);

/// Connect, send `request`, and call `on_line` for each complete body line of
/// the streamed response (newline stripped). Returns when the stream ends
/// (final chunk or EOF) or `on_line` returns an error.
pub fn stream_lines<F>(ep: &Endpoint, request: &str, mut on_line: F) -> Result<(), LlmError>
where
    F: FnMut(&str) -> Result<(), LlmError>,
{
    let addr = format!("{}:{}", ep.host, ep.port);
    let stream = TcpStream::connect(&addr).map_err(|e| LlmError::Connect(e.to_string()))?;
    stream.set_read_timeout(Some(TIMEOUT)).map_err(map_io)?;
    stream.set_write_timeout(Some(TIMEOUT)).map_err(map_io)?;

    {
        let mut w = &stream;
        w.write_all(request.as_bytes()).map_err(map_io)?;
        w.flush().map_err(map_io)?;
    }

    let mut reader = BufReader::new(stream);
    let head = read_head(&mut reader)?;
    if !(200..300).contains(&head.status) {
        let mut buf = [0u8; 512];
        let snippet = match reader.read(&mut buf) {
            Ok(n) => String::from_utf8_lossy(&buf[..n]).into_owned(),
            Err(_) => String::new(),
        };
        return Err(LlmError::HttpStatus {
            code: head.status,
            snippet,
        });
    }

    if head.chunked {
        read_chunked(&mut reader, &mut on_line)
    } else {
        read_plain(&mut reader, &mut on_line)
    }
}

/// Parsed response head: status code and whether the body is chunked.
struct Head {
    status: u16,
    chunked: bool,
}

/// Read the status line and headers up to the blank line.
fn read_head<R: BufRead>(reader: &mut R) -> Result<Head, LlmError> {
    let mut status_line = String::new();
    if reader.read_line(&mut status_line).map_err(map_io)? == 0 {
        return Err(LlmError::ShortResponse);
    }
    let status = status_line
        .split_whitespace()
        .nth(1)
        .and_then(|c| c.parse().ok())
        .ok_or(LlmError::ShortResponse)?;

    let mut chunked = false;
    loop {
        let mut line = String::new();
        if reader.read_line(&mut line).map_err(map_io)? == 0 {
            break; // EOF before blank line
        }
        let trimmed = line.trim_end_matches(['\r', '\n']);
        if trimmed.is_empty() {
            break; // end of headers
        }
        if let Some((name, value)) = trimmed.split_once(':') {
            if name.trim().eq_ignore_ascii_case("transfer-encoding")
                && value.to_ascii_lowercase().contains("chunked")
            {
                chunked = true;
            }
        }
    }
    Ok(Head { status, chunked })
}

/// Read an unchunked body line by line until EOF.
fn read_plain<R, F>(reader: &mut R, on_line: &mut F) -> Result<(), LlmError>
where
    R: BufRead,
    F: FnMut(&str) -> Result<(), LlmError>,
{
    loop {
        let mut line = String::new();
        if reader.read_line(&mut line).map_err(map_io)? == 0 {
            return Ok(());
        }
        on_line(line.trim_end_matches(['\r', '\n']))?;
    }
}

/// Decode `Transfer-Encoding: chunked` framing and feed the payload bytes
/// through a [`LineBuffer`], so SSE lines that span chunk boundaries are
/// reassembled correctly.
fn read_chunked<R, F>(reader: &mut R, on_line: &mut F) -> Result<(), LlmError>
where
    R: BufRead,
    F: FnMut(&str) -> Result<(), LlmError>,
{
    let mut lines = LineBuffer::new();
    loop {
        let mut size_line = String::new();
        if reader.read_line(&mut size_line).map_err(map_io)? == 0 {
            break; // EOF without a terminating zero-chunk; flush what we have.
        }
        let hex = size_line.trim().split(';').next().unwrap_or("").trim();
        if hex.is_empty() {
            continue;
        }
        let size = usize::from_str_radix(hex, 16).map_err(|_| LlmError::ShortResponse)?;
        if size == 0 {
            break; // final chunk
        }
        let mut buf = vec![0u8; size];
        reader.read_exact(&mut buf).map_err(map_io)?;
        lines.push(&buf, on_line)?;
        // Consume the CRLF that terminates the chunk data.
        let mut crlf = [0u8; 2];
        let _ = reader.read_exact(&mut crlf);
    }
    lines.flush(on_line)
}

/// Map an I/O error to a timeout or generic I/O variant.
fn map_io(e: std::io::Error) -> LlmError {
    match e.kind() {
        std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut => LlmError::Timeout,
        _ => LlmError::Io(e.to_string()),
    }
}

/// Reassembles a byte stream (whose boundaries are arbitrary) into complete
/// `\n`-terminated lines, emitting each via a callback. A trailing partial line
/// is held until more bytes arrive or [`flush`](LineBuffer::flush) is called.
#[derive(Default)]
pub struct LineBuffer {
    buf: Vec<u8>,
}

impl LineBuffer {
    pub fn new() -> LineBuffer {
        LineBuffer::default()
    }

    /// Append `data` and emit every complete line it completes.
    pub fn push<F>(&mut self, data: &[u8], on_line: &mut F) -> Result<(), LlmError>
    where
        F: FnMut(&str) -> Result<(), LlmError>,
    {
        self.buf.extend_from_slice(data);
        while let Some(pos) = self.buf.iter().position(|&b| b == b'\n') {
            let line: Vec<u8> = self.buf.drain(..=pos).collect();
            let text = String::from_utf8_lossy(&line);
            on_line(text.trim_end_matches(['\r', '\n']))?;
        }
        Ok(())
    }

    /// Emit any held-back final line that had no trailing newline.
    pub fn flush<F>(&mut self, on_line: &mut F) -> Result<(), LlmError>
    where
        F: FnMut(&str) -> Result<(), LlmError>,
    {
        if !self.buf.is_empty() {
            let text = String::from_utf8_lossy(&self.buf);
            let trimmed = text.trim_end_matches(['\r', '\n']);
            if !trimmed.is_empty() {
                on_line(trimmed)?;
            }
            self.buf.clear();
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn collect(fragments: &[&[u8]]) -> Vec<String> {
        let mut out = Vec::new();
        let mut lb = LineBuffer::new();
        let mut sink = |line: &str| {
            out.push(line.to_string());
            Ok(())
        };
        for f in fragments {
            lb.push(f, &mut sink).unwrap();
        }
        lb.flush(&mut sink).unwrap();
        out
    }

    #[test]
    fn splits_whole_lines() {
        assert_eq!(
            collect(&[b"data: a\ndata: b\n"]),
            vec!["data: a".to_string(), "data: b".to_string()]
        );
    }

    #[test]
    fn reassembles_a_line_split_across_fragments() {
        // The "data: hello" line is delivered in three byte-fragments.
        assert_eq!(
            collect(&[b"data: hel", b"lo\nda", b"ta: x\n"]),
            vec!["data: hello".to_string(), "data: x".to_string()]
        );
    }

    #[test]
    fn strips_crlf_and_keeps_blank_lines() {
        assert_eq!(
            collect(&[b"a\r\n\r\nb\r\n"]),
            vec!["a".to_string(), String::new(), "b".to_string()]
        );
    }

    #[test]
    fn flush_emits_unterminated_tail() {
        assert_eq!(collect(&[b"data: tail"]), vec!["data: tail".to_string()]);
    }

    #[test]
    fn flush_drops_empty_tail() {
        assert_eq!(collect(&[b"only\n"]), vec!["only".to_string()]);
    }
}