|
- use crate::config::ProtocolKind;
- use crate::protocols;
- use anyhow::{Context, Result};
- use chrono::Utc;
- use reqwest::Client;
- use std::fmt;
- use std::path::PathBuf;
- use std::sync::Arc;
- use std::sync::atomic::{AtomicU64, Ordering};
- use std::time::Instant;
-
- #[derive(Clone)]
- pub struct ModelRequest {
- pub base_url: String,
- pub api_token: String,
- pub model: String,
- pub prompt: String,
- pub temperature: f32,
- pub max_tokens: u32,
- pub stream: bool,
- pub raw_debug: Option<RawDebugConfig>,
- pub thinking: Option<ThinkingConfig>,
- }
-
- #[derive(Debug, Clone, PartialEq, Eq)]
- pub struct ThinkingConfig {
- pub enabled: bool,
- pub kind: Option<String>,
- pub budget_tokens: Option<u32>,
- pub effort: Option<String>,
- pub display: Option<String>,
- pub reasoning_effort: Option<String>,
- pub reasoning_summary: Option<String>,
- }
-
- #[derive(Clone)]
- pub struct RawDebugConfig {
- output_dir: PathBuf,
- prefix: String,
- counter: Arc<AtomicU64>,
- }
-
- impl RawDebugConfig {
- pub fn new(output_dir: PathBuf, prefix: String) -> Self {
- Self {
- output_dir,
- prefix: sanitize_filename_component(&prefix),
- counter: Arc::new(AtomicU64::new(0)),
- }
- }
-
- pub async fn write_response(&self, response_kind: &str, contents: &str) -> Result<PathBuf> {
- tokio::fs::create_dir_all(&self.output_dir)
- .await
- .with_context(|| {
- format!(
- "failed to create raw debug dir {}",
- self.output_dir.display()
- )
- })?;
- let sequence = self.counter.fetch_add(1, Ordering::Relaxed) + 1;
- let timestamp = Utc::now().format("%Y%m%dT%H%M%S%.3fZ");
- let filename = format!(
- "{}-{}-{sequence:06}-{}.txt",
- self.prefix,
- timestamp,
- sanitize_filename_component(response_kind)
- );
- let path = self.output_dir.join(filename);
- tokio::fs::write(&path, contents)
- .await
- .with_context(|| format!("failed to write raw debug response {}", path.display()))?;
- Ok(path)
- }
- }
-
- impl fmt::Debug for ModelRequest {
- fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
- formatter
- .debug_struct("ModelRequest")
- .field("base_url", &self.base_url)
- .field("api_token", &"[REDACTED]")
- .field("model", &self.model)
- .field("prompt", &self.prompt)
- .field("temperature", &self.temperature)
- .field("max_tokens", &self.max_tokens)
- .field("stream", &self.stream)
- .field("raw_debug", &self.raw_debug.is_some())
- .field("thinking", &self.thinking)
- .finish()
- }
- }
-
- fn sanitize_filename_component(value: &str) -> String {
- let sanitized = value
- .chars()
- .map(|character| {
- if character.is_ascii_alphanumeric() || matches!(character, '.' | '-' | '_') {
- character
- } else {
- '-'
- }
- })
- .collect::<String>();
- let collapsed = sanitized
- .split('-')
- .filter(|part| !part.is_empty())
- .collect::<Vec<_>>()
- .join("-");
- if collapsed.is_empty() {
- "unknown".to_string()
- } else {
- collapsed
- }
- }
-
- #[derive(Debug, Clone, PartialEq, Eq)]
- pub struct ModelResponse {
- pub text: String,
- pub status: u16,
- pub elapsed_ms: u128,
- pub first_token_ms: Option<u128>,
- }
-
- pub async fn run_model_request(
- protocol: ProtocolKind,
- request: ModelRequest,
- ) -> Result<ModelResponse> {
- let client = Client::new();
- run_model_request_with_client(&client, protocol, &request).await
- }
-
- pub async fn run_model_request_with_client(
- client: &Client,
- protocol: ProtocolKind,
- request: &ModelRequest,
- ) -> Result<ModelResponse> {
- let started = Instant::now();
- let mut response = if request.stream {
- match protocol {
- ProtocolKind::Openai => protocols::openai::send_stream(client, request).await?,
- ProtocolKind::Anthropic => protocols::anthropic::send_stream(client, request).await?,
- ProtocolKind::Google => protocols::google::send_stream(client, request).await?,
- }
- } else {
- match protocol {
- ProtocolKind::Openai => protocols::openai::send(client, request).await?,
- ProtocolKind::Anthropic => protocols::anthropic::send(client, request).await?,
- ProtocolKind::Google => protocols::google::send(client, request).await?,
- }
- };
- response.elapsed_ms = started.elapsed().as_millis();
- Ok(response)
- }
-
- #[cfg(test)]
- mod tests {
- use super::*;
-
- #[test]
- fn model_request_debug_redacts_api_token() {
- let request = ModelRequest {
- base_url: "https://relay.example.com/v1".to_string(),
- api_token: "sk-secret-token".to_string(),
- model: "test-model".to_string(),
- prompt: "hello".to_string(),
- temperature: 0.0,
- max_tokens: 1024,
- stream: false,
- raw_debug: None,
- thinking: None,
- };
-
- let debug = format!("{request:?}");
-
- assert!(debug.contains("api_token"));
- assert!(debug.contains("[REDACTED]"));
- assert!(!debug.contains("sk-secret-token"));
- }
- }
|