agno 2.3.13__py3-none-any.whl → 2.3.15__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (49) hide show
  1. agno/agent/agent.py +1149 -1392
  2. agno/db/migrations/manager.py +3 -3
  3. agno/eval/__init__.py +21 -8
  4. agno/knowledge/embedder/azure_openai.py +0 -1
  5. agno/knowledge/embedder/google.py +1 -1
  6. agno/models/anthropic/claude.py +9 -4
  7. agno/models/base.py +8 -4
  8. agno/models/metrics.py +12 -0
  9. agno/models/openai/chat.py +2 -0
  10. agno/models/openai/responses.py +2 -2
  11. agno/os/app.py +59 -2
  12. agno/os/auth.py +40 -3
  13. agno/os/interfaces/a2a/router.py +619 -9
  14. agno/os/interfaces/a2a/utils.py +31 -32
  15. agno/os/middleware/jwt.py +5 -5
  16. agno/os/router.py +1 -57
  17. agno/os/routers/agents/schema.py +14 -1
  18. agno/os/routers/database.py +150 -0
  19. agno/os/routers/teams/schema.py +14 -1
  20. agno/os/settings.py +3 -0
  21. agno/os/utils.py +61 -53
  22. agno/reasoning/anthropic.py +85 -1
  23. agno/reasoning/azure_ai_foundry.py +93 -1
  24. agno/reasoning/deepseek.py +91 -1
  25. agno/reasoning/gemini.py +81 -1
  26. agno/reasoning/groq.py +103 -1
  27. agno/reasoning/manager.py +1244 -0
  28. agno/reasoning/ollama.py +93 -1
  29. agno/reasoning/openai.py +113 -1
  30. agno/reasoning/vertexai.py +85 -1
  31. agno/run/agent.py +21 -0
  32. agno/run/base.py +20 -1
  33. agno/run/team.py +21 -0
  34. agno/session/team.py +0 -3
  35. agno/team/team.py +1211 -1445
  36. agno/tools/toolkit.py +119 -8
  37. agno/utils/events.py +99 -4
  38. agno/utils/hooks.py +4 -10
  39. agno/utils/print_response/agent.py +26 -0
  40. agno/utils/print_response/team.py +11 -0
  41. agno/utils/prompts.py +8 -6
  42. agno/utils/string.py +46 -0
  43. agno/utils/team.py +1 -1
  44. agno/vectordb/milvus/milvus.py +32 -3
  45. {agno-2.3.13.dist-info → agno-2.3.15.dist-info}/METADATA +3 -2
  46. {agno-2.3.13.dist-info → agno-2.3.15.dist-info}/RECORD +49 -47
  47. {agno-2.3.13.dist-info → agno-2.3.15.dist-info}/WHEEL +0 -0
  48. {agno-2.3.13.dist-info → agno-2.3.15.dist-info}/licenses/LICENSE +0 -0
  49. {agno-2.3.13.dist-info → agno-2.3.15.dist-info}/top_level.txt +0 -0
agno/team/team.py CHANGED
@@ -11,7 +11,6 @@ from dataclasses import dataclass
11
11
  from os import getenv
12
12
  from textwrap import dedent
