mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
refactor: fold ai surfaces into formats
This commit is contained in:
326
crates/aether-ai-formats/src/response/common.rs
Normal file
326
crates/aether-ai-formats/src/response/common.rs
Normal file
@@ -0,0 +1,326 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::contracts::core_success_background_report_kind;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct LocalSyncReportParts {
|
||||
pub trace_id: String,
|
||||
pub report_kind: String,
|
||||
pub report_context: Option<Value>,
|
||||
pub status_code: u16,
|
||||
pub headers: BTreeMap<String, String>,
|
||||
pub body_json: Option<Value>,
|
||||
pub client_body_json: Option<Value>,
|
||||
pub body_base64: Option<String>,
|
||||
}
|
||||
|
||||
pub fn build_generated_tool_call_id(index: usize) -> String {
|
||||
format!("call_auto_{index}")
|
||||
}
|
||||
|
||||
pub fn canonicalize_tool_arguments(value: Option<Value>) -> String {
|
||||
match value {
|
||||
Some(Value::String(text)) => text,
|
||||
Some(other) => serde_json::to_string(&other).unwrap_or_else(|_| "null".to_string()),
|
||||
None => "{}".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn remove_empty_pages_from_tool_arguments(arguments: &str) -> String {
|
||||
let Ok(mut value) = serde_json::from_str::<Value>(arguments) else {
|
||||
return arguments.to_string();
|
||||
};
|
||||
let Some(object) = value.as_object_mut() else {
|
||||
return arguments.to_string();
|
||||
};
|
||||
if object.get("pages").and_then(Value::as_str) != Some("") {
|
||||
return arguments.to_string();
|
||||
}
|
||||
object.remove("pages");
|
||||
serde_json::to_string(&value).unwrap_or_else(|_| arguments.to_string())
|
||||
}
|
||||
|
||||
pub fn prepare_local_success_response_parts(
|
||||
headers: &BTreeMap<String, String>,
|
||||
body_json: &Value,
|
||||
) -> serde_json::Result<(Vec<u8>, BTreeMap<String, String>)> {
|
||||
prepare_local_success_response_parts_owned(headers.clone(), body_json)
|
||||
}
|
||||
|
||||
pub fn prepare_local_success_response_parts_owned(
|
||||
mut headers: BTreeMap<String, String>,
|
||||
body_json: &Value,
|
||||
) -> serde_json::Result<(Vec<u8>, BTreeMap<String, String>)> {
|
||||
headers.remove("content-encoding");
|
||||
headers.remove("content-length");
|
||||
headers.insert("content-type".to_string(), "application/json".to_string());
|
||||
let body_bytes = serde_json::to_vec(body_json)?;
|
||||
headers.insert("content-length".to_string(), body_bytes.len().to_string());
|
||||
Ok((body_bytes, headers))
|
||||
}
|
||||
|
||||
fn should_capture_client_sync_success_body(payload: &LocalSyncReportParts) -> bool {
|
||||
payload
|
||||
.report_context
|
||||
.as_ref()
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|context| context.get("upstream_is_stream"))
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
pub fn build_local_success_background_report(
|
||||
payload: &LocalSyncReportParts,
|
||||
body_json: Value,
|
||||
headers: BTreeMap<String, String>,
|
||||
) -> Option<LocalSyncReportParts> {
|
||||
let report_kind = core_success_background_report_kind(payload.report_kind.as_str())?;
|
||||
let upstream_is_stream = should_capture_client_sync_success_body(payload);
|
||||
let client_body_json = upstream_is_stream.then(|| body_json.clone());
|
||||
let provider_body_json = if upstream_is_stream {
|
||||
payload.body_json.clone()
|
||||
} else {
|
||||
Some(body_json)
|
||||
};
|
||||
let provider_body_base64 = if upstream_is_stream {
|
||||
payload.body_base64.clone()
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
Some(LocalSyncReportParts {
|
||||
trace_id: payload.trace_id.clone(),
|
||||
report_kind: report_kind.to_string(),
|
||||
report_context: payload.report_context.clone(),
|
||||
status_code: payload.status_code,
|
||||
headers,
|
||||
body_json: provider_body_json,
|
||||
client_body_json,
|
||||
body_base64: provider_body_base64,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn build_local_success_conversion_background_report(
|
||||
payload: &LocalSyncReportParts,
|
||||
client_body_json: Value,
|
||||
provider_body_json: Value,
|
||||
) -> Option<LocalSyncReportParts> {
|
||||
let report_kind = core_success_background_report_kind(payload.report_kind.as_str())?;
|
||||
|
||||
Some(LocalSyncReportParts {
|
||||
trace_id: payload.trace_id.clone(),
|
||||
report_kind: report_kind.to_string(),
|
||||
report_context: payload.report_context.clone(),
|
||||
status_code: payload.status_code,
|
||||
headers: payload.headers.clone(),
|
||||
body_json: Some(provider_body_json),
|
||||
client_body_json: Some(client_body_json),
|
||||
body_base64: None,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use base64::Engine as _;
|
||||
use serde_json::Value;
|
||||
|
||||
use super::{
|
||||
build_generated_tool_call_id, build_local_success_background_report,
|
||||
build_local_success_conversion_background_report, canonicalize_tool_arguments,
|
||||
prepare_local_success_response_parts, prepare_local_success_response_parts_owned,
|
||||
remove_empty_pages_from_tool_arguments, LocalSyncReportParts,
|
||||
};
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
#[test]
|
||||
fn generated_tool_call_ids_are_stable() {
|
||||
assert_eq!(build_generated_tool_call_id(3), "call_auto_3");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn canonicalizes_tool_arguments() {
|
||||
assert_eq!(
|
||||
canonicalize_tool_arguments(Some(serde_json::json!({"x": 1}))),
|
||||
"{\"x\":1}"
|
||||
);
|
||||
assert_eq!(canonicalize_tool_arguments(None), "{}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn removes_empty_pages_from_tool_arguments() {
|
||||
assert_eq!(
|
||||
remove_empty_pages_from_tool_arguments(
|
||||
r#"{"file_path":"/tmp/a.txt","offset":1,"limit":20,"pages":""}"#
|
||||
),
|
||||
r#"{"file_path":"/tmp/a.txt","offset":1,"limit":20}"#
|
||||
);
|
||||
assert_eq!(
|
||||
remove_empty_pages_from_tool_arguments(r#"{"pages":"1-2"}"#),
|
||||
r#"{"pages":"1-2"}"#
|
||||
);
|
||||
assert_eq!(
|
||||
remove_empty_pages_from_tool_arguments(r#"{"pages":"#),
|
||||
r#"{"pages":"#
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_local_success_response_parts_normalizes_headers() {
|
||||
let headers = BTreeMap::from([
|
||||
("content-encoding".to_string(), "gzip".to_string()),
|
||||
("content-length".to_string(), "999".to_string()),
|
||||
("x-test".to_string(), "1".to_string()),
|
||||
]);
|
||||
let (body_bytes, normalized_headers) =
|
||||
prepare_local_success_response_parts(&headers, &serde_json::json!({"ok": true}))
|
||||
.expect("response parts should serialize");
|
||||
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<Value>(&body_bytes).expect("json body"),
|
||||
serde_json::json!({"ok": true})
|
||||
);
|
||||
assert_eq!(
|
||||
normalized_headers.get("content-type").map(String::as_str),
|
||||
Some("application/json")
|
||||
);
|
||||
assert!(!normalized_headers.contains_key("content-encoding"));
|
||||
let expected_length = body_bytes.len().to_string();
|
||||
assert_eq!(
|
||||
normalized_headers.get("content-length").map(String::as_str),
|
||||
Some(expected_length.as_str())
|
||||
);
|
||||
assert_eq!(
|
||||
normalized_headers.get("x-test").map(String::as_str),
|
||||
Some("1")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_local_success_response_parts_owned_normalizes_headers() {
|
||||
let headers = BTreeMap::from([
|
||||
("content-encoding".to_string(), "gzip".to_string()),
|
||||
("content-length".to_string(), "999".to_string()),
|
||||
("x-test".to_string(), "1".to_string()),
|
||||
]);
|
||||
let (body_bytes, normalized_headers) =
|
||||
prepare_local_success_response_parts_owned(headers, &serde_json::json!({"ok": true}))
|
||||
.expect("response parts should serialize");
|
||||
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<Value>(&body_bytes).expect("json body"),
|
||||
serde_json::json!({"ok": true})
|
||||
);
|
||||
assert_eq!(
|
||||
normalized_headers.get("content-type").map(String::as_str),
|
||||
Some("application/json")
|
||||
);
|
||||
assert!(!normalized_headers.contains_key("content-encoding"));
|
||||
let expected_length = body_bytes.len().to_string();
|
||||
assert_eq!(
|
||||
normalized_headers.get("content-length").map(String::as_str),
|
||||
Some(expected_length.as_str())
|
||||
);
|
||||
assert_eq!(
|
||||
normalized_headers.get("x-test").map(String::as_str),
|
||||
Some("1")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_local_success_background_report_maps_finalize_kind() {
|
||||
let payload = LocalSyncReportParts {
|
||||
trace_id: "trace-1".to_string(),
|
||||
report_kind: "openai_chat_sync_finalize".to_string(),
|
||||
report_context: Some(serde_json::json!({"request_id": "req-1"})),
|
||||
status_code: 200,
|
||||
headers: BTreeMap::from([("x-test".to_string(), "1".to_string())]),
|
||||
body_json: None,
|
||||
client_body_json: None,
|
||||
body_base64: None,
|
||||
};
|
||||
|
||||
let report = build_local_success_background_report(
|
||||
&payload,
|
||||
serde_json::json!({"id": "resp-1"}),
|
||||
payload.headers.clone(),
|
||||
)
|
||||
.expect("success report should be built");
|
||||
|
||||
assert_eq!(report.report_kind, "openai_chat_sync_success");
|
||||
assert_eq!(report.body_json, Some(serde_json::json!({"id": "resp-1"})));
|
||||
assert_eq!(report.client_body_json, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_local_success_background_report_preserves_provider_stream_for_upstream_stream_sync() {
|
||||
let payload = LocalSyncReportParts {
|
||||
trace_id: "trace-1b".to_string(),
|
||||
report_kind: "openai_chat_sync_finalize".to_string(),
|
||||
report_context: Some(serde_json::json!({
|
||||
"request_id": "req-1b",
|
||||
"upstream_is_stream": true
|
||||
})),
|
||||
status_code: 200,
|
||||
headers: BTreeMap::from([("content-type".to_string(), "text/event-stream".to_string())]),
|
||||
body_json: None,
|
||||
client_body_json: None,
|
||||
body_base64: Some(base64::engine::general_purpose::STANDARD.encode(
|
||||
concat!(
|
||||
"event: response.created\n",
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp-1b\",\"object\":\"response\",\"status\":\"in_progress\",\"output\":[]}}\n\n",
|
||||
"event: response.output_text.delta\n",
|
||||
"data: {\"type\":\"response.output_text.delta\",\"delta\":\"hello\"}\n\n",
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-1b\",\"object\":\"response\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n",
|
||||
)
|
||||
)),
|
||||
};
|
||||
|
||||
let report = build_local_success_background_report(
|
||||
&payload,
|
||||
serde_json::json!({"id": "resp-1b"}),
|
||||
payload.headers.clone(),
|
||||
)
|
||||
.expect("success report should be built");
|
||||
|
||||
assert_eq!(report.body_json, None);
|
||||
assert_eq!(
|
||||
report.client_body_json,
|
||||
Some(serde_json::json!({"id": "resp-1b"}))
|
||||
);
|
||||
assert_eq!(report.body_base64, payload.body_base64);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_local_success_conversion_background_report_maps_provider_body() {
|
||||
let payload = LocalSyncReportParts {
|
||||
trace_id: "trace-2".to_string(),
|
||||
report_kind: "openai_chat_sync_finalize".to_string(),
|
||||
report_context: Some(serde_json::json!({"request_id": "req-2"})),
|
||||
status_code: 200,
|
||||
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
|
||||
body_json: None,
|
||||
client_body_json: None,
|
||||
body_base64: None,
|
||||
};
|
||||
|
||||
let report = build_local_success_conversion_background_report(
|
||||
&payload,
|
||||
serde_json::json!({"client": true}),
|
||||
serde_json::json!({"provider": true}),
|
||||
)
|
||||
.expect("conversion success report should be built");
|
||||
|
||||
assert_eq!(report.report_kind, "openai_chat_sync_success");
|
||||
assert_eq!(
|
||||
report.body_json,
|
||||
Some(serde_json::json!({"provider": true}))
|
||||
);
|
||||
assert_eq!(
|
||||
report.client_body_json,
|
||||
Some(serde_json::json!({"client": true}))
|
||||
);
|
||||
}
|
||||
}
|
||||
168
crates/aether-ai-formats/src/response/error_body.rs
Normal file
168
crates/aether-ai-formats/src/response/error_body.rs
Normal file
@@ -0,0 +1,168 @@
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub enum LocalCoreSyncErrorKind {
|
||||
InvalidRequest,
|
||||
Authentication,
|
||||
PermissionDenied,
|
||||
NotFound,
|
||||
RateLimit,
|
||||
ContextLengthExceeded,
|
||||
Overloaded,
|
||||
ServerError,
|
||||
}
|
||||
|
||||
pub fn is_core_error_finalize_kind(report_kind: &str) -> bool {
|
||||
core_error_default_client_api_format(report_kind).is_some()
|
||||
}
|
||||
|
||||
pub fn core_error_default_client_api_format(report_kind: &str) -> Option<&'static str> {
|
||||
crate::contracts::core_error_default_client_api_format(report_kind)
|
||||
}
|
||||
|
||||
pub fn core_error_background_report_kind(report_kind: &str) -> Option<&'static str> {
|
||||
crate::contracts::core_error_background_report_kind(report_kind)
|
||||
}
|
||||
|
||||
pub fn core_success_background_report_kind(report_kind: &str) -> Option<&'static str> {
|
||||
crate::contracts::core_success_background_report_kind(report_kind)
|
||||
}
|
||||
|
||||
pub fn build_core_error_body_for_client_format(
|
||||
client_api_format: &str,
|
||||
message: &str,
|
||||
code: Option<&str>,
|
||||
kind: LocalCoreSyncErrorKind,
|
||||
) -> Option<Value> {
|
||||
let mut error_object = Map::new();
|
||||
error_object.insert("message".to_string(), Value::String(message.to_string()));
|
||||
|
||||
match aether_ai_formats::normalize_api_format_alias(client_api_format).as_str() {
|
||||
"openai:chat" | "openai:responses" | "openai:responses:compact" => {
|
||||
error_object.insert(
|
||||
"type".to_string(),
|
||||
Value::String(map_local_sync_error_kind_to_openai_type(kind).to_string()),
|
||||
);
|
||||
if let Some(code) = code.filter(|value| !value.is_empty()) {
|
||||
error_object.insert("code".to_string(), Value::String(code.to_string()));
|
||||
}
|
||||
Some(Value::Object(Map::from_iter([(
|
||||
"error".to_string(),
|
||||
Value::Object(error_object),
|
||||
)])))
|
||||
}
|
||||
"claude:messages" => {
|
||||
error_object.insert(
|
||||
"type".to_string(),
|
||||
Value::String(map_local_sync_error_kind_to_claude_type(kind).to_string()),
|
||||
);
|
||||
if let Some(code) = code.filter(|value| !value.is_empty()) {
|
||||
error_object.insert("code".to_string(), Value::String(code.to_string()));
|
||||
}
|
||||
Some(Value::Object(Map::from_iter([
|
||||
("type".to_string(), Value::String("error".to_string())),
|
||||
("error".to_string(), Value::Object(error_object)),
|
||||
])))
|
||||
}
|
||||
"gemini:generate_content" => Some(Value::Object(Map::from_iter([(
|
||||
"error".to_string(),
|
||||
Value::Object(Map::from_iter([
|
||||
(
|
||||
"code".to_string(),
|
||||
Value::from(map_local_sync_error_kind_to_gemini_code(kind)),
|
||||
),
|
||||
("message".to_string(), Value::String(message.to_string())),
|
||||
(
|
||||
"status".to_string(),
|
||||
Value::String(map_local_sync_error_kind_to_gemini_status(kind).to_string()),
|
||||
),
|
||||
])),
|
||||
)]))),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn map_local_sync_error_kind_to_openai_type(kind: LocalCoreSyncErrorKind) -> &'static str {
|
||||
match kind {
|
||||
LocalCoreSyncErrorKind::InvalidRequest => "invalid_request_error",
|
||||
LocalCoreSyncErrorKind::Authentication => "authentication_error",
|
||||
LocalCoreSyncErrorKind::PermissionDenied => "permission_error",
|
||||
LocalCoreSyncErrorKind::NotFound => "not_found_error",
|
||||
LocalCoreSyncErrorKind::RateLimit => "rate_limit_error",
|
||||
LocalCoreSyncErrorKind::ContextLengthExceeded => "context_length_exceeded",
|
||||
LocalCoreSyncErrorKind::Overloaded | LocalCoreSyncErrorKind::ServerError => "server_error",
|
||||
}
|
||||
}
|
||||
|
||||
fn map_local_sync_error_kind_to_claude_type(kind: LocalCoreSyncErrorKind) -> &'static str {
|
||||
match kind {
|
||||
LocalCoreSyncErrorKind::InvalidRequest | LocalCoreSyncErrorKind::ContextLengthExceeded => {
|
||||
"invalid_request_error"
|
||||
}
|
||||
LocalCoreSyncErrorKind::Authentication => "authentication_error",
|
||||
LocalCoreSyncErrorKind::PermissionDenied => "permission_error",
|
||||
LocalCoreSyncErrorKind::NotFound => "not_found_error",
|
||||
LocalCoreSyncErrorKind::RateLimit => "rate_limit_error",
|
||||
LocalCoreSyncErrorKind::Overloaded | LocalCoreSyncErrorKind::ServerError => "api_error",
|
||||
}
|
||||
}
|
||||
|
||||
fn map_local_sync_error_kind_to_gemini_code(kind: LocalCoreSyncErrorKind) -> u16 {
|
||||
match kind {
|
||||
LocalCoreSyncErrorKind::InvalidRequest | LocalCoreSyncErrorKind::ContextLengthExceeded => {
|
||||
400
|
||||
}
|
||||
LocalCoreSyncErrorKind::Authentication => 401,
|
||||
LocalCoreSyncErrorKind::PermissionDenied => 403,
|
||||
LocalCoreSyncErrorKind::NotFound => 404,
|
||||
LocalCoreSyncErrorKind::RateLimit => 429,
|
||||
LocalCoreSyncErrorKind::Overloaded => 503,
|
||||
LocalCoreSyncErrorKind::ServerError => 500,
|
||||
}
|
||||
}
|
||||
|
||||
fn map_local_sync_error_kind_to_gemini_status(kind: LocalCoreSyncErrorKind) -> &'static str {
|
||||
match kind {
|
||||
LocalCoreSyncErrorKind::InvalidRequest | LocalCoreSyncErrorKind::ContextLengthExceeded => {
|
||||
"INVALID_ARGUMENT"
|
||||
}
|
||||
LocalCoreSyncErrorKind::Authentication => "UNAUTHENTICATED",
|
||||
LocalCoreSyncErrorKind::PermissionDenied => "PERMISSION_DENIED",
|
||||
LocalCoreSyncErrorKind::NotFound => "NOT_FOUND",
|
||||
LocalCoreSyncErrorKind::RateLimit => "RESOURCE_EXHAUSTED",
|
||||
LocalCoreSyncErrorKind::Overloaded => "UNAVAILABLE",
|
||||
LocalCoreSyncErrorKind::ServerError => "INTERNAL",
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
build_core_error_body_for_client_format, core_success_background_report_kind,
|
||||
is_core_error_finalize_kind, LocalCoreSyncErrorKind,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn builds_openai_core_error_body() {
|
||||
let body = build_core_error_body_for_client_format(
|
||||
"openai:chat",
|
||||
"bad request",
|
||||
Some("invalid_request"),
|
||||
LocalCoreSyncErrorKind::InvalidRequest,
|
||||
)
|
||||
.expect("body should build");
|
||||
|
||||
assert_eq!(body["error"]["message"], "bad request");
|
||||
assert_eq!(body["error"]["type"], "invalid_request_error");
|
||||
assert_eq!(body["error"]["code"], "invalid_request");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recognizes_finalize_kind_and_success_mapping() {
|
||||
assert!(is_core_error_finalize_kind("openai_chat_sync_finalize"));
|
||||
assert_eq!(
|
||||
core_success_background_report_kind("openai_chat_sync_finalize"),
|
||||
Some("openai_chat_sync_success")
|
||||
);
|
||||
}
|
||||
}
|
||||
47
crates/aether-ai-formats/src/response/mod.rs
Normal file
47
crates/aether-ai-formats/src/response/mod.rs
Normal file
@@ -0,0 +1,47 @@
|
||||
use std::fmt;
|
||||
|
||||
pub use self::sse::{encode_done_sse, encode_json_sse, map_claude_stop_reason};
|
||||
pub use self::standard::stream_core::CanonicalStreamEvent;
|
||||
pub use self::standard::stream_core::CanonicalStreamFrame;
|
||||
pub use self::stream_rewrite::{
|
||||
maybe_build_ai_surface_stream_rewriter, resolve_finalize_stream_rewrite_mode,
|
||||
AiSurfaceStreamRewriter, FinalizeStreamRewriteMode,
|
||||
};
|
||||
|
||||
pub mod common;
|
||||
pub mod error_body;
|
||||
pub mod openai_image_stream;
|
||||
pub mod sse;
|
||||
pub mod standard;
|
||||
pub mod stream_rewrite;
|
||||
pub mod sync_products;
|
||||
pub mod sync_to_stream;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct AiSurfaceFinalizeError(pub String);
|
||||
|
||||
impl AiSurfaceFinalizeError {
|
||||
pub fn new(message: impl Into<String>) -> Self {
|
||||
Self(message.into())
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for AiSurfaceFinalizeError {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f, "AI surface finalize error: {}", self.0)
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for AiSurfaceFinalizeError {}
|
||||
|
||||
impl From<serde_json::Error> for AiSurfaceFinalizeError {
|
||||
fn from(source: serde_json::Error) -> Self {
|
||||
Self(source.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<base64::DecodeError> for AiSurfaceFinalizeError {
|
||||
fn from(source: base64::DecodeError) -> Self {
|
||||
Self(source.to_string())
|
||||
}
|
||||
}
|
||||
714
crates/aether-ai-formats/src/response/openai_image_stream.rs
Normal file
714
crates/aether-ai-formats/src/response/openai_image_stream.rs
Normal file
@@ -0,0 +1,714 @@
|
||||
use base64::Engine as _;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::contracts::OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND;
|
||||
use crate::request::standard::CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT;
|
||||
use crate::response::sse::encode_json_sse;
|
||||
use crate::response::AiSurfaceFinalizeError;
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct OpenAiImageStreamState {
|
||||
buffered: Vec<u8>,
|
||||
latest_image: Option<OpenAiImageFrame>,
|
||||
emitted_partial_count: u64,
|
||||
saw_upstream_partial: bool,
|
||||
emitted_failure: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct OpenAiImageFrame {
|
||||
b64_json: String,
|
||||
}
|
||||
|
||||
impl OpenAiImageStreamState {
|
||||
pub fn push_chunk(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
chunk: &[u8],
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
self.buffered.extend_from_slice(chunk);
|
||||
let mut output = Vec::new();
|
||||
while let Some(block_end) = find_sse_block_end(&self.buffered) {
|
||||
let block = self.buffered.drain(..block_end).collect::<Vec<_>>();
|
||||
output.extend(self.transform_block(report_context, &block)?);
|
||||
drain_sse_separator(&mut self.buffered);
|
||||
}
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
pub fn finish(&mut self, report_context: &Value) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
if self.buffered.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let block = std::mem::take(&mut self.buffered);
|
||||
self.transform_block(report_context, &block)
|
||||
}
|
||||
|
||||
fn transform_block(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
block: &[u8],
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
let text = std::str::from_utf8(block)
|
||||
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?;
|
||||
let mut event_name = None::<String>;
|
||||
let mut data_lines = Vec::new();
|
||||
for raw_line in text.lines() {
|
||||
let line = raw_line.trim_end_matches('\r');
|
||||
if let Some(value) = line.strip_prefix("event:") {
|
||||
event_name = Some(value.trim().to_string());
|
||||
} else if let Some(value) = line.strip_prefix("data:") {
|
||||
data_lines.push(value.trim().to_string());
|
||||
}
|
||||
}
|
||||
let data = data_lines.join("\n");
|
||||
if data.is_empty() || data == "[DONE]" {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let event: Value = serde_json::from_str(&data)?;
|
||||
let event_type = event
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.or(event_name.as_deref())
|
||||
.unwrap_or_default();
|
||||
match event_type {
|
||||
"error" | "response.failed" => self.handle_failed(report_context, &event),
|
||||
"response.image_generation_call.partial_image" => {
|
||||
self.handle_image_generation_partial(report_context, &event)
|
||||
}
|
||||
"response.output_item.done" => self.handle_output_item_done(report_context, &event),
|
||||
"response.completed" => self.handle_completed(report_context, &event),
|
||||
_ => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_image_generation_partial(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
event: &Value,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
if self.emitted_failure {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
if requested_partial_images(report_context) == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let Some(result) = event
|
||||
.get("partial_image_b64")
|
||||
.or_else(|| event.get("b64_json"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
let partial_image_index = event
|
||||
.get("partial_image_index")
|
||||
.or_else(|| event.get("output_index"))
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(self.emitted_partial_count);
|
||||
self.emitted_partial_count = self
|
||||
.emitted_partial_count
|
||||
.max(partial_image_index.saturating_add(1));
|
||||
self.saw_upstream_partial = true;
|
||||
self.latest_image = Some(OpenAiImageFrame {
|
||||
b64_json: result.to_string(),
|
||||
});
|
||||
|
||||
encode_json_sse(
|
||||
Some(image_partial_event_name(report_context)),
|
||||
&serde_json::json!({
|
||||
"type": image_partial_event_name(report_context),
|
||||
"b64_json": result,
|
||||
"partial_image_index": partial_image_index,
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
fn handle_output_item_done(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
event: &Value,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
if self.emitted_failure {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let Some(item) = event.get("item").and_then(Value::as_object) else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
if item.get("type").and_then(Value::as_str) != Some("image_generation_call") {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let Some(result) = item.get("result").and_then(Value::as_str).map(str::trim) else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
if result.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
self.latest_image = Some(OpenAiImageFrame {
|
||||
b64_json: result.to_string(),
|
||||
});
|
||||
|
||||
if requested_partial_images(report_context) == 0 || self.saw_upstream_partial {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let partial_image_index = event
|
||||
.get("output_index")
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(self.emitted_partial_count);
|
||||
self.emitted_partial_count = partial_image_index.saturating_add(1);
|
||||
|
||||
encode_json_sse(
|
||||
Some(image_partial_event_name(report_context)),
|
||||
&serde_json::json!({
|
||||
"type": image_partial_event_name(report_context),
|
||||
"b64_json": result,
|
||||
"partial_image_index": partial_image_index,
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
fn handle_completed(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
event: &Value,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
if self.emitted_failure {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
if self.latest_image.is_none() {
|
||||
if let Some(result) = completed_response_image_result(event) {
|
||||
self.latest_image = Some(OpenAiImageFrame {
|
||||
b64_json: result.to_string(),
|
||||
});
|
||||
}
|
||||
}
|
||||
let Some(latest_image) = self.latest_image.clone() else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
let usage = event
|
||||
.get("response")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|response| {
|
||||
response
|
||||
.get("tool_usage")
|
||||
.and_then(|value| value.get("image_gen"))
|
||||
.cloned()
|
||||
.or_else(|| response.get("usage").cloned())
|
||||
})
|
||||
.unwrap_or(Value::Null);
|
||||
|
||||
encode_json_sse(
|
||||
Some(image_completed_event_name(report_context)),
|
||||
&serde_json::json!({
|
||||
"type": image_completed_event_name(report_context),
|
||||
"b64_json": latest_image.b64_json,
|
||||
"usage": usage,
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
fn handle_failed(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
event: &Value,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
if self.emitted_failure {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
self.emitted_failure = true;
|
||||
let error = image_failure_error(event);
|
||||
encode_json_sse(
|
||||
Some(image_failed_event_name(report_context)),
|
||||
&serde_json::json!({
|
||||
"type": image_failed_event_name(report_context),
|
||||
"error": error,
|
||||
}),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn image_failure_error(event: &Value) -> Value {
|
||||
let mut error = event
|
||||
.get("error")
|
||||
.or_else(|| event.get("response").and_then(|value| value.get("error")))
|
||||
.and_then(Value::as_object)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
|
||||
if !error.contains_key("message") {
|
||||
if let Some(message) = event
|
||||
.get("message")
|
||||
.and_then(Value::as_str)
|
||||
.or_else(|| {
|
||||
event
|
||||
.get("response")
|
||||
.and_then(|value| value.get("error"))
|
||||
.and_then(|value| value.get("message"))
|
||||
.and_then(Value::as_str)
|
||||
})
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
error.insert("message".to_string(), Value::String(message.to_string()));
|
||||
}
|
||||
}
|
||||
if !error.contains_key("code") {
|
||||
if let Some(code) = event
|
||||
.get("code")
|
||||
.or_else(|| {
|
||||
event
|
||||
.get("response")
|
||||
.and_then(|value| value.get("error"))
|
||||
.and_then(|value| value.get("code"))
|
||||
})
|
||||
.cloned()
|
||||
{
|
||||
error.insert("code".to_string(), code);
|
||||
}
|
||||
}
|
||||
if !error.contains_key("type") {
|
||||
let inferred_type = error
|
||||
.get("code")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or("upstream_error");
|
||||
error.insert("type".to_string(), Value::String(inferred_type.to_string()));
|
||||
}
|
||||
if !error.contains_key("message") {
|
||||
error.insert(
|
||||
"message".to_string(),
|
||||
Value::String("Image generation failed".to_string()),
|
||||
);
|
||||
}
|
||||
|
||||
Value::Object(error)
|
||||
}
|
||||
|
||||
fn completed_response_image_result(event: &Value) -> Option<&str> {
|
||||
event
|
||||
.get("response")
|
||||
.and_then(|value| value.get("output"))
|
||||
.and_then(Value::as_array)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter(|item| item.get("type").and_then(Value::as_str) == Some("image_generation_call"))
|
||||
.filter_map(|item| item.get("result").and_then(Value::as_str))
|
||||
.map(str::trim)
|
||||
.find(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn requested_partial_images(report_context: &Value) -> u64 {
|
||||
report_context
|
||||
.get("image_request")
|
||||
.and_then(|value| value.get("partial_images"))
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn image_partial_event_name(report_context: &Value) -> &'static str {
|
||||
if image_request_operation(report_context) == Some("edit") {
|
||||
"image_edit.partial_image"
|
||||
} else {
|
||||
"image_generation.partial_image"
|
||||
}
|
||||
}
|
||||
|
||||
fn image_completed_event_name(report_context: &Value) -> &'static str {
|
||||
if image_request_operation(report_context) == Some("edit") {
|
||||
"image_edit.completed"
|
||||
} else {
|
||||
"image_generation.completed"
|
||||
}
|
||||
}
|
||||
|
||||
fn image_failed_event_name(report_context: &Value) -> &'static str {
|
||||
if image_request_operation(report_context) == Some("edit") {
|
||||
"image_edit.failed"
|
||||
} else {
|
||||
"image_generation.failed"
|
||||
}
|
||||
}
|
||||
|
||||
fn image_request_operation(report_context: &Value) -> Option<&str> {
|
||||
report_context
|
||||
.get("image_request")
|
||||
.and_then(|value| value.get("operation"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn find_sse_block_end(buffer: &[u8]) -> Option<usize> {
|
||||
buffer
|
||||
.windows(2)
|
||||
.position(|window| window == b"\n\n")
|
||||
.map(|index| index + 2)
|
||||
.or_else(|| {
|
||||
buffer
|
||||
.windows(4)
|
||||
.position(|window| window == b"\r\n\r\n")
|
||||
.map(|index| index + 4)
|
||||
})
|
||||
}
|
||||
|
||||
fn drain_sse_separator(buffer: &mut Vec<u8>) {
|
||||
while matches!(buffer.first(), Some(b'\n' | b'\r')) {
|
||||
buffer.remove(0);
|
||||
}
|
||||
}
|
||||
|
||||
pub struct OpenAiImageSyncFinalizeProduct {
|
||||
pub client_body_json: Value,
|
||||
pub provider_body_json: Value,
|
||||
}
|
||||
|
||||
pub fn maybe_build_openai_image_sync_finalize_product(
|
||||
report_kind: &str,
|
||||
status_code: u16,
|
||||
report_context: Option<&Value>,
|
||||
body_base64: Option<&str>,
|
||||
) -> Result<Option<OpenAiImageSyncFinalizeProduct>, AiSurfaceFinalizeError> {
|
||||
if report_kind != OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND || status_code >= 400 {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(report_context) = report_context else {
|
||||
return Ok(None);
|
||||
};
|
||||
if report_context
|
||||
.get("client_api_format")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
!= Some("openai:image")
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(body_base64) = body_base64 else {
|
||||
return Ok(None);
|
||||
};
|
||||
let default_output_format = report_context
|
||||
.get("image_request")
|
||||
.and_then(|value| value.get("output_format"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or(CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT);
|
||||
let body_bytes = base64::engine::general_purpose::STANDARD.decode(body_base64)?;
|
||||
let text = std::str::from_utf8(&body_bytes)
|
||||
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?;
|
||||
|
||||
let mut created = None;
|
||||
let mut completed_response = None;
|
||||
let mut images = Vec::new();
|
||||
|
||||
for raw_block in text.split("\n\n") {
|
||||
let block = raw_block.trim();
|
||||
if block.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let data_line = block
|
||||
.lines()
|
||||
.find_map(|line| line.trim().strip_prefix("data:").map(str::trim));
|
||||
let Some(data_line) = data_line else {
|
||||
continue;
|
||||
};
|
||||
if data_line.is_empty() || data_line == "[DONE]" {
|
||||
continue;
|
||||
}
|
||||
let event: Value = serde_json::from_str(data_line)?;
|
||||
match event
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
{
|
||||
"response.created" => {
|
||||
created = event
|
||||
.get("response")
|
||||
.and_then(|value| value.get("created_at"))
|
||||
.and_then(Value::as_i64)
|
||||
.or(created);
|
||||
}
|
||||
"response.output_item.done" => {
|
||||
let Some(item) = event.get("item").and_then(Value::as_object) else {
|
||||
continue;
|
||||
};
|
||||
if item.get("type").and_then(Value::as_str) != Some("image_generation_call") {
|
||||
continue;
|
||||
}
|
||||
let Some(result) = item.get("result").and_then(Value::as_str) else {
|
||||
continue;
|
||||
};
|
||||
images.push(serde_json::json!({
|
||||
"b64_json": result,
|
||||
"output_format": item.get("output_format").cloned().unwrap_or(Value::String(default_output_format.to_string())),
|
||||
"revised_prompt": item.get("revised_prompt").cloned().unwrap_or(Value::Null),
|
||||
}));
|
||||
}
|
||||
"response.completed" => {
|
||||
completed_response = event.get("response").and_then(Value::as_object).cloned();
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
if images.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let completed_response = completed_response.unwrap_or_default();
|
||||
let provider_usage = completed_response
|
||||
.get("tool_usage")
|
||||
.and_then(|value| value.get("image_gen"))
|
||||
.cloned()
|
||||
.or_else(|| completed_response.get("usage").cloned());
|
||||
let provider_body_json = serde_json::json!({
|
||||
"id": completed_response.get("id").cloned().unwrap_or(Value::Null),
|
||||
"object": "response",
|
||||
"model": completed_response.get("model").cloned().unwrap_or(Value::Null),
|
||||
"status": completed_response.get("status").cloned().unwrap_or(Value::String("completed".to_string())),
|
||||
"usage": provider_usage,
|
||||
"tool_usage": completed_response.get("tool_usage").cloned().unwrap_or(Value::Null),
|
||||
"output": images
|
||||
.iter()
|
||||
.map(|image| serde_json::json!({
|
||||
"type": "image_generation_call",
|
||||
"output_format": image.get("output_format").cloned().unwrap_or(Value::Null),
|
||||
"revised_prompt": image.get("revised_prompt").cloned().unwrap_or(Value::Null),
|
||||
}))
|
||||
.collect::<Vec<_>>(),
|
||||
});
|
||||
let client_images = images
|
||||
.iter()
|
||||
.map(|image| {
|
||||
let revised_prompt = image.get("revised_prompt").cloned().unwrap_or(Value::Null);
|
||||
let b64_json = image
|
||||
.get("b64_json")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
serde_json::json!({
|
||||
"b64_json": b64_json,
|
||||
"revised_prompt": revised_prompt,
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let client_body_json = serde_json::json!({
|
||||
"created": created.unwrap_or_default(),
|
||||
"data": client_images,
|
||||
"usage": provider_body_json.get("usage").cloned().unwrap_or(Value::Null),
|
||||
});
|
||||
|
||||
Ok(Some(OpenAiImageSyncFinalizeProduct {
|
||||
client_body_json,
|
||||
provider_body_json,
|
||||
}))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use base64::Engine as _;
|
||||
use serde_json::json;
|
||||
|
||||
use super::{maybe_build_openai_image_sync_finalize_product, OpenAiImageStreamState};
|
||||
|
||||
fn utf8(bytes: Vec<u8>) -> String {
|
||||
String::from_utf8(bytes).expect("utf8 should decode")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn emits_completed_event_for_generate() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "openai:image",
|
||||
"client_api_format": "openai:image",
|
||||
"needs_conversion": false,
|
||||
"image_request": {
|
||||
"operation": "generate"
|
||||
}
|
||||
});
|
||||
let mut rewriter = OpenAiImageStreamState::default();
|
||||
|
||||
let first = rewriter
|
||||
.push_chunk(
|
||||
&report_context,
|
||||
concat!(
|
||||
"event: response.output_item.done\n",
|
||||
"data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"ig_123\",\"type\":\"image_generation_call\",\"result\":\"aGVsbG8=\"}}\n\n"
|
||||
)
|
||||
.as_bytes(),
|
||||
)
|
||||
.expect("rewrite should succeed");
|
||||
assert!(first.is_empty());
|
||||
|
||||
let second = rewriter
|
||||
.push_chunk(
|
||||
&report_context,
|
||||
concat!(
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"tool_usage\":{\"image_gen\":{\"input_tokens\":1,\"output_tokens\":2,\"total_tokens\":3}}}}\n\n"
|
||||
)
|
||||
.as_bytes(),
|
||||
)
|
||||
.expect("rewrite should succeed");
|
||||
let output_text = utf8(second);
|
||||
assert!(output_text.contains("event: image_generation.completed"));
|
||||
assert!(output_text.contains("\"type\":\"image_generation.completed\""));
|
||||
assert!(output_text.contains("\"b64_json\":\"aGVsbG8=\""));
|
||||
assert!(output_text.contains("\"input_tokens\":1"));
|
||||
assert!(!output_text.contains("data: [DONE]"));
|
||||
assert!(rewriter
|
||||
.finish(&report_context)
|
||||
.expect("finish should succeed")
|
||||
.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_responses_partial_image_events() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "openai:image",
|
||||
"client_api_format": "openai:image",
|
||||
"needs_conversion": false,
|
||||
"image_request": {
|
||||
"operation": "generate",
|
||||
"partial_images": 1
|
||||
}
|
||||
});
|
||||
let mut rewriter = OpenAiImageStreamState::default();
|
||||
|
||||
let partial = rewriter
|
||||
.push_chunk(
|
||||
&report_context,
|
||||
concat!(
|
||||
"event: response.image_generation_call.partial_image\n",
|
||||
"data: {\"type\":\"response.image_generation_call.partial_image\",\"partial_image_index\":0,\"partial_image_b64\":\"cGFydGlhbA==\"}\n\n"
|
||||
)
|
||||
.as_bytes(),
|
||||
)
|
||||
.expect("rewrite should succeed");
|
||||
let partial_text = utf8(partial);
|
||||
assert!(partial_text.contains("event: image_generation.partial_image"));
|
||||
assert!(partial_text.contains("\"type\":\"image_generation.partial_image\""));
|
||||
assert!(partial_text.contains("\"b64_json\":\"cGFydGlhbA==\""));
|
||||
assert!(partial_text.contains("\"partial_image_index\":0"));
|
||||
assert!(!partial_text.contains("response.image_generation_call.partial_image"));
|
||||
|
||||
let done = rewriter
|
||||
.push_chunk(
|
||||
&report_context,
|
||||
concat!(
|
||||
"event: response.output_item.done\n",
|
||||
"data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"ig_123\",\"type\":\"image_generation_call\",\"result\":\"ZmluYWw=\"}}\n\n"
|
||||
)
|
||||
.as_bytes(),
|
||||
)
|
||||
.expect("rewrite should succeed");
|
||||
assert!(done.is_empty());
|
||||
|
||||
let completed = rewriter
|
||||
.push_chunk(
|
||||
&report_context,
|
||||
concat!(
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":4,\"output_tokens\":5,\"total_tokens\":9}}}\n\n"
|
||||
)
|
||||
.as_bytes(),
|
||||
)
|
||||
.expect("rewrite should succeed");
|
||||
let completed_text = utf8(completed);
|
||||
assert!(completed_text.contains("event: image_generation.completed"));
|
||||
assert!(completed_text.contains("\"type\":\"image_generation.completed\""));
|
||||
assert!(completed_text.contains("\"b64_json\":\"ZmluYWw=\""));
|
||||
assert!(completed_text.contains("\"total_tokens\":9"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_upstream_error_to_generation_failed_once() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "openai:image",
|
||||
"client_api_format": "openai:image",
|
||||
"needs_conversion": false,
|
||||
"image_request": {
|
||||
"operation": "generate"
|
||||
}
|
||||
});
|
||||
let mut rewriter = OpenAiImageStreamState::default();
|
||||
|
||||
let output = rewriter
|
||||
.push_chunk(
|
||||
&report_context,
|
||||
concat!(
|
||||
"event: error\n",
|
||||
"data: {\"type\":\"error\",\"error\":{\"type\":\"input-images\",\"code\":\"rate_limit_exceeded\",\"message\":\"Rate limit reached for gpt-image-2\",\"param\":null}}\n\n",
|
||||
"event: response.failed\n",
|
||||
"data: {\"type\":\"response.failed\",\"response\":{\"status\":\"failed\",\"error\":{\"code\":\"rate_limit_exceeded\",\"message\":\"Rate limit reached for gpt-image-2\"}}}\n\n"
|
||||
)
|
||||
.as_bytes(),
|
||||
)
|
||||
.expect("rewrite should succeed");
|
||||
let output_text = utf8(output);
|
||||
assert!(output_text.contains("event: image_generation.failed"));
|
||||
assert_eq!(
|
||||
output_text
|
||||
.matches("event: image_generation.failed")
|
||||
.count(),
|
||||
1
|
||||
);
|
||||
assert!(output_text.contains("\"type\":\"image_generation.failed\""));
|
||||
assert!(output_text.contains("\"type\":\"input-images\""));
|
||||
assert!(output_text.contains("\"code\":\"rate_limit_exceeded\""));
|
||||
assert!(output_text.contains("\"message\":\"Rate limit reached for gpt-image-2\""));
|
||||
assert!(!output_text.contains("response.failed"));
|
||||
assert!(rewriter
|
||||
.finish(&report_context)
|
||||
.expect("finish should succeed")
|
||||
.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sync_finalize_product_maps_stream_response_to_client_and_provider_bodies() {
|
||||
let report_context = json!({
|
||||
"client_api_format": "openai:image",
|
||||
"provider_api_format": "openai:image",
|
||||
"image_request": {
|
||||
"operation": "generate",
|
||||
"output_format": "png"
|
||||
}
|
||||
});
|
||||
let body_base64 = base64::engine::general_purpose::STANDARD.encode(
|
||||
concat!(
|
||||
"event: response.created\n",
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"created_at\":1776839946}}\n\n",
|
||||
"event: response.output_item.done\n",
|
||||
"data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"type\":\"image_generation_call\",\"output_format\":\"png\",\"revised_prompt\":\"revised history prompt\",\"result\":\"aGVsbG8=\"}}\n\n",
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_img_123\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"tool_usage\":{\"image_gen\":{\"input_tokens\":171,\"output_tokens\":1372,\"total_tokens\":1543}}}}\n\n"
|
||||
)
|
||||
.as_bytes(),
|
||||
);
|
||||
|
||||
let product = maybe_build_openai_image_sync_finalize_product(
|
||||
"openai_image_sync_finalize",
|
||||
200,
|
||||
Some(&report_context),
|
||||
Some(&body_base64),
|
||||
)
|
||||
.expect("finalize should succeed")
|
||||
.expect("finalize should match");
|
||||
|
||||
assert_eq!(product.client_body_json["created"], 1776839946);
|
||||
assert_eq!(product.client_body_json["data"][0]["b64_json"], "aGVsbG8=");
|
||||
assert_eq!(
|
||||
product.client_body_json["data"][0]["revised_prompt"],
|
||||
"revised history prompt"
|
||||
);
|
||||
assert_eq!(product.client_body_json["usage"]["input_tokens"], 171);
|
||||
assert_eq!(product.provider_body_json["id"], "resp_img_123");
|
||||
assert_eq!(
|
||||
product.provider_body_json["output"][0]["output_format"],
|
||||
"png"
|
||||
);
|
||||
assert_eq!(
|
||||
product.provider_body_json["output"][0]["revised_prompt"],
|
||||
"revised history prompt"
|
||||
);
|
||||
}
|
||||
}
|
||||
41
crates/aether-ai-formats/src/response/sse.rs
Normal file
41
crates/aether-ai-formats/src/response/sse.rs
Normal file
@@ -0,0 +1,41 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::response::AiSurfaceFinalizeError;
|
||||
|
||||
pub fn map_claude_stop_reason(
|
||||
stop_reason: Option<&str>,
|
||||
has_tool_calls: bool,
|
||||
) -> Option<&'static str> {
|
||||
let mapped = match stop_reason {
|
||||
Some("end_turn") | Some("stop_sequence") => Some("stop"),
|
||||
Some("max_tokens") => Some("length"),
|
||||
Some("tool_use") => Some("tool_calls"),
|
||||
Some("pause_turn") => Some("stop"),
|
||||
_ => None,
|
||||
};
|
||||
if has_tool_calls && mapped.is_none_or(|value| value == "stop") {
|
||||
Some("tool_calls")
|
||||
} else {
|
||||
mapped
|
||||
}
|
||||
}
|
||||
|
||||
pub fn encode_done_sse() -> Vec<u8> {
|
||||
b"data: [DONE]\n\n".to_vec()
|
||||
}
|
||||
|
||||
pub fn encode_json_sse(
|
||||
event: Option<&str>,
|
||||
value: &Value,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
let mut out = Vec::new();
|
||||
if let Some(event) = event.filter(|value| !value.trim().is_empty()) {
|
||||
out.extend_from_slice(b"event: ");
|
||||
out.extend_from_slice(event.as_bytes());
|
||||
out.push(b'\n');
|
||||
}
|
||||
out.extend_from_slice(b"data: ");
|
||||
out.extend(serde_json::to_vec(value).map_err(AiSurfaceFinalizeError::from)?);
|
||||
out.extend_from_slice(b"\n\n");
|
||||
Ok(out)
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
pub mod stream;
|
||||
1302
crates/aether-ai-formats/src/response/standard/claude/stream.rs
Normal file
1302
crates/aether-ai-formats/src/response/standard/claude/stream.rs
Normal file
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1 @@
|
||||
pub mod stream;
|
||||
1077
crates/aether-ai-formats/src/response/standard/gemini/stream.rs
Normal file
1077
crates/aether-ai-formats/src/response/standard/gemini/stream.rs
Normal file
File diff suppressed because it is too large
Load Diff
4
crates/aether-ai-formats/src/response/standard/mod.rs
Normal file
4
crates/aether-ai-formats/src/response/standard/mod.rs
Normal file
@@ -0,0 +1,4 @@
|
||||
pub mod claude;
|
||||
pub mod gemini;
|
||||
pub mod openai;
|
||||
pub mod stream_core;
|
||||
@@ -0,0 +1 @@
|
||||
pub mod stream;
|
||||
3099
crates/aether-ai-formats/src/response/standard/openai/stream.rs
Normal file
3099
crates/aether-ai-formats/src/response/standard/openai/stream.rs
Normal file
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,313 @@
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
pub use aether_ai_formats::protocol::stream::{
|
||||
CanonicalContentPart, CanonicalStreamEvent, CanonicalStreamFrame, CanonicalUsage,
|
||||
};
|
||||
|
||||
pub fn decode_json_data_line(line: &[u8]) -> Option<Value> {
|
||||
let text = std::str::from_utf8(line).ok()?;
|
||||
let trimmed = text.trim_matches('\r').trim();
|
||||
if trimmed.is_empty() || trimmed.starts_with(':') || trimmed.starts_with("event:") {
|
||||
return None;
|
||||
}
|
||||
let data_line = trimmed.strip_prefix("data:")?.trim();
|
||||
if data_line.is_empty() || data_line == "[DONE]" {
|
||||
return None;
|
||||
}
|
||||
serde_json::from_str(data_line).ok()
|
||||
}
|
||||
|
||||
pub fn resolve_identity(
|
||||
response_id: Option<&str>,
|
||||
model: Option<&str>,
|
||||
report_context: &Value,
|
||||
default_id: &str,
|
||||
) -> (String, String) {
|
||||
let id = response_id
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or(default_id)
|
||||
.to_string();
|
||||
let model = model
|
||||
.filter(|value| !value.is_empty())
|
||||
.or_else(|| report_context.get("mapped_model").and_then(Value::as_str))
|
||||
.or_else(|| report_context.get("model").and_then(Value::as_str))
|
||||
.unwrap_or("unknown")
|
||||
.to_string();
|
||||
(id, model)
|
||||
}
|
||||
|
||||
pub fn canonical_usage_from_openai_usage(value: Option<&Value>) -> Option<CanonicalUsage> {
|
||||
let usage = value?.as_object()?;
|
||||
let mut input_tokens = usage
|
||||
.get("input_tokens")
|
||||
.or_else(|| usage.get("prompt_tokens"))
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(0);
|
||||
let output_tokens = usage
|
||||
.get("output_tokens")
|
||||
.or_else(|| usage.get("completion_tokens"))
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(0);
|
||||
let cache_creation_tokens = usage
|
||||
.get("cache_creation_input_tokens")
|
||||
.and_then(Value::as_u64)
|
||||
.or_else(|| {
|
||||
usage
|
||||
.get("input_tokens_details")
|
||||
.or_else(|| usage.get("prompt_tokens_details"))
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|details| details.get("cached_creation_tokens"))
|
||||
.and_then(Value::as_u64)
|
||||
})
|
||||
.unwrap_or(0);
|
||||
let cache_read_tokens = usage
|
||||
.get("cache_read_input_tokens")
|
||||
.and_then(Value::as_u64)
|
||||
.or_else(|| {
|
||||
usage
|
||||
.get("input_tokens_details")
|
||||
.or_else(|| usage.get("prompt_tokens_details"))
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|details| details.get("cached_tokens"))
|
||||
.and_then(Value::as_u64)
|
||||
})
|
||||
.unwrap_or(0);
|
||||
let reasoning_tokens = usage
|
||||
.get("reasoning_tokens")
|
||||
.and_then(Value::as_u64)
|
||||
.or_else(|| {
|
||||
usage
|
||||
.get("output_tokens_details")
|
||||
.or_else(|| usage.get("completion_tokens_details"))
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|details| details.get("reasoning_tokens"))
|
||||
.and_then(Value::as_u64)
|
||||
})
|
||||
.unwrap_or(0);
|
||||
let total_tokens = usage.get("total_tokens").and_then(Value::as_u64).unwrap_or(
|
||||
input_tokens
|
||||
.saturating_add(output_tokens)
|
||||
.saturating_add(cache_creation_tokens)
|
||||
.saturating_add(cache_read_tokens),
|
||||
);
|
||||
if input_tokens == 0 && total_tokens > output_tokens {
|
||||
input_tokens = total_tokens.saturating_sub(output_tokens);
|
||||
}
|
||||
Some(CanonicalUsage {
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
total_tokens,
|
||||
cache_creation_tokens,
|
||||
cache_read_tokens,
|
||||
reasoning_tokens,
|
||||
..CanonicalUsage::default()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn canonical_usage_from_claude_usage(value: Option<&Value>) -> Option<CanonicalUsage> {
|
||||
let usage = value?.as_object()?;
|
||||
let input_tokens = usage
|
||||
.get("input_tokens")
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(0);
|
||||
let output_tokens = usage
|
||||
.get("output_tokens")
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(0);
|
||||
let cache_creation_ephemeral_5m_tokens = usage
|
||||
.get("cache_creation")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|value| value.get("ephemeral_5m_input_tokens"))
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(0);
|
||||
let cache_creation_ephemeral_1h_tokens = usage
|
||||
.get("cache_creation")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|value| value.get("ephemeral_1h_input_tokens"))
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(0);
|
||||
let cache_creation_tokens = usage
|
||||
.get("cache_creation_input_tokens")
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(
|
||||
cache_creation_ephemeral_5m_tokens.saturating_add(cache_creation_ephemeral_1h_tokens),
|
||||
);
|
||||
let cache_read_tokens = usage
|
||||
.get("cache_read_input_tokens")
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(0);
|
||||
let reasoning_tokens = usage
|
||||
.get("reasoning_tokens")
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(0);
|
||||
Some(CanonicalUsage {
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
total_tokens: input_tokens
|
||||
.saturating_add(output_tokens)
|
||||
.saturating_add(cache_creation_tokens)
|
||||
.saturating_add(cache_read_tokens),
|
||||
cache_creation_tokens,
|
||||
cache_creation_ephemeral_5m_tokens,
|
||||
cache_creation_ephemeral_1h_tokens,
|
||||
cache_read_tokens,
|
||||
reasoning_tokens,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn canonical_usage_from_gemini_usage(value: Option<&Value>) -> Option<CanonicalUsage> {
|
||||
let usage = value?.as_object()?;
|
||||
let input_tokens = usage
|
||||
.get("promptTokenCount")
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(0);
|
||||
let output_tokens = usage
|
||||
.get("candidatesTokenCount")
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(0);
|
||||
let reasoning_tokens = usage
|
||||
.get("thoughtsTokenCount")
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(0);
|
||||
let cache_read_tokens = usage
|
||||
.get("cachedContentTokenCount")
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(0);
|
||||
let total_tokens = usage
|
||||
.get("totalTokenCount")
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(
|
||||
input_tokens
|
||||
.saturating_add(output_tokens)
|
||||
.saturating_add(cache_read_tokens),
|
||||
);
|
||||
Some(CanonicalUsage {
|
||||
input_tokens,
|
||||
output_tokens: output_tokens.saturating_add(reasoning_tokens),
|
||||
total_tokens,
|
||||
cache_read_tokens,
|
||||
reasoning_tokens,
|
||||
..CanonicalUsage::default()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn normalize_openai_finish_reason(value: Option<&str>) -> Option<String> {
|
||||
match value {
|
||||
Some("function_call") => Some("tool_calls".to_string()),
|
||||
Some(other) if !other.trim().is_empty() => Some(other.to_string()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn map_openai_finish_reason_to_claude(value: Option<&str>) -> &'static str {
|
||||
match value {
|
||||
Some("length") => "max_tokens",
|
||||
Some("tool_calls") | Some("function_call") => "tool_use",
|
||||
Some("content_filter") => "content_filtered",
|
||||
_ => "end_turn",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn map_openai_finish_reason_to_gemini(value: Option<&str>) -> &'static str {
|
||||
match value {
|
||||
Some("length") => "MAX_TOKENS",
|
||||
Some("content_filter") => "SAFETY",
|
||||
_ => "STOP",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn parse_json_arguments_value(arguments: &str) -> Option<Value> {
|
||||
let trimmed = arguments.trim();
|
||||
if trimmed.is_empty() {
|
||||
return Some(Value::Object(Map::new()));
|
||||
}
|
||||
serde_json::from_str(trimmed).ok()
|
||||
}
|
||||
|
||||
pub fn build_openai_chat_chunk(
|
||||
id: &str,
|
||||
model: &str,
|
||||
text: String,
|
||||
tool_calls: Option<Vec<Value>>,
|
||||
finish_reason: Option<&str>,
|
||||
) -> Value {
|
||||
let mut delta = Map::new();
|
||||
delta.insert("role".to_string(), Value::String("assistant".to_string()));
|
||||
if !text.is_empty() {
|
||||
delta.insert("content".to_string(), Value::String(text));
|
||||
} else if tool_calls.is_none() {
|
||||
delta.insert("content".to_string(), Value::String(String::new()));
|
||||
}
|
||||
if let Some(tool_calls) = tool_calls {
|
||||
delta.insert("tool_calls".to_string(), Value::Array(tool_calls));
|
||||
}
|
||||
|
||||
json!({
|
||||
"id": id,
|
||||
"object": "chat.completion.chunk",
|
||||
"model": model,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": Value::Object(delta),
|
||||
"finish_reason": finish_reason,
|
||||
}]
|
||||
})
|
||||
}
|
||||
|
||||
pub fn build_openai_chat_role_chunk(id: &str, model: &str) -> Value {
|
||||
json!({
|
||||
"id": id,
|
||||
"object": "chat.completion.chunk",
|
||||
"model": model,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"role": "assistant"
|
||||
},
|
||||
"finish_reason": Value::Null
|
||||
}]
|
||||
})
|
||||
}
|
||||
|
||||
pub fn build_openai_chat_finish_chunk(id: &str, model: &str, finish_reason: Option<&str>) -> Value {
|
||||
json!({
|
||||
"id": id,
|
||||
"object": "chat.completion.chunk",
|
||||
"model": model,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {},
|
||||
"finish_reason": finish_reason,
|
||||
}]
|
||||
})
|
||||
}
|
||||
|
||||
pub fn build_openai_chat_usage_chunk(
|
||||
id: &str,
|
||||
model: &str,
|
||||
prompt_tokens: u64,
|
||||
completion_tokens: u64,
|
||||
total_tokens: u64,
|
||||
reasoning_tokens: u64,
|
||||
) -> Value {
|
||||
let mut usage = Map::new();
|
||||
usage.insert("prompt_tokens".to_string(), Value::from(prompt_tokens));
|
||||
usage.insert(
|
||||
"completion_tokens".to_string(),
|
||||
Value::from(completion_tokens),
|
||||
);
|
||||
usage.insert("total_tokens".to_string(), Value::from(total_tokens));
|
||||
if reasoning_tokens > 0 {
|
||||
usage.insert(
|
||||
"completion_tokens_details".to_string(),
|
||||
json!({ "reasoning_tokens": reasoning_tokens }),
|
||||
);
|
||||
}
|
||||
json!({
|
||||
"id": id,
|
||||
"object": "chat.completion.chunk",
|
||||
"model": model,
|
||||
"choices": [],
|
||||
"usage": usage,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,921 @@
|
||||
use aether_ai_formats::FormatId;
|
||||
use aether_contracts::{ExecutionStreamTerminalSummary, StandardizedUsage};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::response::error_body::{
|
||||
build_core_error_body_for_client_format, LocalCoreSyncErrorKind,
|
||||
};
|
||||
use crate::response::sse::encode_json_sse;
|
||||
use crate::response::standard::claude::stream::{ClaudeClientEmitter, ClaudeProviderState};
|
||||
use crate::response::standard::gemini::stream::{GeminiClientEmitter, GeminiProviderState};
|
||||
use crate::response::standard::openai::stream::{
|
||||
OpenAIChatClientEmitter, OpenAIChatProviderState, OpenAIResponsesClientEmitter,
|
||||
OpenAIResponsesProviderState,
|
||||
};
|
||||
use crate::response::standard::stream_core::common::{
|
||||
decode_json_data_line, CanonicalStreamEvent, CanonicalStreamFrame, CanonicalUsage,
|
||||
};
|
||||
use crate::response::AiSurfaceFinalizeError;
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct StreamingStandardFormatMatrix {
|
||||
provider: Option<ProviderStreamParser>,
|
||||
client: Option<ClientStreamEmitter>,
|
||||
terminated: bool,
|
||||
}
|
||||
|
||||
impl StreamingStandardFormatMatrix {
|
||||
pub fn transform_line(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
line: Vec<u8>,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
if self.terminated {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
self.ensure_initialized(report_context);
|
||||
if let Some(error_body) = build_client_error_body_for_line(report_context, &line) {
|
||||
self.terminated = true;
|
||||
return self.emit_error(error_body);
|
||||
}
|
||||
let Some(provider) = self.provider.as_mut() else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
let frames = provider.push_line(report_context, line)?;
|
||||
self.emit_frames(frames)
|
||||
}
|
||||
|
||||
pub fn finish(&mut self, report_context: &Value) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
if self.terminated {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
self.ensure_initialized(report_context);
|
||||
let Some(provider) = self.provider.as_mut() else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
let frames = provider.finish(report_context)?;
|
||||
let mut out = self.emit_frames(frames)?;
|
||||
if let Some(client) = self.client.as_mut() {
|
||||
out.extend(client.finish()?);
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
fn ensure_initialized(&mut self, report_context: &Value) {
|
||||
if self.provider.is_some() && self.client.is_some() {
|
||||
return;
|
||||
}
|
||||
|
||||
let provider_api_format = provider_api_format_for_context(report_context);
|
||||
let client_api_format = client_api_format_for_context(report_context);
|
||||
|
||||
self.provider = ProviderStreamParser::for_api_format(provider_api_format.as_str());
|
||||
self.client = ClientStreamEmitter::for_api_format(client_api_format.as_str());
|
||||
}
|
||||
|
||||
fn emit_frames(
|
||||
&mut self,
|
||||
frames: Vec<CanonicalStreamFrame>,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
let Some(client) = self.client.as_mut() else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
let mut out = Vec::new();
|
||||
for frame in frames {
|
||||
out.extend(client.emit(frame)?);
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
fn emit_error(&mut self, error_body: Value) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
let Some(client) = self.client.as_mut() else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
client.emit_error(error_body)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct StreamingStandardTerminalObserver {
|
||||
provider: Option<ProviderStreamParser>,
|
||||
latest_summary: Option<ExecutionStreamTerminalSummary>,
|
||||
}
|
||||
|
||||
impl StreamingStandardTerminalObserver {
|
||||
pub fn push_line(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
line: Vec<u8>,
|
||||
) -> Result<(), AiSurfaceFinalizeError> {
|
||||
self.ensure_initialized(report_context);
|
||||
let Some(provider) = self.provider.as_mut() else {
|
||||
return Ok(());
|
||||
};
|
||||
let frames = provider.push_line(report_context, line)?;
|
||||
self.observe_frames(frames);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn finish(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
) -> Result<Option<ExecutionStreamTerminalSummary>, AiSurfaceFinalizeError> {
|
||||
self.ensure_initialized(report_context);
|
||||
let Some(provider) = self.provider.as_mut() else {
|
||||
return Ok(self.latest_summary.clone());
|
||||
};
|
||||
let frames = provider.finish(report_context)?;
|
||||
self.observe_frames(frames);
|
||||
Ok(self.latest_summary.clone())
|
||||
}
|
||||
|
||||
pub fn disable_with_error(&mut self, parser_error: impl Into<String>) {
|
||||
let parser_error = parser_error.into();
|
||||
if let Some(summary) = self.latest_summary.as_mut() {
|
||||
if summary.parser_error.is_none() {
|
||||
summary.parser_error = Some(parser_error);
|
||||
}
|
||||
} else {
|
||||
self.latest_summary = Some(ExecutionStreamTerminalSummary {
|
||||
parser_error: Some(parser_error),
|
||||
..ExecutionStreamTerminalSummary::default()
|
||||
});
|
||||
}
|
||||
self.provider = None;
|
||||
}
|
||||
|
||||
pub fn latest_summary(&self) -> Option<&ExecutionStreamTerminalSummary> {
|
||||
self.latest_summary.as_ref()
|
||||
}
|
||||
|
||||
fn ensure_initialized(&mut self, report_context: &Value) {
|
||||
if self.provider.is_some() || self.latest_summary.is_some() {
|
||||
return;
|
||||
}
|
||||
let provider_api_format = provider_api_format_for_context(report_context);
|
||||
self.provider = ProviderStreamParser::for_api_format(provider_api_format.as_str());
|
||||
}
|
||||
|
||||
fn observe_frames(&mut self, frames: Vec<CanonicalStreamFrame>) {
|
||||
for frame in frames {
|
||||
self.observe_frame(frame);
|
||||
}
|
||||
}
|
||||
|
||||
fn observe_frame(&mut self, frame: CanonicalStreamFrame) {
|
||||
let CanonicalStreamFrame { id, model, event } = frame;
|
||||
let summary = self
|
||||
.latest_summary
|
||||
.get_or_insert_with(|| ExecutionStreamTerminalSummary {
|
||||
response_id: Some(id.clone()),
|
||||
model: Some(model.clone()),
|
||||
..ExecutionStreamTerminalSummary::default()
|
||||
});
|
||||
if summary.response_id.is_none() {
|
||||
summary.response_id = Some(id);
|
||||
}
|
||||
if summary.model.is_none() {
|
||||
summary.model = Some(model);
|
||||
}
|
||||
match event {
|
||||
CanonicalStreamEvent::UnknownEvent(_) => {
|
||||
summary.unknown_event_count = summary.unknown_event_count.saturating_add(1);
|
||||
}
|
||||
CanonicalStreamEvent::Finish {
|
||||
finish_reason,
|
||||
usage,
|
||||
} => {
|
||||
summary.finish_reason = finish_reason;
|
||||
summary.standardized_usage = usage.map(standardized_usage_from_canonical);
|
||||
summary.observed_finish = true;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
enum ProviderStreamParser {
|
||||
OpenAIChat(OpenAIChatProviderState),
|
||||
OpenAIResponses(OpenAIResponsesProviderState),
|
||||
Claude(ClaudeProviderState),
|
||||
Gemini(GeminiProviderState),
|
||||
}
|
||||
|
||||
impl ProviderStreamParser {
|
||||
fn for_api_format(provider_api_format: &str) -> Option<Self> {
|
||||
Some(match FormatId::parse(provider_api_format)? {
|
||||
FormatId::OpenAiChat => Self::OpenAIChat(OpenAIChatProviderState::default()),
|
||||
FormatId::OpenAiResponses | FormatId::OpenAiResponsesCompact => {
|
||||
Self::OpenAIResponses(OpenAIResponsesProviderState::default())
|
||||
}
|
||||
FormatId::ClaudeMessages => Self::Claude(ClaudeProviderState::default()),
|
||||
FormatId::GeminiGenerateContent => Self::Gemini(GeminiProviderState::default()),
|
||||
})
|
||||
}
|
||||
|
||||
fn push_line(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
line: Vec<u8>,
|
||||
) -> Result<Vec<CanonicalStreamFrame>, AiSurfaceFinalizeError> {
|
||||
match self {
|
||||
ProviderStreamParser::OpenAIChat(state) => state.push_line(report_context, line),
|
||||
ProviderStreamParser::OpenAIResponses(state) => state.push_line(report_context, line),
|
||||
ProviderStreamParser::Claude(state) => state.push_line(report_context, line),
|
||||
ProviderStreamParser::Gemini(state) => state.push_line(report_context, line),
|
||||
}
|
||||
}
|
||||
|
||||
fn finish(
|
||||
&mut self,
|
||||
report_context: &Value,
|
||||
) -> Result<Vec<CanonicalStreamFrame>, AiSurfaceFinalizeError> {
|
||||
match self {
|
||||
ProviderStreamParser::OpenAIChat(state) => state.finish(report_context),
|
||||
ProviderStreamParser::OpenAIResponses(state) => state.finish(report_context),
|
||||
ProviderStreamParser::Claude(state) => state.finish(report_context),
|
||||
ProviderStreamParser::Gemini(state) => state.finish(report_context),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
enum ClientStreamEmitter {
|
||||
OpenAIChat(OpenAIChatClientEmitter),
|
||||
OpenAIResponses(OpenAIResponsesClientEmitter),
|
||||
Claude(ClaudeClientEmitter),
|
||||
Gemini(GeminiClientEmitter),
|
||||
}
|
||||
|
||||
fn provider_api_format_for_context(report_context: &Value) -> String {
|
||||
string_context_field(report_context, "provider_stream_event_api_format")
|
||||
.or_else(|| string_context_field(report_context, "provider_stream_api_format"))
|
||||
.or_else(|| string_context_field(report_context, "provider_api_format"))
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn string_context_field(report_context: &Value, key: &str) -> Option<String> {
|
||||
let value = report_context.get(key)?.as_str()?.trim();
|
||||
(!value.is_empty()).then(|| value.to_ascii_lowercase())
|
||||
}
|
||||
|
||||
fn client_api_format_for_context(report_context: &Value) -> String {
|
||||
report_context
|
||||
.get("client_api_format")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.trim()
|
||||
.to_ascii_lowercase()
|
||||
}
|
||||
|
||||
fn standardized_usage_from_canonical(usage: CanonicalUsage) -> StandardizedUsage {
|
||||
let mut standardized = StandardizedUsage::new();
|
||||
standardized.input_tokens = usage.input_tokens as i64;
|
||||
standardized.output_tokens = usage.output_tokens as i64;
|
||||
standardized.cache_creation_tokens = usage.cache_creation_tokens as i64;
|
||||
standardized.cache_creation_ephemeral_5m_tokens =
|
||||
usage.cache_creation_ephemeral_5m_tokens as i64;
|
||||
standardized.cache_creation_ephemeral_1h_tokens =
|
||||
usage.cache_creation_ephemeral_1h_tokens as i64;
|
||||
standardized.cache_read_tokens = usage.cache_read_tokens as i64;
|
||||
standardized.reasoning_tokens = usage.reasoning_tokens as i64;
|
||||
standardized.dimensions.insert(
|
||||
"total_tokens".to_string(),
|
||||
serde_json::json!(usage.total_tokens),
|
||||
);
|
||||
standardized.normalize_cache_creation_breakdown()
|
||||
}
|
||||
|
||||
impl ClientStreamEmitter {
|
||||
fn for_api_format(client_api_format: &str) -> Option<Self> {
|
||||
Some(match FormatId::parse(client_api_format)? {
|
||||
FormatId::OpenAiChat => Self::OpenAIChat(OpenAIChatClientEmitter::default()),
|
||||
FormatId::OpenAiResponses | FormatId::OpenAiResponsesCompact => {
|
||||
Self::OpenAIResponses(OpenAIResponsesClientEmitter::default())
|
||||
}
|
||||
FormatId::ClaudeMessages => Self::Claude(ClaudeClientEmitter::default()),
|
||||
FormatId::GeminiGenerateContent => Self::Gemini(GeminiClientEmitter::default()),
|
||||
})
|
||||
}
|
||||
|
||||
fn emit(&mut self, frame: CanonicalStreamFrame) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
match self {
|
||||
ClientStreamEmitter::OpenAIChat(state) => state.emit(frame),
|
||||
ClientStreamEmitter::OpenAIResponses(state) => state.emit(frame),
|
||||
ClientStreamEmitter::Claude(state) => state.emit(frame),
|
||||
ClientStreamEmitter::Gemini(state) => state.emit(frame),
|
||||
}
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
match self {
|
||||
ClientStreamEmitter::OpenAIChat(state) => state.finish(),
|
||||
ClientStreamEmitter::OpenAIResponses(state) => state.finish(),
|
||||
ClientStreamEmitter::Claude(state) => state.finish(),
|
||||
ClientStreamEmitter::Gemini(state) => state.finish(),
|
||||
}
|
||||
}
|
||||
|
||||
fn emit_error(&mut self, error_body: Value) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
match self {
|
||||
ClientStreamEmitter::OpenAIResponses(state) => state.emit_error(error_body),
|
||||
ClientStreamEmitter::Claude(_) => {
|
||||
let event = error_body.get("type").and_then(Value::as_str);
|
||||
encode_json_sse(event, &error_body)
|
||||
}
|
||||
ClientStreamEmitter::OpenAIChat(_) | ClientStreamEmitter::Gemini(_) => {
|
||||
encode_json_sse(None, &error_body)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn build_client_error_body_for_line(report_context: &Value, line: &[u8]) -> Option<Value> {
|
||||
let value = decode_json_data_line(line)?;
|
||||
let provider_api_format = provider_api_format_for_context(report_context);
|
||||
let client_api_format = report_context
|
||||
.get("client_api_format")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.trim()
|
||||
.to_ascii_lowercase();
|
||||
let (message, code, kind) = parse_provider_error(&provider_api_format, &value)?;
|
||||
build_core_error_body_for_client_format(&client_api_format, &message, code.as_deref(), kind)
|
||||
}
|
||||
|
||||
fn parse_provider_error(
|
||||
provider_api_format: &str,
|
||||
payload: &Value,
|
||||
) -> Option<(String, Option<String>, LocalCoreSyncErrorKind)> {
|
||||
match FormatId::parse(provider_api_format)? {
|
||||
FormatId::OpenAiChat | FormatId::OpenAiResponses | FormatId::OpenAiResponsesCompact => {
|
||||
parse_openai_error(payload)
|
||||
}
|
||||
FormatId::ClaudeMessages => parse_claude_error(payload),
|
||||
FormatId::GeminiGenerateContent => parse_gemini_error(payload),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_openai_error(payload: &Value) -> Option<(String, Option<String>, LocalCoreSyncErrorKind)> {
|
||||
let error = payload.get("error")?.as_object()?;
|
||||
let message = error.get("message").and_then(Value::as_str)?.to_string();
|
||||
let code = error
|
||||
.get("code")
|
||||
.and_then(Value::as_str)
|
||||
.map(ToOwned::to_owned);
|
||||
let kind = match error
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
{
|
||||
"invalid_request_error" => LocalCoreSyncErrorKind::InvalidRequest,
|
||||
"authentication_error" => LocalCoreSyncErrorKind::Authentication,
|
||||
"permission_error" => LocalCoreSyncErrorKind::PermissionDenied,
|
||||
"not_found_error" => LocalCoreSyncErrorKind::NotFound,
|
||||
"rate_limit_error" => LocalCoreSyncErrorKind::RateLimit,
|
||||
"context_length_exceeded" => LocalCoreSyncErrorKind::ContextLengthExceeded,
|
||||
"overloaded_error" => LocalCoreSyncErrorKind::Overloaded,
|
||||
_ => LocalCoreSyncErrorKind::ServerError,
|
||||
};
|
||||
Some((message, code, kind))
|
||||
}
|
||||
|
||||
fn parse_claude_error(payload: &Value) -> Option<(String, Option<String>, LocalCoreSyncErrorKind)> {
|
||||
let error = payload.get("error")?.as_object()?;
|
||||
let message = error.get("message").and_then(Value::as_str)?.to_string();
|
||||
let code = error
|
||||
.get("code")
|
||||
.and_then(Value::as_str)
|
||||
.map(ToOwned::to_owned);
|
||||
let kind = match error
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
{
|
||||
"invalid_request_error" => LocalCoreSyncErrorKind::InvalidRequest,
|
||||
"authentication_error" => LocalCoreSyncErrorKind::Authentication,
|
||||
"permission_error" => LocalCoreSyncErrorKind::PermissionDenied,
|
||||
"not_found_error" => LocalCoreSyncErrorKind::NotFound,
|
||||
"rate_limit_error" => LocalCoreSyncErrorKind::RateLimit,
|
||||
"overloaded_error" => LocalCoreSyncErrorKind::Overloaded,
|
||||
_ => LocalCoreSyncErrorKind::ServerError,
|
||||
};
|
||||
Some((message, code, kind))
|
||||
}
|
||||
|
||||
fn parse_gemini_error(payload: &Value) -> Option<(String, Option<String>, LocalCoreSyncErrorKind)> {
|
||||
let error = payload.get("error")?.as_object()?;
|
||||
let message = error.get("message").and_then(Value::as_str)?.to_string();
|
||||
let code = error.get("code").map(|value| match value {
|
||||
Value::String(text) => text.clone(),
|
||||
Value::Number(number) => number.to_string(),
|
||||
_ => String::new(),
|
||||
});
|
||||
let kind = match error
|
||||
.get("status")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
{
|
||||
"INVALID_ARGUMENT" => LocalCoreSyncErrorKind::InvalidRequest,
|
||||
"UNAUTHENTICATED" => LocalCoreSyncErrorKind::Authentication,
|
||||
"PERMISSION_DENIED" => LocalCoreSyncErrorKind::PermissionDenied,
|
||||
"NOT_FOUND" => LocalCoreSyncErrorKind::NotFound,
|
||||
"RESOURCE_EXHAUSTED" => LocalCoreSyncErrorKind::RateLimit,
|
||||
"UNAVAILABLE" => LocalCoreSyncErrorKind::Overloaded,
|
||||
_ => LocalCoreSyncErrorKind::ServerError,
|
||||
};
|
||||
let code = code.filter(|value| !value.is_empty());
|
||||
Some((message, code, kind))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{StreamingStandardFormatMatrix, StreamingStandardTerminalObserver};
|
||||
use serde_json::{json, Value};
|
||||
|
||||
fn report_context(provider_api_format: &str, client_api_format: &str) -> Value {
|
||||
json!({
|
||||
"provider_api_format": provider_api_format,
|
||||
"client_api_format": client_api_format,
|
||||
"mapped_model": "test-model",
|
||||
})
|
||||
}
|
||||
|
||||
fn data_line(value: Value) -> Vec<u8> {
|
||||
format!("data: {}\n", value).into_bytes()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transforms_provider_errors_to_openai_chat_error_bodies() {
|
||||
let cases = [
|
||||
(
|
||||
"openai:chat",
|
||||
data_line(json!({
|
||||
"error": {
|
||||
"message": "bad request",
|
||||
"type": "invalid_request_error",
|
||||
"code": "invalid_request",
|
||||
}
|
||||
})),
|
||||
"\"message\":\"bad request\"",
|
||||
"\"type\":\"invalid_request_error\"",
|
||||
"\"code\":\"invalid_request\"",
|
||||
),
|
||||
(
|
||||
"claude:messages",
|
||||
data_line(json!({
|
||||
"type": "error",
|
||||
"error": {
|
||||
"message": "slow down",
|
||||
"type": "rate_limit_error",
|
||||
"code": "rate_limit",
|
||||
}
|
||||
})),
|
||||
"\"message\":\"slow down\"",
|
||||
"\"type\":\"rate_limit_error\"",
|
||||
"\"code\":\"rate_limit\"",
|
||||
),
|
||||
(
|
||||
"gemini:generate_content",
|
||||
data_line(json!({
|
||||
"error": {
|
||||
"code": 429,
|
||||
"message": "quota exceeded",
|
||||
"status": "RESOURCE_EXHAUSTED",
|
||||
}
|
||||
})),
|
||||
"\"message\":\"quota exceeded\"",
|
||||
"\"type\":\"rate_limit_error\"",
|
||||
"\"code\":\"429\"",
|
||||
),
|
||||
];
|
||||
|
||||
for (provider_api_format, line, message, err_type, code) in cases {
|
||||
let report_context = report_context(provider_api_format, "openai:chat");
|
||||
let mut matrix = StreamingStandardFormatMatrix::default();
|
||||
let output = matrix
|
||||
.transform_line(&report_context, line)
|
||||
.expect("error should convert");
|
||||
let sse = String::from_utf8(output).expect("sse should be utf8");
|
||||
|
||||
assert!(sse.starts_with("data: {\"error\":"));
|
||||
assert!(!sse.contains("event: "));
|
||||
assert!(sse.contains(message));
|
||||
assert!(sse.contains(err_type));
|
||||
assert!(sse.contains(code));
|
||||
assert!(matrix
|
||||
.finish(&report_context)
|
||||
.expect("finish should succeed")
|
||||
.is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transforms_provider_errors_to_claude_error_events() {
|
||||
let cases = [
|
||||
(
|
||||
"openai:chat",
|
||||
data_line(json!({
|
||||
"error": {
|
||||
"message": "bad request",
|
||||
"type": "invalid_request_error",
|
||||
"code": "invalid_request",
|
||||
}
|
||||
})),
|
||||
"\"message\":\"bad request\"",
|
||||
"\"type\":\"invalid_request_error\"",
|
||||
"\"code\":\"invalid_request\"",
|
||||
),
|
||||
(
|
||||
"claude:messages",
|
||||
data_line(json!({
|
||||
"type": "error",
|
||||
"error": {
|
||||
"message": "slow down",
|
||||
"type": "rate_limit_error",
|
||||
"code": "rate_limit",
|
||||
}
|
||||
})),
|
||||
"\"message\":\"slow down\"",
|
||||
"\"type\":\"rate_limit_error\"",
|
||||
"\"code\":\"rate_limit\"",
|
||||
),
|
||||
(
|
||||
"gemini:generate_content",
|
||||
data_line(json!({
|
||||
"error": {
|
||||
"code": 429,
|
||||
"message": "quota exceeded",
|
||||
"status": "RESOURCE_EXHAUSTED",
|
||||
}
|
||||
})),
|
||||
"\"message\":\"quota exceeded\"",
|
||||
"\"type\":\"rate_limit_error\"",
|
||||
"\"code\":\"429\"",
|
||||
),
|
||||
];
|
||||
|
||||
for (provider_api_format, line, message, err_type, code) in cases {
|
||||
let report_context = report_context(provider_api_format, "claude:messages");
|
||||
let mut matrix = StreamingStandardFormatMatrix::default();
|
||||
let output = matrix
|
||||
.transform_line(&report_context, line)
|
||||
.expect("error should convert");
|
||||
let sse = String::from_utf8(output).expect("sse should be utf8");
|
||||
|
||||
assert!(sse.starts_with("event: error\n"));
|
||||
assert!(sse.contains("data: {"));
|
||||
assert!(sse.contains("\"type\":\"error\""));
|
||||
assert!(sse.contains("\"error\":{"));
|
||||
assert!(sse.contains(message));
|
||||
assert!(sse.contains(err_type));
|
||||
assert!(sse.contains(code));
|
||||
assert!(matrix
|
||||
.finish(&report_context)
|
||||
.expect("finish should succeed")
|
||||
.is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transforms_provider_errors_to_gemini_error_bodies() {
|
||||
let cases = [
|
||||
(
|
||||
"openai:chat",
|
||||
data_line(json!({
|
||||
"error": {
|
||||
"message": "bad request",
|
||||
"type": "invalid_request_error",
|
||||
"code": "invalid_request",
|
||||
}
|
||||
})),
|
||||
"\"message\":\"bad request\"",
|
||||
"\"code\":400",
|
||||
"\"status\":\"INVALID_ARGUMENT\"",
|
||||
),
|
||||
(
|
||||
"claude:messages",
|
||||
data_line(json!({
|
||||
"type": "error",
|
||||
"error": {
|
||||
"message": "slow down",
|
||||
"type": "rate_limit_error",
|
||||
"code": "rate_limit",
|
||||
}
|
||||
})),
|
||||
"\"message\":\"slow down\"",
|
||||
"\"code\":429",
|
||||
"\"status\":\"RESOURCE_EXHAUSTED\"",
|
||||
),
|
||||
(
|
||||
"gemini:generate_content",
|
||||
data_line(json!({
|
||||
"error": {
|
||||
"code": 429,
|
||||
"message": "quota exceeded",
|
||||
"status": "RESOURCE_EXHAUSTED",
|
||||
}
|
||||
})),
|
||||
"\"message\":\"quota exceeded\"",
|
||||
"\"code\":429",
|
||||
"\"status\":\"RESOURCE_EXHAUSTED\"",
|
||||
),
|
||||
];
|
||||
|
||||
for (provider_api_format, line, message, code, status) in cases {
|
||||
let report_context = report_context(provider_api_format, "gemini:generate_content");
|
||||
let mut matrix = StreamingStandardFormatMatrix::default();
|
||||
let output = matrix
|
||||
.transform_line(&report_context, line)
|
||||
.expect("error should convert");
|
||||
let sse = String::from_utf8(output).expect("sse should be utf8");
|
||||
|
||||
assert!(sse.starts_with("data: {\"error\":"));
|
||||
assert!(!sse.contains("event: "));
|
||||
assert!(sse.contains(message));
|
||||
assert!(sse.contains(code));
|
||||
assert!(sse.contains(status));
|
||||
assert!(matrix
|
||||
.finish(&report_context)
|
||||
.expect("finish should succeed")
|
||||
.is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transforms_provider_errors_to_openai_responses_failed_events() {
|
||||
let cases = [
|
||||
(
|
||||
"openai:chat",
|
||||
data_line(json!({
|
||||
"error": {
|
||||
"message": "bad request",
|
||||
"type": "invalid_request_error",
|
||||
"code": "invalid_request",
|
||||
}
|
||||
})),
|
||||
"\"message\":\"bad request\"",
|
||||
"\"type\":\"invalid_request_error\"",
|
||||
"\"code\":\"invalid_request\"",
|
||||
),
|
||||
(
|
||||
"claude:messages",
|
||||
data_line(json!({
|
||||
"type": "error",
|
||||
"error": {
|
||||
"message": "slow down",
|
||||
"type": "rate_limit_error",
|
||||
"code": "rate_limit",
|
||||
}
|
||||
})),
|
||||
"\"message\":\"slow down\"",
|
||||
"\"type\":\"rate_limit_error\"",
|
||||
"\"code\":\"rate_limit\"",
|
||||
),
|
||||
(
|
||||
"gemini:generate_content",
|
||||
data_line(json!({
|
||||
"error": {
|
||||
"code": 429,
|
||||
"message": "quota exceeded",
|
||||
"status": "RESOURCE_EXHAUSTED",
|
||||
}
|
||||
})),
|
||||
"\"message\":\"quota exceeded\"",
|
||||
"\"type\":\"rate_limit_error\"",
|
||||
"\"code\":\"429\"",
|
||||
),
|
||||
];
|
||||
|
||||
for (provider_api_format, line, message, err_type, code) in cases {
|
||||
let report_context = report_context(provider_api_format, "openai:responses");
|
||||
let mut matrix = StreamingStandardFormatMatrix::default();
|
||||
let output = matrix
|
||||
.transform_line(&report_context, line)
|
||||
.expect("error should convert");
|
||||
let sse = String::from_utf8(output).expect("sse should be utf8");
|
||||
|
||||
assert!(sse.starts_with("event: response.failed\n"));
|
||||
assert!(sse.contains("\"sequence_number\":1"));
|
||||
assert!(sse.contains(message));
|
||||
assert!(sse.contains(err_type));
|
||||
assert!(sse.contains(code));
|
||||
assert!(matrix
|
||||
.finish(&report_context)
|
||||
.expect("finish should succeed")
|
||||
.is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rewrites_gemini_inline_image_streams_to_claude_image_blocks() {
|
||||
let report_context = report_context("gemini:generate_content", "claude:messages");
|
||||
let mut matrix = StreamingStandardFormatMatrix::default();
|
||||
let output = matrix
|
||||
.transform_line(
|
||||
&report_context,
|
||||
data_line(json!({
|
||||
"responseId": "resp_media_123",
|
||||
"modelVersion": "gemini-2.5-pro",
|
||||
"candidates": [{
|
||||
"index": 0,
|
||||
"content": {
|
||||
"parts": [
|
||||
{ "inlineData": { "mimeType": "image/png", "data": "iVBORw0KGgo=" } }
|
||||
]
|
||||
}
|
||||
}]
|
||||
})),
|
||||
)
|
||||
.expect("image chunk should rewrite");
|
||||
let sse = String::from_utf8(output).expect("sse should be utf8");
|
||||
|
||||
assert!(sse.contains("event: message_start"));
|
||||
assert!(sse.contains("\"type\":\"image\""));
|
||||
assert!(sse.contains("\"media_type\":\"image/png\""));
|
||||
assert!(sse.contains("\"data\":\"iVBORw0KGgo=\""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rewrites_claude_image_blocks_to_gemini_inline_image_streams() {
|
||||
let report_context = report_context("claude:messages", "gemini:generate_content");
|
||||
let mut matrix = StreamingStandardFormatMatrix::default();
|
||||
let output = matrix
|
||||
.transform_line(
|
||||
&report_context,
|
||||
data_line(json!({
|
||||
"type": "content_block_start",
|
||||
"index": 0,
|
||||
"content_block": {
|
||||
"type": "image",
|
||||
"source": {
|
||||
"type": "base64",
|
||||
"media_type": "image/png",
|
||||
"data": "iVBORw0KGgo="
|
||||
}
|
||||
}
|
||||
})),
|
||||
)
|
||||
.expect("image chunk should rewrite");
|
||||
let sse = String::from_utf8(output).expect("sse should be utf8");
|
||||
|
||||
assert!(
|
||||
sse.contains("\"inlineData\":{\"mimeType\":\"image/png\",\"data\":\"iVBORw0KGgo=\"}")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn terminal_observer_preserves_claude_cache_usage() {
|
||||
let report_context = report_context("claude:messages", "openai:chat");
|
||||
let mut observer = StreamingStandardTerminalObserver::default();
|
||||
|
||||
observer
|
||||
.push_line(
|
||||
&report_context,
|
||||
data_line(json!({
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": "msg_cache_123",
|
||||
"model": "claude-sonnet-4-5"
|
||||
}
|
||||
})),
|
||||
)
|
||||
.expect("message_start should parse");
|
||||
observer
|
||||
.push_line(
|
||||
&report_context,
|
||||
data_line(json!({
|
||||
"type": "message_delta",
|
||||
"delta": {
|
||||
"stop_reason": "end_turn"
|
||||
},
|
||||
"usage": {
|
||||
"input_tokens": 6,
|
||||
"output_tokens": 20,
|
||||
"cache_creation_input_tokens": 42262,
|
||||
"cache_read_input_tokens": 0
|
||||
}
|
||||
})),
|
||||
)
|
||||
.expect("message_delta should parse");
|
||||
|
||||
let summary = observer
|
||||
.latest_summary()
|
||||
.cloned()
|
||||
.expect("summary should exist");
|
||||
let usage = summary
|
||||
.standardized_usage
|
||||
.expect("standardized usage should exist");
|
||||
|
||||
assert_eq!(usage.input_tokens, 6);
|
||||
assert_eq!(usage.output_tokens, 20);
|
||||
assert_eq!(usage.cache_creation_tokens, 42_262);
|
||||
assert_eq!(usage.cache_read_tokens, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn terminal_observer_uses_explicit_provider_stream_event_api_format() {
|
||||
let mut report_context = report_context("openai:chat", "openai:responses");
|
||||
report_context["provider_stream_event_api_format"] = json!("openai:responses");
|
||||
let mut observer = StreamingStandardTerminalObserver::default();
|
||||
|
||||
observer
|
||||
.push_line(
|
||||
&report_context,
|
||||
data_line(json!({
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"id": "resp_codex_123",
|
||||
"object": "response",
|
||||
"model": "gpt-5.5",
|
||||
"status": "completed",
|
||||
"output": [],
|
||||
"usage": {
|
||||
"input_tokens": 26,
|
||||
"input_tokens_details": {
|
||||
"cached_tokens": 0,
|
||||
},
|
||||
"output_tokens": 137,
|
||||
"output_tokens_details": {
|
||||
"reasoning_tokens": 10,
|
||||
},
|
||||
"total_tokens": 163,
|
||||
},
|
||||
},
|
||||
"sequence_number": 139,
|
||||
})),
|
||||
)
|
||||
.expect("response.completed should parse");
|
||||
|
||||
let summary = observer
|
||||
.latest_summary()
|
||||
.cloned()
|
||||
.expect("summary should exist");
|
||||
let usage = summary
|
||||
.standardized_usage
|
||||
.expect("standardized usage should exist");
|
||||
|
||||
assert_eq!(usage.input_tokens, 26);
|
||||
assert_eq!(usage.output_tokens, 137);
|
||||
assert_eq!(usage.reasoning_tokens, 10);
|
||||
assert_eq!(usage.cache_read_tokens, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn terminal_observer_does_not_infer_provider_stream_event_api_format() {
|
||||
let report_context = report_context("openai:chat", "openai:responses");
|
||||
let mut observer = StreamingStandardTerminalObserver::default();
|
||||
|
||||
observer
|
||||
.push_line(
|
||||
&report_context,
|
||||
data_line(json!({
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"usage": {
|
||||
"input_tokens": 26,
|
||||
"output_tokens": 137,
|
||||
"total_tokens": 163,
|
||||
},
|
||||
},
|
||||
})),
|
||||
)
|
||||
.expect("line should be ignored by explicitly selected chat parser");
|
||||
|
||||
assert!(
|
||||
observer.latest_summary().is_none(),
|
||||
"provider stream parser selection must come from report context, not event sniffing"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn terminal_observer_counts_unknown_provider_stream_events() {
|
||||
let mut report_context = report_context("openai:chat", "openai:responses");
|
||||
report_context["provider_stream_event_api_format"] = json!("openai:responses");
|
||||
let mut observer = StreamingStandardTerminalObserver::default();
|
||||
|
||||
observer
|
||||
.push_line(
|
||||
&report_context,
|
||||
data_line(json!({
|
||||
"type": "response.future.delta",
|
||||
"response": {
|
||||
"id": "resp_unknown_123",
|
||||
"model": "gpt-5.4",
|
||||
},
|
||||
"payload": {
|
||||
"kept": true,
|
||||
},
|
||||
})),
|
||||
)
|
||||
.expect("unknown stream event should be observed");
|
||||
|
||||
let summary = observer
|
||||
.latest_summary()
|
||||
.cloned()
|
||||
.expect("summary should exist");
|
||||
assert_eq!(summary.response_id.as_deref(), Some("resp_unknown_123"));
|
||||
assert_eq!(summary.model.as_deref(), Some("gpt-5.4"));
|
||||
assert_eq!(summary.unknown_event_count, 1);
|
||||
assert!(!summary.observed_finish);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
pub mod common;
|
||||
pub mod format_matrix;
|
||||
|
||||
pub use common::{CanonicalStreamEvent, CanonicalStreamFrame};
|
||||
pub use format_matrix::{StreamingStandardFormatMatrix, StreamingStandardTerminalObserver};
|
||||
356
crates/aether-ai-formats/src/response/stream_rewrite.rs
Normal file
356
crates/aether-ai-formats/src/response/stream_rewrite.rs
Normal file
@@ -0,0 +1,356 @@
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::provider_compat::kiro_stream::KiroToClaudeCliStreamState;
|
||||
use crate::provider_compat::private_envelope::transform_provider_private_stream_line;
|
||||
use crate::provider_compat::surfaces::{
|
||||
provider_adaptation_should_unwrap_stream_envelope, KIRO_ENVELOPE_NAME,
|
||||
};
|
||||
use crate::response::openai_image_stream::OpenAiImageStreamState;
|
||||
use crate::response::standard::stream_core::StreamingStandardFormatMatrix;
|
||||
use crate::response::AiSurfaceFinalizeError;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum FinalizeStreamRewriteMode {
|
||||
EnvelopeUnwrap,
|
||||
OpenAiImage,
|
||||
Standard,
|
||||
KiroToClaudeCli,
|
||||
KiroToClaudeCliThenStandard,
|
||||
}
|
||||
|
||||
pub fn resolve_finalize_stream_rewrite_mode(
|
||||
report_context: &Value,
|
||||
) -> Option<FinalizeStreamRewriteMode> {
|
||||
let needs_conversion = report_context
|
||||
.get("needs_conversion")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false);
|
||||
let envelope_name = report_context
|
||||
.get("envelope_name")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.trim()
|
||||
.to_ascii_lowercase();
|
||||
let provider_api_format = report_context
|
||||
.get("provider_api_format")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.trim()
|
||||
.to_ascii_lowercase();
|
||||
let client_api_format = report_context
|
||||
.get("client_api_format")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.trim()
|
||||
.to_ascii_lowercase();
|
||||
|
||||
if needs_conversion
|
||||
&& envelope_name.eq_ignore_ascii_case(KIRO_ENVELOPE_NAME)
|
||||
&& provider_api_format == "claude:messages"
|
||||
{
|
||||
return supports_standard_stream_rewrite(
|
||||
provider_api_format.as_str(),
|
||||
client_api_format.as_str(),
|
||||
)
|
||||
.then_some(FinalizeStreamRewriteMode::KiroToClaudeCliThenStandard);
|
||||
}
|
||||
|
||||
if needs_conversion {
|
||||
return supports_standard_stream_rewrite(
|
||||
provider_api_format.as_str(),
|
||||
client_api_format.as_str(),
|
||||
)
|
||||
.then_some(FinalizeStreamRewriteMode::Standard);
|
||||
}
|
||||
|
||||
if provider_api_format == "openai:image" && client_api_format == "openai:image" {
|
||||
return Some(FinalizeStreamRewriteMode::OpenAiImage);
|
||||
}
|
||||
|
||||
if envelope_name.eq_ignore_ascii_case(KIRO_ENVELOPE_NAME) {
|
||||
return (provider_api_format == "claude:messages"
|
||||
&& client_api_format == "claude:messages")
|
||||
.then_some(FinalizeStreamRewriteMode::KiroToClaudeCli);
|
||||
}
|
||||
|
||||
(provider_api_format == client_api_format
|
||||
&& provider_adaptation_should_unwrap_stream_envelope(
|
||||
envelope_name.as_str(),
|
||||
provider_api_format.as_str(),
|
||||
))
|
||||
.then_some(FinalizeStreamRewriteMode::EnvelopeUnwrap)
|
||||
}
|
||||
|
||||
enum AiSurfaceStreamRewriteState {
|
||||
EnvelopeUnwrap,
|
||||
OpenAiImage(Box<OpenAiImageStreamState>),
|
||||
Standard(Box<StreamingStandardFormatMatrix>),
|
||||
KiroToClaudeCli(Box<KiroToClaudeCliStreamState>),
|
||||
KiroToClaudeCliThenStandard {
|
||||
kiro: Box<KiroToClaudeCliStreamState>,
|
||||
standard: Box<StreamingStandardFormatMatrix>,
|
||||
},
|
||||
}
|
||||
|
||||
pub struct AiSurfaceStreamRewriter<'a> {
|
||||
report_context: &'a Value,
|
||||
buffered: Vec<u8>,
|
||||
state: AiSurfaceStreamRewriteState,
|
||||
}
|
||||
|
||||
pub fn maybe_build_ai_surface_stream_rewriter<'a>(
|
||||
report_context: Option<&'a Value>,
|
||||
) -> Option<AiSurfaceStreamRewriter<'a>> {
|
||||
let report_context = report_context?;
|
||||
let state = match resolve_finalize_stream_rewrite_mode(report_context)? {
|
||||
FinalizeStreamRewriteMode::EnvelopeUnwrap => AiSurfaceStreamRewriteState::EnvelopeUnwrap,
|
||||
FinalizeStreamRewriteMode::OpenAiImage => {
|
||||
AiSurfaceStreamRewriteState::OpenAiImage(Box::<OpenAiImageStreamState>::default())
|
||||
}
|
||||
FinalizeStreamRewriteMode::Standard => {
|
||||
AiSurfaceStreamRewriteState::Standard(Box::<StreamingStandardFormatMatrix>::default())
|
||||
}
|
||||
FinalizeStreamRewriteMode::KiroToClaudeCli => AiSurfaceStreamRewriteState::KiroToClaudeCli(
|
||||
Box::new(KiroToClaudeCliStreamState::new(report_context)),
|
||||
),
|
||||
FinalizeStreamRewriteMode::KiroToClaudeCliThenStandard => {
|
||||
AiSurfaceStreamRewriteState::KiroToClaudeCliThenStandard {
|
||||
kiro: Box::new(KiroToClaudeCliStreamState::new(report_context)),
|
||||
standard: Box::<StreamingStandardFormatMatrix>::default(),
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
Some(AiSurfaceStreamRewriter {
|
||||
report_context,
|
||||
buffered: Vec::new(),
|
||||
state,
|
||||
})
|
||||
}
|
||||
|
||||
impl AiSurfaceStreamRewriter<'_> {
|
||||
pub fn push_chunk(&mut self, chunk: &[u8]) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
match &mut self.state {
|
||||
AiSurfaceStreamRewriteState::OpenAiImage(state) => {
|
||||
state.push_chunk(self.report_context, chunk)
|
||||
}
|
||||
AiSurfaceStreamRewriteState::KiroToClaudeCli(state) => {
|
||||
state.push_chunk(self.report_context, chunk)
|
||||
}
|
||||
AiSurfaceStreamRewriteState::KiroToClaudeCliThenStandard { kiro, standard } => {
|
||||
let claude_bytes = kiro.push_chunk(self.report_context, chunk)?;
|
||||
transform_standard_bytes(standard, self.report_context, claude_bytes)
|
||||
}
|
||||
AiSurfaceStreamRewriteState::EnvelopeUnwrap
|
||||
| AiSurfaceStreamRewriteState::Standard(_) => {
|
||||
self.buffered.extend_from_slice(chunk);
|
||||
let mut output = Vec::new();
|
||||
while let Some(line_end) = self.buffered.iter().position(|byte| *byte == b'\n') {
|
||||
let line = self.buffered.drain(..=line_end).collect::<Vec<_>>();
|
||||
output.extend(self.transform_line(line)?);
|
||||
}
|
||||
Ok(output)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn finish(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
match &mut self.state {
|
||||
AiSurfaceStreamRewriteState::OpenAiImage(state) => state.finish(self.report_context),
|
||||
AiSurfaceStreamRewriteState::KiroToClaudeCli(state) => {
|
||||
state.finish(self.report_context)
|
||||
}
|
||||
AiSurfaceStreamRewriteState::KiroToClaudeCliThenStandard { kiro, standard } => {
|
||||
let mut output = transform_standard_bytes(
|
||||
standard,
|
||||
self.report_context,
|
||||
kiro.finish(self.report_context)?,
|
||||
)?;
|
||||
output.extend(standard.finish(self.report_context)?);
|
||||
Ok(output)
|
||||
}
|
||||
AiSurfaceStreamRewriteState::EnvelopeUnwrap
|
||||
| AiSurfaceStreamRewriteState::Standard(_) => {
|
||||
if self.buffered.is_empty() {
|
||||
if let AiSurfaceStreamRewriteState::Standard(state) = &mut self.state {
|
||||
return state.finish(self.report_context);
|
||||
}
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let line = std::mem::take(&mut self.buffered);
|
||||
let mut output = self.transform_line(line)?;
|
||||
if let AiSurfaceStreamRewriteState::Standard(state) = &mut self.state {
|
||||
output.extend(state.finish(self.report_context)?);
|
||||
}
|
||||
Ok(output)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn transform_line(&mut self, line: Vec<u8>) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
match &mut self.state {
|
||||
AiSurfaceStreamRewriteState::EnvelopeUnwrap => {
|
||||
transform_provider_private_stream_line(self.report_context, line)
|
||||
.map_err(AiSurfaceFinalizeError::from)
|
||||
}
|
||||
AiSurfaceStreamRewriteState::Standard(state) => {
|
||||
transform_standard_line(state, self.report_context, line)
|
||||
}
|
||||
AiSurfaceStreamRewriteState::OpenAiImage(_)
|
||||
| AiSurfaceStreamRewriteState::KiroToClaudeCli(_)
|
||||
| AiSurfaceStreamRewriteState::KiroToClaudeCliThenStandard { .. } => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn transform_standard_bytes(
|
||||
standard: &mut StreamingStandardFormatMatrix,
|
||||
report_context: &Value,
|
||||
bytes: Vec<u8>,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
if bytes.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut output = Vec::new();
|
||||
for line in bytes.split_inclusive(|byte| *byte == b'\n') {
|
||||
output.extend(transform_standard_line(
|
||||
standard,
|
||||
report_context,
|
||||
line.to_vec(),
|
||||
)?);
|
||||
}
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn transform_standard_line(
|
||||
standard: &mut StreamingStandardFormatMatrix,
|
||||
report_context: &Value,
|
||||
line: Vec<u8>,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
let line = if should_unwrap_envelope(report_context) {
|
||||
transform_provider_private_stream_line(report_context, line)?
|
||||
} else {
|
||||
line
|
||||
};
|
||||
if line.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
standard.transform_line(report_context, line)
|
||||
}
|
||||
|
||||
fn should_unwrap_envelope(report_context: &Value) -> bool {
|
||||
let envelope_name = report_context
|
||||
.get("envelope_name")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
let provider_api_format = report_context
|
||||
.get("provider_api_format")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
provider_adaptation_should_unwrap_stream_envelope(envelope_name, provider_api_format)
|
||||
}
|
||||
|
||||
fn supports_standard_stream_rewrite(provider_api_format: &str, client_api_format: &str) -> bool {
|
||||
is_standard_provider_api_format(provider_api_format)
|
||||
&& (is_standard_chat_client_api_format(client_api_format)
|
||||
|| is_standard_cli_client_api_format(client_api_format))
|
||||
}
|
||||
|
||||
fn is_standard_provider_api_format(api_format: &str) -> bool {
|
||||
matches!(
|
||||
aether_ai_formats::normalize_api_format_alias(api_format).as_str(),
|
||||
"openai:chat"
|
||||
| "openai:responses"
|
||||
| "openai:responses:compact"
|
||||
| "claude:messages"
|
||||
| "gemini:generate_content"
|
||||
)
|
||||
}
|
||||
|
||||
fn is_standard_chat_client_api_format(api_format: &str) -> bool {
|
||||
matches!(
|
||||
api_format,
|
||||
"openai:chat" | "claude:messages" | "gemini:generate_content"
|
||||
)
|
||||
}
|
||||
|
||||
fn is_standard_cli_client_api_format(api_format: &str) -> bool {
|
||||
matches!(
|
||||
aether_ai_formats::normalize_api_format_alias(api_format).as_str(),
|
||||
"openai:responses"
|
||||
| "openai:responses:compact"
|
||||
| "claude:messages"
|
||||
| "gemini:generate_content"
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::{resolve_finalize_stream_rewrite_mode, FinalizeStreamRewriteMode};
|
||||
|
||||
#[test]
|
||||
fn resolves_standard_mode_for_cross_format_standard_streams() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "claude:messages",
|
||||
"client_api_format": "openai:chat",
|
||||
"needs_conversion": true,
|
||||
});
|
||||
assert_eq!(
|
||||
resolve_finalize_stream_rewrite_mode(&report_context),
|
||||
Some(FinalizeStreamRewriteMode::Standard)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_envelope_unwrap_for_same_format_private_envelopes() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "gemini:generate_content",
|
||||
"client_api_format": "gemini:generate_content",
|
||||
"envelope_name": "antigravity:v1internal",
|
||||
"needs_conversion": false,
|
||||
});
|
||||
assert_eq!(
|
||||
resolve_finalize_stream_rewrite_mode(&report_context),
|
||||
Some(FinalizeStreamRewriteMode::EnvelopeUnwrap)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_kiro_same_format_streams_to_kiro_mode() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "claude:messages",
|
||||
"client_api_format": "claude:messages",
|
||||
"envelope_name": "kiro:generateAssistantResponse",
|
||||
"needs_conversion": false,
|
||||
});
|
||||
assert_eq!(
|
||||
resolve_finalize_stream_rewrite_mode(&report_context),
|
||||
Some(FinalizeStreamRewriteMode::KiroToClaudeCli)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_unsupported_non_conversion_streams() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "openai:chat",
|
||||
"client_api_format": "openai:chat",
|
||||
"needs_conversion": false,
|
||||
});
|
||||
assert_eq!(resolve_finalize_stream_rewrite_mode(&report_context), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_openai_image_mode_for_same_format_image_streams() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "openai:image",
|
||||
"client_api_format": "openai:image",
|
||||
"needs_conversion": false,
|
||||
});
|
||||
assert_eq!(
|
||||
resolve_finalize_stream_rewrite_mode(&report_context),
|
||||
Some(FinalizeStreamRewriteMode::OpenAiImage)
|
||||
);
|
||||
}
|
||||
}
|
||||
4446
crates/aether-ai-formats/src/response/sync_products.rs
Normal file
4446
crates/aether-ai-formats/src/response/sync_products.rs
Normal file
File diff suppressed because it is too large
Load Diff
569
crates/aether-ai-formats/src/response/sync_to_stream.rs
Normal file
569
crates/aether-ai-formats/src/response/sync_to_stream.rs
Normal file
@@ -0,0 +1,569 @@
|
||||
use aether_ai_formats::protocol::conversion::response::{
|
||||
convert_claude_response_to_openai_responses, convert_gemini_response_to_openai_responses,
|
||||
convert_openai_chat_response_to_openai_responses,
|
||||
};
|
||||
use aether_contracts::{ExecutionStreamTerminalSummary, StandardizedUsage};
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use crate::response::sse::encode_json_sse;
|
||||
use crate::response::standard::claude::stream::ClaudeClientEmitter;
|
||||
use crate::response::standard::gemini::stream::GeminiClientEmitter;
|
||||
use crate::response::standard::openai::stream::{
|
||||
OpenAIChatClientEmitter, OpenAIResponsesClientEmitter, OpenAIResponsesProviderState,
|
||||
};
|
||||
use crate::response::standard::stream_core::CanonicalStreamFrame;
|
||||
use crate::response::AiSurfaceFinalizeError;
|
||||
|
||||
pub struct SyncToStreamBridgeOutcome {
|
||||
pub sse_body: Vec<u8>,
|
||||
pub terminal_summary: Option<ExecutionStreamTerminalSummary>,
|
||||
}
|
||||
|
||||
pub fn maybe_bridge_standard_sync_json_to_stream(
|
||||
provider_body_json: &Value,
|
||||
provider_api_format: &str,
|
||||
client_api_format: &str,
|
||||
report_context: Option<&Value>,
|
||||
) -> Result<Option<SyncToStreamBridgeOutcome>, AiSurfaceFinalizeError> {
|
||||
let provider_api_format = normalize_api_format(provider_api_format);
|
||||
let client_api_format = normalize_api_format(client_api_format);
|
||||
if provider_api_format == "openai:image" && client_api_format == "openai:image" {
|
||||
return maybe_bridge_openai_image_sync_json_to_stream(provider_body_json, report_context);
|
||||
}
|
||||
if !is_standard_api_format(provider_api_format.as_str())
|
||||
|| !is_standard_api_format(client_api_format.as_str())
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let bridge_context = build_bridge_report_context(
|
||||
report_context,
|
||||
provider_api_format.as_str(),
|
||||
client_api_format.as_str(),
|
||||
);
|
||||
let Some(openai_responses_response) = convert_provider_sync_response_to_openai_responses(
|
||||
provider_body_json,
|
||||
provider_api_format.as_str(),
|
||||
&bridge_context,
|
||||
) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let terminal_summary =
|
||||
build_terminal_summary_from_openai_responses_response(&openai_responses_response);
|
||||
let canonical_frames = build_canonical_frames_from_openai_responses_response(
|
||||
&openai_responses_response,
|
||||
&bridge_context,
|
||||
)?;
|
||||
let sse_body =
|
||||
emit_client_stream_from_canonical_frames(canonical_frames, client_api_format.as_str())?;
|
||||
|
||||
Ok(Some(SyncToStreamBridgeOutcome {
|
||||
sse_body,
|
||||
terminal_summary,
|
||||
}))
|
||||
}
|
||||
|
||||
fn maybe_bridge_openai_image_sync_json_to_stream(
|
||||
provider_body_json: &Value,
|
||||
report_context: Option<&Value>,
|
||||
) -> Result<Option<SyncToStreamBridgeOutcome>, AiSurfaceFinalizeError> {
|
||||
let Some(response) = provider_body_json.as_object() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(image) = response
|
||||
.get("data")
|
||||
.and_then(Value::as_array)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter_map(Value::as_object)
|
||||
.find_map(extract_openai_image_sync_b64_json)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let usage = response.get("usage").cloned().unwrap_or(Value::Null);
|
||||
let event_name = openai_image_completed_event_name(report_context);
|
||||
let sse_body = encode_json_sse(
|
||||
Some(event_name),
|
||||
&json!({
|
||||
"type": event_name,
|
||||
"b64_json": image,
|
||||
"usage": usage,
|
||||
}),
|
||||
)?;
|
||||
|
||||
Ok(Some(SyncToStreamBridgeOutcome {
|
||||
sse_body,
|
||||
terminal_summary: Some(ExecutionStreamTerminalSummary {
|
||||
standardized_usage: response
|
||||
.get("usage")
|
||||
.and_then(standardized_usage_from_openai_usage),
|
||||
finish_reason: Some("stop".to_string()),
|
||||
response_id: response
|
||||
.get("id")
|
||||
.and_then(Value::as_str)
|
||||
.map(ToOwned::to_owned),
|
||||
model: response
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.map(ToOwned::to_owned)
|
||||
.or_else(|| image_bridge_model(report_context)),
|
||||
observed_finish: true,
|
||||
unknown_event_count: 0,
|
||||
parser_error: None,
|
||||
}),
|
||||
}))
|
||||
}
|
||||
|
||||
fn normalize_api_format(value: &str) -> String {
|
||||
aether_ai_formats::normalize_api_format_alias(value)
|
||||
}
|
||||
|
||||
fn is_standard_api_format(value: &str) -> bool {
|
||||
matches!(
|
||||
value,
|
||||
"openai:chat"
|
||||
| "openai:responses"
|
||||
| "openai:responses:compact"
|
||||
| "claude:messages"
|
||||
| "gemini:generate_content"
|
||||
)
|
||||
}
|
||||
|
||||
fn extract_openai_image_sync_b64_json(item: &serde_json::Map<String, Value>) -> Option<String> {
|
||||
item.get("b64_json")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
.or_else(|| {
|
||||
item.get("url")
|
||||
.and_then(Value::as_str)
|
||||
.and_then(extract_base64_from_data_url)
|
||||
})
|
||||
}
|
||||
|
||||
fn extract_base64_from_data_url(value: &str) -> Option<String> {
|
||||
let trimmed = value.trim();
|
||||
let (metadata, payload) = trimmed.split_once(',')?;
|
||||
if !metadata.starts_with("data:") || !metadata.ends_with(";base64") {
|
||||
return None;
|
||||
}
|
||||
(!payload.trim().is_empty()).then(|| payload.trim().to_string())
|
||||
}
|
||||
|
||||
fn openai_image_completed_event_name(report_context: Option<&Value>) -> &'static str {
|
||||
if openai_image_request_operation(report_context) == Some("edit") {
|
||||
"image_edit.completed"
|
||||
} else {
|
||||
"image_generation.completed"
|
||||
}
|
||||
}
|
||||
|
||||
fn openai_image_request_operation(report_context: Option<&Value>) -> Option<&str> {
|
||||
report_context
|
||||
.and_then(|value| value.get("image_request"))
|
||||
.and_then(|value| value.get("operation"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn image_bridge_model(report_context: Option<&Value>) -> Option<String> {
|
||||
report_context.and_then(|context| {
|
||||
context
|
||||
.get("mapped_model")
|
||||
.or_else(|| context.get("model"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
})
|
||||
}
|
||||
|
||||
fn build_bridge_report_context(
|
||||
report_context: Option<&Value>,
|
||||
provider_api_format: &str,
|
||||
client_api_format: &str,
|
||||
) -> Value {
|
||||
let mut context = report_context
|
||||
.cloned()
|
||||
.filter(Value::is_object)
|
||||
.unwrap_or_else(|| json!({}));
|
||||
let object = context
|
||||
.as_object_mut()
|
||||
.expect("bridge report context should stay object");
|
||||
object
|
||||
.entry("provider_api_format".to_string())
|
||||
.or_insert_with(|| Value::String(provider_api_format.to_string()));
|
||||
object
|
||||
.entry("client_api_format".to_string())
|
||||
.or_insert_with(|| Value::String(client_api_format.to_string()));
|
||||
context
|
||||
}
|
||||
|
||||
fn convert_provider_sync_response_to_openai_responses(
|
||||
provider_body_json: &Value,
|
||||
provider_api_format: &str,
|
||||
report_context: &Value,
|
||||
) -> Option<Value> {
|
||||
match provider_api_format {
|
||||
"openai:responses" | "openai:responses:compact" => Some(provider_body_json.clone()),
|
||||
"openai:chat" => convert_openai_chat_response_to_openai_responses(
|
||||
provider_body_json,
|
||||
report_context,
|
||||
false,
|
||||
),
|
||||
"claude:messages" => {
|
||||
convert_claude_response_to_openai_responses(provider_body_json, report_context)
|
||||
}
|
||||
"gemini:generate_content" => {
|
||||
convert_gemini_response_to_openai_responses(provider_body_json, report_context)
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn build_canonical_frames_from_openai_responses_response(
|
||||
openai_responses_response: &Value,
|
||||
report_context: &Value,
|
||||
) -> Result<Vec<CanonicalStreamFrame>, AiSurfaceFinalizeError> {
|
||||
let mut state = OpenAIResponsesProviderState::default();
|
||||
let line = format!(
|
||||
"data: {}\n",
|
||||
serde_json::to_string(&json!({
|
||||
"type": "response.completed",
|
||||
"response": openai_responses_response,
|
||||
}))
|
||||
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?
|
||||
);
|
||||
let mut frames = state
|
||||
.push_line(report_context, line.into_bytes())
|
||||
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?;
|
||||
frames.extend(
|
||||
state
|
||||
.finish(report_context)
|
||||
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?,
|
||||
);
|
||||
Ok(frames)
|
||||
}
|
||||
|
||||
fn emit_client_stream_from_canonical_frames(
|
||||
canonical_frames: Vec<CanonicalStreamFrame>,
|
||||
client_api_format: &str,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
match client_api_format {
|
||||
"openai:chat" => {
|
||||
let mut emitter = OpenAIChatClientEmitter::default();
|
||||
emit_with_openai_chat_emitter(&mut emitter, canonical_frames)
|
||||
}
|
||||
"openai:responses" | "openai:responses:compact" => {
|
||||
let mut emitter = OpenAIResponsesClientEmitter::default();
|
||||
emit_with_openai_responses_emitter(&mut emitter, canonical_frames)
|
||||
}
|
||||
"claude:messages" => {
|
||||
let mut emitter = ClaudeClientEmitter::default();
|
||||
emit_with_claude_emitter(&mut emitter, canonical_frames)
|
||||
}
|
||||
"gemini:generate_content" => {
|
||||
let mut emitter = GeminiClientEmitter::default();
|
||||
emit_with_gemini_emitter(&mut emitter, canonical_frames)
|
||||
}
|
||||
_ => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
fn emit_with_openai_chat_emitter(
|
||||
emitter: &mut OpenAIChatClientEmitter,
|
||||
canonical_frames: Vec<CanonicalStreamFrame>,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
let mut output = Vec::new();
|
||||
for frame in canonical_frames {
|
||||
output.extend(
|
||||
emitter
|
||||
.emit(frame)
|
||||
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?,
|
||||
);
|
||||
}
|
||||
output.extend(
|
||||
emitter
|
||||
.finish()
|
||||
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?,
|
||||
);
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn emit_with_openai_responses_emitter(
|
||||
emitter: &mut OpenAIResponsesClientEmitter,
|
||||
canonical_frames: Vec<CanonicalStreamFrame>,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
let mut output = Vec::new();
|
||||
for frame in canonical_frames {
|
||||
output.extend(
|
||||
emitter
|
||||
.emit(frame)
|
||||
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?,
|
||||
);
|
||||
}
|
||||
output.extend(
|
||||
emitter
|
||||
.finish()
|
||||
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?,
|
||||
);
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn emit_with_claude_emitter(
|
||||
emitter: &mut ClaudeClientEmitter,
|
||||
canonical_frames: Vec<CanonicalStreamFrame>,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
let mut output = Vec::new();
|
||||
for frame in canonical_frames {
|
||||
output.extend(
|
||||
emitter
|
||||
.emit(frame)
|
||||
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?,
|
||||
);
|
||||
}
|
||||
output.extend(
|
||||
emitter
|
||||
.finish()
|
||||
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?,
|
||||
);
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn emit_with_gemini_emitter(
|
||||
emitter: &mut GeminiClientEmitter,
|
||||
canonical_frames: Vec<CanonicalStreamFrame>,
|
||||
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
||||
let mut output = Vec::new();
|
||||
for frame in canonical_frames {
|
||||
output.extend(
|
||||
emitter
|
||||
.emit(frame)
|
||||
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?,
|
||||
);
|
||||
}
|
||||
output.extend(
|
||||
emitter
|
||||
.finish()
|
||||
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?,
|
||||
);
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn build_terminal_summary_from_openai_responses_response(
|
||||
openai_responses_response: &Value,
|
||||
) -> Option<ExecutionStreamTerminalSummary> {
|
||||
let response = openai_responses_response.as_object()?;
|
||||
let response_id = response
|
||||
.get("id")
|
||||
.and_then(Value::as_str)
|
||||
.map(ToOwned::to_owned);
|
||||
let model = response
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.map(ToOwned::to_owned);
|
||||
let finish_reason = response
|
||||
.get("output")
|
||||
.and_then(Value::as_array)
|
||||
.map(|output| resolve_openai_responses_finish_reason(output))
|
||||
.filter(|value| !value.trim().is_empty());
|
||||
let standardized_usage = response
|
||||
.get("usage")
|
||||
.and_then(standardized_usage_from_openai_usage);
|
||||
Some(ExecutionStreamTerminalSummary {
|
||||
standardized_usage,
|
||||
finish_reason,
|
||||
response_id,
|
||||
model,
|
||||
observed_finish: true,
|
||||
unknown_event_count: 0,
|
||||
parser_error: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn resolve_openai_responses_finish_reason(output: &[Value]) -> String {
|
||||
let has_tool_calls = output.iter().filter_map(Value::as_object).any(|item| {
|
||||
item.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|value| value == "function_call")
|
||||
});
|
||||
if has_tool_calls {
|
||||
"tool_calls".to_string()
|
||||
} else {
|
||||
"stop".to_string()
|
||||
}
|
||||
}
|
||||
|
||||
fn standardized_usage_from_openai_usage(value: &Value) -> Option<StandardizedUsage> {
|
||||
let usage = value.as_object()?;
|
||||
let mut input_tokens = usage
|
||||
.get("input_tokens")
|
||||
.or_else(|| usage.get("prompt_tokens"))
|
||||
.and_then(Value::as_i64)
|
||||
.unwrap_or(0);
|
||||
let output_tokens = usage
|
||||
.get("output_tokens")
|
||||
.or_else(|| usage.get("completion_tokens"))
|
||||
.and_then(Value::as_i64)
|
||||
.unwrap_or(0);
|
||||
let cache_creation_tokens = usage
|
||||
.get("cache_creation_input_tokens")
|
||||
.and_then(Value::as_i64)
|
||||
.or_else(|| {
|
||||
usage
|
||||
.get("input_tokens_details")
|
||||
.or_else(|| usage.get("prompt_tokens_details"))
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|details| details.get("cached_creation_tokens"))
|
||||
.and_then(Value::as_i64)
|
||||
})
|
||||
.unwrap_or(0);
|
||||
let cache_read_tokens = usage
|
||||
.get("cache_read_input_tokens")
|
||||
.and_then(Value::as_i64)
|
||||
.or_else(|| {
|
||||
usage
|
||||
.get("input_tokens_details")
|
||||
.or_else(|| usage.get("prompt_tokens_details"))
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|details| details.get("cached_tokens"))
|
||||
.and_then(Value::as_i64)
|
||||
})
|
||||
.unwrap_or(0);
|
||||
let total_tokens = usage.get("total_tokens").and_then(Value::as_i64).unwrap_or(
|
||||
input_tokens
|
||||
.saturating_add(output_tokens)
|
||||
.saturating_add(cache_creation_tokens)
|
||||
.saturating_add(cache_read_tokens),
|
||||
);
|
||||
if input_tokens == 0 && total_tokens > output_tokens {
|
||||
input_tokens = total_tokens.saturating_sub(output_tokens);
|
||||
}
|
||||
let mut standardized_usage = StandardizedUsage::new();
|
||||
standardized_usage.input_tokens = input_tokens;
|
||||
standardized_usage.output_tokens = output_tokens;
|
||||
standardized_usage.cache_creation_tokens = cache_creation_tokens;
|
||||
standardized_usage.cache_read_tokens = cache_read_tokens;
|
||||
standardized_usage
|
||||
.dimensions
|
||||
.insert("total_tokens".to_string(), json!(total_tokens));
|
||||
Some(standardized_usage.normalize_cache_creation_breakdown())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::{maybe_bridge_standard_sync_json_to_stream, standardized_usage_from_openai_usage};
|
||||
|
||||
fn utf8(bytes: Vec<u8>) -> String {
|
||||
String::from_utf8(bytes).expect("utf8 should decode")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_sync_usage_derives_missing_input_tokens_from_total() {
|
||||
let usage = standardized_usage_from_openai_usage(&json!({
|
||||
"output_tokens": 177,
|
||||
"total_tokens": 20_612,
|
||||
"input_tokens_details": {
|
||||
"cached_tokens": 19_840,
|
||||
},
|
||||
}))
|
||||
.expect("usage should parse");
|
||||
|
||||
assert_eq!(usage.input_tokens, 20_435);
|
||||
assert_eq!(usage.output_tokens, 177);
|
||||
assert_eq!(usage.cache_read_tokens, 19_840);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bridges_openai_image_sync_json_to_generation_completed_sse() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "openai:image",
|
||||
"client_api_format": "openai:image",
|
||||
"mapped_model": "gpt-image-1",
|
||||
"image_request": {
|
||||
"operation": "generate"
|
||||
}
|
||||
});
|
||||
let outcome = maybe_bridge_standard_sync_json_to_stream(
|
||||
&json!({
|
||||
"created": 1776971267,
|
||||
"data": [{
|
||||
"b64_json": "aGVsbG8="
|
||||
}],
|
||||
"usage": {
|
||||
"total_tokens": 100,
|
||||
"input_tokens": 50,
|
||||
"output_tokens": 50,
|
||||
"input_tokens_details": {
|
||||
"text_tokens": 10,
|
||||
"image_tokens": 40
|
||||
}
|
||||
}
|
||||
}),
|
||||
"openai:image",
|
||||
"openai:image",
|
||||
Some(&report_context),
|
||||
)
|
||||
.expect("bridge should succeed")
|
||||
.expect("bridge should produce sse");
|
||||
|
||||
let output = utf8(outcome.sse_body);
|
||||
assert!(output.contains("event: image_generation.completed"));
|
||||
assert!(output.contains("\"type\":\"image_generation.completed\""));
|
||||
assert!(output.contains("\"b64_json\":\"aGVsbG8=\""));
|
||||
assert!(output.contains("\"total_tokens\":100"));
|
||||
|
||||
let summary = outcome
|
||||
.terminal_summary
|
||||
.expect("terminal summary should exist");
|
||||
assert_eq!(summary.model.as_deref(), Some("gpt-image-1"));
|
||||
assert_eq!(summary.finish_reason.as_deref(), Some("stop"));
|
||||
assert_eq!(
|
||||
summary
|
||||
.standardized_usage
|
||||
.as_ref()
|
||||
.and_then(|usage| usage.dimensions.get("total_tokens"))
|
||||
.cloned(),
|
||||
Some(json!(100))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bridges_openai_image_sync_data_url_to_edit_completed_sse() {
|
||||
let report_context = json!({
|
||||
"provider_api_format": "openai:image",
|
||||
"client_api_format": "openai:image",
|
||||
"image_request": {
|
||||
"operation": "edit"
|
||||
}
|
||||
});
|
||||
let outcome = maybe_bridge_standard_sync_json_to_stream(
|
||||
&json!({
|
||||
"created": 1776971267,
|
||||
"data": [{
|
||||
"url": "data:image/webp;base64,d29ybGQ="
|
||||
}],
|
||||
"usage": {
|
||||
"total_tokens": 9,
|
||||
"input_tokens": 4,
|
||||
"output_tokens": 5
|
||||
}
|
||||
}),
|
||||
"openai:image",
|
||||
"openai:image",
|
||||
Some(&report_context),
|
||||
)
|
||||
.expect("bridge should succeed")
|
||||
.expect("bridge should produce sse");
|
||||
|
||||
let output = utf8(outcome.sse_body);
|
||||
assert!(output.contains("event: image_edit.completed"));
|
||||
assert!(output.contains("\"type\":\"image_edit.completed\""));
|
||||
assert!(output.contains("\"b64_json\":\"d29ybGQ=\""));
|
||||
assert!(output.contains("\"total_tokens\":9"));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user