audacity-sdk 0.2.0

//! Types mirroring the Amazon Bedrock Converse SDK surface.

use serde::{Deserialize, Serialize};
use serde_json::Value;

// ── Conversation roles ────────────────────────────────────────────────────────

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ConversationRole {
    User,
    Assistant,
}

// ── Content blocks ────────────────────────────────────────────────────────────

/// A tool-use invocation block (assistant side).
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ToolUseBlock {
    pub tool_use_id: String,
    pub name: String,
    pub input: Value,
}

/// A single item inside a toolResult content list.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ToolResultContent {
    Text(String),
    Json(Value),
}

/// A tool-result block (user side).
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ToolResultBlock {
    pub tool_use_id: String,
    pub content: Vec<ToolResultContent>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub status: Option<String>,
}

/// The format of an image content block.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ImageFormat {
    Png,
    Jpeg,
    Gif,
    Webp,
}

impl ImageFormat {
    /// The lowercase wire spelling, also used as the `image/<subtype>` media subtype.
    pub fn as_str(&self) -> &'static str {
        match self {
            ImageFormat::Png => "png",
            ImageFormat::Jpeg => "jpeg",
            ImageFormat::Gif => "gif",
            ImageFormat::Webp => "webp",
        }
    }
}

/// The source of an image content block.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ImageSource {
    /// Raw image bytes (Bedrock parity) — base64-encoded into a data URL on the wire.
    Bytes(Vec<u8>),
    /// An https or data URL, passed through verbatim (Audacity extension).
    Url(String),
}

/// An image block (user side).
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ImageBlock {
    pub format: ImageFormat,
    pub source: ImageSource,
}

/// The kind of prompt-cache breakpoint.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum CachePointType {
    #[default]
    Default,
}

/// A prompt-cache breakpoint (Bedrock parity). Placed in message or system
/// content after the prefix to cache: the content part immediately preceding
/// it gets an ephemeral `cache_control` marker on the wire; the block itself
/// is never emitted. A cache point with no preceding content part in the same
/// message is silently ignored.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct CachePointBlock {
    #[serde(rename = "type")]
    pub r#type: CachePointType,
}

impl CachePointBlock {
    /// The standard cache point (`{"type": "default"}`).
    pub fn new() -> Self {
        Self::default()
    }
}

/// A content block that can appear in a message.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub enum ContentBlock {
    Text(String),
    Image(ImageBlock),
    ToolUse(ToolUseBlock),
    ToolResult(ToolResultBlock),
    CachePoint(CachePointBlock),
}

impl ContentBlock {
    /// Returns `Ok(&str)` if this is a `Text` block.
    pub fn as_text(&self) -> Result<&str, &Self> {
        match self {
            ContentBlock::Text(s) => Ok(s.as_str()),
            _ => Err(self),
        }
    }

    /// Returns `Ok(&ImageBlock)` if this is an `Image` block.
    pub fn as_image(&self) -> Result<&ImageBlock, &Self> {
        match self {
            ContentBlock::Image(b) => Ok(b),
            _ => Err(self),
        }
    }

    /// Returns `Ok(&ToolUseBlock)` if this is a `ToolUse` block.
    pub fn as_tool_use(&self) -> Result<&ToolUseBlock, &Self> {
        match self {
            ContentBlock::ToolUse(b) => Ok(b),
            _ => Err(self),
        }
    }

    /// Returns `Ok(&ToolResultBlock)` if this is a `ToolResult` block.
    pub fn as_tool_result(&self) -> Result<&ToolResultBlock, &Self> {
        match self {
            ContentBlock::ToolResult(b) => Ok(b),
            _ => Err(self),
        }
    }

    /// Returns `Ok(&CachePointBlock)` if this is a `CachePoint` block.
    pub fn as_cache_point(&self) -> Result<&CachePointBlock, &Self> {
        match self {
            ContentBlock::CachePoint(b) => Ok(b),
            _ => Err(self),
        }
    }
}

// ── Message ───────────────────────────────────────────────────────────────────

