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

chore: Add Analyzer step

Browse source
This commit is contained in:
paulpaliychuk 2024-08-30 15:42:48 -04:00
parent fba24f6b8e
commit f7642f6e6e
3 changed files with 157 additions and 325 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

367
examples/nba/current_nba_roster.json
View file

@ -2,436 +2,205 @@
{
"team_name": "Boston Celtics",
"player_id": 1628369,
"player_name": "Jayson Tatum"
"player_name": "Jayson Tatum",
"last_transfer_price": 2250000
},
{
"team_name": "Boston Celtics",
"player_id": 201950,
"player_name": "Jrue Holiday"
"player_name": "Jrue Holiday",
"last_transfer_price": 3000000
},
{
"team_name": "Boston Celtics",
"player_id": 1627759,
"player_name": "Jaylen Brown"
"player_name": "Jaylen Brown",
"last_transfer_price": 2250000
},
{
"team_name": "Boston Celtics",
"player_id": 204001,
"player_name": "Kristaps Porzingis"
"player_name": "Kristaps Porzingis",
"last_transfer_price": 2000000
},
{
"team_name": "Boston Celtics",
"player_id": 1628401,
"player_name": "Derrick White"
"player_name": "Derrick White",
"last_transfer_price": 2750000
},
{
"team_name": "Boston Celtics",
"player_id": 1630202,
"player_name": "Payton Pritchard"
"player_name": "Payton Pritchard",
"last_transfer_price": 2500000
},
{
"team_name": "Boston Celtics",
"player_id": 1629052,
"player_name": "Oshae Brissett"
"player_name": "Oshae Brissett",
"last_transfer_price": 2000000
},
{
"team_name": "Boston Celtics",
"player_id": 1641809,
"player_name": "Drew Peterson"
"player_name": "Drew Peterson",
"last_transfer_price": 2500000
},
{
"team_name": "Boston Celtics",
"player_id": 1631120,
"player_name": "JD Davison"
"player_name": "JD Davison",
"last_transfer_price": 3000000
},
{
"team_name": "Boston Celtics",
"player_id": 1630214,
"player_name": "Xavier Tillman"
"player_name": "Xavier Tillman",
"last_transfer_price": 2250000
},
{
"team_name": "Boston Celtics",
"player_id": 1641775,
"player_name": "Jordan Walsh"
"player_name": "Jordan Walsh",
"last_transfer_price": 2000000
},
{
"team_name": "Boston Celtics",
"player_id": 1630573,
"player_name": "Sam Hauser"
"player_name": "Sam Hauser",
"last_transfer_price": 2000000
},
{
"team_name": "Boston Celtics",
"player_id": 1628436,
"player_name": "Luke Kornet"
"player_name": "Luke Kornet",
"last_transfer_price": 2250000
},
{
"team_name": "Boston Celtics",
"player_id": 201143,
"player_name": "Al Horford"
"player_name": "Al Horford",
"last_transfer_price": 2500000
},
{
"team_name": "Boston Celtics",
"player_id": 1630531,
"player_name": "Jaden Springer"
"player_name": "Jaden Springer",
"last_transfer_price": 2250000
},
{
"team_name": "Boston Celtics",
"player_id": 1629004,
"player_name": "Svi Mykhailiuk"
"player_name": "Svi Mykhailiuk",
"last_transfer_price": 2250000
},
{
"team_name": "Boston Celtics",
"player_id": 1629674,
"player_name": "Neemias Queta"
},
{
"team_name": "Golden State Warriors",
"player_id": 1627780,
"player_name": "Gary Payton II"
},
{
"team_name": "Golden State Warriors",
"player_id": 1630228,
"player_name": "Jonathan Kuminga"
},
{
"team_name": "Golden State Warriors",
"player_id": 1641764,
"player_name": "Brandin Podziemski"
},
{
"team_name": "Golden State Warriors",
"player_id": 101108,
"player_name": "Chris Paul"
},
{
"team_name": "Golden State Warriors",
"player_id": 1630541,
"player_name": "Moses Moody"
},
{
"team_name": "Golden State Warriors",
"player_id": 1626172,
"player_name": "Kevon Looney"
},
{
"team_name": "Golden State Warriors",
"player_id": 202691,
"player_name": "Klay Thompson"
},
{
"team_name": "Golden State Warriors",
"player_id": 1630586,
"player_name": "Usman Garuba"
},
{
"team_name": "Golden State Warriors",
"player_id": 1630611,
"player_name": "Gui Santos"
},
{
"team_name": "Golden State Warriors",
"player_id": 1629010,
"player_name": "Jerome Robinson"
},
{
"team_name": "Golden State Warriors",
"player_id": 203967,
"player_name": "Dario Saric"
},
{
"team_name": "Golden State Warriors",
"player_id": 203952,
"player_name": "Andrew Wiggins"
},
{
"team_name": "Golden State Warriors",
"player_id": 203110,
"player_name": "Draymond Green"
},
{
"team_name": "Golden State Warriors",
"player_id": 1631311,
"player_name": "Lester Quinones"
},
{
"team_name": "Golden State Warriors",
"player_id": 201939,
"player_name": "Stephen Curry"
},
{
"team_name": "Golden State Warriors",
"player_id": 1631218,
"player_name": "Trayce Jackson-Davis"
},
{
"team_name": "Golden State Warriors",
"player_id": 1630311,
"player_name": "Pat Spencer"
},
{
"team_name": "Los Angeles Lakers",
"player_id": 1641720,
"player_name": "Jalen Hood-Schifino"
},
{
"team_name": "Los Angeles Lakers",
"player_id": 1626156,
"player_name": "D'Angelo Russell"
},
{
"team_name": "Los Angeles Lakers",
"player_id": 1629020,
"player_name": "Jarred Vanderbilt"
},
{
"team_name": "Los Angeles Lakers",
"player_id": 203076,
"player_name": "Anthony Davis"
},
{
"team_name": "Los Angeles Lakers",
"player_id": 1630219,
"player_name": "Skylar Mays"
},
{
"team_name": "Los Angeles Lakers",
"player_id": 1629629,
"player_name": "Cam Reddish"
},
{
"team_name": "Los Angeles Lakers",
"player_id": 1629216,
"player_name": "Gabe Vincent"
},
{
"team_name": "Los Angeles Lakers",
"player_id": 1631108,
"player_name": "Max Christie"
},
{
"team_name": "Los Angeles Lakers",
"player_id": 1629637,
"player_name": "Jaxson Hayes"
},
{
"team_name": "Los Angeles Lakers",
"player_id": 1627752,
"player_name": "Taurean Prince"
},
{
"team_name": "Los Angeles Lakers",
"player_id": 1630658,
"player_name": "Colin Castleton"
},
{
"team_name": "Los Angeles Lakers",
"player_id": 1630559,
"player_name": "Austin Reaves"
},
{
"team_name": "Los Angeles Lakers",
"player_id": 1628385,
"player_name": "Harry Giles III"
},
{
"team_name": "Los Angeles Lakers",
"player_id": 1641721,
"player_name": "Maxwell Lewis"
},
{
"team_name": "Los Angeles Lakers",
"player_id": 2544,
"player_name": "LeBron James"
},
{
"team_name": "Los Angeles Lakers",
"player_id": 203915,
"player_name": "Spencer Dinwiddie"
},
{
"team_name": "Los Angeles Lakers",
"player_id": 1629060,
"player_name": "Rui Hachimura"
},
{
"team_name": "Los Angeles Lakers",
"player_id": 1626174,
"player_name": "Christian Wood"
},
{
"team_name": "Miami Heat",
"player_id": 1626196,
"player_name": "Josh Richardson"
},
{
"team_name": "Miami Heat",
"player_id": 1626179,
"player_name": "Terry Rozier"
},
{
"team_name": "Miami Heat",
"player_id": 1626153,
"player_name": "Delon Wright"
},
{
"team_name": "Miami Heat",
"player_id": 1631107,
"player_name": "Nikola Jovic"
},
{
"team_name": "Miami Heat",
"player_id": 1631288,
"player_name": "Jamal Cain"
},
{
"team_name": "Miami Heat",
"player_id": 201988,
"player_name": "Patty Mills"
},
{
"team_name": "Miami Heat",
"player_id": 1631170,
"player_name": "Jaime Jaquez Jr."
},
{
"team_name": "Miami Heat",
"player_id": 1628389,
"player_name": "Bam Adebayo"
},
{
"team_name": "Miami Heat",
"player_id": 1629639,
"player_name": "Tyler Herro"
},
{
"team_name": "Miami Heat",
"player_id": 1631214,
"player_name": "Alondes Williams"
},
{
"team_name": "Miami Heat",
"player_id": 1628997,
"player_name": "Caleb Martin"
},
{
"team_name": "Miami Heat",
"player_id": 1631306,
"player_name": "Cole Swider"
},
{
"team_name": "Miami Heat",
"player_id": 202710,
"player_name": "Jimmy Butler"
},
{
"team_name": "Miami Heat",
"player_id": 1629312,
"player_name": "Haywood Highsmith"
},
{
"team_name": "Miami Heat",
"player_id": 1631115,
"player_name": "Orlando Robinson"
},
{
"team_name": "Miami Heat",
"player_id": 1628418,
"player_name": "Thomas Bryant"
},
{
"team_name": "Miami Heat",
"player_id": 201567,
"player_name": "Kevin Love"
},
{
"team_name": "Miami Heat",
"player_id": 1629130,
"player_name": "Duncan Robinson"
"player_name": "Neemias Queta",
"last_transfer_price": 2500000
},
{
"team_name": "Toronto Raptors",
"player_id": 1642013,
"player_name": "Malik Williams"
"player_name": "Malik Williams",
"last_transfer_price": 2500000
},
{
"team_name": "Toronto Raptors",
"player_id": 1631241,
"player_name": "Javon Freeman-Liberty"
"player_name": "Javon Freeman-Liberty",
"last_transfer_price": 2500000
},
{
"team_name": "Toronto Raptors",
"player_id": 1641711,
"player_name": "Gradey Dick"
"player_name": "Gradey Dick",
"last_transfer_price": 2750000
},
{
"team_name": "Toronto Raptors",
"player_id": 1629667,
"player_name": "Jalen McDaniels"
"player_name": "Jalen McDaniels",
"last_transfer_price": 2000000
},
{
"team_name": "Toronto Raptors",
"player_id": 1630618,
"player_name": "DJ Carton"
"player_name": "DJ Carton",
"last_transfer_price": 2250000
},
{
"team_name": "Toronto Raptors",
"player_id": 1630567,
"player_name": "Scottie Barnes"
"player_name": "Scottie Barnes",
"last_transfer_price": 3000000
},
{
"team_name": "Toronto Raptors",
"player_id": 1630193,
"player_name": "Immanuel Quickley"
"player_name": "Immanuel Quickley",
"last_transfer_price": 2750000
},
{
"team_name": "Toronto Raptors",
"player_id": 1629628,
"player_name": "RJ Barrett"
"player_name": "RJ Barrett",
"last_transfer_price": 2000000
},
{
"team_name": "Toronto Raptors",
"player_id": 1628971,
"player_name": "Bruce Brown"
"player_name": "Bruce Brown",
"last_transfer_price": 2000000
},
{
"team_name": "Toronto Raptors",
"player_id": 1629670,
"player_name": "Jordan Nwora"
"player_name": "Jordan Nwora",
"last_transfer_price": 2000000
},
{
"team_name": "Toronto Raptors",
"player_id": 1631338,
"player_name": "Mouhamadou Gueye"
"player_name": "Mouhamadou Gueye",
"last_transfer_price": 2750000
},
{
"team_name": "Toronto Raptors",
"player_id": 202066,
"player_name": "Garrett Temple"
"player_name": "Garrett Temple",
"last_transfer_price": 2250000
},
{
"team_name": "Toronto Raptors",
"player_id": 1627751,
"player_name": "Jakob Poeltl"
"player_name": "Jakob Poeltl",
"last_transfer_price": 2750000
},
{
"team_name": "Toronto Raptors",
"player_id": 1628449,
"player_name": "Chris Boucher"
"player_name": "Chris Boucher",
"last_transfer_price": 2000000
},
{
"team_name": "Toronto Raptors",
"player_id": 1630534,
"player_name": "Ochai Agbaji"
"player_name": "Ochai Agbaji",
"last_transfer_price": 2500000
},
{
"team_name": "Toronto Raptors",
"player_id": 1629018,
"player_name": "Gary Trent Jr."
"player_name": "Gary Trent Jr.",
"last_transfer_price": 2000000
},
{
"team_name": "Toronto Raptors",
"player_id": 203482,
"player_name": "Kelly Olynyk"
"player_name": "Kelly Olynyk",
"last_transfer_price": 2000000
}
]

