257 lines
7.8 KiB
Rust
257 lines
7.8 KiB
Rust
use anyhow::{anyhow, Result};
|
|
use reqwest::Client;
|
|
use serde::{Deserialize, Serialize};
|
|
use std::env;
|
|
use std::time::Duration;
|
|
use mem_core::gated_loop::LlmClient;
|
|
|
|
/// Completion response from the model.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct Completion {
|
|
pub text: String,
|
|
pub usage: Usage,
|
|
pub latency_ms: u64,
|
|
}
|
|
|
|
/// Token usage breakdown.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct Usage {
|
|
pub prompt_tokens: u32,
|
|
pub completion_tokens: u32,
|
|
pub total_tokens: u32,
|
|
}
|
|
|
|
/// Chat client for the gateway.
|
|
pub struct ChatClient {
|
|
base_url: String,
|
|
api_key: String,
|
|
model: String,
|
|
http: Client,
|
|
timeout: Duration,
|
|
max_retries: u32,
|
|
}
|
|
|
|
#[derive(Debug, Serialize)]
|
|
struct Message {
|
|
role: String,
|
|
content: String,
|
|
}
|
|
|
|
#[derive(Debug, Serialize)]
|
|
struct CompletionRequest {
|
|
model: String,
|
|
messages: Vec<Message>,
|
|
max_tokens: u32,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct CompletionResponse {
|
|
choices: Vec<Choice>,
|
|
usage: ResponseUsage,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct Choice {
|
|
message: MessageResponse,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct MessageResponse {
|
|
role: String,
|
|
content: String,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct ResponseUsage {
|
|
prompt_tokens: u32,
|
|
completion_tokens: u32,
|
|
total_tokens: u32,
|
|
}
|
|
|
|
impl LlmClient for ChatClient {
|
|
fn complete_blocking(&self, system: &str, user: &str, max_tokens: usize) -> Result<String> {
|
|
ChatClient::complete_blocking(self, system, user, max_tokens as u32)
|
|
}
|
|
}
|
|
|
|
impl ChatClient {
|
|
/// Create a new chat client.
|
|
///
|
|
/// # Arguments
|
|
/// * `base_url` - Gateway base URL (e.g., `https://api.riotpiao.com/v1`)
|
|
/// * `api_key` - Authentication key
|
|
/// * `model` - Model identifier (e.g., `qwen2.5:3b-instruct`)
|
|
pub fn new(base_url: impl Into<String>, api_key: impl Into<String>, model: impl Into<String>) -> Result<Self> {
|
|
Ok(Self {
|
|
base_url: base_url.into(),
|
|
api_key: api_key.into(),
|
|
model: model.into(),
|
|
http: Client::new(),
|
|
timeout: Duration::from_secs(300),
|
|
max_retries: 3,
|
|
})
|
|
}
|
|
|
|
/// Set custom timeout.
|
|
pub fn with_timeout(mut self, timeout: Duration) -> Self {
|
|
self.timeout = timeout;
|
|
self
|
|
}
|
|
|
|
/// Set max retries for 5xx errors (default: 3).
|
|
pub fn with_max_retries(mut self, retries: u32) -> Self {
|
|
self.max_retries = retries;
|
|
self
|
|
}
|
|
|
|
/// Complete synchronously (blocks until response).
|
|
pub fn complete_blocking(&self, system: &str, user: &str, max_tokens: u32) -> Result<String> {
|
|
let rt = tokio::runtime::Runtime::new()?;
|
|
rt.block_on(async {
|
|
let completion = self.complete(system, user, max_tokens).await?;
|
|
Ok(completion.text)
|
|
})
|
|
}
|
|
|
|
/// Complete a prompt.
|
|
pub async fn complete(&self, system: &str, user: &str, max_tokens: u32) -> Result<Completion> {
|
|
let url = format!("{}/chat/completions", self.base_url);
|
|
|
|
let request = CompletionRequest {
|
|
model: self.model.clone(),
|
|
messages: vec![
|
|
Message {
|
|
role: "system".to_string(),
|
|
content: system.to_string(),
|
|
},
|
|
Message {
|
|
role: "user".to_string(),
|
|
content: user.to_string(),
|
|
},
|
|
],
|
|
max_tokens,
|
|
};
|
|
|
|
let body = serde_json::to_string(&request)?;
|
|
|
|
// Record request if MEM_LLM_RECORD is set
|
|
if let Ok(record_dir) = env::var("MEM_LLM_RECORD") {
|
|
let filename = format!("{}/request-{}.json", record_dir, chrono::Local::now().timestamp_millis());
|
|
let _ = std::fs::write(&filename, &body);
|
|
}
|
|
|
|
let start = std::time::Instant::now();
|
|
let mut last_error: Option<anyhow::Error> = None;
|
|
|
|
for attempt in 0..self.max_retries {
|
|
let mut req = self.http.post(&url);
|
|
// Only add apikey header if it's not empty (for backward compatibility)
|
|
if !self.api_key.is_empty() && !self.api_key.starts_with("http") {
|
|
req = req.header("apikey", &self.api_key);
|
|
}
|
|
|
|
let response = req
|
|
.header("Content-Type", "application/json")
|
|
.body(body.clone())
|
|
.timeout(self.timeout)
|
|
.send()
|
|
.await;
|
|
|
|
let response = match response {
|
|
Ok(r) => r,
|
|
Err(e) => {
|
|
last_error = Some(anyhow!("Request failed: {}", e));
|
|
if e.is_timeout() || e.is_status() {
|
|
if attempt < self.max_retries - 1 {
|
|
tokio::time::sleep(Duration::from_millis(100 * 2_u64.pow(attempt))).await;
|
|
continue;
|
|
}
|
|
}
|
|
return Err(last_error.unwrap());
|
|
}
|
|
};
|
|
|
|
let status = response.status();
|
|
let body_text = response.text().await.unwrap_or_default();
|
|
|
|
// Record response if MEM_LLM_RECORD is set
|
|
if let Ok(record_dir) = env::var("MEM_LLM_RECORD") {
|
|
let filename = format!(
|
|
"{}/response-{}-{}.json",
|
|
record_dir,
|
|
chrono::Local::now().timestamp_millis(),
|
|
status
|
|
);
|
|
let _ = std::fs::write(&filename, &body_text);
|
|
}
|
|
|
|
// Handle auth error
|
|
if status == 401 {
|
|
return Err(anyhow!(
|
|
"Auth error (401): check apikey header format. Response: {}",
|
|
body_text
|
|
));
|
|
}
|
|
|
|
// 4xx errors should not be retried
|
|
if status.is_client_error() {
|
|
return Err(anyhow!("Client error ({}): {}", status, body_text));
|
|
}
|
|
|
|
// 5xx errors should be retried
|
|
if status.is_server_error() {
|
|
if attempt < self.max_retries - 1 {
|
|
last_error = Some(anyhow!("Server error ({}): {}", status, body_text));
|
|
tokio::time::sleep(Duration::from_millis(100 * 2_u64.pow(attempt))).await;
|
|
continue;
|
|
} else {
|
|
return Err(anyhow!("Server error ({}): {} (after {} retries)", status, body_text, self.max_retries));
|
|
}
|
|
}
|
|
|
|
// Parse success response
|
|
if status.is_success() {
|
|
let completion_response: CompletionResponse = serde_json::from_str(&body_text)?;
|
|
|
|
if completion_response.choices.is_empty() {
|
|
return Err(anyhow!("No choices in response"));
|
|
}
|
|
|
|
let latency_ms = start.elapsed().as_millis() as u64;
|
|
let text = completion_response.choices[0].message.content.clone();
|
|
let usage = Usage {
|
|
prompt_tokens: completion_response.usage.prompt_tokens,
|
|
completion_tokens: completion_response.usage.completion_tokens,
|
|
total_tokens: completion_response.usage.total_tokens,
|
|
};
|
|
|
|
return Ok(Completion {
|
|
text,
|
|
usage,
|
|
latency_ms,
|
|
});
|
|
}
|
|
|
|
return Err(anyhow!("Unexpected status {}: {}", status, body_text));
|
|
}
|
|
|
|
Err(last_error.unwrap_or_else(|| anyhow!("Max retries exhausted")))
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_completion_structs_serialize() {
|
|
let usage = Usage {
|
|
prompt_tokens: 100,
|
|
completion_tokens: 50,
|
|
total_tokens: 150,
|
|
};
|
|
let json = serde_json::to_string(&usage).unwrap();
|
|
assert!(json.contains("100"));
|
|
}
|
|
}
|