fix: improve LLM caching (#1088)

Author: tusharmath

crates/forge_provider/src/forge_provider/transformers/set_cache.rs

@@ -8,15 +8,21 @@ pub struct SetCache;
 impl Transformer for SetCache {
     type Value = Request;
 
+    /// Caches the last 2 eligible messages to optimize API performance. System
+    /// messages are always eligible, User messages are eligible (but
+    /// consecutive User messages are consolidated to only the last one),
+    /// and Assistant messages are never cached but reset User message
+    /// sequences.
     fn transform(&mut self, mut request: Self::Value) -> Self::Value {
         if let Some(messages) = request.messages.as_mut() {
             let mut last_was_user = false;
             let mut cache_positions = Vec::new();
             for (i, message) in messages.iter().enumerate() {
                 if message.role == Role::User {
-                    if !last_was_user {
-                        cache_positions.push(i);
+                    if last_was_user {
+                        cache_positions.pop();
                     }
+                    cache_positions.push(i);
                     last_was_user = true;
                 } else if message.role == Role::Assistant {
                     last_was_user = false;
@@ -26,7 +32,7 @@ impl Transformer for SetCache {
                 }
             }
 
-            for pos in cache_positions.into_iter().rev().skip(2).take(2) {
+            for pos in cache_positions.into_iter().rev().take(2) {
                 if let Some(ref content) = messages[pos].content {
                     messages[pos].content = Some(content.clone().cached());
                 }
@@ -117,27 +123,31 @@ mod tests {
     #[test]
     fn test_transformation() {
         let actual = create_test_context("suu");
-        let expected = "suu";
+        let expected = "[su[u"; // FIXME
         assert_eq!(actual, expected);
 
         let actual = create_test_context("suua");
-        let expected = "suua";
+        let expected = "[su[ua";
         assert_eq!(actual, expected);
 
         let actual = create_test_context("suuau");
-        let expected = "[suuau";
+        let expected = "su[ua[u";
         assert_eq!(actual, expected);
 
         let actual = create_test_context("suuauu");
-        let expected = "[suuauu";
+        let expected = "su[uau[u";
         assert_eq!(actual, expected);
 
         let actual = create_test_context("suuauuaaau");
-        let expected = "[s[uuauuaaau";
+        let expected = "suuau[uaaa[u";
         assert_eq!(actual, expected);
 
         let actual = create_test_context("suuauuaaauauau");
-        let expected = "suua[uuaaa[uauau";
+        let expected = "suuauuaaaua[ua[u";
+        assert_eq!(actual, expected);
+
+        let actual = create_test_context("suuaaaaaaaaaaa");
+        let expected = "[su[uaaaaaaaaaaa";
         assert_eq!(actual, expected);
     }
 }