limit_llm/
cache.rs

1//! Helper functions for applying cache control to messages.
2
3use crate::config::CacheSettings;
4use crate::types::{CacheControl, Message, Role};
5
6/// Apply cache control to strategic message positions.
7///
8/// Strategy (Anthropic/OpenAI best practices):
9/// 1. System message (if present) - first message if role is System
10/// 2. Last user message - enables caching of conversation context
11///
12/// Returns a new Vec with cache_control applied where appropriate.
13pub fn apply_cache_control(messages: &[Message], settings: &CacheSettings) -> Vec<Message> {
14    if !settings.is_enabled() {
15        return messages.to_vec();
16    }
17
18    let cache_control = if settings.is_long_retention() {
19        CacheControl::ephemeral_long()
20    } else {
21        CacheControl::ephemeral()
22    };
23
24    let mut result = messages.to_vec();
25    let len = result.len();
26
27    // Apply to system message (first message if role is System)
28    if len > 0 && result[0].role == Role::System {
29        result[0].cache_control = Some(cache_control.clone());
30    }
31
32    // Apply to last user message
33    for i in (0..len).rev() {
34        if result[i].role == Role::User {
35            result[i].cache_control = Some(cache_control.clone());
36            break;
37        }
38    }
39
40    result
41}
42
43#[cfg(test)]
44mod tests {
45    use super::*;
46
47    #[test]
48    fn test_no_caching_when_disabled() {
49        let settings = CacheSettings {
50            retention: "none".to_string(),
51        };
52        let messages = vec![Message {
53            role: Role::User,
54            content: Some("Hello".to_string()),
55            tool_calls: None,
56            tool_call_id: None,
57            cache_control: None,
58        }];
59
60        let result = apply_cache_control(&messages, &settings);
61        assert!(result[0].cache_control.is_none());
62    }
63
64    #[test]
65    fn test_cache_applied_to_user_message() {
66        let settings = CacheSettings {
67            retention: "short".to_string(),
68        };
69        let messages = vec![Message {
70            role: Role::User,
71            content: Some("Hello".to_string()),
72            tool_calls: None,
73            tool_call_id: None,
74            cache_control: None,
75        }];
76
77        let result = apply_cache_control(&messages, &settings);
78        assert!(result[0].cache_control.is_some());
79    }
80
81    #[test]
82    fn test_cache_applied_to_system_and_last_user() {
83        let settings = CacheSettings {
84            retention: "short".to_string(),
85        };
86        let messages = vec![
87            Message {
88                role: Role::System,
89                content: Some("System prompt".to_string()),
90                tool_calls: None,
91                tool_call_id: None,
92                cache_control: None,
93            },
94            Message {
95                role: Role::User,
96                content: Some("Hello".to_string()),
97                tool_calls: None,
98                tool_call_id: None,
99                cache_control: None,
100            },
101            Message {
102                role: Role::Assistant,
103                content: Some("Hi!".to_string()),
104                tool_calls: None,
105                tool_call_id: None,
106                cache_control: None,
107            },
108            Message {
109                role: Role::User,
110                content: Some("How are you?".to_string()),
111                tool_calls: None,
112                tool_call_id: None,
113                cache_control: None,
114            },
115        ];
116
117        let result = apply_cache_control(&messages, &settings);
118
119        // System message should have cache
120        assert!(result[0].cache_control.is_some());
121        // First user message should NOT have cache (only last)
122        assert!(result[1].cache_control.is_none());
123        // Assistant should NOT have cache
124        assert!(result[2].cache_control.is_none());
125        // Last user message should have cache
126        assert!(result[3].cache_control.is_some());
127    }
128
129    #[test]
130    fn test_long_retention_ttl() {
131        let settings = CacheSettings {
132            retention: "long".to_string(),
133        };
134        let messages = vec![Message {
135            role: Role::User,
136            content: Some("Hello".to_string()),
137            tool_calls: None,
138            tool_call_id: None,
139            cache_control: None,
140        }];
141
142        let result = apply_cache_control(&messages, &settings);
143        let cc = result[0].cache_control.as_ref().unwrap();
144        assert_eq!(cc.ttl, Some("1h".to_string()));
145    }
146}
limit_llm/cache.rs

limit_llm/
cache.rs
