lean_ctx/core/context_kernel/
provider_traces.rs

1#[cfg(test)]
2mod tests {
3    use super::super::provider_parity::{detect_provider, envelope_from_usage};
4    use super::super::token_envelope::ProviderKind;
5    use serde_json::json;
6
7    #[test]
8    fn openai_standard() {
9        let provider = detect_provider("https://api.openai.com/v1/chat/completions");
10        assert_eq!(provider, ProviderKind::OpenAi);
11
12        let envelope = envelope_from_usage(
13            ProviderKind::OpenAi,
14            "gpt-4",
15            &json!({"prompt_tokens": 150, "completion_tokens": 80}),
16        );
17        assert_eq!(envelope.input_tokens, 150);
18        assert_eq!(envelope.output_tokens, 80);
19    }
20
21    #[test]
22    fn openai_with_cache() {
23        let envelope = envelope_from_usage(
24            ProviderKind::OpenAi,
25            "gpt-4",
26            &json!({
27                "prompt_tokens": 200,
28                "completion_tokens": 50,
29                "prompt_tokens_details": {"cached_tokens": 100}
30            }),
31        );
32
33        assert_eq!(envelope.cache_read_tokens, 100);
34    }
35
36    #[test]
37    fn openai_with_reasoning() {
38        let envelope = envelope_from_usage(
39            ProviderKind::OpenAi,
40            "gpt-4",
41            &json!({
42                "prompt_tokens": 100,
43                "completion_tokens": 500,
44                "completion_tokens_details": {"reasoning_tokens": 400}
45            }),
46        );
47
48        assert_eq!(envelope.reasoning_tokens, 400);
49    }
50
51    #[test]
52    fn anthropic_standard() {
53        let provider = detect_provider("https://api.anthropic.com/v1/messages");
54        assert_eq!(provider, ProviderKind::Anthropic);
55
56        let envelope = envelope_from_usage(
57            ProviderKind::Anthropic,
58            "claude",
59            &json!({"input_tokens": 120, "output_tokens": 60}),
60        );
61        assert_eq!(envelope.input_tokens, 120);
62        assert_eq!(envelope.output_tokens, 60);
63    }
64
65    #[test]
66    fn anthropic_with_cache() {
67        let envelope = envelope_from_usage(
68            ProviderKind::Anthropic,
69            "claude",
70            &json!({
71                "input_tokens": 200,
72                "output_tokens": 50,
73                "cache_read_input_tokens": 80,
74                "cache_creation_input_tokens": 30
75            }),
76        );
77
78        assert_eq!(envelope.cache_read_tokens, 80);
79        assert_eq!(envelope.cache_write_tokens, 30);
80    }
81
82    #[test]
83    fn gemini_usage() {
84        let provider = detect_provider("https://generativelanguage.googleapis.com");
85        assert_eq!(provider, ProviderKind::Gemini);
86
87        let envelope = envelope_from_usage(
88            ProviderKind::Gemini,
89            "gemini",
90            &json!({
91                "promptTokenCount": 100,
92                "candidatesTokenCount": 40,
93                "cachedContentTokenCount": 20
94            }),
95        );
96        assert_eq!(envelope.input_tokens, 100);
97        assert_eq!(envelope.output_tokens, 40);
98        assert_eq!(envelope.cache_read_tokens, 20);
99    }
100
101    #[test]
102    fn bedrock_usage() {
103        let provider = detect_provider("https://bedrock-runtime.us-east-1.amazonaws.com");
104        assert_eq!(provider, ProviderKind::Bedrock);
105
106        let envelope = envelope_from_usage(
107            ProviderKind::Bedrock,
108            "bedrock",
109            &json!({"inputTokens": 90, "outputTokens": 30}),
110        );
111        assert_eq!(envelope.input_tokens, 90);
112        assert_eq!(envelope.output_tokens, 30);
113    }
114
115    #[test]
116    fn azure_openai() {
117        let provider = detect_provider("https://my-resource.openai.azure.com");
118        assert_eq!(provider, ProviderKind::Azure);
119
120        let envelope = envelope_from_usage(
121            ProviderKind::Azure,
122            "gpt-4",
123            &json!({"prompt_tokens": 100, "completion_tokens": 50}),
124        );
125        assert_eq!(envelope.input_tokens, 100);
126        assert_eq!(envelope.output_tokens, 50);
127    }
128
129    #[test]
130    fn empty_usage_no_panic() {
131        let envelope = envelope_from_usage(ProviderKind::Unknown, "x", &json!({}));
132
133        assert_eq!(envelope.provider, ProviderKind::Unknown);
134        assert_eq!(envelope.input_tokens, 0);
135        assert_eq!(envelope.output_tokens, 0);
136        assert_eq!(envelope.cache_read_tokens, 0);
137        assert_eq!(envelope.cache_write_tokens, 0);
138        assert_eq!(envelope.reasoning_tokens, 0);
139    }
140
141    #[test]
142    fn null_fields_safe() {
143        let envelope = envelope_from_usage(
144            ProviderKind::OpenAi,
145            "gpt-4",
146            &json!({
147                "prompt_tokens": null,
148                "completion_tokens": "not_a_number"
149            }),
150        );
151
152        assert_eq!(envelope.input_tokens, 0);
153        assert_eq!(envelope.output_tokens, 0);
154    }
155}
lean_ctx/core/context_kernel/provider_traces.rs

lean_ctx/core/context_kernel/
provider_traces.rs
