api模型检测
No puede seleccionar más de 25 temas Los temas deben comenzar con una letra o número, pueden incluir guiones ('-') y pueden tener hasta 35 caracteres de largo.
 
 

308 líneas
9.2 KiB

  1. use crate::runner::{ModelRequest, ModelResponse};
  2. use anyhow::{Context, Result, bail};
  3. use futures::StreamExt;
  4. use reqwest::Client;
  5. use serde::Deserialize;
  6. use serde_json::json;
  7. use std::time::Instant;
  8. pub async fn send(client: &Client, request: &ModelRequest) -> Result<ModelResponse> {
  9. let url = super::endpoint_url(&request.base_url, "/v1/messages")?;
  10. let response = client
  11. .post(url)
  12. .header("x-api-key", &request.api_token)
  13. .header("anthropic-version", "2023-06-01")
  14. .json(&json!({
  15. "model": request.model,
  16. "messages": [{"role": "user", "content": request.prompt}],
  17. "temperature": request.temperature,
  18. "max_tokens": request.max_tokens
  19. }))
  20. .send()
  21. .await
  22. .context("failed to send Anthropic messages request")?;
  23. let status = response.status();
  24. let status_code = status.as_u16();
  25. let body = response
  26. .text()
  27. .await
  28. .context("failed to read Anthropic response body")?;
  29. if !status.is_success() {
  30. bail!(
  31. "{}",
  32. super::upstream_error_message("Anthropic", status_code, &body)
  33. );
  34. }
  35. let parsed: AnthropicResponse =
  36. serde_json::from_str(&body).context("failed to parse Anthropic response JSON")?;
  37. let text = parsed
  38. .content
  39. .into_iter()
  40. .find_map(|block| match block {
  41. AnthropicContentBlock::Text { text } if !text.is_empty() => Some(text),
  42. _ => None,
  43. })
  44. .context("Anthropic response missing first text content block")?;
  45. Ok(ModelResponse {
  46. text,
  47. status: status_code,
  48. elapsed_ms: 0,
  49. first_token_ms: None,
  50. })
  51. }
  52. pub async fn send_stream(client: &Client, request: &ModelRequest) -> Result<ModelResponse> {
  53. let url = super::endpoint_url(&request.base_url, "/v1/messages")?;
  54. let started = Instant::now();
  55. let response = client
  56. .post(url)
  57. .header("x-api-key", &request.api_token)
  58. .header("anthropic-version", "2023-06-01")
  59. .json(&json!({
  60. "model": request.model,
  61. "messages": [{"role": "user", "content": request.prompt}],
  62. "temperature": request.temperature,
  63. "max_tokens": request.max_tokens,
  64. "stream": true
  65. }))
  66. .send()
  67. .await
  68. .context("failed to send Anthropic streaming request")?;
  69. let status = response.status();
  70. let status_code = status.as_u16();
  71. if !status.is_success() {
  72. let body = response
  73. .text()
  74. .await
  75. .context("failed to read Anthropic error response body")?;
  76. bail!(
  77. "{}",
  78. super::upstream_error_message("Anthropic", status_code, &body)
  79. );
  80. }
  81. let mut stream = response.bytes_stream();
  82. let mut buffer = super::SseLineBuffer::new();
  83. let mut text = String::new();
  84. let mut first_token_ms: Option<u128> = None;
  85. let mut current_event = String::new();
  86. let mut done = false;
  87. while let Some(chunk) = stream.next().await {
  88. let chunk = chunk.context("Anthropic stream interrupted")?;
  89. for line in buffer.feed(&chunk) {
  90. if let Some(event_type) = line.strip_prefix("event: ") {
  91. current_event = event_type.to_string();
  92. continue;
  93. }
  94. if let Some(data) = line.strip_prefix("data: ") {
  95. if current_event == "content_block_delta" {
  96. if let Some(content) = serde_json::from_str::<AnthropicStreamDelta>(data)
  97. .ok()
  98. .and_then(|p| p.delta.text)
  99. .filter(|c| !c.is_empty())
  100. {
  101. if first_token_ms.is_none() {
  102. first_token_ms = Some(started.elapsed().as_millis());
  103. }
  104. text.push_str(&content);
  105. }
  106. } else if current_event == "message_stop" {
  107. done = true;
  108. break;
  109. }
  110. }
  111. if line.is_empty() {
  112. current_event.clear();
  113. }
  114. }
  115. if done {
  116. break;
  117. }
  118. }
  119. if text.is_empty() {
  120. bail!("Anthropic stream completed without producing any content");
  121. }
  122. Ok(ModelResponse {
  123. text,
  124. status: status_code,
  125. elapsed_ms: 0,
  126. first_token_ms,
  127. })
  128. }
  129. #[derive(Debug, Deserialize)]
  130. struct AnthropicResponse {
  131. content: Vec<AnthropicContentBlock>,
  132. }
  133. #[derive(Debug, Deserialize)]
  134. #[serde(tag = "type")]
  135. enum AnthropicContentBlock {
  136. #[serde(rename = "text")]
  137. Text { text: String },
  138. #[serde(other)]
  139. Other,
  140. }
  141. #[derive(Debug, Deserialize)]
  142. struct AnthropicStreamDelta {
  143. delta: AnthropicDeltaContent,
  144. }
  145. #[derive(Debug, Deserialize)]
  146. struct AnthropicDeltaContent {
  147. text: Option<String>,
  148. }
  149. #[cfg(test)]
  150. mod tests {
  151. use crate::runner::ModelRequest;
  152. use reqwest::Client;
  153. use wiremock::matchers::{body_json, header, method, path};
  154. use wiremock::{Mock, MockServer, ResponseTemplate};
  155. #[tokio::test]
  156. async fn extracts_first_text_block() {
  157. let server = MockServer::start().await;
  158. Mock::given(method("POST"))
  159. .and(path("/v1/messages"))
  160. .and(header("x-api-key", "test-token"))
  161. .and(header("anthropic-version", "2023-06-01"))
  162. .and(body_json(serde_json::json!({
  163. "model": "claude-test",
  164. "messages": [{"role": "user", "content": "hello"}],
  165. "temperature": 0.0,
  166. "max_tokens": 1024
  167. })))
  168. .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
  169. "content": [
  170. {"type": "text", "text": "hi from claude"}
  171. ]
  172. })))
  173. .mount(&server)
  174. .await;
  175. let request = ModelRequest {
  176. base_url: server.uri(),
  177. api_token: "test-token".to_string(),
  178. model: "claude-test".to_string(),
  179. prompt: "hello".to_string(),
  180. temperature: 0.0,
  181. max_tokens: 1024,
  182. stream: false
  183. };
  184. let response = super::send(&Client::new(), &request)
  185. .await
  186. .expect("response");
  187. assert_eq!(response.status, 200);
  188. assert_eq!(response.text, "hi from claude");
  189. }
  190. #[tokio::test]
  191. async fn errors_on_non_success_status() {
  192. let server = MockServer::start().await;
  193. Mock::given(method("POST"))
  194. .and(path("/v1/messages"))
  195. .respond_with(ResponseTemplate::new(401).set_body_json(serde_json::json!({
  196. "error": {
  197. "message": "unauthorized"
  198. }
  199. })))
  200. .mount(&server)
  201. .await;
  202. let request = ModelRequest {
  203. base_url: server.uri(),
  204. api_token: "test-token".to_string(),
  205. model: "claude-test".to_string(),
  206. prompt: "hello".to_string(),
  207. temperature: 0.0,
  208. max_tokens: 1024,
  209. stream: false
  210. };
  211. let error = super::send(&Client::new(), &request)
  212. .await
  213. .expect_err("non-success should fail");
  214. assert!(error.to_string().contains("401"));
  215. assert!(error.to_string().contains("unauthorized"));
  216. }
  217. #[tokio::test]
  218. async fn non_success_error_does_not_leak_sensitive_response_body() {
  219. let server = MockServer::start().await;
  220. Mock::given(method("POST"))
  221. .and(path("/v1/messages"))
  222. .respond_with(ResponseTemplate::new(429).set_body_json(serde_json::json!({
  223. "error": {
  224. "message": "rate limited secret prompt sk-leaked-token"
  225. }
  226. })))
  227. .mount(&server)
  228. .await;
  229. let request = ModelRequest {
  230. base_url: server.uri(),
  231. api_token: "test-token".to_string(),
  232. model: "claude-test".to_string(),
  233. prompt: "hello".to_string(),
  234. temperature: 0.0,
  235. max_tokens: 1024,
  236. stream: false
  237. };
  238. let message = super::send(&Client::new(), &request)
  239. .await
  240. .expect_err("non-success should fail")
  241. .to_string();
  242. assert!(message.contains("Anthropic"));
  243. assert!(message.contains("429"));
  244. assert!(message.contains("rate limited"));
  245. assert!(!message.contains("secret prompt"));
  246. assert!(!message.contains("sk-leaked-token"));
  247. }
  248. #[tokio::test]
  249. async fn base_url_with_v1_prefix_does_not_duplicate_messages_path() {
  250. let server = MockServer::start().await;
  251. Mock::given(method("POST"))
  252. .and(path("/v1/messages"))
  253. .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
  254. "content": [
  255. {"type": "text", "text": "prefixed"}
  256. ]
  257. })))
  258. .mount(&server)
  259. .await;
  260. let request = ModelRequest {
  261. base_url: format!("{}/v1", server.uri()),
  262. api_token: "test-token".to_string(),
  263. model: "claude-test".to_string(),
  264. prompt: "hello".to_string(),
  265. temperature: 0.0,
  266. max_tokens: 1024,
  267. stream: false
  268. };
  269. let response = super::send(&Client::new(), &request)
  270. .await
  271. .expect("response");
  272. assert_eq!(response.text, "prefixed");
  273. }
  274. }