Update multi_agent.py

multi_agent.py

@@ -26,7 +26,7 @@ def get_legal_moves() -> Annotated[str, "A list of legal moves in UCI format"]:
        The input should always be an empty string, 
        and this function will always return legal moves in UCI format."""
     try:
-        print("### get_legal_moves")
+        print("## get_legal_moves")
         global legal_moves
         legal_moves = ",".join([str(move) for move in board.legal_moves])
         return legal_moves
@@ -40,13 +40,13 @@ def make_move(move: Annotated[str, "A move in UCI format."]) -> Annotated[str, "
        The input should always be a move in UCI format, 
        and this function will always return the result of the move in UCI format."""
     try:
-        print("### make_move")
+        print("## make_move")
         move = chess.Move.from_uci(move)
         board.push_uci(str(move))
 
         global move_num
         move_num += 1
-        print("### move_num=" + str(move_num))
+        print("## move_num=" + str(move_num))
 
         board_svgs.append(chess.svg.board(
             board,
@@ -94,10 +94,10 @@ def create_agent(llm: ChatOpenAI, tools: list, system_prompt: str):
 
 def agent_node(state, agent, name):
     try:
+        print("## agent_node=" + name)
         result = agent.invoke(state)
-        print("### node=" + name)
-        print("### result=" + str(result))
-        print("### result['output']=" + result["output"])
+        print("## result=" + str(result))
+        print("## result['output']=" + result["output"])
         return {
             "messages": [HumanMessage(content=result["output"], name=name)]
         }
@@ -106,18 +106,16 @@ def agent_node(state, agent, name):
         return {"messages": [HumanMessage(content=f"Error: {e}", name=name)]}
    
 def create_graph():
-    llm_chess_board_proxy = ChatOpenAI(model="gpt-4o")
-    llm_player_white = ChatOpenAI(model="gpt-4o")
-    llm_player_black = ChatOpenAI(model="gpt-4o")
-
+    players = ["player_white", "player_black"]
+    
     system_prompt = (
         "You are a Chess Board Proxy tasked with managing a game of chess "
         "between player_white and player_black. player_white makes the first move, "
         "then the players take turns."
     )
 
-    options = ["player_white", "player_black"]
-    
+    options = players
+
     function_def = {
         "name": "route",
         "description": "Select the next player.",
@@ -147,7 +145,11 @@ def create_graph():
                 "Select one of: {options}.",
             ),
         ]
-    ).partial(options=str(options), members=", ".join(options), verbose=True)
+    ).partial(options=str(options), members=", ".join(players), verbose=True)
+    
+    llm_chess_board_proxy = ChatOpenAI(model="gpt-4o")
+    llm_player_white = ChatOpenAI(model="gpt-4o")
+    llm_player_black = ChatOpenAI(model="gpt-4o")
     
     supervisor_chain = (
         prompt
@@ -190,7 +192,7 @@ def create_graph():
         {"chess_board_proxy": "chess_board_proxy", END: END}
     )
     
-    conditional_map = {k: k for k in options}
+    conditional_map = {k: k for k in players}
     conditional_map["END"] = END
     
     graph.add_conditional_edges("chess_board_proxy", lambda x: x["next"], conditional_map)
@@ -222,10 +224,11 @@ def initialize():
     legal_moves = ""
 
 def run_multi_agent(moves_num):
+    initialize()
+
     global num_moves
-    num_moves = moves_num
 
-    initialize()
+    num_moves = moves_num
     
     graph = create_graph()