openrouter-batch 0.1.2

//! Typed facades over the transport. The engine loop never touches wire
//! types: it submits structured completion batches by item list, and
//! decodes stored embedding payloads into input-ordered vectors.

use super::transport::{
    BatchEndpoint, BatchItem, BatchTransport, EmbeddingResponseBody, SubmitBatch, TransportError,
};
use serde::{Deserialize, Serialize};
use serde_json::Value;

/// One chat-completions message.
#[derive(Debug, Clone, Serialize)]
pub struct Message {
    pub role: String,
    pub content: String,
}

impl Message {
    pub fn system(content: impl Into<String>) -> Self {
        Self {
            role: "system".to_string(),
            content: content.into(),
        }
    }

    pub fn user(content: impl Into<String>) -> Self {
        Self {
            role: "user".to_string(),
            content: content.into(),
        }
    }
}

/// Parameters shared by every request in one structured batch. One schema per
/// batch is mandatory on Google-backed routes (quickstart), so the schema
/// lives here, at batch level, not per item.
#[derive(Debug, Clone)]
pub struct StructuredParams {
    /// The strict JSON Schema for `response_format`.
    pub schema: serde_json::Value,
    /// The `response_format` schema name.
    pub schema_name: String,
    pub max_tokens: u32,
    /// Emit `"reasoning": {"enabled": false}` in every body. Reasoner models
    /// otherwise burn the token budget on thinking and answer with
    /// `content: null` and `finish_reason: "length"` (verified gotcha).
    pub disable_reasoning: bool,
    /// Emit `"provider": {"require_parameters": true}` in every body so
    /// OpenRouter routes only to providers that honor every request
    /// parameter, including `response_format` strict mode. Opt-in: the eval
    /// harness uses it to pin strict-mode behavior; production leaves it off.
    pub require_parameters: bool,
    /// Cap the model's reasoning tokens via `"reasoning": {"max_tokens": N}`,
    /// leaving the rest of `max_tokens` for the JSON body. Counter to the
    /// measured reasoning squeeze where reasoner slugs burn the whole budget
    /// before emitting (isolated replay V1-cap512: wire-verified).
    pub reasoning_max_tokens: Option<u32>,
}

/// One structured chat-completions request body: the exact body the Batch
/// API wraps per item and the synchronous chat-completions endpoint accepts.
/// Shared by both paths so their wire shapes cannot drift.
pub fn structured_body(
    params: &StructuredParams,
    messages: Vec<Message>,
) -> Result<Value, serde_json::Error> {
    let mut schema = serde_json::Map::new();
    schema.insert("name".to_string(), params.schema_name.clone().into());
    schema.insert("strict".to_string(), true.into());
    schema.insert("schema".to_string(), params.schema.clone());

    let mut response_format = serde_json::Map::new();
    response_format.insert("type".to_string(), "json_schema".into());
    response_format.insert("json_schema".to_string(), Value::Object(schema));

    let mut body = serde_json::Map::new();
    body.insert("messages".to_string(), serde_json::to_value(messages)?);
    body.insert("max_tokens".to_string(), params.max_tokens.into());
    body.insert(
        "response_format".to_string(),
        Value::Object(response_format),
    );
    if params.disable_reasoning {
        body.insert(
            "reasoning".to_string(),
            serde_json::json!({"enabled": false}),
        );
    }
    if params.require_parameters {
        body.insert(
            "provider".to_string(),
            serde_json::json!({"require_parameters": true}),
        );
    }
    if let Some(cap) = params.reasoning_max_tokens {
        body.insert(
            "reasoning".to_string(),
            serde_json::json!({"max_tokens": cap}),
        );
    }
    Ok(Value::Object(body))
}

/// Submit one structured chat-completions batch: every item carries the same
/// `response_format` json_schema with `strict: true`.
///
/// `items` pairs each deterministic `custom_id` with its prompt messages.
/// The returned batch object is `202 Accepted`-equivalent: `validating`,
/// queued — never treated as complete.
pub async fn submit_structured(
    transport: &dyn BatchTransport,
    model: &str,
    idempotency_key: &str,
    params: &StructuredParams,
    items: Vec<(String, Vec<Message>)>,
) -> Result<super::transport::BatchObject, TransportError> {
    // The provider rejects the whole batch with 422 when any request key
    // exceeds 64 characters (observed live 2026-08-30). Fail locally with
    // the offending key instead of a remote round-trip.
    for (custom_id, _) in &items {
        if custom_id.len() > 64 {
            return Err(TransportError::CustomIdTooLong {
                custom_id: custom_id.clone(),
            });
        }
    }

    let mut schema = serde_json::Map::new();
    schema.insert("name".to_string(), params.schema_name.clone().into());
    schema.insert("strict".to_string(), true.into());
    schema.insert("schema".to_string(), params.schema.clone());

    let mut response_format = serde_json::Map::new();
    response_format.insert("type".to_string(), "json_schema".into());
    response_format.insert("json_schema".to_string(), Value::Object(schema));

    let requests = items
        .into_iter()
        .map(|(custom_id, messages)| {
            Ok(BatchItem {
                custom_id,
                body: structured_body(params, messages)?,
            })
        })
        .collect::<Result<Vec<_>, serde_json::Error>>()?;

    transport
        .submit(
            &SubmitBatch::new(BatchEndpoint::ChatCompletions, model, requests),
            idempotency_key,
        )
        .await
}

