96 lines
2.8 KiB
Rust
96 lines
2.8 KiB
Rust
//! Minimal OpenAI-compatible chat completions client.
|
|
//! Used for the rewrite step.
|
|
|
|
use std::time::Duration;
|
|
|
|
use serde::{Deserialize, Serialize};
|
|
use serde_json::{json, Value};
|
|
use thiserror::Error;
|
|
|
|
#[derive(Debug, Error)]
|
|
pub enum ChatError {
|
|
#[error("missing api url or key")]
|
|
MissingConfig,
|
|
#[error("http: {0}")]
|
|
Http(#[from] reqwest::Error),
|
|
#[error("api: status={status}, body={body}")]
|
|
Api { status: u16, body: String },
|
|
#[error("api response had no choices")]
|
|
NoChoices,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize)]
|
|
pub struct ChatMessage {
|
|
pub role: String,
|
|
pub content: String,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize)]
|
|
pub struct ChatRequest {
|
|
#[serde(rename = "apiUrl")]
|
|
pub api_url: String,
|
|
#[serde(rename = "apiKey")]
|
|
pub api_key: String,
|
|
pub model: String,
|
|
pub messages: Vec<ChatMessage>,
|
|
#[serde(default)]
|
|
pub temperature: Option<f32>,
|
|
#[serde(default, rename = "max_tokens")]
|
|
pub max_tokens: Option<u32>,
|
|
}
|
|
|
|
#[derive(Debug, Serialize)]
|
|
pub struct ChatResponse {
|
|
pub success: bool,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub content: Option<String>,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub error: Option<String>,
|
|
}
|
|
|
|
pub async fn complete(req: &ChatRequest) -> Result<ChatResponse, ChatError> {
|
|
if req.api_url.trim().is_empty() || req.api_key.trim().is_empty() {
|
|
return Err(ChatError::MissingConfig);
|
|
}
|
|
// Accept both `https://host` (will append /v1/chat/completions) and the
|
|
// full URL.
|
|
let url = if req.api_url.contains("/chat/completions") {
|
|
req.api_url.trim_end_matches('/').to_string()
|
|
} else {
|
|
format!("{}/v1/chat/completions", req.api_url.trim_end_matches('/'))
|
|
};
|
|
|
|
let mut body = json!({
|
|
"model": req.model,
|
|
"messages": req.messages.iter().map(|m| json!({
|
|
"role": m.role,
|
|
"content": m.content,
|
|
})).collect::<Vec<_>>(),
|
|
"stream": false,
|
|
});
|
|
if let Some(t) = req.temperature { body["temperature"] = json!(t); }
|
|
if let Some(mt) = req.max_tokens { body["max_tokens"] = json!(mt); }
|
|
|
|
let client = reqwest::Client::builder()
|
|
.timeout(Duration::from_secs(120))
|
|
.build()?;
|
|
let resp = client
|
|
.post(&url)
|
|
.bearer_auth(&req.api_key)
|
|
.json(&body)
|
|
.send()
|
|
.await?;
|
|
let status = resp.status().as_u16();
|
|
if !(200..300).contains(&status) {
|
|
let txt = resp.text().await.unwrap_or_default();
|
|
return Err(ChatError::Api { status, body: txt });
|
|
}
|
|
let v: Value = resp.json().await?;
|
|
let content = v
|
|
.pointer("/choices/0/message/content")
|
|
.and_then(Value::as_str)
|
|
.ok_or(ChatError::NoChoices)?
|
|
.to_string();
|
|
Ok(ChatResponse { success: true, content: Some(content), error: None })
|
|
}
|