audacity-sdk 0.5.0

//! Provider-native pass-through surfaces (spec §9).
//!
//! Exposes the gateway's OpenAI-compatible and Anthropic-compatible endpoints
//! directly on [`Client`], with the SDK's auth, retry, and §4 error handling
//! but **no shape translation** — requests are sent verbatim and responses
//! come back as raw [`serde_json::Value`]s in the provider's own wire format:
//!
//! ```text
//! client.chat_completions().create(params)        → POST /v1/chat/completions
//! client.chat_completions().create_stream(params) → POST /v1/chat/completions (SSE)
//! client.messages().create(params)                → POST /v1/messages
//! client.messages().create_stream(params)         → POST /v1/messages (SSE)
//! client.messages().count_tokens(params)          → POST /v1/messages/count_tokens
//! ```
//!
//! Params are anything that serializes to a JSON object (`serde_json::json!`
//! literals or your own `Serialize` structs), so any field the gateway
//! supports works — unknown fields flow through with no SDK release needed.

use futures_util::StreamExt;
use serde::Serialize;
use serde_json::{Map, Value};
use tokio::sync::mpsc;

use crate::error::parse_stream_error;
use crate::sse::{SseParser, SseToken};
use crate::{Client, Error};

const OPENAI_PATH: &str = "/v1/chat/completions";
const MESSAGES_PATH: &str = "/v1/messages";
const COUNT_TOKENS_PATH: &str = "/v1/messages/count_tokens";

/// Anthropic API version header sent on the Anthropic-format endpoints. The
/// gateway forwards Anthropic-format bodies as-is; the header keeps parity
/// with the official anthropic SDKs.
const ANTHROPIC_HEADERS: &[(&str, &str)] = &[("anthropic-version", "2023-06-01")];

// ── param preparation ─────────────────────────────────────────────────────────

/// Serialize params to a JSON object and check the two fields every §9
/// endpoint requires (`model`, `messages`) are present — everything else is
/// passed through untouched.
fn prepare(params: impl Serialize, endpoint: &str) -> Result<Map<String, Value>, Error> {
    let value = serde_json::to_value(params)
        .map_err(|e| Error::sdk(format!("failed to serialize {endpoint} params: {e}")))?;
    let body = match value {
        Value::Object(map) => map,
        _ => {
            return Err(Error::client_validation(format!(
                "{endpoint} params must serialize to a JSON object"
            )))
        }
    };
    for field in ["model", "messages"] {
        if matches!(body.get(field), None | Some(Value::Null)) {
            return Err(Error::client_validation(format!(
                "{field} is required for {endpoint}"
            )));
        }
    }
    Ok(body)
}

/// `create` is the non-streaming call; sending `stream: true` through it
/// would get an SSE body the JSON decoder can't handle, so fail fast with a
/// pointer to `create_stream`.
fn reject_stream_flag(body: &Map<String, Value>, endpoint: &str) -> Result<(), Error> {
    if body.get("stream") == Some(&Value::Bool(true)) {
        return Err(Error::client_validation(format!(
            "params set stream=true — use create_stream() for streaming {endpoint} requests"
        )));
    }
    Ok(())
}

// ── OpenAI-format surface ─────────────────────────────────────────────────────

/// `client.chat_completions()` — OpenAI-format pass-through (spec §9).
pub struct ChatCompletions {
    client: Client,
}

impl ChatCompletions {
    pub(crate) fn new(client: Client) -> Self {
        Self { client }
    }

    /// OpenAI-format chat completion (POST /v1/chat/completions).
    ///
    /// Accepts the standard OpenAI Chat Completions request fields (`model`,
    /// `messages`, `max_tokens`, `tools`, …) and passes them through verbatim.
    /// Returns the raw OpenAI-shaped response object. Errors map to the
    /// standard [`Error`] taxonomy (spec §4); retries follow the standard
    /// policy.
    pub async fn create(&self, params: impl Serialize) -> Result<Value, Error> {
        let body = prepare(params, OPENAI_PATH)?;
        reject_stream_flag(&body, OPENAI_PATH)?;
        self.client
            .post_passthrough_json(OPENAI_PATH, &Value::Object(body), &[])
            .await
    }

