tycode_core/ai/
json.rs

1use aws_smithy_types::Document;
2use serde_json::Value;
3use std::collections::{BTreeMap, HashMap};
4
5/// This is a massive hack trying to work around a bug in bedrock prompt cache.
6/// When smithy-rs is convering a hashmap to json string its using hashmap
7/// iteration order, which thanks to rust's DDOS resistent hashing scheme is
8/// effectively random. This causes bedrock to not use the cache since it
9/// detects contents have changed.
10///
11/// The rust SDK needs to own the hashmap (which is probably another bug) so I
12/// can't cache values. So this is the best I got 🤷‍♂️ Perhaps if we retry until
13/// we happen to get sorted order (and hopefully iteration order is stable until
14/// mutation) we will get better prompt caching.
15///
16/// In practice this seems to work. The largest tools have 3 parameters so on
17/// average we'll take 6 attempts to get sorted order.
18fn create_sorted_hashmap(sorted: BTreeMap<String, Document>) -> HashMap<String, Document> {
19    const MAX_RETRIES: usize = 100;
20
21    for _ in 0..MAX_RETRIES {
22        let mut map = HashMap::new();
23        for (k, v) in &sorted {
24            map.insert(k.clone(), v.clone());
25        }
26
27        let keys: Vec<_> = map.keys().collect();
28        let sorted_keys: Vec<_> = sorted.keys().collect();
29
30        if keys == sorted_keys {
31            return map;
32        }
33    }
34
35    let mut map = HashMap::new();
36    for (k, v) in sorted {
37        map.insert(k, v);
38    }
39    map
40}
41
42/// Convert a serde_json::Value to an AWS Document
43pub fn to_doc(value: Value) -> Document {
44    match value {
45        Value::Null => Document::Null,
46        Value::Bool(b) => Document::Bool(b),
47        Value::Number(n) => {
48            if let Some(u) = n.as_u64() {
49                Document::Number(aws_smithy_types::Number::PosInt(u))
50            } else if let Some(i) = n.as_i64() {
51                Document::Number(aws_smithy_types::Number::NegInt(i))
52            } else if let Some(f) = n.as_f64() {
53                Document::Number(aws_smithy_types::Number::Float(f))
54            } else {
55                Document::Null
56            }
57        }
58        Value::String(s) => Document::String(s),
59        Value::Array(arr) => Document::Array(arr.into_iter().map(to_doc).collect()),
60        Value::Object(obj) => {
61            // Sort keys alphabetically to ensure deterministic serialization for Bedrock prompt caching.
62            // Cache keys depend on exact byte-for-byte request equality.
63            // We retry HashMap creation until iteration order matches sorted order (works for small maps).
64            let sorted: BTreeMap<String, Document> =
65                obj.into_iter().map(|(k, v)| (k, to_doc(v))).collect();
66            Document::Object(create_sorted_hashmap(sorted))
67        }
68    }
69}
70
71/// Convert an AWS Document to a serde_json::Value
72pub fn from_doc(doc: Document) -> Value {
73    match doc {
74        Document::Null => Value::Null,
75        Document::Bool(b) => Value::Bool(b),
76        Document::Number(n) => match n {
77            aws_smithy_types::Number::PosInt(u) => Value::Number(u.into()),
78            aws_smithy_types::Number::NegInt(i) => Value::Number(i.into()),
79            aws_smithy_types::Number::Float(f) => serde_json::Number::from_f64(f)
80                .map(Value::Number)
81                .unwrap_or(Value::Null),
82        },
83        Document::String(s) => Value::String(s),
84        Document::Array(arr) => Value::Array(arr.into_iter().map(from_doc).collect()),
85        Document::Object(obj) => {
86            Value::Object(obj.into_iter().map(|(k, v)| (k, from_doc(v))).collect())
87        }
88    }
89}
tycode_core/ai/json.rs

tycode_core/ai/
json.rs
