use crate::{constants::*, error::LLMError};
use serde::de::DeserializeOwned;
pub trait SSEContentExtractor {
fn extract_content(&self) -> Option<&str>;
}
pub fn parse_sse_chunk<T, F>(chunk: &str, parse_fn: F) -> Result<Option<String>, LLMError>
where
T: SSEContentExtractor,
F: Fn(&str) -> Result<T, serde_json::Error>,
{
let mut collected_content = String::new();
for line in chunk.lines() {
let line = line.trim();
if let Some(data) = line.strip_prefix(SSE_DATA_PREFIX) {
if data == SSE_DONE_MARKER {
return if collected_content.is_empty() {
Ok(None)
} else {
Ok(Some(collected_content))
};
}
match parse_fn(data) {
Ok(response) => {
if let Some(content) = response.extract_content() {
collected_content.push_str(content);
}
}
Err(_) => continue,
}
}
}
if collected_content.is_empty() {
Ok(None)
} else {
Ok(Some(collected_content))
}
}
pub fn parse_sse_chunk_json<T>(chunk: &str) -> Result<Option<String>, LLMError>
where
T: DeserializeOwned + SSEContentExtractor,
{
parse_sse_chunk(chunk, |data| serde_json::from_str::<T>(data))
}