api模型检测
Du kan inte välja fler än 25 ämnen Ämnen måste starta med en bokstav eller siffra, kan innehålla bindestreck ('-') och vara max 35 tecken långa.

137 rader
4.1 KiB

  1. use crate::runner::{ModelRequest, ModelResponse};
  2. use anyhow::{Context, Result, bail};
  3. use reqwest::Client;
  4. use serde::Deserialize;
  5. use serde_json::json;
  6. pub async fn send(client: &Client, request: &ModelRequest) -> Result<ModelResponse> {
  7. let url = super::endpoint_url(&request.base_url, "/v1/messages")?;
  8. let response = client
  9. .post(url)
  10. .header("x-api-key", &request.api_token)
  11. .header("anthropic-version", "2023-06-01")
  12. .json(&json!({
  13. "model": request.model,
  14. "messages": [{"role": "user", "content": request.prompt}],
  15. "temperature": request.temperature,
  16. "max_tokens": request.max_tokens
  17. }))
  18. .send()
  19. .await
  20. .context("failed to send Anthropic messages request")?;
  21. let status = response.status();
  22. let status_code = status.as_u16();
  23. let body = response
  24. .text()
  25. .await
  26. .context("failed to read Anthropic response body")?;
  27. if !status.is_success() {
  28. bail!("Anthropic request failed with status {status_code}: {body}");
  29. }
  30. let parsed: AnthropicResponse =
  31. serde_json::from_str(&body).context("failed to parse Anthropic response JSON")?;
  32. let text = parsed
  33. .content
  34. .into_iter()
  35. .find_map(|block| match block {
  36. AnthropicContentBlock::Text { text } if !text.is_empty() => Some(text),
  37. _ => None,
  38. })
  39. .context("Anthropic response missing first text content block")?;
  40. Ok(ModelResponse {
  41. text,
  42. status: status_code,
  43. elapsed_ms: 0,
  44. })
  45. }
  46. #[derive(Debug, Deserialize)]
  47. struct AnthropicResponse {
  48. content: Vec<AnthropicContentBlock>,
  49. }
  50. #[derive(Debug, Deserialize)]
  51. #[serde(tag = "type")]
  52. enum AnthropicContentBlock {
  53. #[serde(rename = "text")]
  54. Text { text: String },
  55. #[serde(other)]
  56. Other,
  57. }
  58. #[cfg(test)]
  59. mod tests {
  60. use crate::runner::ModelRequest;
  61. use reqwest::Client;
  62. use wiremock::matchers::{body_json, header, method, path};
  63. use wiremock::{Mock, MockServer, ResponseTemplate};
  64. #[tokio::test]
  65. async fn extracts_first_text_block() {
  66. let server = MockServer::start().await;
  67. Mock::given(method("POST"))
  68. .and(path("/v1/messages"))
  69. .and(header("x-api-key", "test-token"))
  70. .and(header("anthropic-version", "2023-06-01"))
  71. .and(body_json(serde_json::json!({
  72. "model": "claude-test",
  73. "messages": [{"role": "user", "content": "hello"}],
  74. "temperature": 0.0,
  75. "max_tokens": 1024
  76. })))
  77. .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
  78. "content": [
  79. {"type": "text", "text": "hi from claude"}
  80. ]
  81. })))
  82. .mount(&server)
  83. .await;
  84. let request = ModelRequest {
  85. base_url: server.uri(),
  86. api_token: "test-token".to_string(),
  87. model: "claude-test".to_string(),
  88. prompt: "hello".to_string(),
  89. temperature: 0.0,
  90. max_tokens: 1024,
  91. };
  92. let response = super::send(&Client::new(), &request)
  93. .await
  94. .expect("response");
  95. assert_eq!(response.status, 200);
  96. assert_eq!(response.text, "hi from claude");
  97. }
  98. #[tokio::test]
  99. async fn errors_on_non_success_status() {
  100. let server = MockServer::start().await;
  101. Mock::given(method("POST"))
  102. .and(path("/v1/messages"))
  103. .respond_with(ResponseTemplate::new(401).set_body_string("unauthorized"))
  104. .mount(&server)
  105. .await;
  106. let request = ModelRequest {
  107. base_url: server.uri(),
  108. api_token: "test-token".to_string(),
  109. model: "claude-test".to_string(),
  110. prompt: "hello".to_string(),
  111. temperature: 0.0,
  112. max_tokens: 1024,
  113. };
  114. let error = super::send(&Client::new(), &request)
  115. .await
  116. .expect_err("non-success should fail");
  117. assert!(error.to_string().contains("401"));
  118. assert!(error.to_string().contains("unauthorized"));
  119. }
  120. }