    /// Streaming OpenAI-format chat completion.
    ///
    /// Sets `stream: true` on the body (all other fields pass through
    /// verbatim) and returns a [`RawEventStream`] of raw chunk objects. The
    /// gateway's `data: [DONE]` sentinel ends the stream and is not yielded;
    /// EOF without `[DONE]` surfaces [`Error::ModelStreamError`]; a chunk
    /// carrying `error` aborts with the §4-mapped error. Retries stop at the
    /// first streamed byte.
    pub async fn create_stream(&self, params: impl Serialize) -> Result<RawEventStream, Error> {
        let mut body = prepare(params, OPENAI_PATH)?;
        body.insert("stream".into(), Value::Bool(true));
        let resp = self
            .client
            .post_passthrough_stream(OPENAI_PATH, &Value::Object(body), &[])
            .await?;
        Ok(spawn_raw_event_driver(resp, RawStreamKind::OpenAi))
    }
}

// ── Anthropic-format surface ──────────────────────────────────────────────────

/// `client.messages()` — Anthropic-format pass-through (spec §9).
pub struct Messages {
    client: Client,
}

impl Messages {
    pub(crate) fn new(client: Client) -> Self {
        Self { client }
    }

    /// Anthropic-format message (POST /v1/messages) — the wire format used by
    /// the anthropic SDKs and Claude Code.
    ///
    /// Accepts the standard Anthropic Messages request fields (`model`,
    /// `max_tokens`, `messages`, `system`, `tools`, …) verbatim and returns
    /// the raw Anthropic-shaped response object. Works with every gateway
    /// model, not just Claude — the gateway bridges the wire format.
    pub async fn create(&self, params: impl Serialize) -> Result<Value, Error> {
        let body = prepare(params, MESSAGES_PATH)?;
        reject_stream_flag(&body, MESSAGES_PATH)?;
        self.client
            .post_passthrough_json(MESSAGES_PATH, &Value::Object(body), ANTHROPIC_HEADERS)
            .await
    }

    /// Streaming Anthropic-format message.
    ///
    /// Sets `stream: true` on the body (all other fields pass through
    /// verbatim) and returns a [`RawEventStream`] of raw Anthropic stream
    /// events (`message_start` … `message_stop`). There is no `[DONE]` — a
    /// healthy stream ends with a `message_stop` event followed by EOF; EOF
    /// before `message_stop` surfaces [`Error::ModelStreamError`]; an `error`
    /// event aborts with the §4-mapped error.
    pub async fn create_stream(&self, params: impl Serialize) -> Result<RawEventStream, Error> {
        let mut body = prepare(params, MESSAGES_PATH)?;
        body.insert("stream".into(), Value::Bool(true));
        let resp = self
            .client
            .post_passthrough_stream(MESSAGES_PATH, &Value::Object(body), ANTHROPIC_HEADERS)
            .await?;
        Ok(spawn_raw_event_driver(resp, RawStreamKind::Anthropic))
    }

    /// Anthropic-format token counting (POST /v1/messages/count_tokens).
    ///
    /// Returns the raw response object (`{"input_tokens": …}`). Token
    /// counting is free — no inference happens.
    pub async fn count_tokens(&self, params: impl Serialize) -> Result<Value, Error> {
        let body = prepare(params, COUNT_TOKENS_PATH)?;
        self.client
            .post_passthrough_json(COUNT_TOKENS_PATH, &Value::Object(body), ANTHROPIC_HEADERS)
            .await
    }
}

// ── raw event stream ──────────────────────────────────────────────────────────

/// Termination semantics of a raw pass-through SSE stream.
#[derive(Clone, Copy, PartialEq, Eq)]
enum RawStreamKind {
    /// Ends on `data: [DONE]`; EOF without it is an error.
    OpenAi,
    /// Ends on a `message_stop` event + EOF; EOF before it is an error.
    Anthropic,
}

/// Stream of raw provider events returned by the §9 `create_stream` calls.
///
/// Mirrors [`StreamReceiver`](crate::StreamReceiver):
/// ```rust,ignore
/// while let Some(event) = stream.recv().await? { … }
/// ```
pub struct RawEventStream {
    rx: mpsc::Receiver<Result<Value, Error>>,
}

impl RawEventStream {
    /// Receive the next raw provider event.
    /// Returns `Ok(None)` when the stream has ended normally.
    /// Returns `Err(e)` on stream errors.
    pub async fn recv(&mut self) -> Result<Option<Value>, Error> {
        match self.rx.recv().await {
            Some(Ok(event)) => Ok(Some(event)),
            Some(Err(e)) => Err(e),
            None => Ok(None),
        }
    }
}

