Setting the file. One moment.
Quickstart Complete · MongoDB Search And AI · mongodb/agent-skills · Skills Docs
Back to the top of the page scripts/ quickstart_complete.py
Python · 386 lines · 14 KB
16 #
17 # To run against your own data, change the two values below AND
18 # update the index definitions, `$project` stages, and query strings
19 # to match your own field names.
20
21 DB_NAME = "sample_mflix"
22 COLLECTION_NAME = "movies"
23
24 # Index names. The walkthrough uses these defaults, but creates `<name>_2` when
25 # the default is already taken by an incompatible definition. If that happened,
26 # set the matching constant to the name it created, so this script reuses that
27 # index instead of stopping on the conflicting one.
28 SEMANTIC_INDEX_NAME = "quickstart_semantic"
29 TEXT_INDEX_NAME = "quickstart_text"
30 AUTOCOMPLETE_INDEX_NAME = "quickstart_autocomplete"
31
32 INDEX_WAIT_TIMEOUT_SECONDS = 600
33
34 # ─────────────────────────────────────────────────────────────
35
36 import os
37 import time
38 from pymongo import MongoClient
39 from pymongo.operations import SearchIndexModel
40
41 CONNECTION_STRING = (os.environ.get( "MONGODB_URI" )
42 or os.environ.get( "MDB_MCP_CONNECTION_STRING" ))
43
44 if not CONNECTION_STRING :
45 raise SystemExit ( " \n ERROR: no connection string. Set MONGODB_URI (or"
46 " MDB_MCP_CONNECTION_STRING) to your cluster's SRV"
47 " string, then rerun." )
48
49 client = MongoClient( CONNECTION_STRING )
50 db = client[ DB_NAME ]
51 collection = db[ COLLECTION_NAME ]
52
53 def fail (message):
54 print ( f " \n ERROR: { message } " )
55 client.close()
56 raise SystemExit ( 1 )
57
58 def wait_for_index (name):
59 print ( f " Waiting for index ' { name } ' to be ready..." , end = "" , flush = True )
60 deadline = time.monotonic() + INDEX_WAIT_TIMEOUT_SECONDS
61 while True :
62 indexes = list (collection.list_search_indexes(name))
63 if not indexes:
64 fail( f "Index ' { name } ' does not exist." )
65 status = indexes[ 0 ].get( "status" )
66 if status == "READY" :
67 print ( " ready." )
68 return
69 if status == "STALE" :
70 fail( f "Index ' { name } ' is stale because embedding sync is paused. "
71 "Free disk space, wait for READY, then rerun." )
72 if status == "FAILED" :
73 fail( f "Index ' { name } ' failed to build. Check the tier index cap"
74 " (3 on Free, 10 on Flex) and the index definition." )
75 if time.monotonic() > deadline:
76 fail( f "Index ' { name } ' is still ' { status } ' after"
77 f " { INDEX_WAIT_TIMEOUT_SECONDS } s. Check its status in Atlas,"
78 " then rerun." )
79 print ( "." , end = "" , flush = True )
80 time.sleep( 5 )
81
82 def covers (expected, actual):
83 # True when `actual` contains everything `expected` asks for. Atlas adds
84 # its own defaults to a stored definition, so an exact match is too strict.
85 if isinstance (expected, dict ):
86 # A field can be stored as a list of types (a multi-type field), which
87 # still satisfies a single expected type.
88 if isinstance (actual, list ):
89 return any (covers(expected, candidate) for candidate in actual)
90 return ( isinstance (actual, dict )
91 and all (k in actual and covers(v, actual[k])
92 for k, v in expected.items()))
93 if isinstance (expected, list ):
94 return ( isinstance (actual, list )
95 and all ( any (covers(item, candidate) for candidate in actual)
96 for item in expected))
97 return expected == actual
98
99 def check_definition (name, expected):
100 indexes = list (collection.list_search_indexes(name))
101 if not indexes:
102 fail( f "Index ' { name } ' does not exist." )
103 actual = indexes[ 0 ].get( "latestDefinition" , {})
104 if not covers(expected, actual):
105 fail( f "An index named ' { name } ' already exists, but its definition does"
106 " not match the one this script needs, so the queries below would"
107 " fail. An index created by another route (for example the MCP"
108 " walkthrough) can omit fields this script queries. Drop"
109 f " ' { name } ' in Atlas and rerun, or point the matching"
110 " *_INDEX_NAME constant at the top of this script at a name that"
111 f " is free or already correct. \n Existing definition: { actual } " )
112
113 def ensure_index (model, name, definition):
114 try :
115 collection.create_search_index(model)
116 except Exception as e:
117 if "already exists" not in str (e):
118 fail( f "Could not create index ' { name } ': { e } " )
119 print ( " Index already exists — checking its definition." )
120 check_definition(name, definition)
121 wait_for_index(name)
122
123 def print_results (results, fields = [ "title" , "plot" ]):
124 for i, doc in enumerate (results, 1 ):
125 print ( f " \n { i } . { doc.get( 'title' , 'N/A' ) } " )
126 if "score" in doc:
127 print ( f " Score: { round (doc[ 'score' ], 3 ) } " )
128 if "plot" in fields and doc.get( "plot" ):
129 print ( f " Plot: { doc[ 'plot' ][: 120 ] } ..." )
130 if "genres" in doc:
131 print ( f " Genres: { doc.get( 'genres' , []) } " )
132
133 # ─────────────────────────────────────────────────────────────
134 # 1. CREATE AUTO EMBEDDING INDEX (Semantic Search)
135 # ─────────────────────────────────────────────────────────────
136 print ( " \n ── SEMANTIC SEARCH (Auto Embedding) ──────────────────" )
137 print ( "Creating autoEmbed index on 'plot' field..." )
138
139 semantic_definition = {
140 "fields" : [
141 {
142 "type" : "autoEmbed" ,
143 "modality" : "text" ,
144 "path" : "plot" ,
145 "model" : "voyage-4"
146 },
147 {
148 "type" : "filter" ,
149 "path" : "genres"
150 }
151 ]
152 }
153
154 ensure_index(
155 SearchIndexModel(
156 definition = semantic_definition,
157 name = SEMANTIC_INDEX_NAME ,
158 type = "vectorSearch"
159 ),
160 SEMANTIC_INDEX_NAME ,
161 semantic_definition
162 )
163
164 # Run semantic query
165 print ( " \n Query: 'a story about growing up'" )
166 results = list (collection.aggregate([
167 {
168 "$vectorSearch" : {
169 "index" : SEMANTIC_INDEX_NAME ,
170 "path" : "plot" ,
171 "query" : "a story about growing up" ,
172 "numCandidates" : 100 ,
173 "limit" : 3
174 }
175 },
176 {
177 "$project" : {
178 "_id" : 0 ,
179 "title" : 1 ,
180 "plot" : 1 ,
181 "score" : { "$meta" : "vectorSearchScore" }
182 }
183 }
184 ]))
185 print_results(results)
186
187 # Run filtered semantic query
188 print ( " \n Filtered query: 'a story about growing up' + Drama only" )
189 results = list (collection.aggregate([
190 {
191 "$vectorSearch" : {
192 "index" : SEMANTIC_INDEX_NAME ,
193 "path" : "plot" ,
194 "query" : "a story about growing up" ,
195 "filter" : { "genres" : { "$eq" : "Drama" } },
196 "numCandidates" : 150 ,
197 "limit" : 3
198 }
199 },
200 {
201 "$project" : {
202 "_id" : 0 ,
203 "title" : 1 ,
204 "genres" : 1 ,
205 "score" : { "$meta" : "vectorSearchScore" }
206 }
207 }
208 ]))
209 print_results(results, fields = [ "title" , "genres" ])
210
211 # ─────────────────────────────────────────────────────────────
212 # 2. CREATE TEXT SEARCH INDEX (Keyword Search)
213 # ─────────────────────────────────────────────────────────────
214 print ( " \n ── KEYWORD SEARCH (Atlas Search) ─────────────────────" )
215 print ( "Creating text search index on 'title', 'plot', 'genres'..." )
216
217 text_definition = {
218 "mappings" : {
219 "dynamic" : False ,
220 "fields" : {
221 "title" : { "type" : "string" },
222 "plot" : { "type" : "string" },
223 "genres" : { "type" : "token" }
224 }
225 }
226 }
227
228 ensure_index(
229 SearchIndexModel(
230 definition = text_definition,
231 name = TEXT_INDEX_NAME
232 ),
233 TEXT_INDEX_NAME ,
234 text_definition
235 )
236
237 # Run keyword query
238 print ( " \n Keyword query: 'space adventure'" )
239 results = list (collection.aggregate([
240 {
241 "$search" : {
242 "index" : TEXT_INDEX_NAME ,
243 "text" : {
244 "query" : "space adventure" ,
245 "path" : [ "title" , "plot" ]
246 }
247 }
248 },
249 {
250 "$project" : {
251 "_id" : 0 ,
252 "title" : 1 ,
253 "plot" : 1 ,
254 "score" : { "$meta" : "searchScore" }
255 }
256 },
257 { "$limit" : 3 }
258 ]))
259 print_results(results)
260
261 # Run fuzzy query
262 print ( " \n Fuzzy query: 'sapce adventre' (intentional typos)" )
263 results = list (collection.aggregate([
264 {
265 "$search" : {
266 "index" : TEXT_INDEX_NAME ,
267 "text" : {
268 "query" : "sapce adventre" ,
269 "path" : [ "title" , "plot" ],
270 "fuzzy" : { "maxEdits" : 1 }
271 }
272 }
273 },
274 {
275 "$project" : {
276 "_id" : 0 ,
277 "title" : 1 ,
278 "score" : { "$meta" : "searchScore" }
279 }
280 },
281 { "$limit" : 3 }
282 ]))
283 print_results(results, fields = [ "title" ])
284
285 # Autocomplete needs 'title' typed as autocomplete, which the text index does
286 # not provide. Use a dedicated index, matching the walkthrough's MCP route.
287 print ( " \n Creating autocomplete index on 'title'..." )
288
289 autocomplete_definition = {
290 "mappings" : {
291 "dynamic" : False ,
292 "fields" : {
293 "title" : { "type" : "autocomplete" , "tokenization" : "edgeGram" }
294 }
295 }
296 }
297
298 ensure_index(
299 SearchIndexModel(
300 definition = autocomplete_definition,
301 name = AUTOCOMPLETE_INDEX_NAME
302 ),
303 AUTOCOMPLETE_INDEX_NAME ,
304 autocomplete_definition
305 )
306
307 # Run autocomplete
308 print ( " \n Autocomplete query: 'inc'" )
309 results = list (collection.aggregate([
310 {
311 "$search" : {
312 "index" : AUTOCOMPLETE_INDEX_NAME ,
313 "autocomplete" : { "query" : "inc" , "path" : "title" }
314 }
315 },
316 { "$project" : { "_id" : 0 , "title" : 1 } },
317 { "$limit" : 5 }
318 ]))
319 for doc in results:
320 print ( f " - { doc[ 'title' ] } " )
321
322 # ─────────────────────────────────────────────────────────────
323 # 3. HYBRID SEARCH ($rankFusion)
324 # ─────────────────────────────────────────────────────────────
325 print ( " \n ── HYBRID SEARCH ($rankFusion) ───────────────────────" )
326 print ( "Query: 'a story about growing up' (semantic 70% + keyword 30%)" )
327
328 try :
329 results = list (collection.aggregate([
330 {
331 "$rankFusion" : {
332 "input" : {
333 "pipelines" : {
334 "semanticPipeline" : [
335 {
336 "$vectorSearch" : {
337 "index" : SEMANTIC_INDEX_NAME ,
338 "path" : "plot" ,
339 "query" : "a story about growing up" ,
340 "numCandidates" : 100 ,
341 "limit" : 20
342 }
343 }
344 ],
345 "keywordPipeline" : [
346 {
347 "$search" : {
348 "index" : TEXT_INDEX_NAME ,
349 "text" : {
350 "query" : "coming of age" ,
351 "path" : "plot"
352 }
353 }
354 },
355 { "$limit" : 20 }
356 ]
357 }
358 },
359 "combination" : {
360 "weights" : {
361 "semanticPipeline" : 0.7 ,
362 "keywordPipeline" : 0.3
363 }
364 }
365 }
366 },
367 { "$project" : { "_id" : 0 , "title" : 1 , "plot" : 1 } },
368 { "$limit" : 5 }
369 ]))
370 print_results(results)
371 except Exception as e:
372 message = str (e)
373 if "Unrecognized pipeline stage name: '$rankFusion'" in message:
374 print ( " $rankFusion requires MongoDB 8.0+. This cluster does not"
375 " support it, so hybrid search was skipped." )
376 else :
377 fail( f "Hybrid search query failed: { message } " )
378
379 print ( " \n ── Done! ─────────────────────────────────────────────" )
380 print ( "This script targets the sample_mflix.movies schema. To use your" )
381 print ( "own data, set MONGODB_URI (or MDB_MCP_CONNECTION_STRING), then" )
382 print ( "change DB_NAME and COLLECTION_NAME at the top and adjust the" )
383 print ( "index definitions, $project stages, and query strings to match" )
384 print ( "your own field names." )
385
386 client.close()