llm-memory-graph 0.1.0

Graph-based context-tracking and prompt-lineage database for LLM systems
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
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166

//! Integration tests for Observatory functionality

use llm_memory_graph::{
    engine::AsyncMemoryGraph,
    observatory::{InMemoryPublisher, ObservatoryConfig},
    types::{Config, TokenUsage},
};
use std::sync::Arc;

#[tokio::test]
async fn test_observatory_integration() {
    let dir = tempfile::tempdir().unwrap();
    let config = Config::new(dir.path());
    let publisher = Arc::new(InMemoryPublisher::new());

    let obs_config = ObservatoryConfig::new().enabled().with_metrics(true);

    let graph = AsyncMemoryGraph::with_observatory(config, Some(publisher.clone()), obs_config)
        .await
        .unwrap();

    // Create session
    let session = graph.create_session().await.unwrap();

    // Add prompt
    let prompt_id = graph
        .add_prompt(session.id, "Test prompt".to_string(), None)
        .await
        .unwrap();

    // Add response
    let usage = TokenUsage::new(10, 20);
    let response_id = graph
        .add_response(prompt_id, "Test response".to_string(), usage, None)
        .await
        .unwrap();

    // Give async events time to publish
    tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;

    // Check events were published
    let events = publisher.get_events().await;
    assert!(!events.is_empty(), "Expected events to be published");

    // Should have at least PromptSubmitted event
    let prompt_events = publisher.get_events_by_type("prompt_submitted").await;
    assert_eq!(
        prompt_events.len(),
        1,
        "Expected exactly one PromptSubmitted event"
    );

    // Check metrics
    let metrics = graph.get_metrics().expect("Metrics should be available");
    assert_eq!(metrics.prompts_submitted, 1);
    assert!(
        metrics.nodes_created >= 2,
        "Expected at least 2 nodes created (session + prompt, possibly response)"
    );
}

#[tokio::test]
async fn test_observatory_disabled() {
    let dir = tempfile::tempdir().unwrap();
    let config = Config::new(dir.path());

    let obs_config = ObservatoryConfig::new(); // Not enabled

    let graph = AsyncMemoryGraph::with_observatory(config, None, obs_config)
        .await
        .unwrap();

    // Create session and prompt
    let session = graph.create_session().await.unwrap();
    graph
        .add_prompt(session.id, "Test".to_string(), None)
        .await
        .unwrap();

    // Metrics should still work if enabled
    let metrics = graph.get_metrics();
    assert!(
        metrics.is_some(),
        "Metrics should be available even when events are disabled"
    );
}

#[tokio::test]
async fn test_metrics_collection() {
    let dir = tempfile::tempdir().unwrap();
    let config = Config::new(dir.path());
    let publisher = Arc::new(InMemoryPublisher::new());

    let obs_config = ObservatoryConfig::new().enabled().with_metrics(true);

    let graph = AsyncMemoryGraph::with_observatory(config, Some(publisher), obs_config)
        .await
        .unwrap();

    let session = graph.create_session().await.unwrap();

    // Add multiple prompts
    for i in 0..10 {
        graph
            .add_prompt(session.id, format!("Prompt {}", i), None)
            .await
            .unwrap();
    }

    let metrics = graph.get_metrics().unwrap();
    assert_eq!(metrics.prompts_submitted, 10);
    assert!(
        metrics.avg_write_latency_ms > 0.0,
        "Average write latency should be recorded"
    );
}

#[tokio::test]
async fn test_concurrent_event_publishing() {
    let dir = tempfile::tempdir().unwrap();
    let config = Config::new(dir.path());
    let publisher = Arc::new(InMemoryPublisher::new());

    let obs_config = ObservatoryConfig::new().enabled();

    let graph = Arc::new(
        AsyncMemoryGraph::with_observatory(config, Some(publisher.clone()), obs_config)
            .await
            .unwrap(),
    );

    let session = graph.create_session().await.unwrap();

    // Create 50 prompts concurrently
    let mut handles = vec![];
    for i in 0..50 {
        let graph_clone = Arc::clone(&graph);
        let session_id = session.id;
        let handle = tokio::spawn(async move {
            graph_clone
                .add_prompt(session_id, format!("Concurrent prompt {}", i), None)
                .await
        });
        handles.push(handle);
    }

    // Wait for all
    for handle in handles {
        handle.await.unwrap().unwrap();
    }

    // Give events time to publish
    tokio::time::sleep(tokio::time::Duration::from_millis(200)).await;

    // Check metrics
    let metrics = graph.get_metrics().unwrap();
    assert_eq!(metrics.prompts_submitted, 50);

    // Check events
    let events = publisher.get_events().await;
    let prompt_events: Vec<_> = events
        .iter()
        .filter(|e| e.event_type() == "prompt_submitted")
        .collect();
    assert_eq!(prompt_events.len(), 50);
}