Update interim.py

interim.py

@@ -54,13 +54,15 @@ class RAGWorkflowBasic(Workflow):
 class RAGWorkflowRerank(RAGWorkflowBasic):
     @step
     def rerank(self, ctx: Context, ev: RetrieverEvent) -> RerankEvent:
-        # Use query directly as a mock QueryBundle
-        query_bundle = {"query_str": ctx.get("query")}
-        # Initialize the reranker
+        # Construct QueryBundle properly using the query string
+        query_str = ctx.get("query")  # Retrieve the query from the context
+        query_bundle = QueryBundle(query_str=query_str)
+        # Initialize the LLM reranker
         llm_rerank = LLMRerank(top_n=2, llm=Groq(model="llama3-70b-8192"))
-        # Rerank nodes
+        # Perform the reranking using QueryBundle
         reranked_nodes = llm_rerank.postprocess_nodes(ev.nodes, query_bundle)
         return RerankEvent(nodes=reranked_nodes)
+
         
 
 # Function to Visualize Workflows
@@ -96,24 +98,30 @@ elif data_source == "Provide PDF URL":
 if st.button("Run Workflow"):
     if os.listdir(temp_dir):
         st.write("### Step 1: Ingesting Documents...")
+        # Step 1: Ingest Documents
         embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
         documents = SimpleDirectoryReader(temp_dir).load_data()
         index = VectorStoreIndex.from_documents(documents, embed_model=embed_model)
 
         st.write("### Step 2: Retrieving Documents...")
+        # Step 2: Retrieve Documents
         retriever = index.as_retriever(similarity_top_k=2)
         nodes = retriever.retrieve(query)
 
         if workflow_choice == "Workflow with Reranking":
             st.write("### Step 3: Reranking Results...")
+            # Step 3: Wrap query into QueryBundle and rerank
+            query_bundle = QueryBundle(query_str=query)  # Wrap query into QueryBundle
             reranker = LLMRerank(top_n=2, llm=Groq(model="llama3-70b-8192"))
-            nodes = reranker.postprocess_nodes(nodes, query)
+            nodes = reranker.postprocess_nodes(nodes, query_bundle)
 
         st.write("### Step 4: Synthesizing Response...")
+        # Step 4: Synthesize Response
         summarizer = CompactAndRefine(llm=Groq(model="llama3-70b-8192"))
         response = summarizer.synthesize(query, nodes=nodes)
         st.markdown(f"### **Response:** {response}")
 
+        # Workflow Visualization
         st.write("### Workflow Visualization")
         workflow_class = RAGWorkflowRerank if workflow_choice == "Workflow with Reranking" else RAGWorkflowBasic
         visualize_workflow(workflow_class, "workflow.html")