atheneum 0.2.1

Agent coordination graph database - episodic and semantic memory for multi-agent workflows
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
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239

//! Tests for graph navigation — neighbors, subgraph, navigate, stats.
//!
//! These tests drive TDD for:
//!   - get_entity / get_edge (already exist)
//!   - get_neighbors (NEW)
//!   - get_subgraph (NEW)
//!   - navigate (semantic search → subgraph) (NEW)
//!   - graph_stats (NEW)
//!   - auto-index on discovery write (NEW)

use atheneum::graph::{AtheneumGraph, EdgeType};
use serde_json::json;

fn make_chain_graph() -> anyhow::Result<(AtheneumGraph, i64, i64, i64)> {
    let g = AtheneumGraph::open_in_memory()?;

    let agent = g.insert_agent("claude1", json!({}))?;
    let discovery = g.store_discovery(
        "claude1",
        "Symbol",
        "build_router",
        json!({"file":"src/http.rs","summary":"builds axum router"}),
    )?;
    let event = g.insert_event(
        "discovery-stored",
        json!({"agent":"claude1","discovery_id":discovery}),
    )?;

    // event --Created--> discovery
    g.insert_edge(event, discovery, EdgeType::Created, json!({}))?;
    // event --PerformedBy--> agent
    g.insert_edge(event, agent, EdgeType::PerformedBy, json!({}))?;

    Ok((g, agent, discovery, event))
}

// ============================================================================
// Basics (already exist — regression smoke)
// ============================================================================

#[test]
fn test_get_entity() {
    let (g, _agent, discovery, _event) = make_chain_graph().expect("setup");

    let e = g.get_entity(discovery).expect("get_entity");
    assert_eq!(e.kind, "Discovery");
    assert_eq!(e.name, "claude1: build_router");
}

#[test]
fn test_get_edge() {
    let g = AtheneumGraph::open_in_memory().expect("open");
    let a = g.insert_agent("a", json!({})).expect("insert");
    let b = g.insert_agent("b", json!({})).expect("insert");
    let eid = g
        .insert_edge(a, b, EdgeType::RelatedTo, json!({}))
        .expect("edge");

    let edge = g.get_edge(eid).expect("get_edge");
    assert_eq!(edge.from_id, a);
    assert_eq!(edge.to_id, b);
    assert_eq!(edge.edge_type, "related_to");
}

// ============================================================================
// Neighbors (new)
// ============================================================================

#[test]
fn test_get_neighbors_returns_outgoing_and_incoming() {
    let (g, _agent, discovery, _event) = make_chain_graph().expect("setup");

    let (outgoing, incoming) = g.get_neighbors(discovery).expect("get_neighbors");

    // discovery has 1 incoming edge from event (Created)
    assert_eq!(incoming.len(), 1, "discovery should have 1 incoming edge");
    assert_eq!(incoming[0].edge_type, "created");

    // discovery has no outgoing edges
    assert_eq!(outgoing.len(), 0, "discovery should have 0 outgoing edges");
}

#[test]
fn test_get_neighbors_for_event() {
    let (g, _agent, _discovery, event) = make_chain_graph().expect("setup");

    let (outgoing, incoming) = g.get_neighbors(event).expect("get_neighbors");

    // Event has 2 outgoing edges: Created → discovery, PerformedBy → agent
    assert_eq!(outgoing.len(), 2, "event should have 2 outgoing edges");
    let types: Vec<&str> = outgoing.iter().map(|e| e.edge_type.as_str()).collect();
    assert!(types.contains(&"created"));
    assert!(types.contains(&"performed_by"));

    // Event has no incoming edges
    assert_eq!(incoming.len(), 0);
}

// ============================================================================
// Subgraph (new)
// ============================================================================

#[test]
fn test_get_subgraph_depth_0() {
    let (g, _agent, discovery, _event) = make_chain_graph().expect("setup");

    let sg = g.get_subgraph(discovery, 0).expect("get_subgraph");
    assert_eq!(sg.entry.id, discovery);
    assert_eq!(sg.depth, 0);
    assert!(sg.edges.is_empty());
    assert!(sg.entities.is_empty() || sg.entities.len() == 1); // entry-only at depth 0
}

