agent-runtime-kernel 0.1.0

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

use runtime_ports::{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],
    max_output_tokens: Option<u32>,
) -> 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(max_output_tokens.unwrap_or_default().to_be_bytes());
    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], Some(8)),
            model_request_digest(&messages, &[two], Some(8))
        );
    }
}