wip

examples/nba/ingest_current_roster.py

@@ -77,7 +77,12 @@ def fetch_current_roster():
     for t in all_teams:
         name = t['full_name']
         print(name)
-        if name == 'Golden State Warriors' or name == 'Boston Celtics' or name == 'Toronto Raptors':
+        if (
+            name == 'Golden State Warriors'
+            or name == 'Boston Celtics'
+            or name == 'Toronto Raptors'
+            or name == 'Los Angeles Lakers'
+        ):
             roster = commonteamroster.CommonTeamRoster(team_id=t['id']).get_dict()
             players_data = roster['resultSets'][0]
             headers = players_data['headers']

examples/nba/nba_agent.py

@@ -1,20 +1,27 @@
 import asyncio
+import json
 import logging
+import operator
 import os
-from typing import TypedDict, Dict, List, Optional, Any
+import re
+from datetime import datetime
+from typing import Annotated, Any, Dict, List, TypedDict
 
 from dotenv import load_dotenv
 from langchain.agents import AgentExecutor, create_openai_functions_agent
-from langchain.prompts import PromptTemplate
+from langchain.prompts import ChatPromptTemplate, PromptTemplate
-from langchain.schema import AIMessage, HumanMessage, SystemMessage
+from langchain.schema import HumanMessage
 from langchain_core.tools import tool
 from langchain_openai import ChatOpenAI
-from langgraph.graph import END, StateGraph
+from langgraph.graph import END, START, StateGraph
-from langgraph.prebuilt.tool_executor import ToolExecutor
+
-from datetime import datetime
 from graphiti_core import Graphiti
 from graphiti_core.nodes import EpisodeType
 
