diff --git a/examples/nba/nba_agent.py b/examples/nba/nba_agent.py index c931a781..9fcfd5b4 100644 --- a/examples/nba/nba_agent.py +++ b/examples/nba/nba_agent.py @@ -1,18 +1,17 @@ import asyncio import logging import os -from datetime import datetime -from typing import Annotated, Any, Dict, List, Literal, Optional, TypedDict, Union +from typing import TypedDict, Dict, List, Optional, Any from dotenv import load_dotenv from langchain.agents import AgentExecutor, create_openai_functions_agent -from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder -from langchain_core.messages import AIMessage, HumanMessage, SystemMessage +from langchain.prompts import PromptTemplate +from langchain.schema import AIMessage, HumanMessage, SystemMessage from langchain_core.tools import tool from langchain_openai import ChatOpenAI from langgraph.graph import END, StateGraph -from langgraph.graph.message import add_messages - +from langgraph.prebuilt.tool_executor import ToolExecutor +from datetime import datetime from graphiti_core import Graphiti from graphiti_core.nodes import EpisodeType @@ -39,9 +38,13 @@ if not openai_api_key: MAX_NEGOTIATION_ROUNDS = 5 -class State(TypedDict): - messages: List[HumanMessage] - teams: Dict[str, Any] +def format_step_result(messages: List[str], **kwargs) -> Dict[str, Any]: + return {'messages': messages, **kwargs} + + +class SimulationState(TypedDict): + messages: List[str] # Changed from HumanMessage to str for simplicity + teams: Dict[str, Dict[str, Any]] # Store team data as a dictionary event: str team_actions: Dict[str, str] transfer_offers: List[Dict[str, Any]] @@ -50,49 +53,93 @@ class State(TypedDict): negotiation_complete: bool -def format_step_result(messages: List[str], **kwargs) -> Dict[str, Any]: - return {'messages': messages, **kwargs} - - class TeamAgent: - def __init__( - self, team_name: str, tools: List[Any], budget: int = 100000000, roster: List[str] = [] - ): - # logger.debug(f'Initializing TeamAgent for {team_name}') - self.team_name = team_name + def __init__(self, name: str, tools: List[Any]): + self.name = name + self.roster: List[str] = [] + self.budget: int = 100_000_000 self.tools = tools - self.budget = budget - self.roster = [] - self.proposed_transfers = [] - self.prompt = ChatPromptTemplate.from_messages( - [ - ( - 'system', - f"""You are an AI assistant managing the {team_name} NBA team. - Your role is to make strategic decisions for your team, react to events, and interact with other teams. - Use the available tools to gather information and perform actions. - When an event occurs, decide how to react. You can: - 1. Use tools to gather more information about players or team situations. - 2. Propose player transfers (buy or sell) based on events and your team's needs. - 3. Set transfer prices based on player performance and your available budget. - 4. Negotiate with other teams on transfer prices. - Always consider what's best for your team in the long term. Be strategic and competitive. - The only teams in this simulation are: Toronto Raptors, Boston Celtics, and Golden State Warriors. - Only interact with these teams and do not mention or propose transfers to any other teams.""", - ), - ('human', '{input}'), - MessagesPlaceholder(variable_name='agent_scratchpad'), - ] + self.last_proposed_transfer: Optional[Dict[str, Any]] = None + + # Create the language model + llm = ChatOpenAI(temperature=0.3) + + # Create the agent executor + template = """You are the manager of the {team_name} NBA team. Make decisions to improve your team. + +Current event: {event} + +Your task is to decide on an action based on the event. Use the available tools to gather information and make decisions. Do not ask for further input. Instead, take action based on the information you have. + +If you decide to propose a transfer, use the propose_transfer tool and include the exact output from the tool in your response, prefixed with "TRANSFER PROPOSAL:". + +{agent_scratchpad}""" + + prompt = PromptTemplate( + input_variables=['team_name', 'event', 'agent_scratchpad'], template=template ) - self.llm = ChatOpenAI(temperature=0.2) - self.agent = create_openai_functions_agent(self.llm, self.tools, self.prompt) - self.executor = AgentExecutor(agent=self.agent, tools=self.tools, verbose=True) + + agent = create_openai_functions_agent(llm, self.tools, prompt) + self.executor = AgentExecutor(agent=agent, tools=self.tools, verbose=True) + + async def process_event(self, event: str, relevant_offers: List[Dict[str, Any]]) -> str: + logger.debug(f'{self.name}: Processing event: {event}') + logger.debug(f'{self.name}: Current roster: {self.roster}') + logger.debug(f'{self.name}: Current budget: {self.budget}') + + try: + result = await self.executor.ainvoke( + { + 'team_name': self.name, + 'event': event, + 'roster': self.roster, + 'budget': self.budget, + 'relevant_offers': relevant_offers, + } + ) + logger.debug(f"{self.name}: Agent output: {result['output']}") + + # Check if a transfer was proposed + if 'TRANSFER PROPOSAL:' in result['output']: + # Parse the transfer details and set last_proposed_transfer + transfer_details = result['output'].split('TRANSFER PROPOSAL:')[-1].strip() + self.last_proposed_transfer = self.parse_transfer_proposal(transfer_details) + logger.debug(f'{self.name}: Proposed transfer: {self.last_proposed_transfer}') + + return result['output'] + except Exception as e: + logger.error(f'{self.name}: Error processing event: {str(e)}') + return f'Error processing event: {str(e)}' + + def parse_transfer_proposal(self, proposal: str) -> Dict[str, Any]: + # More flexible parsing + parts = proposal.lower().replace(',', '').split() + to_team = next(parts[i - 1] for i, word in enumerate(parts) if word == 'to') + from_team = next(parts[i - 1] for i, word in enumerate(parts) if word == 'from') + player_name = next(parts[i + 1] for i, word in enumerate(parts) if word == 'buy') + proposed_price = int(''.join(filter(str.isdigit, parts[-1]))) + + return { + 'to_team': to_team, + 'from_team': from_team, + 'player_name': player_name, + 'proposed_price': proposed_price, + } + + async def propose_transfer(self, player_name: str, to_team: str, price: int): + self.last_proposed_transfer = { + 'from_team': self.name, + 'to_team': to_team, + 'player_name': player_name, + 'proposed_price': price, + } + return f'Proposed transfer of {player_name} to {to_team} for ${price:,}' async def update_budget(self, amount: int): """Update the team's budget.""" self.budget += amount - logger.debug(f"{self.team_name}'s new budget: ${self.budget:,}") - return f"{self.team_name}'s new budget: ${self.budget:,}" + logger.debug(f"{self.name}'s new budget: ${self.budget:,}") + return f"{self.name}'s new budget: ${self.budget:,}" async def propose_transfers(self) -> List[Dict[str, Any]]: """Propose transfer offers based on the agent's strategy.""" @@ -104,24 +151,24 @@ class TeamAgent: async def update_roster(self): """Update the team's roster using the get_team_roster tool.""" roster_tool = next(tool for tool in self.tools if tool.name == 'get_team_roster') - roster_string = await roster_tool.ainvoke(self.team_name) + roster_string = await roster_tool.ainvoke(self.name) self.roster = roster_string.split(': ')[1].split(', ') def remove_player(self, player_name: str): """Remove a player from the team's roster.""" if player_name in self.roster: self.roster.remove(player_name) - logger.debug(f"{player_name} removed from {self.team_name}'s roster") + logger.debug(f"{player_name} removed from {self.name}'s roster") else: - logger.warning(f"{player_name} not found in {self.team_name}'s roster") + logger.warning(f"{player_name} not found in {self.name}'s roster") def add_player(self, player_name: str): """Add a player to the team's roster.""" if player_name not in self.roster: self.roster.append(player_name) - logger.debug(f"{player_name} added to {self.team_name}'s roster") + logger.debug(f"{player_name} added to {self.name}'s roster") else: - logger.warning(f"{player_name} already in {self.team_name}'s roster") + logger.warning(f"{player_name} already in {self.name}'s roster") async def get_transfer_offers(self): """Get the list of proposed transfers for this team.""" @@ -130,7 +177,7 @@ class TeamAgent: async def submit_transfer_offer(self, player_name: str, to_team: str, proposed_price: int): """Submit a transfer offer for a player.""" offer = { - 'from_team': self.team_name, + 'from_team': self.name, 'to_team': to_team, 'player_name': player_name, 'proposed_price': proposed_price, @@ -139,37 +186,6 @@ class TeamAgent: logger.debug(f'Transfer offer submitted: {offer}') return offer - async def process_event( - self, state: Dict[str, Any], event: str, relevant_offers: List[Dict[str, Any]] = [] - ): - logger.debug(f'{self.team_name}: Processing event: {event}') - logger.debug(f'{self.team_name}: Current roster: {self.roster}') - logger.debug(f'{self.team_name}: Current budget: {self.budget}') - - try: - result = await self.executor.ainvoke( - { - 'input': f'Event: {event}\n\nCurrent team: {self.team_name}\nCurrent budget: ${self.budget:,}\nCurrent roster: {", ".join(self.roster)}\nRelevant transfer offers: {relevant_offers}\n\nReact to this event and consider the transfer offers. Make decisions and take actions without asking for confirmation.', - 'agent_scratchpad': [], - } - ) - logger.debug(f"{self.team_name}: Finished processing event. Result: {result['output']}") - - # Check if propose_transfer was called - if 'propose_transfer' in result['output']: - # Extract transfer details - transfer_details = result['output'].split('propose_transfer(')[1].split(')')[0] - transfer_dict = eval(f'dict({transfer_details})') - self.last_proposed_transfer = transfer_dict - logger.debug( - f'{self.team_name}: Set last_proposed_transfer to {self.last_proposed_transfer}' - ) - - return result['output'] - except Exception as e: - logger.error(f'{self.team_name}: Error processing event: {str(e)}', exc_info=True) - raise - async def react_to_others(self, other_actions: List[str]) -> str: reaction_prompt = f"""Other teams have taken the following actions: {' '.join(other_actions)} @@ -185,16 +201,16 @@ class TeamAgent: return result['output'] async def decide_on_transfer(self, offer: Dict[str, Any]) -> Dict[str, Any]: - logger.debug(f'{self.team_name}: Deciding on transfer offer: {offer}') + logger.debug(f'{self.name}: Deciding on transfer offer: {offer}') decision = await self.executor.ainvoke( { - 'input': f"Transfer offer received:\nPlayer: {offer['player_name']}\nFrom: {offer['from_team']}\nTo: {self.team_name}\nPrice: ${offer['proposed_price']:,}\n\nMake a decision to accept, reject, or counter-offer. Respond with a dictionary containing 'action' (accept/reject/counter) and 'counter_offer' (if applicable).", + 'input': f"Transfer offer received:\nPlayer: {offer['player_name']}\nFrom: {offer['from_team']}\nTo: {self.name}\nPrice: ${offer['proposed_price']:,}\n\nMake a decision to accept, reject, or counter-offer. Respond with a dictionary containing 'action' (accept/reject/counter) and 'counter_offer' (if applicable).", 'agent_scratchpad': [], } ) - logger.debug(f"{self.team_name}: Decision on transfer offer: {decision['output']}") + logger.debug(f"{self.name}: Decision on transfer offer: {decision['output']}") return eval( decision['output'] ) # Convert the string representation of the dictionary to an actual dictionary @@ -216,6 +232,14 @@ class TeamAgent: return f'Error executing tool {action}: {e}' return None + def to_dict(self) -> Dict[str, Any]: + return { + 'name': self.name, + 'roster': self.roster, + 'budget': self.budget, + 'last_proposed_transfer': self.last_proposed_transfer, + } + async def add_episode(event_description: str): """Add a new episode to the Graphiti client.""" @@ -264,7 +288,7 @@ async def search_player_info(player_name: str): @tool async def propose_transfer(player_name: str, from_team: str, to_team: str, proposed_price: int): """Propose a player transfer from one team to another with a proposed price.""" - return f'Transfer proposal: {to_team} wants to buy {player_name} from {from_team} for ${proposed_price:,}.' + return f'TRANSFER PROPOSAL: {to_team} wants to buy {player_name} from {from_team} for ${proposed_price:,}.' @tool @@ -339,354 +363,192 @@ tools = [ ] -def root(state: dict) -> dict: - logger.debug(f'Entering root node. State: {state}') - return { - **state, - 'messages': state['messages'] - + [HumanMessage(content=f"Processing event: {state['event']}")], - } - - -async def add_episode_node(state: dict) -> dict: - logger.debug(f'Entering add_episode_node. State: {state}') - logger.debug(f"Adding episode: {state['event']}") - result = await add_episode(state['event']) - logger.debug(f'Episode added successfully: {result}') - return {**state, 'messages': state['messages'] + [HumanMessage(content=result)]} - - -async def process_event(state: dict) -> dict: +def process_event(state: SimulationState) -> SimulationState: logger.debug('Entering process_event') + new_message = f"Event processed: {state['event']}" return { **state, - 'messages': state['messages'] - + [HumanMessage(content=f"Event processed: {state['event']}")], + 'messages': state.get('messages', []) + [new_message], } -async def parallel_agent_processing(state: Dict[str, Any]) -> Dict[str, Any]: +async def parallel_agent_processing(state: SimulationState) -> SimulationState: logger.debug('Entering parallel_agent_processing') - logger.debug(f'Current state: {state}') tasks = [] - for team_name, agent in state['teams'].items(): - logger.debug(f"{team_name}: Starting to process event: {state['event']}") - relevant_offers = [ - offer for offer in state.get('transfer_offers', []) if offer['to_team'] == team_name - ] - logger.debug(f'{team_name}: Relevant offers: {relevant_offers}') - tasks.append(agent.process_event(state, state['event'], relevant_offers)) - results = await asyncio.gather(*tasks) - logger.debug(f'Parallel processing results: {results}') + team_agents = {} + for team_name, team_data in state['teams'].items(): + team_agent = TeamAgent(team_data['name'], tools) + team_agents[team_name] = team_agent + tasks.append( + asyncio.create_task(team_agent.process_event(state['event'], state['transfer_offers'])) + ) - # Merge transfer offers from all agents - all_transfer_offers = state.get('transfer_offers', []) - for agent in state['teams'].values(): - if hasattr(agent, 'last_proposed_transfer'): - all_transfer_offers.append(agent.last_proposed_transfer) - logger.debug(f'Added transfer offer: {agent.last_proposed_transfer}') + results = await asyncio.gather(*tasks, return_exceptions=True) + + updated_state = state.copy() + updated_state['transfer_offers'] = [] + for i, (team_name, team_data) in enumerate(state['teams'].items()): + if isinstance(results[i], Exception): + logger.error(f'Error processing event for {team_name}: {str(results[i])}') + updated_state['team_actions'][team_name] = f'Error: {str(results[i])}' + else: + logger.debug(f'Team {team_name} action: {results[i]}') + updated_state['team_actions'][team_name] = results[i] + if team_agents[team_name].last_proposed_transfer: + updated_state['transfer_offers'].append( + team_agents[team_name].last_proposed_transfer + ) + updated_state['teams'][team_name] = team_agents[team_name].to_dict() - updated_state = { - **state, - 'team_actions': {team: action for team, action in zip(state['teams'].keys(), results)}, - 'transfer_offers': all_transfer_offers, - } logger.debug(f'Updated state after parallel processing: {updated_state}') return updated_state -async def agent_reaction_phase(state: dict) -> dict: - logger.debug('Entering agent_reaction_phase') - messages = [] - for team_name, action in state['team_actions'].items(): - messages.append(f'{team_name}: {action}') - return { - **state, - 'messages': state['messages'] + [HumanMessage(content=msg) for msg in messages], - } - - -async def collect_transfer_offers(state: State) -> State: +async def collect_transfer_offers(state: SimulationState) -> SimulationState: logger.debug('Entering collect_transfer_offers') - transfer_offers = [] - for team_name, agent in state['teams'].items(): - offers = await agent.propose_transfers() - transfer_offers.extend(offers) - return {**state, 'transfer_offers': state['transfer_offers'] + transfer_offers} + updated_state = state.copy() + logger.debug(f'Collected transfer offers: {updated_state}') + # The transfer offers are already collected in parallel_agent_processing + logger.debug(f"Collected transfer offers: {updated_state['transfer_offers']}") + return updated_state -async def select_negotiation(state: Dict[str, Any]) -> Dict[str, Any]: +def select_negotiation(state: SimulationState) -> SimulationState: logger.debug('Entering select_negotiation') - logger.debug(f'Current state in select_negotiation: {state}') - if not state.get('transfer_offers', []): - return { - **state, - 'messages': state['messages'] - + [HumanMessage(content='No transfer offers to negotiate.')], - 'current_negotiation': None, - } - - selected_offer = state['transfer_offers'][0] - return { - **state, - 'messages': state['messages'] - + [ - HumanMessage( - content=f"Selected negotiation: {selected_offer['player_name']} from {selected_offer['from_team']} to {selected_offer['to_team']}" - ) - ], - 'current_negotiation': selected_offer, - } - - -async def negotiate_transfer(state: dict) -> dict: - logger.debug('Entering negotiate_transfer') if not state['transfer_offers']: + return {**state, 'current_negotiation': None, 'negotiation_complete': True} + return {**state, 'current_negotiation': state['transfer_offers'][0]} + + +async def negotiate_transfer(state: SimulationState) -> SimulationState: + logger.debug('Entering negotiate_transfer') + updated_state = state.copy() + + if not updated_state['transfer_offers']: logger.debug('No transfer offers to negotiate') - return { - **state, - 'messages': state['messages'] - + [HumanMessage(content='No transfer offers to negotiate.')], - 'negotiation_complete': True, - } + return updated_state - offer = state['transfer_offers'][0] # Take the first offer to negotiate - from_agent = state['teams'][offer['from_team']] - to_agent = state['teams'][offer['to_team']] + # Sort offers by proposed price (highest first) + sorted_offers = sorted( + updated_state['transfer_offers'], key=lambda x: x['proposed_price'], reverse=True + ) + best_offer = sorted_offers[0] - logger.debug(f'Negotiating transfer: {offer}') + # Simulate negotiation + from_team = updated_state['teams'][best_offer['from_team']] + to_team = updated_state['teams'][best_offer['to_team']] - # Ask the receiving team to make a decision - decision = await to_agent.decide_on_transfer(offer) - logger.debug(f'Decision from {to_agent.team_name}: {decision}') - - if decision['action'] == 'accept': - # Execute the transfer - execute_transfer_tool = next(tool for tool in tools if tool.name == 'execute_transfer') - transfer_result = await execute_transfer_tool.ainvoke( - player_name=offer['player_name'], - from_team=offer['from_team'], - to_team=offer['to_team'], - price=offer['proposed_price'], + # Simple negotiation logic: accept if the price is above a threshold + threshold = 50 # This can be adjusted + if best_offer['proposed_price'] > threshold: + logger.info( + f"Transfer accepted: {best_offer['player_name']} from {best_offer['from_team']} to {best_offer['to_team']} for ${best_offer['proposed_price']}" ) - # Update team rosters and budgets - from_agent.remove_player(offer['player_name']) - to_agent.add_player(offer['player_name']) - await from_agent.update_budget(offer['proposed_price']) - await to_agent.update_budget(-offer['proposed_price']) - - transfer_message = f"{offer['player_name']} transferred from {offer['from_team']} to {offer['to_team']} for ${offer['proposed_price']:,}" - logger.debug(transfer_message) - - return { - **state, - 'messages': state['messages'] + [HumanMessage(content=transfer_message)], - 'transfer_offers': state['transfer_offers'][1:], # Remove the processed offer - 'negotiation_complete': True, - } - elif decision['action'] == 'reject': - reject_message = f"{offer['to_team']} rejected the offer for {offer['player_name']}." - logger.debug(reject_message) - return { - **state, - 'messages': state['messages'] + [HumanMessage(content=reject_message)], - 'transfer_offers': state['transfer_offers'][1:], # Remove the rejected offer - 'negotiation_complete': True, - } - elif decision['action'] == 'counter': - counter_offer = decision['counter_offer'] - counter_message = ( - f"{offer['to_team']} counter-offered for {offer['player_name']} at ${counter_offer:,}." + from_team['roster'].remove(best_offer['player_name']) + to_team['roster'].append(best_offer['player_name']) + from_team['budget'] += best_offer['proposed_price'] + to_team['budget'] -= best_offer['proposed_price'] + updated_state['negotiation_complete'] = True + else: + logger.info( + f"Transfer rejected: {best_offer['player_name']} from {best_offer['from_team']} to {best_offer['to_team']} for ${best_offer['proposed_price']}" ) - logger.debug(counter_message) - # Update the offer with the new price - updated_offer = {**offer, 'proposed_price': counter_offer} - return { - **state, - 'messages': state['messages'] + [HumanMessage(content=counter_message)], - 'transfer_offers': [updated_offer] + state['transfer_offers'][1:], - 'negotiation_complete': False, # Continue negotiation - } - return state + updated_state['current_negotiation'] = None + updated_state['transfer_offers'] = [] + return updated_state -async def handle_tool_use(state: Dict[str, Any]) -> Dict[str, Any]: - logger.debug('Entering handle_tool_use') - logger.debug(f'Current state in handle_tool_use: {state}') +async def execute_transfer(state: SimulationState, offer: Dict[str, Any]) -> None: + from_agent = TeamAgent( + state['teams'][offer['from_team']]['name'], tools + ) # Recreate TeamAgent from data + to_agent = TeamAgent( + state['teams'][offer['to_team']]['name'], tools + ) # Recreate TeamAgent from data - executed_transfers = [] - for team_name, agent in state['teams'].items(): - logger.debug(f'Checking for last_proposed_transfer on {team_name}') - if hasattr(agent, 'last_proposed_transfer'): - transfer = agent.last_proposed_transfer - logger.debug(f'Processing transfer: {transfer}') + from_agent.roster.remove(offer['player_name']) + to_agent.roster.append(offer['player_name']) + from_agent.budget += offer['proposed_price'] + to_agent.budget -= offer['proposed_price'] - from_agent = state['teams'][transfer['from_team']] - to_agent = state['teams'][transfer['to_team']] - - # Execute the transfer - execute_transfer_tool = next(tool for tool in tools if tool.name == 'execute_transfer') - transfer_result = await execute_transfer_tool.ainvoke( - player_name=transfer['player_name'], - from_team=transfer['from_team'], - to_team=transfer['to_team'], - price=transfer['proposed_price'], - ) - - logger.debug(f'Transfer result: {transfer_result}') - - # Update team rosters and budgets - logger.debug(f'Before transfer - {from_agent.team_name} roster: {from_agent.roster}') - logger.debug(f'Before transfer - {to_agent.team_name} roster: {to_agent.roster}') - - from_agent.remove_player(transfer['player_name']) - to_agent.add_player(transfer['player_name']) - - logger.debug(f'After transfer - {from_agent.team_name} roster: {from_agent.roster}') - logger.debug(f'After transfer - {to_agent.team_name} roster: {to_agent.roster}') - - await from_agent.update_budget(transfer['proposed_price']) - await to_agent.update_budget(-transfer['proposed_price']) - - logger.debug(f'Updated {from_agent.team_name} budget: {from_agent.budget}') - logger.debug(f'Updated {to_agent.team_name} budget: {to_agent.budget}') - - # Add transfer result to messages - state['messages'].append(HumanMessage(content=transfer_result['messages'][0].content)) - executed_transfers.append(transfer) - - # Clear the last proposed transfer - delattr(agent, 'last_proposed_transfer') - else: - logger.debug(f'No last_proposed_transfer found for {team_name}') - - # Add executed transfers to the state - state['executed_transfers'] = executed_transfers - logger.debug(f'Executed transfers: {executed_transfers}') - - return state + transfer_message = f"{offer['player_name']} transferred from {offer['from_team']} to {offer['to_team']} for ${offer['proposed_price']:,}" + state['messages'].append(transfer_message) -def should_continue(x: dict): - if x.get('negotiation_complete', False): - return END - if len(x.get('transfer_offers', [])) == 0 and x.get('current_negotiation') is None: - return END - return 'collect_transfer_offers' +def should_continue(state: SimulationState) -> List[str]: + if state['negotiation_complete'] and not state['transfer_offers']: + return [END] + return ['select_negotiation'] -# Define the workflow -workflow = StateGraph(State) +# Define the graph +workflow = StateGraph(SimulationState) -workflow.add_node('root', root) -workflow.add_node('add_episode', add_episode_node) +# Add nodes workflow.add_node('process_event', process_event) workflow.add_node('parallel_agent_processing', parallel_agent_processing) -workflow.add_node('agent_reaction_phase', agent_reaction_phase) workflow.add_node('collect_transfer_offers', collect_transfer_offers) workflow.add_node('select_negotiation', select_negotiation) workflow.add_node('negotiate_transfer', negotiate_transfer) -workflow.add_node('handle_tool_use', handle_tool_use) -workflow.set_entry_point('root') - -workflow.add_edge('root', 'add_episode') -workflow.add_edge('add_episode', 'process_event') +# Add edges workflow.add_edge('process_event', 'parallel_agent_processing') -workflow.add_edge('parallel_agent_processing', 'agent_reaction_phase') -workflow.add_edge('agent_reaction_phase', 'collect_transfer_offers') +workflow.add_edge('parallel_agent_processing', 'collect_transfer_offers') workflow.add_edge('collect_transfer_offers', 'select_negotiation') workflow.add_edge('select_negotiation', 'negotiate_transfer') -workflow.add_edge('negotiate_transfer', 'handle_tool_use') +# Add conditional edge workflow.add_conditional_edges( - 'handle_tool_use', - should_continue, - {END: END, 'collect_transfer_offers': 'collect_transfer_offers'}, + 'negotiate_transfer', should_continue, {'select_negotiation': 'select_negotiation', END: END} ) +# Set the entrypoint +workflow.set_entry_point('process_event') + +# Compile the graph app = workflow.compile() -teams = { - 'Toronto Raptors': TeamAgent('Toronto Raptors', tools), - 'Boston Celtics': TeamAgent('Boston Celtics', tools), - 'Golden State Warriors': TeamAgent('Golden State Warriors', tools), -} - -# Update the initial_state -initial_state: State = { - 'messages': [], - 'teams': teams, - 'event': '', - 'team_actions': {}, - 'transfer_offers': [], - 'current_negotiation': None, - 'negotiation_rounds': 0, - 'negotiation_complete': False, -} - async def run_simulation(): - state = { - 'messages': [], - 'teams': teams, - 'event': None, - 'team_actions': {}, - 'transfer_offers': [], - 'current_negotiation': None, - 'negotiation_rounds': 0, - 'negotiation_complete': False, - } - while True: event = input("Enter an event (or 'quit' to exit): ") if event.lower() == 'quit': break - state['event'] = event - state['messages'] = [] - state['team_actions'] = {} - state['transfer_offers'] = [] - state['current_negotiation'] = None - state['negotiation_rounds'] = 0 - state['negotiation_complete'] = False + initial_state = SimulationState( + messages=[], + teams={ + 'Toronto Raptors': TeamAgent('Toronto Raptors', tools).to_dict(), + 'Boston Celtics': TeamAgent('Boston Celtics', tools).to_dict(), + 'Golden State Warriors': TeamAgent('Golden State Warriors', tools).to_dict(), + }, + event=event, + team_actions={}, + transfer_offers=[], + current_negotiation=None, + negotiation_rounds=0, + negotiation_complete=False, + ) - # Process the event - state = await process_event(state) + async for state in app.astream(initial_state): + if 'messages' in state: + for message in state['messages']: + print(message) - # Parallel agent processing - state = await parallel_agent_processing(state) + if 'transfer_offers' in state: + print(f"Current transfer offers: {state['transfer_offers']}") - # Collect transfer offers - state = await collect_transfer_offers(state) - - # Negotiate transfers - while ( - not state['negotiation_complete'] - and state['negotiation_rounds'] < MAX_NEGOTIATION_ROUNDS - ): - state = await negotiate_transfer(state) - state['negotiation_rounds'] += 1 - - # Handle tool use (execute transfers, etc.) - state = await handle_tool_use(state) - - # Display results - for message in state['messages']: - print(message.content) + if 'current_negotiation' in state: + print(f"Current negotiation: {state['current_negotiation']}") print('\nFinal team states:') - for team_name, team_agent in state['teams'].items(): - print(f'{team_name} - Roster: {team_agent.roster}, Budget: ${team_agent.budget:,}') + for team_name, team_data in initial_state['teams'].items(): + print(f"{team_name} - Roster: {team_data['roster']}, Budget: ${team_data['budget']:,}") - print('Simulation ended.') - - -def main(): - asyncio.run(run_simulation()) + print('\n' + '=' * 50 + '\n') if __name__ == '__main__': - main() + asyncio.run(run_simulation())