Setting the file. One moment.
Agentic Memory · 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
214
def hybrid_recall
— line 214
This file
Number 67.15
Position 15 of 19
Type Python
Size 8 KB
Lines 256 scripts/ agentic_memory.py
Python · 256 lines · 8 KB
from
typing
import
Any, Dict, List, Optional
17
18 import boto3
19
20 # =============================================================================
21 # Short-Term Memory (DynamoDB)
22 # =============================================================================
23
24 dynamodb = boto3.resource( "dynamodb" )
25 SHORT_TERM_TABLE = os.environ.get( "SHORT_TERM_TABLE" , "agent-short-term-memory" )
26
27
28 def get_short_term_table ():
29 return dynamodb.Table( SHORT_TERM_TABLE )
30
31
32 def store_message (session_id: str , role: str , content: str ):
33 """Store a message in short-term memory with 24h TTL."""
34 get_short_term_table().put_item(
35 Item = {
36 "session_id" : session_id,
37 "timestamp" : int (time.time() * 1000 ),
38 "role" : role,
39 "content" : content,
40 "ttl" : int (time.time()) + 86400 ,
41 }
42 )
43
44
45 def get_recent_messages (session_id: str , limit: int = 20 ) -> List[Dict]:
46 """Retrieve recent messages (most recent first, then reversed)."""
47 response = get_short_term_table().query(
48 KeyConditionExpression = "session_id = :sid" ,
49 ExpressionAttributeValues = { ":sid" : session_id},
50 ScanIndexForward = False ,
51 Limit = limit,
52 )
53 return list ( reversed (response[ "Items" ]))
54
55
56 # =============================================================================
57 # Long-Term Memory — Neptune Database (Gremlin)
58 # =============================================================================
59
60 NEPTUNE_ENDPOINT = os.environ.get( "NEPTUNE_ENDPOINT" , "localhost" )
61 NEPTUNE_PORT = int (os.environ.get( "NEPTUNE_PORT" , "8182" ))
62
63
64 def get_gremlin_client ():
65 """Create a Gremlin client for Neptune Database."""
66 from gremlin_python.driver import client, serializer
67
68 return client.Client(
69 f "wss:// { NEPTUNE_ENDPOINT } : { NEPTUNE_PORT } /gremlin" ,
70 "g" ,
71 message_serializer = serializer.GraphSONSerializersV2d0(),
72 )
73
74
75 def remember_entity_gremlin (gremlin_client, name: str , entity_type: str , ** properties):
76 """Store or update an entity in Neptune Database."""
77 bindings = { "entity_name" : name, "entity_type" : entity_type}
78 query = """
79 g.V().has('Entity', 'name', entity_name).fold()
80 .coalesce(unfold(), addV('Entity').property('name', entity_name))
81 .property('type', entity_type)
82 .property('updated_at', new Date().getTime())
83 """
84 for key, value in properties.items():
85 binding_key = f "prop_ { key } "
86 query += f " .property(' { key } ', { binding_key } ) \n "
87 bindings[binding_key] = value
88 gremlin_client.submit(query, bindings = bindings).all().result()
89
90
91 def remember_relationship_gremlin (
92 gremlin_client, entity_a: str , entity_b: str , rel_type: str , confidence: float = 1.0
93 ):
94 """Store a relationship between entities in Neptune Database."""
95 bindings = {
96 "name_a" : entity_a,
97 "name_b" : entity_b,
98 "rel_type" : rel_type,
99 "confidence" : confidence,
100 }
101 query = """
102 g.V().has('Entity', 'name', name_a).as('a')
103 .V().has('Entity', 'name', name_b).as('b')
104 .coalesce(
105 select('a').outE('RELATED_TO').where(inV().as('b')),
106 select('a').addE('RELATED_TO').to(select('b'))
107 )
108 .property('type', rel_type)
109 .property('confidence', confidence)
110 .property('updated_at', new Date().getTime())
111 """
112 gremlin_client.submit(query, bindings = bindings).all().result()
113
114
115 def recall_entity_gremlin (gremlin_client, entity_name: str ) -> List:
116 """Recall everything about an entity from Neptune Database."""
117 bindings = { "entity_name" : entity_name}
118 query = """
119 g.V().has('Entity', 'name', entity_name)
120 .project('entity', 'relationships', 'conversations')
121 .by(valueMap())
122 .by(bothE('RELATED_TO').project('type', 'target', 'confidence')
123 .by(values('type')).by(otherV().values('name')).by(values('confidence'))
124 .fold())
125 .by(in('ABOUT').hasLabel('Conversation')
126 .order().by('date', desc).limit(5)
127 .valueMap('summary', 'date').fold())
128 """
129 return gremlin_client.submit(query, bindings = bindings).all().result()
130
131
132 # =============================================================================
133 # Long-Term Memory — Neptune Analytics (openCypher + Vector)
134 # =============================================================================
135
136 analytics_client = boto3.client( "neptune-graph" )
137 GRAPH_ID = os.environ.get( "GRAPH_ID" , "g-xxxxxxxxxx" )
138
139
140 def run_query (graph_id: str , query: str , parameters: Optional[Dict[ str , Any]] = None ) -> List[Dict]:
141 """Execute openCypher query against Neptune Analytics."""
142 kwargs: Dict[ str , Any] = {
143 "graphIdentifier" : graph_id,
144 "queryString" : query,
145 "language" : "OPEN_CYPHER" ,
146 }
147 if parameters:
148 kwargs[ "parameters" ] = parameters
149 response = analytics_client.execute_query( ** kwargs)
150 payload = json.loads(response[ "payload" ].read())
151 return payload.get( "results" , [])
152
153
154 def create_memory_graph (
155 graph_name: str = "agent-memory" ,
156 memory_gb: int = 16 ,
157 embedding_dim: int = 1536 ,
158 generation_model: str = "unknown" ,
159 ) -> str :
160 """Create Neptune Analytics graph with vector search for agent memory.
161
162 Secure defaults: deletionProtection=True and the mandatory skill tags.
163 """
164 response = analytics_client.create_graph(
165 graphName = graph_name,
166 provisionedMemory = memory_gb,
167 publicConnectivity = False ,
168 vectorSearchConfiguration = { "dimension" : embedding_dim},
169 deletionProtection = True ,
170 tags = { "created_by" : "neptune-skill" , "generation_model" : generation_model},
171 )
172 return response[ "id" ]
173
174
175 def remember_with_embedding (
176 graph_id: str , entity_name: str , description: str , embedding: List[ float ]
177 ):
178 """Store entity with embedding for semantic recall."""
179 run_query(
180 graph_id,
181 """
182 MERGE (e:Entity {name: $name} )
183 SET e.description = $description, e.updated_at = timestamp()
184 """ ,
185 parameters = { "name" : entity_name, "description" : description},
186 )
187
188 run_query(
189 graph_id,
190 """
191 MATCH (e:Entity {name: $name} )
192 CALL neptune.algo.vectors.upsert(e, $embedding)
193 YIELD node RETURN node.name
194 """ ,
195 parameters = { "name" : entity_name, "embedding" : embedding},
196 )
197
198
199 def semantic_recall (graph_id: str , query_embedding: List[ float ], top_k: int = 10 ) -> List[Dict]:
200 """Recall memories by vector similarity."""
201 return run_query(
202 graph_id,
203 """
204 CALL neptune.algo.vectors.topKByEmbedding($embedding, {topK: $top_k} )
205 YIELD node, score
206 RETURN node.name AS name, node.type AS type,
207 node.description AS description, score
208 ORDER BY score DESC
209 """ ,
210 parameters = { "embedding" : query_embedding, "top_k" : top_k},
211 )
212
213
214 def hybrid_recall (graph_id: str , entity_name: str , query_embedding: List[ float ]) -> Dict:
215 """Combine graph traversal + vector search for comprehensive recall."""
216 graph_results = run_query(
217 graph_id,
218 """
219 MATCH (e:Entity {name: $name} )-[r:RELATED_TO]-(related)
220 RETURN related.name AS name, related.type AS type,
221 r.type AS relationship, r.confidence AS confidence
222 ORDER BY r.confidence DESC LIMIT 20
223 """ ,
224 parameters = { "name" : entity_name},
225 )
226
227 vector_results = run_query(
228 graph_id,
229 """
230 CALL neptune.algo.vectors.topKByEmbedding($embedding, {topK: 10} )
231 YIELD node, score
232 WHERE node.name <> $name
233 RETURN node.name AS name, node.type AS type, score
234 """ ,
235 parameters = { "embedding" : query_embedding, "name" : entity_name},
236 )
237
238 return { "graph_recall" : graph_results, "semantic_recall" : vector_results}
239
240
241 def consolidate_memories (graph_id: str ):
242 """Decay old memories and remove low-confidence facts."""
243 run_query(
244 graph_id,
245 """
246 MATCH (f:Fact) WHERE f.updated_at < timestamp() - 7776000000
247 SET f.confidence = f.confidence * 0.9
248 """ ,
249 )
250 run_query(
251 graph_id,
252 """
253 MATCH (f:Fact) WHERE f.confidence < 0.1
254 DETACH DELETE f
255 """ ,
256 )