feat: add compact dictionary codebook

src/dictionary.rs

@@ -1,6 +1,7 @@
-use anyhow::{Context, Result};
+use anyhow::{Context, Result, bail};
 use serde::{Deserialize, Serialize};
-use std::collections::BTreeMap;
+use sha2::{Digest, Sha256};
+use std::collections::{BTreeMap, BTreeSet};
 use std::fs::{self, File};
 use std::io::Write;
 #[cfg(test)]
@@ -9,6 +10,8 @@ use std::path::{Path, PathBuf};
 
 use crate::registry::ActionDefinition;
 
+pub const DICTIONARY_SCHEMA_VERSION: &str = "1.0";
+
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 pub struct ActionDictionaryRecord {
     pub action_id: String,
@@ -22,14 +25,56 @@ pub struct ActionDictionaryRecord {
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 pub struct ActionIndex {
     pub schema_version: String,
+    pub dictionary_version: String,
+    pub registry_checksum: String,
+    pub static_base_checksum: String,
+    pub codebook_checksum: String,
     pub actions: BTreeMap<String, ActionDictionaryRecord>,
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
+pub struct StaticBaseDictionary {
+    pub dictionary_version: String,
+    pub keys: BTreeMap<String, String>,
+    pub types: BTreeMap<String, String>,
+    pub errors: BTreeMap<String, String>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
+pub struct AdaptiveCodebook {
+    pub actions: BTreeMap<String, String>,
+    pub fields: BTreeMap<String, String>,
+    pub modules: BTreeMap<String, String>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
+pub struct ModelCodebook {
+    pub schema_version: String,
+    pub dictionary_version: String,
+    pub registry_checksum: String,
+    pub static_base_checksum: String,
+    pub codebook_checksum: String,
+    pub static_base: StaticBaseDictionary,
+    pub adaptive: AdaptiveCodebook,
+}
+
+#[derive(Debug, Clone, Serialize)]
+struct ModelCodebookForChecksum<'a> {
+    schema_version: &'a str,
+    dictionary_version: &'a str,
+    registry_checksum: &'a str,
+    static_base_checksum: &'a str,
+    static_base: &'a StaticBaseDictionary,
+    adaptive: &'a AdaptiveCodebook,
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct DictionarySummary {
     pub action_count: usize,
     pub jsonl_path: PathBuf,
     pub index_path: PathBuf,
+    pub codebook_path: PathBuf,
+    pub examples_path: PathBuf,
 }
 
 pub fn default_dictionary_dir() -> PathBuf {
@@ -43,37 +88,49 @@ pub fn generate_dictionary_files(
     fs::create_dir_all(out_dir)
         .with_context(|| format!("creating dictionary directory {}", out_dir.display()))?;
 
+    let static_base = load_static_base_dictionary(out_dir)?;
+    validate_static_base_dictionary(&static_base)?;
+
+    let records = sorted_records(actions);
+    let registry_checksum = checksum_json(&records)?;
+    let static_base_checksum = checksum_json(&static_base)?;
+    let adaptive = build_adaptive_codebook(&records);
+    let codebook = build_model_codebook(
+        static_base,
+        adaptive,
+        registry_checksum,
+        static_base_checksum,
+    )?;
+
     let jsonl_path = out_dir.join("actions.jsonl");
     let index_path = out_dir.join("actions.index.json");
-    let records = actions
-        .iter()
-        .cloned()
-        .map(ActionDictionaryRecord::from)
-        .collect::<Vec<_>>();
+    let codebook_path = out_dir.join("model.codebook.json");
+    let examples_path = out_dir.join("model.examples.jsonl");
 
-    let mut jsonl =
-        File::create(&jsonl_path).with_context(|| format!("creating {}", jsonl_path.display()))?;
-    for record in &records {
-        serde_json::to_writer(&mut jsonl, record)?;
-        jsonl.write_all(b"\n")?;
-    }
+    write_actions_jsonl(&jsonl_path, &records)?;
 
     let index = ActionIndex {
-        schema_version: "1.0".to_string(),
+        schema_version: DICTIONARY_SCHEMA_VERSION.to_string(),
+        dictionary_version: codebook.dictionary_version.clone(),
+        registry_checksum: codebook.registry_checksum.clone(),
+        static_base_checksum: codebook.static_base_checksum.clone(),
+        codebook_checksum: codebook.codebook_checksum.clone(),
         actions: records
-            .into_iter()
+            .iter()
+            .cloned()
             .map(|record| (record.action_id.clone(), record))
             .collect(),
     };
-
-    let pretty = serde_json::to_string_pretty(&index)?;
-    fs::write(&index_path, format!("{pretty}\n"))
-        .with_context(|| format!("writing {}", index_path.display()))?;
+    write_pretty_json(&index_path, &index)?;
+    write_pretty_json(&codebook_path, &codebook)?;
+    write_model_examples_jsonl(&examples_path, &codebook)?;
 
     Ok(DictionarySummary {
         action_count: index.actions.len(),
         jsonl_path,
         index_path,
+        codebook_path,
+        examples_path,
     })
 }
 
@@ -90,6 +147,236 @@ pub fn load_action_index(dictionary_dir: &Path) -> Result<ActionIndex> {
     Ok(index)
 }
 
+pub fn load_model_codebook(dictionary_dir: &Path) -> Result<ModelCodebook> {
+    let path = dictionary_dir.join("model.codebook.json");
+    let raw = fs::read_to_string(&path).with_context(|| {
+        format!(
+            "reading {}; run `cargo run -- dictionary generate` first",
+            path.display()
+        )
+    })?;
+    let codebook = serde_json::from_str::<ModelCodebook>(&raw)
+        .with_context(|| format!("parsing {}", path.display()))?;
+    Ok(codebook)
+}
+
+pub fn resolve_action_code<'a>(codebook: &'a ModelCodebook, code: &str) -> Option<&'a str> {
+    codebook.adaptive.actions.get(code).map(String::as_str)
+}
+
+pub fn resolve_field_code<'a>(codebook: &'a ModelCodebook, code: &str) -> Option<&'a str> {
+    codebook.adaptive.fields.get(code).map(String::as_str)
+}
+
+fn load_static_base_dictionary(dictionary_dir: &Path) -> Result<StaticBaseDictionary> {
+    let path = dictionary_dir.join("static-base.json");
+    let raw = fs::read_to_string(&path).with_context(|| {
+        format!(
+            "reading tracked static base dictionary {}; create this file before generating dictionaries",
+            path.display()
+        )
+    })?;
+    serde_json::from_str::<StaticBaseDictionary>(&raw)
+        .with_context(|| format!("parsing {}", path.display()))
+}
+
+fn validate_static_base_dictionary(static_base: &StaticBaseDictionary) -> Result<()> {
+    let expected_keys = [("K001", "a"), ("K002", "p"), ("K003", "v"), ("K004", "c")];
+    let expected_types = [
+        ("T001", "action"),
+        ("T002", "field"),
+        ("T003", "error"),
+        ("T004", "module"),
+    ];
+    let expected_errors = [
+        ("E001", "unknown_action"),
+        ("E002", "invalid_payload"),
+        ("E003", "checksum_mismatch"),
+    ];
+
+    for (code, value) in expected_keys {
+        if static_base.keys.get(code).map(String::as_str) != Some(value) {
+            bail!("static base dictionary must reserve {code}={value}");
+        }
+    }
+    for (code, value) in expected_types {
+        if static_base.types.get(code).map(String::as_str) != Some(value) {
+            bail!("static base dictionary must reserve {code}={value}");
+        }
+    }
+    for (code, value) in expected_errors {
+        if static_base.errors.get(code).map(String::as_str) != Some(value) {
+            bail!("static base dictionary must reserve {code}={value}");
+        }
+    }
+    Ok(())
+}
+
+fn sorted_records(actions: &[ActionDefinition]) -> Vec<ActionDictionaryRecord> {
+    let mut records = actions
+        .iter()
+        .cloned()
+        .map(ActionDictionaryRecord::from)
+        .collect::<Vec<_>>();
+    records.sort_by(|left, right| left.action_id.cmp(&right.action_id));
+    records
+}
+
+fn build_adaptive_codebook(records: &[ActionDictionaryRecord]) -> AdaptiveCodebook {
+    let actions = records
+        .iter()
+        .enumerate()
+        .map(|(index, record)| (format_code("A", index), record.action_id.clone()))
+        .collect::<BTreeMap<_, _>>();
+
+    let fields = records
+        .iter()
+        .flat_map(payload_fields)
+        .collect::<BTreeSet<_>>()
+        .into_iter()
+        .enumerate()
+        .map(|(index, field)| (format_code("F", index), field))
+        .collect::<BTreeMap<_, _>>();
+
+    let modules = records
+        .iter()
+        .map(|record| record.executable.script.clone())
+        .collect::<BTreeSet<_>>()
+        .into_iter()
+        .enumerate()
+        .map(|(index, module)| (format_code("M", index), module))
+        .collect::<BTreeMap<_, _>>();
+
+    AdaptiveCodebook {
+        actions,
+        fields,
+        modules,
+    }
+}
+
+fn build_model_codebook(
+    static_base: StaticBaseDictionary,
+    adaptive: AdaptiveCodebook,
+    registry_checksum: String,
+    static_base_checksum: String,
+) -> Result<ModelCodebook> {
+    let mut codebook = ModelCodebook {
+        schema_version: DICTIONARY_SCHEMA_VERSION.to_string(),
+        dictionary_version: static_base.dictionary_version.clone(),
+        registry_checksum,
+        static_base_checksum,
+        codebook_checksum: String::new(),
+        static_base,
+        adaptive,
+    };
+
+    codebook.codebook_checksum = checksum_codebook_without_self(&codebook)?;
+    Ok(codebook)
+}
+
+fn checksum_codebook_without_self(codebook: &ModelCodebook) -> Result<String> {
+    checksum_json(&ModelCodebookForChecksum {
+        schema_version: &codebook.schema_version,
+        dictionary_version: &codebook.dictionary_version,
+        registry_checksum: &codebook.registry_checksum,
+        static_base_checksum: &codebook.static_base_checksum,
+        static_base: &codebook.static_base,
+        adaptive: &codebook.adaptive,
+    })
+}
+
+fn checksum_json<T: Serialize>(value: &T) -> Result<String> {
+    let encoded = serde_json::to_vec(value)?;
+    let digest = Sha256::digest(encoded);
+    Ok(format!("{digest:x}"))
+}
+
+fn payload_fields(record: &ActionDictionaryRecord) -> Vec<String> {
+    record
+        .payload_schema
+        .get("properties")
+        .and_then(serde_json::Value::as_object)
+        .map(|properties| properties.keys().cloned().collect())
+        .unwrap_or_default()
+}
+
+fn format_code(prefix: &str, zero_based_index: usize) -> String {
+    format!("{prefix}{:03}", zero_based_index + 1)
+}
+
+fn write_actions_jsonl(path: &Path, records: &[ActionDictionaryRecord]) -> Result<()> {
+    let mut jsonl = File::create(path).with_context(|| format!("creating {}", path.display()))?;
+    for record in records {
+        serde_json::to_writer(&mut jsonl, record)?;
+        jsonl.write_all(b"\n")?;
+    }
+    Ok(())
+}
+
+fn write_model_examples_jsonl(path: &Path, codebook: &ModelCodebook) -> Result<()> {
+    let mut jsonl = File::create(path).with_context(|| format!("creating {}", path.display()))?;
+    let examples = [
+        serde_json::json!({
+            "kind": "verbose_example",
+            "input": "log hello from codex",
+            "target": {
+                "action_id": "logger",
+                "payload": {
+                    "message": "hello from codex"
+                }
+            }
+        }),
+        serde_json::json!({
+            "kind": "compact_example",
+            "input": "log hello from codex",
+            "target": {
+                "v": codebook.dictionary_version,
+                "c": codebook.codebook_checksum,
+                "a": "A001",
+                "p": {
+                    "F001": "hello from codex"
+                }
+            }
+        }),
+        serde_json::json!({
+            "kind": "invalid_checksum_example",
+            "input": "log hello with a stale dictionary",
+            "target": {
+                "v": codebook.dictionary_version,
+                "c": "0000000000000000000000000000000000000000000000000000000000000000",
+                "a": "A001",
+                "p": {
+                    "F001": "hello"
+                }
+            },
+            "expected_runtime_result": "reject_checksum_mismatch"
+        }),
+        serde_json::json!({
+            "kind": "synthetic_template",
+            "input": "{{natural_language_request}}",
+            "target": {
+                "v": codebook.dictionary_version,
+                "c": codebook.codebook_checksum,
+                "a": "{{action_code}}",
+                "p": {
+                    "{{field_code}}": "{{field_value}}"
+                }
+            }
+        }),
+    ];
+
+    for example in examples {
+        serde_json::to_writer(&mut jsonl, &example)?;
+        jsonl.write_all(b"\n")?;
+    }
+    Ok(())
+}
+
+fn write_pretty_json<T: Serialize>(path: &Path, value: &T) -> Result<()> {
+    let pretty = serde_json::to_string_pretty(value)?;
+    fs::write(path, format!("{pretty}\n")).with_context(|| format!("writing {}", path.display()))
+}
+
 #[cfg(test)]
 fn read_jsonl_records(path: &Path) -> Result<Vec<ActionDictionaryRecord>> {
     let file = File::open(path).with_context(|| format!("opening {}", path.display()))?;
@@ -126,13 +413,57 @@ mod tests {
     use crate::registry::registered_actions;
 
     #[test]
-    fn generates_jsonl_and_index() {
-        let temp = tempfile::tempdir().expect("temp dir");
+    fn static_base_dictionary_reserves_expected_codes() {
+        let temp = temp_dictionary_dir();
+        let static_base = load_static_base_dictionary(temp.path()).unwrap();
+        validate_static_base_dictionary(&static_base).unwrap();
+        assert_eq!(static_base.keys["K001"], "a");
+        assert_eq!(static_base.types["T002"], "field");
+        assert_eq!(static_base.errors["E003"], "checksum_mismatch");
+    }
+
+    #[test]
+    fn adaptive_layer_maps_logger_and_message_deterministically() {
+        let records = sorted_records(&registered_actions());
+        let adaptive = build_adaptive_codebook(&records);
+        assert_eq!(adaptive.actions["A001"], "logger");
+        assert_eq!(adaptive.fields["F001"], "message");
+        assert_eq!(adaptive.modules["M001"], "logger");
+    }
+
+    #[test]
+    fn checksums_are_stable_for_unchanged_inputs() {
+        let temp = temp_dictionary_dir();
+        let first = generate_dictionary_files(&registered_actions(), temp.path()).unwrap();
+        let first_codebook = load_model_codebook(temp.path()).unwrap();
+        let second = generate_dictionary_files(&registered_actions(), temp.path()).unwrap();
+        let second_codebook = load_model_codebook(temp.path()).unwrap();
+
+        assert_eq!(first.action_count, second.action_count);
+        assert_eq!(
+            first_codebook.registry_checksum,
+            second_codebook.registry_checksum
+        );
+        assert_eq!(
+            first_codebook.static_base_checksum,
+            second_codebook.static_base_checksum
+        );
+        assert_eq!(
+            first_codebook.codebook_checksum,
+            second_codebook.codebook_checksum
+        );
+    }
+
+    #[test]
+    fn generates_all_dictionary_outputs() {
+        let temp = temp_dictionary_dir();
         let summary = generate_dictionary_files(&registered_actions(), temp.path()).unwrap();
 
         assert_eq!(summary.action_count, 1);
         assert!(summary.jsonl_path.exists());
         assert!(summary.index_path.exists());
+        assert!(summary.codebook_path.exists());
+        assert!(summary.examples_path.exists());
 
         let records = read_jsonl_records(&summary.jsonl_path).unwrap();
         assert_eq!(records.len(), 1);
@@ -141,5 +472,20 @@ mod tests {
         let index = load_action_index(temp.path()).unwrap();
         assert!(index.actions.contains_key("logger"));
         assert!(!index.actions.contains_key("log"));
+
+        let codebook = load_model_codebook(temp.path()).unwrap();
+        assert_eq!(codebook.adaptive.actions["A001"], "logger");
+        assert_eq!(codebook.adaptive.fields["F001"], "message");
+        assert!(!codebook.codebook_checksum.is_empty());
+    }
+
+    fn temp_dictionary_dir() -> tempfile::TempDir {
+        let temp = tempfile::tempdir().expect("temp dir");
+        fs::write(
+            temp.path().join("static-base.json"),
+            include_str!("../dictionary/static-base.json"),
+        )
+        .unwrap();
+        temp
     }
 }