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

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

//! Parsing the SSE response stream of `data: {chat.completion.chunk}` lines.
//!
//! Two pure helpers, both unit-tested against sample strings:
//! - [`sse_data_payload`] extracts the payload from a `data:` line.
//! - [`parse_chunk`] turns one payload into a [`Chunk`] (text delta + tool-call
//!   fragments, or the `[DONE]` sentinel).

use akurai_json::{parse, Value};

use crate::error::LlmError;
use crate::types::{Chunk, ToolCallDelta};

/// Extract the payload of an SSE `data:` line, or `None` for any other line
/// (comments, `event:`/`id:` fields, blank lines). A single optional space
/// after the colon is stripped, per the SSE spec.
pub fn sse_data_payload(line: &str) -> Option<&str> {
    let rest = line.strip_prefix("data:")?;
    Some(rest.strip_prefix(' ').unwrap_or(rest))
}

/// Parse one `data:` payload into a [`Chunk`].
///
/// The `[DONE]` sentinel returns [`Chunk::Done`]. Otherwise the payload is
/// parsed as a `chat.completion.chunk` and the first choice's `delta` is read
/// for `content` text and `tool_calls[]` fragments. Chunks with no `choices`
/// (e.g. a usage-only trailer) parse to an empty [`Chunk::Delta`] rather than an
/// error. Invalid JSON returns [`LlmError::MalformedChunk`].
pub fn parse_chunk(payload: &str) -> Result<Chunk, LlmError> {
    let p = payload.trim();
    if p == "[DONE]" {
        return Ok(Chunk::Done);
    }
    let value = parse(p).map_err(|e| LlmError::MalformedChunk(e.to_string()))?;

    let empty = Chunk::Delta {
        content: None,
        tool_calls: Vec::new(),
    };
    let choices = match value.get("choices") {
        Some(Value::Array(items)) => items,
        _ => return Ok(empty),
    };
    let Some(delta) = choices.first().and_then(|c| c.get("delta")) else {
        return Ok(empty);
    };

    let content = delta
        .get("content")
        .and_then(Value::as_str)
        .map(str::to_string);

    let mut tool_calls = Vec::new();
    if let Some(Value::Array(items)) = delta.get("tool_calls") {
        for (i, tc) in items.iter().enumerate() {
            let index = tc
                .get("index")
                .and_then(Value::as_i64)
                .map(|n| n.max(0) as usize)
                .unwrap_or(i);
            let id = tc.get("id").and_then(Value::as_str).map(str::to_string);
            let func = tc.get("function");
            let name = func
                .and_then(|f| f.get("name"))
                .and_then(Value::as_str)
                .map(str::to_string);
            let arguments = func
                .and_then(|f| f.get("arguments"))
                .and_then(Value::as_str)
                .map(str::to_string);
            tool_calls.push(ToolCallDelta {
                index,
                id,
                name,
                arguments,
            });
        }
    }

    Ok(Chunk::Delta {
        content,
        tool_calls,
    })
}

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

    #[test]
    fn extracts_data_payload_with_and_without_space() {
        assert_eq!(sse_data_payload("data: hello"), Some("hello"));
        assert_eq!(sse_data_payload("data:hello"), Some("hello"));
        assert_eq!(sse_data_payload("data:  two"), Some(" two")); // only one space stripped
    }

    #[test]
    fn ignores_non_data_lines() {
        assert_eq!(sse_data_payload(""), None);
        assert_eq!(sse_data_payload(": keep-alive"), None);
        assert_eq!(sse_data_payload("event: delta"), None);
    }

    #[test]
    fn done_sentinel() {
        assert_eq!(parse_chunk("[DONE]").unwrap(), Chunk::Done);
        assert_eq!(parse_chunk("  [DONE]  ").unwrap(), Chunk::Done);
    }

    #[test]
    fn parses_a_content_delta() {
        let payload =
            r#"{"choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]}"#;
        assert_eq!(
            parse_chunk(payload).unwrap(),
            Chunk::Delta {
                content: Some("Hello".into()),
                tool_calls: vec![],
            }
        );
    }

    #[test]
    fn parses_a_tool_call_fragment() {
        let payload = r#"{"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"get_weather","arguments":"{\"loc"}}]}}]}"#;
        assert_eq!(
            parse_chunk(payload).unwrap(),
            Chunk::Delta {
                content: None,
                tool_calls: vec![ToolCallDelta {
                    index: 0,
                    id: Some("call_1".into()),
                    name: Some("get_weather".into()),
                    arguments: Some("{\"loc".into()),
                }],
            }
        );
    }

    #[test]
    fn tool_call_argument_only_fragment_defaults_index_to_position() {
        // No explicit "index"; falls back to array position (0 here).
        let payload =
            r#"{"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"ation\":1}"}}]}}]}"#;
        let Chunk::Delta { tool_calls, .. } = parse_chunk(payload).unwrap() else {
            panic!("expected delta");
        };
        assert_eq!(tool_calls.len(), 1);
        assert_eq!(tool_calls[0].index, 0);
        assert_eq!(tool_calls[0].id, None);
        assert_eq!(tool_calls[0].name, None);
        assert_eq!(tool_calls[0].arguments.as_deref(), Some("ation\":1}"));
    }

    #[test]
    fn role_only_first_chunk_yields_empty_delta() {
        let payload = r#"{"choices":[{"delta":{"role":"assistant"}}]}"#;
        assert_eq!(
            parse_chunk(payload).unwrap(),
            Chunk::Delta {
                content: None,
                tool_calls: vec![],
            }
        );
    }

    #[test]
    fn empty_choices_is_empty_delta_not_error() {
        let payload = r#"{"choices":[],"usage":{"total_tokens":7}}"#;
        assert_eq!(
            parse_chunk(payload).unwrap(),
            Chunk::Delta {
                content: None,
                tool_calls: vec![],
            }
        );
    }

    #[test]
    fn malformed_json_errs() {
        assert!(matches!(
            parse_chunk("{not json"),
            Err(LlmError::MalformedChunk(_))
        ));
    }
}