tea-context 0.1.0

Deterministic prompt modules, context providers, and compiler for tea-rs
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
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

use crate::common;

use common::{module, segment};
use tea_context::{
    BudgetBehavior, ContextErrorCode, PromptAuthority, PromptBudget, PromptCompiler,
    PromptDiagnosticCode, SegmentDisposition, TRUNCATION_MARKER, estimate_tokens,
};

#[test]
fn required_overflow_fails_and_omit_is_diagnostic() {
    let required = PromptCompiler
        .compile(
            [module(
                "required",
                PromptAuthority::Kernel,
                0,
                vec![segment("required", "12345", BudgetBehavior::Required)],
            )],
            PromptBudget::new(4, 100).unwrap(),
        )
        .unwrap_err();
    assert_eq!(required.code(), ContextErrorCode::BudgetExceeded);

    let prompt = PromptCompiler
        .compile(
            [module(
                "mixed",
                PromptAuthority::Product,
                0,
                vec![
                    segment("keep", "ok", BudgetBehavior::Required),
                    segment("omit", "too long", BudgetBehavior::Omit),
                ],
            )],
            PromptBudget::new(4, 100).unwrap(),
        )
        .unwrap();
    assert_eq!(prompt.text(), "ok");
    assert_eq!(
        prompt.diagnostics()[0].code(),
        PromptDiagnosticCode::OmittedForBudget
    );
}

#[test]
fn truncation_is_utf8_safe_and_uses_fixed_marker() {
    let prompt = PromptCompiler
        .compile(
            [module(
                "truncate",
                PromptAuthority::Workspace,
                0,
                vec![segment(
                    "unicode",
                    "你好abcdefghijklmnop",
                    BudgetBehavior::Truncate,
                )],
            )],
            PromptBudget::new(TRUNCATION_MARKER.len() + 5, 100).unwrap(),
        )
        .unwrap();
    assert!(prompt.text().ends_with(TRUNCATION_MARKER));
    assert!(std::str::from_utf8(prompt.text().as_bytes()).is_ok());
    assert_eq!(
        prompt.diagnostics()[0].code(),
        PromptDiagnosticCode::TruncatedForBudget
    );
    assert_eq!(
        prompt.inspection()[0].disposition(),
        SegmentDisposition::Truncated
    );
}

#[test]
fn separators_count_toward_byte_and_token_budgets() {
    let modules = [module(
        "two",
        PromptAuthority::Product,
        0,
        vec![
            segment("one", "abc", BudgetBehavior::Required),
            segment("two", "def", BudgetBehavior::Required),
        ],
    )];
    let prompt = PromptCompiler
        .compile(modules.clone(), PromptBudget::new(8, 3).unwrap())
        .unwrap();
    assert_eq!(prompt.text(), "abc\n\ndef");
    assert_eq!(prompt.bytes(), 8);
    assert_eq!(prompt.estimated_tokens(), 3);
    assert_eq!(estimate_tokens(8), 3);
    assert_eq!(
        PromptCompiler
            .compile(modules, PromptBudget::new(7, 3).unwrap())
            .unwrap_err()
            .code(),
        ContextErrorCode::BudgetExceeded
    );
}

#[test]
fn fully_optional_prompt_may_compile_empty_with_diagnostics() {
    let prompt = PromptCompiler
        .compile(
            [module(
                "optional",
                PromptAuthority::Session,
                0,
                vec![segment("optional", "too long", BudgetBehavior::Omit)],
            )],
            PromptBudget::new(1, 1).unwrap(),
        )
        .unwrap();
    assert!(prompt.text().is_empty());
    assert_eq!(prompt.estimated_tokens(), 0);
    assert_eq!(
        prompt.diagnostics()[0].code(),
        PromptDiagnosticCode::OmittedForBudget
    );
}

#[test]
fn truncation_without_marker_space_omits_deterministically() {
    let prompt = PromptCompiler
        .compile(
            [module(
                "mixed",
                PromptAuthority::Product,
                0,
                vec![
                    segment("keep", "x", BudgetBehavior::Required),
                    segment("truncate", "abcdef", BudgetBehavior::Truncate),
                ],
            )],
            PromptBudget::new(4, 100).unwrap(),
        )
        .unwrap();
    assert_eq!(prompt.text(), "x");
    assert_eq!(
        prompt.inspection()[1].disposition(),
        SegmentDisposition::OmittedForBudget
    );
}