Skip to main content

codex_context_fragments/
additional_context.rs

1use codex_utils_string::truncate_middle_with_token_budget;
2
3use crate::ContextualUserFragment;
4
5const MAX_ADDITIONAL_CONTEXT_VALUE_TOKENS: usize = 1_000;
6const ADDITIONAL_CONTEXT_END_MARKER_SUFFIX: &str = ">";
7const ADDITIONAL_CONTEXT_START_MARKER_PREFIX: &str = "<external_";
8
9#[derive(Debug, Clone, PartialEq, Eq)]
10pub struct AdditionalContextUserFragment {
11    key: String,
12    value: String,
13}
14
15impl AdditionalContextUserFragment {
16    pub fn new(key: String, value: String) -> Self {
17        Self { key, value }
18    }
19}
20
21impl ContextualUserFragment for AdditionalContextUserFragment {
22    fn role(&self) -> &'static str {
23        "user"
24    }
25
26    fn markers(&self) -> (&'static str, &'static str) {
27        Self::type_markers()
28    }
29
30    fn type_markers() -> (&'static str, &'static str) {
31        (
32            ADDITIONAL_CONTEXT_START_MARKER_PREFIX,
33            ADDITIONAL_CONTEXT_END_MARKER_SUFFIX,
34        )
35    }
36
37    fn matches_text(text: &str) -> bool {
38        let trimmed = text.trim();
39        let Some(rest) = trimmed.strip_prefix(ADDITIONAL_CONTEXT_START_MARKER_PREFIX) else {
40            return false;
41        };
42        let Some((key, value_and_close)) = rest.split_once(ADDITIONAL_CONTEXT_END_MARKER_SUFFIX)
43        else {
44            return false;
45        };
46
47        value_and_close.ends_with(&format!("</external_{key}>"))
48    }
49
50    fn body(&self) -> String {
51        additional_context_body(&self.key, &self.value)
52    }
53}
54
55#[derive(Debug, Clone, PartialEq, Eq)]
56pub struct AdditionalContextDeveloperFragment {
57    key: String,
58    value: String,
59}
60
61impl AdditionalContextDeveloperFragment {
62    pub fn new(key: String, value: String) -> Self {
63        Self { key, value }
64    }
65}
66
67impl ContextualUserFragment for AdditionalContextDeveloperFragment {
68    fn role(&self) -> &'static str {
69        "developer"
70    }
71
72    fn markers(&self) -> (&'static str, &'static str) {
73        Self::type_markers()
74    }
75
76    fn type_markers() -> (&'static str, &'static str) {
77        ("", "")
78    }
79
80    fn body(&self) -> String {
81        additional_context_developer_body(&self.key, &self.value)
82    }
83}
84
85fn additional_context_body(key: &str, value: &str) -> String {
86    let value = truncate_middle_with_token_budget(value, MAX_ADDITIONAL_CONTEXT_VALUE_TOKENS).0;
87    format!("{key}>{value}</external_{key}")
88}
89
90fn additional_context_developer_body(key: &str, value: &str) -> String {
91    let value = truncate_middle_with_token_budget(value, MAX_ADDITIONAL_CONTEXT_VALUE_TOKENS).0;
92    format!("<{key}>{value}</{key}>")
93}