feat: add cache-aligned prompt builder for LLM API cost savings
PROBLEM: - PromptBuilder.build() puts everything in a single user message - System + query + memory + chunk all change together - LLM prompt caching gets 0% hits (entire message differs per call) - For a 50-chunk ingestion run, we pay full input price 50 times SOLUTION: PromptBuilder.build_cache_aligned() - Splits prompt into 3 separate messages: 1. SYSTEM: instructions (stable across ALL calls) → CACHED 2. USER[0]: query/problem (stable per run) → CACHED 3. USER[1]: memory + chunk (varies per call) → not cached - Cache prefix (system + query) reused across all chunks in a run - Estimated 30-70% cache hit ratio depending on chunk sizes - ~50% input token cost savings for multi-chunk ingestion TEMPLATES: - templates/gru-mem-system.txt (instructions only, 840B) - templates/gru-mem-query.txt (problem wrapper, 29B) - templates/gru-mem-turn.txt (memory + section, 57B) - templates/gru-mem.txt (legacy, unchanged) API: - PromptBuilder::build() — legacy, backward compatible - PromptBuilder::build_cache_aligned() → PromptMessages - PromptMessages.cache_prefix_tokens() — cacheable token count - PromptMessages.total_tokens() — total estimated tokens - PromptMessages.headroom() — tokens available for response TESTS: 11 unit + 3 integration = 14 new tests - test_cache_aligned_produces_two_user_messages - test_cache_prefix_is_stable_across_chunks - test_cache_prefix_is_stable_across_memory_changes - test_cache_prefix_tokens_positive - test_headroom_positive_under_budget - test_legacy_build_still_works - test_cache_aligned_contains_query - test_cache_aligned_memory/chunk_budget_exceeded - a8_cache_prefix_stable_across_50_chunks - a9_cache_aligned_headroom - a10_cache_savings_estimate TOTAL: 64 mem-core tests passing (52 unit + 12 integration)
This commit is contained in:
+111
-1
@@ -1,5 +1,5 @@
|
||||
use mem_core::prompt::PromptBuilder;
|
||||
use mem_core::{Chunk, Query, Record, Role, Provenance};
|
||||
use mem_core::{Chunk, PromptMessages, Query, Record, Role, Provenance};
|
||||
use time::OffsetDateTime;
|
||||
|
||||
fn make_chunk(records: Vec<(Role, &str)>) -> Chunk {
|
||||
@@ -206,3 +206,113 @@ fn a7_budget_headroom() {
|
||||
30720
|
||||
);
|
||||
}
|
||||
|
||||
// ── Cache alignment integration tests ──────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn a8_cache_prefix_stable_across_50_chunks() {
|
||||
// Simulate a real ingestion run: same query, 50 different chunks.
|
||||
// The cache prefix (system + query) must be identical every time.
|
||||
let query = Query {
|
||||
id: "arch-decisions".to_string(),
|
||||
question: "What architectural decisions were made and why?".to_string(),
|
||||
exit_gate: false,
|
||||
};
|
||||
|
||||
let mut prefixes: Vec<(String, String)> = Vec::new();
|
||||
|
||||
for i in 0..50 {
|
||||
let content = format!("Evidence chunk number {} with unique content {}", i, "x".repeat(i * 10));
|
||||
let chunk = make_chunk(vec![(Role::User, &content)]);
|
||||
let memory = if i == 0 {
|
||||
None
|
||||
} else {
|
||||
Some("Accumulated memory from previous turns")
|
||||
};
|
||||
|
||||
let msgs = PromptBuilder::build_cache_aligned(&query, memory, &chunk)
|
||||
.expect("Should build cache-aligned prompt");
|
||||
|
||||
prefixes.push((msgs.system.clone(), msgs.user_messages[0].clone()));
|
||||
}
|
||||
|
||||
// All 50 system messages must be identical
|
||||
for (i, (system, _)) in prefixes.iter().enumerate() {
|
||||
assert_eq!(
|
||||
system, &prefixes[0].0,
|
||||
"System message must be stable (chunk {})", i
|
||||
);
|
||||
}
|
||||
|
||||
// All 50 query messages must be identical
|
||||
for (i, (_, query_msg)) in prefixes.iter().enumerate() {
|
||||
assert_eq!(
|
||||
query_msg, &prefixes[0].1,
|
||||
"Query message must be stable (chunk {})", i
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a9_cache_aligned_headroom() {
|
||||
// Cache-aligned prompts should also have positive headroom
|
||||
let query = Query {
|
||||
id: "test".to_string(),
|
||||
question: "Test question?".to_string(),
|
||||
exit_gate: false,
|
||||
};
|
||||
|
||||
let chunk_content = "x".repeat(4000);
|
||||
let chunk = make_chunk(vec![(Role::User, &chunk_content)]);
|
||||
|
||||
let msgs = PromptBuilder::build_cache_aligned(&query, None, &chunk)
|
||||
.expect("Should build cache-aligned prompt");
|
||||
|
||||
assert!(
|
||||
msgs.headroom() > 0,
|
||||
"Cache-aligned prompt should have positive headroom: {} tokens",
|
||||
msgs.headroom()
|
||||
);
|
||||
|
||||
// Cache prefix should be meaningful (not zero)
|
||||
assert!(
|
||||
msgs.cache_prefix_tokens() > 50,
|
||||
"Cache prefix should be >50 tokens, got {}",
|
||||
msgs.cache_prefix_tokens()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a10_cache_savings_estimate() {
|
||||
// Demonstrate the savings: compare total tokens vs cached tokens
|
||||
// across a simulated 20-chunk run
|
||||
let query = Query {
|
||||
id: "test".to_string(),
|
||||
question: "What are the deployment patterns?".to_string(),
|
||||
exit_gate: false,
|
||||
};
|
||||
|
||||
let mut total_input_tokens = 0usize;
|
||||
let mut total_cached_tokens = 0usize;
|
||||
|
||||
for i in 0..20 {
|
||||
let content = format!("Deployment evidence chunk {} with details", i);
|
||||
let chunk = make_chunk(vec![(Role::User, &content)]);
|
||||
|
||||
let msgs = PromptBuilder::build_cache_aligned(&query, None, &chunk)
|
||||
.expect("Should build");
|
||||
|
||||
total_input_tokens += msgs.total_tokens();
|
||||
total_cached_tokens += msgs.cache_prefix_tokens();
|
||||
}
|
||||
|
||||
// Cache prefix should be a significant portion of total input
|
||||
let cache_ratio = total_cached_tokens as f64 / total_input_tokens as f64;
|
||||
assert!(
|
||||
cache_ratio > 0.3,
|
||||
"Cache ratio should be >30%, got {:.1}% ({}/{} tokens)",
|
||||
cache_ratio * 100.0,
|
||||
total_cached_tokens,
|
||||
total_input_tokens
|
||||
);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user