Setting the file. One moment.
Seed Test Data · Amazon Neptune · aws/agent-toolkit-for-aws · Skills Docs
Back to the top of the page 67.2
Agentic Memory · references
92
Routing Traffic With Route53 And CloudFront
Resilience Program Design
Debugging Lambda Timeouts
scripts/ seed_test_data.py
Python · 195 lines · 6 KB
17 import sys
18 import time
19
20 from gremlin_python.driver import client, serializer
21
22 NEPTUNE_ENDPOINT = os.environ.get( "NEPTUNE_ENDPOINT" , "localhost" )
23 NEPTUNE_PORT = 8182
24
25
26 def get_client ():
27 return client.Client(
28 f "wss:// { NEPTUNE_ENDPOINT } : { NEPTUNE_PORT } /gremlin" ,
29 "g" ,
30 message_serializer = serializer.GraphSONSerializersV2d0(),
31 )
32
33
34 def clear_graph (c, confirmed: bool = False ):
35 """Wipe ALL vertices and edges. Destructive and irreversible.
36
37 `g.V().drop()` is on the skill's Never-Auto-Execute list (see
38 references/action-safety.md). It runs ONLY when the caller passes an
39 explicit confirmation — either --confirm-wipe on the command line or
40 NEPTUNE_SEED_CONFIRM_WIPE=1 in the environment.
41 """
42 if not confirmed:
43 raise RuntimeError (
44 "Refusing to wipe the graph: g.V().drop() requires explicit "
45 "confirmation. Re-run with --confirm-wipe or set "
46 "NEPTUNE_SEED_CONFIRM_WIPE=1. Never run this against a graph "
47 "holding data you need."
48 )
49 print ( f "Clearing ALL graph data at { NEPTUNE_ENDPOINT } ..." )
50 c.submit( "g.V().drop()" ).all().result()
51 time.sleep( 1 )
52
53
54 def seed_social_graph (c):
55 """Social network for QUERY and DEC evals."""
56 print ( "Seeding social graph..." )
57
58 # Create users
59 users = [
60 ( "U1" , "Alice" , 30 ),
61 ( "U2" , "Bob" , 28 ),
62 ( "U3" , "Carol" , 35 ),
63 ( "U4" , "Dave" , 25 ),
64 ( "U5" , "Eve" , 32 ),
65 ]
66
67 for uid, name, age in users:
68 c.submit(
69 "g.addV('Person').property('id', uid).property('name', name).property('age', age)" ,
70 { "uid" : uid, "name" : name, "age" : age},
71 ).all().result()
72
73 # Create follows edges
74 follows = [( "U1" , "U2" ), ( "U1" , "U3" ), ( "U2" , "U4" ), ( "U3" , "U4" ), ( "U4" , "U5" )]
75 for src, dst in follows:
76 c.submit(
77 "g.V().has('Person','id',src).addE('FOLLOWS').to(g.V().has('Person','id',dst))" ,
78 { "src" : src, "dst" : dst},
79 ).all().result()
80
81 print ( f " Created { len (users) } users, { len (follows) } follows edges" )
82
83
84 def seed_fraud_graph (c):
85 """Fraud detection graph for DEC-02 eval."""
86 print ( "Seeding fraud graph..." )
87
88 # Accounts
89 accounts = [( "A1" , "clean" ), ( "A2" , "flagged" ), ( "A3" , "unknown" ), ( "A4" , "clean" )]
90 for aid, status in accounts:
91 c.submit(
92 "g.addV('Account').property('id', aid).property('status', status)" ,
93 { "aid" : aid, "status" : status},
94 ).all().result()
95
96 # Shared identifiers
97 identifiers = [
98 ( "E1" , "email" , "shared@example.com" ),
99 ( "P1" , "phone" , "+15551234567" ),
100 ( "D1" , "device" , "device-abc-123" ),
101 ]
102 for iid, itype, value in identifiers:
103 c.submit(
104 "g.addV('Identifier').property('id', iid).property('type', itype).property('value', value)" ,
105 { "iid" : iid, "itype" : itype, "value" : value},
106 ).all().result()
107
108 # Account → Identifier edges (shared identifiers create fraud ring)
109 uses = [( "A1" , "E1" ), ( "A2" , "E1" ), ( "A2" , "P1" ), ( "A3" , "P1" ), ( "A4" , "D1" )]
110 for aid, iid in uses:
111 c.submit(
112 "g.V().has('Account','id',aid).addE('USES').to(g.V().has('Identifier','id',iid))" ,
113 { "aid" : aid, "iid" : iid},
114 ).all().result()
115
116 print (
117 f " Created { len (accounts) } accounts, { len (identifiers) } identifiers, { len (uses) } USES edges"
118 )
119
120
121 def seed_product_graph (c):
122 """Product recommendation graph."""
123 print ( "Seeding product graph..." )
124
125 products = [( "P1" , "Widget" , 29.99 ), ( "P2" , "Gadget" , 49.99 ), ( "P3" , "Doohickey" , 9.99 )]
126 for pid, name, price in products:
127 c.submit(
128 "g.addV('Product').property('id', pid).property('name', name).property('price', price)" ,
129 { "pid" : pid, "name" : name, "price" : price},
130 ).all().result()
131
132 # User purchases (links social and product graphs)
133 purchases = [( "U1" , "P1" ), ( "U2" , "P1" ), ( "U2" , "P2" ), ( "U3" , "P3" )]
134 for uid, pid in purchases:
135 c.submit(
136 "g.V().has('Person','id',uid).addE('PURCHASED').to(g.V().has('Product','id',pid))" ,
137 { "uid" : uid, "pid" : pid},
138 ).all().result()
139
140 print ( f " Created { len (products) } products, { len (purchases) } PURCHASED edges" )
141
142
143 def verify_seed (c):
144 """Quick verification of seeded data."""
145 print ( " \n Verifying seeded data..." )
146 v_count = c.submit( "g.V().count()" ).all().result()[ 0 ]
147 e_count = c.submit( "g.E().count()" ).all().result()[ 0 ]
148 print ( f " Vertices: { v_count } " )
149 print ( f " Edges: { e_count } " )
150
151 # Test a fraud ring query
152 ring = (
153 c.submit(
154 "g.V().has('Account','id','A2').out('USES').in('USES').dedup().values('id').toList()"
155 )
156 .all()
157 .result()
158 )
159 print ( f " Accounts sharing identifiers with A2: { ring } " )
160 assert len (ring) >= 2 , "Fraud ring query returned unexpected results"
161
162 # Test a 2-hop social traversal
163 fof = (
164 c.submit(
165 "g.V().has('Person','id','U1').out('FOLLOWS').out('FOLLOWS').dedup().values('name').toList()"
166 )
167 .all()
168 .result()
169 )
170 print ( f " Friends-of-friends for Alice: { fof } " )
171 assert len (fof) >= 1 , "Friends-of-friends query returned unexpected results"
172
173 print ( " \n ✅ Seed data verified successfully" )
174
175
176 if __name__ == "__main__" :
177 # g.V().drop() is destructive and requires explicit confirmation.
178 wipe_confirmed = (
179 "--confirm-wipe" in sys.argv or os.environ.get( "NEPTUNE_SEED_CONFIRM_WIPE" ) == "1"
180 )
181
182 print ( f "Connecting to Neptune at { NEPTUNE_ENDPOINT } : { NEPTUNE_PORT } ..." )
183 c = get_client()
184
185 try :
186 clear_graph(c, confirmed = wipe_confirmed)
187 seed_social_graph(c)
188 seed_fraud_graph(c)
189 seed_product_graph(c)
190 verify_seed(c)
191 finally :
192 c.close()
193
194 print ( " \n Seed complete. Neptune test cluster is ready for eval runs." )
195 print ( f "Endpoint: wss:// { NEPTUNE_ENDPOINT } : { NEPTUNE_PORT } /gremlin" )