13
13
  from typing import (
14
- TYPE_CHECKING,
15
14
  Any,
16
15
  AsyncIterator,
17
16
  Callable,
@@ -33,12 +32,10 @@ from uuid import uuid4
33
32
 
34
33
  from pydantic import BaseModel
35
34
 
36
- if TYPE_CHECKING:
37
- from agno.eval.base import BaseEval
38
-
39
35
  from agno.agent import Agent
40
36
  from agno.compression.manager import CompressionManager
41
37
  from agno.db.base import AsyncBaseDb, BaseDb, SessionType, UserMemory
38
+ from agno.eval.base import BaseEval
42
39
  from agno.exceptions import (
43
40
  InputCheckError,
44
41
  OutputCheckError,
@@ -67,7 +64,12 @@ from agno.run.cancel import (
67
64
  register_run,
68
65
  )
69
66
  from agno.run.messages import RunMessages
70
- from agno.run.team import TeamRunEvent, TeamRunInput, TeamRunOutput, TeamRunOutputEvent
67
+ from agno.run.team import (
68
+ TeamRunEvent,
69
+ TeamRunInput,
70
+ TeamRunOutput,
71
+ TeamRunOutputEvent,
72
+ )
71
73
  from agno.session import SessionSummaryManager, TeamSession, WorkflowSession
72
74
  from agno.session.summary import SessionSummary
73
75
  from agno.tools import Toolkit
@@ -107,6 +109,7 @@ from agno.utils.agent import (
107
109
  )
108
110
  from agno.utils.common import is_typed_dict, validate_typed_dict
109
111
  from agno.utils.events import (
112
+ add_team_error_event,
110
113
  create_team_parser_model_response_completed_event,
111
114
  create_team_parser_model_response_started_event,
112
115
  create_team_post_hook_completed_event,
@@ -114,16 +117,19 @@ from agno.utils.events import (
114
117
  create_team_pre_hook_completed_event,
115
118
  create_team_pre_hook_started_event,
116
119
  create_team_reasoning_completed_event,
120
+ create_team_reasoning_content_delta_event,
117
121
  create_team_reasoning_started_event,
118
122
  create_team_reasoning_step_event,
119
123
  create_team_run_cancelled_event,
120
124
  create_team_run_completed_event,
121
125
  create_team_run_content_completed_event,
126
+ create_team_run_error_event,
122
127
  create_team_run_output_content_event,
123
128
  create_team_run_started_event,
124
129
  create_team_session_summary_completed_event,
125
130
  create_team_session_summary_started_event,
126
131
  create_team_tool_call_completed_event,
132
+ create_team_tool_call_error_event,
127
133
  create_team_tool_call_started_event,
128
134
  handle_event,
129
135
  )
@@ -161,12 +167,11 @@ from agno.utils.reasoning import (
161
167
  update_run_output_with_reasoning,
162
168
  )
163
169
  from agno.utils.response import (
164
- async_generator_wrapper,
165
170
  check_if_run_cancelled,
166
171
  generator_wrapper,
167
172
  )
168
173
  from agno.utils.safe_formatter import SafeFormatter
169
- from agno.utils.string import generate_id_from_name, parse_response_model_str
174
+ from agno.utils.string import generate_id_from_name, parse_response_dict_str, parse_response_model_str
170
175
  from agno.utils.team import (
171
176
  add_interaction_to_team_run_context,
172
177
  format_member_agent_task,
@@ -354,17 +359,18 @@ class Team:
354
359
 
355
360
  # --- Team Hooks ---
356
361
  # Functions called right after team session is loaded, before processing starts
357
- pre_hooks: Optional[List[Union[Callable[..., Any], BaseGuardrail, "BaseEval"]]] = None
362
+ pre_hooks: Optional[List[Union[Callable[..., Any], BaseGuardrail, BaseEval]]] = None
358
363
  # Functions called after output is generated but before the response is returned
359
- post_hooks: Optional[List[Union[Callable[..., Any], BaseGuardrail, "BaseEval"]]] = None
364
+ post_hooks: Optional[List[Union[Callable[..., Any], BaseGuardrail, BaseEval]]] = None
360
365
  # If True, run hooks as FastAPI background tasks (non-blocking). Set by AgentOS.
361
366
  _run_hooks_in_background: Optional[bool] = None
362
367
 
363
368
  # --- Structured output ---
364
369
  # Input schema for validating input
365
370
  input_schema: Optional[Type[BaseModel]] = None
366
- # Output schema for the team response
367
- output_schema: Optional[Type[BaseModel]] = None
371
+ # Provide a response model to get the response in the implied format.
372
+ # You can use a Pydantic model or a JSON fitting the provider's expected schema.
373
+ output_schema: Optional[Union[Type[BaseModel], Dict[str, Any]]] = None
368
374
  # Provide a secondary model to parse the response from the primary model
369
375
  parser_model: Optional[Model] = None
370
376
  # Provide a prompt for the parser model
@@ -525,10 +531,10 @@ class Team:
525
531
  tool_call_limit: Optional[int] = None,
526
532
  tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
527
533
  tool_hooks: Optional[List[Callable]] = None,
528
- pre_hooks: Optional[List[Union[Callable[..., Any], BaseGuardrail, "BaseEval"]]] = None,
529
- post_hooks: Optional[List[Union[Callable[..., Any], BaseGuardrail, "BaseEval"]]] = None,
534
+ pre_hooks: Optional[List[Union[Callable[..., Any], BaseGuardrail, BaseEval]]] = None,
535
+ post_hooks: Optional[List[Union[Callable[..., Any], BaseGuardrail, BaseEval]]] = None,
530
536
  input_schema: Optional[Type[BaseModel]] = None,
531
- output_schema: Optional[Type[BaseModel]] = None,
537
+ output_schema: Optional[Union[Type[BaseModel], Dict[str, Any]]] = None,
532
538
  parser_model: Optional[Union[Model, str]] = None,
533
539
  parser_model_prompt: Optional[str] = None,
534
540
  output_model: Optional[Union[Model, str]] = None,
@@ -728,7 +734,7 @@ class Team:
728
734
  self._tool_instructions: Optional[List[str]] = None
729
735
 
730
736
  # True if we should parse a member response model
731
- self._member_response_model: Optional[Type[BaseModel]] = None
737
+ self._member_response_model: Optional[Union[Type[BaseModel], Dict[str, Any]]] = None
732
738
 
733
739
  self._formatter: Optional[SafeFormatter] = None
734
740
 
@@ -1484,7 +1490,6 @@ class Team:
1484
1490
  **kwargs: Any,
1485
1491
  ) -> TeamRunOutput:
1486
1492
  """Run the Team and return the response.
1487
-
1488
1493
  Steps:
1489
1494
  1. Execute pre-hooks
1490
1495
  2. Determine tools for model
@@ -1500,7 +1505,6 @@ class Team:
1500
1505
  12. Create session summary
1501
1506
  13. Cleanup and store (scrub, stop timer, add to session, calculate metrics, save session)
1502
1507
  """
1503
-
1504
1508
  # 1. Execute pre-hooks
1505
1509
  run_input = cast(TeamRunInput, run_response.input)
1506
1510
  self.model = cast(Model, self.model)
@@ -1580,108 +1584,97 @@ class Team:
1580
1584
  self._make_memories, run_messages=run_messages, user_id=user_id
1581
1585
  )
1582
1586
 
1583
- try:
1584
- raise_if_cancelled(run_response.run_id) # type: ignore
1585
-
1586
- # 5. Reason about the task if reasoning is enabled
1587
- self._handle_reasoning(run_response=run_response, run_messages=run_messages)
1587
+ raise_if_cancelled(run_response.run_id) # type: ignore
1588
1588
 
1589
- # Check for cancellation before model call
1590
- raise_if_cancelled(run_response.run_id) # type: ignore
1589
+ # 5. Reason about the task if reasoning is enabled
1590
+ self._handle_reasoning(run_response=run_response, run_messages=run_messages)
1591
1591
 
1592
- # 6. Get the model response for the team leader
1593
- self.model = cast(Model, self.model)
1594
- model_response: ModelResponse = self.model.response(
1595
- messages=run_messages.messages,
1596
- response_format=response_format,
1597
- tools=_tools,
1598
- tool_choice=self.tool_choice,
1599
- tool_call_limit=self.tool_call_limit,
1600
- send_media_to_model=self.send_media_to_model,
1601
- compression_manager=self.compression_manager if self.compress_tool_results else None,
1602
- )
1592
+ # Check for cancellation before model call
1593
+ raise_if_cancelled(run_response.run_id) # type: ignore
1603
1594
 
1604
- # Check for cancellation after model call
1605
- raise_if_cancelled(run_response.run_id) # type: ignore
1595
+ # 6. Get the model response for the team leader
1596
+ self.model = cast(Model, self.model)
1597
+ model_response: ModelResponse = self.model.response(
1598
+ messages=run_messages.messages,
1599
+ response_format=response_format,
1600
+ tools=_tools,
1601
+ tool_choice=self.tool_choice,
1602
+ tool_call_limit=self.tool_call_limit,
1603
+ send_media_to_model=self.send_media_to_model,
1604
+ compression_manager=self.compression_manager if self.compress_tool_results else None,
1605
+ )
1606
1606
 
1607
- # If an output model is provided, generate output using the output model
1608
- self._parse_response_with_output_model(model_response, run_messages)
1607
+ # Check for cancellation after model call
1608
+ raise_if_cancelled(run_response.run_id) # type: ignore
1609
1609
 
1610
- # If a parser model is provided, structure the response separately
1611
- self._parse_response_with_parser_model(model_response, run_messages, run_context=run_context)
1610
+ # If an output model is provided, generate output using the output model
1611
+ self._parse_response_with_output_model(model_response, run_messages)
1612
1612
 
1613
- # 7. Update TeamRunOutput with the model response
1614
- self._update_run_response(
1615
- model_response=model_response,
1616
- run_response=run_response,
1617
- run_messages=run_messages,
1618
- run_context=run_context,
1619
- )
1613
+ # If a parser model is provided, structure the response separately
1614
+ self._parse_response_with_parser_model(model_response, run_messages, run_context=run_context)
1620
1615
 
1621
- # 8. Store media if enabled
1622
- if self.store_media:
1623
- store_media_util(run_response, model_response)
1616
+ # 7. Update TeamRunOutput with the model response
1617
+ self._update_run_response(
1618
+ model_response=model_response,
1619
+ run_response=run_response,
1620
+ run_messages=run_messages,
1621
+ run_context=run_context,
1622
+ )
1624
1623
 
1625
- # 9. Convert response to structured format
1626
- self._convert_response_to_structured_format(run_response=run_response, run_context=run_context)
1624
+ # 8. Store media if enabled
1625
+ if self.store_media:
1626
+ store_media_util(run_response, model_response)
1627
1627
 
1628
- # 10. Execute post-hooks after output is generated but before response is returned
1629
- if self.post_hooks is not None:
1630
- iterator = self._execute_post_hooks(
1631
- hooks=self.post_hooks, # type: ignore
1632
- run_output=run_response,
1633
- run_context=run_context,
1634
- session=session,
1635
- user_id=user_id,
1636
- debug_mode=debug_mode,
1637
- background_tasks=background_tasks,
1638
- **kwargs,
1639
- )
1640
- deque(iterator, maxlen=0)
1641
- raise_if_cancelled(run_response.run_id) # type: ignore
1628
+ # 9. Convert response to structured format
1629
+ self._convert_response_to_structured_format(run_response=run_response, run_context=run_context)
1642
1630
 
1643
- # 11. Wait for background memory creation
1644
- wait_for_open_threads(memory_future=memory_future)
1631
+ # 10. Execute post-hooks after output is generated but before response is returned
1632
+ if self.post_hooks is not None:
1633
+ iterator = self._execute_post_hooks(
1634
+ hooks=self.post_hooks, # type: ignore
1635
+ run_output=run_response,
1636
+ run_context=run_context,
1637
+ session=session,
1638
+ user_id=user_id,
1639
+ debug_mode=debug_mode,
1640
+ background_tasks=background_tasks,
1641
+ **kwargs,
1642
+ )
1643
+ deque(iterator, maxlen=0)
1644
+ raise_if_cancelled(run_response.run_id) # type: ignore
1645
1645
 
1646
- raise_if_cancelled(run_response.run_id) # type: ignore
1646
+ # 11. Wait for background memory creation
1647
+ wait_for_open_threads(memory_future=memory_future)
1647
1648
 
1648
- # 12. Create session summary
1649
- if self.session_summary_manager is not None:
1650
- # Upsert the RunOutput to Team Session before creating the session summary
1651
- session.upsert_run(run_response=run_response)
1652
- try:
1653
- self.session_summary_manager.create_session_summary(session=session)
1654
- except Exception as e:
1655
- log_warning(f"Error in session summary creation: {str(e)}")
1649
+ raise_if_cancelled(run_response.run_id) # type: ignore
1656
1650
 
1657
- raise_if_cancelled(run_response.run_id) # type: ignore
1651
+ # 12. Create session summary
1652
+ if self.session_summary_manager is not None:
1653
+ # Upsert the RunOutput to Team Session before creating the session summary
1654
+ session.upsert_run(run_response=run_response)
1655
+ try:
1656
+ self.session_summary_manager.create_session_summary(session=session)
1657
+ except Exception as e:
1658
+ log_warning(f"Error in session summary creation: {str(e)}")
1658
1659
 
1659
- # Set the run status to completed
1660
- run_response.status = RunStatus.completed
1660
+ raise_if_cancelled(run_response.run_id) # type: ignore
1661
1661
 
1662
- # 13. Cleanup and store the run response
1663
- self._cleanup_and_store(run_response=run_response, session=session)
1662
+ # Set the run status to completed
1663
+ run_response.status = RunStatus.completed
1664
1664
 
1665
- # Log Team Telemetry
1666
- self._log_team_telemetry(session_id=session.session_id, run_id=run_response.run_id)
1665
+ # 13. Cleanup and store the run response
1666
+ self._cleanup_and_store(run_response=run_response, session=session)
1667
1667
 
1668
- log_debug(f"Team Run End: {run_response.run_id}", center=True, symbol="*")
1668
+ # Log Team Telemetry
1669
+ self._log_team_telemetry(session_id=session.session_id, run_id=run_response.run_id)
1669
1670
 
1670
- return run_response
1671
+ log_debug(f"Team Run End: {run_response.run_id}", center=True, symbol="*")
1671
1672
 
1672
- except RunCancelledException as e:
1673
- # Handle run cancellation during streaming
1674
- log_info(f"Team run {run_response.run_id} was cancelled")
1675
- run_response.status = RunStatus.cancelled
1676
- run_response.content = str(e)
1673
+ # Disconnect tools and clean up run tracking
1674
+ self._disconnect_connectable_tools()
1675
+ cleanup_run(run_response.run_id) # type: ignore
1677
1676
 
1678
- # Add the RunOutput to Team Session even when cancelled
1679
- self._cleanup_and_store(run_response=run_response, session=session)
1680
- return run_response
1681
- finally:
1682
- # Always disconnect connectable tools
1683
- self._disconnect_connectable_tools()
1684
- cleanup_run(run_response.run_id) # type: ignore
1677
+ return run_response
1685
1678
 
1686
1679
  def _run_stream(
1687
1680
  self,
@@ -1700,7 +1693,6 @@ class Team:
1700
1693
  **kwargs: Any,
1701
1694
  ) -> Iterator[Union[TeamRunOutputEvent, RunOutputEvent, TeamRunOutput]]:
1702
1695
  """Run the Team and return the response iterator.
1703
-
1704
1696
  Steps:
1705
1697
  1. Execute pre-hooks
1706
1698
  2. Determine tools for model
@@ -1794,190 +1786,171 @@ class Team:
1794
1786
  self._make_memories, run_messages=run_messages, user_id=user_id
1795
1787
  )
1796
1788
 
1797
- try:
1798
- # Start the Run by yielding a RunStarted event
1799
- if stream_events:
1800
- yield handle_event( # type: ignore
1801
- create_team_run_started_event(run_response),
1802
- run_response,
1803
- events_to_skip=self.events_to_skip,
1804
- store_events=self.store_events,
1805
- )
1789
+ # Start the Run by yielding a RunStarted event
1790
+ if stream_events:
1791
+ yield handle_event( # type: ignore
1792
+ create_team_run_started_event(run_response),
1793
+ run_response,
1794
+ events_to_skip=self.events_to_skip,
1795
+ store_events=self.store_events,
1796
+ )
1797
+
1798
+ raise_if_cancelled(run_response.run_id) # type: ignore
1799
+
1800
+ # 5. Reason about the task if reasoning is enabled
1801
+ yield from self._handle_reasoning_stream(
1802
+ run_response=run_response,
1803
+ run_messages=run_messages,
1804
+ stream_events=stream_events,
1805
+ )
1806
1806
 
1807
- raise_if_cancelled(run_response.run_id) # type: ignore
1807
+ # Check for cancellation before model processing
1808
+ raise_if_cancelled(run_response.run_id) # type: ignore
1808
1809
 
1809
- # 5. Reason about the task if reasoning is enabled
1810
- yield from self._handle_reasoning_stream(
1810
+ # 6. Get a response from the model
1811
+ if self.output_model is None:
1812
+ for event in self._handle_model_response_stream(
1813
+ session=session,
1811
1814
  run_response=run_response,
1812
1815
  run_messages=run_messages,
1816
+ tools=_tools,
1817
+ response_format=response_format,
1813
1818
  stream_events=stream_events,
1814
- )
1815
-
1816
- # Check for cancellation before model processing
1817
- raise_if_cancelled(run_response.run_id) # type: ignore
1819
+ session_state=run_context.session_state,
1820
+ run_context=run_context,
1821
+ ):
1822
+ raise_if_cancelled(run_response.run_id) # type: ignore
1823
+ yield event
1824
+ else:
1825
+ for event in self._handle_model_response_stream(
1826
+ session=session,
1827
+ run_response=run_response,
1828
+ run_messages=run_messages,
1829
+ tools=_tools,
1830
+ response_format=response_format,
1831
+ stream_events=stream_events,
1832
+ session_state=run_context.session_state,
1833
+ run_context=run_context,
1834
+ ):
1835
+ raise_if_cancelled(run_response.run_id) # type: ignore
1836
+ from agno.run.team import IntermediateRunContentEvent, RunContentEvent
1818
1837
 
1819
- # 6. Get a response from the model
1820
- if self.output_model is None:
1821
- for event in self._handle_model_response_stream(
1822
- session=session,
1823
- run_response=run_response,
1824
- run_messages=run_messages,
1825
- tools=_tools,
1826
- response_format=response_format,
1827
- stream_events=stream_events,
1828
- session_state=run_context.session_state,
1829
- run_context=run_context,
1830
- ):
1831
- raise_if_cancelled(run_response.run_id) # type: ignore
1838
+ if isinstance(event, RunContentEvent):
1839
+ if stream_events:
1840
+ yield IntermediateRunContentEvent(
1841
+ content=event.content,
1842
+ content_type=event.content_type,
1843
+ )
1844
+ else:
1832
1845
  yield event
1833
- else:
1834
- for event in self._handle_model_response_stream(
1835
- session=session,
1836
- run_response=run_response,
1837
- run_messages=run_messages,
1838
- tools=_tools,
1839
- response_format=response_format,
1840
- stream_events=stream_events,
1841
- session_state=run_context.session_state,
1842
- run_context=run_context,
1843
- ):
1844
- raise_if_cancelled(run_response.run_id) # type: ignore
1845
- from agno.run.team import IntermediateRunContentEvent, RunContentEvent
1846
1846
 
1847
- if isinstance(event, RunContentEvent):
1848
- if stream_events:
1849
- yield IntermediateRunContentEvent(
1850
- content=event.content,
1851
- content_type=event.content_type,
1852
- )
1853
- else:
1854
- yield event
1847
+ for event in self._generate_response_with_output_model_stream(
1848
+ session=session,
1849
+ run_response=run_response,
1850
+ run_messages=run_messages,
1851
+ stream_events=stream_events,
1852
+ ):
1853
+ raise_if_cancelled(run_response.run_id) # type: ignore
1854
+ yield event
1855
1855
 
1856
- for event in self._generate_response_with_output_model_stream(
1857
- session=session,
1858
- run_response=run_response,
1859
- run_messages=run_messages,
1860
- stream_events=stream_events,
1861
- ):
1862
- raise_if_cancelled(run_response.run_id) # type: ignore
1863
- yield event
1856
+ # Check for cancellation after model processing
1857
+ raise_if_cancelled(run_response.run_id) # type: ignore
1864
1858
 
1865
- # Check for cancellation after model processing
1866
- raise_if_cancelled(run_response.run_id) # type: ignore
1859
+ # 7. Parse response with parser model if provided
1860
+ yield from self._parse_response_with_parser_model_stream(
1861
+ session=session, run_response=run_response, stream_events=stream_events, run_context=run_context
1862
+ )
1867
1863
 
1868
- # 7. Parse response with parser model if provided
1869
- yield from self._parse_response_with_parser_model_stream(
1870
- session=session, run_response=run_response, stream_events=stream_events, run_context=run_context
1864
+ # Yield RunContentCompletedEvent
1865
+ if stream_events:
1866
+ yield handle_event( # type: ignore
1867
+ create_team_run_content_completed_event(from_run_response=run_response),
1868
+ run_response,
1869
+ events_to_skip=self.events_to_skip,
1870
+ store_events=self.store_events,
1871
+ )
1872
+ # Execute post-hooks after output is generated but before response is returned
1873
+ if self.post_hooks is not None:
1874
+ yield from self._execute_post_hooks(
1875
+ hooks=self.post_hooks, # type: ignore
1876
+ run_output=run_response,
1877
+ run_context=run_context,
1878
+ session=session,
1879
+ user_id=user_id,
1880
+ debug_mode=debug_mode,
1881
+ stream_events=stream_events,
1882
+ background_tasks=background_tasks,
1883
+ **kwargs,
1871
1884
  )
1885
+ raise_if_cancelled(run_response.run_id) # type: ignore
1886
+
1887
+ # 8. Wait for background memory creation
1888
+ yield from wait_for_thread_tasks_stream(
1889
+ run_response=run_response,
1890
+ memory_future=memory_future,
1891
+ stream_events=stream_events,
1892
+ events_to_skip=self.events_to_skip, # type: ignore
1893
+ store_events=self.store_events,
1894
+ )
1895
+
1896
+ raise_if_cancelled(run_response.run_id) # type: ignore
1897
+ # 9. Create session summary
1898
+ if self.session_summary_manager is not None:
1899
+ # Upsert the RunOutput to Team Session before creating the session summary
1900
+ session.upsert_run(run_response=run_response)
1872
1901
 
1873
- # Yield RunContentCompletedEvent
1874
1902
  if stream_events:
1875
1903
  yield handle_event( # type: ignore
1876
- create_team_run_content_completed_event(from_run_response=run_response),
1904
+ create_team_session_summary_started_event(from_run_response=run_response),
1877
1905
  run_response,
1878
1906
  events_to_skip=self.events_to_skip,
1879
1907
  store_events=self.store_events,
1880
1908
  )
1881
- # Execute post-hooks after output is generated but before response is returned
1882
- if self.post_hooks is not None:
1883
- yield from self._execute_post_hooks(
1884
- hooks=self.post_hooks, # type: ignore
1885
- run_output=run_response,
1886
- run_context=run_context,
1887
- session=session,
1888
- user_id=user_id,
1889
- debug_mode=debug_mode,
1890
- stream_events=stream_events,
1891
- background_tasks=background_tasks,
1892
- **kwargs,
1909
+ try:
1910
+ self.session_summary_manager.create_session_summary(session=session)
1911
+ except Exception as e:
1912
+ log_warning(f"Error in session summary creation: {str(e)}")
1913
+ if stream_events:
1914
+ yield handle_event( # type: ignore
1915
+ create_team_session_summary_completed_event(
1916
+ from_run_response=run_response, session_summary=session.summary
1917
+ ),
1918
+ run_response,
1919
+ events_to_skip=self.events_to_skip,
1920
+ store_events=self.store_events,
1893
1921
  )
1894
- raise_if_cancelled(run_response.run_id) # type: ignore
1895
-
1896
- # 8. Wait for background memory creation
1897
- yield from wait_for_thread_tasks_stream(
1898
- run_response=run_response,
1899
- memory_future=memory_future,
1900
- stream_events=stream_events,
1901
- events_to_skip=self.events_to_skip, # type: ignore
1902
- store_events=self.store_events,
1903
- )
1904
-
1905
- raise_if_cancelled(run_response.run_id) # type: ignore
1906
- # 9. Create session summary
1907
- if self.session_summary_manager is not None:
1908
- # Upsert the RunOutput to Team Session before creating the session summary
1909
- session.upsert_run(run_response=run_response)
1910
1922
 
1911
- if stream_events:
1912
- yield handle_event( # type: ignore
1913
- create_team_session_summary_started_event(from_run_response=run_response),
1914
- run_response,
1915
- events_to_skip=self.events_to_skip,
1916
- store_events=self.store_events,
1917
- )
1918
- try:
1919
- self.session_summary_manager.create_session_summary(session=session)
1920
- except Exception as e:
1921
- log_warning(f"Error in session summary creation: {str(e)}")
1922
- if stream_events:
1923
- yield handle_event( # type: ignore
1924
- create_team_session_summary_completed_event(
1925
- from_run_response=run_response, session_summary=session.summary
1926
- ),
1927
- run_response,
1928
- events_to_skip=self.events_to_skip,
1929
- store_events=self.store_events,
1930
- )
1931
-
1932
- raise_if_cancelled(run_response.run_id) # type: ignore
1933
- # Create the run completed event
1934
- completed_event = handle_event(
1935
- create_team_run_completed_event(
1936
- from_run_response=run_response,
1937
- ),
1938
- run_response,
1939
- events_to_skip=self.events_to_skip,
1940
- store_events=self.store_events,
1941
- )
1942
-
1943
- # Set the run status to completed
1944
- run_response.status = RunStatus.completed
1945
-
1946
- # 10. Cleanup and store the run response
1947
- self._cleanup_and_store(run_response=run_response, session=session)
1923
+ raise_if_cancelled(run_response.run_id) # type: ignore
1924
+ # Create the run completed event
1925
+ completed_event = handle_event(
1926
+ create_team_run_completed_event(
1927
+ from_run_response=run_response,
1928
+ ),
1929
+ run_response,
1930
+ events_to_skip=self.events_to_skip,
1931
+ store_events=self.store_events,
1932
+ )
1948
1933
 
1949
- if stream_events:
1950
- yield completed_event
1934
+ # Set the run status to completed
1935
+ run_response.status = RunStatus.completed
1951
1936
 
1952
- if yield_run_output:
1953
- yield run_response
1937
+ # 10. Cleanup and store the run response
1938
+ self._cleanup_and_store(run_response=run_response, session=session)
1954
1939
 
1955
- # Log Team Telemetry
1956
- self._log_team_telemetry(session_id=session.session_id, run_id=run_response.run_id)
1940
+ if stream_events:
1941
+ yield completed_event
1957
1942
 
1958
- log_debug(f"Team Run End: {run_response.run_id}", center=True, symbol="*")
1943
+ if yield_run_output:
1944
+ yield run_response
1959
1945
 
1960
- except RunCancelledException as e:
1961
- # Handle run cancellation during streaming
1962
- log_info(f"Team run {run_response.run_id} was cancelled during streaming")
1963
- run_response.status = RunStatus.cancelled
1964
- run_response.content = str(e)
1946
+ # Log Team Telemetry
1947
+ self._log_team_telemetry(session_id=session.session_id, run_id=run_response.run_id)
1965
1948
 
1966
- # Yield the cancellation event
1967
- yield handle_event( # type: ignore
1968
- create_team_run_cancelled_event(from_run_response=run_response, reason=str(e)),
1969
- run_response,
1970
- events_to_skip=self.events_to_skip,
1971
- store_events=self.store_events,
1972
- )
1949
+ log_debug(f"Team Run End: {run_response.run_id}", center=True, symbol="*")
1973
1950
 
1974
- # Add the RunOutput to Team Session even when cancelled
1975
- self._cleanup_and_store(run_response=run_response, session=session)
1976
- finally:
1977
- # Always disconnect connectable tools
1978
- self._disconnect_connectable_tools()
1979
- # Always clean up the run tracking
1980
- cleanup_run(run_response.run_id) # type: ignore
1951
+ # Disconnect tools and clean up run tracking
1952
+ self._disconnect_connectable_tools()
1953
+ cleanup_run(run_response.run_id) # type: ignore
1981
1954
 
1982
1955
  @overload
1983
1956
  def run(
@@ -2002,7 +1975,7 @@ class Team:
2002
1975
  dependencies: Optional[Dict[str, Any]] = None,
2003
1976
  metadata: Optional[Dict[str, Any]] = None,
2004
1977
  debug_mode: Optional[bool] = None,
2005
- output_schema: Optional[Type[BaseModel]] = None,
1978
+ output_schema: Optional[Union[Type[BaseModel], Dict[str, Any]]] = None,
2006
1979
  **kwargs: Any,
2007
1980
  ) -> TeamRunOutput: ...
2008
1981
 
@@ -2032,7 +2005,7 @@ class Team:
2032
2005
  debug_mode: Optional[bool] = None,
2033
2006
  yield_run_response: Optional[bool] = None, # To be deprecated: use yield_run_output instead
2034
2007
  yield_run_output: bool = False,
2035
- output_schema: Optional[Type[BaseModel]] = None,
2008
+ output_schema: Optional[Union[Type[BaseModel], Dict[str, Any]]] = None,
2036
2009
  **kwargs: Any,
2037
2010
  ) -> Iterator[Union[RunOutputEvent, TeamRunOutputEvent]]: ...
2038
2011
 
@@ -2061,16 +2034,15 @@ class Team:
2061
2034
  debug_mode: Optional[bool] = None,
2062
2035
  yield_run_response: Optional[bool] = None, # To be deprecated: use yield_run_output instead
2063
2036
  yield_run_output: bool = False,
2064
- output_schema: Optional[Type[BaseModel]] = None,
2037
+ output_schema: Optional[Union[Type[BaseModel], Dict[str, Any]]] = None,
2065
2038
  **kwargs: Any,
2066
2039
  ) -> Union[TeamRunOutput, Iterator[Union[RunOutputEvent, TeamRunOutputEvent]]]:
2067
2040
  """Run the Team and return the response."""
2068
2041
  if self._has_async_db():
2069
2042
  raise Exception("run() is not supported with an async DB. Please use arun() instead.")
2070
2043
 
2071
- # Set the id for the run and register it immediately for cancellation tracking
2044
+ # Set the id for the run
2072
2045
  run_id = run_id or str(uuid4())
2073
- register_run(run_id)
2074
2046
 
2075
2047
  # Initialize Team
2076
2048
  self.initialize_team(debug_mode=debug_mode)
@@ -2088,148 +2060,155 @@ class Team:
2088
2060
  )
2089
2061
  yield_run_output = yield_run_output or yield_run_response # For backwards compatibility
2090
2062
 
2091
- background_tasks = kwargs.pop("background_tasks", None)
2092
- if background_tasks is not None:
2093
- from fastapi import BackgroundTasks
2094
-
2095
- background_tasks: BackgroundTasks = background_tasks # type: ignore
2096
-
2097
- # Validate input against input_schema if provided
2098
- validated_input = self._validate_input(input)
2063
+ # Set up retry logic
2064
+ num_attempts = self.retries + 1
2065
+ for attempt in range(num_attempts):
2066
+ if num_attempts > 1:
2067
+ log_debug(f"Retrying Team run {run_id}. Attempt {attempt + 1} of {num_attempts}...")
2099
2068
 
2100
- # Normalise hook & guardails
2101
- if not self._hooks_normalised:
2102
- if self.pre_hooks:
2103
- self.pre_hooks = normalize_pre_hooks(self.pre_hooks) # type: ignore
2104
- if self.post_hooks:
2105
- self.post_hooks = normalize_post_hooks(self.post_hooks) # type: ignore
2106
- self._hooks_normalised = True
2069
+ try:
2070
+ # Register run for cancellation tracking
2071
+ register_run(run_id) # type: ignore
2107
2072
 
2108
- session_id, user_id = self._initialize_session(session_id=session_id, user_id=user_id)
2073
+ background_tasks = kwargs.pop("background_tasks", None)
2074
+ if background_tasks is not None:
2075
+ from fastapi import BackgroundTasks
2109
2076
 
2110
- image_artifacts, video_artifacts, audio_artifacts, file_artifacts = validate_media_object_id(
2111
- images=images, videos=videos, audios=audio, files=files
2112
- )
2077
+ background_tasks: BackgroundTasks = background_tasks # type: ignore
2113
2078
 
2114
- # Create RunInput to capture the original user input
2115
- run_input = TeamRunInput(
2116
- input_content=validated_input,
2117
- images=image_artifacts,
2118
- videos=video_artifacts,
2119
- audios=audio_artifacts,
2120
- files=file_artifacts,
2121
- )
2079
+ # Validate input against input_schema if provided
2080
+ validated_input = self._validate_input(input)
2122
2081
 
2123
- # Read existing session from database
2124
- team_session = self._read_or_create_session(session_id=session_id, user_id=user_id)
2125
- self._update_metadata(session=team_session)
2082
+ # Normalise hook & guardails
2083
+ if not self._hooks_normalised:
2084
+ if self.pre_hooks:
2085
+ self.pre_hooks = normalize_pre_hooks(self.pre_hooks) # type: ignore
2086
+ if self.post_hooks:
2087
+ self.post_hooks = normalize_post_hooks(self.post_hooks) # type: ignore
2088
+ self._hooks_normalised = True
2126
2089
 
2127
- # Initialize session state
2128
- session_state = self._initialize_session_state(
2129
- session_state=session_state if session_state is not None else {},
2130
- user_id=user_id,
2131
- session_id=session_id,
2132
- run_id=run_id,
2133
- )
2134
- # Update session state from DB
2135
- session_state = self._load_session_state(session=team_session, session_state=session_state)
2090
+ session_id, user_id = self._initialize_session(session_id=session_id, user_id=user_id)
2136
2091
 
2137
- # Determine runtime dependencies
2138
- dependencies = dependencies if dependencies is not None else self.dependencies
2092
+ image_artifacts, video_artifacts, audio_artifacts, file_artifacts = validate_media_object_id(
2093
+ images=images, videos=videos, audios=audio, files=files
2094
+ )
2139
2095
 
2140
- # Resolve output_schema parameter takes precedence, then fall back to self.output_schema
2141
- if output_schema is None:
2142
- output_schema = self.output_schema
2096
+ # Create RunInput to capture the original user input
2097
+ run_input = TeamRunInput(
2098
+ input_content=validated_input,
2099
+ images=image_artifacts,
2100
+ videos=video_artifacts,
2101
+ audios=audio_artifacts,
2102
+ files=file_artifacts,
2103
+ )
2143
2104
 
2144
- # Initialize run context
2145
- run_context = run_context or RunContext(
2146
- run_id=run_id,
2147
- session_id=session_id,
2148
- user_id=user_id,
2149
- session_state=session_state,
2150
- dependencies=dependencies,
2151
- output_schema=output_schema,
2152
- )
2153
- # output_schema parameter takes priority, even if run_context was provided
2154
- run_context.output_schema = output_schema
2105
+ # Read existing session from database
2106
+ team_session = self._read_or_create_session(session_id=session_id, user_id=user_id)
2107
+ self._update_metadata(session=team_session)
2155
2108
 
2156
- # Resolve callable dependencies if present
2157
- if run_context.dependencies is not None:
2158
- self._resolve_run_dependencies(run_context=run_context)
2109
+ # Initialize session state
2110
+ session_state = self._initialize_session_state(
2111
+ session_state=session_state if session_state is not None else {},
2112
+ user_id=user_id,
2113
+ session_id=session_id,
2114
+ run_id=run_id,
2115
+ )
2116
+ # Update session state from DB
2117
+ session_state = self._load_session_state(session=team_session, session_state=session_state)
2159
2118
 
2160
- # Determine runtime context parameters
2161
- add_dependencies = (
2162
- add_dependencies_to_context if add_dependencies_to_context is not None else self.add_dependencies_to_context
2163
- )
2164
- add_session_state = (
2165
- add_session_state_to_context
2166
- if add_session_state_to_context is not None
2167
- else self.add_session_state_to_context
2168
- )
2169
- add_history = add_history_to_context if add_history_to_context is not None else self.add_history_to_context
2119
+ # Determine runtime dependencies
2120
+ dependencies = dependencies if dependencies is not None else self.dependencies
2170
2121
 
2171
- # When filters are passed manually
2172
- if self.knowledge_filters or knowledge_filters:
2173
- run_context.knowledge_filters = self._get_effective_filters(knowledge_filters)
2122
+ # Resolve output_schema parameter takes precedence, then fall back to self.output_schema
2123
+ if output_schema is None:
2124
+ output_schema = self.output_schema
2174
2125
 
2175
- # Use stream override value when necessary
2176
- if stream is None:
2177
- stream = False if self.stream is None else self.stream
2126
+ # Initialize run context
2127
+ run_context = run_context or RunContext(
2128
+ run_id=run_id,
2129
+ session_id=session_id,
2130
+ user_id=user_id,
2131
+ session_state=session_state,
2132
+ dependencies=dependencies,
2133
+ output_schema=output_schema,
2134
+ )
2135
+ # output_schema parameter takes priority, even if run_context was provided
2136
+ run_context.output_schema = output_schema
2137
+
2138
+ # Resolve callable dependencies if present
2139
+ if run_context.dependencies is not None:
2140
+ self._resolve_run_dependencies(run_context=run_context)
2141
+
2142
+ # Determine runtime context parameters
2143
+ add_dependencies = (
2144
+ add_dependencies_to_context
2145
+ if add_dependencies_to_context is not None
2146
+ else self.add_dependencies_to_context
2147
+ )
2148
+ add_session_state = (
2149
+ add_session_state_to_context
2150
+ if add_session_state_to_context is not None
2151
+ else self.add_session_state_to_context
2152
+ )
2153
+ add_history = (
2154
+ add_history_to_context if add_history_to_context is not None else self.add_history_to_context
2155
+ )
2178
2156
 
2179
- # Considering both stream_events and stream_intermediate_steps (deprecated)
2180
- stream_events = stream_events or stream_intermediate_steps
2157
+ # When filters are passed manually
2158
+ if self.knowledge_filters or knowledge_filters:
2159
+ run_context.knowledge_filters = self._get_effective_filters(knowledge_filters)
2181
2160
 
2182
- # Can't stream events if streaming is disabled
2183
- if stream is False:
2184
- stream_events = False
2161
+ # Use stream override value when necessary
2162
+ if stream is None:
2163
+ stream = False if self.stream is None else self.stream
2185
2164
 
2186
- if stream_events is None:
2187
- stream_events = False if self.stream_events is None else self.stream_events
2165
+ # Considering both stream_events and stream_intermediate_steps (deprecated)
2166
+ stream_events = stream_events or stream_intermediate_steps
2188
2167
 
2189
- self.model = cast(Model, self.model)
2168
+ # Can't stream events if streaming is disabled
2169
+ if stream is False:
2170
+ stream_events = False
2190
2171
 
2191
- if self.metadata is not None:
2192
- if metadata is None:
2193
- metadata = self.metadata
2194
- else:
2195
- merge_dictionaries(metadata, self.metadata)
2172
+ if stream_events is None:
2173
+ stream_events = False if self.stream_events is None else self.stream_events
2196
2174
 
2197
- if metadata:
2198
- run_context.metadata = metadata
2175
+ self.model = cast(Model, self.model)
2199
2176
 
2200
- # Configure the model for runs
2201
- response_format: Optional[Union[Dict, Type[BaseModel]]] = (
2202
- self._get_response_format(run_context=run_context) if self.parser_model is None else None
2203
- )
2177
+ if self.metadata is not None:
2178
+ if metadata is None:
2179
+ metadata = self.metadata
2180
+ else:
2181
+ merge_dictionaries(metadata, self.metadata)
2204
2182
 
2205
- # Create a new run_response for this attempt
2206
- run_response = TeamRunOutput(
2207
- run_id=run_id,
2208
- session_id=session_id,
2209
- user_id=user_id,
2210
- team_id=self.id,
2211
- team_name=self.name,
2212
- metadata=run_context.metadata,
2213
- session_state=run_context.session_state,
2214
- input=run_input,
2215
- )
2183
+ if metadata:
2184
+ run_context.metadata = metadata
2216
2185
 
2217
- run_response.model = self.model.id if self.model is not None else None
2218
- run_response.model_provider = self.model.provider if self.model is not None else None
2186
+ # Configure the model for runs
2187
+ response_format: Optional[Union[Dict, Type[BaseModel]]] = (
2188
+ self._get_response_format(run_context=run_context) if self.parser_model is None else None
2189
+ )
2219
2190
 
2220
- # Start the run metrics timer, to calculate the run duration
2221
- run_response.metrics = Metrics()
2222
- run_response.metrics.start_timer()
2191
+ # Create a new run_response for this attempt
2192
+ run_response = TeamRunOutput(
2193
+ run_id=run_id,
2194
+ session_id=session_id,
2195
+ user_id=user_id,
2196
+ team_id=self.id,
2197
+ team_name=self.name,
2198
+ metadata=run_context.metadata,
2199
+ session_state=run_context.session_state,
2200
+ input=run_input,
2201
+ )
2223
2202
 
2224
- # Set up retry logic
2225
- num_attempts = self.retries + 1
2203
+ run_response.model = self.model.id if self.model is not None else None
2204
+ run_response.model_provider = self.model.provider if self.model is not None else None
2205
+
2206
+ # Start the run metrics timer, to calculate the run duration
2207
+ run_response.metrics = Metrics()
2208
+ run_response.metrics.start_timer()
2226
2209
 
2227
- for attempt in range(num_attempts):
2228
- log_debug(f"Retrying Team run {run_id}. Attempt {attempt + 1} of {num_attempts}...")
2229
- # Run the team
2230
- try:
2231
2210
  if stream:
2232
- response_iterator = self._run_stream(
2211
+ return self._run_stream(
2233
2212
  run_response=run_response,
2234
2213
  run_context=run_context,
2235
2214
  session=team_session,
@@ -2243,9 +2222,8 @@ class Team:
2243
2222
  debug_mode=debug_mode,
2244
2223
  background_tasks=background_tasks,
2245
2224
  **kwargs,
2246
- )
2225
+ ) # type: ignore
2247
2226
 
2248
- return response_iterator # type: ignore
2249
2227
  else:
2250
2228
  return self._run(
2251
2229
  run_response=run_response,
@@ -2260,24 +2238,67 @@ class Team:
2260
2238
  background_tasks=background_tasks,
2261
2239
  **kwargs,
2262
2240
  )
2241
+ except InputCheckError as e:
2242
+ run_response.status = RunStatus.error
2243
+ if stream:
2244
+ run_error = create_team_run_error_event(
2245
+ run_response,
2246
+ error=str(e),
2247
+ error_id=e.error_id,
2248
+ error_type=e.type,
2249
+ additional_data=e.additional_data,
2250
+ )
2251
+ run_response.events = add_team_error_event(error=run_error, events=run_response.events)
2252
+ if run_response.content is None:
2253
+ run_response.content = str(e)
2263
2254
 
2264
- except (InputCheckError, OutputCheckError) as e:
2265
2255
  log_error(f"Validation failed: {str(e)} | Check: {e.check_trigger}")
2266
- raise e
2267
- except KeyboardInterrupt:
2268
- run_response.content = "Operation cancelled by user"
2256
+
2257
+ if stream:
2258
+ return generator_wrapper(run_error) # type: ignore
2259
+ else:
2260
+ return run_response
2261
+ except RunCancelledException as e:
2262
+ # Handle run cancellation during streaming
2263
+ log_info(f"Team run {run_response.run_id} was cancelled during streaming")
2269
2264
  run_response.status = RunStatus.cancelled
2265
+ run_response.content = str(e)
2270
2266
 
2267
+ # Yield the cancellation event
2271
2268
  if stream:
2272
- return generator_wrapper( # type: ignore
2273
- create_team_run_cancelled_event(
2274
- from_run_response=run_response, reason="Operation cancelled by user"
2275
- )
2269
+ cancelled_run_error = handle_event(
2270
+ create_team_run_cancelled_event(from_run_response=run_response, reason=str(e)),
2271
+ run_response,
2272
+ events_to_skip=self.events_to_skip,
2273
+ store_events=self.store_events,
2276
2274
  )
2275
+ return generator_wrapper(cancelled_run_error) # type: ignore
2276
+ else:
2277
+ return run_response
2278
+ except (InputCheckError, OutputCheckError) as e:
2279
+ run_response.status = RunStatus.error
2280
+
2281
+ if stream:
2282
+ # Add error event to list of events
2283
+ run_error = create_team_run_error_event(
2284
+ run_response,
2285
+ error=str(e),
2286
+ error_id=e.error_id,
2287
+ error_type=e.type,
2288
+ additional_data=e.additional_data,
2289
+ )
2290
+ run_response.events = add_team_error_event(error=run_error, events=run_response.events)
2291
+
2292
+ if run_response.content is None:
2293
+ run_response.content = str(e)
2294
+
2295
+ log_error(f"Validation failed: {str(e)} | Check: {e.check_trigger}")
2296
+
2297
+ if stream:
2298
+ return generator_wrapper(run_error) # type: ignore
2277
2299
  else:
2278
2300
  return run_response
2279
2301
  except Exception as e:
2280
- # Check if this is the last attempt
2281
2302
  if attempt < num_attempts - 1:
2282
2303
  # Calculate delay with exponential backoff if enabled
2283
2304
  if self.exponential_backoff:
@@ -2287,12 +2308,23 @@ class Team:
2287
2308
 
2288
2309
  log_warning(f"Attempt {attempt + 1}/{num_attempts} failed: {str(e)}. Retrying in {delay}s...")
2289
2310
  time.sleep(delay)
2311
+ continue
2312
+
2313
+ run_response.status = RunStatus.error
2314
+ if stream:
2315
+ run_error = create_team_run_error_event(run_response, error=str(e))
2316
+ run_response.events = add_team_error_event(error=run_error, events=run_response.events)
2317
+ if run_response.content is None:
2318
+ run_response.content = str(e)
2319
+
2320
+ log_error(f"Error in Team run: {str(e)}")
2321
+
2322
+ if stream:
2323
+ return generator_wrapper(run_error) # type: ignore
2290
2324
  else:
2291
- # Final attempt failed - re-raise the exception
2292
- log_error(f"All {num_attempts} attempts failed. Final error: {str(e)}")
2293
- raise e
2325
+ return run_response
2294
2326
 
2295
- # If we get here, all retries failed
2327
+ # If we get here, all retries failed (shouldn't happen with current logic)
2296
2328
  raise Exception(f"Failed after {num_attempts} attempts.")
2297
2329
 
2298
2330
  async def _arun(
@@ -2329,219 +2361,276 @@ class Team:
2329
2361
  15. Cleanup and store (scrub, add to session, calculate metrics, save session)
2330
2362
  """
2331
2363
  log_debug(f"Team Run Start: {run_response.run_id}", center=True)
2364
+ memory_task = None
2332
2365
 
2333
- if run_context.dependencies is not None:
2334
- await self._aresolve_run_dependencies(run_context=run_context)
2366
+ # Set up retry logic
2367
+ num_attempts = self.retries + 1
2368
+ for attempt in range(num_attempts):
2369
+ if num_attempts > 1:
2370
+ log_debug(f"Retrying Team run {run_response.run_id}. Attempt {attempt + 1} of {num_attempts}...")
2335
2371
 
2336
- # 1. Read or create session. Reads from the database if provided.
2337
- if self._has_async_db():
2338
- team_session = await self._aread_or_create_session(session_id=session_id, user_id=user_id)
2339
- else:
2340
- team_session = self._read_or_create_session(session_id=session_id, user_id=user_id)
2372
+ try:
2373
+ if run_context.dependencies is not None:
2374
+ await self._aresolve_run_dependencies(run_context=run_context)
2341
2375
 
2342
- # 2. Update metadata and session state
2343
- self._update_metadata(session=team_session)
2344
- # Initialize session state
2345
- run_context.session_state = self._initialize_session_state(
2346
- session_state=run_context.session_state if run_context.session_state is not None else {},
2347
- user_id=user_id,
2348
- session_id=session_id,
2349
- run_id=run_response.run_id,
2350
- )
2351
- # Update session state from DB
2352
- if run_context.session_state is not None:
2353
- run_context.session_state = self._load_session_state(
2354
- session=team_session, session_state=run_context.session_state
2355
- )
2376
+ # 1. Read or create session. Reads from the database if provided.
2377
+ if self._has_async_db():
2378
+ team_session = await self._aread_or_create_session(session_id=session_id, user_id=user_id)
2379
+ else:
2380
+ team_session = self._read_or_create_session(session_id=session_id, user_id=user_id)
2356
2381
 
2357
- run_input = cast(TeamRunInput, run_response.input)
2382
+ # 2. Update metadata and session state
2383
+ self._update_metadata(session=team_session)
2384
+ # Initialize session state
2385
+ run_context.session_state = self._initialize_session_state(
2386
+ session_state=run_context.session_state if run_context.session_state is not None else {},
2387
+ user_id=user_id,
2388
+ session_id=session_id,
2389
+ run_id=run_response.run_id,
2390
+ )
2391
+ # Update session state from DB
2392
+ if run_context.session_state is not None:
2393
+ run_context.session_state = self._load_session_state(
2394
+ session=team_session, session_state=run_context.session_state
2395
+ )
2358
2396
 
2359
- # 3. Execute pre-hooks after session is loaded but before processing starts
2360
- if self.pre_hooks is not None:
2361
- pre_hook_iterator = self._aexecute_pre_hooks(
2362
- hooks=self.pre_hooks, # type: ignore
2363
- run_response=run_response,
2364
- run_context=run_context,
2365
- run_input=run_input,
2366
- session=team_session,
2367
- user_id=user_id,
2368
- debug_mode=debug_mode,
2369
- background_tasks=background_tasks,
2370
- **kwargs,
2371
- )
2397
+ run_input = cast(TeamRunInput, run_response.input)
2372
2398
 
2373
- # Consume the async iterator without yielding
2374
- async for _ in pre_hook_iterator:
2375
- pass
2399
+ # 3. Execute pre-hooks after session is loaded but before processing starts
2400
+ if self.pre_hooks is not None:
2401
+ pre_hook_iterator = self._aexecute_pre_hooks(
2402
+ hooks=self.pre_hooks, # type: ignore
2403
+ run_response=run_response,
2404
+ run_context=run_context,
2405
+ run_input=run_input,
2406
+ session=team_session,
2407
+ user_id=user_id,
2408
+ debug_mode=debug_mode,
2409
+ background_tasks=background_tasks,
2410
+ **kwargs,
2411
+ )
2376
2412
 
2377
- # 4. Determine tools for model
2378
- team_run_context: Dict[str, Any] = {}
2379
- self.model = cast(Model, self.model)
2380
- await self._check_and_refresh_mcp_tools()
2381
- _tools = self._determine_tools_for_model(
2382
- model=self.model,
2383
- run_response=run_response,
2384
- run_context=run_context,
2385
- team_run_context=team_run_context,
2386
- session=team_session,
2387
- user_id=user_id,
2388
- async_mode=True,
2389
- input_message=run_input.input_content,
2390
- images=run_input.images,
2391
- videos=run_input.videos,
2392
- audio=run_input.audios,
2393
- files=run_input.files,
2394
- debug_mode=debug_mode,
2395
- add_history_to_context=add_history_to_context,
2396
- add_dependencies_to_context=add_dependencies_to_context,
2397
- add_session_state_to_context=add_session_state_to_context,
2398
- stream=False,
2399
- stream_events=False,
2400
- )
2413
+ # Consume the async iterator without yielding
2414
+ async for _ in pre_hook_iterator:
2415
+ pass
2401
2416
 
2402
- # 5. Prepare run messages
2403
- run_messages = await self._aget_run_messages(
2404
- run_response=run_response,
2405
- run_context=run_context,
2406
- session=team_session, # type: ignore
2407
- user_id=user_id,
2408
- input_message=run_input.input_content,
2409
- audio=run_input.audios,
2410
- images=run_input.images,
2411
- videos=run_input.videos,
2412
- files=run_input.files,
2413
- add_history_to_context=add_history_to_context,
2414
- add_dependencies_to_context=add_dependencies_to_context,
2415
- add_session_state_to_context=add_session_state_to_context,
2416
- tools=_tools,
2417
- **kwargs,
2418
- )
2417
+ # 4. Determine tools for model
2418
+ team_run_context: Dict[str, Any] = {}
2419
+ self.model = cast(Model, self.model)
2420
+ await self._check_and_refresh_mcp_tools()
2421
+ _tools = self._determine_tools_for_model(
2422
+ model=self.model,
2423
+ run_response=run_response,
2424
+ run_context=run_context,
2425
+ team_run_context=team_run_context,
2426
+ session=team_session,
2427
+ user_id=user_id,
2428
+ async_mode=True,
2429
+ input_message=run_input.input_content,
2430
+ images=run_input.images,
2431
+ videos=run_input.videos,
2432
+ audio=run_input.audios,
2433
+ files=run_input.files,
2434
+ debug_mode=debug_mode,
2435
+ add_history_to_context=add_history_to_context,
2436
+ add_dependencies_to_context=add_dependencies_to_context,
2437
+ add_session_state_to_context=add_session_state_to_context,
2438
+ stream=False,
2439
+ stream_events=False,
2440
+ )
2419
2441
 
2420
- self.model = cast(Model, self.model)
2421
- log_debug(f"Team Run Start: {run_response.run_id}", center=True)
2442
+ # 5. Prepare run messages
2443
+ run_messages = await self._aget_run_messages(
2444
+ run_response=run_response,
2445
+ run_context=run_context,
2446
+ session=team_session, # type: ignore
2447
+ user_id=user_id,
2448
+ input_message=run_input.input_content,
2449
+ audio=run_input.audios,
2450
+ images=run_input.images,
2451
+ videos=run_input.videos,
2452
+ files=run_input.files,
2453
+ add_history_to_context=add_history_to_context,
2454
+ add_dependencies_to_context=add_dependencies_to_context,
2455
+ add_session_state_to_context=add_session_state_to_context,
2456
+ tools=_tools,
2457
+ **kwargs,
2458
+ )
2422
2459
 
2423
- # 6. Start memory creation in background task
2424
- memory_task = None
2425
- if (
2426
- run_messages.user_message is not None
2427
- and self.memory_manager is not None
2428
- and self.enable_user_memories
2429
- and not self.enable_agentic_memory
2430
- ):
2431
- log_debug("Starting memory creation in background task.")
2432
- memory_task = asyncio.create_task(self._amake_memories(run_messages=run_messages, user_id=user_id))
2460
+ self.model = cast(Model, self.model)
2461
+ log_debug(f"Team Run Start: {run_response.run_id}", center=True)
2433
2462
 
2434
- try:
2435
- raise_if_cancelled(run_response.run_id) # type: ignore
2436
- # 7. Reason about the task if reasoning is enabled
2437
- await self._ahandle_reasoning(run_response=run_response, run_messages=run_messages)
2463
+ # 6. Start memory creation in background task
2464
+ memory_task = None
2465
+ if (
2466
+ run_messages.user_message is not None
2467
+ and self.memory_manager is not None
2468
+ and self.enable_user_memories
2469
+ and not self.enable_agentic_memory
2470
+ ):
2471
+ log_debug("Starting memory creation in background task.")
2472
+ memory_task = asyncio.create_task(self._amake_memories(run_messages=run_messages, user_id=user_id))
2438
2473
 
2439
- # Check for cancellation before model call
2440
- raise_if_cancelled(run_response.run_id) # type: ignore
2474
+ raise_if_cancelled(run_response.run_id) # type: ignore
2475
+ # 7. Reason about the task if reasoning is enabled
2476
+ await self._ahandle_reasoning(run_response=run_response, run_messages=run_messages)
2441
2477
 
2442
- # 8. Get the model response for the team leader
2443
- model_response = await self.model.aresponse(
2444
- messages=run_messages.messages,
2445
- tools=_tools,
2446
- tool_choice=self.tool_choice,
2447
- tool_call_limit=self.tool_call_limit,
2448
- response_format=response_format,
2449
- send_media_to_model=self.send_media_to_model,
2450
- run_response=run_response,
2451
- compression_manager=self.compression_manager if self.compress_tool_results else None,
2452
- ) # type: ignore
2478
+ # Check for cancellation before model call
2479
+ raise_if_cancelled(run_response.run_id) # type: ignore
2453
2480
 
2454
- # Check for cancellation after model call
2455
- raise_if_cancelled(run_response.run_id) # type: ignore
2481
+ # 8. Get the model response for the team leader
2482
+ model_response = await self.model.aresponse(
2483
+ messages=run_messages.messages,
2484
+ tools=_tools,
2485
+ tool_choice=self.tool_choice,
2486
+ tool_call_limit=self.tool_call_limit,
2487
+ response_format=response_format,
2488
+ send_media_to_model=self.send_media_to_model,
2489
+ run_response=run_response,
2490
+ compression_manager=self.compression_manager if self.compress_tool_results else None,
2491
+ ) # type: ignore
2456
2492
 
2457
- # If an output model is provided, generate output using the output model
2458
- await self._agenerate_response_with_output_model(model_response=model_response, run_messages=run_messages)
2493
+ # Check for cancellation after model call
2494
+ raise_if_cancelled(run_response.run_id) # type: ignore
2459
2495
 
2460
- # If a parser model is provided, structure the response separately
2461
- await self._aparse_response_with_parser_model(
2462
- model_response=model_response, run_messages=run_messages, run_context=run_context
2463
- )
2496
+ # If an output model is provided, generate output using the output model
2497
+ await self._agenerate_response_with_output_model(
2498
+ model_response=model_response, run_messages=run_messages
2499
+ )
2464
2500
 
2465
- # 9. Update TeamRunOutput with the model response
2466
- self._update_run_response(
2467
- model_response=model_response,
2468
- run_response=run_response,
2469
- run_messages=run_messages,
2470
- run_context=run_context,
2471
- )
2501
+ # If a parser model is provided, structure the response separately
2502
+ await self._aparse_response_with_parser_model(
2503
+ model_response=model_response, run_messages=run_messages, run_context=run_context
2504
+ )
2472
2505
 
2473
- # 10. Store media if enabled
2474
- if self.store_media:
2475
- store_media_util(run_response, model_response)
2506
+ # 9. Update TeamRunOutput with the model response
2507
+ self._update_run_response(
2508
+ model_response=model_response,
2509
+ run_response=run_response,
2510
+ run_messages=run_messages,
2511
+ run_context=run_context,
2512
+ )
2476
2513
 
2477
- # 11. Convert response to structured format
2478
- self._convert_response_to_structured_format(run_response=run_response, run_context=run_context)
2514
+ # 10. Store media if enabled
2515
+ if self.store_media:
2516
+ store_media_util(run_response, model_response)
2479
2517
 
2480
- # 12. Execute post-hooks after output is generated but before response is returned
2481
- if self.post_hooks is not None:
2482
- async for _ in self._aexecute_post_hooks(
2483
- hooks=self.post_hooks, # type: ignore
2484
- run_output=run_response,
2485
- run_context=run_context,
2486
- session=team_session,
2487
- user_id=user_id,
2488
- debug_mode=debug_mode,
2489
- background_tasks=background_tasks,
2490
- **kwargs,
2491
- ):
2492
- pass
2518
+ # 11. Convert response to structured format
2519
+ self._convert_response_to_structured_format(run_response=run_response, run_context=run_context)
2493
2520
 
2494
- raise_if_cancelled(run_response.run_id) # type: ignore
2521
+ # 12. Execute post-hooks after output is generated but before response is returned
2522
+ if self.post_hooks is not None:
2523
+ async for _ in self._aexecute_post_hooks(
2524
+ hooks=self.post_hooks, # type: ignore
2525
+ run_output=run_response,
2526
+ run_context=run_context,
2527
+ session=team_session,
2528
+ user_id=user_id,
2529
+ debug_mode=debug_mode,
2530
+ background_tasks=background_tasks,
2531
+ **kwargs,
2532
+ ):
2533
+ pass
2495
2534
 
2496
- # 13. Wait for background memory creation
2497
- await await_for_open_threads(memory_task=memory_task)
2535
+ raise_if_cancelled(run_response.run_id) # type: ignore
2498
2536
 
2499
- raise_if_cancelled(run_response.run_id) # type: ignore
2500
- # 14. Create session summary
2501
- if self.session_summary_manager is not None:
2502
- # Upsert the RunOutput to Team Session before creating the session summary
2503
- team_session.upsert_run(run_response=run_response)
2504
- try:
2505
- await self.session_summary_manager.acreate_session_summary(session=team_session)
2506
- except Exception as e:
2507
- log_warning(f"Error in session summary creation: {str(e)}")
2537
+ # 13. Wait for background memory creation
2538
+ await await_for_open_threads(memory_task=memory_task)
2508
2539
 
2509
- raise_if_cancelled(run_response.run_id) # type: ignore
2510
- run_response.status = RunStatus.completed
2540
+ raise_if_cancelled(run_response.run_id) # type: ignore
2541
+ # 14. Create session summary
2542
+ if self.session_summary_manager is not None:
2543
+ # Upsert the RunOutput to Team Session before creating the session summary
2544
+ team_session.upsert_run(run_response=run_response)
2545
+ try:
2546
+ await self.session_summary_manager.acreate_session_summary(session=team_session)
2547
+ except Exception as e:
2548
+ log_warning(f"Error in session summary creation: {str(e)}")
2511
2549
 
2512
- # 15. Cleanup and store the run response and session
2513
- await self._acleanup_and_store(run_response=run_response, session=team_session)
2550
+ raise_if_cancelled(run_response.run_id) # type: ignore
2551
+ run_response.status = RunStatus.completed
2514
2552
 
2515
- # Log Team Telemetry
2516
- await self._alog_team_telemetry(session_id=team_session.session_id, run_id=run_response.run_id)
2553
+ # 15. Cleanup and store the run response and session
2554
+ await self._acleanup_and_store(run_response=run_response, session=team_session)
2517
2555
 
2518
- log_debug(f"Team Run End: {run_response.run_id}", center=True, symbol="*")
2556
+ # Log Team Telemetry
2557
+ await self._alog_team_telemetry(session_id=team_session.session_id, run_id=run_response.run_id)
2519
2558
 
2520
- return run_response
2521
- except RunCancelledException as e:
2522
- # Handle run cancellation
2523
- log_info(f"Run {run_response.run_id} was cancelled")
2524
- run_response.content = str(e)
2525
- run_response.status = RunStatus.cancelled
2559
+ log_debug(f"Team Run End: {run_response.run_id}", center=True, symbol="*")
2526
2560
 
2527
- # Cleanup and store the run response and session
2528
- await self._acleanup_and_store(run_response=run_response, session=team_session)
2561
+ return run_response
2529
2562
 
2530
- return run_response
2531
- finally:
2532
- # Always disconnect connectable tools
2533
- self._disconnect_connectable_tools()
2534
- await self._disconnect_mcp_tools()
2535
- # Cancel the memory task if it's still running
2536
- if memory_task is not None and not memory_task.done():
2537
- memory_task.cancel()
2538
- try:
2539
- await memory_task
2540
- except asyncio.CancelledError:
2541
- pass
2563
+ except RunCancelledException as e:
2564
+ # Handle run cancellation
2565
+ log_info(f"Run {run_response.run_id} was cancelled")
2566
+ run_response.content = str(e)
2567
+ run_response.status = RunStatus.cancelled
2568
+
2569
+ # Cleanup and store the run response and session
2570
+ await self._acleanup_and_store(run_response=run_response, session=team_session)
2571
+
2572
+ return run_response
2573
+
2574
+ except (InputCheckError, OutputCheckError) as e:
2575
+ run_response.status = RunStatus.error
2576
+ run_error = create_team_run_error_event(
2577
+ run_response,
2578
+ error=str(e),
2579
+ error_id=e.error_id,
2580
+ error_type=e.type,
2581
+ additional_data=e.additional_data,
2582
+ )
2583
+ run_response.events = add_team_error_event(error=run_error, events=run_response.events)
2584
+ if run_response.content is None:
2585
+ run_response.content = str(e)
2542
2586
 
2543
- # Always clean up the run tracking
2544
- cleanup_run(run_response.run_id) # type: ignore
2587
+ log_error(f"Validation failed: {str(e)} | Check: {e.check_trigger}")
2588
+
2589
+ await self._acleanup_and_store(run_response=run_response, session=team_session)
2590
+
2591
+ return run_response
2592
+
2593
+ except Exception as e:
2594
+ if attempt < num_attempts - 1:
2595
+ # Calculate delay with exponential backoff if enabled
2596
+ if self.exponential_backoff:
2597
+ delay = self.delay_between_retries * (2**attempt)
2598
+ else:
2599
+ delay = self.delay_between_retries
2600
+
2601
+ log_warning(f"Attempt {attempt + 1}/{num_attempts} failed: {str(e)}. Retrying in {delay}s...")
2602
+ time.sleep(delay)
2603
+ continue
2604
+
2605
+ run_error = create_team_run_error_event(run_response, error=str(e))
2606
+ run_response.events = add_team_error_event(error=run_error, events=run_response.events)
2607
+
2608
+ if run_response.content is None:
2609
+ run_response.content = str(e)
2610
+
2611
+ log_error(f"Error in Team run: {str(e)}")
2612
+
2613
+ # Cleanup and store the run response and session
2614
+ await self._acleanup_and_store(run_response=run_response, session=team_session)
2615
+
2616
+ return run_response
2617
+
2618
+ finally:
2619
+ # Always disconnect connectable tools
2620
+ self._disconnect_connectable_tools()
2621
+ await self._disconnect_mcp_tools()
2622
+ # Cancel the memory task if it's still running
2623
+ if memory_task is not None and not memory_task.done():
2624
+ memory_task.cancel()
2625
+ try:
2626
+ await memory_task
2627
+ except asyncio.CancelledError:
2628
+ pass
2629
+
2630
+ # Always clean up the run tracking
2631
+ cleanup_run(run_response.run_id) # type: ignore
2632
+
2633
+ return run_response
2545
2634
 
2546
2635
  async def _arun_stream(
2547
2636
  self,
@@ -2578,307 +2667,364 @@ class Team:
2578
2667
  13. Cleanup and store (scrub, add to session, calculate metrics, save session)
2579
2668
  """
2580
2669
 
2581
- # 1. Resolve dependencies
2582
- if run_context.dependencies is not None:
2583
- await self._aresolve_run_dependencies(run_context=run_context)
2584
-
2585
- # 2. Read or create session. Reads from the database if provided.
2586
- if self._has_async_db():
2587
- team_session = await self._aread_or_create_session(session_id=session_id, user_id=user_id)
2588
- else:
2589
- team_session = self._read_or_create_session(session_id=session_id, user_id=user_id)
2590
-
2591
- # 3. Update metadata and session state
2592
- self._update_metadata(session=team_session)
2593
- # Initialize session state
2594
- run_context.session_state = self._initialize_session_state(
2595
- session_state=run_context.session_state if run_context.session_state is not None else {},
2596
- user_id=user_id,
2597
- session_id=session_id,
2598
- run_id=run_response.run_id,
2599
- )
2600
- # Update session state from DB
2601
- if run_context.session_state is not None:
2602
- run_context.session_state = self._load_session_state(
2603
- session=team_session, session_state=run_context.session_state
2604
- ) # type: ignore
2670
+ memory_task = None
2605
2671
 
2606
- # 4. Execute pre-hooks
2607
- run_input = cast(TeamRunInput, run_response.input)
2608
- self.model = cast(Model, self.model)
2609
- if self.pre_hooks is not None:
2610
- pre_hook_iterator = self._aexecute_pre_hooks(
2611
- hooks=self.pre_hooks, # type: ignore
2612
- run_response=run_response,
2613
- run_context=run_context,
2614
- run_input=run_input,
2615
- session=team_session,
2616
- user_id=user_id,
2617
- debug_mode=debug_mode,
2618
- stream_events=stream_events,
2619
- background_tasks=background_tasks,
2620
- **kwargs,
2621
- )
2622
- async for pre_hook_event in pre_hook_iterator:
2623
- yield pre_hook_event
2672
+ # Set up retry logic
2673
+ num_attempts = self.retries + 1
2674
+ for attempt in range(num_attempts):
2675
+ if num_attempts > 1:
2676
+ log_debug(f"Retrying Team run {run_response.run_id}. Attempt {attempt + 1} of {num_attempts}...")
2624
2677
 
2625
- # 5. Determine tools for model
2626
- team_run_context: Dict[str, Any] = {}
2627
- self.model = cast(Model, self.model)
2628
- await self._check_and_refresh_mcp_tools()
2629
- _tools = self._determine_tools_for_model(
2630
- model=self.model,
2631
- run_response=run_response,
2632
- run_context=run_context,
2633
- team_run_context=team_run_context,
2634
- session=team_session, # type: ignore
2635
- user_id=user_id,
2636
- async_mode=True,
2637
- input_message=run_input.input_content,
2638
- images=run_input.images,
2639
- videos=run_input.videos,
2640
- audio=run_input.audios,
2641
- files=run_input.files,
2642
- debug_mode=debug_mode,
2643
- add_history_to_context=add_history_to_context,
2644
- add_dependencies_to_context=add_dependencies_to_context,
2645
- add_session_state_to_context=add_session_state_to_context,
2646
- stream=True,
2647
- stream_events=stream_events,
2648
- )
2678
+ try:
2679
+ # 1. Resolve dependencies
2680
+ if run_context.dependencies is not None:
2681
+ await self._aresolve_run_dependencies(run_context=run_context)
2649
2682
 
2650
- # 6. Prepare run messages
2651
- run_messages = await self._aget_run_messages(
2652
- run_response=run_response,
2653
- run_context=run_context,
2654
- session=team_session, # type: ignore
2655
- user_id=user_id,
2656
- input_message=run_input.input_content,
2657
- audio=run_input.audios,
2658
- images=run_input.images,
2659
- videos=run_input.videos,
2660
- files=run_input.files,
2661
- add_history_to_context=add_history_to_context,
2662
- add_dependencies_to_context=add_dependencies_to_context,
2663
- add_session_state_to_context=add_session_state_to_context,
2664
- tools=_tools,
2665
- **kwargs,
2666
- )
2683
+ # 2. Read or create session. Reads from the database if provided.
2684
+ if self._has_async_db():
2685
+ team_session = await self._aread_or_create_session(session_id=session_id, user_id=user_id)
2686
+ else:
2687
+ team_session = self._read_or_create_session(session_id=session_id, user_id=user_id)
2667
2688
 
2668
- log_debug(f"Team Run Start: {run_response.run_id}", center=True)
2689
+ # 3. Update metadata and session state
2690
+ self._update_metadata(session=team_session)
2691
+ # Initialize session state
2692
+ run_context.session_state = self._initialize_session_state(
2693
+ session_state=run_context.session_state if run_context.session_state is not None else {},
2694
+ user_id=user_id,
2695
+ session_id=session_id,
2696
+ run_id=run_response.run_id,
2697
+ )
2698
+ # Update session state from DB
2699
+ if run_context.session_state is not None:
2700
+ run_context.session_state = self._load_session_state(
2701
+ session=team_session, session_state=run_context.session_state
2702
+ ) # type: ignore
2703
+
2704
+ # 4. Execute pre-hooks
2705
+ run_input = cast(TeamRunInput, run_response.input)
2706
+ self.model = cast(Model, self.model)
2707
+ if self.pre_hooks is not None:
2708
+ pre_hook_iterator = self._aexecute_pre_hooks(
2709
+ hooks=self.pre_hooks, # type: ignore
2710
+ run_response=run_response,
2711
+ run_context=run_context,
2712
+ run_input=run_input,
2713
+ session=team_session,
2714
+ user_id=user_id,
2715
+ debug_mode=debug_mode,
2716
+ stream_events=stream_events,
2717
+ background_tasks=background_tasks,
2718
+ **kwargs,
2719
+ )
2720
+ async for pre_hook_event in pre_hook_iterator:
2721
+ yield pre_hook_event
2722
+
2723
+ # 5. Determine tools for model
2724
+ team_run_context: Dict[str, Any] = {}
2725
+ self.model = cast(Model, self.model)
2726
+ await self._check_and_refresh_mcp_tools()
2727
+ _tools = self._determine_tools_for_model(
2728
+ model=self.model,
2729
+ run_response=run_response,
2730
+ run_context=run_context,
2731
+ team_run_context=team_run_context,
2732
+ session=team_session, # type: ignore
2733
+ user_id=user_id,
2734
+ async_mode=True,
2735
+ input_message=run_input.input_content,
2736
+ images=run_input.images,
2737
+ videos=run_input.videos,
2738
+ audio=run_input.audios,
2739
+ files=run_input.files,
2740
+ debug_mode=debug_mode,
2741
+ add_history_to_context=add_history_to_context,
2742
+ add_dependencies_to_context=add_dependencies_to_context,
2743
+ add_session_state_to_context=add_session_state_to_context,
2744
+ stream=True,
2745
+ stream_events=stream_events,
2746
+ )
2669
2747
 
2670
- # 7. Start memory creation in background task
2671
- memory_task = None
2672
- if (
2673
- run_messages.user_message is not None
2674
- and self.memory_manager is not None
2675
- and self.enable_user_memories
2676
- and not self.enable_agentic_memory
2677
- ):
2678
- log_debug("Starting memory creation in background task.")
2679
- memory_task = asyncio.create_task(self._amake_memories(run_messages=run_messages, user_id=user_id))
2748
+ # 6. Prepare run messages
2749
+ run_messages = await self._aget_run_messages(
2750
+ run_response=run_response,
2751
+ run_context=run_context,
2752
+ session=team_session, # type: ignore
2753
+ user_id=user_id,
2754
+ input_message=run_input.input_content,
2755
+ audio=run_input.audios,
2756
+ images=run_input.images,
2757
+ videos=run_input.videos,
2758
+ files=run_input.files,
2759
+ add_history_to_context=add_history_to_context,
2760
+ add_dependencies_to_context=add_dependencies_to_context,
2761
+ add_session_state_to_context=add_session_state_to_context,
2762
+ tools=_tools,
2763
+ **kwargs,
2764
+ )
2680
2765
 
2681
- try:
2682
- # Considering both stream_events and stream_intermediate_steps (deprecated)
2683
- stream_events = stream_events or stream_intermediate_steps
2766
+ log_debug(f"Team Run Start: {run_response.run_id}", center=True)
2684
2767
 
2685
- # Yield the run started event
2686
- if stream_events:
2687
- yield handle_event( # type: ignore
2688
- create_team_run_started_event(from_run_response=run_response),
2689
- run_response,
2690
- events_to_skip=self.events_to_skip,
2691
- store_events=self.store_events,
2692
- )
2768
+ # 7. Start memory creation in background task
2769
+ memory_task = None
2770
+ if (
2771
+ run_messages.user_message is not None
2772
+ and self.memory_manager is not None
2773
+ and self.enable_user_memories
2774
+ and not self.enable_agentic_memory
2775
+ ):
2776
+ log_debug("Starting memory creation in background task.")
2777
+ memory_task = asyncio.create_task(self._amake_memories(run_messages=run_messages, user_id=user_id))
2693
2778
 
2694
- # 8. Reason about the task if reasoning is enabled
2695
- async for item in self._ahandle_reasoning_stream(
2696
- run_response=run_response,
2697
- run_messages=run_messages,
2698
- stream_events=stream_events,
2699
- ):
2700
- raise_if_cancelled(run_response.run_id) # type: ignore
2701
- yield item
2779
+ # Considering both stream_events and stream_intermediate_steps (deprecated)
2780
+ stream_events = stream_events or stream_intermediate_steps
2702
2781
 
2703
- # Check for cancellation before model processing
2704
- raise_if_cancelled(run_response.run_id) # type: ignore
2782
+ # Yield the run started event
2783
+ if stream_events:
2784
+ yield handle_event( # type: ignore
2785
+ create_team_run_started_event(from_run_response=run_response),
2786
+ run_response,
2787
+ events_to_skip=self.events_to_skip,
2788
+ store_events=self.store_events,
2789
+ )
2705
2790
 
2706
- # 9. Get a response from the model
2707
- if self.output_model is None:
2708
- async for event in self._ahandle_model_response_stream(
2709
- session=team_session,
2791
+ # 8. Reason about the task if reasoning is enabled
2792
+ async for item in self._ahandle_reasoning_stream(
2710
2793
  run_response=run_response,
2711
2794
  run_messages=run_messages,
2712
- tools=_tools,
2713
- response_format=response_format,
2714
2795
  stream_events=stream_events,
2715
- session_state=run_context.session_state,
2716
- run_context=run_context,
2717
2796
  ):
2718
2797
  raise_if_cancelled(run_response.run_id) # type: ignore
2719
- yield event
2720
- else:
2721
- async for event in self._ahandle_model_response_stream(
2798
+ yield item
2799
+
2800
+ # Check for cancellation before model processing
2801
+ raise_if_cancelled(run_response.run_id) # type: ignore
2802
+
2803
+ # 9. Get a response from the model
2804
+ if self.output_model is None:
2805
+ async for event in self._ahandle_model_response_stream(
2806
+ session=team_session,
2807
+ run_response=run_response,
2808
+ run_messages=run_messages,
2809
+ tools=_tools,
2810
+ response_format=response_format,
2811
+ stream_events=stream_events,
2812
+ session_state=run_context.session_state,
2813
+ run_context=run_context,
2814
+ ):
2815
+ raise_if_cancelled(run_response.run_id) # type: ignore
2816
+ yield event
2817
+ else:
2818
+ async for event in self._ahandle_model_response_stream(
2819
+ session=team_session,
2820
+ run_response=run_response,
2821
+ run_messages=run_messages,
2822
+ tools=_tools,
2823
+ response_format=response_format,
2824
+ stream_events=stream_events,
2825
+ session_state=run_context.session_state,
2826
+ run_context=run_context,
2827
+ ):
2828
+ raise_if_cancelled(run_response.run_id) # type: ignore
2829
+ from agno.run.team import IntermediateRunContentEvent, RunContentEvent
2830
+
2831
+ if isinstance(event, RunContentEvent):
2832
+ if stream_events:
2833
+ yield IntermediateRunContentEvent(
2834
+ content=event.content,
2835
+ content_type=event.content_type,
2836
+ )
2837
+ else:
2838
+ yield event
2839
+
2840
+ async for event in self._agenerate_response_with_output_model_stream(
2841
+ session=team_session,
2842
+ run_response=run_response,
2843
+ run_messages=run_messages,
2844
+ stream_events=stream_events,
2845
+ ):
2846
+ raise_if_cancelled(run_response.run_id) # type: ignore
2847
+ yield event
2848
+
2849
+ # Check for cancellation after model processing
2850
+ raise_if_cancelled(run_response.run_id) # type: ignore
2851
+
2852
+ # 10. Parse response with parser model if provided
2853
+ async for event in self._aparse_response_with_parser_model_stream(
2722
2854
  session=team_session,
2723
2855
  run_response=run_response,
2724
- run_messages=run_messages,
2725
- tools=_tools,
2726
- response_format=response_format,
2727
2856
  stream_events=stream_events,
2728
- session_state=run_context.session_state,
2729
2857
  run_context=run_context,
2730
2858
  ):
2731
- raise_if_cancelled(run_response.run_id) # type: ignore
2732
- from agno.run.team import IntermediateRunContentEvent, RunContentEvent
2859
+ yield event
2733
2860
 
2734
- if isinstance(event, RunContentEvent):
2735
- if stream_events:
2736
- yield IntermediateRunContentEvent(
2737
- content=event.content,
2738
- content_type=event.content_type,
2739
- )
2740
- else:
2861
+ # Yield RunContentCompletedEvent
2862
+ if stream_events:
2863
+ yield handle_event( # type: ignore
2864
+ create_team_run_content_completed_event(from_run_response=run_response),
2865
+ run_response,
2866
+ events_to_skip=self.events_to_skip,
2867
+ store_events=self.store_events,
2868
+ )
2869
+
2870
+ # Execute post-hooks after output is generated but before response is returned
2871
+ if self.post_hooks is not None:
2872
+ async for event in self._aexecute_post_hooks(
2873
+ hooks=self.post_hooks, # type: ignore
2874
+ run_output=run_response,
2875
+ run_context=run_context,
2876
+ session=team_session,
2877
+ user_id=user_id,
2878
+ debug_mode=debug_mode,
2879
+ stream_events=stream_events,
2880
+ background_tasks=background_tasks,
2881
+ **kwargs,
2882
+ ):
2741
2883
  yield event
2742
2884
 
2743
- async for event in self._agenerate_response_with_output_model_stream(
2744
- session=team_session,
2885
+ raise_if_cancelled(run_response.run_id) # type: ignore
2886
+ # 11. Wait for background memory creation
2887
+ async for event in await_for_thread_tasks_stream(
2745
2888
  run_response=run_response,
2746
- run_messages=run_messages,
2889
+ memory_task=memory_task,
2747
2890
  stream_events=stream_events,
2891
+ events_to_skip=self.events_to_skip, # type: ignore
2892
+ store_events=self.store_events,
2748
2893
  ):
2749
- raise_if_cancelled(run_response.run_id) # type: ignore
2750
2894
  yield event
2751
2895
 
2752
- # Check for cancellation after model processing
2753
- raise_if_cancelled(run_response.run_id) # type: ignore
2896
+ raise_if_cancelled(run_response.run_id) # type: ignore
2754
2897
 
2755
- # 10. Parse response with parser model if provided
2756
- async for event in self._aparse_response_with_parser_model_stream(
2757
- session=team_session, run_response=run_response, stream_events=stream_events, run_context=run_context
2758
- ):
2759
- yield event
2898
+ # 12. Create session summary
2899
+ if self.session_summary_manager is not None:
2900
+ # Upsert the RunOutput to Team Session before creating the session summary
2901
+ team_session.upsert_run(run_response=run_response)
2760
2902
 
2761
- # Yield RunContentCompletedEvent
2762
- if stream_events:
2763
- yield handle_event( # type: ignore
2764
- create_team_run_content_completed_event(from_run_response=run_response),
2903
+ if stream_events:
2904
+ yield handle_event( # type: ignore
2905
+ create_team_session_summary_started_event(from_run_response=run_response),
2906
+ run_response,
2907
+ events_to_skip=self.events_to_skip,
2908
+ store_events=self.store_events,
2909
+ )
2910
+ try:
2911
+ await self.session_summary_manager.acreate_session_summary(session=team_session)
2912
+ except Exception as e:
2913
+ log_warning(f"Error in session summary creation: {str(e)}")
2914
+ if stream_events:
2915
+ yield handle_event( # type: ignore
2916
+ create_team_session_summary_completed_event(
2917
+ from_run_response=run_response, session_summary=team_session.summary
2918
+ ),
2919
+ run_response,
2920
+ events_to_skip=self.events_to_skip,
2921
+ store_events=self.store_events,
2922
+ )
2923
+
2924
+ raise_if_cancelled(run_response.run_id) # type: ignore
2925
+
2926
+ # Create the run completed event
2927
+ completed_event = handle_event(
2928
+ create_team_run_completed_event(from_run_response=run_response),
2765
2929
  run_response,
2766
2930
  events_to_skip=self.events_to_skip,
2767
2931
  store_events=self.store_events,
2768
2932
  )
2769
2933
 
2770
- # Execute post-hooks after output is generated but before response is returned
2771
- if self.post_hooks is not None:
2772
- async for event in self._aexecute_post_hooks(
2773
- hooks=self.post_hooks, # type: ignore
2774
- run_output=run_response,
2775
- run_context=run_context,
2776
- session=team_session,
2777
- user_id=user_id,
2778
- debug_mode=debug_mode,
2779
- stream_events=stream_events,
2780
- background_tasks=background_tasks,
2781
- **kwargs,
2782
- ):
2783
- yield event
2934
+ # Set the run status to completed
2935
+ run_response.status = RunStatus.completed
2784
2936
 
2785
- raise_if_cancelled(run_response.run_id) # type: ignore
2786
- # 11. Wait for background memory creation
2787
- async for event in await_for_thread_tasks_stream(
2788
- run_response=run_response,
2789
- memory_task=memory_task,
2790
- stream_events=stream_events,
2791
- events_to_skip=self.events_to_skip, # type: ignore
2792
- store_events=self.store_events,
2793
- ):
2794
- yield event
2937
+ # 13. Cleanup and store the run response and session
2938
+ await self._acleanup_and_store(run_response=run_response, session=team_session)
2795
2939
 
2796
- raise_if_cancelled(run_response.run_id) # type: ignore
2940
+ if stream_events:
2941
+ yield completed_event
2797
2942
 
2798
- # 12. Create session summary
2799
- if self.session_summary_manager is not None:
2800
- # Upsert the RunOutput to Team Session before creating the session summary
2801
- team_session.upsert_run(run_response=run_response)
2943
+ if yield_run_output:
2944
+ yield run_response
2802
2945
 
2803
- if stream_events:
2804
- yield handle_event( # type: ignore
2805
- create_team_session_summary_started_event(from_run_response=run_response),
2806
- run_response,
2807
- events_to_skip=self.events_to_skip,
2808
- store_events=self.store_events,
2809
- )
2810
- try:
2811
- await self.session_summary_manager.acreate_session_summary(session=team_session)
2812
- except Exception as e:
2813
- log_warning(f"Error in session summary creation: {str(e)}")
2814
- if stream_events:
2815
- yield handle_event( # type: ignore
2816
- create_team_session_summary_completed_event(
2817
- from_run_response=run_response, session_summary=team_session.summary
2818
- ),
2819
- run_response,
2820
- events_to_skip=self.events_to_skip,
2821
- store_events=self.store_events,
2822
- )
2946
+ # Log Team Telemetry
2947
+ await self._alog_team_telemetry(session_id=team_session.session_id, run_id=run_response.run_id)
2823
2948
 
2824
- raise_if_cancelled(run_response.run_id) # type: ignore
2949
+ log_debug(f"Team Run End: {run_response.run_id}", center=True, symbol="*")
2825
2950
 
2826
- # Create the run completed event
2827
- completed_event = handle_event(
2828
- create_team_run_completed_event(from_run_response=run_response),
2829
- run_response,
2830
- events_to_skip=self.events_to_skip,
2831
- store_events=self.store_events,
2832
- )
2951
+ except RunCancelledException as e:
2952
+ # Handle run cancellation during async streaming
2953
+ log_info(f"Team run {run_response.run_id} was cancelled during async streaming")
2954
+ run_response.status = RunStatus.cancelled
2955
+ run_response.content = str(e)
2833
2956
 
2834
- # Set the run status to completed
2835
- run_response.status = RunStatus.completed
2957
+ # Yield the cancellation event
2958
+ yield handle_event( # type: ignore
2959
+ create_team_run_cancelled_event(from_run_response=run_response, reason=str(e)),
2960
+ run_response,
2961
+ events_to_skip=self.events_to_skip,
2962
+ store_events=self.store_events,
2963
+ )
2836
2964
 
2837
- # 13. Cleanup and store the run response and session
2838
- await self._acleanup_and_store(run_response=run_response, session=team_session)
2965
+ # Cleanup and store the run response and session
2966
+ await self._acleanup_and_store(run_response=run_response, session=team_session)
2839
2967
 
2840
- if stream_events:
2841
- yield completed_event
2968
+ except (InputCheckError, OutputCheckError) as e:
2969
+ run_response.status = RunStatus.error
2970
+ run_error = create_team_run_error_event(
2971
+ run_response,
2972
+ error=str(e),
2973
+ error_id=e.error_id,
2974
+ error_type=e.type,
2975
+ additional_data=e.additional_data,
2976
+ )
2977
+ run_response.events = add_team_error_event(error=run_error, events=run_response.events)
2978
+ if run_response.content is None:
2979
+ run_response.content = str(e)
2842
2980
 
2843
- if yield_run_output:
2844
- yield run_response
2981
+ log_error(f"Validation failed: {str(e)} | Check: {e.check_trigger}")
2845
2982
 
2846
- # Log Team Telemetry
2847
- await self._alog_team_telemetry(session_id=team_session.session_id, run_id=run_response.run_id)
2983
+ await self._acleanup_and_store(run_response=run_response, session=team_session)
2848
2984
 
2849
- log_debug(f"Team Run End: {run_response.run_id}", center=True, symbol="*")
2985
+ yield run_error
2850
2986
 
2851
- except RunCancelledException as e:
2852
- # Handle run cancellation during async streaming
2853
- log_info(f"Team run {run_response.run_id} was cancelled during async streaming")
2854
- run_response.status = RunStatus.cancelled
2855
- run_response.content = str(e)
2987
+ break
2856
2988
 
2857
- # Yield the cancellation event
2858
- yield handle_event( # type: ignore
2859
- create_team_run_cancelled_event(from_run_response=run_response, reason=str(e)),
2860
- run_response,
2861
- events_to_skip=self.events_to_skip,
2862
- store_events=self.store_events,
2863
- )
2989
+ except Exception as e:
2990
+ if attempt < num_attempts - 1:
2991
+ # Calculate delay with exponential backoff if enabled
2992
+ if self.exponential_backoff:
2993
+ delay = self.delay_between_retries * (2**attempt)
2994
+ else:
2995
+ delay = self.delay_between_retries
2996
+
2997
+ log_warning(f"Attempt {attempt + 1}/{num_attempts} failed: {str(e)}. Retrying in {delay}s...")
2998
+ time.sleep(delay)
2999
+ continue
2864
3000
 
2865
- # Cleanup and store the run response and session
2866
- await self._acleanup_and_store(run_response=run_response, session=team_session)
3001
+ run_response.status = RunStatus.error
3002
+ run_error = create_team_run_error_event(run_response, error=str(e))
3003
+ run_response.events = add_team_error_event(error=run_error, events=run_response.events)
3004
+ if run_response.content is None:
3005
+ run_response.content = str(e)
2867
3006
 
2868
- finally:
2869
- # Always disconnect connectable tools
2870
- self._disconnect_connectable_tools()
2871
- await self._disconnect_mcp_tools()
2872
- # Cancel the memory task if it's still running
2873
- if memory_task is not None and not memory_task.done():
2874
- memory_task.cancel()
2875
- try:
2876
- await memory_task
2877
- except asyncio.CancelledError:
2878
- pass
3007
+ log_error(f"Error in Team run: {str(e)}")
3008
+
3009
+ # Cleanup and store the run response and session
3010
+ await self._acleanup_and_store(run_response=run_response, session=team_session)
3011
+
3012
+ yield run_error
3013
+
3014
+ finally:
3015
+ # Always disconnect connectable tools
3016
+ self._disconnect_connectable_tools()
3017
+ await self._disconnect_mcp_tools()
3018
+ # Cancel the memory task if it's still running
3019
+ if memory_task is not None and not memory_task.done():
3020
+ memory_task.cancel()
3021
+ try:
3022
+ await memory_task
3023
+ except asyncio.CancelledError:
3024
+ pass
2879
3025
 
2880
- # Always clean up the run tracking
2881
- cleanup_run(run_response.run_id) # type: ignore
3026
+ # Always clean up the run tracking
3027
+ cleanup_run(run_response.run_id) # type: ignore
2882
3028
 
2883
3029
  @overload
2884
3030
  async def arun(
@@ -2904,7 +3050,7 @@ class Team:
2904
3050
  dependencies: Optional[Dict[str, Any]] = None,
2905
3051
  metadata: Optional[Dict[str, Any]] = None,
2906
3052
  debug_mode: Optional[bool] = None,
2907
- output_schema: Optional[Type[BaseModel]] = None,
3053
+ output_schema: Optional[Union[Type[BaseModel], Dict[str, Any]]] = None,
2908
3054
  **kwargs: Any,
2909
3055
  ) -> TeamRunOutput: ...
2910
3056
 
@@ -2934,7 +3080,7 @@ class Team:
2934
3080
  debug_mode: Optional[bool] = None,
2935
3081
  yield_run_response: Optional[bool] = None, # To be deprecated: use yield_run_output instead
2936
3082
  yield_run_output: bool = False,
2937
- output_schema: Optional[Type[BaseModel]] = None,
3083
+ output_schema: Optional[Union[Type[BaseModel], Dict[str, Any]]] = None,
2938
3084
  **kwargs: Any,
2939
3085
  ) -> AsyncIterator[Union[RunOutputEvent, TeamRunOutputEvent]]: ...
2940
3086
 
@@ -2963,7 +3109,7 @@ class Team:
2963
3109
  debug_mode: Optional[bool] = None,
2964
3110
  yield_run_response: Optional[bool] = None, # To be deprecated: use yield_run_output instead
2965
3111
  yield_run_output: bool = False,
2966
- output_schema: Optional[Type[BaseModel]] = None,
3112
+ output_schema: Optional[Union[Type[BaseModel], Dict[str, Any]]] = None,
2967
3113
  **kwargs: Any,
2968
3114
  ) -> Union[TeamRunOutput, AsyncIterator[Union[RunOutputEvent, TeamRunOutputEvent]]]:
2969
3115
  """Run the Team asynchronously and return the response."""
@@ -3110,79 +3256,38 @@ class Team:
3110
3256
 
3111
3257
  yield_run_output = bool(yield_run_output or yield_run_response) # For backwards compatibility
3112
3258
 
3113
- # Resolve retry parameters
3114
- num_attempts = self.retries + 1
3115
-
3116
- for attempt in range(num_attempts):
3117
- # Run the team
3118
- try:
3119
- if stream:
3120
- return self._arun_stream( # type: ignore
3121
- input=validated_input,
3122
- run_response=run_response,
3123
- run_context=run_context,
3124
- session_id=session_id,
3125
- user_id=user_id,
3126
- add_history_to_context=add_history,
3127
- add_dependencies_to_context=add_dependencies,
3128
- add_session_state_to_context=add_session_state,
3129
- response_format=response_format,
3130
- stream_events=stream_events,
3131
- yield_run_output=yield_run_output,
3132
- debug_mode=debug_mode,
3133
- background_tasks=background_tasks,
3134
- **kwargs,
3135
- )
3136
- else:
3137
- return self._arun( # type: ignore
3138
- input=validated_input,
3139
- run_response=run_response,
3140
- run_context=run_context,
3141
- session_id=session_id,
3142
- user_id=user_id,
3143
- add_history_to_context=add_history,
3144
- add_dependencies_to_context=add_dependencies,
3145
- add_session_state_to_context=add_session_state,
3146
- response_format=response_format,
3147
- debug_mode=debug_mode,
3148
- background_tasks=background_tasks,
3149
- **kwargs,
3150
- )
3151
-
3152
- except (InputCheckError, OutputCheckError) as e:
3153
- log_error(f"Validation failed: {str(e)} | Check: {e.check_trigger}")
3154
- raise e
3155
- except KeyboardInterrupt:
3156
- run_response.content = "Operation cancelled by user"
3157
- run_response.status = RunStatus.cancelled
3158
-
3159
- if stream:
3160
- return async_generator_wrapper(
3161
- create_team_run_cancelled_event(
3162
- from_run_response=run_response, reason="Operation cancelled by user"
3163
- )
3164
- )
3165
- else:
3166
- return run_response
3167
- except Exception as e:
3168
- # Check if this is the last attempt
3169
- if attempt < num_attempts - 1:
3170
- # Calculate delay with exponential backoff if enabled
3171
- if self.exponential_backoff:
3172
- delay = self.delay_between_retries * (2**attempt)
3173
- else:
3174
- delay = self.delay_between_retries
3175
-
3176
- log_warning(f"Attempt {attempt + 1}/{num_attempts} failed: {str(e)}. Retrying in {delay}s...")
3177
- time.sleep(delay)
3178
- continue
3179
- else:
3180
- # Final attempt failed - re-raise the exception
3181
- log_error(f"All {num_attempts} attempts failed. Final error: {str(e)}")
3182
- raise e
3183
-
3184
- # If we get here, all retries failed
3185
- raise Exception(f"Failed after {num_attempts} attempts.")
3259
+ if stream:
3260
+ return self._arun_stream( # type: ignore
3261
+ input=validated_input,
3262
+ run_response=run_response,
3263
+ run_context=run_context,
3264
+ session_id=session_id,
3265
+ user_id=user_id,
3266
+ add_history_to_context=add_history,
3267
+ add_dependencies_to_context=add_dependencies,
3268
+ add_session_state_to_context=add_session_state,
3269
+ response_format=response_format,
3270
+ stream_events=stream_events,
3271
+ yield_run_output=yield_run_output,
3272
+ debug_mode=debug_mode,
3273
+ background_tasks=background_tasks,
3274
+ **kwargs,
3275
+ )
3276
+ else:
3277
+ return self._arun( # type: ignore
3278
+ input=validated_input,
3279
+ run_response=run_response,
3280
+ run_context=run_context,
3281
+ session_id=session_id,
3282
+ user_id=user_id,
3283
+ add_history_to_context=add_history,
3284
+ add_dependencies_to_context=add_dependencies,
3285
+ add_session_state_to_context=add_session_state,
3286
+ response_format=response_format,
3287
+ debug_mode=debug_mode,
3288
+ background_tasks=background_tasks,
3289
+ **kwargs,
3290
+ )
3186
3291
 
3187
3292
  def _update_run_response(
3188
3293
  self,
@@ -3199,7 +3304,7 @@ class Team:
3199
3304
  # Update the run_response content with the structured output
3200
3305
  run_response.content = model_response.parsed
3201
3306
  # Update the run_response content_type with the structured output class name
3202
- run_response.content_type = output_schema.__name__
3307
+ run_response.content_type = "dict" if isinstance(output_schema, dict) else output_schema.__name__
3203
3308
  else:
3204
3309
  # Update the run_response content with the model response content
3205
3310
  if not run_response.content:
@@ -3498,12 +3603,16 @@ class Team:
3498
3603
  self._convert_response_to_structured_format(full_model_response, run_context=run_context)
3499
3604
  # Get output_schema from run_context
3500
3605
  output_schema = run_context.output_schema if run_context else None
3501
- content_type = output_schema.__name__ # type: ignore
3606
+ content_type = "dict" if isinstance(output_schema, dict) else output_schema.__name__ # type: ignore
3502
3607
  run_response.content_type = content_type
3503
3608
  elif self._member_response_model is not None:
3504
3609
  full_model_response.content = model_response_event.content
3505
3610
  self._convert_response_to_structured_format(full_model_response, run_context=run_context)
3506
- content_type = self._member_response_model.__name__ # type: ignore
3611
+ content_type = (
3612
+ "dict"
3613
+ if isinstance(self._member_response_model, dict)
3614
+ else self._member_response_model.__name__
3615
+ ) # type: ignore
3507
3616
  run_response.content_type = content_type
3508
3617
  elif isinstance(model_response_event.content, str):
3509
3618
  full_model_response.content = (full_model_response.content or "") + model_response_event.content
@@ -3723,6 +3832,15 @@ class Team:
3723
3832
  events_to_skip=self.events_to_skip,
3724
3833
  store_events=self.store_events,
3725
3834
  )
3835
+ if tool_call.tool_call_error:
3836
+ yield handle_event( # type: ignore
3837
+ create_team_tool_call_error_event(
3838
+ from_run_response=run_response, tool=tool_call, error=str(tool_call.result)
3839
+ ),
3840
+ run_response,
3841
+ events_to_skip=self.events_to_skip,
3842
+ store_events=self.store_events,
3843
+ )
3726
3844
 
3727
3845
  if stream_events:
3728
3846
  if reasoning_step is not None:
@@ -3755,41 +3873,71 @@ class Team:
3755
3873
  output_schema = run_context.output_schema if run_context else None
3756
3874
 
3757
3875
  # Convert the response to the structured format if needed
3758
- if output_schema is not None and not isinstance(run_response.content, output_schema):
3759
- if isinstance(run_response.content, str) and self.parse_response:
3760
- try:
3761
- parsed_response_content = parse_response_model_str(run_response.content, output_schema)
3762
-
3763
- # Update TeamRunOutput
3764
- if parsed_response_content is not None:
3765
- run_response.content = parsed_response_content
3876
+ if output_schema is not None:
3877
+ # If the output schema is a dict, do not convert it into a BaseModel
3878
+ if isinstance(output_schema, dict):
3879
+ if isinstance(run_response.content, dict):
3880
+ # Content is already a dict - just set content_type
3881
+ if hasattr(run_response, "content_type"):
3882
+ run_response.content_type = "dict"
3883
+ elif isinstance(run_response.content, str):
3884
+ parsed_dict = parse_response_dict_str(run_response.content)
3885
+ if parsed_dict is not None:
3886
+ run_response.content = parsed_dict
3766
3887
  if hasattr(run_response, "content_type"):
3767
- run_response.content_type = output_schema.__name__
3888
+ run_response.content_type = "dict"
3768
3889
  else:
3769
- log_warning("Failed to convert response to output_schema")
3770
- except Exception as e:
3771
- log_warning(f"Failed to convert response to output model: {e}")
3772
- else:
3773
- log_warning("Something went wrong. Team run response content is not a string")
3774
- elif self._member_response_model is not None and not isinstance(
3775
- run_response.content, self._member_response_model
3776
- ):
3777
- if isinstance(run_response.content, str):
3778
- try:
3779
- parsed_response_content = parse_response_model_str(
3780
- run_response.content, self._member_response_model
3781
- )
3782
- # Update TeamRunOutput
3783
- if parsed_response_content is not None:
3784
- run_response.content = parsed_response_content
3890
+ log_warning("Failed to parse JSON response")
3891
+ # If the output schema is a Pydantic model and parse_response is True, parse it into a BaseModel
3892
+ elif not isinstance(run_response.content, output_schema):
3893
+ if isinstance(run_response.content, str) and self.parse_response:
3894
+ try:
3895
+ parsed_response_content = parse_response_model_str(run_response.content, output_schema)
3896
+
3897
+ # Update TeamRunOutput
3898
+ if parsed_response_content is not None:
3899
+ run_response.content = parsed_response_content
3900
+ if hasattr(run_response, "content_type"):
3901
+ run_response.content_type = output_schema.__name__
3902
+ else:
3903
+ log_warning("Failed to convert response to output_schema")
3904
+ except Exception as e:
3905
+ log_warning(f"Failed to convert response to output model: {e}")
3906
+ else:
3907
+ log_warning("Something went wrong. Team run response content is not a string")
3908
+ elif self._member_response_model is not None:
3909
+ # Handle dict schema from member
3910
+ if isinstance(self._member_response_model, dict):
3911
+ if isinstance(run_response.content, dict):
3912
+ # Content is already a dict - just set content_type
3913
+ if hasattr(run_response, "content_type"):
3914
+ run_response.content_type = "dict"
3915
+ elif isinstance(run_response.content, str):
3916
+ parsed_dict = parse_response_dict_str(run_response.content)
3917
+ if parsed_dict is not None:
3918
+ run_response.content = parsed_dict
3785
3919
  if hasattr(run_response, "content_type"):
3786
- run_response.content_type = self._member_response_model.__name__
3920
+ run_response.content_type = "dict"
3787
3921
  else:
3788
- log_warning("Failed to convert response to output_schema")
3789
- except Exception as e:
3790
- log_warning(f"Failed to convert response to output model: {e}")
3791
- else:
3792
- log_warning("Something went wrong. Member run response content is not a string")
3922
+ log_warning("Failed to parse JSON response")
3923
+ # Handle Pydantic schema from member
3924
+ elif not isinstance(run_response.content, self._member_response_model):
3925
+ if isinstance(run_response.content, str):
3926
+ try:
3927
+ parsed_response_content = parse_response_model_str(
3928
+ run_response.content, self._member_response_model
3929
+ )
3930
+ # Update TeamRunOutput
3931
+ if parsed_response_content is not None:
3932
+ run_response.content = parsed_response_content
3933
+ if hasattr(run_response, "content_type"):
3934
+ run_response.content_type = self._member_response_model.__name__
3935
+ else:
3936
+ log_warning("Failed to convert response to output_schema")
3937
+ except Exception as e:
3938
+ log_warning(f"Failed to convert response to output model: {e}")
3939
+ else:
3940
+ log_warning("Something went wrong. Member run response content is not a string")
3793
3941
 
3794
3942
  def _cleanup_and_store(self, run_response: TeamRunOutput, session: TeamSession) -> None:
3795
3943
  # Scrub the stored run based on storage flags
@@ -3892,6 +4040,10 @@ class Team:
3892
4040
  elif model.supports_json_schema_outputs:
3893
4041
  if self.use_json_mode:
3894
4042
  log_debug("Setting Model.response_format to JSON response mode")
4043
+ # Handle JSON schema - pass through directly (user provides full provider format)
4044
+ if isinstance(output_schema, dict):
4045
+ return output_schema
4046
+ # Handle Pydantic schema
3895
4047
  return {
3896
4048
  "type": "json_schema",
3897
4049
  "json_schema": {
@@ -4750,281 +4902,77 @@ class Team:
4750
4902
 
4751
4903
  return updated_reasoning_content
4752
4904
 
4753
- def _reason(
4905
+ def _handle_reasoning_event(
4754
4906
  self,
4907
+ event: "ReasoningEvent", # type: ignore # noqa: F821
4755
4908
  run_response: TeamRunOutput,
4756
- run_messages: RunMessages,
4757
4909
  stream_events: bool,
4758
4910
  ) -> Iterator[TeamRunOutputEvent]:
4759
- if stream_events:
4760
- yield handle_event( # type: ignore
4761
- create_team_reasoning_started_event(from_run_response=run_response),
4762
- run_response,
4763
- events_to_skip=self.events_to_skip,
4764
- store_events=self.store_events,
4765
- )
4766
-
4767
- use_default_reasoning = False
4768
-
4769
- # Get the reasoning model
4770
- reasoning_model: Optional[Model] = self.reasoning_model
4771
- reasoning_model_provided = reasoning_model is not None
4772
- if reasoning_model is None and self.model is not None:
4773
- from copy import deepcopy
4774
-
4775
- reasoning_model = deepcopy(self.model)
4776
- if reasoning_model is None:
4777
- log_warning("Reasoning error. Reasoning model is None, continuing regular session...")
4778
- return
4779
-
4780
- # If a reasoning model is provided, use it to generate reasoning
4781
- if reasoning_model_provided:
4782
- from agno.reasoning.anthropic import is_anthropic_reasoning_model
4783
- from agno.reasoning.azure_ai_foundry import is_ai_foundry_reasoning_model
4784
- from agno.reasoning.deepseek import is_deepseek_reasoning_model
4785
- from agno.reasoning.gemini import is_gemini_reasoning_model
4786
- from agno.reasoning.groq import is_groq_reasoning_model
4787
- from agno.reasoning.helpers import get_reasoning_agent
4788
- from agno.reasoning.ollama import is_ollama_reasoning_model
4789
- from agno.reasoning.openai import is_openai_reasoning_model
4790
- from agno.reasoning.vertexai import is_vertexai_reasoning_model
4791
-
4792
- reasoning_agent = self.reasoning_agent or get_reasoning_agent(
4793
- reasoning_model=reasoning_model,
4794
- session_state=self.session_state,
4795
- dependencies=self.dependencies,
4796
- metadata=self.metadata,
4797
- )
4798
- is_deepseek = is_deepseek_reasoning_model(reasoning_model)
4799
- is_groq = is_groq_reasoning_model(reasoning_model)
4800
- is_openai = is_openai_reasoning_model(reasoning_model)
4801
- is_ollama = is_ollama_reasoning_model(reasoning_model)
4802
- is_ai_foundry = is_ai_foundry_reasoning_model(reasoning_model)
4803
- is_gemini = is_gemini_reasoning_model(reasoning_model)
4804
- is_anthropic = is_anthropic_reasoning_model(reasoning_model)
4805
- is_vertexai = is_vertexai_reasoning_model(reasoning_model)
4806
-
4807
- if (
4808
- is_deepseek
4809
- or is_groq
4810
- or is_openai
4811
- or is_ollama
4812
- or is_ai_foundry
4813
- or is_gemini
4814
- or is_anthropic
4815
- or is_vertexai
4816
- ):
4817
- reasoning_message: Optional[Message] = None
4818
- if is_deepseek:
4819
- from agno.reasoning.deepseek import get_deepseek_reasoning
4820
-
4821
- log_debug("Starting DeepSeek Reasoning", center=True, symbol="=")
4822
- reasoning_message = get_deepseek_reasoning(
4823
- reasoning_agent=reasoning_agent, messages=run_messages.get_input_messages()
4824
- )
4825
- elif is_groq:
4826
- from agno.reasoning.groq import get_groq_reasoning
4827
-
4828
- log_debug("Starting Groq Reasoning", center=True, symbol="=")
4829
- reasoning_message = get_groq_reasoning(
4830
- reasoning_agent=reasoning_agent, messages=run_messages.get_input_messages()
4831
- )
4832
- elif is_openai:
4833
- from agno.reasoning.openai import get_openai_reasoning
4834
-
4835
- log_debug("Starting OpenAI Reasoning", center=True, symbol="=")
4836
- reasoning_message = get_openai_reasoning(
4837
- reasoning_agent=reasoning_agent, messages=run_messages.get_input_messages()
4838
- )
4839
- elif is_ollama:
4840
- from agno.reasoning.ollama import get_ollama_reasoning
4841
-
4842
- log_debug("Starting Ollama Reasoning", center=True, symbol="=")
4843
- reasoning_message = get_ollama_reasoning(
4844
- reasoning_agent=reasoning_agent, messages=run_messages.get_input_messages()
4845
- )
4846
- elif is_ai_foundry:
4847
- from agno.reasoning.azure_ai_foundry import get_ai_foundry_reasoning
4848
-
4849
- log_debug("Starting Azure AI Foundry Reasoning", center=True, symbol="=")
4850
- reasoning_message = get_ai_foundry_reasoning(
4851
- reasoning_agent=reasoning_agent, messages=run_messages.get_input_messages()
4852
- )
4853
- elif is_gemini:
4854
- from agno.reasoning.gemini import get_gemini_reasoning
4855
-
4856
- log_debug("Starting Gemini Reasoning", center=True, symbol="=")
4857
- reasoning_message = get_gemini_reasoning(
4858
- reasoning_agent=reasoning_agent, messages=run_messages.get_input_messages()
4859
- )
4860
- elif is_anthropic:
4861
- from agno.reasoning.anthropic import get_anthropic_reasoning
4911
+ """
4912
+ Convert a ReasoningEvent from the ReasoningManager to Team-specific TeamRunOutputEvents.
4862
4913
 
4863
- log_debug("Starting Anthropic Claude Reasoning", center=True, symbol="=")
4864
- reasoning_message = get_anthropic_reasoning(
4865
- reasoning_agent=reasoning_agent, messages=run_messages.get_input_messages()
4866
- )
4867
- elif is_vertexai:
4868
- from agno.reasoning.vertexai import get_vertexai_reasoning
4914
+ This method handles the conversion of generic reasoning events to Team events,
4915
+ keeping the Team._reason() method clean and simple.
4916
+ """
4917
+ from agno.reasoning.manager import ReasoningEventType
4869
4918
 
4870
- log_debug("Starting VertexAI Reasoning", center=True, symbol="=")
4871
- reasoning_message = get_vertexai_reasoning(
4872
- reasoning_agent=reasoning_agent, messages=run_messages.get_input_messages()
4873
- )
4919
+ if event.event_type == ReasoningEventType.started:
4920
+ if stream_events:
4921
+ yield handle_event( # type: ignore
4922
+ create_team_reasoning_started_event(from_run_response=run_response),
4923
+ run_response,
4924
+ events_to_skip=self.events_to_skip,
4925
+ store_events=self.store_events,
4926
+ )
4874
4927
 
4875
- if reasoning_message is None:
4876
- log_warning("Reasoning error. Reasoning response is None, continuing regular session...")
4877
- return
4928
+ elif event.event_type == ReasoningEventType.content_delta:
4929
+ if stream_events and event.reasoning_content:
4930
+ yield handle_event( # type: ignore
4931
+ create_team_reasoning_content_delta_event(
4932
+ from_run_response=run_response,
4933
+ reasoning_content=event.reasoning_content,
4934
+ ),
4935
+ run_response,
4936
+ events_to_skip=self.events_to_skip,
4937
+ store_events=self.store_events,
4938
+ )
4878
4939
 
4879
- run_messages.messages.append(reasoning_message)
4880
- # Add reasoning step to the Agent's run_response
4940
+ elif event.event_type == ReasoningEventType.step:
4941
+ if event.reasoning_step:
4942
+ # Update run_response with this step
4881
4943
  update_run_output_with_reasoning(
4882
4944
  run_response=run_response,
4883
- reasoning_steps=[ReasoningStep(result=reasoning_message.content)],
4884
- reasoning_agent_messages=[reasoning_message],
4945
+ reasoning_steps=[event.reasoning_step],
4946
+ reasoning_agent_messages=[],
4885
4947
  )
4886
4948
  if stream_events:
4949
+ updated_reasoning_content = self._format_reasoning_step_content(
4950
+ run_response=run_response,
4951
+ reasoning_step=event.reasoning_step,
4952
+ )
4887
4953
  yield handle_event( # type: ignore
4888
- create_team_reasoning_completed_event(
4954
+ create_team_reasoning_step_event(
4889
4955
  from_run_response=run_response,
4890
- content=ReasoningSteps(reasoning_steps=[ReasoningStep(result=reasoning_message.content)]),
4891
- content_type=ReasoningSteps.__name__,
4956
+ reasoning_step=event.reasoning_step,
4957
+ reasoning_content=updated_reasoning_content,
4892
4958
  ),
4893
4959
  run_response,
4894
4960
  events_to_skip=self.events_to_skip,
4895
4961
  store_events=self.store_events,
4896
4962
  )
4897
- else:
4898
- log_info(
4899
- f"Reasoning model: {reasoning_model.__class__.__name__} is not a native reasoning model, defaulting to manual Chain-of-Thought reasoning"
4900
- )
4901
- use_default_reasoning = True
4902
- # If no reasoning model is provided, use default reasoning
4903
- else:
4904
- use_default_reasoning = True
4905
-
4906
- if use_default_reasoning:
4907
- from agno.reasoning.default import get_default_reasoning_agent
4908
- from agno.reasoning.helpers import get_next_action, update_messages_with_reasoning
4909
-
4910
- # Get default reasoning agent
4911
- use_json_mode: bool = self.use_json_mode
4912
4963
 
4913
- reasoning_agent: Optional[Agent] = self.reasoning_agent # type: ignore
4914
- if reasoning_agent is None:
4915
- reasoning_agent = get_default_reasoning_agent(
4916
- reasoning_model=reasoning_model,
4917
- min_steps=self.reasoning_min_steps,
4918
- max_steps=self.reasoning_max_steps,
4919
- tool_call_limit=self.tool_call_limit,
4920
- telemetry=self.telemetry,
4921
- debug_mode=self.debug_mode,
4922
- debug_level=self.debug_level,
4923
- use_json_mode=use_json_mode,
4924
- session_state=self.session_state,
4925
- dependencies=self.dependencies,
4926
- metadata=self.metadata,
4964
+ elif event.event_type == ReasoningEventType.completed:
4965
+ if event.message and event.reasoning_steps:
4966
+ update_run_output_with_reasoning(
4967
+ run_response=run_response,
4968
+ reasoning_steps=event.reasoning_steps,
4969
+ reasoning_agent_messages=event.reasoning_messages,
4927
4970
  )
4928
-
4929
- # Validate reasoning agent
4930
- if reasoning_agent is None:
4931
- log_warning("Reasoning error. Reasoning agent is None, continuing regular session...")
4932
- return
4933
- # Ensure the reasoning agent response model is ReasoningSteps
4934
- if (
4935
- reasoning_agent.output_schema is not None
4936
- and not isinstance(reasoning_agent.output_schema, type)
4937
- and not issubclass(reasoning_agent.output_schema, ReasoningSteps)
4938
- ):
4939
- log_warning("Reasoning agent response model should be `ReasoningSteps`, continuing regular session...")
4940
- return
4941
- # Ensure the reasoning model and agent do not show tool calls
4942
-
4943
- step_count = 1
4944
- next_action = NextAction.CONTINUE
4945
- reasoning_messages: List[Message] = []
4946
- all_reasoning_steps: List[ReasoningStep] = []
4947
- log_debug("Starting Reasoning", center=True, symbol="=")
4948
- while next_action == NextAction.CONTINUE and step_count < self.reasoning_max_steps:
4949
- log_debug(f"Step {step_count}", center=True, symbol="-")
4950
- step_count += 1
4951
- try:
4952
- # Run the reasoning agent
4953
- reasoning_agent_response: RunOutput = reasoning_agent.run( # type: ignore
4954
- input=run_messages.get_input_messages()
4955
- )
4956
- if reasoning_agent_response.content is None or reasoning_agent_response.messages is None:
4957
- log_warning("Reasoning error. Reasoning response is empty, continuing regular session...")
4958
- break
4959
-
4960
- if isinstance(reasoning_agent_response.content, str):
4961
- log_warning(
4962
- "Reasoning error. Content is a string, not structured output. Continuing regular session..."
4963
- )
4964
- break
4965
-
4966
- if reasoning_agent_response.content.reasoning_steps is None:
4967
- log_warning("Reasoning error. Reasoning steps are empty, continuing regular session...")
4968
- break
4969
-
4970
- reasoning_steps: List[ReasoningStep] = reasoning_agent_response.content.reasoning_steps
4971
- all_reasoning_steps.extend(reasoning_steps)
4972
- # Yield reasoning steps
4973
- if stream_events:
4974
- for reasoning_step in reasoning_steps:
4975
- updated_reasoning_content = self._format_reasoning_step_content(
4976
- run_response, reasoning_step
4977
- )
4978
-
4979
- yield handle_event( # type: ignore
4980
- create_team_reasoning_step_event(
4981
- from_run_response=run_response,
4982
- reasoning_step=reasoning_step,
4983
- reasoning_content=updated_reasoning_content,
4984
- ),
4985
- run_response,
4986
- events_to_skip=self.events_to_skip,
4987
- store_events=self.store_events,
4988
- )
4989
-
4990
- # Find the index of the first assistant message
4991
- first_assistant_index = next(
4992
- (i for i, m in enumerate(reasoning_agent_response.messages) if m.role == "assistant"),
4993
- len(reasoning_agent_response.messages),
4994
- )
4995
- # Extract reasoning messages starting from the message after the first assistant message
4996
- reasoning_messages = reasoning_agent_response.messages[first_assistant_index:]
4997
-
4998
- # Add reasoning step to the Agent's run_response
4999
- update_run_output_with_reasoning(
5000
- run_response=run_response,
5001
- reasoning_steps=reasoning_steps,
5002
- reasoning_agent_messages=reasoning_agent_response.messages,
5003
- )
5004
-
5005
- # Get the next action
5006
- next_action = get_next_action(reasoning_steps[-1])
5007
- if next_action == NextAction.FINAL_ANSWER:
5008
- break
5009
- except Exception as e:
5010
- log_error(f"Reasoning error: {e}")
5011
- break
5012
-
5013
- log_debug(f"Total Reasoning steps: {len(all_reasoning_steps)}")
5014
- log_debug("Reasoning finished", center=True, symbol="=")
5015
-
5016
- # Update the messages_for_model to include reasoning messages
5017
- update_messages_with_reasoning(
5018
- run_messages=run_messages,
5019
- reasoning_messages=reasoning_messages,
5020
- )
5021
-
5022
- # Yield the final reasoning completed event
5023
4971
  if stream_events:
5024
4972
  yield handle_event( # type: ignore
5025
4973
  create_team_reasoning_completed_event(
5026
4974
  from_run_response=run_response,
5027
- content=ReasoningSteps(reasoning_steps=all_reasoning_steps),
4975
+ content=ReasoningSteps(reasoning_steps=event.reasoning_steps),
5028
4976
  content_type=ReasoningSteps.__name__,
5029
4977
  ),
5030
4978
  run_response,
@@ -5032,285 +4980,97 @@ class Team:
5032
4980
  store_events=self.store_events,
5033
4981
  )
5034
4982
 
5035
- async def _areason(
4983
+ elif event.event_type == ReasoningEventType.error:
4984
+ log_warning(f"Reasoning error. {event.error}, continuing regular session...")
4985
+
4986
+ def _reason(
5036
4987
  self,
5037
4988
  run_response: TeamRunOutput,
5038
4989
  run_messages: RunMessages,
5039
4990
  stream_events: bool,
5040
- ) -> AsyncIterator[TeamRunOutputEvent]:
5041
- if stream_events:
5042
- yield handle_event( # type: ignore
5043
- create_team_reasoning_started_event(from_run_response=run_response),
5044
- run_response,
5045
- events_to_skip=self.events_to_skip,
5046
- store_events=self.store_events,
5047
- )
4991
+ ) -> Iterator[TeamRunOutputEvent]:
4992
+ """
4993
+ Run reasoning using the ReasoningManager.
5048
4994
 
5049
- use_default_reasoning = False
4995
+ Handles both native reasoning models (DeepSeek, Anthropic, etc.) and
4996
+ default Chain-of-Thought reasoning with a clean, unified interface.
4997
+ """
4998
+ from agno.reasoning.manager import ReasoningConfig, ReasoningManager
5050
4999
 
5051
- # Get the reasoning model
5000
+ # Get the reasoning model (use copy of main model if not provided)
5052
5001
  reasoning_model: Optional[Model] = self.reasoning_model
5053
- reasoning_model_provided = reasoning_model is not None
5054
5002
  if reasoning_model is None and self.model is not None:
5055
5003
  from copy import deepcopy
5056
5004
 
5057
5005
  reasoning_model = deepcopy(self.model)
5058
- if reasoning_model is None:
5059
- log_warning("Reasoning error. Reasoning model is None, continuing regular session...")
5060
- return
5061
5006
 
5062
- # If a reasoning model is provided, use it to generate reasoning
5063
- if reasoning_model_provided:
5064
- from agno.reasoning.anthropic import is_anthropic_reasoning_model
5065
- from agno.reasoning.azure_ai_foundry import is_ai_foundry_reasoning_model
5066
- from agno.reasoning.deepseek import is_deepseek_reasoning_model
5067
- from agno.reasoning.gemini import is_gemini_reasoning_model
5068
- from agno.reasoning.groq import is_groq_reasoning_model
5069
- from agno.reasoning.helpers import get_reasoning_agent
5070
- from agno.reasoning.ollama import is_ollama_reasoning_model
5071
- from agno.reasoning.openai import is_openai_reasoning_model
5072
- from agno.reasoning.vertexai import is_vertexai_reasoning_model
5073
-
5074
- reasoning_agent = self.reasoning_agent or get_reasoning_agent(
5007
+ # Create reasoning manager with config
5008
+ manager = ReasoningManager(
5009
+ ReasoningConfig(
5075
5010
  reasoning_model=reasoning_model,
5011
+ reasoning_agent=self.reasoning_agent,
5012
+ min_steps=self.reasoning_min_steps,
5013
+ max_steps=self.reasoning_max_steps,
5014
+ tools=self.tools,
5015
+ tool_call_limit=self.tool_call_limit,
5016
+ use_json_mode=self.use_json_mode,
5017
+ telemetry=self.telemetry,
5018
+ debug_mode=self.debug_mode,
5019
+ debug_level=self.debug_level,
5076
5020
  session_state=self.session_state,
5077
5021
  dependencies=self.dependencies,
5078
5022
  metadata=self.metadata,
5079
5023
  )
5080
- is_deepseek = is_deepseek_reasoning_model(reasoning_model)
5081
- is_groq = is_groq_reasoning_model(reasoning_model)
5082
- is_openai = is_openai_reasoning_model(reasoning_model)
5083
- is_ollama = is_ollama_reasoning_model(reasoning_model)
5084
- is_ai_foundry = is_ai_foundry_reasoning_model(reasoning_model)
5085
- is_gemini = is_gemini_reasoning_model(reasoning_model)
5086
- is_anthropic = is_anthropic_reasoning_model(reasoning_model)
5087
- is_vertexai = is_vertexai_reasoning_model(reasoning_model)
5088
-
5089
- if (
5090
- is_deepseek
5091
- or is_groq
5092
- or is_openai
5093
- or is_ollama
5094
- or is_ai_foundry
5095
- or is_gemini
5096
- or is_anthropic
5097
- or is_vertexai
5098
- ):
5099
- reasoning_message: Optional[Message] = None
5100
- if is_deepseek:
5101
- from agno.reasoning.deepseek import aget_deepseek_reasoning
5102
-
5103
- log_debug("Starting DeepSeek Reasoning", center=True, symbol="=")
5104
- reasoning_message = await aget_deepseek_reasoning(
5105
- reasoning_agent=reasoning_agent, messages=run_messages.get_input_messages()
5106
- )
5107
- elif is_groq:
5108
- from agno.reasoning.groq import aget_groq_reasoning
5109
-
5110
- log_debug("Starting Groq Reasoning", center=True, symbol="=")
5111
- reasoning_message = await aget_groq_reasoning(
5112
- reasoning_agent=reasoning_agent, messages=run_messages.get_input_messages()
5113
- )
5114
- elif is_openai:
5115
- from agno.reasoning.openai import aget_openai_reasoning
5116
-
5117
- log_debug("Starting OpenAI Reasoning", center=True, symbol="=")
5118
- reasoning_message = await aget_openai_reasoning(
5119
- reasoning_agent=reasoning_agent, messages=run_messages.get_input_messages()
5120
- )
5121
- elif is_ollama:
5122
- from agno.reasoning.ollama import get_ollama_reasoning
5123
-
5124
- log_debug("Starting Ollama Reasoning", center=True, symbol="=")
5125
- reasoning_message = get_ollama_reasoning(
5126
- reasoning_agent=reasoning_agent, messages=run_messages.get_input_messages()
5127
- )
5128
- elif is_ai_foundry:
5129
- from agno.reasoning.azure_ai_foundry import get_ai_foundry_reasoning
5130
-
5131
- log_debug("Starting Azure AI Foundry Reasoning", center=True, symbol="=")
5132
- reasoning_message = get_ai_foundry_reasoning(
5133
- reasoning_agent=reasoning_agent, messages=run_messages.get_input_messages()
5134
- )
5135
- elif is_gemini:
5136
- from agno.reasoning.gemini import aget_gemini_reasoning
5137
-
5138
- log_debug("Starting Gemini Reasoning", center=True, symbol="=")
5139
- reasoning_message = await aget_gemini_reasoning(
5140
- reasoning_agent=reasoning_agent, messages=run_messages.get_input_messages()
5141
- )
5142
- elif is_anthropic:
5143
- from agno.reasoning.anthropic import aget_anthropic_reasoning
5144
-
5145
- log_debug("Starting Anthropic Claude Reasoning", center=True, symbol="=")
5146
- reasoning_message = await aget_anthropic_reasoning(
5147
- reasoning_agent=reasoning_agent, messages=run_messages.get_input_messages()
5148
- )
5149
- elif is_vertexai:
5150
- from agno.reasoning.vertexai import aget_vertexai_reasoning
5151
-
5152
- log_debug("Starting VertexAI Reasoning", center=True, symbol="=")
5153
- reasoning_message = await aget_vertexai_reasoning(
5154
- reasoning_agent=reasoning_agent, messages=run_messages.get_input_messages()
5155
- )
5156
-
5157
- if reasoning_message is None:
5158
- log_warning("Reasoning error. Reasoning response is None, continuing regular session...")
5159
- return
5160
- run_messages.messages.append(reasoning_message)
5161
- # Add reasoning step to the Agent's run_response
5162
- update_run_output_with_reasoning(
5163
- run_response=run_response,
5164
- reasoning_steps=[ReasoningStep(result=reasoning_message.content)],
5165
- reasoning_agent_messages=[reasoning_message],
5166
- )
5167
- if stream_events:
5168
- yield handle_event( # type: ignore
5169
- create_team_reasoning_completed_event(
5170
- from_run_response=run_response,
5171
- content=ReasoningSteps(reasoning_steps=[ReasoningStep(result=reasoning_message.content)]),
5172
- content_type=ReasoningSteps.__name__,
5173
- ),
5174
- run_response,
5175
- events_to_skip=self.events_to_skip,
5176
- store_events=self.store_events,
5177
- )
5178
- else:
5179
- log_info(
5180
- f"Reasoning model: {reasoning_model.__class__.__name__} is not a native reasoning model, defaulting to manual Chain-of-Thought reasoning"
5181
- )
5182
- use_default_reasoning = True
5183
- # If no reasoning model is provided, use default reasoning
5184
- else:
5185
- use_default_reasoning = True
5186
-
5187
- if use_default_reasoning:
5188
- from agno.reasoning.default import get_default_reasoning_agent
5189
- from agno.reasoning.helpers import get_next_action, update_messages_with_reasoning
5190
-
5191
- # Get default reasoning agent
5192
- use_json_mode: bool = self.use_json_mode
5193
- reasoning_agent: Optional[Agent] = self.reasoning_agent # type: ignore
5194
- if reasoning_agent is None:
5195
- reasoning_agent = get_default_reasoning_agent( # type: ignore
5196
- reasoning_model=reasoning_model,
5197
- min_steps=self.reasoning_min_steps,
5198
- max_steps=self.reasoning_max_steps,
5199
- telemetry=self.telemetry,
5200
- debug_mode=self.debug_mode,
5201
- debug_level=self.debug_level,
5202
- use_json_mode=use_json_mode,
5203
- session_state=self.session_state,
5204
- dependencies=self.dependencies,
5205
- metadata=self.metadata,
5206
- )
5207
-
5208
- # Validate reasoning agent
5209
- if reasoning_agent is None:
5210
- log_warning("Reasoning error. Reasoning agent is None, continuing regular session...")
5211
- return
5212
- # Ensure the reasoning agent response model is ReasoningSteps
5213
- if (
5214
- reasoning_agent.output_schema is not None
5215
- and not isinstance(reasoning_agent.output_schema, type)
5216
- and not issubclass(reasoning_agent.output_schema, ReasoningSteps)
5217
- ):
5218
- log_warning("Reasoning agent response model should be `ReasoningSteps`, continuing regular session...")
5219
- return
5220
-
5221
- # Ensure the reasoning model and agent do not show tool calls
5222
-
5223
- step_count = 1
5224
- next_action = NextAction.CONTINUE
5225
- reasoning_messages: List[Message] = []
5226
- all_reasoning_steps: List[ReasoningStep] = []
5227
- log_debug("Starting Reasoning", center=True, symbol="=")
5228
- while next_action == NextAction.CONTINUE and step_count < self.reasoning_max_steps:
5229
- log_debug(f"Step {step_count}", center=True, symbol="-")
5230
- step_count += 1
5231
- try:
5232
- # Run the reasoning agent
5233
- reasoning_agent_response: RunOutput = await reasoning_agent.arun( # type: ignore
5234
- input=run_messages.get_input_messages()
5235
- )
5236
- if reasoning_agent_response.content is None or reasoning_agent_response.messages is None:
5237
- log_warning("Reasoning error. Reasoning response is empty, continuing regular session...")
5238
- break
5239
-
5240
- if isinstance(reasoning_agent_response.content, str):
5241
- log_warning(
5242
- "Reasoning error. Content is a string, not structured output. Continuing regular session..."
5243
- )
5244
- break
5245
-
5246
- if reasoning_agent_response.content.reasoning_steps is None:
5247
- log_warning("Reasoning error. Reasoning steps are empty, continuing regular session...")
5248
- break
5249
-
5250
- reasoning_steps: List[ReasoningStep] = reasoning_agent_response.content.reasoning_steps
5251
- all_reasoning_steps.extend(reasoning_steps)
5252
- # Yield reasoning steps
5253
- if stream_events:
5254
- for reasoning_step in reasoning_steps:
5255
- updated_reasoning_content = self._format_reasoning_step_content(
5256
- run_response, reasoning_step
5257
- )
5024
+ )
5258
5025
 
5259
- yield handle_event( # type: ignore
5260
- create_team_reasoning_step_event(
5261
- from_run_response=run_response,
5262
- reasoning_step=reasoning_step,
5263
- reasoning_content=updated_reasoning_content,
5264
- ),
5265
- run_response,
5266
- events_to_skip=self.events_to_skip,
5267
- store_events=self.store_events,
5268
- )
5026
+ # Use the unified reason() method and convert events
5027
+ for event in manager.reason(run_messages, stream=stream_events):
5028
+ yield from self._handle_reasoning_event(event, run_response, stream_events)
5269
5029
 
5270
- # Find the index of the first assistant message
5271
- first_assistant_index = next(
5272
- (i for i, m in enumerate(reasoning_agent_response.messages) if m.role == "assistant"),
5273
- len(reasoning_agent_response.messages),
5274
- )
5275
- # Extract reasoning messages starting from the message after the first assistant message
5276
- reasoning_messages = reasoning_agent_response.messages[first_assistant_index:]
5030
+ async def _areason(
5031
+ self,
5032
+ run_response: TeamRunOutput,
5033
+ run_messages: RunMessages,
5034
+ stream_events: bool,
5035
+ ) -> AsyncIterator[TeamRunOutputEvent]:
5036
+ """
5037
+ Run reasoning asynchronously using the ReasoningManager.
5277
5038
 
5278
- # Add reasoning step to the Agent's run_response
5279
- update_run_output_with_reasoning(
5280
- run_response=run_response,
5281
- reasoning_steps=reasoning_steps,
5282
- reasoning_agent_messages=reasoning_agent_response.messages,
5283
- )
5039
+ Handles both native reasoning models (DeepSeek, Anthropic, etc.) and
5040
+ default Chain-of-Thought reasoning with a clean, unified interface.
5041
+ """
5042
+ from agno.reasoning.manager import ReasoningConfig, ReasoningManager
5284
5043
 
5285
- # Get the next action
5286
- next_action = get_next_action(reasoning_steps[-1])
5287
- if next_action == NextAction.FINAL_ANSWER:
5288
- break
5289
- except Exception as e:
5290
- log_error(f"Reasoning error: {e}")
5291
- break
5044
+ # Get the reasoning model (use copy of main model if not provided)
5045
+ reasoning_model: Optional[Model] = self.reasoning_model
5046
+ if reasoning_model is None and self.model is not None:
5047
+ from copy import deepcopy
5292
5048
 
5293
- log_debug(f"Total Reasoning steps: {len(all_reasoning_steps)}")
5294
- log_debug("Reasoning finished", center=True, symbol="=")
5049
+ reasoning_model = deepcopy(self.model)
5295
5050
 
5296
- # Update the messages_for_model to include reasoning messages
5297
- update_messages_with_reasoning(
5298
- run_messages=run_messages,
5299
- reasoning_messages=reasoning_messages,
5051
+ # Create reasoning manager with config
5052
+ manager = ReasoningManager(
5053
+ ReasoningConfig(
5054
+ reasoning_model=reasoning_model,
5055
+ reasoning_agent=self.reasoning_agent,
5056
+ min_steps=self.reasoning_min_steps,
5057
+ max_steps=self.reasoning_max_steps,
5058
+ tools=self.tools,
5059
+ tool_call_limit=self.tool_call_limit,
5060
+ use_json_mode=self.use_json_mode,
5061
+ telemetry=self.telemetry,
5062
+ debug_mode=self.debug_mode,
5063
+ debug_level=self.debug_level,
5064
+ session_state=self.session_state,
5065
+ dependencies=self.dependencies,
5066
+ metadata=self.metadata,
5300
5067
  )
5068
+ )
5301
5069
 
5302
- # Yield the final reasoning completed event
5303
- if stream_events:
5304
- yield handle_event( # type: ignore # type: ignore
5305
- create_team_reasoning_completed_event(
5306
- from_run_response=run_response,
5307
- content=ReasoningSteps(reasoning_steps=all_reasoning_steps),
5308
- content_type=ReasoningSteps.__name__,
5309
- ),
5310
- run_response,
5311
- events_to_skip=self.events_to_skip,
5312
- store_events=self.store_events,
5313
- )
5070
+ # Use the unified areason() method and convert events
5071
+ async for event in manager.areason(run_messages, stream=stream_events):
5072
+ for output_event in self._handle_reasoning_event(event, run_response, stream_events):
5073
+ yield output_event
5314
5074
 
5315
5075
  def _resolve_run_dependencies(self, run_context: RunContext) -> None:
5316
5076
  from inspect import signature
@@ -7021,7 +6781,7 @@ class Team:
7021
6781
  log_error(f"Failed to convert sanitized context to JSON: {e}")
7022
6782
  return str(context)
7023
6783
 
7024
- def _get_json_output_prompt(self, output_schema: Optional[Type[BaseModel]] = None) -> str:
6784
+ def _get_json_output_prompt(self, output_schema: Optional[Union[Type[BaseModel], Dict[str, Any]]] = None) -> str:
7025
6785
  """Return the JSON output prompt for the Agent.
7026
6786
 
7027
6787
  This is added to the system prompt when the output_schema is set and structured_outputs is False.
@@ -7038,7 +6798,11 @@ class Team:
7038
6798
  json_output_prompt += "\n<json_fields>"
7039
6799
  json_output_prompt += f"\n{json.dumps(output_schema)}"
7040
6800
  json_output_prompt += "\n</json_fields>"
7041
- elif issubclass(output_schema, BaseModel):
6801
+ elif isinstance(output_schema, dict):
6802
+ json_output_prompt += "\n<json_fields>"
6803
+ json_output_prompt += f"\n{json.dumps(output_schema)}"
6804
+ json_output_prompt += "\n</json_fields>"
6805
+ elif isinstance(output_schema, type) and issubclass(output_schema, BaseModel):
7042
6806
  json_schema = output_schema.model_json_schema()
7043
6807
  if json_schema is not None:
7044
6808
  response_model_properties = {}
@@ -7552,7 +7316,7 @@ class Team:
7552
7316
  member_agent_run_response.parent_run_id = run_response.run_id # type: ignore
7553
7317
 
7554
7318
  # Update the top-level team run_response tool call to have the run_id of the member run
7555
- if run_response.tools is not None:
7319
+ if run_response.tools is not None and member_agent_run_response is not None:
7556
7320
  for tool in run_response.tools:
7557
7321
  if tool.tool_name and tool.tool_name.lower() == "delegate_task_to_member":
7558
7322
  tool.child_run_id = member_agent_run_response.run_id # type: ignore
@@ -7782,9 +7546,9 @@ class Team:
7782
7546
  check_if_run_cancelled(member_agent_run_response_event)
7783
7547
 
7784
7548
  # Yield the member event directly
7785
- member_agent_run_response_event.parent_run_id = (
7786
- getattr(member_agent_run_response_event, "parent_run_id", None) or run_response.run_id
7787
- )
7549
+ member_agent_run_response_event.parent_run_id = getattr(
7550
+ member_agent_run_response_event, "parent_run_id", None
7551
+ ) or (run_response.run_id if run_response is not None else None)
7788
7552
  yield member_agent_run_response_event # type: ignore
7789
7553
  else:
7790
7554
  member_agent_run_response = await member_agent.arun( # type: ignore
@@ -7895,7 +7659,8 @@ class Team:
7895
7659
 
7896
7660
  # Yield the member event directly
7897
7661
  member_agent_run_response_chunk.parent_run_id = (
7898
- member_agent_run_response_chunk.parent_run_id or run_response.run_id
7662
+ member_agent_run_response_chunk.parent_run_id
7663
+ or (run_response.run_id if run_response is not None else None)
7899
7664
  )
7900
7665
  yield member_agent_run_response_chunk # type: ignore
7901
7666
 
@@ -8005,7 +7770,8 @@ class Team:
8005
7770
 
8006
7771
  check_if_run_cancelled(member_agent_run_output_event)
8007
7772
  member_agent_run_output_event.parent_run_id = (
8008
- member_agent_run_output_event.parent_run_id or run_response.run_id
7773
+ member_agent_run_output_event.parent_run_id
7774
+ or (run_response.run_id if run_response is not None else None)
8009
7775
  )
8010
7776
  await queue.put(member_agent_run_output_event)
8011
7777
  finally: