api模型检测
Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.

995 строки
33 KiB

  1. use crate::runner::{
  2. HttpDebugRequest, HttpDebugResponse, ModelRequest, ModelResponse, request_headers_for_debug,
  3. response_headers_for_debug,
  4. };
  5. use anyhow::{Context, Result, bail};
  6. use futures::StreamExt;
  7. use reqwest::Client;
  8. use serde::Deserialize;
  9. use serde_json::{Value, json};
  10. use std::time::Instant;
  11. pub async fn send(client: &Client, request: &ModelRequest) -> Result<ModelResponse> {
  12. let url = super::endpoint_url(&request.base_url, "/v1/messages")?;
  13. let debug_url = url.to_string();
  14. let request_body = request_body(request, false);
  15. let response = client
  16. .post(url)
  17. .header("x-api-key", &request.api_token)
  18. .header("anthropic-version", "2023-06-01")
  19. .json(&request_body)
  20. .send()
  21. .await
  22. .map_err(|error| {
  23. write_request_error_debug_blocking(
  24. request,
  25. &debug_url,
  26. request_body.clone(),
  27. format!("failed to send Anthropic messages request: {error}"),
  28. "anthropic-request-error",
  29. );
  30. error
  31. })
  32. .context("failed to send Anthropic messages request")?;
  33. let status = response.status();
  34. let status_code = status.as_u16();
  35. let response_headers = response.headers().clone();
  36. let body = response
  37. .text()
  38. .await
  39. .map_err(|error| {
  40. write_request_error_debug_blocking(
  41. request,
  42. &debug_url,
  43. request_body.clone(),
  44. format!("failed to read Anthropic response body: {error}"),
  45. "anthropic-request-error",
  46. );
  47. error
  48. })
  49. .context("failed to read Anthropic response body")?;
  50. if status.is_success()
  51. && let Some(raw_debug) = &request.raw_debug
  52. && raw_debug.write_success_raw()
  53. {
  54. raw_debug
  55. .write_response("anthropic-json", &body)
  56. .await
  57. .context("failed to write Anthropic raw debug response")?;
  58. }
  59. if !status.is_success() {
  60. write_error_debug(
  61. request,
  62. &debug_url,
  63. request_body,
  64. status_code,
  65. &response_headers,
  66. &body,
  67. "anthropic-error-http",
  68. )
  69. .await?;
  70. bail!(
  71. "{}",
  72. super::upstream_error_message("Anthropic", status_code, &body)
  73. );
  74. }
  75. let parsed: AnthropicResponse = match serde_json::from_str(&body) {
  76. Ok(parsed) => parsed,
  77. Err(error) => {
  78. write_response_request_error_debug(
  79. request,
  80. &debug_url,
  81. request_body,
  82. status_code,
  83. &response_headers,
  84. &body,
  85. format!("failed to parse Anthropic response JSON: {error}"),
  86. )
  87. .await?;
  88. return Err(error).context("failed to parse Anthropic response JSON");
  89. }
  90. };
  91. let Some(text) = parsed.content.into_iter().find_map(|block| match block {
  92. AnthropicContentBlock::Text { text } if !text.is_empty() => Some(text),
  93. _ => None,
  94. }) else {
  95. write_response_request_error_debug(
  96. request,
  97. &debug_url,
  98. request_body,
  99. status_code,
  100. &response_headers,
  101. &body,
  102. "Anthropic response missing first text content block".to_string(),
  103. )
  104. .await?;
  105. bail!("Anthropic response missing first text content block");
  106. };
  107. Ok(ModelResponse {
  108. text,
  109. status: status_code,
  110. elapsed_ms: 0,
  111. first_token_ms: None,
  112. })
  113. }
  114. pub async fn send_stream(client: &Client, request: &ModelRequest) -> Result<ModelResponse> {
  115. let url = super::endpoint_url(&request.base_url, "/v1/messages")?;
  116. let debug_url = url.to_string();
  117. let started = Instant::now();
  118. let request_body = request_body(request, true);
  119. let response = client
  120. .post(url)
  121. .header("x-api-key", &request.api_token)
  122. .header("anthropic-version", "2023-06-01")
  123. .json(&request_body)
  124. .send()
  125. .await
  126. .map_err(|error| {
  127. write_request_error_debug_blocking(
  128. request,
  129. &debug_url,
  130. request_body.clone(),
  131. format!("failed to send Anthropic streaming request: {error}"),
  132. "anthropic-request-error",
  133. );
  134. error
  135. })
  136. .context("failed to send Anthropic streaming request")?;
  137. let status = response.status();
  138. let status_code = status.as_u16();
  139. let response_headers = response.headers().clone();
  140. if !status.is_success() {
  141. let body = response
  142. .text()
  143. .await
  144. .map_err(|error| {
  145. write_request_error_debug_blocking(
  146. request,
  147. &debug_url,
  148. request_body.clone(),
  149. format!("failed to read Anthropic error response body: {error}"),
  150. "anthropic-request-error",
  151. );
  152. error
  153. })
  154. .context("failed to read Anthropic error response body")?;
  155. if let Some(raw_debug) = &request.raw_debug {
  156. raw_debug
  157. .write_http_error(
  158. "anthropic-error-http",
  159. debug_request(request, &debug_url, request_body),
  160. HttpDebugResponse {
  161. status: Some(status_code),
  162. headers: response_headers_for_debug(&response_headers),
  163. body: body.clone(),
  164. error_kind: None,
  165. error: None,
  166. },
  167. )
  168. .await
  169. .context("failed to write Anthropic raw debug error response")?;
  170. }
  171. bail!(
  172. "{}",
  173. super::upstream_error_message("Anthropic", status_code, &body)
  174. );
  175. }
  176. let mut stream = response.bytes_stream();
  177. let mut buffer = super::SseLineBuffer::new();
  178. let mut text = String::new();
  179. let mut raw_stream = String::new();
  180. let mut first_token_ms: Option<u128> = None;
  181. let mut current_event = String::new();
  182. let mut done = false;
  183. while let Some(chunk) = stream.next().await {
  184. let chunk = match chunk {
  185. Ok(chunk) => chunk,
  186. Err(error) => {
  187. write_response_request_error_debug(
  188. request,
  189. &debug_url,
  190. request_body.clone(),
  191. status_code,
  192. &response_headers,
  193. &raw_stream,
  194. format!("Anthropic stream interrupted: {error}"),
  195. )
  196. .await?;
  197. return Err(error).context("Anthropic stream interrupted");
  198. }
  199. };
  200. for line in buffer.feed(&chunk) {
  201. raw_stream.push_str(&line);
  202. raw_stream.push('\n');
  203. if let Some(event_type) = line.strip_prefix("event: ") {
  204. current_event = event_type.to_string();
  205. continue;
  206. }
  207. if let Some(data) = line.strip_prefix("data: ") {
  208. if current_event == "content_block_delta" {
  209. if let Some(content) = serde_json::from_str::<AnthropicStreamDelta>(data)
  210. .ok()
  211. .and_then(|p| p.delta.text)
  212. .filter(|c| !c.is_empty())
  213. {
  214. if first_token_ms.is_none() {
  215. first_token_ms = Some(started.elapsed().as_millis());
  216. }
  217. text.push_str(&content);
  218. }
  219. } else if current_event == "message_stop" {
  220. done = true;
  221. break;
  222. }
  223. }
  224. if line.is_empty() {
  225. current_event.clear();
  226. }
  227. }
  228. if done {
  229. break;
  230. }
  231. }
  232. if let Some(raw_debug) = &request.raw_debug
  233. && raw_debug.write_success_raw()
  234. {
  235. raw_debug
  236. .write_response("anthropic-sse", &raw_stream)
  237. .await
  238. .context("failed to write Anthropic raw debug stream")?;
  239. }
  240. if text.is_empty() {
  241. write_response_request_error_debug(
  242. request,
  243. &debug_url,
  244. request_body,
  245. status_code,
  246. &response_headers,
  247. &raw_stream,
  248. "Anthropic stream completed without producing any content".to_string(),
  249. )
  250. .await?;
  251. bail!("Anthropic stream completed without producing any content");
  252. }
  253. Ok(ModelResponse {
  254. text,
  255. status: status_code,
  256. elapsed_ms: 0,
  257. first_token_ms,
  258. })
  259. }
  260. #[derive(Debug, Deserialize)]
  261. struct AnthropicResponse {
  262. content: Vec<AnthropicContentBlock>,
  263. }
  264. #[derive(Debug, Deserialize)]
  265. #[serde(tag = "type")]
  266. enum AnthropicContentBlock {
  267. #[serde(rename = "text")]
  268. Text { text: String },
  269. #[serde(other)]
  270. Other,
  271. }
  272. #[derive(Debug, Deserialize)]
  273. struct AnthropicStreamDelta {
  274. delta: AnthropicDeltaContent,
  275. }
  276. #[derive(Debug, Deserialize)]
  277. struct AnthropicDeltaContent {
  278. text: Option<String>,
  279. }
  280. fn request_body(request: &ModelRequest, stream: bool) -> Value {
  281. let mut body = json!({
  282. "model": request.model,
  283. "messages": [{"role": "user", "content": request.prompt}],
  284. "temperature": request.temperature,
  285. "max_tokens": request.max_tokens
  286. });
  287. let object = body
  288. .as_object_mut()
  289. .expect("Anthropic request body is an object");
  290. if stream {
  291. object.insert("stream".to_string(), json!(true));
  292. }
  293. if let Some(thinking) = &request.thinking
  294. && thinking.enabled
  295. {
  296. object.remove("temperature");
  297. let mut thinking_body = serde_json::Map::new();
  298. let thinking_type = thinking.kind.as_deref().unwrap_or("enabled");
  299. thinking_body.insert("type".to_string(), json!(thinking_type));
  300. if thinking_type == "adaptive" {
  301. if let Some(effort) = &thinking.effort {
  302. thinking_body.insert("effort".to_string(), json!(effort));
  303. }
  304. } else if let Some(budget_tokens) = thinking.budget_tokens {
  305. thinking_body.insert("budget_tokens".to_string(), json!(budget_tokens));
  306. }
  307. if let Some(display) = &thinking.display {
  308. thinking_body.insert("display".to_string(), json!(display));
  309. }
  310. object.insert("thinking".to_string(), Value::Object(thinking_body));
  311. }
  312. body
  313. }
  314. async fn write_error_debug(
  315. request: &ModelRequest,
  316. url: &str,
  317. request_body: Value,
  318. status: u16,
  319. response_headers: &reqwest::header::HeaderMap,
  320. response_body: &str,
  321. response_kind: &str,
  322. ) -> Result<()> {
  323. if let Some(raw_debug) = &request.raw_debug {
  324. raw_debug
  325. .write_http_error(
  326. response_kind,
  327. debug_request(request, url, request_body),
  328. HttpDebugResponse {
  329. status: Some(status),
  330. headers: response_headers_for_debug(response_headers),
  331. body: response_body.to_string(),
  332. error_kind: None,
  333. error: None,
  334. },
  335. )
  336. .await
  337. .context("failed to write Anthropic raw debug error response")?;
  338. }
  339. Ok(())
  340. }
  341. async fn write_response_request_error_debug(
  342. request: &ModelRequest,
  343. url: &str,
  344. request_body: Value,
  345. status: u16,
  346. response_headers: &reqwest::header::HeaderMap,
  347. response_body: &str,
  348. error: String,
  349. ) -> Result<()> {
  350. if let Some(raw_debug) = &request.raw_debug {
  351. raw_debug
  352. .write_http_error(
  353. "anthropic-request-error",
  354. debug_request(request, url, request_body),
  355. HttpDebugResponse {
  356. status: Some(status),
  357. headers: response_headers_for_debug(response_headers),
  358. body: response_body.to_string(),
  359. error_kind: Some("request_error".to_string()),
  360. error: Some(error),
  361. },
  362. )
  363. .await
  364. .context("failed to write Anthropic raw debug request error response")?;
  365. }
  366. Ok(())
  367. }
  368. fn debug_request(request: &ModelRequest, url: &str, body: Value) -> HttpDebugRequest {
  369. HttpDebugRequest {
  370. method: "POST".to_string(),
  371. url: url.to_string(),
  372. headers: request_headers_for_debug(&[
  373. ("x-api-key", request.api_token.clone()),
  374. ("anthropic-version", "2023-06-01".to_string()),
  375. ]),
  376. body,
  377. }
  378. }
  379. fn write_request_error_debug_blocking(
  380. request: &ModelRequest,
  381. url: &str,
  382. request_body: Value,
  383. error: String,
  384. response_kind: &str,
  385. ) {
  386. if let Some(raw_debug) = &request.raw_debug {
  387. let _ = futures::executor::block_on(raw_debug.write_request_error(
  388. response_kind,
  389. debug_request(request, url, request_body),
  390. error,
  391. ));
  392. }
  393. }
  394. #[cfg(test)]
  395. mod tests {
  396. use crate::runner::{ModelRequest, RawDebugConfig, ThinkingConfig};
  397. use reqwest::Client;
  398. use std::path::{Path, PathBuf};
  399. use wiremock::matchers::{body_json, header, method, path};
  400. use wiremock::{Mock, MockServer, ResponseTemplate};
  401. fn read_single_debug_file(root: &Path) -> String {
  402. let mut files = Vec::new();
  403. collect_debug_files(root, &mut files);
  404. assert_eq!(files.len(), 1);
  405. std::fs::read_to_string(&files[0]).expect("read raw debug file")
  406. }
  407. fn collect_debug_files(dir: &Path, files: &mut Vec<PathBuf>) {
  408. for entry in std::fs::read_dir(dir)
  409. .expect("read debug dir")
  410. .collect::<Result<Vec<_>, _>>()
  411. .expect("debug entries")
  412. {
  413. let path = entry.path();
  414. if path.is_dir() {
  415. collect_debug_files(&path, files);
  416. } else {
  417. files.push(path);
  418. }
  419. }
  420. }
  421. #[tokio::test]
  422. async fn extracts_first_text_block() {
  423. let server = MockServer::start().await;
  424. Mock::given(method("POST"))
  425. .and(path("/v1/messages"))
  426. .and(header("x-api-key", "test-token"))
  427. .and(header("anthropic-version", "2023-06-01"))
  428. .and(body_json(serde_json::json!({
  429. "model": "claude-test",
  430. "messages": [{"role": "user", "content": "hello"}],
  431. "temperature": 0.0,
  432. "max_tokens": 1024
  433. })))
  434. .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
  435. "content": [
  436. {"type": "text", "text": "hi from claude"}
  437. ]
  438. })))
  439. .mount(&server)
  440. .await;
  441. let request = ModelRequest {
  442. base_url: server.uri(),
  443. api_token: "test-token".to_string(),
  444. model: "claude-test".to_string(),
  445. prompt: "hello".to_string(),
  446. temperature: 0.0,
  447. max_tokens: 1024,
  448. stream: false,
  449. raw_debug: None,
  450. thinking: None,
  451. };
  452. let response = super::send(&Client::new(), &request)
  453. .await
  454. .expect("response");
  455. assert_eq!(response.status, 200);
  456. assert_eq!(response.text, "hi from claude");
  457. }
  458. #[tokio::test]
  459. async fn sends_thinking_and_omits_temperature_when_enabled() {
  460. let server = MockServer::start().await;
  461. Mock::given(method("POST"))
  462. .and(path("/v1/messages"))
  463. .and(body_json(serde_json::json!({
  464. "model": "claude-test",
  465. "messages": [{"role": "user", "content": "hello"}],
  466. "max_tokens": 1024,
  467. "thinking": {
  468. "type": "enabled",
  469. "budget_tokens": 10000,
  470. "display": "omitted"
  471. }
  472. })))
  473. .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
  474. "content": [
  475. {"type": "text", "text": "hi from claude"}
  476. ]
  477. })))
  478. .mount(&server)
  479. .await;
  480. let mut request = anthropic_request(server.uri());
  481. request.thinking = Some(ThinkingConfig {
  482. enabled: true,
  483. kind: Some("enabled".to_string()),
  484. budget_tokens: Some(10000),
  485. effort: None,
  486. display: Some("omitted".to_string()),
  487. reasoning_effort: None,
  488. reasoning_summary: None,
  489. });
  490. let response = super::send(&Client::new(), &request)
  491. .await
  492. .expect("response");
  493. assert_eq!(response.text, "hi from claude");
  494. }
  495. #[tokio::test]
  496. async fn sends_adaptive_thinking_effort_without_budget_tokens() {
  497. let server = MockServer::start().await;
  498. Mock::given(method("POST"))
  499. .and(path("/v1/messages"))
  500. .and(body_json(serde_json::json!({
  501. "model": "claude-test",
  502. "messages": [{"role": "user", "content": "hello"}],
  503. "max_tokens": 1024,
  504. "thinking": {
  505. "type": "adaptive",
  506. "effort": "high",
  507. "display": "omitted"
  508. }
  509. })))
  510. .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
  511. "content": [
  512. {"type": "text", "text": "hi from claude"}
  513. ]
  514. })))
  515. .mount(&server)
  516. .await;
  517. let mut request = anthropic_request(server.uri());
  518. request.thinking = Some(ThinkingConfig {
  519. enabled: true,
  520. kind: Some("adaptive".to_string()),
  521. budget_tokens: Some(10000),
  522. effort: Some("high".to_string()),
  523. display: Some("omitted".to_string()),
  524. reasoning_effort: None,
  525. reasoning_summary: None,
  526. });
  527. let response = super::send(&Client::new(), &request)
  528. .await
  529. .expect("response");
  530. assert_eq!(response.text, "hi from claude");
  531. }
  532. #[tokio::test]
  533. async fn errors_on_non_success_status() {
  534. let server = MockServer::start().await;
  535. Mock::given(method("POST"))
  536. .and(path("/v1/messages"))
  537. .respond_with(ResponseTemplate::new(401).set_body_json(serde_json::json!({
  538. "error": {
  539. "message": "unauthorized"
  540. }
  541. })))
  542. .mount(&server)
  543. .await;
  544. let request = ModelRequest {
  545. base_url: server.uri(),
  546. api_token: "test-token".to_string(),
  547. model: "claude-test".to_string(),
  548. prompt: "hello".to_string(),
  549. temperature: 0.0,
  550. max_tokens: 1024,
  551. stream: false,
  552. raw_debug: None,
  553. thinking: None,
  554. };
  555. let error = super::send(&Client::new(), &request)
  556. .await
  557. .expect_err("non-success should fail");
  558. assert!(error.to_string().contains("401"));
  559. assert!(error.to_string().contains("unauthorized"));
  560. }
  561. #[tokio::test]
  562. async fn non_success_error_does_not_leak_sensitive_response_body() {
  563. let server = MockServer::start().await;
  564. Mock::given(method("POST"))
  565. .and(path("/v1/messages"))
  566. .respond_with(ResponseTemplate::new(429).set_body_json(serde_json::json!({
  567. "error": {
  568. "message": "rate limited secret prompt sk-leaked-token"
  569. }
  570. })))
  571. .mount(&server)
  572. .await;
  573. let request = ModelRequest {
  574. base_url: server.uri(),
  575. api_token: "test-token".to_string(),
  576. model: "claude-test".to_string(),
  577. prompt: "hello".to_string(),
  578. temperature: 0.0,
  579. max_tokens: 1024,
  580. stream: false,
  581. raw_debug: None,
  582. thinking: None,
  583. };
  584. let message = super::send(&Client::new(), &request)
  585. .await
  586. .expect_err("non-success should fail")
  587. .to_string();
  588. assert!(message.contains("Anthropic"));
  589. assert!(message.contains("429"));
  590. assert!(message.contains("rate limited"));
  591. assert!(!message.contains("secret prompt"));
  592. assert!(!message.contains("sk-leaked-token"));
  593. }
  594. #[tokio::test]
  595. async fn non_success_error_debug_records_request_and_response_without_token() {
  596. let server = MockServer::start().await;
  597. let temp_dir = tempfile::tempdir().expect("create temp dir");
  598. Mock::given(method("POST"))
  599. .and(path("/v1/messages"))
  600. .respond_with(
  601. ResponseTemplate::new(504)
  602. .insert_header("x-request-id", "anthropic-req")
  603. .set_body_json(serde_json::json!({
  604. "error": {
  605. "message": "gateway timeout"
  606. }
  607. })),
  608. )
  609. .mount(&server)
  610. .await;
  611. let request = ModelRequest {
  612. base_url: server.uri(),
  613. api_token: "anthropic-secret-token".to_string(),
  614. model: "claude-test".to_string(),
  615. prompt: "full anthropic prompt".to_string(),
  616. temperature: 0.0,
  617. max_tokens: 1024,
  618. stream: false,
  619. raw_debug: Some(RawDebugConfig::new(
  620. temp_dir.path().to_path_buf(),
  621. "anthropic-claude-test".to_string(),
  622. )),
  623. thinking: None,
  624. };
  625. let _ = super::send(&Client::new(), &request)
  626. .await
  627. .expect_err("non-success should fail");
  628. let raw = read_single_debug_file(temp_dir.path());
  629. let debug: serde_json::Value = serde_json::from_str(&raw).expect("debug json");
  630. assert_eq!(debug["request"]["headers"]["x-api-key"], "[REDACTED]");
  631. assert_eq!(
  632. debug["request"]["body"]["messages"][0]["content"],
  633. "full anthropic prompt"
  634. );
  635. assert_eq!(debug["response"]["status"], 504);
  636. assert_eq!(
  637. debug["response"]["headers"]["x-request-id"],
  638. "anthropic-req"
  639. );
  640. assert!(
  641. debug["response"]["body"]
  642. .as_str()
  643. .expect("response body")
  644. .contains("gateway timeout")
  645. );
  646. assert!(!raw.contains("anthropic-secret-token"));
  647. }
  648. #[tokio::test]
  649. async fn request_error_debug_records_request_and_local_error_without_token() {
  650. let temp_dir = tempfile::tempdir().expect("create temp dir");
  651. let request = ModelRequest {
  652. base_url: "http://127.0.0.1:9".to_string(),
  653. api_token: "anthropic-secret-token".to_string(),
  654. model: "claude-test".to_string(),
  655. prompt: "prompt before anthropic connect error".to_string(),
  656. temperature: 0.0,
  657. max_tokens: 1024,
  658. stream: false,
  659. raw_debug: Some(RawDebugConfig::new(
  660. temp_dir.path().to_path_buf(),
  661. "anthropic-claude-test".to_string(),
  662. )),
  663. thinking: None,
  664. };
  665. let _ = super::send(&Client::new(), &request)
  666. .await
  667. .expect_err("connection should fail");
  668. let raw = read_single_debug_file(temp_dir.path());
  669. let debug: serde_json::Value = serde_json::from_str(&raw).expect("debug json");
  670. assert_eq!(debug["request"]["method"], "POST");
  671. assert!(
  672. debug["request"]["url"]
  673. .as_str()
  674. .expect("request url")
  675. .ends_with("/v1/messages")
  676. );
  677. assert_eq!(debug["request"]["headers"]["x-api-key"], "[REDACTED]");
  678. assert_eq!(
  679. debug["request"]["body"]["messages"][0]["content"],
  680. "prompt before anthropic connect error"
  681. );
  682. assert_eq!(debug["response"]["status"], serde_json::Value::Null);
  683. assert_eq!(debug["response"]["error_kind"], "request_error");
  684. assert!(
  685. debug["response"]["error"]
  686. .as_str()
  687. .expect("local error")
  688. .contains("failed to send Anthropic messages request")
  689. );
  690. assert!(!raw.contains("anthropic-secret-token"));
  691. }
  692. #[tokio::test]
  693. async fn base_url_with_v1_prefix_does_not_duplicate_messages_path() {
  694. let server = MockServer::start().await;
  695. Mock::given(method("POST"))
  696. .and(path("/v1/messages"))
  697. .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
  698. "content": [
  699. {"type": "text", "text": "prefixed"}
  700. ]
  701. })))
  702. .mount(&server)
  703. .await;
  704. let request = ModelRequest {
  705. base_url: format!("{}/v1", server.uri()),
  706. api_token: "test-token".to_string(),
  707. model: "claude-test".to_string(),
  708. prompt: "hello".to_string(),
  709. temperature: 0.0,
  710. max_tokens: 1024,
  711. stream: false,
  712. raw_debug: None,
  713. thinking: None,
  714. };
  715. let response = super::send(&Client::new(), &request)
  716. .await
  717. .expect("response");
  718. assert_eq!(response.text, "prefixed");
  719. }
  720. #[tokio::test]
  721. async fn extracts_streaming_messages_text_and_ttft() {
  722. let server = MockServer::start().await;
  723. let temp_dir = tempfile::tempdir().expect("create temp dir");
  724. Mock::given(method("POST"))
  725. .and(path("/v1/messages"))
  726. .and(header("x-api-key", "test-token"))
  727. .and(header("anthropic-version", "2023-06-01"))
  728. .and(body_json(serde_json::json!({
  729. "model": "claude-test",
  730. "messages": [{"role": "user", "content": "hello"}],
  731. "temperature": 0.0,
  732. "max_tokens": 1024,
  733. "stream": true
  734. })))
  735. .respond_with(
  736. ResponseTemplate::new(200)
  737. .insert_header("content-type", "text/event-stream")
  738. .set_body_string(
  739. "event: content_block_delta\n\
  740. data: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"hi \"}}\n\n\
  741. event: content_block_delta\n\
  742. data: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"there\"}}\n\n\
  743. event: message_stop\n\
  744. data: {\"type\":\"message_stop\"}\n\n",
  745. ),
  746. )
  747. .mount(&server)
  748. .await;
  749. let request = ModelRequest {
  750. base_url: server.uri(),
  751. api_token: "test-token".to_string(),
  752. model: "claude-test".to_string(),
  753. prompt: "hello".to_string(),
  754. temperature: 0.0,
  755. max_tokens: 1024,
  756. stream: true,
  757. raw_debug: Some(RawDebugConfig::new(
  758. temp_dir.path().to_path_buf(),
  759. "anthropic-claude-test".to_string(),
  760. )),
  761. thinking: None,
  762. };
  763. let response = super::send_stream(&Client::new(), &request)
  764. .await
  765. .expect("stream response");
  766. assert_eq!(response.status, 200);
  767. assert_eq!(response.text, "hi there");
  768. assert!(response.first_token_ms.is_some());
  769. let raw = read_single_debug_file(temp_dir.path());
  770. assert!(raw.contains("event: content_block_delta"));
  771. assert!(raw.contains("\"text\":\"hi \""));
  772. }
  773. #[tokio::test]
  774. async fn empty_stream_request_error_debug_records_request_and_raw_stream() {
  775. let server = MockServer::start().await;
  776. let temp_dir = tempfile::tempdir().expect("create temp dir");
  777. Mock::given(method("POST"))
  778. .and(path("/v1/messages"))
  779. .respond_with(
  780. ResponseTemplate::new(200)
  781. .insert_header("content-type", "text/event-stream")
  782. .set_body_string(
  783. "event: message_start\n\
  784. data: {\"type\":\"message_start\"}\n\n\
  785. event: message_stop\n\
  786. data: {\"type\":\"message_stop\"}\n\n",
  787. ),
  788. )
  789. .mount(&server)
  790. .await;
  791. let request = ModelRequest {
  792. base_url: server.uri(),
  793. api_token: "anthropic-secret-token".to_string(),
  794. model: "claude-test".to_string(),
  795. prompt: "prompt before empty stream".to_string(),
  796. temperature: 0.0,
  797. max_tokens: 1024,
  798. stream: true,
  799. raw_debug: Some(
  800. RawDebugConfig::new(
  801. temp_dir.path().to_path_buf(),
  802. "anthropic-claude-test".to_string(),
  803. )
  804. .with_success_raw(false),
  805. ),
  806. thinking: None,
  807. };
  808. let error = super::send_stream(&Client::new(), &request)
  809. .await
  810. .expect_err("empty stream should fail");
  811. assert!(
  812. error
  813. .to_string()
  814. .contains("completed without producing any content")
  815. );
  816. let raw = read_single_debug_file(temp_dir.path());
  817. let debug: serde_json::Value = serde_json::from_str(&raw).expect("debug json");
  818. assert_eq!(
  819. debug["request"]["body"]["messages"][0]["content"],
  820. "prompt before empty stream"
  821. );
  822. assert_eq!(debug["response"]["status"], 200);
  823. assert_eq!(debug["response"]["error_kind"], "request_error");
  824. assert!(
  825. debug["response"]["body"]
  826. .as_str()
  827. .expect("response body")
  828. .contains("event: message_stop")
  829. );
  830. assert!(!raw.contains("anthropic-secret-token"));
  831. }
  832. #[tokio::test]
  833. async fn interrupted_stream_debug_records_status_headers_and_partial_stream() {
  834. use tokio::io::{AsyncReadExt, AsyncWriteExt};
  835. use tokio::net::TcpListener;
  836. let listener = TcpListener::bind("127.0.0.1:0")
  837. .await
  838. .expect("bind listener");
  839. let addr = listener.local_addr().expect("listener addr");
  840. let server = tokio::spawn(async move {
  841. let (mut socket, _) = listener.accept().await.expect("accept request");
  842. let mut buffer = [0_u8; 4096];
  843. let _ = socket.read(&mut buffer).await.expect("read request");
  844. let chunk = "event: ping\ndata: {\"type\":\"ping\"}\n\n";
  845. let response = format!(
  846. "HTTP/1.1 200 OK\r\n\
  847. content-type: text/event-stream\r\n\
  848. transfer-encoding: chunked\r\n\
  849. x-request-id: broken-stream\r\n\
  850. \r\n\
  851. {:x}\r\n{}\r\n\
  852. zz\r\nbroken\r\n",
  853. chunk.len(),
  854. chunk
  855. );
  856. socket
  857. .write_all(response.as_bytes())
  858. .await
  859. .expect("write malformed response");
  860. });
  861. let temp_dir = tempfile::tempdir().expect("create temp dir");
  862. let request = ModelRequest {
  863. base_url: format!("http://{addr}"),
  864. api_token: "anthropic-secret-token".to_string(),
  865. model: "claude-test".to_string(),
  866. prompt: "prompt before broken stream".to_string(),
  867. temperature: 0.0,
  868. max_tokens: 1024,
  869. stream: true,
  870. raw_debug: Some(
  871. RawDebugConfig::new(
  872. temp_dir.path().to_path_buf(),
  873. "anthropic-claude-test".to_string(),
  874. )
  875. .with_success_raw(false),
  876. ),
  877. thinking: None,
  878. };
  879. let error = super::send_stream(&Client::new(), &request)
  880. .await
  881. .expect_err("malformed chunked stream should fail");
  882. server.await.expect("server task");
  883. assert!(error.to_string().contains("stream interrupted"));
  884. let raw = read_single_debug_file(temp_dir.path());
  885. let debug: serde_json::Value = serde_json::from_str(&raw).expect("debug json");
  886. assert_eq!(debug["response"]["status"], 200);
  887. assert_eq!(
  888. debug["response"]["headers"]["x-request-id"],
  889. "broken-stream"
  890. );
  891. assert_eq!(debug["response"]["error_kind"], "request_error");
  892. assert!(
  893. debug["response"]["body"]
  894. .as_str()
  895. .expect("partial stream")
  896. .contains("event: ping")
  897. );
  898. assert!(
  899. debug["response"]["error"]
  900. .as_str()
  901. .expect("stream error")
  902. .contains("Anthropic stream interrupted")
  903. );
  904. assert!(!raw.contains("anthropic-secret-token"));
  905. }
  906. fn anthropic_request(base_url: String) -> ModelRequest {
  907. ModelRequest {
  908. base_url,
  909. api_token: "test-token".to_string(),
  910. model: "claude-test".to_string(),
  911. prompt: "hello".to_string(),
  912. temperature: 0.0,
  913. max_tokens: 1024,
  914. stream: false,
  915. raw_debug: None,
  916. thinking: None,
  917. }
  918. }
  919. }