jev_bridge/

prompt.rs

//! The direct decision prompt contract and the single-token answer-slot check.
//!
//! This mirrors fastjev's `direct-options-v1` contract: one system instruction,
//! one JSON payload of evidence / criterion / lettered options, rendered
//! through the model's own chat template with thinking disabled, then scored
//! by reading the next-token logits of the fixed uppercase answer letters.
//!
//! Rendering can come from two places. By default the serving runtime renders
//! it (`/apply-template` on llama.cpp), which guarantees the prompt is exactly
//! what that runtime feeds the model. When the runtime exposes no such endpoint
//! — vLLM, for instance — [`crate::render::LocalRenderer`] renders the same
//! template in-process instead.

use anyhow::{bail, Context, Result};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use sha2::{Digest, Sha256};

use crate::render::LocalRenderer;
use crate::wire::{to_python_json, RowOption};
#[cfg(feature = "local-tokenizer")]
use crate::tokenizer::LocalTokenizer;

/// The uppercase answer letters, in option order.
pub const LETTERS: &str = "ABCDEFGHIJKLMNOP";
/// The system instruction the direct contract freezes.
pub const DIRECT_SYSTEM: &str = "Apply the supplied criterion to the supplied evidence. Choose exactly one listed option. Respond with only its uppercase letter, with no explanation or reasoning.";
/// The prompt contract version recorded with every result.
pub const PROMPT_VERSION: &str = "direct-options-v1";

/// One chat message, in the shape every chat template expects.
#[derive(Debug, Clone, Serialize)]
pub struct ChatMessage {
    /// Who speaks: `system` or `user`.
    pub role: &'static str,
    /// The message text, already serialised by the caller.
    pub content: String,
}

/// The serialised evidence text for one state.
///
/// A batch of questions about the same state serialises it once and reuses the
/// text; the bytes are exactly what [`direct_messages`] embeds, which a test
/// pins.
pub fn evidence_json(state: &Value) -> String {
    to_python_json(state)
}

/// Build the chat messages for one decision from an already-serialised
/// evidence text.
///
/// `json.dumps` serialises a nested value exactly as it would serialise that
/// value alone, so assembling the payload from independently serialised parts
/// is byte-identical to serialising the whole object — a test pins that, since
/// the byte-for-byte alignment with the reference implementation rests on it.
pub fn direct_messages_reusing(
    evidence: &str,
    question: &str,
    options: &[RowOption],
) -> Vec<ChatMessage> {
    let entries: Vec<Value> = options
        .iter()
        .enumerate()
        .map(|(index, option)| {
            json!({
                "letter": &LETTERS[index..index + 1],
                "description": option.description,
            })
        })
        .collect();
    let payload = format!(
        "{{\"evidence\": {evidence}, \"criterion\": {}, \"options\": {}}}",
        to_python_json(&Value::String(question.to_string())),
        to_python_json(&Value::Array(entries)),
    );
    vec![
        ChatMessage {
            role: "system",
            content: DIRECT_SYSTEM.to_string(),
        },
        ChatMessage {
            role: "user",
            content: payload,
        },
    ]
}

/// Build the chat messages for one decision.
///
/// The payload uses Python's JSON layout so the text matches the reference
/// implementation byte for byte.
pub fn direct_messages(state: &Value, question: &str, options: &[RowOption]) -> Vec<ChatMessage> {
    direct_messages_reusing(&evidence_json(state), question, options)
}

/// SHA-256 of the rendered prompt.
///
/// This is the strongest available check that a bridged run used exactly the
/// same prompt as a reference run: equal hashes mean equal bytes.
pub fn prompt_sha256(prompt: &str) -> String {
    let mut hasher = Sha256::new();
    hasher.update(prompt.as_bytes());
    format!("{:x}", hasher.finalize())
}

/// Softmax over the selected option logits, matching fastjev's helper.
pub fn softmax(values: &[f64]) -> Vec<f64> {
    let maximum = values.iter().copied().fold(f64::NEG_INFINITY, f64::max);
    let weights: Vec<f64> = values.iter().map(|value| (value - maximum).exp()).collect();
    let total: f64 = weights.iter().sum();
    weights.iter().map(|weight| weight / total).collect()
}

