agent-runtime-kernel 0.1.1

Internal model, context, and Tool execution Kernel for Agent Runtime
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

use runtime_ports::{ModelGenerationOptions, ModelMessage, ModelToolDefinition};
use serde::Serialize;
use sha2::{Digest, Sha256};

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct ContextMeasurement {
    pub(crate) message_tokens: u64,
    pub(crate) tool_schema_tokens: u64,
    pub(crate) protocol_overhead_tokens: u64,
    pub(crate) reserved_output_tokens: u64,
    pub(crate) total_tokens: u64,
}

pub(crate) fn measure_context(
    messages: &[ModelMessage],
    tools: &[ModelToolDefinition],
    protocol_overhead_tokens: u64,
    reserved_output_tokens: u32,
) -> ContextMeasurement {
    let message_tokens = estimated_serialized_tokens(messages);
    let tool_schema_tokens = estimated_serialized_tokens(tools);
    let reserved_output_tokens = u64::from(reserved_output_tokens);
    ContextMeasurement {
        message_tokens,
        tool_schema_tokens,
        protocol_overhead_tokens,
        reserved_output_tokens,
        total_tokens: message_tokens
            .saturating_add(tool_schema_tokens)
            .saturating_add(protocol_overhead_tokens)
            .saturating_add(reserved_output_tokens),
    }
}

pub(crate) fn model_request_digest(
    messages: &[ModelMessage],
    tools: &[ModelToolDefinition],
    generation: &ModelGenerationOptions,
) -> String {
    let mut digest = Sha256::new();
    digest.update(serde_json::to_vec(messages).unwrap_or_default());
    digest.update([0]);
    digest.update(serde_json::to_vec(tools).unwrap_or_default());
    digest.update([0]);
    digest.update(serde_json::to_vec(generation).unwrap_or_default());
    format!("sha256:{:x}", digest.finalize())
}

fn estimated_serialized_tokens<T: Serialize + ?Sized>(value: &T) -> u64 {
    let bytes = serde_json::to_vec(value)
        .map(|encoded| encoded.len() as u64)
        .unwrap_or_default();
    bytes.saturating_add(3) / 4
}

#[cfg(test)]
mod tests {
    use runtime_ports::{ModelMessage, ModelRole, ModelToolDefinition};
    use serde_json::json;

    use super::*;

    #[test]
    fn complete_measurement_counts_tools_protocol_and_reserved_output() {
        let messages = vec![ModelMessage::text(ModelRole::User, "task")];
        let tools = vec![ModelToolDefinition {
            name: "run".into(),
            description: "execute a command".into(),
            input_schema: json!({"type": "object"}),
        }];

        let without_tools = measure_context(&messages, &[], 0, 0);
        let complete = measure_context(&messages, &tools, 17, 31);

        assert!(complete.tool_schema_tokens > 0);
        assert_eq!(complete.protocol_overhead_tokens, 17);
        assert_eq!(complete.reserved_output_tokens, 31);
        assert!(complete.total_tokens > without_tools.total_tokens + 47);
    }

    #[test]
    fn request_digest_changes_with_tool_projection() {
        let messages = vec![ModelMessage::text(ModelRole::User, "task")];
        let one = ModelToolDefinition {
            name: "one".into(),
            description: String::new(),
            input_schema: json!({}),
        };
        let two = ModelToolDefinition {
            name: "two".into(),
            ..one.clone()
        };
        assert_ne!(
            model_request_digest(
                &messages,
                &[one],
                &ModelGenerationOptions {
                    max_output_tokens: Some(8),
                    ..Default::default()
                },
            ),
            model_request_digest(
                &messages,
                &[two],
                &ModelGenerationOptions {
                    max_output_tokens: Some(8),
                    ..Default::default()
                },
            )
        );
    }

    #[test]
    fn request_digest_changes_with_generation_controls() {
        let messages = vec![ModelMessage::text(ModelRole::User, "task")];
        let baseline = ModelGenerationOptions::default();
        let sampled = ModelGenerationOptions {
            temperature: Some(0.7),
            ..Default::default()
        };
        assert_ne!(
            model_request_digest(&messages, &[], &baseline),
            model_request_digest(&messages, &[], &sampled)
        );
    }
}