spectracost 0.1.0

AI cost observability SDK - see the full spectrum of your AI spend
Documentation
//! OpenAI-compatible chat completion surface.
//!
//! A minimal request/response surface covering the fields most apps
//! actually use. For power users that need the full API, bring your own
//! HTTP client — this module is the batteries-included path.

use std::time::Instant;

use serde::{Deserialize, Serialize};

use crate::{detect_provider, Attribution, Error, Spectracost};

/// A single chat message.
#[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(),
        }
    }
}

/// Request body for a chat completion.
#[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>,
}

/// Response body from a chat completion.
#[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 {
    /// Convenience: the first assistant message, if any.
    pub fn first_message(&self) -> Option<String> {
        self.choices.first().map(|c| c.message.content.clone())
    }
}

impl Spectracost {
    /// Send a chat completion to OpenAI (or any OpenAI-compatible endpoint
    /// configured via [`crate::Options::openai_base_url`]) and emit a
    /// telemetry event.
    pub async fn openai_chat(&self, request: ChatRequest) -> Result<ChatResponse, Error> {
        self.openai_chat_with(request, None).await
    }

    /// Like [`Spectracost::openai_chat`] but accepts per-call attribution
    /// overrides for the emitted event.
    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)
    }
}