/// The body of a `/apply-template` call.
#[derive(Debug, Serialize)]
struct ApplyTemplateRequest<'a> {
    model: &'a str,
    messages: &'a [ChatMessage],
    /// Template variables such as `enable_thinking`, forwarded verbatim.
    chat_template_kwargs: &'a Value,
}

#[derive(Debug, Deserialize)]
struct ApplyTemplateResponse {
    prompt: String,
}

/// The body of a `/tokenize` call.
#[derive(Debug, Serialize)]
struct TokenizeRequest<'a> {
    model: &'a str,
    content: &'a str,
    /// The template already emits BOS, so tokenizing must not add another.
    add_special: bool,
}

#[derive(Debug, Deserialize)]
struct TokenizeResponse {
    tokens: Vec<u32>,
}

#[derive(Debug, Serialize)]
struct DetokenizeRequest<'a> {
    model: &'a str,
    tokens: &'a [u32],
}

#[derive(Debug, Deserialize)]
struct DetokenizeResponse {
    content: String,
}

/// HTTP access to the serving runtime.
///
/// Tokenization always goes through the runtime unless a local tokenizer is
/// configured, so the bridge never needs a copy of the model's vocabulary.
/// Rendering goes through it too unless a local renderer is configured.
pub struct RuntimeClient {
    client: reqwest::Client,
    base: String,
    model: String,
    api_key: Option<String>,
}

impl RuntimeClient {
    /// Point a client at the runtime's native base URL.
    pub fn new(
        client: reqwest::Client,
        base: impl Into<String>,
        model: impl Into<String>,
        api_key: Option<String>,
    ) -> Self {
        Self {
            client,
            base: base.into().trim_end_matches('/').to_string(),
            model: model.into(),
            api_key,
        }
    }

    async fn post<T: serde::de::DeserializeOwned>(
        &self,
        path: &str,
        body: &(impl serde::Serialize + ?Sized),
    ) -> Result<T> {
        let url = format!("{}{path}", self.base);
        let mut request = self.client.post(&url).json(body);
        if let Some(key) = &self.api_key {
            request = request.bearer_auth(key);
        }
        let response = request
            .send()
            .await
            .with_context(|| format!("POST {url} failed"))?;
        let status = response.status();
        let text = response
            .text()
            .await
            .with_context(|| format!("reading {url} response failed"))?;
        if !status.is_success() {
            bail!("POST {url} returned HTTP {status}: {}", truncate(&text));
        }
        serde_json::from_str(&text)
            .with_context(|| format!("POST {url} returned an unexpected body: {}", truncate(&text)))
    }

    /// Render messages with the runtime's own chat template.
    pub async fn apply_template(
        &self,
        messages: &[ChatMessage],
        chat_template_kwargs: &Value,
    ) -> Result<String> {
        let body = ApplyTemplateRequest {
            model: &self.model,
            messages,
            chat_template_kwargs,
        };
        let response: ApplyTemplateResponse = self
            .post("/apply-template", &body)
            .await
            .context("rendering the chat template on the serving runtime failed")?;
        Ok(response.prompt)
    }

    /// Tokenize without special tokens, matching fastjev's `encode_prompt`.
    pub async fn tokenize(&self, content: &str) -> Result<Vec<u32>> {
        let body = TokenizeRequest {
            model: &self.model,
            content,
            add_special: false,
        };
        let response: TokenizeResponse = self.post("/tokenize", &body).await?;
        Ok(response.tokens)
    }

    /// Decode tokens for the answer-slot round-trip check.
    pub async fn detokenize(&self, tokens: &[u32]) -> Result<String> {
        let body = DetokenizeRequest {
            model: &self.model,
            tokens,
        };
        let response: DetokenizeResponse = self.post("/detokenize", &body).await?;
        Ok(response.content)
    }
}

/// The work the bridge can do without asking the serving runtime.
///
/// Both fields are optional and independent: a runtime may render templates but
/// not tokenize, or the other way round. Anything left unset falls back to the
/// runtime's own endpoints.
#[derive(Default)]
pub struct LocalComponents {
    /// Renders the template in-process instead of calling `/apply-template`.
    pub renderer: Option<LocalRenderer>,
    #[cfg(feature = "local-tokenizer")]
    /// Tokenizes in-process instead of calling `/tokenize`.
    pub tokenizer: Option<LocalTokenizer>,
}

/// The prompt contract: rendering, tokenization, and the startup checks.
pub struct ChatTemplate {
    runtime: RuntimeClient,
    local: LocalComponents,
    chat_template_kwargs: Value,
}

impl ChatTemplate {
    /// Combine a runtime with whichever in-process halves are configured.
    pub fn new(runtime: RuntimeClient, local: LocalComponents, chat_template_kwargs: Value) -> Self {
        Self {
            runtime,
            local,
            chat_template_kwargs,
        }
    }

    /// Whether the prompt is rendered in this process rather than by the runtime.
    pub fn renders_locally(&self) -> bool {
        self.local.renderer.is_some()
    }

    /// Whether tokenization happens in this process rather than by the runtime.
    pub fn tokenizes_locally(&self) -> bool {
        #[cfg(feature = "local-tokenizer")]
        {
            self.local.tokenizer.is_some()
        }
        #[cfg(not(feature = "local-tokenizer"))]
        {
            false
        }
    }

    /// Render the decision prompt for one set of messages.
    pub async fn render(&self, messages: &[ChatMessage]) -> Result<String> {
        match &self.local.renderer {
            Some(local) => local.render(messages, &self.chat_template_kwargs),
            None => {
                self.runtime
                    .apply_template(messages, &self.chat_template_kwargs)
                    .await
            }
        }
    }

    /// Tokenize locally when configured, otherwise through the runtime.
    pub async fn tokenize(&self, content: &str) -> Result<Vec<u32>> {
        #[cfg(feature = "local-tokenizer")]
        if let Some(local) = &self.local.tokenizer {
            return local.tokenize(content);
        }
        self.runtime.tokenize(content).await
    }

    /// Detokenize locally when configured, otherwise through the runtime.
    pub async fn detokenize(&self, tokens: &[u32]) -> Result<String> {
        #[cfg(feature = "local-tokenizer")]
        if let Some(local) = &self.local.tokenizer {
            return local.detokenize(tokens);
        }
        self.runtime.detokenize(tokens).await
    }

    /// Resolve the sixteen answer slots and prove they are exact single tokens.
    ///
    /// A slot that is not one round-tripping token would make the logit readout
    /// meaningless, so this fails loudly instead of scoring the wrong ids.
    pub async fn resolve_slots(&self) -> Result<Vec<u32>> {
        let mut slots = Vec::with_capacity(LETTERS.len());
        for letter in LETTERS.chars() {
            let ids = self.tokenize(&letter.to_string()).await?;
            if ids.len() != 1 {
                bail!(
                    "answer slot {letter:?} is not a single token ({} tokens); this model cannot \
                     serve the direct option-logit contract",
                    ids.len()
                );
            }
            slots.push(ids[0]);
        }
        let mut unique = slots.clone();
        unique.sort_unstable();
        unique.dedup();
        if unique.len() != slots.len() {
            bail!("answer-slot tokens collide");
        }
        let decoded = self
            .detokenize(&slots)
            .await
            .context("verifying answer-slot round trip failed")?;
        if decoded != LETTERS {
            bail!("answer slots decode to {decoded:?} instead of {LETTERS:?}");
        }
        Ok(slots)
    }

    /// Prove that appending an answer letter cannot change tokenization.
    ///
    /// Without this, the last prompt token could merge with the letter and the
    /// scored logit would belong to a different token than the chosen option.
    pub async fn verify_boundary(&self, prompt: &str, slots: &[u32]) -> Result<()> {
        let base = self.tokenize(prompt).await?;
        if base.is_empty() {
            bail!("rendered prompt is empty");
        }
        for (index, letter) in LETTERS.chars().enumerate() {
            let combined = self.tokenize(&format!("{prompt}{letter}")).await?;
            let mut expected = base.clone();
            expected.push(slots[index]);
            if combined != expected {
                bail!(
                    "answer boundary changes tokenization for slot {letter:?}: the prompt does \
                     not end in a position where {letter:?} stays one token"
                );
            }
        }
        Ok(())
    }

    /// The template variables forwarded on every render.
    pub fn chat_template_kwargs(&self) -> Value {
        self.chat_template_kwargs.clone()
    }
}