/// A conversation message with a role and one or more content blocks.
#[derive(Debug, Clone)]
pub struct Message {
    pub role: ConversationRole,
    pub content: Vec<ContentBlock>,
}

impl Message {
    pub fn builder() -> MessageBuilder {
        MessageBuilder::default()
    }

    pub fn role(&self) -> &ConversationRole {
        &self.role
    }

    pub fn content(&self) -> &[ContentBlock] {
        &self.content
    }
}

#[derive(Default)]
pub struct MessageBuilder {
    role: Option<ConversationRole>,
    content: Vec<ContentBlock>,
}

impl MessageBuilder {
    pub fn role(mut self, role: ConversationRole) -> Self {
        self.role = Some(role);
        self
    }

    pub fn content(mut self, block: ContentBlock) -> Self {
        self.content.push(block);
        self
    }

    pub fn build(self) -> Result<Message, crate::Error> {
        let role = self
            .role
            .ok_or_else(|| crate::Error::client_validation("Message.role is required"))?;
        if self.content.is_empty() {
            return Err(crate::Error::client_validation(
                "Message must have at least one content block",
            ));
        }
        Ok(Message {
            role,
            content: self.content,
        })
    }
}

// ── System content ─────────────────────────────────────────────────────────────

/// A system-prompt content entry: either text or a prompt-cache breakpoint.
/// Construct with [`SystemContentBlock::text`] or
/// [`SystemContentBlock::cache_point`].
#[derive(Debug, Clone)]
pub struct SystemContentBlock {
    pub text: String,
    /// When `Some`, this entry is a prompt-cache breakpoint (`text` is ignored).
    pub cache_point: Option<CachePointBlock>,
}

impl SystemContentBlock {
    pub fn text(text: impl Into<String>) -> Self {
        Self {
            text: text.into(),
            cache_point: None,
        }
    }

    /// A prompt-cache breakpoint marking the end of the cacheable system prefix.
    pub fn cache_point() -> Self {
        Self {
            text: String::new(),
            cache_point: Some(CachePointBlock::new()),
        }
    }
}

// ── Inference configuration ───────────────────────────────────────────────────

#[derive(Debug, Clone, Default)]
pub struct InferenceConfiguration {
    pub max_tokens: Option<u32>,
    pub temperature: Option<f64>,
    pub top_p: Option<f64>,
    pub stop_sequences: Option<Vec<String>>,
}

impl InferenceConfiguration {
    pub fn builder() -> InferenceConfigurationBuilder {
        InferenceConfigurationBuilder::default()
    }
}

#[derive(Default)]
pub struct InferenceConfigurationBuilder {
    inner: InferenceConfiguration,
}

impl InferenceConfigurationBuilder {
    pub fn max_tokens(mut self, v: u32) -> Self {
        self.inner.max_tokens = Some(v);
        self
    }
    pub fn temperature(mut self, v: f64) -> Self {
        self.inner.temperature = Some(v);
        self
    }
    pub fn top_p(mut self, v: f64) -> Self {
        self.inner.top_p = Some(v);
        self
    }
    pub fn stop_sequences(mut self, v: Vec<String>) -> Self {
        self.inner.stop_sequences = Some(v);
        self
    }
    pub fn build(self) -> InferenceConfiguration {
        self.inner
    }
}

// ── Tool configuration ────────────────────────────────────────────────────────

#[derive(Debug, Clone)]
pub struct ToolInputSchema {
    pub json: Value,
}

#[derive(Debug, Clone)]
pub struct ToolSpecification {
    pub name: String,
    pub description: Option<String>,
    pub input_schema: ToolInputSchema,
}

#[derive(Debug, Clone)]
pub struct Tool {
    pub tool_spec: ToolSpecification,
}

#[derive(Debug, Clone)]
pub enum ToolChoice {
    Auto,
    Any,
    Tool { name: String },
}

#[derive(Debug, Clone)]
pub struct ToolConfiguration {
    pub tools: Vec<Tool>,
    pub tool_choice: Option<ToolChoice>,
}

