bstraehle/multi-agent-ai-langgraph-chess

1import chess, chess.svg, math2import functools, operator3 4from datetime import date5 6from typing import Annotated, Any, Dict, List, Optional, Sequence, Tuple, TypedDict, Union7 8from langchain.agents import AgentExecutor, create_openai_tools_agent9from langchain_community.tools.tavily_search import TavilySearchResults10from langchain_core.messages import BaseMessage, HumanMessage11from langchain_core.output_parsers.openai_functions import JsonOutputFunctionsParser12from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder13from langchain_core.tools import tool14from langchain_openai import ChatOpenAI15 16from langgraph.graph import StateGraph, END17        18board = None19board_svgs = None20 21num_moves = 022move_num = 023 24legal_moves = ""25 26class AgentState(TypedDict):27    messages: Annotated[Sequence[BaseMessage], operator.add]28    next: str29 30def create_agent(llm: ChatOpenAI, tools: list, system_prompt: str):31    prompt = ChatPromptTemplate.from_messages(32        [33            ("system", system_prompt),34            MessagesPlaceholder(variable_name="messages"),35            MessagesPlaceholder(variable_name="agent_scratchpad"),36        ]37    )38    39    agent = create_openai_tools_agent(llm, tools, prompt)40    41    return AgentExecutor(agent=agent, 42                         tools=tools,43                         handle_parsing_errors=True,44                         return_intermediate_steps=True,45                         verbose=True,46                         max_iterations=5)47 48def agent_node(state, agent, name):49    try:50        #print(f"agent node: {name}")51        result = agent.invoke(state)52        return {"messages": [HumanMessage(content=result["output"], name=name)]}53    except Exception as e:54        print(f"An error occurred in agent_node: {e}")55        return {"messages": [HumanMessage(content=f"Error: {e}", name=name)]}56 57@tool58def get_legal_moves() -> Annotated[str, "A list of legal moves in UCI format"]:59    """Returns a list of legal moves in UCI format. 60       The input should always be an empty string, 61       and this function will always return legal moves in UCI format."""62    try:63        global legal_moves64        legal_moves = ",".join([str(move) for move in board.legal_moves])65        return legal_moves66    except Exception as e:67        print(f"An error occurred in get_legal_moves: {e}")68        return "Error: unable to get legal moves"69 70@tool71def make_move(move: Annotated[str, "A move in UCI format."]) -> Annotated[str, "Result of the move."]:72    """Makes a move. 73       The input should always be a move in UCI format, 74       and this function will always return the result of the move."""75    try:76        move = chess.Move.from_uci(move)77        board.push_uci(str(move))78 79        global move_num80        move_num += 181        print(f"move_num: {str(move_num)}")82        83        board_svgs.append(chess.svg.board(84            board,85            arrows=[(move.from_square, move.to_square)],86            fill={move.from_square: "gray"},87            size=60088        ))89 90        piece = board.piece_at(move.to_square)91        piece_symbol = piece.unicode_symbol()92        piece_name = (93            chess.piece_name(piece.piece_type).capitalize()94            if piece_symbol.isupper()95            else chess.piece_name(piece.piece_type)96        )97        98        return f"Moved {piece_name} ({piece_symbol}) from "\99               f"{chess.SQUARE_NAMES[move.from_square]} to "\100               f"{chess.SQUARE_NAMES[move.to_square]}."101    except Exception as e:102        print(f"An error occurred in make_move: {e}")103        return f"Error: unable to make move {move}"104    105def create_graph(llm_board, llm_white, llm_black):106    players = ["player_white", "player_black"]107    options = players108 109    llm_board_proxy = ChatOpenAI(model=llm_board)110    llm_player_white = ChatOpenAI(model=llm_white)111    llm_player_black = ChatOpenAI(model=llm_black)112    113    system_prompt = (114        "You are a Chess Board Proxy tasked with managing a game of chess "115        "between player_white and player_black. player_white makes the first move, "116        "then the players take turns."117    )118 119    function_def = {120        "name": "route",121        "description": "Select the next player.",122        "parameters": {123            "title": "routeSchema",124            "type": "object",125            "properties": {126                "next": {127                    "title": "Next",128                    "anyOf": [129                        {"enum": options},130                    ],131                }132            },133            "required": ["next"],134        },135    }136    137    prompt = ChatPromptTemplate.from_messages(138        [139            ("system", system_prompt),140            MessagesPlaceholder(variable_name="messages"),141            (142                "system",143                "If player_white made a move, player_black must make the next move. "144                "If player_black made a move, player_white must make the next move. "145                "Select one of: {options}.",146            ),147        ]148    ).partial(options=str(options), members=", ".join(players), verbose=True)149    150    supervisor_chain = (151        prompt152        | llm_board_proxy.bind_functions(functions=[function_def], function_call="route")153        | JsonOutputFunctionsParser()154    )155 156    player_white_agent = create_agent(llm_player_white, [get_legal_moves, make_move], system_prompt=157                                     "You are a chess Grandmaster and you play as white. "158                                     "First call get_legal_moves() to get a list of legal moves. "159                                     "Then study the returned moves and call make_move(move) to make the best move. "160                                     "Finally analyze the move: **Analysis:** move in UCI format, emoji of piece, unordered list of 3 items.")161    player_white_node = functools.partial(agent_node, agent=player_white_agent, name="player_white")162 163    player_black_agent = create_agent(llm_player_black, [get_legal_moves, make_move], system_prompt=164                                     "You are a chess Grandmaster and you play as black. "165                                     "First call get_legal_moves() to get a list of legal moves. "166                                     "Then study the returned moves and call make_move(move) to make the best move. "167                                     "Finally analyze the move: **Analysis:** move in UCI format, emoji of piece, unordered list of 3 items.")168    player_black_node = functools.partial(agent_node, agent=player_black_agent, name="player_black")169    170    graph = StateGraph(AgentState)171 172    graph.add_node("chess_board_proxy", supervisor_chain)173    graph.add_node("player_white", player_white_node)174    graph.add_node("player_black", player_black_node)175 176    conditional_map = {k: k for k in players}177    graph.add_conditional_edges("chess_board_proxy", lambda x: x["next"], conditional_map)178 179    graph.add_conditional_edges(180        "player_white", 181        should_continue, 182        {"chess_board_proxy": "chess_board_proxy", END: END}183    )184 185    graph.add_conditional_edges(186        "player_black", 187        should_continue, 188        {"chess_board_proxy": "chess_board_proxy", END: END}189    )190    191    graph.set_entry_point("chess_board_proxy")192    193    return graph.compile()194 195def should_continue(state):196    global move_num, num_moves, legal_moves197    198    if move_num == num_moves:199        return END # max moves reached200    201    if not legal_moves:202        return END # checkmate or stalemate203 204    return "chess_board_proxy"205 206def initialize():207    global board, board_svgs, num_moves, move_num, legal_moves208 209    board = chess.Board()210    board_svgs = []211 212    num_moves = 0213    move_num = 0214    215    legal_moves = ""216 217def run_multi_agent(llm_board, llm_white, llm_black, moves_num):218    initialize()219    220    global num_moves221    num_moves = moves_num222 223    graph = create_graph(llm_board, llm_white, llm_black)224 225    result = ""226    227    try:228        config = {"recursion_limit": 100}229        230        result = graph.invoke({231            "messages": [232                HumanMessage(content="Let's play chess, player_white starts.")233            ]234        }, config=config)235    except Exception as e:236        print(f"An error occurred: {e}")237 238    result_md = ""239    num_move = 0240 241    if "messages" in result:242        for message in result["messages"]:243            player = ""244            245            if num_move % 2 == 0:246                player = "Player Black"247            else:248                player = "Player White"249    250            if num_move > 0:251                result_md += f"**{player}, Move {num_move}**\n{message.content}\n{board_svgs[num_move - 1]}\n\n"252            253            num_move += 1254    255            if num_moves % 2 == 0 and num_move == num_moves + 1:256                break257    258    return result_md