umi-memory 0.1.0

Memory library for AI agents with deterministic simulation testing
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

//! Retrieval Prompts - LLM Prompt Templates
//!
//! TigerStyle: Structured prompts with clear output format.

/// Query rewrite prompt template.
///
/// Placeholder:
/// - `{query}` - The search query to rewrite
pub const QUERY_REWRITE_PROMPT: &str = r#"Rewrite this search query into 2-3 variations that would help find relevant memories.

Query: {query}

Return as JSON array of strings. Example: ["variation 1", "variation 2", "variation 3"]
Only return the JSON array, nothing else."#;

/// Build the query rewrite prompt.
///
/// # Arguments
/// - `query` - The search query to rewrite
///
/// # Returns
/// The formatted prompt string.
#[must_use]
pub fn build_query_rewrite_prompt(query: &str) -> String {
    QUERY_REWRITE_PROMPT.replace("{query}", query)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_build_prompt() {
        let prompt = build_query_rewrite_prompt("Who works at Acme?");

        assert!(prompt.contains("Who works at Acme?"));
        assert!(prompt.contains("Query:"));
        assert!(prompt.contains("JSON array"));
    }

    #[test]
    fn test_prompt_structure() {
        let prompt = build_query_rewrite_prompt("test");

        // Verify key structural elements
        assert!(prompt.contains("Rewrite this search query"));
        assert!(prompt.contains("Query: test"));
        assert!(prompt.contains("2-3 variations"));
        assert!(prompt.contains("Only return the JSON"));
    }
}