lean_ctx/proxy/
holdout.rs

1//! Deterministic output-savings holdout (#895 Track B).
2//!
3//! To honestly measure how much our **output-shaping** (cache-safe effort
4//! control #834 + the optional verbosity steer) reduces *output* tokens, we A/B
5//! it against a control arm. The cohort is a pure function of conversation
6//! identity (system prompt + first user message), so:
7//! * the SAME conversation is always in the SAME arm — every turn — which keeps
8//!   the decision (and therefore the request body) stable across turns, and
9//! * a fraction `f` of conversations land in the control arm, which is metered
10//!   but NOT output-shaped, giving a clean paired comparison of average output
11//!   tokens per arm.
12//!
13//! `output_holdout = 0` (default) puts everyone in `Treatment` → no behaviour
14//! change. The hash is content-addressed (blake3), not random, so it is stable
15//! across processes and machines.
16
17use serde_json::Value;
18
19/// Number of cohort buckets; a conversation maps to one bucket in `0..BUCKETS`.
20const BUCKETS: u64 = 10_000;
21
22/// Separator between key fields, chosen so it cannot occur in normal prose.
23const FIELD_SEP: char = '\u{1}';
24
25/// Which experiment arm a conversation is assigned to.
26#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub enum Arm {
28    /// Output-shaping skipped (the measurement baseline); still metered.
29    Control,
30    /// Output-shaping applied (effort control + verbosity steer).
31    Treatment,
32}
33
34impl Arm {
35    #[must_use]
36    pub fn as_str(self) -> &'static str {
37        match self {
38            Arm::Control => "control",
39            Arm::Treatment => "treatment",
40        }
41    }
42}
43
44/// Deterministic bucket `0..BUCKETS` from conversation key material.
45#[must_use]
46pub fn bucket(key: &str) -> u64 {
47    let hash = blake3::hash(key.as_bytes());
48    let b = hash.as_bytes();
49    let n = u64::from_le_bytes([b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]]);
50    n % BUCKETS
51}
52
53/// Assign an [`Arm`] for `key` given the control fraction `holdout` in `[0,1]`.
54/// `holdout <= 0` ⇒ everyone is [`Arm::Treatment`] (no control arm).
55#[must_use]
56pub fn assign(key: &str, holdout: f64) -> Arm {
57    let h = holdout.clamp(0.0, 1.0);
58    if h <= 0.0 {
59        return Arm::Treatment;
60    }
61    #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
62    let threshold = (h * BUCKETS as f64).round() as u64;
63    if bucket(key) < threshold {
64        Arm::Control
65    } else {
66        Arm::Treatment
67    }
68}
69
70/// Flatten a JSON content value (string / array of text blocks / `{text}` /
71/// `{content}`) into plain text for cohort keying. Order-preserving and pure.
72fn flatten_text(v: &Value) -> String {
73    match v {
74        Value::String(s) => s.clone(),
75        Value::Array(items) => items.iter().map(flatten_text).collect::<Vec<_>>().join(" "),
76        Value::Object(map) => {
77            if let Some(Value::String(t)) = map.get("text") {
78                t.clone()
79            } else if let Some(inner) = map.get("content") {
80                flatten_text(inner)
81            } else if let Some(parts) = map.get("parts") {
82                flatten_text(parts)
83            } else {
84                String::new()
85            }
86        }
87        _ => String::new(),
88    }
89}
90
91/// Text of the first message in `messages` whose `role` matches, flattening the
92/// `content_field` (e.g. `"content"` for OpenAI/Anthropic, `"parts"` for
93/// Gemini). Empty when absent.
94fn first_message_text(messages: Option<&Value>, role: &str, content_field: &str) -> String {
95    messages
96        .and_then(Value::as_array)
97        .and_then(|arr| {
98            arr.iter()
99                .find(|m| m.get("role").and_then(Value::as_str) == Some(role))
100        })
101        .and_then(|m| m.get(content_field))
102        .map(flatten_text)
103        .unwrap_or_default()
104}
105
106/// Cohort key for an Anthropic `/v1/messages` body: `system` + first user turn.
107#[must_use]
108pub fn anthropic_key(doc: &Value) -> String {
109    let system = doc.get("system").map(flatten_text).unwrap_or_default();
110    let first_user = first_message_text(doc.get("messages"), "user", "content");
111    format!("{system}{FIELD_SEP}{first_user}")
112}
113
114/// Cohort key for an OpenAI Chat Completions body: first `system`/`developer`
115/// message + first user message.
116#[must_use]
117pub fn openai_chat_key(doc: &Value) -> String {
118    let messages = doc.get("messages");
119    let mut system = first_message_text(messages, "system", "content");
120    if system.is_empty() {
121        system = first_message_text(messages, "developer", "content");
122    }
123    let first_user = first_message_text(messages, "user", "content");
124    format!("{system}{FIELD_SEP}{first_user}")
125}
126
127/// Cohort key for an OpenAI Responses body: `instructions` + first user `input`.
128#[must_use]
129pub fn openai_responses_key(doc: &Value) -> String {
130    let system = doc
131        .get("instructions")
132        .map(flatten_text)
133        .unwrap_or_default();
134    // `input` is either a plain string or an array of role/content items.
135    let first_user = match doc.get("input") {
136        Some(Value::String(s)) => s.clone(),
137        other => first_message_text(other, "user", "content"),
138    };
139    format!("{system}{FIELD_SEP}{first_user}")
140}
141
142/// Cohort key for a Gemini `generateContent` body: `systemInstruction` + first
143/// user `contents` turn.
144#[must_use]
145pub fn google_key(doc: &Value) -> String {
146    let system = doc
147        .get("systemInstruction")
148        .map(flatten_text)
149        .unwrap_or_default();
150    let first_user = first_message_text(doc.get("contents"), "user", "parts");
151    format!("{system}{FIELD_SEP}{first_user}")
152}
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157    use serde_json::json;
158
159    #[test]
160    fn holdout_zero_is_all_treatment() {
161        assert_eq!(assign("any key", 0.0), Arm::Treatment);
162        assert_eq!(assign("any key", -1.0), Arm::Treatment);
163    }
164
165    #[test]
166    fn holdout_one_is_all_control() {
167        assert_eq!(assign("any key", 1.0), Arm::Control);
168        assert_eq!(assign("another", 2.0), Arm::Control);
169    }
170
171    #[test]
172    fn assignment_is_deterministic_and_stable() {
173        // Same key → same arm, every call (cross-turn stability).
174        let k = "system\u{1}first user message";
175        let a = assign(k, 0.5);
176        for _ in 0..20 {
177            assert_eq!(assign(k, 0.5), a);
178        }
179    }
180
181    #[test]
182    fn fraction_is_approximately_honoured() {
183        // ~30% of distinct conversations land in control.
184        let control = (0..5000)
185            .filter(|i| assign(&format!("conv-{i}"), 0.3) == Arm::Control)
186            .count();
187        let frac = control as f64 / 5000.0;
188        assert!((0.27..0.33).contains(&frac), "got {frac}");
189    }
190
191    #[test]
192    fn anthropic_key_uses_system_and_first_user() {
193        let doc = json!({
194            "system": "You are helpful.",
195            "messages": [
196                {"role": "user", "content": "Hello there"},
197                {"role": "assistant", "content": "Hi"},
198                {"role": "user", "content": "later turn"}
199            ]
200        });
201        assert_eq!(anthropic_key(&doc), "You are helpful.\u{1}Hello there");
202    }
203
204    #[test]
205    fn anthropic_key_flattens_block_arrays() {
206        let doc = json!({
207            "system": [{"type": "text", "text": "Sys A"}, {"type": "text", "text": "Sys B"}],
208            "messages": [{"role": "user", "content": [{"type": "text", "text": "U1"}]}]
209        });
210        assert_eq!(anthropic_key(&doc), "Sys A Sys B\u{1}U1");
211    }
212
213    #[test]
214    fn openai_chat_key_prefers_system_then_developer() {
215        let doc = json!({
216            "messages": [
217                {"role": "developer", "content": "Dev rules"},
218                {"role": "user", "content": "Q"}
219            ]
220        });
221        assert_eq!(openai_chat_key(&doc), "Dev rules\u{1}Q");
222    }
223
224    #[test]
225    fn google_key_uses_system_instruction_and_contents() {
226        let doc = json!({
227            "systemInstruction": {"parts": [{"text": "Be terse"}]},
228            "contents": [{"role": "user", "parts": [{"text": "First Q"}]}]
229        });
230        assert_eq!(google_key(&doc), "Be terse\u{1}First Q");
231    }
232
233    #[test]
234    fn responses_key_handles_string_and_array_input() {
235        let s = json!({"instructions": "Sys", "input": "just a string"});
236        assert_eq!(openai_responses_key(&s), "Sys\u{1}just a string");
237        let a = json!({
238            "instructions": "Sys",
239            "input": [{"role": "user", "content": "arr q"}]
240        });
241        assert_eq!(openai_responses_key(&a), "Sys\u{1}arr q");
242    }
243}
lean_ctx/proxy/holdout.rs

lean_ctx/proxy/
holdout.rs
