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

302 строки
8.7 KiB

  1. use crate::config::ProtocolKind;
  2. use crate::protocols;
  3. use anyhow::{Context, Result};
  4. use chrono::Utc;
  5. use reqwest::Client;
  6. use reqwest::header::HeaderMap;
  7. use serde::Serialize;
  8. use serde_json::Value;
  9. use std::collections::BTreeMap;
  10. use std::fmt;
  11. use std::path::PathBuf;
  12. use std::sync::Arc;
  13. use std::sync::atomic::{AtomicU64, Ordering};
  14. use std::time::Instant;
  15. #[derive(Clone)]
  16. pub struct ModelRequest {
  17. pub base_url: String,
  18. pub api_token: String,
  19. pub model: String,
  20. pub prompt: String,
  21. pub temperature: f32,
  22. pub max_tokens: u32,
  23. pub stream: bool,
  24. pub raw_debug: Option<RawDebugConfig>,
  25. pub thinking: Option<ThinkingConfig>,
  26. }
  27. #[derive(Debug, Clone, PartialEq, Eq)]
  28. pub struct ThinkingConfig {
  29. pub enabled: bool,
  30. pub kind: Option<String>,
  31. pub budget_tokens: Option<u32>,
  32. pub effort: Option<String>,
  33. pub display: Option<String>,
  34. pub reasoning_effort: Option<String>,
  35. pub reasoning_summary: Option<String>,
  36. }
  37. #[derive(Clone)]
  38. pub struct RawDebugConfig {
  39. output_dir: PathBuf,
  40. prefix: String,
  41. counter: Arc<AtomicU64>,
  42. write_success_raw: bool,
  43. }
  44. impl RawDebugConfig {
  45. pub fn new(output_dir: PathBuf, prefix: String) -> Self {
  46. Self {
  47. output_dir,
  48. prefix: sanitize_filename_component(&prefix),
  49. counter: Arc::new(AtomicU64::new(0)),
  50. write_success_raw: true,
  51. }
  52. }
  53. pub fn with_success_raw(mut self, write_success_raw: bool) -> Self {
  54. self.write_success_raw = write_success_raw;
  55. self
  56. }
  57. pub fn write_success_raw(&self) -> bool {
  58. self.write_success_raw
  59. }
  60. pub async fn write_response(&self, response_kind: &str, contents: &str) -> Result<PathBuf> {
  61. self.write_debug_file(response_kind, contents).await
  62. }
  63. pub async fn write_http_error(
  64. &self,
  65. response_kind: &str,
  66. request: HttpDebugRequest,
  67. response: HttpDebugResponse,
  68. ) -> Result<PathBuf> {
  69. let envelope = HttpDebugEnvelope { request, response };
  70. let contents = serde_json::to_string_pretty(&envelope)
  71. .context("failed to serialize raw debug HTTP error")?;
  72. self.write_debug_file(response_kind, &contents).await
  73. }
  74. pub async fn write_request_error(
  75. &self,
  76. response_kind: &str,
  77. request: HttpDebugRequest,
  78. error: String,
  79. ) -> Result<PathBuf> {
  80. let envelope = HttpDebugEnvelope {
  81. request,
  82. response: HttpDebugResponse {
  83. status: None,
  84. headers: BTreeMap::new(),
  85. body: String::new(),
  86. error_kind: Some("request_error".to_string()),
  87. error: Some(error),
  88. },
  89. };
  90. let contents = serde_json::to_string_pretty(&envelope)
  91. .context("failed to serialize raw debug request error")?;
  92. self.write_debug_file(response_kind, &contents).await
  93. }
  94. async fn write_debug_file(&self, response_kind: &str, contents: &str) -> Result<PathBuf> {
  95. tokio::fs::create_dir_all(&self.output_dir)
  96. .await
  97. .with_context(|| {
  98. format!(
  99. "failed to create raw debug dir {}",
  100. self.output_dir.display()
  101. )
  102. })?;
  103. let sequence = self.counter.fetch_add(1, Ordering::Relaxed) + 1;
  104. let timestamp = Utc::now().format("%Y%m%dT%H%M%S%.3fZ");
  105. let filename = format!(
  106. "{}-{}-{sequence:06}-{}.txt",
  107. self.prefix,
  108. timestamp,
  109. sanitize_filename_component(response_kind)
  110. );
  111. let path = self.output_dir.join(filename);
  112. tokio::fs::write(&path, contents)
  113. .await
  114. .with_context(|| format!("failed to write raw debug response {}", path.display()))?;
  115. Ok(path)
  116. }
  117. }
  118. #[derive(Debug, Serialize)]
  119. pub struct HttpDebugEnvelope {
  120. pub request: HttpDebugRequest,
  121. pub response: HttpDebugResponse,
  122. }
  123. #[derive(Debug, Serialize)]
  124. pub struct HttpDebugRequest {
  125. pub method: String,
  126. pub url: String,
  127. pub headers: BTreeMap<String, String>,
  128. pub body: Value,
  129. }
  130. #[derive(Debug, Serialize)]
  131. pub struct HttpDebugResponse {
  132. pub status: Option<u16>,
  133. pub headers: BTreeMap<String, String>,
  134. pub body: String,
  135. #[serde(skip_serializing_if = "Option::is_none")]
  136. pub error_kind: Option<String>,
  137. #[serde(skip_serializing_if = "Option::is_none")]
  138. pub error: Option<String>,
  139. }
  140. pub fn request_headers_for_debug(headers: &[(&str, String)]) -> BTreeMap<String, String> {
  141. headers
  142. .iter()
  143. .map(|(name, value)| {
  144. let name = name.to_ascii_lowercase();
  145. let value = if is_sensitive_header(&name) {
  146. redact_header_value(&name, value)
  147. } else {
  148. value.clone()
  149. };
  150. (name, value)
  151. })
  152. .collect()
  153. }
  154. pub fn response_headers_for_debug(headers: &HeaderMap) -> BTreeMap<String, String> {
  155. headers
  156. .iter()
  157. .map(|(name, value)| {
  158. (
  159. name.as_str().to_ascii_lowercase(),
  160. value.to_str().unwrap_or("<non-utf8>").to_string(),
  161. )
  162. })
  163. .collect()
  164. }
  165. fn is_sensitive_header(name: &str) -> bool {
  166. matches!(
  167. name,
  168. "authorization" | "x-api-key" | "x-goog-api-key" | "api-key"
  169. )
  170. }
  171. fn redact_header_value(name: &str, value: &str) -> String {
  172. if name == "authorization"
  173. && let Some((scheme, _)) = value.split_once(' ')
  174. {
  175. format!("{scheme} [REDACTED]")
  176. } else {
  177. "[REDACTED]".to_string()
  178. }
  179. }
  180. impl fmt::Debug for ModelRequest {
  181. fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
  182. formatter
  183. .debug_struct("ModelRequest")
  184. .field("base_url", &self.base_url)
  185. .field("api_token", &"[REDACTED]")
  186. .field("model", &self.model)
  187. .field("prompt", &self.prompt)
  188. .field("temperature", &self.temperature)
  189. .field("max_tokens", &self.max_tokens)
  190. .field("stream", &self.stream)
  191. .field("raw_debug", &self.raw_debug.is_some())
  192. .field("thinking", &self.thinking)
  193. .finish()
  194. }
  195. }
  196. fn sanitize_filename_component(value: &str) -> String {
  197. let sanitized = value
  198. .chars()
  199. .map(|character| {
  200. if character.is_ascii_alphanumeric() || matches!(character, '.' | '-' | '_') {
  201. character
  202. } else {
  203. '-'
  204. }
  205. })
  206. .collect::<String>();
  207. let collapsed = sanitized
  208. .split('-')
  209. .filter(|part| !part.is_empty())
  210. .collect::<Vec<_>>()
  211. .join("-");
  212. if collapsed.is_empty() {
  213. "unknown".to_string()
  214. } else {
  215. collapsed
  216. }
  217. }
  218. #[derive(Debug, Clone, PartialEq, Eq)]
  219. pub struct ModelResponse {
  220. pub text: String,
  221. pub status: u16,
  222. pub elapsed_ms: u128,
  223. pub first_token_ms: Option<u128>,
  224. }
  225. pub async fn run_model_request(
  226. protocol: ProtocolKind,
  227. request: ModelRequest,
  228. ) -> Result<ModelResponse> {
  229. let client = Client::new();
  230. run_model_request_with_client(&client, protocol, &request).await
  231. }
  232. pub async fn run_model_request_with_client(
  233. client: &Client,
  234. protocol: ProtocolKind,
  235. request: &ModelRequest,
  236. ) -> Result<ModelResponse> {
  237. let started = Instant::now();
  238. let mut response = if request.stream {
  239. match protocol {
  240. ProtocolKind::Openai => protocols::openai::send_stream(client, request).await?,
  241. ProtocolKind::Anthropic => protocols::anthropic::send_stream(client, request).await?,
  242. ProtocolKind::Google => protocols::google::send_stream(client, request).await?,
  243. }
  244. } else {
  245. match protocol {
  246. ProtocolKind::Openai => protocols::openai::send(client, request).await?,
  247. ProtocolKind::Anthropic => protocols::anthropic::send(client, request).await?,
  248. ProtocolKind::Google => protocols::google::send(client, request).await?,
  249. }
  250. };
  251. response.elapsed_ms = started.elapsed().as_millis();
  252. Ok(response)
  253. }
  254. #[cfg(test)]
  255. mod tests {
  256. use super::*;
  257. #[test]
  258. fn model_request_debug_redacts_api_token() {
  259. let request = ModelRequest {
  260. base_url: "https://relay.example.com/v1".to_string(),
  261. api_token: "sk-secret-token".to_string(),
  262. model: "test-model".to_string(),
  263. prompt: "hello".to_string(),
  264. temperature: 0.0,
  265. max_tokens: 1024,
  266. stream: false,
  267. raw_debug: None,
  268. thinking: None,
  269. };
  270. let debug = format!("{request:?}");
  271. assert!(debug.contains("api_token"));
  272. assert!(debug.contains("[REDACTED]"));
  273. assert!(!debug.contains("sk-secret-token"));
  274. }
  275. }