IronMind/test_graph_manager.py at main · stealthwhizz/IronMind · GitHub

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
"""
Quick test of the GraphManager with Neo4j
Tests creating sessions and prompt nodes
"""
import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from src.core.graph_manager import GraphManager
from config.settings import settings

print("=" * 60)
print("🧪 Testing IronMind GraphManager")
print("=" * 60)

try:
    # Initialize GraphManager
    print("\n[1/5] Initializing GraphManager...")
    gm = GraphManager(
        uri=settings.neo4j_uri,
        user=settings.neo4j_user,
        password=settings.neo4j_password,
        database=settings.neo4j_database
    )
    print("✓ GraphManager initialized")

    # Create a session
    print("\n[2/5] Creating test session...")
    session_id = gm.create_session()
    print(f"✓ Session created: {session_id[:8]}...")

    # Add first prompt node
    print("\n[3/5] Adding first prompt node...")
    node_id_1 = gm.add_prompt_node(
        session_id=session_id,
        turn_id=1,
        prompt_text="What is the weather today?",
        token_count=6,
        injection_score=0.05,
        metadata={"test": True}
    )
    print(f"✓ Node 1 created: {node_id_1[:8]}...")

    # Add second prompt node
    print("\n[4/5] Adding second prompt node...")
    node_id_2 = gm.add_prompt_node(
        session_id=session_id,
        turn_id=2,
        prompt_text="Tell me about security best practices",
        token_count=6,
        injection_score=0.1,
        metadata={"test": True}
    )
    print(f"✓ Node 2 created: {node_id_2[:8]}...")

    # Calculate drift score
    print("\n[5/5] Calculating drift score...")
    drift_score = gm.calculate_drift_score(session_id, 2)
    print(f"✓ Drift score: {drift_score:.4f}")

    # Get session history
    print("\n[BONUS] Retrieving session history...")
    history = gm.get_session_history(session_id)
    print(f"✓ History retrieved: {len(history)} nodes")
    for i, node in enumerate(history, 1):
        print(f"  Turn {i}: {node['prompt_text'][:30]}... (score: {node['injection_score']:.2f})")

    # Get stats
    print("\n[STATS] Database statistics...")
    stats = gm.get_graph_stats()
    print(f"  Sessions: {stats['session_count']}")
    print(f"  Total prompts: {stats['total_prompts']}")
    avg_score = stats['avg_injection_score'] if stats['avg_injection_score'] is not None else 0.0
    print(f"  Avg injection score: {avg_score:.4f}")

    # Cleanup
    print("\n[CLEANUP] Cleaning up test data...")
    gm.delete_node(node_id_1)
    gm.delete_node(node_id_2)
    print("✓ Test nodes deleted")

    gm.close()

    print("\n" + "=" * 60)
    print("✅ ALL TESTS PASSED!")
    print("=" * 60)
    print("\nYour Neo4j setup is fully operational for IronMind.")
    print("The Sentinel-Graph (IDG) is ready to track multi-turn attacks!")

except Exception as e:
    print("\n" + "=" * 60)
    print("❌ TEST FAILED")
    print("=" * 60)
    print(f"\nError: {e}")
    import traceback
    traceback.print_exc()