spectracost 0.1.0

//! Anthropic Messages API surface.

use std::time::Instant;

use serde::{Deserialize, Serialize};

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

/// A single message in the Anthropic messages array.
#[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 the Messages API.
#[derive(Debug, Clone, Serialize)]
pub struct MessagesRequest {
    pub model: String,
    pub max_tokens: u32,
    pub messages: Vec<Message>,
}

/// Response body from the Messages API.
#[derive(Debug, Clone, Deserialize)]
pub struct MessagesResponse {
    pub id: Option<String>,
    pub model: Option<String>,
    #[serde(default)]
    pub content: Vec<ContentBlock>,
    pub usage: Option<MessagesUsage>,
}

#[derive(Debug, Clone, Deserialize)]
pub struct ContentBlock {
    #[serde(rename = "type")]
    pub block_type: String,
    #[serde(default)]
    pub text: Option<String>,
}

#[derive(Debug, Clone, Deserialize)]
pub struct MessagesUsage {
    pub input_tokens: u32,
    pub output_tokens: u32,
    #[serde(default)]
    pub cache_read_input_tokens: u32,
}

impl MessagesResponse {
    /// First text block in the response, if any.
    pub fn first_text(&self) -> Option<String> {
        self.content
            .iter()
            .find(|b| b.block_type == "text")
            .and_then(|b| b.text.clone())
    }
}

impl Spectracost {
    /// Send a messages request to Anthropic and emit a telemetry event.
    pub async fn anthropic_messages(
        &self,
        request: MessagesRequest,
    ) -> Result<MessagesResponse, Error> {
        self.anthropic_messages_with(request, None).await
    }

    /// Like [`Spectracost::anthropic_messages`] but accepts per-call
    /// attribution overrides.
    pub async fn anthropic_messages_with(
        &self,
        request: MessagesRequest,
        attribution: Option<&Attribution>,
    ) -> Result<MessagesResponse, Error> {
        let opts = self.options();
        let api_key = opts
            .anthropic_api_key
            .as_deref()
            .ok_or(Error::MissingCredential("anthropic_api_key"))?;

        let base_url = opts
            .anthropic_base_url
            .clone()
            .unwrap_or_else(|| "https://api.anthropic.com".to_string());
        let url = format!("{}/v1/messages", base_url.trim_end_matches('/'));

        let start = Instant::now();
        let result = self
            .http()
            .post(&url)
            .header("x-api-key", api_key)
            .header("anthropic-version", "2023-06-01")
            .json(&request)
            .send()
            .await;
        let latency_ms = start.elapsed().as_millis() as u32;
        let provider = detect_provider(opts.anthropic_base_url.as_deref(), "anthropic");

        let resp = match result {
            Ok(r) => r,
            Err(err) => {
                self.emit(self.build_event(
                    &provider,
                    &request.model,
                    "messages",
                    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,
                "messages",
                0,
                0,
                latency_ms,
                "error",
                attribution,
            ));
            return Err(Error::Provider { status: status.as_u16(), body });
        }

        let parsed: MessagesResponse = resp.json().await?;
        let (input, output) = match &parsed.usage {
            Some(u) => (u.input_tokens, u.output_tokens),
            None => (0, 0),
        };
        let mut event = self.build_event(
            &provider,
            &request.model,
            "messages",
            input,
            output,
            latency_ms,
            "success",
            attribution,
        );
        if let Some(u) = &parsed.usage {
            event.cached_tokens = u.cache_read_input_tokens;
        }
        self.emit(event);
        Ok(parsed)
    }
}