gmakstutis/graphiti
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

wip

Browse source
This commit is contained in:
paulpaliychuk 2024-08-28 23:06:31 -04:00
parent 6ce6802796
commit b773c281c9
1 changed files with 232 additions and 370 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

602
examples/nba/nba_agent.py
View file

@ -1,18 +1,17 @@
import asyncio import asyncio
import logging import logging
import os import os
from datetime import datetime from typing import TypedDict, Dict, List, Optional, Any
from typing import Annotated, Any, Dict, List, Literal, Optional, TypedDict, Union
from dotenv import load_dotenv from dotenv import load_dotenv
from langchain.agents import AgentExecutor, create_openai_functions_agent from langchain.agents import AgentExecutor, create_openai_functions_agent
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder from langchain.prompts import PromptTemplate
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage from langchain.schema import AIMessage, HumanMessage, SystemMessage
from langchain_core.tools import tool from langchain_core.tools import tool
from langchain_openai import ChatOpenAI from langchain_openai import ChatOpenAI
from langgraph.graph import END, StateGraph 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 import Graphiti
from graphiti_core.nodes import EpisodeType from graphiti_core.nodes import EpisodeType
@ -39,9 +38,13 @@ if not openai_api_key:
MAX_NEGOTIATION_ROUNDS = 5 MAX_NEGOTIATION_ROUNDS = 5
class State(TypedDict): def format_step_result(messages: List[str], **kwargs) -> Dict[str, Any]:
messages: List[HumanMessage] return {'messages': messages, **kwargs}
teams: Dict[str, Any]
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 event: str
team_actions: Dict[str, str] team_actions: Dict[str, str]
transfer_offers: List[Dict[str, Any]] transfer_offers: List[Dict[str, Any]]
@ -50,49 +53,93 @@ class State(TypedDict):
negotiation_complete: bool negotiation_complete: bool
def format_step_result(messages: List[str], **kwargs) -> Dict[str, Any]:
return {'messages': messages, **kwargs}
class TeamAgent: class TeamAgent:
def __init__( def __init__(self, name: str, tools: List[Any]):
self, team_name: str, tools: List[Any], budget: int = 100000000, roster: List[str] = [] self.name = name
): self.roster: List[str] = []
# logger.debug(f'Initializing TeamAgent for {team_name}') self.budget: int = 100_000_000
self.team_name = team_name
self.tools = tools self.tools = tools
self.budget = budget self.last_proposed_transfer: Optional[Dict[str, Any]] = None
self.roster = []
self.proposed_transfers = [] # Create the language model
self.prompt = ChatPromptTemplate.from_messages( llm = ChatOpenAI(temperature=0.3)
[
( # Create the agent executor
'system', template = """You are the manager of the {team_name} NBA team. Make decisions to improve your team.
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. Current event: {event}
Use the available tools to gather information and perform actions.
When an event occurs, decide how to react. You can: 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.
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. 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:".
3. Set transfer prices based on player performance and your available budget.
4. Negotiate with other teams on transfer prices. {agent_scratchpad}"""
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. prompt = PromptTemplate(
Only interact with these teams and do not mention or propose transfers to any other teams.""", input_variables=['team_name', 'event', 'agent_scratchpad'], template=template
),
('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) agent = create_openai_functions_agent(llm, self.tools, prompt)
self.executor = AgentExecutor(agent=self.agent, tools=self.tools, verbose=True) 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): async def update_budget(self, amount: int):
"""Update the team's budget.""" """Update the team's budget."""
self.budget += amount self.budget += amount
logger.debug(f"{self.team_name}'s new budget: ${self.budget:,}") logger.debug(f"{self.name}'s new budget: ${self.budget:,}")
return f"{self.team_name}'s new budget: ${self.budget:,}" return f"{self.name}'s new budget: ${self.budget:,}"
async def propose_transfers(self) -> List[Dict[str, Any]]: async def propose_transfers(self) -> List[Dict[str, Any]]:
"""Propose transfer offers based on the agent's strategy.""" """Propose transfer offers based on the agent's strategy."""
@ -104,24 +151,24 @@ class TeamAgent:
async def update_roster(self): async def update_roster(self):
"""Update the team's roster using the get_team_roster tool.""" """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_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(', ') self.roster = roster_string.split(': ')[1].split(', ')
def remove_player(self, player_name: str): def remove_player(self, player_name: str):
"""Remove a player from the team's roster.""" """Remove a player from the team's roster."""
if player_name in self.roster: if player_name in self.roster:
self.roster.remove(player_name) 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: 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): def add_player(self, player_name: str):
"""Add a player to the team's roster.""" """Add a player to the team's roster."""
if player_name not in self.roster: if player_name not in self.roster:
self.roster.append(player_name) 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: 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): async def get_transfer_offers(self):
"""Get the list of proposed transfers for this team.""" """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): async def submit_transfer_offer(self, player_name: str, to_team: str, proposed_price: int):
"""Submit a transfer offer for a player.""" """Submit a transfer offer for a player."""
offer = { offer = {
'from_team': self.team_name, 'from_team': self.name,
'to_team': to_team, 'to_team': to_team,
'player_name': player_name, 'player_name': player_name,
'proposed_price': proposed_price, 'proposed_price': proposed_price,
@ -139,37 +186,6 @@ class TeamAgent:
logger.debug(f'Transfer offer submitted: {offer}') logger.debug(f'Transfer offer submitted: {offer}')
return 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: async def react_to_others(self, other_actions: List[str]) -> str:
reaction_prompt = f"""Other teams have taken the following actions: reaction_prompt = f"""Other teams have taken the following actions:
{' '.join(other_actions)} {' '.join(other_actions)}
@ -185,16 +201,16 @@ class TeamAgent:
return result['output'] return result['output']
async def decide_on_transfer(self, offer: Dict[str, Any]) -> Dict[str, Any]: 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( 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': [], '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( return eval(
decision['output'] decision['output']
) # Convert the string representation of the dictionary to an actual dictionary ) # 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 f'Error executing tool {action}: {e}'
return None 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): async def add_episode(event_description: str):
"""Add a new episode to the Graphiti client.""" """Add a new episode to the Graphiti client."""
@ -264,7 +288,7 @@ async def search_player_info(player_name: str):
@tool @tool
async def propose_transfer(player_name: str, from_team: str, to_team: str, proposed_price: int): 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.""" """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 @tool
@ -339,354 +363,192 @@ tools = [
] ]
def root(state: dict) -> dict: def process_event(state: SimulationState) -> SimulationState:
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') logger.debug('Entering process_event')
new_message = f"Event processed: {state['event']}"
return { return {
**state, **state,
'messages': state['messages'] 'messages': state.get('messages', []) + [new_message],
+ [HumanMessage(content=f"Event processed: {state['event']}")],
} }
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('Entering parallel_agent_processing')
logger.debug(f'Current state: {state}')
tasks = [] tasks = []
for team_name, agent in state['teams'].items(): team_agents = {}
logger.debug(f"{team_name}: Starting to process event: {state['event']}") for team_name, team_data in state['teams'].items():
relevant_offers = [ team_agent = TeamAgent(team_data['name'], tools)
offer for offer in state.get('transfer_offers', []) if offer['to_team'] == team_name team_agents[team_name] = team_agent
] tasks.append(
logger.debug(f'{team_name}: Relevant offers: {relevant_offers}') asyncio.create_task(team_agent.process_event(state['event'], state['transfer_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 results = await asyncio.gather(*tasks, return_exceptions=True)
all_transfer_offers = state.get('transfer_offers', [])
for agent in state['teams'].values(): updated_state = state.copy()
if hasattr(agent, 'last_proposed_transfer'): updated_state['transfer_offers'] = []
all_transfer_offers.append(agent.last_proposed_transfer) for i, (team_name, team_data) in enumerate(state['teams'].items()):
logger.debug(f'Added transfer offer: {agent.last_proposed_transfer}') 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}') logger.debug(f'Updated state after parallel processing: {updated_state}')
return updated_state return updated_state
async def agent_reaction_phase(state: dict) -> dict: async def collect_transfer_offers(state: SimulationState) -> SimulationState:
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') logger.debug('Entering collect_transfer_offers')
transfer_offers = [] updated_state = state.copy()
for team_name, agent in state['teams'].items(): logger.debug(f'Collected transfer offers: {updated_state}')
offers = await agent.propose_transfers() # The transfer offers are already collected in parallel_agent_processing
transfer_offers.extend(offers) logger.debug(f"Collected transfer offers: {updated_state['transfer_offers']}")
return {**state, 'transfer_offers': state['transfer_offers'] + 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('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']: 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') logger.debug('No transfer offers to negotiate')
return { return updated_state
**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 # Sort offers by proposed price (highest first)
from_agent = state['teams'][offer['from_team']] sorted_offers = sorted(
to_agent = state['teams'][offer['to_team']] 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 # Simple negotiation logic: accept if the price is above a threshold
decision = await to_agent.decide_on_transfer(offer) threshold = 50 # This can be adjusted
logger.debug(f'Decision from {to_agent.team_name}: {decision}') if best_offer['proposed_price'] > threshold:
logger.info(
if decision['action'] == 'accept': f"Transfer accepted: {best_offer['player_name']} from {best_offer['from_team']} to {best_offer['to_team']} for ${best_offer['proposed_price']}"
# 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 # Update team rosters and budgets
from_agent.remove_player(offer['player_name']) from_team['roster'].remove(best_offer['player_name'])
to_agent.add_player(offer['player_name']) to_team['roster'].append(best_offer['player_name'])
await from_agent.update_budget(offer['proposed_price']) from_team['budget'] += best_offer['proposed_price']
await to_agent.update_budget(-offer['proposed_price']) to_team['budget'] -= best_offer['proposed_price']
updated_state['negotiation_complete'] = True
transfer_message = f"{offer['player_name']} transferred from {offer['from_team']} to {offer['to_team']} for ${offer['proposed_price']:,}" else:
logger.debug(transfer_message) logger.info(
f"Transfer rejected: {best_offer['player_name']} from {best_offer['from_team']} to {best_offer['to_team']} for ${best_offer['proposed_price']}"
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 updated_state['current_negotiation'] = None
updated_state['transfer_offers'] = []
return updated_state
async def handle_tool_use(state: Dict[str, Any]) -> Dict[str, Any]: async def execute_transfer(state: SimulationState, offer: Dict[str, Any]) -> None:
logger.debug('Entering handle_tool_use') from_agent = TeamAgent(
logger.debug(f'Current state in handle_tool_use: {state}') 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 = [] from_agent.roster.remove(offer['player_name'])
for team_name, agent in state['teams'].items(): to_agent.roster.append(offer['player_name'])
logger.debug(f'Checking for last_proposed_transfer on {team_name}') from_agent.budget += offer['proposed_price']
if hasattr(agent, 'last_proposed_transfer'): to_agent.budget -= offer['proposed_price']
transfer = agent.last_proposed_transfer
logger.debug(f'Processing transfer: {transfer}')
from_agent = state['teams'][transfer['from_team']] transfer_message = f"{offer['player_name']} transferred from {offer['from_team']} to {offer['to_team']} for ${offer['proposed_price']:,}"
to_agent = state['teams'][transfer['to_team']] state['messages'].append(transfer_message)
# 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): def should_continue(state: SimulationState) -> List[str]:
if x.get('negotiation_complete', False): if state['negotiation_complete'] and not state['transfer_offers']:
return END return [END]
if len(x.get('transfer_offers', [])) == 0 and x.get('current_negotiation') is None: return ['select_negotiation']
return END
return 'collect_transfer_offers'
# Define the workflow # Define the graph
workflow = StateGraph(State) workflow = StateGraph(SimulationState)
workflow.add_node('root', root) # Add nodes
workflow.add_node('add_episode', add_episode_node)
workflow.add_node('process_event', process_event) workflow.add_node('process_event', process_event)
workflow.add_node('parallel_agent_processing', parallel_agent_processing) 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('collect_transfer_offers', collect_transfer_offers)
workflow.add_node('select_negotiation', select_negotiation) workflow.add_node('select_negotiation', select_negotiation)
workflow.add_node('negotiate_transfer', negotiate_transfer) workflow.add_node('negotiate_transfer', negotiate_transfer)
workflow.add_node('handle_tool_use', handle_tool_use)
workflow.set_entry_point('root') # Add edges
workflow.add_edge('root', 'add_episode')
workflow.add_edge('add_episode', 'process_event')
workflow.add_edge('process_event', 'parallel_agent_processing') workflow.add_edge('process_event', 'parallel_agent_processing')
workflow.add_edge('parallel_agent_processing', 'agent_reaction_phase') workflow.add_edge('parallel_agent_processing', 'collect_transfer_offers')
workflow.add_edge('agent_reaction_phase', 'collect_transfer_offers')
workflow.add_edge('collect_transfer_offers', 'select_negotiation') workflow.add_edge('collect_transfer_offers', 'select_negotiation')
workflow.add_edge('select_negotiation', 'negotiate_transfer') workflow.add_edge('select_negotiation', 'negotiate_transfer')
workflow.add_edge('negotiate_transfer', 'handle_tool_use')
# Add conditional edge
workflow.add_conditional_edges( workflow.add_conditional_edges(
'handle_tool_use', 'negotiate_transfer', should_continue, {'select_negotiation': 'select_negotiation', END: END}
should_continue,
{END: END, 'collect_transfer_offers': 'collect_transfer_offers'},
) )
# Set the entrypoint
workflow.set_entry_point('process_event')
# Compile the graph
app = workflow.compile() 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(): 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: while True:
event = input("Enter an event (or 'quit' to exit): ") event = input("Enter an event (or 'quit' to exit): ")
if event.lower() == 'quit': if event.lower() == 'quit':
break break
state['event'] = event initial_state = SimulationState(
state['messages'] = [] messages=[],
state['team_actions'] = {} teams={
state['transfer_offers'] = [] 'Toronto Raptors': TeamAgent('Toronto Raptors', tools).to_dict(),
state['current_negotiation'] = None 'Boston Celtics': TeamAgent('Boston Celtics', tools).to_dict(),
state['negotiation_rounds'] = 0 'Golden State Warriors': TeamAgent('Golden State Warriors', tools).to_dict(),
state['negotiation_complete'] = False },
event=event,
team_actions={},
transfer_offers=[],
current_negotiation=None,
negotiation_rounds=0,
negotiation_complete=False,
)
# Process the event async for state in app.astream(initial_state):
state = await process_event(state) if 'messages' in state:
for message in state['messages']:
print(message)
# Parallel agent processing if 'transfer_offers' in state:
state = await parallel_agent_processing(state) print(f"Current transfer offers: {state['transfer_offers']}")
# Collect transfer offers if 'current_negotiation' in state:
state = await collect_transfer_offers(state) print(f"Current negotiation: {state['current_negotiation']}")
# 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:') print('\nFinal team states:')
for team_name, team_agent in state['teams'].items(): for team_name, team_data in initial_state['teams'].items():
print(f'{team_name} - Roster: {team_agent.roster}, Budget: ${team_agent.budget:,}') print(f"{team_name} - Roster: {team_data['roster']}, Budget: ${team_data['budget']:,}")
print('Simulation ended.') print('\n' + '=' * 50 + '\n')
def main():
asyncio.run(run_simulation())
if __name__ == '__main__': if __name__ == '__main__':
main() asyncio.run(run_simulation())