use std::time::Instant;
use serde::{Deserialize, Serialize};
use crate::{detect_provider, Attribution, Error, Spectracost};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
pub role: String,
pub content: String,
}
impl From<(String, String)> for Message {
fn from((role, content): (String, String)) -> Self {
Self { role, content }
}
}
impl From<(&str, &str)> for Message {
fn from((role, content): (&str, &str)) -> Self {
Self {
role: role.to_string(),
content: content.to_string(),
}
}
}
#[derive(Debug, Clone, Serialize)]
pub struct ChatRequest {
pub model: String,
pub messages: Vec<Message>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_tokens: Option<u32>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct ChatResponse {
pub id: Option<String>,
pub model: Option<String>,
pub choices: Vec<ChatChoice>,
pub usage: Option<ChatUsage>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct ChatChoice {
pub index: Option<u32>,
pub message: Message,
pub finish_reason: Option<String>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct ChatUsage {
pub prompt_tokens: u32,
pub completion_tokens: u32,
pub total_tokens: u32,
}
impl ChatResponse {
pub fn first_message(&self) -> Option<String> {
self.choices.first().map(|c| c.message.content.clone())
}
}
impl Spectracost {
pub async fn openai_chat(&self, request: ChatRequest) -> Result<ChatResponse, Error> {
self.openai_chat_with(request, None).await
}
pub async fn openai_chat_with(
&self,
request: ChatRequest,
attribution: Option<&Attribution>,
) -> Result<ChatResponse, Error> {
let opts = self.options();
let api_key = opts
.openai_api_key
.as_deref()
.ok_or(Error::MissingCredential("openai_api_key"))?;
let base_url = opts
.openai_base_url
.clone()
.unwrap_or_else(|| "https://api.openai.com".to_string());
let url = format!("{}/v1/chat/completions", base_url.trim_end_matches('/'));
let start = Instant::now();
let result = self
.http()
.post(&url)
.bearer_auth(api_key)
.json(&request)
.send()
.await;
let latency_ms = start.elapsed().as_millis() as u32;
let provider = detect_provider(opts.openai_base_url.as_deref(), "openai");
let resp = match result {
Ok(r) => r,
Err(err) => {
self.emit(self.build_event(
&provider,
&request.model,
"chat.completions",
0,
0,
latency_ms,
"error",
attribution,
));
return Err(Error::Http(err));
}
};
let status = resp.status();
if !status.is_success() {
let body = resp.text().await.unwrap_or_default();
self.emit(self.build_event(
&provider,
&request.model,
"chat.completions",
0,
0,
latency_ms,
"error",
attribution,
));
return Err(Error::Provider { status: status.as_u16(), body });
}
let parsed: ChatResponse = resp.json().await?;
let (input, output) = match &parsed.usage {
Some(u) => (u.prompt_tokens, u.completion_tokens),
None => (0, 0),
};
self.emit(self.build_event(
&provider,
&request.model,
"chat.completions",
input,
output,
latency_ms,
"success",
attribution,
));
Ok(parsed)
}
}