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 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())