Update multi_agent.py

multi_agent.py

@@ -28,9 +28,6 @@ class AgentState(TypedDict):
     next: str
 
 def create_agent(llm: ChatOpenAI, tools: list, system_prompt: str):
-    print("## create_agent")
-    global num_moves
-    
     prompt = ChatPromptTemplate.from_messages(
         [
             ("system", system_prompt),
@@ -46,12 +43,11 @@ def create_agent(llm: ChatOpenAI, tools: list, system_prompt: str):
                          handle_parsing_errors=True,
                          return_intermediate_steps=True,
                          verbose=True,
-                         max_iterations=num_moves)
+                         max_iterations=2) # get_legal_moves & make_move
 
 def agent_node(state, agent, name):
     try:
-        print("## agent_node=" + name)
-        print("## state=" + str(state))
+        print("### node=" + name)
         result = agent.invoke(state)
         return {"messages": [HumanMessage(content=result["output"], name=name)]}
     except Exception as e:
@@ -64,7 +60,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
@@ -78,9 +74,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))
         
         board_svgs.append(chess.svg.board(
             board,
@@ -96,10 +96,6 @@ def make_move(move: Annotated[str, "A move in UCI format."]) -> Annotated[str, "
             if piece_symbol.isupper()
             else chess.piece_name(piece.piece_type)
         )
-
-        global move_num
-        move_num += 1
-        print("## move_num=" + str(move_num))
         
         return f"Moved {piece_name} ({piece_symbol}) from "\
                f"{chess.SQUARE_NAMES[move.from_square]} to "\
@@ -109,9 +105,12 @@ def make_move(move: Annotated[str, "A move in UCI format."]) -> Annotated[str, "
         return f"Error: unable to make move {move}"
     
 def create_graph():
-    print("## create_graph")
-    
     players = ["player_white", "player_black"]
+    options = players
+
+    llm_board_proxy = ChatOpenAI(model="gpt-4o")
+    llm_player_white = ChatOpenAI(model="gpt-4o")
+    llm_player_black = ChatOpenAI(model="gpt-4o")
     
     system_prompt = (
         "You are a Chess Board Proxy tasked with managing a game of chess "
@@ -119,9 +118,6 @@ def create_graph():
         "then the players take turns."
     )
 
-    #options = ["FINISH"] + players
-    options = players
-
     function_def = {
         "name": "route",
         "description": "Select the next player.",
@@ -153,34 +149,35 @@ def create_graph():
         ]
     ).partial(options=str(options), members=", ".join(players), verbose=True)
     
-    llm_1 = ChatOpenAI(model="gpt-4o")
-    llm_2 = ChatOpenAI(model="gpt-4o")
-    llm_3 = ChatOpenAI(model="gpt-4o")
-    
     supervisor_chain = (
         prompt
-        | llm_1.bind_functions(functions=[function_def], function_call="route")
+        | llm_board_proxy.bind_functions(functions=[function_def], function_call="route")
         | JsonOutputFunctionsParser()
     )
 
-    player_white_agent = create_agent(llm_2, [get_legal_moves, make_move], system_prompt=
+    player_white_agent = create_agent(llm_player_white, [get_legal_moves, make_move], system_prompt=
                                      "You are a chess Grandmaster and you play as white. "
-                                     "1. First call get_legal_moves(), to get a list of legal moves. "
-                                     "2. Then call make_move(move) to make a move. ONLY make a move in the list returned by step 1.")
-                                     #"3. Finally analyze the move in format: **Analysis:** move in UCI format, emoji of piece emoji, unordered list.")
+                                     "First call get_legal_moves() to get a list of legal moves. "
+                                     "Then study the returned moves and call make_move(move) to make the best move. "
+                                     "Finally analyze the move in format: **Analysis:** move in UCI format, emoji of piece, unordered list of 3 items.")
     player_white_node = functools.partial(agent_node, agent=player_white_agent, name="player_white")
 
-    player_black_agent = create_agent(llm_3, [get_legal_moves, make_move], system_prompt=
+    player_black_agent = create_agent(llm_player_black, [get_legal_moves, make_move], system_prompt=
                                      "You are a chess Grandmaster and you play as black. "
-                                     "1. First call get_legal_moves(), to get a list of legal moves. "
-                                     "2. Then call make_move(move) to make a move. ONLY make a move in the list returned by step 1.")
-                                     #"3. Finally analyze the move in format: **Analysis:** move in UCI format, emoji of piece emoji, unordered list.")
+                                     "First call get_legal_moves() to get a list of legal moves. "
+                                     "Then study the returned moves and call make_move(move) to make the best move. "
+                                     "Finally analyze the move in format: **Analysis:** move in UCI format, emoji of piece, unordered list of 3 items.")
     player_black_node = functools.partial(agent_node, agent=player_black_agent, name="player_black")
     
     graph = StateGraph(AgentState)
+
+    graph.add_node("chess_board_proxy", supervisor_chain)
     graph.add_node("player_white", player_white_node)
     graph.add_node("player_black", player_black_node)
-    graph.add_node("chess_board_proxy", supervisor_chain)
+
+    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)
 
     graph.add_conditional_edges(
         "player_white", 
@@ -194,10 +191,6 @@ def create_graph():
         {"chess_board_proxy": "chess_board_proxy", END: END}
     )
     
-    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)
     graph.set_entry_point("chess_board_proxy")
     
     return graph.compile()
@@ -205,13 +198,13 @@ def create_graph():
 def should_continue(state):
     print("#### should_continue")
     global move_num, num_moves, legal_moves
+
     if move_num == num_moves:
-        print("False (move_num == num_moves)")
         return END # max moves reached
+    
     if not legal_moves:
-        print("False (not legal_moves)")
         return END # checkmate or stalemate
-    print("True")
+    
     return "chess_board_proxy"
 
 def initialize():
@@ -226,21 +219,17 @@ def initialize():
     legal_moves = ""
 
 def run_multi_agent(moves_num):
-    initialize()
-
     global num_moves
-
     num_moves = moves_num
-    
-    print("## START")
-    print("## num_moves=" + str(num_moves))
-    
+
+    initialize()
+
     graph = create_graph()
 
     result = ""
     
     try:
-        config = {"recursion_limit": 500}
+        config = {"recursion_limit": 100}
         
         result = graph.invoke({
             "messages": [
@@ -250,22 +239,10 @@ def run_multi_agent(moves_num):
     except Exception as e:
         print(f"An error occurred: {e}")
 
-    ###
-    
-    result2 = ""
+    result_md = ""
     num_move = 0
 
-    """
-    for message in result["messages"]:
-            if message.name:
-                print(f"{message.name}: {message.content}")
-            else:
-                print(message.content)
-    """
 
-    print("### "+ str(type(result)))
-    print("### "+ str(len(result["messages"])))
-    
     if "messages" in result:
         for message in result["messages"]:
             player = ""
@@ -276,7 +253,7 @@ def run_multi_agent(moves_num):
                 player = "Player White"
     
             if num_move > 0:
-                result2 += f"**{player}, Move {num_move}**\n{message.content}\n{board_svgs[num_move - 1]}\n\n"
+                result_md += f"**{player}, Move {num_move}**\n{message.content}\n{board_svgs[num_move - 1]}\n\n"
             
             num_move += 1
     
@@ -284,9 +261,7 @@ def run_multi_agent(moves_num):
                 break
     
     print("===")
-    print(str(result))
-    print("===")
-    print(str(result2))
+    print(result_md)
     print("===")
     
-    return str(result2)
+    return result_md