use axum::body::Body;
use axum::http::header;
use axum::response::{IntoResponse, Response};
use serde_json::json;
use crate::error::ApiError;
pub(crate) fn sse_events<'a>(
url: &'a str,
title: Option<&'a str>,
markdown: &'a str,
) -> impl Iterator<Item = String> + 'a {
let progress = json!({
"status": "processing",
"url": url,
"title": title.unwrap_or(""),
});
std::iter::once(format!("event: progress\ndata: {progress}\n\n"))
.chain(
markdown
.split("\n\n")
.filter(|p| !p.trim().is_empty())
.map(|p| format!("event: content\ndata: {}\n\n", json!({ "chunk": p }))),
)
.chain(std::iter::once_with(move || {
let done = json!({
"status": "completed",
"url": url,
});
format!("event: done\ndata: {done}\n\n")
}))
}
pub(crate) fn error_event(e: &ApiError) -> String {
let mut body = json!({
"success": false,
"code": e.code.as_str(),
"error": e.message,
});
if let Some(d) = &e.details {
body["details"] = d.clone();
}
format!("event: error\ndata: {body}\n\n")
}
pub(crate) fn sse_response(body: Body) -> Response {
Response::builder()
.header(header::CONTENT_TYPE, "text/event-stream; charset=utf-8")
.header(header::CACHE_CONTROL, "no-cache")
.header(header::CONNECTION, "keep-alive")
.body(body)
.unwrap_or_else(|_| {
(
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
"SSE stream error",
)
.into_response()
})
}
pub fn stream_markdown_sse(url: &str, title: Option<&str>, markdown: &str) -> Response {
sse_response(Body::from(
sse_events(url, title, markdown).collect::<String>(),
))
}