fn truncate(text: &str) -> String {
    let mut owned = text.chars().take(400).collect::<String>();
    if text.chars().count() > 400 {
        owned.push('…');
    }
    owned
}

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

    fn options() -> Vec<RowOption> {
        vec![
            RowOption {
                id: "billing".into(),
                description: "Billing and refunds.".into(),
            },
            RowOption {
                id: "sales".into(),
                description: "Pricing and contracts.".into(),
            },
        ]
    }

    #[test]
    fn payload_matches_python_json_layout() {
        let messages = direct_messages(&json!("charged twice"), "Which queue?", &options());
        let user = messages[1].content.as_str();
        assert_eq!(
            user,
            r#"{"evidence": "charged twice", "criterion": "Which queue?", "options": [{"letter": "A", "description": "Billing and refunds."}, {"letter": "B", "description": "Pricing and contracts."}]}"#
        );
        assert_eq!(messages[0].content, DIRECT_SYSTEM);
    }

    #[test]
    fn structured_state_is_embedded_as_json() {
        let messages = direct_messages(&json!({"message": "hi", "count": 2}), "q", &options());
        let user = messages[1].content.as_str();
        assert!(
            user.contains(r#""evidence": {"message": "hi", "count": 2}"#),
            "{user}"
        );
    }

    #[test]
    fn reusing_the_serialised_evidence_is_byte_identical() {
        // A batch serialises the shared state once and assembles each payload
        // from parts; the result must be exactly what the whole-object path
        // produces, or the byte-for-byte alignment with the reference
        // implementation would quietly break for batched requests only.
        let state = json!({
            "message": "已扣款两次 charged twice",
            "count": 2,
            "nested": {"tags": ["a", "b"], "ok": true, "missing": null, "ratio": 0.5},
        });
        for question in ["Which queue?", "哪一个队列？", "quotes \" and \\ backslash"] {
            let whole = direct_messages(&state, question, &options());
            let reused = direct_messages_reusing(&evidence_json(&state), question, &options());
            assert_eq!(whole[0].content, reused[0].content);
            assert_eq!(whole[1].content, reused[1].content, "question {question:?}");
        }
    }

    #[test]
    fn softmax_is_stable_for_large_magnitudes() {
        let probabilities = softmax(&[-1000.0, -1001.0]);
        assert!((probabilities[0] - 0.7310585786300049).abs() < 1e-12);
        assert!((probabilities.iter().sum::<f64>() - 1.0).abs() < 1e-12);
    }

    #[test]
    fn prompt_hash_is_stable_and_distinguishing() {
        assert_eq!(prompt_sha256("abc").len(), 64);
        assert_ne!(prompt_sha256("abc"), prompt_sha256("abd"));
    }

    #[test]
    fn softmax_matches_the_measured_llamacpp_distribution() {
        // Measured on MiniCPM5-2B-Q8_0 for a billing question: the option
        // logits were -12.32 (A), -2.1e-05 (B), -15.61 (C).
        let probabilities = softmax(&[-12.320333, -0.0000213, -15.617624]);
        assert!(probabilities[1] > 0.9999, "{probabilities:?}");
        assert!(probabilities[0] < 1e-5);
    }

    #[tokio::test]
    async fn a_local_renderer_takes_precedence_over_the_runtime() {
        let runtime = RuntimeClient::new(
            reqwest::Client::new(),
            "http://127.0.0.1:1".to_string(),
            "unreachable".to_string(),
            None,
        );
        let local = LocalRenderer::new(
            "local:{{ messages[0].content }}:{{ enable_thinking }}".to_string(),
            Map::new(),
        )
        .unwrap();
        let template = ChatTemplate::new(
            runtime,
            LocalComponents {
                renderer: Some(local),
                ..Default::default()
            },
            json!({"enable_thinking": false}),
        );

        assert!(template.renders_locally());
        let messages = direct_messages(&json!("state"), "question", &options());
        // No HTTP call happens, so the unreachable base URL never matters.
        let rendered = template.render(&messages).await.unwrap();
        assert!(rendered.starts_with("local:"), "{rendered}");
        assert!(rendered.ends_with(":False"), "{rendered}");
    }
}