import asyncio import logging import os from datetime import datetime from typing import Annotated, Any, Dict, List, Literal, Optional, TypedDict, Union 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_core.tools import tool from langchain_openai import ChatOpenAI from langgraph.graph import END, StateGraph from langgraph.graph.message import add_messages from graphiti_core import Graphiti from graphiti_core.nodes import EpisodeType load_dotenv() logging.basicConfig( level=logging.DEBUG, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' ) logger = logging.getLogger('nba_agent') for name in logging.root.manager.loggerDict: if name != 'nba_agent': logging.getLogger(name).setLevel(logging.WARNING) # Initialize Graphiti client neo4j_uri = os.environ.get('NEO4J_URI', 'bolt://localhost:7687') neo4j_user = os.environ.get('NEO4J_USER', 'neo4j') neo4j_password = os.environ.get('NEO4J_PASSWORD', 'password') graphiti_client = Graphiti(neo4j_uri, neo4j_user, neo4j_password) openai_api_key = os.getenv('OPENAI_API_KEY') if not openai_api_key: logger.error('OPENAI_API_KEY is not set in the environment variables.') raise ValueError('OPENAI_API_KEY is not set') MAX_NEGOTIATION_ROUNDS = 5 class State(TypedDict): messages: List[HumanMessage] teams: Dict[str, Any] event: str team_actions: Dict[str, str] transfer_offers: List[Dict[str, Any]] current_negotiation: Optional[Dict[str, Any]] negotiation_rounds: int 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 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.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) 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:,}" async def propose_transfers(self) -> List[Dict[str, Any]]: """Propose transfer offers based on the agent's strategy.""" # This is a placeholder implementation. In a real scenario, this would involve more complex logic. if hasattr(self, 'last_proposed_transfer'): return [self.last_proposed_transfer] return [] 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) 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") else: logger.warning(f"{player_name} not found in {self.team_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") else: logger.warning(f"{player_name} already in {self.team_name}'s roster") async def get_transfer_offers(self): """Get the list of proposed transfers for this team.""" return self.proposed_transfers 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, 'to_team': to_team, 'player_name': player_name, 'proposed_price': proposed_price, } self.proposed_transfers.append(offer) 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)} How do you want to react to these actions? Consider if you need to adjust your strategy or make counter-moves.""" result = await self.executor.ainvoke( { 'input': reaction_prompt, 'agent_scratchpad': [], } ) 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}') 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).", 'agent_scratchpad': [], } ) logger.debug(f"{self.team_name}: Decision on transfer offer: {decision['output']}") return eval( decision['output'] ) # Convert the string representation of the dictionary to an actual dictionary async def handle_tool_use(self, response): if 'Action:' in response and 'Action Input:' in response: action = response.split('Action:')[1].split('Action Input:')[0].strip() action_input = response.split('Action Input:')[1].strip() try: tool = next(t for t in self.tools if t.name.lower() == action.lower()) if tool.name == 'execute_transfer': # Parse the action_input for execute_transfer inputs = eval(action_input) result = await tool.ainvoke(**inputs) else: result = await tool.ainvoke(input=action_input) return f'Tool execution result: {result}' except Exception as e: return f'Error executing tool {action}: {e}' return None async def add_episode(event_description: str): """Add a new episode to the Graphiti client.""" result = await graphiti_client.add_episode( name='New Event', episode_body=event_description, source_description='User Input', reference_time=datetime.now(), source=EpisodeType.message, ) return f"Episode '{event_description}' added successfully." async def invoke_tool(tool_name: str, **kwargs): tool = next(t for t in tools if t.name == tool_name) return await tool.ainvoke(input=kwargs) def get_fact_string(edge): return f'{edge.fact} {edge.valid_at or edge.created_at}' @tool async def get_team_roster(team_name: str): """Get the current roster for a specific team.""" search_result = await graphiti_client.search(f'plays for {team_name}', num_results=30) roster = [ edge.fact.split(' plays for ')[0] for edge in search_result if 'plays for' in edge.fact.lower() ] return f"{team_name}'s roster: {', '.join(roster)}" @tool async def search_player_info(player_name: str): """Search for information about a specific player.""" search_result = await graphiti_client.search(f'{player_name}', num_results=30) player_info = { 'name': player_name, 'facts': [get_fact_string(edge) for edge in search_result], } return player_info @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:,}.' @tool async def respond_to_transfer( player_name: str, from_team: str, to_team: str, response: str, counter_offer: int = None ): """Respond to a transfer proposal with an accept, reject, or counter-offer.""" response_message = f'{from_team} {response}s the transfer of {player_name} to {to_team}' if counter_offer: response_message += f' with a counter-offer of ${counter_offer:,}' return f'Transfer response: {response_message}.' @tool async def execute_transfer( player_name: str, from_team: str, to_team: str, price: int ) -> Dict[str, Any]: """Execute a transfer between two teams.""" # This is a simplified version. In a real scenario, you'd need to handle this more robustly. return { 'messages': [ HumanMessage( content=f'Transfer executed: {player_name} moved from {from_team} to {to_team} for ${price:,}' ) ], } @tool async def check_team_budget(team_name: str) -> Dict[str, Any]: """Check the current budget of a team.""" # This is a placeholder. In a real scenario, you'd fetch the actual budget. return { 'messages': [HumanMessage(content=f"Checking {team_name}'s budget...")], } @tool async def submit_transfer_offer( player_name: str, from_team: str, to_team: str, proposed_price: int ) -> Dict[str, Any]: """Submit a transfer offer for a player.""" logger.debug( f'submit_transfer_offer called with args: player_name={player_name}, from_team={from_team}, to_team={to_team}, proposed_price={proposed_price}' ) offer = { 'from_team': from_team, 'to_team': to_team, 'player_name': player_name, 'proposed_price': proposed_price, } logger.debug(f'Transfer offer created: {offer}') return { 'messages': [ HumanMessage( content=f'Transfer offer submitted: {from_team} offers to sell {player_name} to {to_team} for ${proposed_price:,}.' ) ], 'transfer_offers': [offer], } # Update the tools list tools = [ get_team_roster, search_player_info, propose_transfer, respond_to_transfer, execute_transfer, check_team_budget, submit_transfer_offer, ] 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: logger.debug('Entering process_event') return { **state, 'messages': state['messages'] + [HumanMessage(content=f"Event processed: {state['event']}")], } async def parallel_agent_processing(state: Dict[str, Any]) -> Dict[str, Any]: 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}') # 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}') 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: 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} async def select_negotiation(state: Dict[str, Any]) -> Dict[str, Any]: 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']: logger.debug('No transfer offers to negotiate') return { **state, 'messages': state['messages'] + [HumanMessage(content='No transfer offers to negotiate.')], 'negotiation_complete': True, } 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']] logger.debug(f'Negotiating transfer: {offer}') # 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'], ) # 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:,}." ) 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 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}') 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 = 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 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' # Define the workflow workflow = StateGraph(State) workflow.add_node('root', root) workflow.add_node('add_episode', add_episode_node) 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') 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('collect_transfer_offers', 'select_negotiation') workflow.add_edge('select_negotiation', 'negotiate_transfer') workflow.add_edge('negotiate_transfer', 'handle_tool_use') workflow.add_conditional_edges( 'handle_tool_use', should_continue, {END: END, 'collect_transfer_offers': 'collect_transfer_offers'}, ) 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 # Process the event state = await process_event(state) # Parallel agent processing state = await parallel_agent_processing(state) # 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) 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:,}') print('Simulation ended.') def main(): asyncio.run(run_simulation()) if __name__ == '__main__': main()