373 lines
13 KiB
Rust
373 lines
13 KiB
Rust
//! Preprocessing FFI functions for chat requests
|
|
//!
|
|
//! This module provides C-compatible functions for preprocessing chat completion requests:
|
|
//! - Apply chat_template to messages
|
|
//! - Tokenize the processed text
|
|
//! - Generate tool constraints
|
|
//!
|
|
//! These functions are designed to be called once per request, reducing FFI overhead.
|
|
|
|
use std::ffi::{CStr, CString};
|
|
use std::os::raw::{c_char, c_int};
|
|
use std::ptr;
|
|
use std::os::raw::c_uint;
|
|
|
|
use smg::tokenizer::create_tokenizer_from_file;
|
|
use smg::protocols::chat::ChatCompletionRequest;
|
|
use smg::routers::grpc::utils::{process_chat_messages, generate_tool_constraints};
|
|
|
|
use super::error::{SglErrorCode, set_error_message};
|
|
use super::memory::{sgl_free_string, sgl_free_token_ids};
|
|
use super::tokenizer::TokenizerHandle;
|
|
|
|
/// Handle for preprocessed request
|
|
#[repr(C)]
|
|
pub struct PreprocessedRequestHandle {
|
|
pub(crate) prompt_text: CString,
|
|
pub(crate) token_ids: Vec<i32>,
|
|
pub(crate) tool_constraints_json: Option<CString>,
|
|
pub(crate) prompt_tokens: i32,
|
|
}
|
|
|
|
/// Preprocess a chat completion request
|
|
///
|
|
/// This function:
|
|
/// 1. Applies chat_template to messages
|
|
/// 2. Tokenizes the processed text
|
|
/// 3. Generates tool constraints (if tools are present)
|
|
///
|
|
/// # Arguments
|
|
/// * `request_json` - OpenAI ChatCompletionRequest as JSON string
|
|
/// * `tokenizer_path` - Path to tokenizer directory
|
|
/// * `prompt_text_out` - Pointer to receive prompt text (C string, must be freed with sgl_free_string)
|
|
/// * `token_ids_out` - Pointer to receive token IDs array (must be freed with sgl_free_token_ids)
|
|
/// * `token_ids_len_out` - Pointer to receive token IDs array length
|
|
/// * `tool_constraints_json_out` - Optional pointer to receive tool constraints JSON (must be freed with sgl_free_string)
|
|
/// * `prompt_tokens_out` - Pointer to receive prompt token count
|
|
/// * `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_preprocess_chat_request(
|
|
request_json: *const c_char,
|
|
tokenizer_path: *const c_char,
|
|
prompt_text_out: *mut *mut c_char,
|
|
token_ids_out: *mut *mut c_uint,
|
|
token_ids_len_out: *mut usize,
|
|
tool_constraints_json_out: *mut *mut c_char,
|
|
prompt_tokens_out: *mut c_int,
|
|
error_out: *mut *mut c_char,
|
|
) -> SglErrorCode {
|
|
if request_json.is_null()
|
|
|| tokenizer_path.is_null()
|
|
|| prompt_text_out.is_null()
|
|
|| token_ids_out.is_null()
|
|
|| token_ids_len_out.is_null()
|
|
|| prompt_tokens_out.is_null()
|
|
{
|
|
set_error_message(error_out, "Invalid arguments: null pointer");
|
|
return SglErrorCode::InvalidArgument;
|
|
}
|
|
|
|
// Parse input strings
|
|
let request_str = match CStr::from_ptr(request_json).to_str() {
|
|
Ok(s) => s,
|
|
Err(_) => {
|
|
set_error_message(error_out, "Invalid UTF-8 in request_json");
|
|
return SglErrorCode::InvalidArgument;
|
|
}
|
|
};
|
|
|
|
let tokenizer_path_str = match CStr::from_ptr(tokenizer_path).to_str() {
|
|
Ok(s) => s,
|
|
Err(_) => {
|
|
set_error_message(error_out, "Invalid UTF-8 in tokenizer_path");
|
|
return SglErrorCode::InvalidArgument;
|
|
}
|
|
};
|
|
|
|
// Parse ChatCompletionRequest
|
|
let chat_request: ChatCompletionRequest = match serde_json::from_str(request_str) {
|
|
Ok(req) => req,
|
|
Err(e) => {
|
|
set_error_message(error_out, &format!("Failed to parse request JSON: {}", e));
|
|
return SglErrorCode::ParsingError;
|
|
}
|
|
};
|
|
|
|
// Create tokenizer
|
|
let tokenizer = match create_tokenizer_from_file(tokenizer_path_str) {
|
|
Ok(t) => t,
|
|
Err(e) => {
|
|
set_error_message(error_out, &format!("Failed to create tokenizer: {}", e));
|
|
return SglErrorCode::TokenizationError;
|
|
}
|
|
};
|
|
|
|
// Process chat messages (apply chat_template)
|
|
let processed_messages = match process_chat_messages(&chat_request, tokenizer.as_ref()) {
|
|
Ok(msgs) => msgs,
|
|
Err(e) => {
|
|
set_error_message(error_out, &format!("Failed to process chat messages: {}", e));
|
|
return SglErrorCode::ParsingError;
|
|
}
|
|
};
|
|
|
|
// Tokenize the processed text
|
|
let encoding = match tokenizer.encode(&processed_messages.text, false) {
|
|
Ok(enc) => enc,
|
|
Err(e) => {
|
|
set_error_message(error_out, &format!("Tokenization failed: {}", e));
|
|
return SglErrorCode::TokenizationError;
|
|
}
|
|
};
|
|
|
|
let token_ids_vec: Vec<i32> = encoding
|
|
.token_ids()
|
|
.iter()
|
|
.map(|&id| id as i32)
|
|
.collect();
|
|
|
|
let prompt_tokens = token_ids_vec.len() as i32;
|
|
|
|
// Generate tool constraints if tools are present
|
|
let tool_constraints_json = if let Some(tools) = chat_request.tools.as_ref() {
|
|
match generate_tool_constraints(tools, &chat_request.tool_choice, &chat_request.model) {
|
|
Ok(Some(constraints)) => {
|
|
match serde_json::to_string(&constraints) {
|
|
Ok(json_str) => Some(CString::new(json_str).unwrap()),
|
|
Err(e) => {
|
|
set_error_message(
|
|
error_out,
|
|
&format!("Failed to serialize tool constraints: {}", e),
|
|
);
|
|
return SglErrorCode::ParsingError;
|
|
}
|
|
}
|
|
}
|
|
Ok(None) => None,
|
|
Err(e) => {
|
|
set_error_message(
|
|
error_out,
|
|
&format!("Failed to generate tool constraints: {}", e),
|
|
);
|
|
return SglErrorCode::ParsingError;
|
|
}
|
|
}
|
|
} else {
|
|
None
|
|
};
|
|
|
|
// Allocate memory for outputs
|
|
let prompt_text_cstr = match CString::new(processed_messages.text) {
|
|
Ok(s) => s,
|
|
Err(e) => {
|
|
set_error_message(error_out, &format!("Failed to create C string: {}", e));
|
|
return SglErrorCode::MemoryError;
|
|
}
|
|
};
|
|
|
|
let token_ids_len = token_ids_vec.len();
|
|
// Convert i32 to u32 for token IDs (as expected by the memory management functions)
|
|
let token_ids_u32: Vec<u32> = token_ids_vec.iter().map(|&id| id as u32).collect();
|
|
let token_ids_ptr = if token_ids_u32.is_empty() {
|
|
ptr::null_mut()
|
|
} else {
|
|
let boxed = token_ids_u32.into_boxed_slice();
|
|
Box::into_raw(boxed) as *mut c_uint
|
|
};
|
|
|
|
// Set output values
|
|
*prompt_text_out = prompt_text_cstr.into_raw();
|
|
*token_ids_out = token_ids_ptr;
|
|
*token_ids_len_out = token_ids_len;
|
|
*prompt_tokens_out = prompt_tokens;
|
|
|
|
if !tool_constraints_json_out.is_null() {
|
|
if let Some(constraints) = tool_constraints_json {
|
|
*tool_constraints_json_out = constraints.into_raw();
|
|
} else {
|
|
*tool_constraints_json_out = ptr::null_mut();
|
|
}
|
|
}
|
|
|
|
SglErrorCode::Success
|
|
}
|
|
|
|
/// Preprocess a chat completion request using an existing tokenizer handle
|
|
///
|
|
/// This function is similar to sgl_preprocess_chat_request, but accepts a TokenizerHandle
|
|
/// instead of creating a new tokenizer. This allows reusing a cached tokenizer instance,
|
|
/// significantly reducing initialization overhead in concurrent scenarios.
|
|
///
|
|
/// # Arguments
|
|
/// * `request_json` - OpenAI ChatCompletionRequest as JSON string
|
|
/// * `tokenizer_handle` - Existing tokenizer handle (must be valid)
|
|
/// * `prompt_text_out` - Pointer to receive prompt text (C string, must be freed with sgl_free_string)
|
|
/// * `token_ids_out` - Pointer to receive token IDs array (must be freed with sgl_free_token_ids)
|
|
/// * `token_ids_len_out` - Pointer to receive token IDs array length
|
|
/// * `tool_constraints_json_out` - Optional pointer to receive tool constraints JSON (must be freed with sgl_free_string)
|
|
/// * `prompt_tokens_out` - Pointer to receive prompt token count
|
|
/// * `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_preprocess_chat_request_with_tokenizer(
|
|
request_json: *const c_char,
|
|
tokenizer_handle: *mut TokenizerHandle,
|
|
prompt_text_out: *mut *mut c_char,
|
|
token_ids_out: *mut *mut c_uint,
|
|
token_ids_len_out: *mut usize,
|
|
tool_constraints_json_out: *mut *mut c_char,
|
|
prompt_tokens_out: *mut c_int,
|
|
error_out: *mut *mut c_char,
|
|
) -> SglErrorCode {
|
|
if request_json.is_null()
|
|
|| tokenizer_handle.is_null()
|
|
|| prompt_text_out.is_null()
|
|
|| token_ids_out.is_null()
|
|
|| token_ids_len_out.is_null()
|
|
|| prompt_tokens_out.is_null()
|
|
{
|
|
set_error_message(error_out, "Invalid arguments: null pointer");
|
|
return SglErrorCode::InvalidArgument;
|
|
}
|
|
|
|
// Parse input string
|
|
let request_str = match CStr::from_ptr(request_json).to_str() {
|
|
Ok(s) => s,
|
|
Err(_) => {
|
|
set_error_message(error_out, "Invalid UTF-8 in request_json");
|
|
return SglErrorCode::InvalidArgument;
|
|
}
|
|
};
|
|
|
|
// Parse ChatCompletionRequest
|
|
let chat_request: ChatCompletionRequest = match serde_json::from_str(request_str) {
|
|
Ok(req) => req,
|
|
Err(e) => {
|
|
set_error_message(error_out, &format!("Failed to parse request JSON: {}", e));
|
|
return SglErrorCode::ParsingError;
|
|
}
|
|
};
|
|
|
|
// Use existing tokenizer from handle (no need to create new one!)
|
|
let handle_ref = &*tokenizer_handle;
|
|
let tokenizer = &handle_ref.tokenizer;
|
|
|
|
// Process chat messages (apply chat_template)
|
|
let processed_messages = match process_chat_messages(&chat_request, tokenizer.as_ref()) {
|
|
Ok(msgs) => msgs,
|
|
Err(e) => {
|
|
set_error_message(error_out, &format!("Failed to process chat messages: {}", e));
|
|
return SglErrorCode::ParsingError;
|
|
}
|
|
};
|
|
|
|
// Tokenize the processed text
|
|
let encoding = match tokenizer.encode(&processed_messages.text, false) {
|
|
Ok(enc) => enc,
|
|
Err(e) => {
|
|
set_error_message(error_out, &format!("Tokenization failed: {}", e));
|
|
return SglErrorCode::TokenizationError;
|
|
}
|
|
};
|
|
|
|
let token_ids_vec: Vec<i32> = encoding
|
|
.token_ids()
|
|
.iter()
|
|
.map(|&id| id as i32)
|
|
.collect();
|
|
|
|
let prompt_tokens = token_ids_vec.len() as i32;
|
|
|
|
// Generate tool constraints if tools are present
|
|
let tool_constraints_json = if let Some(tools) = chat_request.tools.as_ref() {
|
|
match generate_tool_constraints(tools, &chat_request.tool_choice, &chat_request.model) {
|
|
Ok(Some(constraints)) => {
|
|
match serde_json::to_string(&constraints) {
|
|
Ok(json_str) => Some(CString::new(json_str).unwrap()),
|
|
Err(e) => {
|
|
set_error_message(
|
|
error_out,
|
|
&format!("Failed to serialize tool constraints: {}", e),
|
|
);
|
|
return SglErrorCode::ParsingError;
|
|
}
|
|
}
|
|
}
|
|
Ok(None) => None,
|
|
Err(e) => {
|
|
set_error_message(
|
|
error_out,
|
|
&format!("Failed to generate tool constraints: {}", e),
|
|
);
|
|
return SglErrorCode::ParsingError;
|
|
}
|
|
}
|
|
} else {
|
|
None
|
|
};
|
|
|
|
// Allocate memory for outputs
|
|
let prompt_text_cstr = match CString::new(processed_messages.text) {
|
|
Ok(s) => s,
|
|
Err(e) => {
|
|
set_error_message(error_out, &format!("Failed to create C string: {}", e));
|
|
return SglErrorCode::MemoryError;
|
|
}
|
|
};
|
|
|
|
let token_ids_len = token_ids_vec.len();
|
|
// Convert i32 to u32 for token IDs (as expected by the memory management functions)
|
|
let token_ids_u32: Vec<u32> = token_ids_vec.iter().map(|&id| id as u32).collect();
|
|
let token_ids_ptr = if token_ids_u32.is_empty() {
|
|
ptr::null_mut()
|
|
} else {
|
|
let boxed = token_ids_u32.into_boxed_slice();
|
|
Box::into_raw(boxed) as *mut c_uint
|
|
};
|
|
|
|
// Set output values
|
|
*prompt_text_out = prompt_text_cstr.into_raw();
|
|
*token_ids_out = token_ids_ptr;
|
|
*token_ids_len_out = token_ids_len;
|
|
*prompt_tokens_out = prompt_tokens;
|
|
|
|
if !tool_constraints_json_out.is_null() {
|
|
if let Some(constraints) = tool_constraints_json {
|
|
*tool_constraints_json_out = constraints.into_raw();
|
|
} else {
|
|
*tool_constraints_json_out = ptr::null_mut();
|
|
}
|
|
}
|
|
|
|
SglErrorCode::Success
|
|
}
|
|
|
|
/// Free a preprocessed request handle (cleanup function)
|
|
///
|
|
/// This function frees the memory allocated by sgl_preprocess_chat_request.
|
|
/// It should be called after the preprocessed data is no longer needed.
|
|
#[no_mangle]
|
|
pub unsafe extern "C" fn sgl_preprocessed_request_free(
|
|
prompt_text: *mut c_char,
|
|
token_ids: *mut c_uint,
|
|
token_ids_len: usize,
|
|
tool_constraints_json: *mut c_char,
|
|
) {
|
|
if !prompt_text.is_null() {
|
|
sgl_free_string(prompt_text);
|
|
}
|
|
|
|
if !token_ids.is_null() && token_ids_len > 0 {
|
|
sgl_free_token_ids(token_ids, token_ids_len);
|
|
}
|
|
|
|
if !tool_constraints_json.is_null() {
|
|
sgl_free_string(tool_constraints_json);
|
|
}
|
|
}
|