// ── Converse output ───────────────────────────────────────────────────────────

/// Output wrapper — mirrors Bedrock's `ConverseOutput` enum.
#[derive(Debug, Clone)]
pub enum ConverseOutputEnum {
    Message(Message),
}

impl ConverseOutputEnum {
    /// Returns `Ok(&Message)` if this is the `Message` variant.
    pub fn as_message(&self) -> Result<&Message, &Self> {
        match self {
            ConverseOutputEnum::Message(m) => Ok(m),
            // future variants would fall through here
        }
    }
}

#[derive(Debug, Clone)]
pub struct TokenUsage {
    pub input_tokens: u32,
    pub output_tokens: u32,
    pub total_tokens: u32,
    /// Prompt tokens served from the provider's prompt cache (Bedrock name).
    pub cache_read_input_tokens: u32,
    /// Prompt tokens written to the provider's prompt cache (Bedrock name).
    pub cache_write_input_tokens: u32,
}

#[derive(Debug, Clone)]
pub struct Metrics {
    pub latency_ms: u64,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum StopReason {
    EndTurn,
    MaxTokens,
    ToolUse,
    StopSequence,
    ContentFiltered,
}

impl StopReason {
    pub(crate) fn from_finish_reason(s: &str) -> Self {
        match s {
            "stop" => StopReason::EndTurn,
            "length" => StopReason::MaxTokens,
            "tool_calls" | "function_call" => StopReason::ToolUse,
            "content_filter" => StopReason::ContentFiltered,
            _ => StopReason::EndTurn,
        }
    }
}

/// The top-level response from a `converse()` call.
#[derive(Debug)]
pub struct ConverseOutput {
    pub(crate) output: Option<ConverseOutputEnum>,
    pub(crate) stop_reason: StopReason,
    pub(crate) usage: TokenUsage,
    pub(crate) metrics: Metrics,
}

impl ConverseOutput {
    pub fn output(&self) -> Option<&ConverseOutputEnum> {
        self.output.as_ref()
    }
    pub fn stop_reason(&self) -> &StopReason {
        &self.stop_reason
    }
    pub fn usage(&self) -> &TokenUsage {
        &self.usage
    }
    pub fn metrics(&self) -> &Metrics {
        &self.metrics
    }
}

// ── Stream event types ────────────────────────────────────────────────────────

#[derive(Debug, Clone)]
pub struct MessageStartEvent {
    pub role: ConversationRole,
}

#[derive(Debug, Clone)]
pub struct ContentBlockStartToolUse {
    pub tool_use_id: String,
    pub name: String,
}

#[derive(Debug, Clone)]
pub enum ContentBlockStartPayload {
    ToolUse(ContentBlockStartToolUse),
}

#[derive(Debug, Clone)]
pub struct ContentBlockStartEvent {
    pub content_block_index: u32,
    pub start: ContentBlockStartPayload,
}

#[derive(Debug, Clone)]
pub enum ContentBlockDeltaPayload {
    Text(String),
    ToolUse { input: String },
}

#[derive(Debug, Clone)]
pub struct ContentBlockDeltaEvent {
    pub content_block_index: u32,
    pub delta: ContentBlockDeltaPayload,
}

#[derive(Debug, Clone)]
pub struct ContentBlockStopEvent {
    pub content_block_index: u32,
}

#[derive(Debug, Clone)]
pub struct MessageStopEvent {
    pub stop_reason: StopReason,
}

#[derive(Debug, Clone)]
pub struct MetadataEvent {
    pub usage: TokenUsage,
    pub metrics: Metrics,
}

/// All possible stream events — one variant per Bedrock stream event type.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub enum ConverseStreamOutput {
    MessageStart(MessageStartEvent),
    ContentBlockStart(ContentBlockStartEvent),
    ContentBlockDelta(ContentBlockDeltaEvent),
    ContentBlockStop(ContentBlockStopEvent),
    MessageStop(MessageStopEvent),
    Metadata(MetadataEvent),
}