Setting the file. One moment.
Graphrag Pipeline · 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
Creating API Gateway Stage
def graphrag_retrieve
— line 253
This file
Number 67.17
Position 17 of 19
Type Python
Size 9 KB
Lines 302 scripts/ graphrag_pipeline.py
Python · 302 lines · 9 KB
from
typing
import
Any, Dict, List, Optional
17
18 import boto3
19
20 # Neptune Analytics client
21 analytics_client = boto3.client( "neptune-graph" )
22 GRAPH_ID = os.environ.get( "GRAPH_ID" , "g-xxxxxxxxxx" )
23
24
25 # =============================================================================
26 # Core SDK Helper
27 # =============================================================================
28
29
30 def run_query (graph_id: str , query: str , parameters: Optional[Dict[ str , Any]] = None ) -> List[Dict]:
31 """
32 Execute an openCypher query against a Neptune Analytics graph.
33
34 Uses boto3 neptune-graph client (not WebSocket, not port 8182).
35
36 Args:
37 graph_id: Neptune Analytics graph identifier (e.g., 'g-xxxxxxxxxx')
38 query: openCypher query. Use $param_name for parameterized values.
39 parameters: Query parameters dict (prevents injection).
40
41 Returns:
42 List of result dictionaries.
43 """
44 import re
45
46 if not re.match( r " ^ g- [ a-z0-9 ] {10,} $ " , graph_id):
47 raise ValueError (
48 f "Invalid graph_id format: { graph_id } . Expected 'g-' followed by lowercase alphanumeric."
49 )
50
51 kwargs: Dict[ str , Any] = {
52 "graphIdentifier" : graph_id,
53 "queryString" : query,
54 "language" : "OPEN_CYPHER" ,
55 }
56 if parameters:
57 kwargs[ "parameters" ] = parameters
58
59 response = analytics_client.execute_query( ** kwargs)
60 payload = json.loads(response[ "payload" ].read())
61 return payload.get( "results" , [])
62
63
64 def escape_cypher_string (value: str ) -> str :
65 """
66 Escape a string for openCypher. Prefer parameterized queries ($param) instead.
67 Use only when parameters are not supported (e.g., dynamic label names).
68 """
69 return value.replace( " \\ " , " \\\\ " ).replace( "'" , " \\ '" ).replace( '"' , ' \\ "' )
70
71
72 # =============================================================================
73 # Graph Creation
74 # =============================================================================
75
76
77 def create_graphrag_graph (
78 graph_name: str ,
79 memory_gb: int = 32 ,
80 embedding_dimension: int = 1536 ,
81 generation_model: str = "unknown" ,
82 ) -> str :
83 """Create a Neptune Analytics graph with vector search enabled.
84
85 Secure defaults: deletionProtection=True and the mandatory skill tags. See
86 the CreateGraph API reference for the valid provisionedMemory range.
87 """
88 response = analytics_client.create_graph(
89 graphName = graph_name,
90 provisionedMemory = memory_gb,
91 publicConnectivity = False ,
92 vectorSearchConfiguration = { "dimension" : embedding_dimension},
93 replicaCount = 0 ,
94 deletionProtection = True ,
95 tags = { "created_by" : "neptune-skill" , "generation_model" : generation_model},
96 )
97 return response[ "id" ]
98
99
100 # =============================================================================
101 # Document Processing
102 # =============================================================================
103
104
105 def chunk_document (text: str , chunk_size: int = 512 , overlap: int = 50 ) -> List[ str ]:
106 """Split document into overlapping chunks."""
107 chunks = []
108 start = 0
109 while start < len (text):
110 end = start + chunk_size
111 chunks.append(text[start:end])
112 start = end - overlap
113 return chunks
114
115
116 # =============================================================================
117 # Vector Operations
118 # =============================================================================
119
120
121 def store_embedding (graph_id: str , vertex_id: str , embedding: List[ float ]):
122 """Store embedding on a vertex in Neptune Analytics."""
123 run_query(
124 graph_id,
125 """
126 MATCH (n {id: $vertex_id} )
127 CALL neptune.algo.vectors.upsert(n, $embedding)
128 YIELD node
129 RETURN node.id
130 """ ,
131 parameters = { "vertex_id" : vertex_id, "embedding" : embedding},
132 )
133
134
135 def vector_search (
136 graph_id: str , query_embedding: List[ float ], top_k: int = 5 , vertex_label: str = "Chunk"
137 ) -> List[Dict]:
138 """Neptune Analytics vector similarity search."""
139 return run_query(
140 graph_id,
141 """
142 CALL neptune.algo.vectors.topKByEmbedding($embedding, {topK: $top_k} )
143 YIELD node, score
144 WHERE $label IN labels(node)
145 RETURN node.id AS id, node.text AS text, score
146 ORDER BY score DESC
147 """ ,
148 parameters = { "embedding" : query_embedding, "top_k" : top_k, "label" : vertex_label},
149 )
150
151
152 # =============================================================================
153 # Graph Construction
154 # =============================================================================
155
156
157 def ingest_document (graph_id: str , doc_id: str , title: str , source: str ):
158 """Create a Document vertex."""
159 run_query(
160 graph_id,
161 """
162 CREATE (d:Document {id: $doc_id, title: $title, source: $source} )
163 """ ,
164 parameters = { "doc_id" : doc_id, "title" : title, "source" : source},
165 )
166
167
168 def ingest_chunk (
169 graph_id: str ,
170 chunk_id: str ,
171 doc_id: str ,
172 text: str ,
173 sequence: int ,
174 embedding: List[ float ],
175 prev_chunk_id: Optional[ str ] = None ,
176 ):
177 """Create a Chunk vertex, link to document, optionally link to previous chunk."""
178 run_query(
179 graph_id,
180 """
181 CREATE (c:Chunk {id: $chunk_id, text: $text, sequence: $seq} )
182 """ ,
183 parameters = { "chunk_id" : chunk_id, "text" : text, "seq" : sequence},
184 )
185
186 store_embedding(graph_id, chunk_id, embedding)
187
188 run_query(
189 graph_id,
190 """
191 MATCH (d:Document {id: $doc_id} ), (c:Chunk {id: $chunk_id} )
192 CREATE (d)-[:HAS_CHUNK]->(c)
193 """ ,
194 parameters = { "doc_id" : doc_id, "chunk_id" : chunk_id},
195 )
196
197 if prev_chunk_id:
198 run_query(
199 graph_id,
200 """
201 MATCH (prev:Chunk {id: $prev_id} ), (curr:Chunk {id: $curr_id} )
202 CREATE (prev)-[:NEXT]->(curr)
203 """ ,
204 parameters = { "prev_id" : prev_chunk_id, "curr_id" : chunk_id},
205 )
206
207
208 def ingest_entity (
209 graph_id: str , name: str , entity_type: str , description: str , embedding: List[ float ]
210 ):
211 """Merge an Entity vertex with embedding."""
212 entity_id = hashlib.sha256(name.encode()).hexdigest()[: 32 ]
213 run_query(
214 graph_id,
215 """
216 MERGE (e:Entity {name: $name} )
217 ON CREATE SET e.type = $type, e.description = $desc, e.id = $entity_id
218 """ ,
219 parameters = { "name" : name, "type" : entity_type, "desc" : description, "entity_id" : entity_id},
220 )
221 store_embedding(graph_id, entity_id, embedding)
222
223
224 def link_chunk_entity (graph_id: str , chunk_id: str , entity_name: str ):
225 """Link a chunk to an entity it mentions."""
226 run_query(
227 graph_id,
228 """
229 MATCH (c:Chunk {id: $chunk_id} ), (e:Entity {name: $name} )
230 CREATE (c)-[:MENTIONS]->(e)
231 """ ,
232 parameters = { "chunk_id" : chunk_id, "name" : entity_name},
233 )
234
235
236 def link_entities (graph_id: str , source: str , target: str , rel_type: str ):
237 """Create a relationship between two entities."""
238 run_query(
239 graph_id,
240 """
241 MATCH (s:Entity {name: $source} ), (t:Entity {name: $target} )
242 MERGE (s)-[:RELATED_TO {type: $rel_type} ]->(t)
243 """ ,
244 parameters = { "source" : source, "target" : target, "rel_type" : rel_type},
245 )
246
247
248 # =============================================================================
249 # Retrieval
250 # =============================================================================
251
252
253 def graphrag_retrieve (graph_id: str , query_embedding: List[ float ], top_k: int = 5 ) -> List[Dict]:
254 """
255 Two-phase retrieval:
256 1. Vector search for similar chunks
257 2. Graph traversal to expand context
258 """
259 similar_chunks = vector_search(graph_id, query_embedding, top_k = top_k)
260
261 expanded_context = []
262 for chunk in similar_chunks:
263 chunk_id = chunk[ "id" ]
264
265 entities = run_query(
266 graph_id,
267 """
268 MATCH (c:Chunk {id: $cid} )-[:MENTIONS]->(e:Entity)
269 RETURN e.name AS name, e.type AS type, e.description AS description
270 """ ,
271 parameters = { "cid" : chunk_id},
272 )
273
274 related = run_query(
275 graph_id,
276 """
277 MATCH (c:Chunk {id: $cid} )-[:MENTIONS]->(e:Entity)-[:RELATED_TO]-(r:Entity)
278 RETURN DISTINCT r.name AS name, r.type AS type, r.description AS description
279 """ ,
280 parameters = { "cid" : chunk_id},
281 )
282
283 neighbors = run_query(
284 graph_id,
285 """
286 MATCH (c:Chunk {id: $cid} )-[:NEXT]-(n:Chunk)
287 RETURN n.text AS text, n.id AS id
288 """ ,
289 parameters = { "cid" : chunk_id},
290 )
291
292 expanded_context.append(
293 {
294 "chunk_text" : chunk[ "text" ],
295 "score" : chunk[ "score" ],
296 "entities" : entities,
297 "related_entities" : related,
298 "neighboring_chunks" : [n[ "text" ] for n in neighbors],
299 }
300 )
301
302 return expanded_context