Update rag_langgraph.py

rag_langgraph.py

@@ -45,8 +45,8 @@ class AgentState(TypedDict):
     messages: Annotated[Sequence[BaseMessage], operator.add]
     # The 'next' field indicates where to route to next
     next: str
-    
-def run_multi_agent(prompt):
+
+def create_graph(topic, word_count):
     tavily_tool = TavilySearchResults(max_results=5)
     python_repl_tool = PythonREPLTool()
 
@@ -126,19 +126,9 @@ def run_multi_agent(prompt):
     # Finally, add entrypoint
     workflow.set_entry_point("supervisor")
     
-    graph = workflow.compile()
+    return workflow.compile()
 
-    result = "DONE"
+def run_multi_agent(topic, word_count):
+    graph = create_graph(topic, word_count)
 
-    for s in graph.stream(
-        {
-            "messages": [
-                HumanMessage(content="Code hello world and print it to the terminal")
-            ]
-        }
-    ):
-        if "__end__" not in s:
-            print(s)
-            print("----")
-        
-    return result
+    return graph.invoke(new HumanMessage("Code hello world and print it to the terminal"));