mistralai-client-rs/src/v1/chat_stream.rs

use serde::{Deserialize, Serialize};
use serde_json::from_str;

use crate::v1::{chat, common, constants, error, tool};

// -----------------------------------------------------------------------------
// Response

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct ChatStreamChunk {
    pub id: String,
    pub object: String,
    /// Unix timestamp (in seconds).
    pub created: u64,
    pub model: constants::Model,
    pub choices: Vec<ChatStreamChunkChoice>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub usage: Option<common::ResponseUsage>,
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct ChatStreamChunkChoice {
    pub index: u32,
    pub delta: ChatStreamChunkChoiceDelta,
    pub finish_reason: Option<String>,
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct ChatStreamChunkChoiceDelta {
    pub role: Option<chat::ChatMessageRole>,
    #[serde(default)]
    pub content: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<Vec<tool::ToolCall>>,
}

/// Extracts serialized chunks from a stream message.
pub fn get_chunk_from_stream_message_line(
    line: &str,
) -> Result<Option<Vec<ChatStreamChunk>>, error::ApiError> {
    if line.trim() == "data: [DONE]" {
        return Ok(None);
    }

    let chunk_as_json = line.trim_start_matches("data: ").trim();
    if chunk_as_json.is_empty() {
        return Ok(Some(vec![]));
    }

    match from_str::<ChatStreamChunk>(chunk_as_json) {
        Ok(chunk) => Ok(Some(vec![chunk])),
        Err(e) => Err(error::ApiError {
            message: e.to_string(),
        }),
    }
}