14
examples/nba/ingest_current_roster.py
View file

@ -22,7 +22,7 @@ import sys
from datetime import datetime
from pathlib import Path
from typing import TypedDict
import random
from dotenv import load_dotenv
from nba_api.stats.endpoints import commonteamroster, teamdetails
from nba_api.stats.static import players, teams
@ -77,13 +77,7 @@ 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'
or name == 'Los Angeles Lakers'
or name == 'Miami Heat'
):
if name == 'Boston Celtics' or name == 'Toronto Raptors':
roster = commonteamroster.CommonTeamRoster(team_id=t['id']).get_dict()
players_data = roster['resultSets'][0]
headers = players_data['headers']
@ -93,10 +87,14 @@ def fetch_current_roster():
player_dict = dict(zip(headers, row))
player_dict['team_name'] = name
print(player_dict)
random_number_from_list = random.choice(
[2_000_000, 2_250_000, 2_500_000, 2_750_000, 3_000_000]
)
meaningful_data = {
'team_name': name,
'player_id': player_dict['PLAYER_ID'],
'player_name': player_dict['PLAYER'],
'last_transfer_price': random_number_from_list,
# 'player_number': player_dict['NUM'],
# 'player_position': player_dict['POSITION'],
# 'player_school': player_dict['SCHOOL'],

101
examples/nba/nba_agent.py
View file

@ -24,9 +24,9 @@ DEFAULT_MODEL = 'gpt-4o-mini'
VALID_TEAMS = [
'Toronto Raptors',
'Boston Celtics',
'Golden State Warriors',
'Miami Heat',
'Los Angeles Lakers',
# 'Golden State Warriors',
# 'Miami Heat',
# 'Los Angeles Lakers',
]
load_dotenv()
logging.basicConfig(
@ -59,6 +59,7 @@ class SimulationState(TypedDict):
current_iteration: int
all_events: List[str]
max_iterations: int
teams_context: Dict[str, List[dict]]
@tool
@ -99,7 +100,7 @@ async def fetch_all_teams_context(teams: List[str]):
"players": [
{{
"name": "Player Name",
"summary": "Brief summary of the player"
"summary": "Brief summary of the player. Make sure you include all the information you can get from the roster facts about the player. Do not include any information that is not in the roster facts. If dates are mentioned, make sure to integrate them well in the summary. If applicable it can tell the course of events"
}},
...
]
@ -190,7 +191,7 @@ async def search_player_info(player_name: str):
@tool
async def execute_transfer(
player_name: str, from_team: str, to_team: str, price: int
player_name: str, from_team: str, to_team: str, price: int, reason: str
) -> Dict[str, Any]:
"""Execute a transfer between two teams."""
await graphiti_client.add_episode(
@ -203,7 +204,7 @@ async def execute_transfer(
return {
'messages': [
HumanMessage(
content=f'Transfer executed: {player_name} moved from {from_team} to {to_team} for ${price:,}'
content=f'Transfer executed: {player_name} moved from {from_team} to {to_team} for ${price:,} Reason: {reason}'
)
],
}
@ -211,6 +212,7 @@ async def execute_transfer(
async def add_episode(event_description: str):
"""Add a new episode to the Graphiti client."""
logger.info(f'Adding episode: {event_description}')
await graphiti_client.add_episode(
name='New Event',
episode_body=event_description,
@ -244,12 +246,14 @@ def create_team_agent(team_name: str):
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.
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. Make sure to get familiar with the entire transfer history of a given player
Ensure that you use the current budget info and the current state of your team (use an appropriate tool to get the current state of your team) to make the best decision.
Current budget: ${budget}
Valid teams for transfers: {valid_teams}
Do not propose transfers you cannot afford.
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:
@ -258,11 +262,15 @@ IMPORTANT: After gathering information, you MUST make a decision. Your options a
"to_team": "team_name",
"from_team": "team_name",
"player_name": "player_name",
"proposed_price": price
"proposed_price": price,
"reason": "reason for the proposed transfer"
}}
}}
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)
2. Do nothing (output the following JSON object)
{{
"no_transfer": "reason for not proposing a transfer"
}}
Do not ask for more information or clarification. Make a decision based on what you know.
@ -285,25 +293,33 @@ Do not ask for more information or clarification. Make a decision based on what
)
json_result = json.loads(result['output'])
messages = []
transfer_offer = None
if 'transfer_proposal' in json_result:
transfer_offer = json_result['transfer_proposal']
logger.info(f'Transfer proposal made by {team_name}: {transfer_offer["reason"]}')
messages.append(f'Transfer proposal made by {team_name}: {transfer_offer["reason"]}')
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
if 'no_transfer' in json_result:
logger.info(f'No transfer proposal made by {team_name}: {json_result["no_transfer"]}')
messages.append(
f'No transfer proposal made by {team_name}: {json_result["no_transfer"]}'
)
return {
'transfer_offers': [transfer_offer] if transfer_offer else [],
'messages': messages,
}
return team_agent_function
async def process_event(state: SimulationState) -> SimulationState:
# await add_episode(state['event'])
await add_episode(state['event'])
return {
**state,
'messages': [f"Event processed: {state['event']}"],
@ -338,10 +354,11 @@ async def process_transfers(state: SimulationState) -> SimulationState:
'from_team': best_offer['from_team'],
'to_team': best_offer['to_team'],
'price': best_offer['proposed_price'],
'reason': best_offer['reason'],
}
)
# Add the transfer result message to the state
state['messages'].extend(transfer_result['messages'])
state['messages'] = [transfer_result['messages']]
# Update team rosters and budgets
from_team = best_offer['from_team']
@ -365,8 +382,11 @@ def create_simulator_agent():
temperature=0.7, model=DEFAULT_MODEL
) # Higher temperature for more creative events
prompt = ChatPromptTemplate.from_template("""
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.
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.
Ensure that you use as much as possible from the teams_context to create the event.
Use the existing events to get a sense of the narrative unfolding.
Current NBA landscape:
{teams_context}
@ -383,9 +403,35 @@ def create_simulator_agent():
simulator_prompt, simulator_llm = create_simulator_agent()
def create_analyzer_agent():
llm = ChatOpenAI(temperature=0.7, model='gpt-4o') # Higher temperature for more creative events
prompt = ChatPromptTemplate.from_template("""
You are an NBA Simulation Turn analyzer.
Your task is to compare the previous state of the graph (before the last turn) and the current state of the graph (after the last turn) and provide a brief analysis of the changes that occurred.
Make sure to put more emphasis on the transfers or reasons why transfers were not made.
Previous state of the graph:
{previous_state}
Current state of the graph:
{current_state}
Changes that occurred during the last turn:
{changes}
Output the analysis as a brief summary of the changes that occurred.
Analysis:
""")
return prompt, llm
analyzer_prompt, analyzer_llm = create_analyzer_agent()
async def simulate_event(state: SimulationState) -> SimulationState:
teams_context = await fetch_all_teams_context.ainvoke({'teams': VALID_TEAMS})
result = await simulator_llm.ainvoke(
simulator_prompt.format_prompt(teams_context=json.dumps(teams_context, indent=2))
)
@ -396,12 +442,29 @@ async def simulate_event(state: SimulationState) -> SimulationState:
return {
**state,
'event': new_event,
'teams_context': teams_context,
'all_events': existing_events,
'transfer_offers': [],
'current_iteration': state['current_iteration'] + 1,
}
async def analyze_simulation_turn(state: SimulationState) -> SimulationState:
current_teams_context = state['teams_context']
updated_team_context = await fetch_all_teams_context.ainvoke({'teams': VALID_TEAMS})
result = await analyzer_llm.ainvoke(
analyzer_prompt.format_prompt(
previous_state=json.dumps(current_teams_context, indent=2),
current_state=json.dumps(updated_team_context, indent=2),
changes=state['messages'],
)
)
summary = result.content
logger.info(f'Analysis of the last turn: {summary}')
return state
# Create the graph
workflow = StateGraph(SimulationState)
@ -411,7 +474,7 @@ workflow.add_node('process_event', process_event)
for team in VALID_TEAMS:
workflow.add_node(f'agent_{team}', create_team_agent(team))
workflow.add_node('process_transfers', process_transfers)
workflow.add_node('analyze_simulation_turn', analyze_simulation_turn)
# Add edges
workflow.add_edge(START, 'simulate_event')
workflow.add_edge('simulate_event', 'process_event')
@ -431,8 +494,10 @@ def routing_function(state: SimulationState) -> str:
return 'simulate_event'
workflow.add_edge('process_transfers', 'analyze_simulation_turn')
workflow.add_conditional_edges(
'process_transfers',
'analyze_simulation_turn',
routing_function,
)