llmtrim 0.1.4

//! Neutral request/response model.
//!
//! JSON is the source of truth: a [`Request`] wraps the original provider body as
//! a [`serde_json::Value`], and stage mutations are applied at JSON-pointer
//! addresses. Any field a stage does not touch is reproduced byte-for-byte on
//! serialization — lossless passthrough by construction, which is stronger than a
//! typed struct with a `#[serde(flatten)]` catch-all (no field can be silently
//! re-typed or dropped).

use std::str::FromStr;

use anyhow::{Context, Result, bail};
use serde_json::Value;

/// The closed LLM provider a request targets. llmtrim is API-only and
/// provider-agnostic; each kind has its own wire shape and tokenizer.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProviderKind {
    OpenAi,
    Anthropic,
    Google,
}

impl ProviderKind {
    pub fn as_str(self) -> &'static str {
        match self {
            ProviderKind::OpenAi => "openai",
            ProviderKind::Anthropic => "anthropic",
            ProviderKind::Google => "google",
        }
    }
}

impl FromStr for ProviderKind {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> Result<Self> {
        match s.to_ascii_lowercase().as_str() {
            "openai" | "oai" | "gpt" => Ok(ProviderKind::OpenAi),
            "anthropic" | "claude" => Ok(ProviderKind::Anthropic),
            "google" | "gemini" | "googleai" => Ok(ProviderKind::Google),
            other => bail!("unknown provider '{other}' (expected openai|anthropic|google)"),
        }
    }
}

/// A provider request body. `raw` is the source of truth for serialization.
#[derive(Debug, Clone)]
pub struct Request {
    kind: ProviderKind,
    raw: Value,
}

impl Request {
    /// Parse a request body string.
    pub fn parse(kind: ProviderKind, body: &str) -> Result<Self> {
        let raw: Value = serde_json::from_str(body).context("request body is not valid JSON")?;
        Ok(Self { kind, raw })
    }

    /// Build a request from an already-parsed value (avoids a re-parse).
    pub fn from_value(kind: ProviderKind, raw: Value) -> Self {
        Self { kind, raw }
    }

    pub fn kind(&self) -> ProviderKind {
        self.kind
    }

    pub fn raw(&self) -> &Value {
        &self.raw
    }

    pub fn raw_mut(&mut self) -> &mut Value {
        &mut self.raw
    }

    pub fn to_json_string(&self) -> Result<String> {
        serde_json::to_string(&self.raw).context("failed to serialize request")
    }

    /// Read a string at a JSON pointer (RFC 6901), e.g. `/messages/0/content`.
    pub fn get_str(&self, pointer: &str) -> Option<&str> {
        self.raw.pointer(pointer).and_then(Value::as_str)
    }

    /// Replace the value at a JSON pointer. Returns `false` if the pointer is absent.
    pub fn set(&mut self, pointer: &str, value: Value) -> bool {
        match self.raw.pointer_mut(pointer) {
            Some(slot) => {
                *slot = value;
                true
            }
            None => false,
        }
    }
}

/// A provider response body (source of truth = `raw`), used for rehydration.
#[derive(Debug, Clone)]
pub struct Response {
    kind: ProviderKind,
    raw: Value,
}

impl Response {
    pub fn parse(kind: ProviderKind, body: &str) -> Result<Self> {
        let raw: Value = serde_json::from_str(body).context("response body is not valid JSON")?;
        Ok(Self { kind, raw })
    }

    pub fn kind(&self) -> ProviderKind {
        self.kind
    }

    pub fn raw(&self) -> &Value {
        &self.raw
    }

    pub fn raw_mut(&mut self) -> &mut Value {
        &mut self.raw
    }

    pub fn to_json_string(&self) -> Result<String> {
        serde_json::to_string(&self.raw).context("failed to serialize response")
    }
}

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

    #[test]
    fn provider_kind_parses_aliases() {
        assert_eq!(
            "openai".parse::<ProviderKind>().unwrap(),
            ProviderKind::OpenAi
        );
        assert_eq!(
            "Claude".parse::<ProviderKind>().unwrap(),
            ProviderKind::Anthropic
        );
        assert!("bogus".parse::<ProviderKind>().is_err());
    }

    #[test]
    fn request_round_trip_is_lossless() {
        let body = r#"{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}],"max_tokens":7,"unknown_field":{"keep":[1,2,3]}}"#;
        let req = Request::parse(ProviderKind::OpenAi, body).unwrap();
        let out = req.to_json_string().unwrap();
        let a: Value = serde_json::from_str(body).unwrap();
        let b: Value = serde_json::from_str(&out).unwrap();
        assert_eq!(a, b, "unknown fields must survive round-trip");
    }

    #[test]
    fn pointer_get_and_set() {
        let mut req = Request::parse(
            ProviderKind::OpenAi,
            r#"{"messages":[{"role":"user","content":"hi"}]}"#,
        )
        .unwrap();
        assert_eq!(req.get_str("/messages/0/content"), Some("hi"));
        assert!(req.set("/messages/0/content", Value::String("bye".into())));
        assert_eq!(req.get_str("/messages/0/content"), Some("bye"));
        assert!(!req.set("/nope/0", Value::Null));
    }
}