Update app.py

app.py

@@ -406,7 +406,8 @@ class DocumentRAG:
     
 
 
-    def run_multiagent_storygraph(self, topic: str, context: str):
+    def run_multiagent_storygraph(self, topic: str, context: str, language: str = "English"):
+
         if self.embedding_choice == "OpenAI":
             self.llm = ChatOpenAI(model_name="gpt-4", temperature=0.7, api_key=self.api_key)
         elif self.embedding_choice == "Cohere":
@@ -421,7 +422,7 @@ class DocumentRAG:
         story_graph = StateGraph(StoryState)
         story_graph.add_node("Retrieve", self.retrieve_node)
         story_graph.add_node("Rerank", self.rerank_node)
-        story_graph.add_node("Generate", self.generate_story_node)
+        story_graph.add_node("Generate", lambda state: self.generate_story_node(state, language=state.get("language", "English")))
         story_graph.set_entry_point("Retrieve")
         story_graph.add_edge("Retrieve", "Rerank")
         story_graph.add_edge("Rerank", "Generate")
@@ -442,7 +443,7 @@ class DocumentRAG:
         graph.add_edge("advanced_topic", "topic_extractor")
         graph.add_conditional_edges(
             "topic_extractor",
-            lambda state: [Send("story_generator", {"story_topic": t}) for t in state.sub_topic_list],
+            lambda state: [Send("story_generator", {"story_topic": t, "language": language}) for t in state.sub_topic_list],
             ["story_generator"]
         )
         graph.add_edge("story_generator", END)