+logging.getLogger('langchain.callbacks.tracers.langchain').setLevel(logging.WARNING)
+logging.getLogger('urllib3.connectionpool').setLevel(logging.ERROR)
+DEFAULT_MODEL = 'gpt-4o-mini'
+
 load_dotenv()
 logging.basicConfig(
     level=logging.DEBUG, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
@@ -28,6 +35,7 @@ for name in logging.root.manager.loggerDict:
 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')
@@ -35,254 +43,143 @@ 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
-
-
-def format_step_result(messages: List[str], **kwargs) -> Dict[str, Any]:
-    return {'messages': messages, **kwargs}
-
 
+# Define the SimulationState
 class SimulationState(TypedDict):
-    messages: List[str]  # Changed from HumanMessage to str for simplicity
+    messages: Annotated[List[str], operator.add]
-    teams: Dict[str, Dict[str, Any]]  # Store team data as a dictionary
+    teams: Dict[str, Dict[str, int]]  # Changed to only store budget
     event: str
-    team_actions: Dict[str, str]
+    transfer_offers: Annotated[List[Dict], operator.add]
-    transfer_offers: List[Dict[str, Any]]
+    current_iteration: int
-    current_negotiation: Optional[Dict[str, Any]]
+    all_events: List[str]
-    negotiation_rounds: int
+    max_iterations: int
-    negotiation_complete: bool
 
 
-class TeamAgent:
+@tool
-    def __init__(self, name: str, tools: List[Any]):
+async def fetch_all_teams_context(teams: List[str]):
-        self.name = name
+    """Get the current roster and player summaries for specified teams."""
-        self.roster: List[str] = []
+    teams_with_players_dict = {}
-        self.budget: int = 100_000_000
+    llm = ChatOpenAI(temperature=0.2, model=DEFAULT_MODEL).bind(
-        self.tools = tools
+        response_format={'type': 'json_object'}
-        self.last_proposed_transfer: Optional[Dict[str, Any]] = None
+    )
 
-        # Create the language model
+    for team in teams:
-        llm = ChatOpenAI(temperature=0.3)
+        team_nodes = await graphiti_client.get_nodes_by_query(team, 1)
+        if not team_nodes:
+            logger.warning(f'No nodes found for team: {team}')
+            continue
 
-        # Create the agent executor
+        team_node = team_nodes[0]
-        template = """You are the manager of the {team_name} NBA team. Make decisions to improve your team.
+        search_result = await graphiti_client.search(
-
+            f'plays for {team_node.name}',
-Current event: {event}
+            center_node_uuid=team_node.uuid,
-
+            num_results=30,
-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
         )
 
-        agent = create_openai_functions_agent(llm, self.tools, prompt)
+        # Include all facts and timestamps (expired at if exists)
-        self.executor = AgentExecutor(agent=agent, tools=self.tools, verbose=True)
+        roster_facts = [get_fact_string(edge) for edge in search_result]
 
-    async def process_event(self, event: str, relevant_offers: List[Dict[str, Any]]) -> str:
+        prompt = PromptTemplate.from_template("""
-        logger.debug(f'{self.name}: Processing event: {event}')
+        Given the following list of facts about players and their teams, extract only the names and provide brief summaries for players who currently play for {team_name}. Follow these guidelines:
-        logger.debug(f'{self.name}: Current roster: {self.roster}')
+
-        logger.debug(f'{self.name}: Current budget: {self.budget}')
+        1. Only include players who are currently on the team.
+        2. Discard any information about players who are no longer on the team or were never on the team.
+        3. Use the 'expired_at' field to determine if a fact is still current. If 'expired_at' is not null, the fact is no longer current.
+        4. If there are conflicting facts, use the most recent one based on the 'valid_at' or 'created_at' timestamps.
+
+        Return the information as a JSON object containing a "players" field, which is an array of objects, each containing 'name' and 'summary' fields
+        example output:
+        {{
+            "players": [
+                {{
+                    "name": "Player Name",
+                    "summary": "Brief summary of the player"
+                }},
+                ...
+            ]
+        }}
+        Facts:
+        {facts}
+
+        Current players for {team_name}:
+        """)
+
+        llm_response = await llm.ainvoke(
+            prompt.format(
+                team_name=team_node.name,
+                facts='\n'.join(roster_facts),
+            )
+        )
 
         try:
-            result = await self.executor.ainvoke(
+            result = json.loads(llm_response.content)
-                {
+            players = result.get('players', [])
-                    'team_name': self.name,
+            if not isinstance(players, list):
-                    'event': event,
+                raise ValueError('Expected a JSON array')
-                    'roster': self.roster,
+        except json.JSONDecodeError:
-                    'budget': self.budget,
+            logger.error(f'Failed to parse JSON from LLM response for {team_node.name}')
-                    'relevant_offers': relevant_offers,
+            players = []
-                }
+        except ValueError as e:
-            )
+            logger.error(f'Invalid data structure in LLM response for {team_node.name}: {e}')
-            logger.debug(f"{self.name}: Agent output: {result['output']}")
+            players = []
 
-            # Check if a transfer was proposed
+        teams_with_players_dict[team_node.name] = players
-            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']
+    return teams_with_players_dict
-        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.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."""
-        # 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.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.name}'s roster")
-        else:
-            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.name}'s roster")
-        else:
-            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."""
-        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.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 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.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.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.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
-
-    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."""
-    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}'
 
 
+# Existing tools
 @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 = [
+    roster_facts = [get_fact_string(edge) for edge in search_result]
-        edge.fact.split(' plays for ')[0]
+
-        for edge in search_result
+    # Use LLM to extract player names
-        if 'plays for' in edge.fact.lower()
+    llm = ChatOpenAI(temperature=0, model=DEFAULT_MODEL)
-    ]
+    prompt = PromptTemplate.from_template("""
-    return f"{team_name}'s roster: {', '.join(roster)}"
+    Given the following list of facts about players and their teams, extract only the names of players who play for {team_name}. Return the names as a comma-separated list.
+
+    Facts:
+    {facts}
+
+    Players who play for {team_name}:
+    """)
+
+    llm_response = await llm.ainvoke(
+        prompt.format(team_name=team_name, facts='\n'.join(roster_facts))
+    )
+
+    player_names = [name.strip() for name in llm_response.content.split(',')]
+
+    return f"{team_name}'s roster: {', '.join(player_names)}"
 
 
 @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 = {
+    all_facts = [get_fact_string(edge) for edge in search_result]
-        'name': player_name,
+
-        'facts': [get_fact_string(edge) for edge in search_result],
+    # Use LLM to extract relevant player information
-    }
+    llm = ChatOpenAI(temperature=0, model=DEFAULT_MODEL)
-    return player_info
+    prompt = PromptTemplate.from_template("""
+    Given the following list of facts, extract only the information that is relevant to {player_name}. 
+    Return the relevant facts as a list, with each fact on a new line.
+
+    Facts:
+    {facts}
+
+    Relevant facts about {player_name}:
+    """)
+
+    llm_response = await llm.ainvoke(
+        prompt.format(player_name=player_name, facts='\n'.join(all_facts))
+    )
+
+    relevant_facts = llm_response.content.strip().split('\n')
+
+    return {'name': player_name, 'facts': relevant_facts}
 
 
 @tool
