| 1 | //! Moonshot/Kimi native search adapters. |
| 2 | |
| 3 | use anyhow::{Context, Result, bail}; |
| 4 | use reqwest::header::{HeaderName, HeaderValue}; |
| 5 | use serde_json::{Map, Value, json}; |
| 6 | use uuid::Uuid; |
| 7 | |
| 8 | use super::{ |
| 9 | ProviderNativeSearchClient, ProviderNativeSearchRequest, ProviderNativeSearchResponse, |
| 10 | citation_from_url, citations_from_text, finish_reason_is_length, joined_answer, push_citation, |
| 11 | }; |
| 12 | use crate::{ |
| 13 | client::api_url, |
| 14 | config::{MOONSHOT_KIMI_K3_MODEL, moonshot_base_url_is_exact_kimi_code}, |
| 15 | }; |
| 16 | |
| 17 | const MAX_NATIVE_SEARCH_ROUNDS: usize = 4; |
| 18 | const MAX_NATIVE_SEARCH_TOOL_CALLS: usize = 8; |
| 19 | /// What the Kimi adapters requested before #6508, kept for a model whose |
| 20 | /// output ceiling the catalogue does not document. |
| 21 | pub(super) const PRIOR_NATIVE_SEARCH_MAX_COMPLETION_TOKENS: u32 = 4_096; |
| 22 | const WEB_SEARCH_FORMULA_URI: &str = "moonshot/web-search:latest"; |
| 23 | const WEB_SEARCH_FORMULA_FUNCTION: &str = "web_search"; |
| 24 | |
| 25 | pub(super) async fn search( |
| 26 | client: &ProviderNativeSearchClient, |
| 27 | request: &ProviderNativeSearchRequest, |
| 28 | ) -> Result<ProviderNativeSearchResponse> { |
| 29 | if moonshot_base_url_is_exact_kimi_code(&client.inner.base_url) { |
| 30 | // Exact Kimi Code membership endpoint only: the structured `/search` |
| 31 | // service is a first-party contract, so differently-cased or adjacent |
| 32 | // Kimi-hosted paths must not inherit it. |
| 33 | search_kimi_code(client, request).await |
| 34 | } else if client |
| 35 | .inner |
| 36 | .default_model |
| 37 | .trim() |
| 38 | .eq_ignore_ascii_case(MOONSHOT_KIMI_K3_MODEL) |
| 39 | { |
| 40 | search_formula(client, request).await |
| 41 | } else { |
| 42 | search_builtin(client, request).await |
| 43 | } |
| 44 | } |
| 45 | |
| 46 | async fn search_kimi_code( |
| 47 | client: &ProviderNativeSearchClient, |
| 48 | request: &ProviderNativeSearchRequest, |
| 49 | ) -> Result<ProviderNativeSearchResponse> { |
| 50 | let call_id = HeaderValue::from_str(&Uuid::new_v4().to_string()) |
| 51 | .context("failed to build Kimi search call id")?; |
| 52 | let url = format!("{}/search", client.inner.base_url.trim_end_matches('/')); |
| 53 | let payload = client |
| 54 | .post_json( |
| 55 | &url, |
| 56 | &json!({ "text_query": request.query }), |
| 57 | &[(HeaderName::from_static("x-msh-tool-call-id"), call_id)], |
| 58 | ) |
| 59 | .await?; |
| 60 | Ok(parse_kimi_code(&payload)) |
| 61 | } |
| 62 | |
| 63 | async fn search_builtin( |
| 64 | client: &ProviderNativeSearchClient, |
| 65 | request: &ProviderNativeSearchRequest, |
| 66 | ) -> Result<ProviderNativeSearchResponse> { |
| 67 | let tools = builtin_search_tools(); |
| 68 | let mut messages = vec![json!({ |
| 69 | "role": "user", |
| 70 | "content": super::search_prompt(request), |
| 71 | })]; |
| 72 | let mut tool_calls_executed = 0; |
| 73 | let url = api_url(&client.inner.base_url, "chat/completions"); |
| 74 | let max_completion_tokens = client |
| 75 | .requested_answer_output_tokens() |
| 76 | .unwrap_or(PRIOR_NATIVE_SEARCH_MAX_COMPLETION_TOKENS); |
| 77 | |
| 78 | for _ in 0..MAX_NATIVE_SEARCH_ROUNDS { |
| 79 | let body = json!({ |
| 80 | "model": client.inner.default_model, |
| 81 | "messages": &messages, |
| 82 | "tools": &tools, |
| 83 | "max_completion_tokens": max_completion_tokens, |
| 84 | "stream": false, |
| 85 | "thinking": { "type": "disabled" }, |
| 86 | }); |
| 87 | let payload = client.post_json(&url, &body, &[]).await?; |
| 88 | let choice = payload |
| 89 | .pointer("/choices/0") |
| 90 | .context("Kimi web search response omitted choices[0]")?; |
| 91 | let message = choice |
| 92 | .get("message") |
| 93 | .and_then(Value::as_object) |
| 94 | .context("Kimi web search response omitted assistant message")?; |
| 95 | if choice.get("finish_reason").and_then(Value::as_str) != Some("tool_calls") { |
| 96 | return Ok(parse_final_message(message, choice)); |
| 97 | } |
| 98 | |
| 99 | messages.push(Value::Object(message.clone())); |
| 100 | let tool_calls = message |
| 101 | .get("tool_calls") |
| 102 | .and_then(Value::as_array) |
| 103 | .context("Kimi returned tool_calls finish reason without tool calls")?; |
| 104 | if tool_calls.is_empty() { |
| 105 | bail!("Kimi returned an empty native web-search tool call list"); |
| 106 | } |
| 107 | reserve_native_search_tool_calls(&mut tool_calls_executed, tool_calls.len())?; |
| 108 | for tool_call in tool_calls { |
| 109 | if tool_call.pointer("/function/name").and_then(Value::as_str) != Some("$web_search") { |
| 110 | bail!("Kimi native search requested an unexpected tool"); |
| 111 | } |
| 112 | let id = tool_call |
| 113 | .get("id") |
| 114 | .and_then(Value::as_str) |
| 115 | .context("Kimi native web-search call omitted id")?; |
| 116 | let arguments = tool_call |
| 117 | .pointer("/function/arguments") |
| 118 | .and_then(Value::as_str) |
| 119 | .context("Kimi native web-search call omitted arguments")?; |
| 120 | let _: Value = serde_json::from_str(arguments) |
| 121 | .context("Kimi native web-search arguments were not valid JSON")?; |
| 122 | messages.push(json!({ |
| 123 | "role": "tool", |
| 124 | "tool_call_id": id, |
| 125 | "name": "$web_search", |
| 126 | "content": arguments, |
| 127 | })); |
| 128 | } |
| 129 | } |
| 130 | |
| 131 | bail!("Kimi native web search exceeded the bounded tool-call loop") |
| 132 | } |
| 133 | |
| 134 | async fn search_formula( |
| 135 | client: &ProviderNativeSearchClient, |
| 136 | request: &ProviderNativeSearchRequest, |
| 137 | ) -> Result<ProviderNativeSearchResponse> { |
| 138 | let formula_path = format!("formulas/{WEB_SEARCH_FORMULA_URI}"); |
| 139 | let tools_payload = client |
| 140 | .get_json(&api_url( |
| 141 | &client.inner.base_url, |
| 142 | &format!("{formula_path}/tools"), |
| 143 | )) |
| 144 | .await?; |
| 145 | let tools = formula_web_search_tools(&tools_payload)?; |
| 146 | let mut messages = vec![json!({ |
| 147 | "role": "user", |
| 148 | "content": super::search_prompt(request), |
| 149 | })]; |
| 150 | let mut tool_calls_executed = 0; |
| 151 | let chat_url = api_url(&client.inner.base_url, "chat/completions"); |
| 152 | let max_completion_tokens = client |
| 153 | .requested_answer_output_tokens() |
| 154 | .unwrap_or(PRIOR_NATIVE_SEARCH_MAX_COMPLETION_TOKENS); |
| 155 | let fiber_url = api_url(&client.inner.base_url, &format!("{formula_path}/fibers")); |
| 156 | |
| 157 | for _ in 0..MAX_NATIVE_SEARCH_ROUNDS { |
| 158 | let body = json!({ |
| 159 | "model": client.inner.default_model, |
| 160 | "messages": &messages, |
| 161 | "tools": &tools, |
| 162 | "max_completion_tokens": max_completion_tokens, |
| 163 | "stream": false, |
| 164 | }); |
| 165 | let payload = client.post_json(&chat_url, &body, &[]).await?; |
| 166 | let choice = payload |
| 167 | .pointer("/choices/0") |
| 168 | .context("Kimi Formula web search response omitted choices[0]")?; |
| 169 | let message = choice |
| 170 | .get("message") |
| 171 | .and_then(Value::as_object) |
| 172 | .context("Kimi Formula web search response omitted assistant message")?; |
| 173 | let Some(tool_calls) = message |
| 174 | .get("tool_calls") |
| 175 | .and_then(Value::as_array) |
| 176 | .filter(|calls| !calls.is_empty()) |
| 177 | else { |
| 178 | return Ok(parse_final_message(message, choice)); |
| 179 | }; |
| 180 | |
| 181 | reserve_native_search_tool_calls(&mut tool_calls_executed, tool_calls.len())?; |
| 182 | messages.push(Value::Object(message.clone())); |
| 183 | for tool_call in tool_calls { |
| 184 | let id = tool_call |
| 185 | .get("id") |
| 186 | .and_then(Value::as_str) |
| 187 | .context("Kimi Formula web-search call omitted id")?; |
| 188 | let function = tool_call |
| 189 | .get("function") |
| 190 | .and_then(Value::as_object) |
| 191 | .context("Kimi Formula web-search call omitted function")?; |
| 192 | let name = function |
| 193 | .get("name") |
| 194 | .and_then(Value::as_str) |
| 195 | .context("Kimi Formula web-search call omitted function name")?; |
| 196 | if name != WEB_SEARCH_FORMULA_FUNCTION { |
| 197 | bail!("Kimi Formula web search requested an unexpected tool"); |
| 198 | } |
| 199 | let arguments = function |
| 200 | .get("arguments") |
| 201 | .and_then(Value::as_str) |
| 202 | .context("Kimi Formula web-search call omitted arguments")?; |
| 203 | let _: Value = serde_json::from_str(arguments) |
| 204 | .context("Kimi Formula web-search arguments were not valid JSON")?; |
| 205 | let fiber = client |
| 206 | .post_json( |
| 207 | &fiber_url, |
| 208 | &json!({ "name": name, "arguments": arguments }), |
| 209 | &[], |
| 210 | ) |
| 211 | .await?; |
| 212 | messages.push(json!({ |
| 213 | "role": "tool", |
| 214 | "tool_call_id": id, |
| 215 | "content": formula_fiber_result(&fiber)?, |
| 216 | })); |
| 217 | } |
| 218 | } |
| 219 | |
| 220 | bail!("Kimi Formula web search exceeded the bounded tool-call loop") |
| 221 | } |
| 222 | |
| 223 | fn reserve_native_search_tool_calls(executed: &mut usize, additional: usize) -> Result<()> { |
| 224 | let total = executed |
| 225 | .checked_add(additional) |
| 226 | .context("Kimi native web search tool-call count overflowed")?; |
| 227 | if total > MAX_NATIVE_SEARCH_TOOL_CALLS { |
| 228 | bail!( |
| 229 | "Kimi native web search exceeded the {MAX_NATIVE_SEARCH_TOOL_CALLS}-call safety limit" |
| 230 | ); |
| 231 | } |
| 232 | *executed = total; |
| 233 | Ok(()) |
| 234 | } |
| 235 | |
| 236 | fn formula_web_search_tools(payload: &Value) -> Result<Value> { |
| 237 | let tools = payload |
| 238 | .get("tools") |
| 239 | .and_then(Value::as_array) |
| 240 | .context("Kimi web-search Formula omitted tools")?; |
| 241 | if tools.len() != 1 |
| 242 | || tools[0].get("type").and_then(Value::as_str) != Some("function") |
| 243 | || tools[0].pointer("/function/name").and_then(Value::as_str) |
| 244 | != Some(WEB_SEARCH_FORMULA_FUNCTION) |
| 245 | { |
| 246 | bail!("Kimi web-search Formula returned an unexpected tool declaration"); |
| 247 | } |
| 248 | Ok(Value::Array(tools.clone())) |
| 249 | } |
| 250 | |
| 251 | fn formula_fiber_result(payload: &Value) -> Result<&str> { |
| 252 | if payload.get("status").and_then(Value::as_str) != Some("succeeded") { |
| 253 | bail!("Kimi web-search Formula fiber did not succeed"); |
| 254 | } |
| 255 | payload |
| 256 | .pointer("/context/output") |
| 257 | .or_else(|| payload.pointer("/context/encrypted_output")) |
| 258 | .and_then(Value::as_str) |
| 259 | .map(str::trim) |
| 260 | .filter(|result| !result.is_empty()) |
| 261 | .context("Kimi web-search Formula fiber omitted its result") |
| 262 | } |
| 263 | |
| 264 | fn builtin_search_tools() -> Value { |
| 265 | json!([{ |
| 266 | "type": "builtin_function", |
| 267 | "function": { "name": "$web_search" } |
| 268 | }]) |
| 269 | } |
| 270 | |
| 271 | fn parse_kimi_code(payload: &Value) -> ProviderNativeSearchResponse { |
| 272 | let mut citations = Vec::new(); |
| 273 | if let Some(results) = payload.get("search_results").and_then(Value::as_array) { |
| 274 | for result in results { |
| 275 | let Some(url) = result.get("url").and_then(Value::as_str) else { |
| 276 | continue; |
| 277 | }; |
| 278 | let title = result |
| 279 | .get("title") |
| 280 | .and_then(Value::as_str) |
| 281 | .map(str::to_string); |
| 282 | let snippet = result |
| 283 | .get("snippet") |
| 284 | .and_then(Value::as_str) |
| 285 | .map(str::to_string); |
| 286 | let published = result |
| 287 | .get("date") |
| 288 | .and_then(Value::as_str) |
| 289 | .filter(|value| !value.is_empty()) |
| 290 | .map(str::to_string); |
| 291 | push_citation( |
| 292 | &mut citations, |
| 293 | citation_from_url(url, title, snippet, published), |
| 294 | ); |
| 295 | } |
| 296 | } |
| 297 | ProviderNativeSearchResponse { |
| 298 | answer: None, |
| 299 | citations, |
| 300 | truncated: false, |
| 301 | } |
| 302 | } |
| 303 | |
| 304 | /// The final answer of a Kimi search loop. `finish_reason: "length"` means |
| 305 | /// Kimi stopped it at the output limit, so it is marked as cut (#6508). |
| 306 | fn parse_final_message( |
| 307 | message: &Map<String, Value>, |
| 308 | choice: &Value, |
| 309 | ) -> ProviderNativeSearchResponse { |
| 310 | let answer = message |
| 311 | .get("content") |
| 312 | .and_then(Value::as_str) |
| 313 | .map(str::trim) |
| 314 | .filter(|text| !text.is_empty()) |
| 315 | .map(str::to_string); |
| 316 | let citations = answer |
| 317 | .as_deref() |
| 318 | .map(citations_from_text) |
| 319 | .unwrap_or_default(); |
| 320 | ProviderNativeSearchResponse { |
| 321 | answer: joined_answer(answer.into_iter().collect()), |
| 322 | citations, |
| 323 | truncated: finish_reason_is_length(Some(choice)), |
| 324 | } |
| 325 | } |
| 326 | |
| 327 | #[cfg(test)] |
| 328 | mod tests { |
| 329 | use super::*; |
| 330 | use crate::config::{Config, ProviderConfig, ProvidersConfig}; |
| 331 | use wiremock::matchers::{body_partial_json, body_string_contains, header, method, path}; |
| 332 | use wiremock::{Mock, MockServer, ResponseTemplate}; |
| 333 | |
| 334 | fn request() -> ProviderNativeSearchRequest { |
| 335 | ProviderNativeSearchRequest { |
| 336 | query: "current release".to_string(), |
| 337 | max_results: 3, |
| 338 | domains: Vec::new(), |
| 339 | } |
| 340 | } |
| 341 | |
| 342 | #[test] |
| 343 | fn kimi_code_request_and_structured_response_contract() { |
| 344 | let body = json!({ "text_query": request().query }); |
| 345 | assert_eq!(body["text_query"], "current release"); |
| 346 | assert_eq!(body.as_object().map(serde_json::Map::len), Some(1)); |
| 347 | |
| 348 | let parsed = parse_kimi_code(&json!({ |
| 349 | "search_results": [{ |
| 350 | "title": "Kimi", |
| 351 | "url": "https://example.com/kimi", |
| 352 | "snippet": "Summary", |
| 353 | "date": "2026-08-28" |
| 354 | }] |
| 355 | })); |
| 356 | assert_eq!(parsed.citations.len(), 1); |
| 357 | assert_eq!(parsed.citations[0].snippet.as_deref(), Some("Summary")); |
| 358 | assert_eq!(parsed.citations[0].published.as_deref(), Some("2026-08-28")); |
| 359 | } |
| 360 | |
| 361 | #[test] |
| 362 | fn direct_search_contracts_are_bounded() { |
| 363 | let tools = builtin_search_tools(); |
| 364 | assert_eq!(tools[0]["function"]["name"], "$web_search"); |
| 365 | assert_eq!(PRIOR_NATIVE_SEARCH_MAX_COMPLETION_TOKENS, 4_096); |
| 366 | |
| 367 | let formula_tools = formula_web_search_tools(&json!({ |
| 368 | "tools": [{ |
| 369 | "type": "function", |
| 370 | "function": { "name": "web_search" } |
| 371 | }] |
| 372 | })) |
| 373 | .expect("formula tools"); |
| 374 | assert_eq!(formula_tools[0]["function"]["name"], "web_search"); |
| 375 | assert_eq!( |
| 376 | formula_fiber_result(&json!({ |
| 377 | "status": "succeeded", |
| 378 | "context": { "encrypted_output": "encrypted result" } |
| 379 | })) |
| 380 | .expect("formula result"), |
| 381 | "encrypted result" |
| 382 | ); |
| 383 | } |
| 384 | |
| 385 | #[test] |
| 386 | fn native_search_tool_call_limit_is_total_not_per_round() { |
| 387 | let mut executed = 0; |
| 388 | reserve_native_search_tool_calls(&mut executed, 4).expect("first rounds"); |
| 389 | reserve_native_search_tool_calls(&mut executed, 4).expect("final allowed round"); |
| 390 | assert_eq!(executed, MAX_NATIVE_SEARCH_TOOL_CALLS); |
| 391 | assert!(reserve_native_search_tool_calls(&mut executed, 1).is_err()); |
| 392 | } |
| 393 | |
| 394 | #[test] |
| 395 | fn kimi_code_dispatch_reuses_the_exact_route_matcher() { |
| 396 | for route in [ |
| 397 | "https://api.kimi.com/coding/v1", |
| 398 | "https://api.kimi.com/coding/v1/", |
| 399 | "HTTPS://API.KIMI.COM/coding/v1", |
| 400 | ] { |
| 401 | assert!( |
| 402 | moonshot_base_url_is_exact_kimi_code(route), |
| 403 | "{route} is the membership endpoint" |
| 404 | ); |
| 405 | } |
| 406 | for neighboring_route in [ |
| 407 | // A case-variant path is a different route, not the official one. |
| 408 | "https://API.KIMI.COM/CODING/V1", |
| 409 | "https://api.kimi.com/coding/v2", |
| 410 | "https://api.kimi.com/coding", |
| 411 | "http://api.kimi.com/coding/v1", |
| 412 | "https://api.moonshot.ai/v1", |
| 413 | ] { |
| 414 | assert!( |
| 415 | !moonshot_base_url_is_exact_kimi_code(neighboring_route), |
| 416 | "{neighboring_route} must not reach the Kimi Code /search service" |
| 417 | ); |
| 418 | } |
| 419 | } |
| 420 | |
| 421 | #[test] |
| 422 | fn finish_reason_length_marks_the_final_answer_as_cut() { |
| 423 | // Both Kimi loops end in parse_final_message. A reply stopped at the |
| 424 | // output limit used to be returned as a finished answer (#6508). |
| 425 | let message = json!({ "role": "assistant", "content": "Partial answer" }); |
| 426 | let message = message.as_object().expect("object"); |
| 427 | let cut = parse_final_message(message, &json!({ "finish_reason": "length" })); |
| 428 | assert!(cut.truncated); |
| 429 | assert_eq!(cut.answer.as_deref(), Some("Partial answer")); |
| 430 | let whole = parse_final_message(message, &json!({ "finish_reason": "stop" })); |
| 431 | assert!(!whole.truncated); |
| 432 | } |
| 433 | |
| 434 | #[test] |
| 435 | fn model_without_a_documented_ceiling_keeps_4096() { |
| 436 | let config = Config { |
| 437 | provider: Some("moonshot".to_string()), |
| 438 | providers: Some(ProvidersConfig { |
| 439 | moonshot: ProviderConfig { |
| 440 | api_key: Some("moonshot-test-key".to_string()), |
| 441 | base_url: Some("https://api.moonshot.ai/v1".to_string()), |
| 442 | model: Some("kimi-unlisted-test-model".to_string()), |
| 443 | ..ProviderConfig::default() |
| 444 | }, |
| 445 | ..ProvidersConfig::default() |
| 446 | }), |
| 447 | ..Config::default() |
| 448 | }; |
| 449 | let client = ProviderNativeSearchClient::new( |
| 450 | crate::client::CodewhaleClient::new(&config).expect("test Moonshot client"), |
| 451 | ) |
| 452 | .expect("Moonshot native adapter"); |
| 453 | let route_cap = client |
| 454 | .inner |
| 455 | .effective_max_output_tokens("kimi-unlisted-test-model"); |
| 456 | assert_eq!( |
| 457 | client.answer_output_tokens(PRIOR_NATIVE_SEARCH_MAX_COMPLETION_TOKENS), |
| 458 | PRIOR_NATIVE_SEARCH_MAX_COMPLETION_TOKENS.min(route_cap) |
| 459 | ); |
| 460 | } |
| 461 | |
| 462 | #[tokio::test] |
| 463 | async fn k3_formula_executes_tool_fiber_and_returns_citations() { |
| 464 | let server = MockServer::start().await; |
| 465 | let config = Config { |
| 466 | provider: Some("moonshot".to_string()), |
| 467 | providers: Some(ProvidersConfig { |
| 468 | moonshot: ProviderConfig { |
| 469 | api_key: Some("moonshot-test-key".to_string()), |
| 470 | base_url: Some(format!("{}/v1", server.uri())), |
| 471 | model: Some("kimi-k3".to_string()), |
| 472 | ..ProviderConfig::default() |
| 473 | }, |
| 474 | ..ProvidersConfig::default() |
| 475 | }), |
| 476 | ..Config::default() |
| 477 | }; |
| 478 | let client = ProviderNativeSearchClient::new( |
| 479 | crate::client::CodewhaleClient::new(&config).expect("test Moonshot client"), |
| 480 | ) |
| 481 | .expect("Moonshot native adapter"); |
| 482 | let expected_completion_tokens = client |
| 483 | .requested_answer_output_tokens() |
| 484 | .expect("Kimi requests an explicit answer length"); |
| 485 | Mock::given(method("GET")) |
| 486 | .and(path("/v1/formulas/moonshot/web-search:latest/tools")) |
| 487 | .and(header("authorization", "Bearer moonshot-test-key")) |
| 488 | .respond_with(ResponseTemplate::new(200).set_body_json(json!({ |
| 489 | "tools": [{ |
| 490 | "type": "function", |
| 491 | "function": { |
| 492 | "name": "web_search", |
| 493 | "description": "Search the web", |
| 494 | "parameters": { |
| 495 | "type": "object", |
| 496 | "properties": { "query": { "type": "string" } }, |
| 497 | "required": ["query"] |
| 498 | } |
| 499 | } |
| 500 | }] |
| 501 | }))) |
| 502 | .expect(1) |
| 503 | .mount(&server) |
| 504 | .await; |
| 505 | Mock::given(method("POST")) |
| 506 | .and(path("/v1/chat/completions")) |
| 507 | .and(header("authorization", "Bearer moonshot-test-key")) |
| 508 | .and(body_partial_json(json!({ |
| 509 | "model": "kimi-k3", |
| 510 | "max_completion_tokens": expected_completion_tokens, |
| 511 | "tools": [{ |
| 512 | "type": "function", |
| 513 | "function": { "name": "web_search" } |
| 514 | }] |
| 515 | }))) |
| 516 | .respond_with(ResponseTemplate::new(200).set_body_json(json!({ |
| 517 | "choices": [{ |
| 518 | "finish_reason": "tool_calls", |
| 519 | "message": { |
| 520 | "role": "assistant", |
| 521 | "content": "", |
| 522 | "tool_calls": [{ |
| 523 | "id": "web_search:0", |
| 524 | "type": "function", |
| 525 | "function": { |
| 526 | "name": "web_search", |
| 527 | "arguments": "{\"query\":\"current release\"}" |
| 528 | } |
| 529 | }] |
| 530 | } |
| 531 | }] |
| 532 | }))) |
| 533 | .up_to_n_times(1) |
| 534 | .expect(1) |
| 535 | .mount(&server) |
| 536 | .await; |
| 537 | Mock::given(method("POST")) |
| 538 | .and(path("/v1/formulas/moonshot/web-search:latest/fibers")) |
| 539 | .and(header("authorization", "Bearer moonshot-test-key")) |
| 540 | .and(body_partial_json(json!({ |
| 541 | "name": "web_search", |
| 542 | "arguments": "{\"query\":\"current release\"}" |
| 543 | }))) |
| 544 | .respond_with(ResponseTemplate::new(200).set_body_json(json!({ |
| 545 | "status": "succeeded", |
| 546 | "context": { "encrypted_output": "encrypted-search-result" } |
| 547 | }))) |
| 548 | .expect(1) |
| 549 | .mount(&server) |
| 550 | .await; |
| 551 | Mock::given(method("POST")) |
| 552 | .and(path("/v1/chat/completions")) |
| 553 | .and(header("authorization", "Bearer moonshot-test-key")) |
| 554 | .and(body_string_contains("encrypted-search-result")) |
| 555 | .respond_with(ResponseTemplate::new(200).set_body_json(json!({ |
| 556 | "choices": [{ |
| 557 | "finish_reason": "stop", |
| 558 | "message": { |
| 559 | "role": "assistant", |
| 560 | "content": "See https://example.com/kimi for the current result." |
| 561 | } |
| 562 | }] |
| 563 | }))) |
| 564 | .expect(1) |
| 565 | .mount(&server) |
| 566 | .await; |
| 567 | |
| 568 | let response = search_formula(&client, &request()) |
| 569 | .await |
| 570 | .expect("K3 Formula search"); |
| 571 | |
| 572 | assert_eq!(response.citations.len(), 1); |
| 573 | assert_eq!(response.citations[0].url, "https://example.com/kimi"); |
| 574 | } |
| 575 | } |
| 576 |