Update custom_utils.py

custom_utils.py

@@ -193,9 +193,7 @@ def vector_search(user_query, db, collection, additional_stages=[], vector_index
 
     return list(results)
 
-def handle_user_prompt(openai_api_key, prompt, db, collection, stages=[], vector_index="vector_index"):
-    openai.api_key = openai_api_key
-    
+def rag_retrieval(db, collection, stages=[], vector_index="vector_index"):
     # Assuming vector_search returns a list of dictionaries with keys 'title' and 'plot'
     get_knowledge = vector_search(prompt, db, collection, stages, vector_index)
 
@@ -212,6 +210,11 @@ def handle_user_prompt(openai_api_key, prompt, db, collection, stages=[], vector
     # Convert search results into a DataFrame for better rendering in Jupyter
     search_results_df = pd.DataFrame([item.dict() for item in search_results_models])
 
+    return search_results_df
+
+def rag_inference(openai_api_key, prompt, search_results_df):
+    openai.api_key = openai_api_key
+    
     # Generate system response using OpenAI's completion
     content = f"Answer this user question: {prompt} with the following context:\n{search_results_df}"