466 lines
16 KiB
Rust
466 lines
16 KiB
Rust
//! Postprocessing FFI functions for gRPC stream chunks
|
|
//!
|
|
//! This module provides C-compatible functions for postprocessing gRPC stream chunks:
|
|
//! - Parse tool calls from model output
|
|
//! - Convert proto format to OpenAI format
|
|
//! - Handle reasoning content parsing
|
|
//!
|
|
//! These functions are designed to be called for each stream chunk, but can be optimized
|
|
//! with batching in the future.
|
|
|
|
use std::ffi::{CStr, CString};
|
|
use std::os::raw::{c_char, c_int};
|
|
use std::ptr;
|
|
use std::sync::Arc;
|
|
use serde_json::Value;
|
|
|
|
use smg_grpc_client::sglang_proto as proto;
|
|
|
|
use super::error::{SglErrorCode, set_error_message};
|
|
use super::grpc_converter::GrpcResponseConverterHandle;
|
|
|
|
use tokio::runtime::Runtime;
|
|
use once_cell::sync::Lazy;
|
|
|
|
/// Global tokio runtime for async operations
|
|
static RUNTIME: Lazy<Runtime> = Lazy::new(|| {
|
|
Runtime::new().expect("Failed to create tokio runtime for postprocessor FFI")
|
|
});
|
|
|
|
/// Postprocess a gRPC stream chunk to OpenAI format
|
|
///
|
|
/// This function:
|
|
/// 1. Parses the proto chunk from JSON
|
|
/// 2. Converts it to OpenAI format using the converter handle
|
|
/// 3. Returns the OpenAI format JSON
|
|
///
|
|
/// # Arguments
|
|
/// * `converter_handle` - Converter handle (created with sgl_grpc_response_converter_create)
|
|
/// * `proto_chunk_json` - JSON string of proto.GenerateResponse
|
|
/// * `openai_json_out` - Pointer to receive OpenAI format JSON (must be freed with sgl_free_string)
|
|
/// * `is_done_out` - Pointer to receive is_done flag (1 if stream is complete, 0 otherwise)
|
|
/// * `error_out` - Optional pointer to receive error message
|
|
///
|
|
/// # Returns
|
|
/// * SglErrorCode::Success on success, error code on failure
|
|
#[no_mangle]
|
|
pub unsafe extern "C" fn sgl_postprocess_stream_chunk(
|
|
converter_handle: *mut GrpcResponseConverterHandle,
|
|
proto_chunk_json: *const c_char,
|
|
openai_json_out: *mut *mut c_char,
|
|
is_done_out: *mut c_int,
|
|
error_out: *mut *mut c_char,
|
|
) -> SglErrorCode {
|
|
if converter_handle.is_null()
|
|
|| proto_chunk_json.is_null()
|
|
|| openai_json_out.is_null()
|
|
|| is_done_out.is_null()
|
|
{
|
|
set_error_message(error_out, "Invalid arguments: null pointer");
|
|
return SglErrorCode::InvalidArgument;
|
|
}
|
|
|
|
let proto_chunk_str = match CStr::from_ptr(proto_chunk_json).to_str() {
|
|
Ok(s) => s,
|
|
Err(_) => {
|
|
set_error_message(error_out, "Invalid UTF-8 in proto_chunk_json");
|
|
return SglErrorCode::InvalidArgument;
|
|
}
|
|
};
|
|
|
|
// Parse proto.GenerateResponse from JSON
|
|
let json_value: Value = match serde_json::from_str(proto_chunk_str) {
|
|
Ok(v) => v,
|
|
Err(e) => {
|
|
set_error_message(error_out, &format!("Failed to parse proto chunk JSON: {}", e));
|
|
return SglErrorCode::ParsingError;
|
|
}
|
|
};
|
|
|
|
// Build proto::GenerateResponse from JSON value
|
|
let mut proto_response = proto::GenerateResponse {
|
|
request_id: json_value
|
|
.get("request_id")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string(),
|
|
response: None,
|
|
};
|
|
|
|
// Parse the response oneof field
|
|
let is_done = if let Some(chunk_json) = json_value.get("chunk") {
|
|
let chunk = proto::GenerateStreamChunk {
|
|
token_ids: chunk_json
|
|
.get("token_ids")
|
|
.and_then(|v| v.as_array())
|
|
.map(|arr| {
|
|
arr.iter()
|
|
.filter_map(|v| v.as_u64().map(|n| n as u32))
|
|
.collect()
|
|
})
|
|
.unwrap_or_default(),
|
|
prompt_tokens: chunk_json
|
|
.get("prompt_tokens")
|
|
.and_then(|v| v.as_i64())
|
|
.map(|n| n as i32)
|
|
.unwrap_or(0),
|
|
completion_tokens: chunk_json
|
|
.get("completion_tokens")
|
|
.and_then(|v| v.as_i64())
|
|
.map(|n| n as i32)
|
|
.unwrap_or(0),
|
|
cached_tokens: chunk_json
|
|
.get("cached_tokens")
|
|
.and_then(|v| v.as_i64())
|
|
.map(|n| n as i32)
|
|
.unwrap_or(0),
|
|
output_logprobs: None,
|
|
hidden_states: vec![],
|
|
input_logprobs: None,
|
|
index: 0,
|
|
};
|
|
proto_response.response = Some(proto::generate_response::Response::Chunk(chunk));
|
|
false
|
|
} else if let Some(complete_json) = json_value.get("complete") {
|
|
let complete = proto::GenerateComplete {
|
|
output_ids: complete_json
|
|
.get("output_ids")
|
|
.and_then(|v| v.as_array())
|
|
.map(|arr| {
|
|
arr.iter()
|
|
.filter_map(|v| v.as_u64().map(|n| n as u32))
|
|
.collect()
|
|
})
|
|
.unwrap_or_default(),
|
|
finish_reason: complete_json
|
|
.get("finish_reason")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string(),
|
|
prompt_tokens: complete_json
|
|
.get("prompt_tokens")
|
|
.and_then(|v| v.as_i64())
|
|
.map(|n| n as i32)
|
|
.unwrap_or(0),
|
|
completion_tokens: complete_json
|
|
.get("completion_tokens")
|
|
.and_then(|v| v.as_i64())
|
|
.map(|n| n as i32)
|
|
.unwrap_or(0),
|
|
cached_tokens: complete_json
|
|
.get("cached_tokens")
|
|
.and_then(|v| v.as_i64())
|
|
.map(|n| n as i32)
|
|
.unwrap_or(0),
|
|
output_logprobs: None,
|
|
all_hidden_states: vec![],
|
|
input_logprobs: None,
|
|
matched_stop: None,
|
|
index: 0,
|
|
};
|
|
proto_response.response = Some(proto::generate_response::Response::Complete(complete));
|
|
true
|
|
} else if let Some(error_json) = json_value.get("error") {
|
|
let error = proto::GenerateError {
|
|
message: error_json
|
|
.get("message")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string(),
|
|
http_status_code: error_json
|
|
.get("http_status_code")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("500")
|
|
.to_string(),
|
|
details: error_json
|
|
.get("details")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string(),
|
|
};
|
|
proto_response.response = Some(proto::generate_response::Response::Error(error));
|
|
true
|
|
} else {
|
|
set_error_message(
|
|
error_out,
|
|
"Proto chunk JSON must contain 'chunk', 'complete', or 'error' field",
|
|
);
|
|
return SglErrorCode::ParsingError;
|
|
};
|
|
|
|
// Convert proto chunk to OpenAI format using the converter's convert_chunk function
|
|
// We'll use the existing converter API instead of calling the internal function directly
|
|
let proto_chunk_json_cstr = match CString::new(proto_chunk_str) {
|
|
Ok(s) => s,
|
|
Err(e) => {
|
|
set_error_message(error_out, &format!("Failed to create C string: {}", e));
|
|
return SglErrorCode::MemoryError;
|
|
}
|
|
};
|
|
|
|
// Use the existing converter API
|
|
let mut openai_json_ptr: *mut c_char = ptr::null_mut();
|
|
let result = super::grpc_converter::sgl_grpc_response_converter_convert_chunk(
|
|
converter_handle,
|
|
proto_chunk_json_cstr.as_ptr(),
|
|
&mut openai_json_ptr,
|
|
error_out,
|
|
);
|
|
|
|
if result == SglErrorCode::Success {
|
|
*openai_json_out = openai_json_ptr;
|
|
*is_done_out = if is_done { 1 } else { 0 };
|
|
SglErrorCode::Success
|
|
} else {
|
|
*openai_json_out = ptr::null_mut();
|
|
*is_done_out = if is_done { 1 } else { 0 };
|
|
result
|
|
}
|
|
}
|
|
|
|
/// Postprocess multiple gRPC stream chunks in batch (reduces FFI overhead)
|
|
///
|
|
/// This function processes multiple chunks in a single FFI call, significantly reducing
|
|
/// FFI overhead in streaming scenarios.
|
|
///
|
|
/// # Arguments
|
|
/// * `converter_handle` - Converter handle (created with sgl_grpc_response_converter_create)
|
|
/// * `proto_chunks_json_array` - JSON array string of proto.GenerateResponse chunks
|
|
/// * `max_chunks` - Maximum number of chunks to process (for safety)
|
|
/// * `openai_chunks_json_array_out` - Pointer to receive JSON array of OpenAI format chunks (must be freed with sgl_free_string)
|
|
/// * `chunks_count_out` - Pointer to receive number of processed chunks
|
|
/// * `error_out` - Optional pointer to receive error message
|
|
///
|
|
/// # Returns
|
|
/// * SglErrorCode::Success on success, error code on failure
|
|
#[no_mangle]
|
|
pub unsafe extern "C" fn sgl_postprocess_stream_chunks_batch(
|
|
converter_handle: *mut GrpcResponseConverterHandle,
|
|
proto_chunks_json_array: *const c_char,
|
|
max_chunks: c_int,
|
|
openai_chunks_json_array_out: *mut *mut c_char,
|
|
chunks_count_out: *mut c_int,
|
|
error_out: *mut *mut c_char,
|
|
) -> SglErrorCode {
|
|
if converter_handle.is_null()
|
|
|| proto_chunks_json_array.is_null()
|
|
|| openai_chunks_json_array_out.is_null()
|
|
|| chunks_count_out.is_null()
|
|
{
|
|
set_error_message(error_out, "Invalid arguments: null pointer");
|
|
return SglErrorCode::InvalidArgument;
|
|
}
|
|
|
|
let chunks_array_str = match CStr::from_ptr(proto_chunks_json_array).to_str() {
|
|
Ok(s) => s,
|
|
Err(_) => {
|
|
set_error_message(error_out, "Invalid UTF-8 in proto_chunks_json_array");
|
|
return SglErrorCode::InvalidArgument;
|
|
}
|
|
};
|
|
|
|
// Parse JSON array of chunks
|
|
let chunks_array: Vec<Value> = match serde_json::from_str(chunks_array_str) {
|
|
Ok(arr) => arr,
|
|
Err(e) => {
|
|
set_error_message(
|
|
error_out,
|
|
&format!("Failed to parse chunks JSON array: {}", e),
|
|
);
|
|
return SglErrorCode::ParsingError;
|
|
}
|
|
};
|
|
|
|
// Limit batch size for safety
|
|
let max_chunks_usize = max_chunks as usize;
|
|
let chunks_to_process = if chunks_array.len() > max_chunks_usize {
|
|
&chunks_array[..max_chunks_usize]
|
|
} else {
|
|
&chunks_array
|
|
};
|
|
|
|
let handle_ref = &mut *converter_handle;
|
|
let tokenizer = Arc::clone(&handle_ref.tokenizer);
|
|
let model = handle_ref.model.clone();
|
|
let request_id = handle_ref.request_id.clone();
|
|
let created = handle_ref.created;
|
|
let system_fingerprint = handle_ref.system_fingerprint.clone();
|
|
|
|
// Process chunks in batch
|
|
let mut results = Vec::new();
|
|
let mut has_error = false;
|
|
let mut error_msg = String::new();
|
|
|
|
for chunk_json in chunks_to_process {
|
|
// Parse proto.GenerateResponse from JSON
|
|
let mut proto_response = proto::GenerateResponse {
|
|
request_id: chunk_json
|
|
.get("request_id")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string(),
|
|
response: None,
|
|
};
|
|
|
|
// Parse the response oneof field (same logic as single chunk processing)
|
|
let _is_done = if let Some(chunk_json) = chunk_json.get("chunk") {
|
|
let chunk = proto::GenerateStreamChunk {
|
|
token_ids: chunk_json
|
|
.get("token_ids")
|
|
.and_then(|v| v.as_array())
|
|
.map(|arr| {
|
|
arr.iter()
|
|
.filter_map(|v| v.as_u64().map(|n| n as u32))
|
|
.collect()
|
|
})
|
|
.unwrap_or_default(),
|
|
prompt_tokens: chunk_json
|
|
.get("prompt_tokens")
|
|
.and_then(|v| v.as_i64())
|
|
.map(|n| n as i32)
|
|
.unwrap_or(0),
|
|
completion_tokens: chunk_json
|
|
.get("completion_tokens")
|
|
.and_then(|v| v.as_i64())
|
|
.map(|n| n as i32)
|
|
.unwrap_or(0),
|
|
cached_tokens: chunk_json
|
|
.get("cached_tokens")
|
|
.and_then(|v| v.as_i64())
|
|
.map(|n| n as i32)
|
|
.unwrap_or(0),
|
|
output_logprobs: None,
|
|
hidden_states: vec![],
|
|
input_logprobs: None,
|
|
index: 0,
|
|
};
|
|
proto_response.response = Some(proto::generate_response::Response::Chunk(chunk));
|
|
false
|
|
} else if let Some(complete_json) = chunk_json.get("complete") {
|
|
let complete = proto::GenerateComplete {
|
|
output_ids: complete_json
|
|
.get("output_ids")
|
|
.and_then(|v| v.as_array())
|
|
.map(|arr| {
|
|
arr.iter()
|
|
.filter_map(|v| v.as_u64().map(|n| n as u32))
|
|
.collect()
|
|
})
|
|
.unwrap_or_default(),
|
|
finish_reason: complete_json
|
|
.get("finish_reason")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string(),
|
|
prompt_tokens: complete_json
|
|
.get("prompt_tokens")
|
|
.and_then(|v| v.as_i64())
|
|
.map(|n| n as i32)
|
|
.unwrap_or(0),
|
|
completion_tokens: complete_json
|
|
.get("completion_tokens")
|
|
.and_then(|v| v.as_i64())
|
|
.map(|n| n as i32)
|
|
.unwrap_or(0),
|
|
cached_tokens: complete_json
|
|
.get("cached_tokens")
|
|
.and_then(|v| v.as_i64())
|
|
.map(|n| n as i32)
|
|
.unwrap_or(0),
|
|
output_logprobs: None,
|
|
all_hidden_states: vec![],
|
|
input_logprobs: None,
|
|
matched_stop: None,
|
|
index: 0,
|
|
};
|
|
proto_response.response = Some(proto::generate_response::Response::Complete(complete));
|
|
true
|
|
} else if let Some(error_json) = chunk_json.get("error") {
|
|
let error = proto::GenerateError {
|
|
message: error_json
|
|
.get("message")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string(),
|
|
http_status_code: error_json
|
|
.get("http_status_code")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("500")
|
|
.to_string(),
|
|
details: error_json
|
|
.get("details")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string(),
|
|
};
|
|
proto_response.response = Some(proto::generate_response::Response::Error(error));
|
|
true
|
|
} else {
|
|
error_msg = format!(
|
|
"Chunk JSON must contain 'chunk', 'complete', or 'error' field: {}",
|
|
chunk_json
|
|
);
|
|
has_error = true;
|
|
break;
|
|
};
|
|
|
|
// Convert proto chunk to OpenAI format
|
|
let result = RUNTIME.block_on(async {
|
|
super::grpc_converter::convert_proto_chunk_to_openai(
|
|
proto_response,
|
|
handle_ref,
|
|
&tokenizer,
|
|
&model,
|
|
&request_id,
|
|
created,
|
|
system_fingerprint.as_deref(),
|
|
)
|
|
.await
|
|
});
|
|
|
|
match result {
|
|
Ok(Some(openai_response)) => {
|
|
results.push(openai_response);
|
|
}
|
|
Ok(None) => {
|
|
// Empty response, skip
|
|
}
|
|
Err(e) => {
|
|
error_msg = format!("Postprocessing failed for chunk: {}", e);
|
|
has_error = true;
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
if has_error {
|
|
set_error_message(error_out, &error_msg);
|
|
return SglErrorCode::ParsingError;
|
|
}
|
|
|
|
// Serialize results to JSON array
|
|
let results_json = match serde_json::to_string(&results) {
|
|
Ok(s) => s,
|
|
Err(e) => {
|
|
set_error_message(
|
|
error_out,
|
|
&format!("Failed to serialize results JSON array: {}", e),
|
|
);
|
|
return SglErrorCode::ParsingError;
|
|
}
|
|
};
|
|
|
|
let results_cstr = match CString::new(results_json) {
|
|
Ok(s) => s,
|
|
Err(e) => {
|
|
set_error_message(error_out, &format!("Failed to create C string: {}", e));
|
|
return SglErrorCode::MemoryError;
|
|
}
|
|
};
|
|
|
|
*openai_chunks_json_array_out = results_cstr.into_raw();
|
|
*chunks_count_out = results.len() as c_int;
|
|
|
|
SglErrorCode::Success
|
|
}
|