refactor generate code

This commit is contained in:
jhqxxx
2026-04-02 22:22:52 +08:00
parent b254d21efc
commit bd8eee6520
59 changed files with 1745 additions and 1800 deletions
+3 -307
View File
@@ -1,6 +1,7 @@
pub mod audio_utils;
pub mod img_utils;
pub mod interpolate;
pub mod response_utils;
pub mod tensor_utils;
pub mod video_utils;
@@ -10,14 +11,8 @@ use std::time::{SystemTime, UNIX_EPOCH};
use std::{collections::HashMap, fs, path::PathBuf, process::Command, time::Duration};
use crate::models::common::model_mapping::WhichModel;
use crate::params::{
chat::{
AudioUrlType, ChatCompletionChoice, ChatCompletionChunkChoice, ChatCompletionChunkResponse,
ChatCompletionParameters, ChatCompletionResponse, ChatMessage, ChatMessageAudioContentPart,
ChatMessageContent, ChatMessageContentPart, ChatMessageImageContentPart, DeltaChatMessage,
DeltaFunction, DeltaToolCall, Function, ImageUrlType, ToolCall,
},
shared::{FinishReason, Usage},
use crate::params::chat::{
ChatCompletionParameters, ChatMessage, ChatMessageContent, ChatMessageContentPart,
};
use anyhow::{Result, anyhow};
use byteorder::{LittleEndian, ReadBytesExt};
@@ -409,305 +404,6 @@ pub fn ceil_by_factor(num: f32, factor: u32) -> u32 {
ceil * factor
}
pub fn build_img_completion_response(
base64vec: &Vec<String>,
model_name: &str,
) -> ChatCompletionResponse {
let id = uuid::Uuid::new_v4().to_string();
let mut response = ChatCompletionResponse {
id: Some(id),
choices: vec![],
// created: chrono::Utc::now().timestamp() as u32,
created: timestamp() as u32,
model: model_name.to_string(),
service_tier: None,
system_fingerprint: None,
object: "chat.completion".to_string(),
usage: None,
};
let mut conten_part_vec = vec![];
for img_bas64 in base64vec {
let img_base64_prefix = "data:image/png;base64,".to_string() + img_bas64;
let part = ChatMessageContentPart::Image(ChatMessageImageContentPart {
r#type: "image".to_string(),
image_url: ImageUrlType {
url: img_base64_prefix,
detail: None,
},
});
conten_part_vec.push(part);
}
let choice = ChatCompletionChoice {
index: 0,
message: ChatMessage::Assistant {
content: Some(ChatMessageContent::ContentPart(conten_part_vec)),
reasoning_content: None,
refusal: None,
name: None,
audio: None,
tool_calls: None,
},
finish_reason: Some(FinishReason::StopSequenceReached),
logprobs: None,
};
response.choices.push(choice);
response
}
pub fn build_audio_completion_response(
base64_audio: &String,
model_name: &str,
) -> ChatCompletionResponse {
let id = uuid::Uuid::new_v4().to_string();
let mut response = ChatCompletionResponse {
id: Some(id),
choices: vec![],
created: timestamp() as u32,
model: model_name.to_string(),
service_tier: None,
system_fingerprint: None,
object: "chat.completion".to_string(),
usage: None,
};
let base64_audio = format!("data:audio/wav;base64,{}", base64_audio);
let conten_part_vec = vec![ChatMessageContentPart::Audio(ChatMessageAudioContentPart {
r#type: "audio".to_string(),
audio_url: AudioUrlType {
url: base64_audio.to_string(),
},
})];
let choice = ChatCompletionChoice {
index: 0,
message: ChatMessage::Assistant {
content: Some(ChatMessageContent::ContentPart(conten_part_vec)),
reasoning_content: None,
refusal: None,
name: None,
audio: None,
tool_calls: None,
},
finish_reason: Some(FinishReason::StopSequenceReached),
logprobs: None,
};
response.choices.push(choice);
response
}
fn build_response(res: String, model_name: &str, usage: Option<Usage>) -> ChatCompletionResponse {
let id = uuid::Uuid::new_v4().to_string();
let mut response = ChatCompletionResponse {
id: Some(id),
choices: vec![],
created: timestamp() as u32,
model: model_name.to_string(),
service_tier: None,
system_fingerprint: None,
object: "chat.completion".to_string(),
usage,
};
let choice = if res.contains("<tool_call>") {
let mes: Vec<&str> = res.split("<tool_call>").collect();
let content = mes[0].to_string();
let mut tool_vec = Vec::new();
for (i, m) in mes.iter().enumerate().skip(1) {
let tool_mes = m.replace("</tool_call>", "");
let function = match serde_json::from_str::<serde_json::Value>(&tool_mes) {
Ok(json_value) => {
let name = json_value
.get("name")
.and_then(|v| v.as_str())
.map(|s| s.to_string())
.unwrap_or_default();
let arguments = json_value
.get("arguments")
.map(|v| v.to_string())
.unwrap_or_default();
Function { name, arguments }
}
Err(_) => Function {
name: "".to_string(),
arguments: "".to_string(),
},
};
let tool_call = ToolCall {
id: (i - 1).to_string(),
r#type: "function".to_string(),
function,
};
tool_vec.push(tool_call);
}
ChatCompletionChoice {
index: 0,
message: ChatMessage::Assistant {
content: Some(ChatMessageContent::Text(content)),
reasoning_content: None,
refusal: None,
name: None,
audio: None,
tool_calls: Some(tool_vec),
},
finish_reason: Some(FinishReason::ToolCalls),
logprobs: None,
}
} else {
ChatCompletionChoice {
index: 0,
message: ChatMessage::Assistant {
content: Some(ChatMessageContent::Text(res)),
reasoning_content: None,
refusal: None,
name: None,
audio: None,
tool_calls: None,
},
finish_reason: Some(FinishReason::StopSequenceReached),
logprobs: None,
}
};
response.choices.push(choice);
response
}
pub fn build_completion_response(
res: String,
model_name: &str,
completion_tokens: Option<u32>,
prompt_tokens: Option<u32>,
) -> ChatCompletionResponse {
let usage = if prompt_tokens.is_none() && completion_tokens.is_none() {
None
} else {
Some(Usage {
prompt_tokens,
prompt_secs: None,
completion_tokens,
completion_secs: None,
completion_per_token_secs: None,
completion_tps: None,
total_tokens: prompt_tokens.unwrap_or(0) + completion_tokens.unwrap_or(0),
prompt_tokens_details: None,
completion_tokens_details: None,
})
};
build_response(res, model_name, usage)
}
pub fn build_completion_response_with_time(
res: String,
model_name: &str,
completion_tokens: Option<u32>,
completion_secs: Option<f64>,
prompt_tokens: Option<u32>,
prompt_secs: Option<f64>,
) -> ChatCompletionResponse {
let usage = if prompt_tokens.is_none() && completion_tokens.is_none() {
None
} else {
let (completion_per_token_secs, completion_tps) = if let Some(completion_tokens) =
completion_tokens
&& let Some(completion_secs) = completion_secs
{
let per_token_secs = completion_secs / completion_tokens as f64;
let tps = completion_tokens as f64 / completion_secs;
(Some(per_token_secs), Some(tps))
} else {
(None, None)
};
Some(Usage {
prompt_tokens,
prompt_secs,
completion_tokens,
completion_secs,
completion_per_token_secs,
completion_tps,
total_tokens: prompt_tokens.unwrap_or(0) + completion_tokens.unwrap_or(0),
prompt_tokens_details: None,
completion_tokens_details: None,
})
};
build_response(res, model_name, usage)
}
pub fn build_completion_chunk_response(
res: String,
model_name: &str,
tool_call_id: Option<String>,
tool_call_content: Option<String>,
) -> ChatCompletionChunkResponse {
let id = uuid::Uuid::new_v4().to_string();
let mut response = ChatCompletionChunkResponse {
id: Some(id),
choices: vec![],
created: timestamp() as u32,
model: model_name.to_string(),
system_fingerprint: None,
object: "chat.completion.chunk".to_string(),
usage: None,
};
let choice = if let Some(tool_call_id) = tool_call_id {
let function = if let Some(content) = tool_call_content {
match serde_json::from_str::<serde_json::Value>(&content) {
Ok(json_value) => {
let name = json_value
.get("name")
.and_then(|v| v.as_str())
.map(|s| s.to_string());
let arguments = json_value.get("arguments").map(|v| v.to_string());
DeltaFunction { name, arguments }
}
Err(_) => DeltaFunction {
name: None,
arguments: Some(content),
},
}
} else {
DeltaFunction {
name: None,
arguments: None,
}
};
ChatCompletionChunkChoice {
index: Some(0),
delta: DeltaChatMessage::Assistant {
content: None,
reasoning_content: None,
refusal: None,
name: None,
tool_calls: Some(vec![DeltaToolCall {
index: Some(0),
id: Some(tool_call_id),
r#type: Some("function".to_string()),
function,
}]),
},
finish_reason: None,
logprobs: None,
}
} else {
ChatCompletionChunkChoice {
index: Some(0),
delta: DeltaChatMessage::Assistant {
content: Some(ChatMessageContent::Text(res)),
reasoning_content: None,
refusal: None,
name: None,
tool_calls: None,
},
finish_reason: None,
logprobs: None,
}
};
response.choices.push(choice);
response
}
pub fn extract_mes(mes: &ChatCompletionParameters) -> Result<Vec<(String, String)>> {
let mut mes_vec = Vec::new();
for chat_mes in mes.messages.clone() {