bstraehle commited on
Commit
c413f9d
·
verified ·
1 Parent(s): 83045be

Update custom_utils.py

Browse files
Files changed (1) hide show
  1. custom_utils.py +6 -8
custom_utils.py CHANGED
@@ -61,8 +61,6 @@ def rag_retrieval_advanced(openai_api_key,
61
  # "bedrooms": { "$eq": 1}
62
  # }
63
  #}
64
-
65
- #additional_stages = [match_stage]
66
 
67
  # 2) Average review score and review count boost, sorted in descending order
68
 
@@ -188,7 +186,7 @@ def vector_search_naive(openai_api_key,
188
 
189
  pipeline = [vector_search_stage, get_remove_embedding_stage()]
190
 
191
- return invoke_search(collection, pipeline)
192
 
193
  def vector_search_advanced(openai_api_key,
194
  user_query,
@@ -221,28 +219,28 @@ def vector_search_advanced(openai_api_key,
221
 
222
  pipeline = [vector_search_stage, get_remove_embedding_stage()] + additional_stages
223
 
224
- return invoke_search(collection, pipeline)
225
 
226
  def get_remove_embedding_stage():
227
  return {
228
  "$unset": "description_embedding"
229
  }
230
 
231
- def invoke_search(collection, pipeline):
232
  results = collection.aggregate(pipeline)
233
 
234
- print(f"Vector search millis elapsed: {get_millis_elapsed()}")
235
 
236
  return list(results)
237
 
238
- def get_millis_elapsed():
239
  explain_query_execution = db.command(
240
  "explain", {
241
  "aggregate": collection.name,
242
  "pipeline": pipeline,
243
  "cursor": {}
244
  },
245
- verbosity='executionStats')
246
 
247
  explain_vector_search = explain_query_execution["stages"][0]["$vectorSearch"]
248
 
 
61
  # "bedrooms": { "$eq": 1}
62
  # }
63
  #}
 
 
64
 
65
  # 2) Average review score and review count boost, sorted in descending order
66
 
 
186
 
187
  pipeline = [vector_search_stage, get_remove_embedding_stage()]
188
 
189
+ return invoke_search(db, collection, pipeline)
190
 
191
  def vector_search_advanced(openai_api_key,
192
  user_query,
 
219
 
220
  pipeline = [vector_search_stage, get_remove_embedding_stage()] + additional_stages
221
 
222
+ return invoke_search(db, collection, pipeline)
223
 
224
  def get_remove_embedding_stage():
225
  return {
226
  "$unset": "description_embedding"
227
  }
228
 
229
+ def invoke_search(db, collection, pipeline):
230
  results = collection.aggregate(pipeline)
231
 
232
+ print(f"Vector search millis elapsed: {get_millis_elapsed(db)}")
233
 
234
  return list(results)
235
 
236
+ def get_millis_elapsed(db):
237
  explain_query_execution = db.command(
238
  "explain", {
239
  "aggregate": collection.name,
240
  "pipeline": pipeline,
241
  "cursor": {}
242
  },
243
+ verbosity="executionStats")
244
 
245
  explain_vector_search = explain_query_execution["stages"][0]["$vectorSearch"]
246