Skip to content

Commit 6ed2d2b

Browse files
pritipsinghdirkbrnd
pritipsingh
and
dirkbrnd
authored
feat: add team_session_state state management in teams (#3404)
## Summary This PR introduces a distinction between two types of state management in teams: `team_session_state` (Shared State) Purpose: Data that needs to be shared and synchronized across all team members Scope: Automatically propagated to all nested teams and agents within the team hierarchy Use Case: Shared resources like shopping lists, project data, or any information that multiple team members need to access and modify Example: A shopping list that both the main team and sub-teams can add/remove items from `session_state` (Team-Specific State) Purpose: Data that belongs exclusively to a specific team Scope: Private to the individual team, not shared with members or parent teams Use Case: Team-specific configuration, preferences, history, or metadata Example: Team budget, shopping history, store preferences, or team-specific settings (If applicable, issue number: #____) ## Type of change - [ ] Bug fix - [x] New feature - [x] Breaking change - [ ] Improvement - [ ] Model update - [ ] Other: --- ## Checklist - [ ] Code complies with style guidelines - [ ] Ran format/validation scripts (`./scripts/format.sh` and `./scripts/validate.sh`) - [ ] Self-review completed - [ ] Documentation updated (comments, docstrings) - [ ] Examples and guides: Relevant cookbook examples have been included or updated (if applicable) - [ ] Tested in clean environment - [ ] Tests added/updated (if applicable) --- ## Additional Notes Add any important context (deployment instructions, screenshots, security considerations, etc.) --------- Co-authored-by: Dirk Brand <dirkbrnd@gmail.com>
1 parent d353e07 commit 6ed2d2b

File tree

5 files changed

+476
-27
lines changed

5 files changed

+476
-27
lines changed

‎cookbook/teams/team_with_nested_shared_state.py‎

Lines changed: 219 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,219 @@
1+
from agno.agent.agent import Agent
2+
from agno.models.openai.chat import OpenAIChat
3+
from agno.team import Team
4+
5+
6+
# Define tools to manage our shopping list
7+
def add_item(agent: Agent, item: str) -> str:
8+
"""Add an item to the shopping list and return confirmation.
9+
10+
Args:
11+
item (str): The item to add to the shopping list.
12+
"""
13+
# Add the item if it's not already in the list
14+
if item.lower() not in [
15+
i.lower() for i in agent.team_session_state["shopping_list"]
16+
]:
17+
agent.team_session_state["shopping_list"].append(item)
18+
return f"Added '{item}' to the shopping list"
19+
else:
20+
return f"'{item}' is already in the shopping list"
21+
22+
23+
def remove_item(agent: Agent, item: str) -> str:
24+
"""Remove an item from the shopping list by name.
25+
26+
Args:
27+
item (str): The item to remove from the shopping list.
28+
"""
29+
# Case-insensitive search
30+
for i, list_item in enumerate(agent.team_session_state["shopping_list"]):
31+
if list_item.lower() == item.lower():
32+
agent.team_session_state["shopping_list"].pop(i)
33+
return f"Removed '{list_item}' from the shopping list"
34+
35+
return f"'{item}' was not found in the shopping list. Current shopping list: {agent.team_session_state['shopping_list']}"
36+
37+
38+
def remove_all_items(agent: Agent) -> str:
39+
"""Remove all items from the shopping list."""
40+
agent.team_session_state["shopping_list"] = []
41+
return "All items removed from the shopping list"
42+
43+
44+
shopping_list_agent = Agent(
45+
name="Shopping List Agent",
46+
role="Manage the shopping list",
47+
agent_id="shopping_list_manager",
48+
model=OpenAIChat(id="gpt-4o-mini"),
49+
tools=[add_item, remove_item, remove_all_items],
50+
instructions=[
51+
"Manage the shopping list by adding and removing items",
52+
"Always confirm when items are added or removed",
53+
"If the task is done, update the session state to log the changes & chores you've performed",
54+
],
55+
)
56+
57+
58+
# Shopping management team - new layer for handling all shopping list modifications
59+
shopping_mgmt_team = Team(
60+
name="Shopping Management Team",
61+
team_id="shopping_management",
62+
mode="coordinate",
63+
model=OpenAIChat(id="gpt-4o-mini"),
64+
show_tool_calls=True,
65+
members=[shopping_list_agent],
66+
instructions=[
67+
"Manage adding and removing items from the shopping list using the Shopping List Agent",
68+
"Forward requests to add or remove items to the Shopping List Agent",
69+
],
70+
)
71+
72+
73+
def get_ingredients(agent: Agent) -> str:
74+
"""Retrieve ingredients from the shopping list to use for recipe suggestions.
75+
76+
Args:
77+
meal_type (str): Type of meal to suggest (breakfast, lunch, dinner, snack, or any)
78+
"""
79+
shopping_list = agent.team_session_state["shopping_list"]
80+
81+
if not shopping_list:
82+
return "The shopping list is empty. Add some ingredients first to get recipe suggestions."
83+
84+
# Just return the ingredients - the agent will create recipes
85+
return f"Available ingredients from shopping list: {', '.join(shopping_list)}"
86+
87+
88+
recipe_agent = Agent(
89+
name="Recipe Suggester",
90+
agent_id="recipe_suggester",
91+
role="Suggest recipes based on available ingredients",
92+
model=OpenAIChat(id="gpt-4o-mini"),
93+
tools=[get_ingredients],
94+
instructions=[
95+
"First, use the get_ingredients tool to get the current ingredients from the shopping list",
96+
"After getting the ingredients, create detailed recipe suggestions based on those ingredients",
97+
"Create at least 3 different recipe ideas using the available ingredients",
98+
"For each recipe, include: name, ingredients needed (highlighting which ones are from the shopping list), and brief preparation steps",
99+
"Be creative but practical with recipe suggestions",
100+
"Consider common pantry items that people usually have available in addition to shopping list items",
101+
"Consider dietary preferences if mentioned by the user",
102+
"If no meal type is specified, suggest a variety of options (breakfast, lunch, dinner, snacks)",
103+
],
104+
)
105+
106+
107+
def list_items(team: Team) -> str:
108+
"""List all items in the shopping list."""
109+
shopping_list = team.team_session_state["shopping_list"]
110+
111+
if not shopping_list:
112+
return "The shopping list is empty."
113+
114+
items_text = "\n".join([f"- {item}" for item in shopping_list])
115+
return f"Current shopping list:\n{items_text}"
116+
117+
118+
# Create meal planning subteam
119+
meal_planning_team = Team(
120+
name="Meal Planning Team",
121+
team_id="meal_planning",
122+
mode="coordinate",
123+
model=OpenAIChat(id="gpt-4o-mini"),
124+
members=[recipe_agent],
125+
instructions=[
126+
"You are a meal planning team that suggests recipes based on shopping list items.",
127+
"IMPORTANT: When users ask 'What can I make with these ingredients?' or any recipe-related questions, IMMEDIATELY forward the EXACT SAME request to the recipe_agent WITHOUT asking for further information.",
128+
"DO NOT ask the user for ingredients - the recipe_agent will work with what's already in the shopping list.",
129+
"Your primary job is to forward recipe requests directly to the recipe_agent without modification.",
130+
],
131+
)
132+
133+
134+
def add_chore(team: Team, chore: str, priority: str = "medium") -> str:
135+
"""Add a chore to the list with priority level.
136+
137+
Args:
138+
chore (str): The chore to add to the list
139+
priority (str): Priority level of the chore (low, medium, high)
140+
141+
Returns:
142+
str: Confirmation message
143+
"""
144+
# Initialize chores list if it doesn't exist
145+
if "chores" not in team.session_state:
146+
team.session_state["chores"] = []
147+
148+
# Validate priority
149+
valid_priorities = ["low", "medium", "high"]
150+
if priority.lower() not in valid_priorities:
151+
priority = "medium" # Default to medium if invalid
152+
153+
# Add the chore with timestamp and priority
154+
from datetime import datetime
155+
156+
chore_entry = {
157+
"description": chore,
158+
"priority": priority.lower(),
159+
"added_at": datetime.now().strftime("%Y-%m-%d %H:%M"),
160+
}
161+
162+
team.session_state["chores"].append(chore_entry)
163+
164+
return f"Added chore: '{chore}' with {priority} priority"
165+
166+
167+
shopping_team = Team(
168+
name="Shopping List Team",
169+
mode="coordinate",
170+
model=OpenAIChat(id="gpt-4o-mini"),
171+
team_session_state={"shopping_list": []},
172+
tools=[list_items, add_chore],
173+
session_state={"chores": []},
174+
team_id="shopping_list_team",
175+
members=[
176+
shopping_mgmt_team,
177+
meal_planning_team,
178+
],
179+
show_tool_calls=True,
180+
markdown=True,
181+
instructions=[
182+
"You are a team that manages a shopping list & helps plan meals using that list.",
183+
"If you need to add or remove items from the shopping list, forward the full request to the Shopping Management Team.",
184+
"IMPORTANT: If the user asks about recipes or what they can make with ingredients, IMMEDIATELY forward the EXACT request to the meal_planning_team with NO additional questions.",
185+
"Example: When user asks 'What can I make with these ingredients?', you should simply forward this exact request to meal_planning_team without asking for more information.",
186+
"If you need to list the items in the shopping list, use the list_items tool.",
187+
"If the user got something from the shopping list, it means it can be removed from the shopping list.",
188+
"After each completed task, use the add_chore tool to log exactly what was done with high priority.",
189+
],
190+
show_members_responses=True,
191+
)
192+
193+
# Example usage
194+
shopping_team.print_response(
195+
"Add milk, eggs, and bread to the shopping list", stream=True
196+
)
197+
print(f"Session state: {shopping_team.team_session_state}")
198+
199+
shopping_team.print_response("I got bread", stream=True)
200+
print(f"Session state: {shopping_team.team_session_state}")
201+
202+
shopping_team.print_response("I need apples and oranges", stream=True)
203+
print(f"Session state: {shopping_team.team_session_state}")
204+
205+
shopping_team.print_response("whats on my list?", stream=True)
206+
print(f"Session state: {shopping_team.team_session_state}")
207+
208+
# Try the meal planning feature
209+
shopping_team.print_response("What can I make with these ingredients?", stream=True)
210+
print(f"Session state: {shopping_team.team_session_state}")
211+
212+
shopping_team.print_response(
213+
"Clear everything from my list and start over with just bananas and yogurt",
214+
stream=True,
215+
)
216+
print(f"Shared Session state: {shopping_team.team_session_state}")
217+
218+
219+
print(f"Team session state: {shopping_team.session_state}")

‎cookbook/teams/team_with_shared_state.py‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ def remove_all_items(agent: Agent) -> str:
5454

5555
def list_items(team: Team) -> str:
5656
"""List all items in the shopping list."""
57-
shopping_list = team.session_state["shopping_list"]
57+
shopping_list = team.team_session_state["shopping_list"]
5858

5959
if not shopping_list:
6060
return "The shopping list is empty."
@@ -67,7 +67,7 @@ def list_items(team: Team) -> str:
6767
name="Shopping List Team",
6868
mode="coordinate",
6969
model=OpenAIChat(id="gpt-4o-mini"),
70-
session_state={"shopping_list": []},
70+
team_session_state={"shopping_list": []},
7171
tools=[list_items],
7272
members=[
7373
shopping_list_agent,
@@ -87,19 +87,19 @@ def list_items(team: Team) -> str:
8787
shopping_team.print_response(
8888
"Add milk, eggs, and bread to the shopping list", stream=True
8989
)
90-
print(f"Session state: {shopping_team.session_state}")
90+
print(f"Session state: {shopping_team.team_session_state}")
9191

9292
shopping_team.print_response("I got bread", stream=True)
93-
print(f"Session state: {shopping_team.session_state}")
93+
print(f"Session state: {shopping_team.team_session_state}")
9494

9595
shopping_team.print_response("I need apples and oranges", stream=True)
96-
print(f"Session state: {shopping_team.session_state}")
96+
print(f"Session state: {shopping_team.team_session_state}")
9797

9898
shopping_team.print_response("whats on my list?", stream=True)
99-
print(f"Session state: {shopping_team.session_state}")
99+
print(f"Session state: {shopping_team.team_session_state}")
100100

101101
shopping_team.print_response(
102102
"Clear everything from my list and start over with just bananas and yogurt",
103103
stream=True,
104104
)
105-
print(f"Session state: {shopping_team.session_state}")
105+
print(f"Session state: {shopping_team.team_session_state}")

‎libs/agno/agno/agent/agent.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3975,6 +3975,7 @@ def format_message_with_state_variables(self, msg: Any) -> Any:
39753975

39763976
format_variables = ChainMap(
39773977
self.session_state or {},
3978+
self.team_session_state or {},
39783979
self.context or {},
39793980
self.extra_data or {},
39803981
{"user_id": self.user_id} if self.user_id is not None else {},

‎libs/agno/agno/team/team.py‎

Lines changed: 38 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,9 @@ class Team:
107107
session_name: Optional[str] = None
108108
# Session state (stored in the database to persist across runs)
109109
session_state: Optional[Dict[str, Any]] = None
110+
111+
# Team session state (shared between team leaders and team members)
112+
team_session_state: Optional[Dict[str, Any]] = None
110113
# If True, add the session state variables in the user and system messages
111114
add_state_in_messages: bool = False
112115

@@ -247,6 +250,7 @@ def __init__(
247250
session_id: Optional[str] = None,
248251
session_name: Optional[str] = None,
249252
session_state: Optional[Dict[str, Any]] = None,
253+
team_session_state: Optional[Dict[str, Any]] = None,
250254
add_state_in_messages: bool = False,
251255
description: Optional[str] = None,
252256
instructions: Optional[Union[str, List[str], Callable]] = None,
@@ -309,6 +313,7 @@ def __init__(
309313
self.session_id = session_id
310314
self.session_name = session_name
311315
self.session_state = session_state
316+
self.team_session_state = team_session_state
312317
self.add_state_in_messages = add_state_in_messages
313318

314319
self.description = description
@@ -439,17 +444,11 @@ def _initialize_member(self, member: Union["Team", Agent], session_id: Optional[
439444
member.team_session_id = session_id
440445

441446
# Set the team session state on members
442-
if self.session_state is not None:
443-
if isinstance(member, Agent):
444-
if member.team_session_state is None:
445-
member.team_session_state = self.session_state
446-
else:
447-
merge_dictionaries(member.team_session_state, self.session_state)
448-
elif isinstance(member, Team):
449-
if member.session_state is None:
450-
member.session_state = self.session_state
451-
else:
452-
merge_dictionaries(member.session_state, self.session_state)
447+
if self.team_session_state is not None:
448+
if member.team_session_state is None:
449+
member.team_session_state = self.team_session_state
450+
else:
451+
merge_dictionaries(member.team_session_state, self.team_session_state)
453452

454453
if isinstance(member, Agent):
455454
member.team_id = self.team_id
@@ -1912,10 +1911,16 @@ def _handle_model_response_chunk(
19121911

19131912
def _initialize_session_state(self, user_id: Optional[str] = None, session_id: Optional[str] = None) -> None:
19141913
self.session_state = self.session_state or {}
1914+
19151915
if user_id is not None:
19161916
self.session_state["current_user_id"] = user_id
1917+
if self.team_session_state is not None:
1918+
self.team_session_state["current_user_id"] = user_id
1919+
19171920
if session_id is not None:
19181921
self.session_state["current_session_id"] = session_id
1922+
if self.team_session_state is not None:
1923+
self.team_session_state["current_session_id"] = session_id
19191924

19201925
def _make_memories_and_summaries(
19211926
self, run_messages: RunMessages, session_id: str, user_id: Optional[str] = None
@@ -4752,6 +4757,7 @@ def _format_message_with_state_variables(self, message: str, user_id: Optional[s
47524757
"""Format a message with the session state variables."""
47534758
format_variables = ChainMap(
47544759
self.session_state or {},
4760+
self.team_session_state or {},
47554761
self.context or {},
47564762
self.extra_data or {},
47574763
{"user_id": user_id} if user_id is not None else {},
@@ -5001,16 +5007,12 @@ def set_shared_context(state: Union[str, dict]) -> str:
50015007
return set_shared_context
50025008

50035009
def _update_team_session_state(self, member_agent: Union[Agent, "Team"]) -> None:
5004-
if isinstance(member_agent, Agent) and member_agent.team_session_state is not None:
5005-
if self.session_state is None:
5006-
self.session_state = member_agent.team_session_state
5007-
else:
5008-
merge_dictionaries(self.session_state, member_agent.team_session_state)
5009-
elif isinstance(member_agent, Team) and member_agent.session_state is not None:
5010-
if self.session_state is None:
5011-
self.session_state = member_agent.session_state
5010+
"""Update team session state from either an Agent or nested Team member"""
5011+
if member_agent.team_session_state is not None:
5012+
if self.team_session_state is None:
5013+
self.team_session_state = member_agent.team_session_state
50125014
else:
5013-
merge_dictionaries(self.session_state, member_agent.session_state)
5015+
merge_dictionaries(self.team_session_state, member_agent.team_session_state)
50145016

50155017
def get_run_member_agents_function(
50165018
self,
@@ -6147,6 +6149,20 @@ def load_team_session(self, session: TeamSession):
61476149
# Update the current session_state
61486150
self.session_state = session_state_from_db
61496151

6152+
if "team_session_state" in session.session_data:
6153+
team_session_state_from_db = session.session_data.get("team_session_state")
6154+
if (
6155+
team_session_state_from_db is not None
6156+
and isinstance(team_session_state_from_db, dict)
6157+
and len(team_session_state_from_db) > 0
6158+
):
6159+
# If the team_session_state is already set, merge the team_session_state from the database with the current team_session_state
6160+
if self.team_session_state is not None and len(self.team_session_state) > 0:
6161+
# This updates team_session_state_from_db
6162+
merge_dictionaries(team_session_state_from_db, self.team_session_state)
6163+
# Update the current team_session_state
6164+
self.team_session_state = team_session_state_from_db
6165+
61506166
# Get the session_metrics from the database
61516167
if "session_metrics" in session.session_data:
61526168
session_metrics_from_db = session.session_data.get("session_metrics")
@@ -6905,6 +6921,8 @@ def _get_session_data(self) -> Dict[str, Any]:
69056921
session_data["session_name"] = self.session_name
69066922
if self.session_state is not None and len(self.session_state) > 0:
69076923
session_data["session_state"] = self.session_state
6924+
if self.team_session_state is not None and len(self.team_session_state) > 0:
6925+
session_data["team_session_state"] = self.team_session_state
69086926
if self.session_metrics is not None:
69096927
session_data["session_metrics"] = asdict(self.session_metrics) if self.session_metrics is not None else None
69106928
if self.images is not None:

0 commit comments

Comments
 (0)