@@ -291,23 +188,18 @@ async def propose_transfer(player_name: str, from_team: str, to_team: str, propo
     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.
+    await graphiti_client.add_episode(
+        name=f'Transfer {player_name}',
+        episode_body=f'{player_name} transferred from {from_team} to {to_team} for ${price:,}',
+        source_description='Player Transfer',
+        reference_time=datetime.now(),
+        source=EpisodeType.message,
+    )
     return {
         'messages': [
             HumanMessage(
@@ -317,237 +209,290 @@ async def execute_transfer(
     }
 
 
-@tool
+async def add_episode(event_description: str):
-async def check_team_budget(team_name: str) -> Dict[str, Any]:
+    """Add a new episode to the Graphiti client."""
-    """Check the current budget of a team."""
+    await graphiti_client.add_episode(
-    # This is a placeholder. In a real scenario, you'd fetch the actual budget.
+        name='New Event',
-    return {
+        episode_body=event_description,
-        'messages': [HumanMessage(content=f"Checking {team_name}'s budget...")],
+        source_description='User Input',
-    }
+        reference_time=datetime.now(),
-
+        source=EpisodeType.message,
-
-@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 = {
+    return f"Episode '{event_description}' added successfully."
-        'from_team': from_team,
+
-        'to_team': to_team,
+
-        'player_name': player_name,
+def get_fact_string(edge):
-        'proposed_price': proposed_price,
+    fact_string = f'{edge.fact} Valid At: {edge.valid_at or edge.created_at}'
-    }
+    if edge.expired_at:
-    logger.debug(f'Transfer offer created: {offer}')
+        fact_string += f' Expired At: {edge.expired_at}'
-    return {
+    return fact_string
-        '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 process_event(state: SimulationState) -> SimulationState:
+# Define the team agent function
-    logger.debug('Entering process_event')
+def create_team_agent(team_name: str, valid_teams: List[str]):
-    new_message = f"Event processed: {state['event']}"
+    llm = ChatOpenAI(temperature=0.3, model=DEFAULT_MODEL).bind(
+        response_format={'type': 'json_object'}
+    )
+    prompt = ChatPromptTemplate.from_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, but focus on making a decision quickly. If you think a player transfer would benefit your team, propose one following the guidelines below.
+Ensure that you use the current budget info and the current state of your team to make the best decision.
+Current budget: ${budget}
+
+Valid teams for transfers: {valid_teams}
+
+IMPORTANT: After gathering information, you MUST make a decision. Your options are:
+1. Propose a transfer
+    Note: if you are proposing a transfer make sure to output JSON in the following format:
+    {{
+        "transfer_proposal": {{
+            "to_team": "team_name",
+            "from_team": "team_name",
+            "player_name": "player_name",
+            "proposed_price": price
+        }}
+    }}
+    IMPORTANT: Only propose transfers to teams in the valid teams list. Make sure that the player_name is a valid player on the from_team. Ensure that the the from_team name is a valid team name.
+2. Do nothing (output an empty JSON object)
+
+Do not ask for more information or clarification. Make a decision based on what you know.
+
+{agent_scratchpad}""")
+
+    async def team_agent_function(state: SimulationState) -> Dict:
+        agent = create_openai_functions_agent(llm, tools, prompt)
+        executor = AgentExecutor(
+            agent=agent, tools=[get_team_roster, search_player_info], verbose=True
+        )
+        team_data = state['teams'][team_name]
+
+        result = await executor.ainvoke(
+            {
+                'team_name': team_name,
+                'event': state['event'],
+                'budget': team_data['budget'],
+                'valid_teams': ', '.join(valid_teams),
+            }
+        )
+
+        json_result = json.loads(result['output'])
+        transfer_offer = None
+        if 'transfer_proposal' in json_result:
+            transfer_offer = json_result['transfer_proposal']
+            if (
+                transfer_offer['to_team'] not in valid_teams
+                or transfer_offer['from_team'] not in valid_teams
+            ):
+                logger.warning(f'Invalid transfer proposal: {transfer_offer}. Ignoring.')
+                transfer_offer = None
+
+        return {
+            'transfer_offers': [transfer_offer] if transfer_offer else [],
+        }
+
+    return team_agent_function
+
+
+def parse_transfer_proposal(proposal: str) -> Dict[str, Any]:
+    # Use regex to extract information
+    to_team_match = re.search(r'(.*?) wants to buy', proposal)
+    player_match = re.search(r'buy (.*?) from', proposal)
+    from_team_match = re.search(r'from (.*?) for', proposal)
+    price_match = re.search(r'\$([0-9,]+)', proposal)
+
+    if not all([to_team_match, player_match, from_team_match, price_match]):
+        raise ValueError(f'Unable to parse transfer proposal: {proposal}')
+
+    to_team = to_team_match.group(1)
+    player_name = player_match.group(1)
+    from_team = from_team_match.group(1)
+    proposed_price = int(price_match.group(1).replace(',', ''))
+
     return {
-        **state,
+        'to_team': to_team,
-        'messages': state.get('messages', []) + [new_message],
+        'from_team': from_team,
+        'player_name': player_name,
+        'proposed_price': proposed_price,
     }
 
 
-async def parallel_agent_processing(state: SimulationState) -> SimulationState:
+async def process_event(state: SimulationState) -> SimulationState:
-    logger.debug('Entering parallel_agent_processing')
+    # await add_episode(state['event'])
-    tasks = []
+    return {
-    team_agents = {}
+        **state,
-    for team_name, team_data in state['teams'].items():
+        'messages': [f"Event processed: {state['event']}"],
-        team_agent = TeamAgent(team_data['name'], tools)
+        'transfer_offers': [],
-        team_agents[team_name] = team_agent
+    }
-        tasks.append(
-            asyncio.create_task(team_agent.process_event(state['event'], state['transfer_offers']))
-        )
-
-    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()
-
-    logger.debug(f'Updated state after parallel processing: {updated_state}')
-    return updated_state
 
 
-async def collect_transfer_offers(state: SimulationState) -> SimulationState:
+async def process_transfers(state: SimulationState) -> SimulationState:
-    logger.debug('Entering collect_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
-
-
-def select_negotiation(state: SimulationState) -> SimulationState:
-    logger.debug('Entering select_negotiation')
     if not state['transfer_offers']:
-        return {**state, 'current_negotiation': None, 'negotiation_complete': True}
+        return state
-    return {**state, 'current_negotiation': state['transfer_offers'][0]}
 
+    # Group offers by player
+    offers_by_player = {}
+    for offer in state['transfer_offers']:
+        player = offer['player_name']
+        if player not in offers_by_player:
+            offers_by_player[player] = []
+        offers_by_player[player].append(offer)
 
-async def negotiate_transfer(state: SimulationState) -> SimulationState:
+    for player, offers in offers_by_player.items():
-    logger.debug('Entering negotiate_transfer')
+        # Sort offers by price, highest first
-    updated_state = state.copy()
+        best_offer = max(offers, key=lambda x: x['proposed_price'])
 
-    if not updated_state['transfer_offers']:
-        logger.debug('No transfer offers to negotiate')
-        return updated_state
-
-    # 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]
-
-    # Simulate negotiation
-    from_team = updated_state['teams'][best_offer['from_team']]
-    to_team = updated_state['teams'][best_offer['to_team']]
-
-    # 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']}"
+            f"Best offer for {player}: {best_offer['to_team']} wants to buy from {best_offer['from_team']} for ${best_offer['proposed_price']:,}"
         )
+
+        # Execute the transfer
+        transfer_result = await execute_transfer.ainvoke(
+            {
+                'player_name': best_offer['player_name'],
+                'from_team': best_offer['from_team'],
+                'to_team': best_offer['to_team'],
+                'price': best_offer['proposed_price'],
+            }
+        )
+        # Add the transfer result message to the state
+        state['messages'].extend(transfer_result['messages'])
+
         # Update team rosters and budgets
-        from_team['roster'].remove(best_offer['player_name'])
+        from_team = best_offer['from_team']
-        to_team['roster'].append(best_offer['player_name'])
+        to_team = best_offer['to_team']
-        from_team['budget'] += best_offer['proposed_price']
+        price = 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']}"
-        )
 
-    updated_state['current_negotiation'] = None
+        if from_team in state['teams'] and to_team in state['teams']:
-    updated_state['transfer_offers'] = []
+            state['teams'][from_team]['budget'] += price
-    return updated_state
+            state['teams'][to_team]['budget'] -= price
+        else:
+            logger.warning(f'Cannot process transfer: {from_team} or {to_team} not in simulation.')
+
+    # Clear all processed offers
+    state['transfer_offers'] = []
+
+    return state
 
 
-async def execute_transfer(state: SimulationState, offer: Dict[str, Any]) -> None:
+def create_simulator_agent():
-    from_agent = TeamAgent(
+    llm = ChatOpenAI(
-        state['teams'][offer['from_team']]['name'], tools
+        temperature=0.7, model=DEFAULT_MODEL
-    )  # Recreate TeamAgent from data
+    )  # Higher temperature for more creative events
-    to_agent = TeamAgent(
+    prompt = ChatPromptTemplate.from_template("""
-        state['teams'][offer['to_team']]['name'], tools
+    You are an NBA event simulator. Your role is to generate realistic events based on the current state of NBA teams and players. Use the provided team and player information to create engaging and plausible scenarios.
-    )  # Recreate TeamAgent from data
 
-    from_agent.roster.remove(offer['player_name'])
+    Current NBA landscape:
-    to_agent.roster.append(offer['player_name'])
+    {teams_context}
-    from_agent.budget += offer['proposed_price']
-    to_agent.budget -= offer['proposed_price']
 
-    transfer_message = f"{offer['player_name']} transferred from {offer['from_team']} to {offer['to_team']} for ${offer['proposed_price']:,}"
+    Generate a single, specific event involving one or more teams or players. The event should be impactful enough to potentially influence team decisions. Examples include outstanding performances, injuries, trade rumors, or off-court incidents.
-    state['messages'].append(transfer_message)
+
+    Output the event as a brief, news-like statement.
+
+    Event:
+    """)
+
+    return prompt, llm
 
 
-def should_continue(state: SimulationState) -> List[str]:
+simulator_prompt, simulator_llm = create_simulator_agent()
-    if state['negotiation_complete'] and not state['transfer_offers']:
-        return [END]
-    return ['select_negotiation']
 
 
-# Define the graph
+async def simulate_event(state: SimulationState) -> SimulationState:
+    teams = ['Toronto Raptors', 'Boston Celtics', 'Golden State Warriors']
+    teams_context = await fetch_all_teams_context.ainvoke({'teams': teams})
+
+    result = await simulator_llm.ainvoke(
+        simulator_prompt.format_prompt(teams_context=json.dumps(teams_context, indent=2))
+    )
+
+    new_event = result.content
+    existing_events = state['all_events'] or []
+    existing_events.append(new_event)
+    return {
+        **state,
+        'event': new_event,
+        'all_events': existing_events,
+        'transfer_offers': [],
+        'current_iteration': state['current_iteration'] + 1,
+    }
+
+
+# Create the graph
 workflow = StateGraph(SimulationState)
 
 # Add nodes
+workflow.add_node('simulate_event', simulate_event)
 workflow.add_node('process_event', process_event)
-workflow.add_node('parallel_agent_processing', parallel_agent_processing)
+valid_teams = ['Toronto Raptors', 'Boston Celtics', 'Golden State Warriors']
-workflow.add_node('collect_transfer_offers', collect_transfer_offers)
+for team in valid_teams:
-workflow.add_node('select_negotiation', select_negotiation)
+    workflow.add_node(f'agent_{team}', create_team_agent(team, valid_teams))
-workflow.add_node('negotiate_transfer', negotiate_transfer)
+workflow.add_node('process_transfers', process_transfers)
 
 # Add edges
-workflow.add_edge('process_event', 'parallel_agent_processing')
+workflow.add_edge(START, 'simulate_event')
-workflow.add_edge('parallel_agent_processing', 'collect_transfer_offers')
+workflow.add_edge('simulate_event', 'process_event')
-workflow.add_edge('collect_transfer_offers', 'select_negotiation')
+
-workflow.add_edge('select_negotiation', 'negotiate_transfer')
+# Add edges from process_event to all agent nodes
+for team in valid_teams:
+    workflow.add_edge('process_event', f'agent_{team}')
+
+for team in valid_teams:
+    workflow.add_edge(f'agent_{team}', 'process_transfers')
+
+
+def routing_function(state: SimulationState) -> str:
+    if state['current_iteration'] >= state['max_iterations']:
+        return END
+    else:
+        return 'simulate_event'
+
 
-# Add conditional edge
 workflow.add_conditional_edges(
-    'negotiate_transfer', should_continue, {'select_negotiation': 'select_negotiation', END: END}
+    'process_transfers',
+    routing_function,
 )
 
-# Set the entrypoint
-workflow.set_entry_point('process_event')
-
 # Compile the graph
 app = workflow.compile()
 
+print(app.get_graph().draw_mermaid())
+
 
 async def run_simulation():
-    while True:
+    num_iterations = int(input('Enter the number of simulation iterations: '))
-        event = input("Enter an event (or 'quit' to exit): ")
-        if event.lower() == 'quit':
-            break
 
-        initial_state = SimulationState(
+    initial_state = SimulationState(
-            messages=[],
+        messages=[],
-            teams={
+        teams={
-                'Toronto Raptors': TeamAgent('Toronto Raptors', tools).to_dict(),
+            'Toronto Raptors': {'budget': 100000000},
-                'Boston Celtics': TeamAgent('Boston Celtics', tools).to_dict(),
+            'Boston Celtics': {'budget': 100000000},
-                'Golden State Warriors': TeamAgent('Golden State Warriors', tools).to_dict(),
+            'Golden State Warriors': {'budget': 100000000},
-            },
+        },
-            event=event,
+        event='',
-            team_actions={},
+        transfer_offers=[],
-            transfer_offers=[],
+        current_iteration=0,
-            current_negotiation=None,
+        max_iterations=num_iterations,
-            negotiation_rounds=0,
+    )
-            negotiation_complete=False,
-        )
 
-        async for state in app.astream(initial_state):
+    final_state = await app.ainvoke(initial_state, {'recursion_limit': 200})
-            if 'messages' in state:
-                for message in state['messages']:
-                    print(message)
 
-            if 'transfer_offers' in state:
+    print('\nFinal team states:')
-                print(f"Current transfer offers: {state['transfer_offers']}")
+    for team_name, team_data in final_state['teams'].items():
-
+        print(f"{team_name} - Budget: ${team_data['budget']:,}")
-            if 'current_negotiation' in state:
+    print(f'Steps taken: {final_state["current_iteration"]}')
-                print(f"Current negotiation: {state['current_negotiation']}")
+    for event in final_state['all_events']:
-
+        print('/n')
-        print('\nFinal team states:')
+        print(event)
-        for team_name, team_data in initial_state['teams'].items():
+        print('\n')
-            print(f"{team_name} - Roster: {team_data['roster']}, Budget: ${team_data['budget']:,}")
-
-        print('\n' + '=' * 50 + '\n')
 
 
 if __name__ == '__main__':