use agent_base::{ChatMessage, LlmClient, OpenAiClient, StreamChunk};
use futures_util::StreamExt;
use serde_json::{Value, json};
#[tokio::test]
async fn test_openai_client_request_body_format() {
let client = OpenAiClient::new(
"test-api-key".to_string(),
"qwen-flash".to_string(),
Some("https://test.example.com/v1".to_string()),
);
let messages = vec![
ChatMessage::system("你是一个助手"),
ChatMessage::user("你好"),
];
let tools: Vec<Value> = vec![];
let reasoning = agent_base::ReasoningConfig {
enabled: Some(true),
budget_tokens: Some(128),
effort: None,
};
let result = client
.chat_stream(&messages, &tools, Some(&reasoning), None)
.await;
assert!(result.is_err(), "Expected error due to invalid URL");
}
#[test]
fn test_request_body_json_structure() {
let mut request_body = json!({
"model": "qwen-flash",
"messages": [
{
"role": "system",
"content": "你是一个助手"
},
{
"role": "user",
"content": "你好"
}
],
"tools": [],
"stream": true,
"stream_options": {
"include_usage": true
}
});
if let Some(obj) = request_body.as_object_mut() {
obj.insert("enable_thinking".to_string(), json!(true));
}
let mut extra_body = serde_json::Map::new();
extra_body.insert("thinking_budget".to_string(), json!(128));
if let Some(obj) = request_body.as_object_mut() {
obj.insert("extra_body".to_string(), Value::Object(extra_body));
}
let json_str = serde_json::to_string_pretty(&request_body).unwrap();
println!("Expected request body:\n{}", json_str);
assert_eq!(request_body["enable_thinking"], true);
assert!(request_body["extra_body"].is_object());
assert_eq!(request_body["extra_body"]["thinking_budget"], 128);
}
#[test]
fn test_thinking_budget_values() {
let budgets = vec![100, 128, 256, 512, 1024, 2048, 4096];
for budget in budgets {
let mut request_body = json!({
"model": "qwen-flash",
"messages": [],
"tools": [],
"stream": true,
});
let mut extra_body = serde_json::Map::new();
extra_body.insert("thinking_budget".to_string(), json!(budget));
if let Some(obj) = request_body.as_object_mut() {
obj.insert("extra_body".to_string(), Value::Object(extra_body));
}
assert_eq!(
request_body["extra_body"]["thinking_budget"], budget,
"thinking_budget should be {}",
budget
);
}
}