zesdex/src/app/runtime/stream/mod.rs

pub mod turn;
pub mod tools;

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

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum StreamEvent {
    Token(String),
    Reasoning(String),
    ToolCallDelta {
        index: usize,
        id: Option<String>,
        name: Option<String>,
        arguments_delta: String,
    },
    Usage {
        prompt_tokens: u64,
        completion_tokens: u64,
        total_tokens: u64,
    },
    Done,
    Error(String),
}

pub struct SseParser {
    buffer: String,
    event_type: Option<String>,
    data_lines: Vec<String>,
}

impl SseParser {
    pub fn new() -> Self {
        SseParser {
            buffer: String::new(),
            event_type: None,
            data_lines: Vec::new(),
        }
    }

    pub fn feed(&mut self, chunk: &str) -> Vec<StreamEvent> {
        self.buffer.push_str(chunk);
        let mut events = Vec::new();
        while let Some(line_end) = self.buffer.find('\n') {
            let line = self.buffer[..line_end].trim_end_matches('\r').to_string();
            self.buffer = self.buffer[line_end + 1..].to_string();
            if line.is_empty() {
                if let Some(event) = self.flush_event() {
                    events.push(event);
                }
            } else if let Some(ty) = line.strip_prefix("event: ") {
                self.event_type = Some(ty.trim().to_string());
            } else if let Some(data) = line.strip_prefix("data: ") {
                self.data_lines.push(data.to_string());
            } else if line.starts_with("data:") {
                self.data_lines.push(String::new());
            }
        }
        events
    }

    fn flush_event(&mut self) -> Option<StreamEvent> {
        let data = self.data_lines.join("\n");
        self.data_lines.clear();
        let event_type = self.event_type.take().unwrap_or_default();
        if data.is_empty() || data == "[DONE]" {
            if data == "[DONE]" {
                return Some(StreamEvent::Done);
            }
            return None;
        }
        let value: Value = serde_json::from_str(&data).ok()?;
        match event_type.as_str() {
            "message.stop" => Some(StreamEvent::Done),
            "message.start" => None,
            "message.delta" | "" => {
                let delta = value.get("delta").or_else(|| value.get("choices"))?;
                if let Some(choices) = delta.as_array() {
                    let choice = choices.first()?;
                    let delta = choice.get("delta")?;
                    if let Some(content) = delta.get("content").and_then(|c| c.as_str()) {
                        return Some(StreamEvent::Token(content.to_string()));
                    }
                    if let Some(reasoning) = delta.get("reasoning_content").and_then(|r| r.as_str()) {
                        return Some(StreamEvent::Reasoning(reasoning.to_string()));
                    }
                    if let Some(tool_calls) = delta.get("tool_calls").and_then(|tc| tc.as_array()) {
                        for tc in tool_calls {
                            let index = tc.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as usize;
                            let id = tc.get("id").and_then(|i| i.as_str()).map(|s| s.to_string());
                            let name = tc.get("function")
                                .and_then(|f| f.get("name"))
                                .and_then(|n| n.as_str())
                                .map(|s| s.to_string());
                            let args_delta = tc.get("function")
                                .and_then(|f| f.get("arguments"))
                                .and_then(|a| a.as_str())
                                .unwrap_or("")
                                .to_string();
                            return Some(StreamEvent::ToolCallDelta {
                                index,
                                id,
                                name,
                                arguments_delta: args_delta,
                            });
                        }
                    }
                    let finish = choice.get("finish_reason");
                    if let Some(reason) = finish.and_then(|r| r.as_str()) {
                        if reason == "stop" || reason == "tool_calls" {
                            return Some(StreamEvent::Done);
                        }
                    }
                }
                if let Some(content) = delta.get("content").and_then(|c| c.as_str()) {
                    return Some(StreamEvent::Token(content.to_string()));
                }
                None
            }
            _ => None,
        }
    }

    pub fn reset(&mut self) {
        self.buffer.clear();
        self.event_type = None;
        self.data_lines.clear();
    }
}

pub fn parse_stream_chunk(data: &str) -> Option<StreamEvent> {
    let value: Value = serde_json::from_str(data).ok()?;
    if value == Value::Null {
        return None;
    }
    let choices = value.get("choices")?.as_array()?;
    let choice = choices.first()?;
    let delta = choice.get("delta")?;
    if let Some(content) = delta.get("content").and_then(|c| c.as_str()) {
        return Some(StreamEvent::Token(content.to_string()));
    }
    if let Some(reasoning) = delta.get("reasoning_content").and_then(|r| r.as_str()) {
        return Some(StreamEvent::Reasoning(reasoning.to_string()));
    }
    if let Some(finish) = choice.get("finish_reason").and_then(|r| r.as_str()) {
        if finish == "stop" || finish == "tool_calls" {
            return Some(StreamEvent::Done);
        }
    }
    if let Some(tool_calls) = delta.get("tool_calls").and_then(|tc| tc.as_array()) {
        if let Some(tc) = tool_calls.first() {
            let index = tc.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as usize;
            let id = tc.get("id").and_then(|i| i.as_str()).map(|s| s.to_string());
            let name = tc.get("function")
                .and_then(|f| f.get("name"))
                .and_then(|n| n.as_str())
                .map(|s| s.to_string());
            let args = tc.get("function")
                .and_then(|f| f.get("arguments"))
                .and_then(|a| a.as_str())
                .unwrap_or("")
                .to_string();
            return Some(StreamEvent::ToolCallDelta {
                index,
                id,
                name,
                arguments_delta: args,
            });
        }
    }
    None
}