api模型检测
Você não pode selecionar mais de 25 tópicos Os tópicos devem começar com uma letra ou um número, podem incluir traços ('-') e podem ter até 35 caracteres.

181 linhas
5.3 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 std::fmt;
  7. use std::path::PathBuf;
  8. use std::sync::Arc;
  9. use std::sync::atomic::{AtomicU64, Ordering};
  10. use std::time::Instant;
  11. #[derive(Clone)]
  12. pub struct ModelRequest {
  13. pub base_url: String,
  14. pub api_token: String,
  15. pub model: String,
  16. pub prompt: String,
  17. pub temperature: f32,
  18. pub max_tokens: u32,
  19. pub stream: bool,
  20. pub raw_debug: Option<RawDebugConfig>,
  21. pub thinking: Option<ThinkingConfig>,
  22. }
  23. #[derive(Debug, Clone, PartialEq, Eq)]
  24. pub struct ThinkingConfig {
  25. pub enabled: bool,
  26. pub kind: Option<String>,
  27. pub budget_tokens: Option<u32>,
  28. pub effort: Option<String>,
  29. pub display: Option<String>,
  30. pub reasoning_effort: Option<String>,
  31. pub reasoning_summary: Option<String>,
  32. }
  33. #[derive(Clone)]
  34. pub struct RawDebugConfig {
  35. output_dir: PathBuf,
  36. prefix: String,
  37. counter: Arc<AtomicU64>,
  38. }
  39. impl RawDebugConfig {
  40. pub fn new(output_dir: PathBuf, prefix: String) -> Self {
  41. Self {
  42. output_dir,
  43. prefix: sanitize_filename_component(&prefix),
  44. counter: Arc::new(AtomicU64::new(0)),
  45. }
  46. }
  47. pub async fn write_response(&self, response_kind: &str, contents: &str) -> Result<PathBuf> {
  48. tokio::fs::create_dir_all(&self.output_dir)
  49. .await
  50. .with_context(|| {
  51. format!(
  52. "failed to create raw debug dir {}",
  53. self.output_dir.display()
  54. )
  55. })?;
  56. let sequence = self.counter.fetch_add(1, Ordering::Relaxed) + 1;
  57. let timestamp = Utc::now().format("%Y%m%dT%H%M%S%.3fZ");
  58. let filename = format!(
  59. "{}-{}-{sequence:06}-{}.txt",
  60. self.prefix,
  61. timestamp,
  62. sanitize_filename_component(response_kind)
  63. );
  64. let path = self.output_dir.join(filename);
  65. tokio::fs::write(&path, contents)
  66. .await
  67. .with_context(|| format!("failed to write raw debug response {}", path.display()))?;
  68. Ok(path)
  69. }
  70. }
  71. impl fmt::Debug for ModelRequest {
  72. fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
  73. formatter
  74. .debug_struct("ModelRequest")
  75. .field("base_url", &self.base_url)
  76. .field("api_token", &"[REDACTED]")
  77. .field("model", &self.model)
  78. .field("prompt", &self.prompt)
  79. .field("temperature", &self.temperature)
  80. .field("max_tokens", &self.max_tokens)
  81. .field("stream", &self.stream)
  82. .field("raw_debug", &self.raw_debug.is_some())
  83. .field("thinking", &self.thinking)
  84. .finish()
  85. }
  86. }
  87. fn sanitize_filename_component(value: &str) -> String {
  88. let sanitized = value
  89. .chars()
  90. .map(|character| {
  91. if character.is_ascii_alphanumeric() || matches!(character, '.' | '-' | '_') {
  92. character
  93. } else {
  94. '-'
  95. }
  96. })
  97. .collect::<String>();
  98. let collapsed = sanitized
  99. .split('-')
  100. .filter(|part| !part.is_empty())
  101. .collect::<Vec<_>>()
  102. .join("-");
  103. if collapsed.is_empty() {
  104. "unknown".to_string()
  105. } else {
  106. collapsed
  107. }
  108. }
  109. #[derive(Debug, Clone, PartialEq, Eq)]
  110. pub struct ModelResponse {
  111. pub text: String,
  112. pub status: u16,
  113. pub elapsed_ms: u128,
  114. pub first_token_ms: Option<u128>,
  115. }
  116. pub async fn run_model_request(
  117. protocol: ProtocolKind,
  118. request: ModelRequest,
  119. ) -> Result<ModelResponse> {
  120. let client = Client::new();
  121. run_model_request_with_client(&client, protocol, &request).await
  122. }
  123. pub async fn run_model_request_with_client(
  124. client: &Client,
  125. protocol: ProtocolKind,
  126. request: &ModelRequest,
  127. ) -> Result<ModelResponse> {
  128. let started = Instant::now();
  129. let mut response = if request.stream {
  130. match protocol {
  131. ProtocolKind::Openai => protocols::openai::send_stream(client, request).await?,
  132. ProtocolKind::Anthropic => protocols::anthropic::send_stream(client, request).await?,
  133. ProtocolKind::Google => protocols::google::send_stream(client, request).await?,
  134. }
  135. } else {
  136. match protocol {
  137. ProtocolKind::Openai => protocols::openai::send(client, request).await?,
  138. ProtocolKind::Anthropic => protocols::anthropic::send(client, request).await?,
  139. ProtocolKind::Google => protocols::google::send(client, request).await?,
  140. }
  141. };
  142. response.elapsed_ms = started.elapsed().as_millis();
  143. Ok(response)
  144. }
  145. #[cfg(test)]
  146. mod tests {
  147. use super::*;
  148. #[test]
  149. fn model_request_debug_redacts_api_token() {
  150. let request = ModelRequest {
  151. base_url: "https://relay.example.com/v1".to_string(),
  152. api_token: "sk-secret-token".to_string(),
  153. model: "test-model".to_string(),
  154. prompt: "hello".to_string(),
  155. temperature: 0.0,
  156. max_tokens: 1024,
  157. stream: false,
  158. raw_debug: None,
  159. thinking: None,
  160. };
  161. let debug = format!("{request:?}");
  162. assert!(debug.contains("api_token"));
  163. assert!(debug.contains("[REDACTED]"));
  164. assert!(!debug.contains("sk-secret-token"));
  165. }
  166. }