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

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

//! The public streaming entry point: [`stream_chat`].
//!
//! Endpoint and model are always parameters — never hardcoded. The CLI reads
//! `AKURAI_LLM_URL` / `AKURAI_LLM_MODEL` from the environment and passes them
//! here, relaying each [`ChatEvent`] as an SSE event on `/api/assistant`.

use akurai_json::Value;

use crate::error::LlmError;
use crate::parse::{parse_chunk, sse_data_payload};
use crate::request::{build_request, parse_endpoint};
use crate::stream::stream_lines;
use crate::types::{ChatEvent, Chunk, Message, ToolCallAccumulator};

/// Stream a chat completion from an OpenAI-compatible `/v1/chat/completions`
/// endpoint, invoking `on_event` per event as the response arrives.
///
/// Events: zero or more [`ChatEvent::Delta`] text tokens, then zero or more
/// [`ChatEvent::ToolCall`]s (assembled from streamed fragments), then exactly
/// one terminal [`ChatEvent::Done`]. Tool calls and `Done` are emitted once the
/// server sends `data: [DONE]` or the connection closes.
///
/// `tools` is the optional OpenAI `tools` array (passed straight through to the
/// request body). Returns [`LlmError::TlsUnsupported`] for an `https://`
/// endpoint (std has no TLS); other variants cover connect/timeout/HTTP-status
/// failures. Individual malformed SSE chunks are skipped rather than aborting
/// the stream.
pub fn stream_chat<F>(
    endpoint: &str,
    model: &str,
    messages: &[Message],
    tools: Option<&[Value]>,
    mut on_event: F,
) -> Result<(), LlmError>
where
    F: FnMut(ChatEvent),
{
    let ep = parse_endpoint(endpoint)?;
    let request = build_request(&ep.host, ep.port, model, messages, tools);

    let mut acc = ToolCallAccumulator::new();
    let mut finished = false;

    stream_lines(&ep, &request, |line| {
        if finished {
            return Ok(());
        }
        let Some(payload) = sse_data_payload(line) else {
            return Ok(());
        };
        match parse_chunk(payload) {
            Ok(Chunk::Done) => {
                for event in acc.take() {
                    on_event(event);
                }
                on_event(ChatEvent::Done);
                finished = true;
            }
            Ok(Chunk::Delta {
                content,
                tool_calls,
            }) => {
                if let Some(text) = content {
                    if !text.is_empty() {
                        on_event(ChatEvent::Delta(text));
                    }
                }
                acc.merge(&tool_calls);
            }
            // A stray unparseable line shouldn't kill an otherwise-good stream.
            Err(_) => {}
        }
        Ok(())
    })?;

    // The server closed the stream without an explicit `[DONE]` — flush
    // whatever we assembled and emit the terminal event ourselves.
    if !finished {
        for event in acc.take() {
            on_event(event);
        }
        on_event(ChatEvent::Done);
    }
    Ok(())
}