/// Drive the SSE byte stream in a background task, forwarding raw provider
/// events under the §1 transport rules (32 MiB line cap, mid-stream transport
/// failures → ModelStreamError) with per-format termination handling.
fn spawn_raw_event_driver(response: reqwest::Response, kind: RawStreamKind) -> RawEventStream {
    let (tx, rx) = mpsc::channel::<Result<Value, Error>>(64);

    tokio::spawn(async move {
        let mut byte_stream = response.bytes_stream();
        let mut parser = SseParser::new();

        let mut done = false; // OpenAI: saw the [DONE] sentinel
        let mut saw_message_stop = false; // Anthropic: saw a message_stop event
        let mut aborted = false; // already surfaced an error / receiver gone
        let mut eof = false; // byte stream ended

        'read: while !done && !eof {
            match byte_stream.next().await {
                Some(Ok(chunk)) => {
                    parser.push(&chunk);
                    if parser.overflowed() {
                        let _ = tx
                            .send(Err(Error::model_stream_error(
                                "SSE line exceeds the 32 MiB limit",
                            )))
                            .await;
                        aborted = true;
                        break 'read;
                    }
                }
                Some(Err(e)) => {
                    let _ = tx
                        .send(Err(Error::model_stream_error(format!(
                            "stream read error: {e}"
                        ))))
                        .await;
                    aborted = true;
                    break 'read;
                }
                None => {
                    parser.flush();
                    eof = true;
                }
            }

            for token in parser.drain() {
                let payload = match token {
                    SseToken::Done => match kind {
                        RawStreamKind::OpenAi => {
                            done = true;
                            break;
                        }
                        // Anthropic streams carry no [DONE]; a stray one is
                        // not a valid event — skip it.
                        RawStreamKind::Anthropic => continue,
                    },
                    SseToken::Data(payload) => payload,
                };

                // Malformed / non-object payloads are skipped (bad wire data).
                let Ok(event) = serde_json::from_str::<Value>(&payload) else {
                    continue;
                };
                if !event.is_object() {
                    continue;
                }

                let is_error = match kind {
                    RawStreamKind::OpenAi => event.get("error").is_some(),
                    RawStreamKind::Anthropic => {
                        event.get("type").and_then(Value::as_str) == Some("error")
                    }
                };
                if is_error {
                    let _ = tx.send(Err(parse_stream_error(&payload))).await;
                    aborted = true;
                    break 'read;
                }

                if kind == RawStreamKind::Anthropic
                    && event.get("type").and_then(Value::as_str) == Some("message_stop")
                {
                    saw_message_stop = true;
                }

                if tx.send(Ok(event)).await.is_err() {
                    aborted = true;
                    break 'read;
                }
            }
        }

        // Terminal rule (spec §9): OpenAI streams must end with [DONE];
        // Anthropic streams must have seen message_stop before EOF.
        let healthy = match kind {
            RawStreamKind::OpenAi => done,
            RawStreamKind::Anthropic => saw_message_stop,
        };
        if !healthy && !aborted {
            let message = match kind {
                RawStreamKind::OpenAi => "connection closed before [DONE]",
                RawStreamKind::Anthropic => "stream ended unexpectedly before message_stop",
            };
            let _ = tx.send(Err(Error::model_stream_error(message))).await;
        }
    });

    RawEventStream { rx }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    #[test]
    fn prepare_accepts_object_with_required_fields() {
        let body = prepare(
            json!({"model": "m", "messages": [], "custom_field": {"x": 1}}),
            OPENAI_PATH,
        )
        .unwrap();
        // unknown fields survive untouched
        assert_eq!(body["custom_field"], json!({"x": 1}));
    }

    #[test]
    fn prepare_rejects_non_object_params() {
        let err = prepare(json!(["not", "an", "object"]), OPENAI_PATH).unwrap_err();
        assert!(matches!(err, Error::Validation(_)));
    }

    #[test]
    fn prepare_rejects_missing_model() {
        let err = prepare(json!({"messages": []}), MESSAGES_PATH).unwrap_err();
        match err {
            Error::Validation(d) => assert!(d.message.contains("model")),
            other => panic!("expected Validation, got {other:?}"),
        }
    }

    #[test]
    fn prepare_rejects_null_messages() {
        let err = prepare(json!({"model": "m", "messages": null}), OPENAI_PATH).unwrap_err();
        match err {
            Error::Validation(d) => assert!(d.message.contains("messages")),
            other => panic!("expected Validation, got {other:?}"),
        }
    }

    #[test]
    fn reject_stream_flag_only_fires_on_true() {
        let on = prepare(
            json!({"model": "m", "messages": [], "stream": true}),
            OPENAI_PATH,
        )
        .unwrap();
        assert!(reject_stream_flag(&on, OPENAI_PATH).is_err());
        let off = prepare(
            json!({"model": "m", "messages": [], "stream": false}),
            OPENAI_PATH,
        )
        .unwrap();
        assert!(reject_stream_flag(&off, OPENAI_PATH).is_ok());
    }
}