/// A decode failure of one stored embedding payload. The `Display` strings
/// are the durable `last_error` texts operators see in Postgres.
#[derive(Debug, thiserror::Error)]
pub enum EmbeddingDecodeError {
    #[error("embedding response decode: {source}")]
    Envelope {
        #[from]
        source: serde_json::Error,
    },
    #[error("embedding response count: expected {expected}, received {actual}")]
    CountMismatch { expected: usize, actual: usize },
    #[error("embedding response index {0} is negative")]
    NegativeIndex(i64),
    #[error("embedding response index {0} exceeds input count {1}")]
    IndexOutOfRange(usize, usize),
    #[error("embedding response has duplicate index {0}")]
    DuplicateIndex(usize),
    #[error("embedding response has no index {0}")]
    MissingIndex(usize),
}

/// Decode one stored raw embedding payload and return the vectors ordered
/// to match the request inputs: `vectors[i]` belongs to input `i`. Validates
/// the envelope, the count, the index range, duplicates, and completeness.
///
/// (`raw_provider_results.raw_line`), so a crash between submit and
/// finalize never needs the provider again.
pub fn embedding_vectors(
    raw: &serde_json::Value,
    expected: usize,
) -> Result<Vec<Vec<f32>>, EmbeddingDecodeError> {
    let response = EmbeddingResponseBody::deserialize(raw)?;
    if response.data.len() != expected {
        return Err(EmbeddingDecodeError::CountMismatch {
            expected,
            actual: response.data.len(),
        });
    }
    let mut indexed: Vec<Option<Vec<f32>>> = vec![None; expected];
    for data in response.data {
        let index = usize::try_from(data.index)
            .map_err(|_| EmbeddingDecodeError::NegativeIndex(data.index))?;
        if index >= expected {
            return Err(EmbeddingDecodeError::IndexOutOfRange(index, expected));
        }
        if indexed[index].replace(data.embedding).is_some() {
            return Err(EmbeddingDecodeError::DuplicateIndex(index));
        }
    }
    indexed
        .into_iter()
        .enumerate()
        .map(|(index, vector)| vector.ok_or(EmbeddingDecodeError::MissingIndex(index)))
        .collect()
}

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

    #[test]
    fn structured_body_caps_reasoning_tokens_only_when_set() {
        let base = StructuredParams {
            schema: serde_json::json!({"type": "object"}),
            schema_name: "reply".to_string(),
            max_tokens: 100,
            disable_reasoning: false,
            require_parameters: false,
            reasoning_max_tokens: Some(1024),
        };
        let body = structured_body(&base, vec![Message::user("hi")]).unwrap();
        assert_eq!(
            body.get("reasoning"),
            Some(&serde_json::json!({"max_tokens": 1024}))
        );

        let off = StructuredParams {
            reasoning_max_tokens: None,
            ..base
        };
        let body = structured_body(&off, vec![Message::user("hi")]).unwrap();
        assert!(body.get("reasoning").is_none());
    }

    #[test]
    fn structured_body_requires_provider_parameters_only_when_opted_in() {
        let base = StructuredParams {
            schema: serde_json::json!({"type": "object"}),
            schema_name: "reply".to_string(),
            max_tokens: 100,
            disable_reasoning: false,
            require_parameters: true,
            reasoning_max_tokens: None,
        };
        let body = structured_body(&base, vec![Message::user("hi")]).unwrap();
        assert_eq!(
            body.get("provider"),
            Some(&serde_json::json!({"require_parameters": true}))
        );

        let off = StructuredParams {
            require_parameters: false,
            ..base
        };
        let body = structured_body(&off, vec![Message::user("hi")]).unwrap();
        assert!(body.get("provider").is_none());
    }

    #[test]
    fn vectors_are_reordered_to_match_input_indexes() {
        let raw = serde_json::json!({
            "id": "embed-request",
            "data": [
                {"index": 1, "embedding": [2.0]},
                {"index": 0, "embedding": [1.0]},
            ]
        });
        let vectors = embedding_vectors(&raw, 2).unwrap();
        assert_eq!(vectors, vec![vec![1.0], vec![2.0]]);
    }

    #[test]
    fn vector_decode_rejects_count_and_index_defects() {
        let short = serde_json::json!({"data": [{"index": 0, "embedding": [1.0]}]});
        assert!(matches!(
            embedding_vectors(&short, 2),
            Err(EmbeddingDecodeError::CountMismatch {
                expected: 2,
                actual: 1
            })
        ));
        let duplicate = serde_json::json!({
            "data": [
                {"index": 0, "embedding": [1.0]},
                {"index": 0, "embedding": [2.0]},
            ]
        });
        assert!(matches!(
            embedding_vectors(&duplicate, 2),
            Err(EmbeddingDecodeError::DuplicateIndex(0))
        ));
    }
}