Update app.py

app.py

@@ -308,16 +308,18 @@ class DocumentRAG:
     def run_multiagent_storygraph(self, topic: str, context: str):
         self.llm = ChatOpenAI(model_name="gpt-4", temperature=0.7, api_key=self.api_key)
 
-        # Story subgraph
+        # Define the story subgraph with reranking
         story_graph = StateGraph(StoryState)
         story_graph.add_node("Retrieve", self.retrieve_docs)
+        story_graph.add_node("Rerank", self.rerank_docs)  # Add rerank step
         story_graph.add_node("Generate", self.generate_story)
         story_graph.set_entry_point("Retrieve")
-        story_graph.add_edge("Retrieve", "Generate")
+        story_graph.add_edge("Retrieve", "Rerank")
+        story_graph.add_edge("Rerank", "Generate")
         story_graph.set_finish_point("Generate")
         story_subgraph = story_graph.compile()
 
-        # Main graph
+        # Main graph setup
         graph = StateGraph(MultiAgentState)
         graph.add_node("beginner_topic", self.beginner_topic)
         graph.add_node("middle_topic", self.middle_topic)
@@ -334,17 +336,24 @@ class DocumentRAG:
             ["story_generator"])
         graph.add_edge("story_generator", END)
 
-
         compiled = graph.compile(checkpointer=MemorySaver())
         thread = {"configurable": {"thread_id": "storygraph-session"}}
 
+        # Initial run to extract subtopics
         result = compiled.invoke({"topic": [topic], "context": [context]}, thread)
 
+        # Fallback if no subtopics were extracted
+        if not result.get("sub_topic_list"):
+            fallback_subs = ["Neural Networks", "Reinforcement Learning", "Supervised vs Unsupervised"]
+            compiled.update_state(thread, {"sub_topic_list": fallback_subs})
+            result = compiled.invoke(None, thread, stream_mode="values")
+
         return result
 
 
 
 
+
 # Initialize RAG system in session state
 if "rag_system" not in st.session_state:
     st.session_state.rag_system = DocumentRAG()