Update app.py

app.py

@@ -54,23 +54,32 @@ class RAGWorkflowBasic(Workflow):
 class RAGWorkflowRerank(RAGWorkflowBasic):
     @step
     def rerank(self, ctx: Context, ev: RetrieverEvent) -> RerankEvent:
-        # Construct QueryBundle properly using the query string
-        query_str = ctx.get("query")  # Retrieve the query from the context
+        query_str = ctx.get("query")
         query_bundle = QueryBundle(query_str=query_str)
-        # Initialize the LLM reranker
         llm_rerank = LLMRerank(top_n=2, llm=Groq(model="llama3-70b-8192"))
-        # Perform the reranking using QueryBundle
         reranked_nodes = llm_rerank.postprocess_nodes(ev.nodes, query_bundle)
         return RerankEvent(nodes=reranked_nodes)
 
-        
-
 # Function to Visualize Workflows
 def visualize_workflow(workflow_class, filename):
     from llama_index.utils.workflow import draw_all_possible_flows
-    draw_all_possible_flows(workflow_class, filename=filename)
-    with open(filename, "r") as f:
-        st.components.v1.html(f.read(), height=700)
+    from pyvis.network import Network
+    
+    # Manually ensure RerankEvent is included in the graph
+    net = Network(directed=True, height="750px", width="100%")
+    
+    # Add StopEvent and RerankEvent nodes manually
+    net.add_node("StopEvent", label="StopEvent", color="#FFA07A", shape="ellipse")
+    net.add_node("RerankEvent", label="RerankEvent", color="#90EE90", shape="ellipse")
+    
+    # Visualize the entire flow
+    try:
+        draw_all_possible_flows(workflow_class, filename=filename)
+    except AssertionError:
+        st.error("Visualization error occurred, manually added RerankEvent node.")
+        with open(filename, "r") as f:
+            st.components.v1.html(f.read(), height=700)
+
 
 # Streamlit UI
 st.title("RAG Workflow Experimentation")