#[test]
fn test_get_subgraph_depth_1() {
    let (g, agent, discovery, event) = make_chain_graph().expect("setup");

    let sg = g.get_subgraph(discovery, 1).expect("get_subgraph");

    // depth 1 from discovery: should reach the event
    let entity_ids: Vec<i64> = sg.entities.iter().map(|e| e.id).collect();
    assert!(
        entity_ids.contains(&event),
        "depth 1 should include event (the event that created discovery)"
    );
    // depth 1 should NOT reach agent (agent is 2 hops from discovery)
    assert!(
        !entity_ids.contains(&agent),
        "depth 1 should NOT include agent"
    );

    // should have at least 1 edge
    assert!(!sg.edges.is_empty(), "subgraph should have edges");
    let edge_types: Vec<&str> = sg.edges.iter().map(|e| e.edge_type.as_str()).collect();
    assert!(edge_types.contains(&"created") || edge_types.contains(&"performed_by"));
}

#[test]
fn test_get_subgraph_depth_2() {
    let (g, agent, discovery, _event) = make_chain_graph().expect("setup");

    let sg = g.get_subgraph(discovery, 2).expect("get_subgraph");

    let entity_ids: Vec<i64> = sg.entities.iter().map(|e| e.id).collect();
    assert!(
        entity_ids.contains(&agent),
        "depth 2 should reach agent (discovery→event→agent)"
    );
}

// ============================================================================
// Navigate — semantic search → subgraph (new)
// ============================================================================

#[test]
fn test_navigate_finds_entry_points_and_walks() {
    let (g, _agent, _discovery, _event) = make_chain_graph().expect("setup");

    // Navigate on "router" should find build_router → walk graph
    let views = g
        .navigate("router construction axum", 5, 2, None)
        .expect("navigate");

    assert!(
        !views.is_empty(),
        "navigate should return at least one subgraph view"
    );
    let sg = &views[0];
    // The entry should be a discovery about build_router
    assert!(
        sg.entry.name.contains("build_router"),
        "entry should be build_router: {}",
        sg.entry.name
    );
    // Depth 2 should have walked to the event and agent
    let names: Vec<&str> = sg.entities.iter().map(|e| e.name.as_str()).collect();
    assert!(
        names.iter().any(|n| n.contains("claude1: build_router")),
        "subgraph should contain discovery entity"
    );
}

// ============================================================================
// Stats (new)
// ============================================================================

#[test]
fn test_graph_stats_returns_counts() {
    let g = AtheneumGraph::open_in_memory().expect("open");

    // Fresh DB: stats should exist but all counts are 0.
    let stats = g.graph_stats().expect("graph_stats");
    assert_eq!(stats.total_entities, 0, "fresh DB has 0 entities");
    assert_eq!(stats.entity_counts.len(), 0, "fresh DB has no kinds");

    g.insert_agent("new_agent", json!({"status":"active"}))
        .expect("insert");
    let stats_after = g.graph_stats().expect("graph_stats");

    let agent_count = stats_after
        .entity_counts
        .iter()
        .find(|(k, _)| k == "Agent")
        .map(|(_, c)| *c)
        .unwrap_or(0);
    assert_eq!(agent_count, 1, "Agent count should be 1 after insert");
    assert_eq!(stats_after.total_entities, 1, "total_entities should be 1");
}

// ============================================================================
// Auto-index on discovery write (new)
// ============================================================================

#[test]
fn test_discovery_auto_indexed() {
    let g = AtheneumGraph::open_in_memory().expect("open");

    g.store_discovery(
        "agent",
        "Symbol",
        "semantic_navigation",
        json!({"summary":"navigate with HNSW and graph traversals"}),
    )
    .expect("store_discovery");

    // After store_discovery we should be able to search WITHOUT manually calling build_search_index()
    let results = g
        .lexical_search("semantic navigation HNSW traversal", 5, None)
        .expect("semantic_search");

    assert!(
        !results.is_empty(),
        "auto-indexed discovery should be searchable immediately (got 0 results)"
    );
    assert_eq!(
        results[0].name, "agent: semantic_navigation",
        "best match should be the discovery we just stored"
    );
}