agno 2.3.4__py3-none-any.whl → 2.3.5__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 (112) hide show
  1. agno/agent/agent.py +177 -41
  2. agno/culture/manager.py +2 -2
  3. agno/db/base.py +330 -8
  4. agno/db/dynamo/dynamo.py +722 -2
  5. agno/db/dynamo/schemas.py +127 -0
  6. agno/db/firestore/firestore.py +573 -1
  7. agno/db/firestore/schemas.py +40 -0
  8. agno/db/gcs_json/gcs_json_db.py +446 -1
  9. agno/db/in_memory/in_memory_db.py +143 -1
  10. agno/db/json/json_db.py +438 -1
  11. agno/db/mongo/async_mongo.py +522 -0
  12. agno/db/mongo/mongo.py +523 -1
  13. agno/db/mongo/schemas.py +29 -0
  14. agno/db/mysql/mysql.py +536 -3
  15. agno/db/mysql/schemas.py +38 -0
  16. agno/db/postgres/async_postgres.py +541 -13
  17. agno/db/postgres/postgres.py +535 -2
  18. agno/db/postgres/schemas.py +38 -0
  19. agno/db/redis/redis.py +468 -1
  20. agno/db/redis/schemas.py +32 -0
  21. agno/db/singlestore/schemas.py +38 -0
  22. agno/db/singlestore/singlestore.py +523 -1
  23. agno/db/sqlite/async_sqlite.py +548 -9
  24. agno/db/sqlite/schemas.py +38 -0
  25. agno/db/sqlite/sqlite.py +537 -5
  26. agno/db/sqlite/utils.py +6 -8
  27. agno/db/surrealdb/models.py +25 -0
  28. agno/db/surrealdb/surrealdb.py +548 -1
  29. agno/eval/accuracy.py +10 -4
  30. agno/eval/performance.py +10 -4
  31. agno/eval/reliability.py +22 -13
  32. agno/exceptions.py +11 -0
  33. agno/hooks/__init__.py +3 -0
  34. agno/hooks/decorator.py +164 -0
  35. agno/knowledge/chunking/semantic.py +2 -2
  36. agno/models/aimlapi/aimlapi.py +2 -3
  37. agno/models/anthropic/claude.py +18 -13
  38. agno/models/aws/bedrock.py +3 -4
  39. agno/models/aws/claude.py +5 -1
  40. agno/models/azure/ai_foundry.py +2 -2
  41. agno/models/azure/openai_chat.py +8 -0
  42. agno/models/cerebras/cerebras.py +63 -11
  43. agno/models/cerebras/cerebras_openai.py +2 -3
  44. agno/models/cohere/chat.py +1 -5
  45. agno/models/cometapi/cometapi.py +2 -3
  46. agno/models/dashscope/dashscope.py +2 -3
  47. agno/models/deepinfra/deepinfra.py +2 -3
  48. agno/models/deepseek/deepseek.py +2 -3
  49. agno/models/fireworks/fireworks.py +2 -3
  50. agno/models/google/gemini.py +9 -7
  51. agno/models/groq/groq.py +2 -3
  52. agno/models/huggingface/huggingface.py +1 -5
  53. agno/models/ibm/watsonx.py +1 -5
  54. agno/models/internlm/internlm.py +2 -3
  55. agno/models/langdb/langdb.py +6 -4
  56. agno/models/litellm/chat.py +2 -2
  57. agno/models/litellm/litellm_openai.py +2 -3
  58. agno/models/meta/llama.py +1 -5
  59. agno/models/meta/llama_openai.py +4 -5
  60. agno/models/mistral/mistral.py +1 -5
  61. agno/models/nebius/nebius.py +2 -3
  62. agno/models/nvidia/nvidia.py +4 -5
  63. agno/models/openai/chat.py +14 -3
  64. agno/models/openai/responses.py +14 -3
  65. agno/models/openrouter/openrouter.py +4 -5
  66. agno/models/perplexity/perplexity.py +2 -3
  67. agno/models/portkey/portkey.py +7 -6
  68. agno/models/requesty/requesty.py +4 -5
  69. agno/models/response.py +2 -1
  70. agno/models/sambanova/sambanova.py +4 -5
  71. agno/models/siliconflow/siliconflow.py +3 -4
  72. agno/models/together/together.py +4 -5
  73. agno/models/vercel/v0.py +4 -5
  74. agno/models/vllm/vllm.py +19 -14
  75. agno/models/xai/xai.py +4 -5
  76. agno/os/app.py +104 -0
  77. agno/os/config.py +13 -0
  78. agno/os/interfaces/whatsapp/router.py +0 -1
  79. agno/os/mcp.py +1 -0
  80. agno/os/router.py +31 -0
  81. agno/os/routers/traces/__init__.py +3 -0
  82. agno/os/routers/traces/schemas.py +414 -0
  83. agno/os/routers/traces/traces.py +499 -0
  84. agno/os/schema.py +10 -1
  85. agno/os/utils.py +57 -0
  86. agno/run/agent.py +1 -0
  87. agno/run/base.py +17 -0
  88. agno/run/team.py +4 -0
  89. agno/session/team.py +1 -0
  90. agno/table.py +10 -0
  91. agno/team/team.py +214 -65
  92. agno/tools/function.py +10 -8
  93. agno/tools/nano_banana.py +1 -1
  94. agno/tracing/__init__.py +12 -0
  95. agno/tracing/exporter.py +157 -0
  96. agno/tracing/schemas.py +276 -0
  97. agno/tracing/setup.py +111 -0
  98. agno/utils/agent.py +4 -4
  99. agno/utils/hooks.py +56 -1
  100. agno/vectordb/qdrant/qdrant.py +22 -22
  101. agno/workflow/condition.py +8 -0
  102. agno/workflow/loop.py +8 -0
  103. agno/workflow/parallel.py +8 -0
  104. agno/workflow/router.py +8 -0
  105. agno/workflow/step.py +20 -0
  106. agno/workflow/steps.py +8 -0
  107. agno/workflow/workflow.py +83 -17
  108. {agno-2.3.4.dist-info → agno-2.3.5.dist-info}/METADATA +2 -2
  109. {agno-2.3.4.dist-info → agno-2.3.5.dist-info}/RECORD +112 -102
  110. {agno-2.3.4.dist-info → agno-2.3.5.dist-info}/WHEEL +0 -0
  111. {agno-2.3.4.dist-info → agno-2.3.5.dist-info}/licenses/LICENSE +0 -0
  112. {agno-2.3.4.dist-info → agno-2.3.5.dist-info}/top_level.txt +0 -0
agno/team/team.py CHANGED
@@ -1,3 +1,5 @@
1
+ from __future__ import annotations
2
+
1
3
  import asyncio
2
4
  import contextlib
3
5
  import json
@@ -76,8 +78,8 @@ from agno.utils.agent import (
76
78
  aget_session_state_util,
77
79
  aset_session_name_util,
78
80
  aupdate_session_state_util,
79
- await_for_background_tasks,
80
- await_for_background_tasks_stream,
81
+ await_for_open_threads,
82
+ await_for_thread_tasks_stream,
81
83
  collect_joint_audios,
82
84
  collect_joint_files,
83
85
  collect_joint_images,
@@ -96,8 +98,8 @@ from agno.utils.agent import (
96
98
  store_media_util,
97
99
  update_session_state_util,
98
100
  validate_media_object_id,
99
- wait_for_background_tasks,
100
- wait_for_background_tasks_stream,
101
+ wait_for_open_threads,
102
+ wait_for_thread_tasks_stream,
101
103
  )
102
104
  from agno.utils.common import is_typed_dict, validate_typed_dict
103
105
  from agno.utils.events import (
@@ -122,7 +124,7 @@ from agno.utils.events import (
122
124
  create_team_tool_call_started_event,
123
125
  handle_event,
124
126
  )
125
- from agno.utils.hooks import filter_hook_args, normalize_hooks
127
+ from agno.utils.hooks import copy_args_for_background, filter_hook_args, normalize_hooks, should_run_hook_in_background
126
128
  from agno.utils.knowledge import get_agentic_or_user_search_filters
127
129
  from agno.utils.log import (
128
130
  log_debug,
@@ -344,6 +346,8 @@ class Team:
344
346
  pre_hooks: Optional[List[Union[Callable[..., Any], BaseGuardrail]]] = None
345
347
  # Functions called after output is generated but before the response is returned
346
348
  post_hooks: Optional[List[Union[Callable[..., Any], BaseGuardrail]]] = None
349
+ # If True, run hooks as FastAPI background tasks (non-blocking). Set by AgentOS.
350
+ _run_hooks_in_background: Optional[bool] = None
347
351
 
348
352
  # --- Structured output ---
349
353
  # Input schema for validating input
@@ -636,7 +640,7 @@ class Team:
636
640
  self.tool_call_limit = tool_call_limit
637
641
  self.tool_hooks = tool_hooks
638
642
 
639
- # Initialize hooks with backward compatibility
643
+ # Initialize hooks
640
644
  self.pre_hooks = pre_hooks
641
645
  self.post_hooks = post_hooks
642
646
 
@@ -752,8 +756,13 @@ class Team:
752
756
  self.id = generate_id_from_name(self.name)
753
757
 
754
758
  def _set_debug(self, debug_mode: Optional[bool] = None) -> None:
759
+ # Get the debug level from the environment variable or the default debug level
760
+ debug_level: Literal[1, 2] = (
761
+ cast(Literal[1, 2], int(env)) if (env := getenv("AGNO_DEBUG_LEVEL")) in ("1", "2") else self.debug_level
762
+ )
763
+ # If the default debug mode is set, or passed on run, or via environment variable, set the debug mode to True
755
764
  if self.debug_mode or debug_mode or getenv("AGNO_DEBUG", "false").lower() == "true":
756
- set_log_level_to_debug(source_type="team", level=self.debug_level)
765
+ set_log_level_to_debug(source_type="team", level=debug_level)
757
766
  else:
758
767
  set_log_level_to_info(source_type="team")
759
768
 
@@ -837,6 +846,26 @@ class Team:
837
846
  for sub_member in member.members:
838
847
  member._initialize_member(sub_member, debug_mode=debug_mode)
839
848
 
849
+ def propagate_run_hooks_in_background(self, run_in_background: bool = True) -> None:
850
+ """
851
+ Propagate _run_hooks_in_background setting to this team and all nested members recursively.
852
+
853
+ This method sets _run_hooks_in_background on the team and all its members (agents and nested teams).
854
+ For nested teams, it recursively propagates the setting to their members as well.
855
+
856
+ Args:
857
+ run_in_background: Whether hooks should run in background. Defaults to True.
858
+ """
859
+ self._run_hooks_in_background = run_in_background
860
+
861
+ for member in self.members:
862
+ if hasattr(member, "_run_hooks_in_background"):
863
+ member._run_hooks_in_background = run_in_background
864
+
865
+ # If it's a nested team, recursively propagate to its members
866
+ if isinstance(member, Team):
867
+ member.propagate_run_hooks_in_background(run_in_background)
868
+
840
869
  def _set_default_model(self) -> None:
841
870
  # Set the default model
842
871
  if self.model is None:
@@ -1037,13 +1066,14 @@ class Team:
1037
1066
  user_id: Optional[str] = None,
1038
1067
  debug_mode: Optional[bool] = None,
1039
1068
  stream_events: bool = False,
1069
+ background_tasks: Optional[Any] = None,
1040
1070
  **kwargs: Any,
1041
1071
  ) -> Iterator[TeamRunOutputEvent]:
1042
1072
  """Execute multiple pre-hook functions in succession."""
1043
1073
  if hooks is None:
1044
1074
  return
1045
1075
 
1046
- # Prepare all possible arguments once
1076
+ # Prepare arguments for hooks
1047
1077
  all_args = {
1048
1078
  "run_input": run_input,
1049
1079
  "run_context": run_context,
@@ -1055,9 +1085,32 @@ class Team:
1055
1085
  "dependencies": run_context.dependencies,
1056
1086
  "debug_mode": debug_mode or self.debug_mode,
1057
1087
  }
1088
+
1089
+ # Check if background_tasks is available and ALL hooks should run in background
1090
+ # Note: Pre-hooks running in background may not be able to modify run_input
1091
+ if self._run_hooks_in_background is True and background_tasks is not None:
1092
+ # Schedule ALL pre_hooks as background tasks
1093
+ # Copy args to prevent race conditions
1094
+ bg_args = copy_args_for_background(all_args)
1095
+ for hook in hooks:
1096
+ # Filter arguments to only include those that the hook accepts
1097
+ filtered_args = filter_hook_args(hook, bg_args)
1098
+
1099
+ # Add to background tasks
1100
+ background_tasks.add_task(hook, **filtered_args)
1101
+ return
1102
+
1058
1103
  all_args.update(kwargs)
1059
1104
 
1060
1105
  for i, hook in enumerate(hooks):
1106
+ # Check if this specific hook should run in background (via @hook decorator)
1107
+ if should_run_hook_in_background(hook) and background_tasks is not None:
1108
+ # Copy args to prevent race conditions
1109
+ bg_args = copy_args_for_background(all_args)
1110
+ filtered_args = filter_hook_args(hook, bg_args)
1111
+ background_tasks.add_task(hook, **filtered_args)
1112
+ continue
1113
+
1061
1114
  if stream_events:
1062
1115
  yield handle_event( # type: ignore
1063
1116
  run_response=run_response,
@@ -1105,13 +1158,14 @@ class Team:
1105
1158
  user_id: Optional[str] = None,
1106
1159
  debug_mode: Optional[bool] = None,
1107
1160
  stream_events: bool = False,
1161
+ background_tasks: Optional[Any] = None,
1108
1162
  **kwargs: Any,
1109
1163
  ) -> AsyncIterator[TeamRunOutputEvent]:
1110
1164
  """Execute multiple pre-hook functions in succession (async version)."""
1111
1165
  if hooks is None:
1112
1166
  return
1113
1167
 
1114
- # Prepare all possible arguments once
1168
+ # Prepare arguments for hooks
1115
1169
  all_args = {
1116
1170
  "run_input": run_input,
1117
1171
  "run_context": run_context,
@@ -1123,9 +1177,32 @@ class Team:
1123
1177
  "metadata": run_context.metadata,
1124
1178
  "debug_mode": debug_mode or self.debug_mode,
1125
1179
  }
1180
+
1181
+ # Check if background_tasks is available and ALL hooks should run in background
1182
+ # Note: Pre-hooks running in background may not be able to modify run_input
1183
+ if self._run_hooks_in_background is True and background_tasks is not None:
1184
+ # Schedule ALL pre_hooks as background tasks
1185
+ # Copy args to prevent race conditions
1186
+ bg_args = copy_args_for_background(all_args)
1187
+ for hook in hooks:
1188
+ # Filter arguments to only include those that the hook accepts
1189
+ filtered_args = filter_hook_args(hook, bg_args)
1190
+
1191
+ # Add to background tasks (both sync and async hooks supported)
1192
+ background_tasks.add_task(hook, **filtered_args)
1193
+ return
1194
+
1126
1195
  all_args.update(kwargs)
1127
1196
 
1128
1197
  for i, hook in enumerate(hooks):
1198
+ # Check if this specific hook should run in background (via @hook decorator)
1199
+ if should_run_hook_in_background(hook) and background_tasks is not None:
1200
+ # Copy args to prevent race conditions
1201
+ bg_args = copy_args_for_background(all_args)
1202
+ filtered_args = filter_hook_args(hook, bg_args)
1203
+ background_tasks.add_task(hook, **filtered_args)
1204
+ continue
1205
+
1129
1206
  if stream_events:
1130
1207
  yield handle_event( # type: ignore
1131
1208
  run_response=run_response,
@@ -1178,13 +1255,14 @@ class Team:
1178
1255
  user_id: Optional[str] = None,
1179
1256
  debug_mode: Optional[bool] = None,
1180
1257
  stream_events: bool = False,
1258
+ background_tasks: Optional[Any] = None,
1181
1259
  **kwargs: Any,
1182
1260
  ) -> Iterator[TeamRunOutputEvent]:
1183
1261
  """Execute multiple post-hook functions in succession."""
1184
1262
  if hooks is None:
1185
1263
  return
1186
1264
 
1187
- # Prepare all possible arguments once
1265
+ # Prepare arguments for hooks
1188
1266
  all_args = {
1189
1267
  "run_output": run_output,
1190
1268
  "run_context": run_context,
@@ -1196,9 +1274,31 @@ class Team:
1196
1274
  "metadata": run_context.metadata,
1197
1275
  "debug_mode": debug_mode or self.debug_mode,
1198
1276
  }
1277
+
1278
+ # Check if background_tasks is available and ALL hooks should run in background
1279
+ if self._run_hooks_in_background is True and background_tasks is not None:
1280
+ # Schedule ALL post_hooks as background tasks
1281
+ # Copy args to prevent race conditions
1282
+ bg_args = copy_args_for_background(all_args)
1283
+ for hook in hooks:
1284
+ # Filter arguments to only include those that the hook accepts
1285
+ filtered_args = filter_hook_args(hook, bg_args)
1286
+
1287
+ # Add to background tasks
1288
+ background_tasks.add_task(hook, **filtered_args)
1289
+ return
1290
+
1199
1291
  all_args.update(kwargs)
1200
1292
 
1201
1293
  for i, hook in enumerate(hooks):
1294
+ # Check if this specific hook should run in background (via @hook decorator)
1295
+ if should_run_hook_in_background(hook) and background_tasks is not None:
1296
+ # Copy args to prevent race conditions
1297
+ bg_args = copy_args_for_background(all_args)
1298
+ filtered_args = filter_hook_args(hook, bg_args)
1299
+ background_tasks.add_task(hook, **filtered_args)
1300
+ continue
1301
+
1202
1302
  if stream_events:
1203
1303
  yield handle_event( # type: ignore
1204
1304
  run_response=run_output,
@@ -1241,13 +1341,14 @@ class Team:
1241
1341
  user_id: Optional[str] = None,
1242
1342
  debug_mode: Optional[bool] = None,
1243
1343
  stream_events: bool = False,
1344
+ background_tasks: Optional[Any] = None,
1244
1345
  **kwargs: Any,
1245
1346
  ) -> AsyncIterator[TeamRunOutputEvent]:
1246
1347
  """Execute multiple post-hook functions in succession (async version)."""
1247
1348
  if hooks is None:
1248
1349
  return
1249
1350
 
1250
- # Prepare all possible arguments once
1351
+ # Prepare arguments for hooks
1251
1352
  all_args = {
1252
1353
  "run_output": run_output,
1253
1354
  "run_context": run_context,
@@ -1259,9 +1360,31 @@ class Team:
1259
1360
  "metadata": run_context.metadata,
1260
1361
  "debug_mode": debug_mode or self.debug_mode,
1261
1362
  }
1363
+
1364
+ # Check if background_tasks is available and ALL hooks should run in background
1365
+ if self._run_hooks_in_background is True and background_tasks is not None:
1366
+ # Schedule ALL post_hooks as background tasks
1367
+ # Copy args to prevent race conditions
1368
+ bg_args = copy_args_for_background(all_args)
1369
+ for hook in hooks:
1370
+ # Filter arguments to only include those that the hook accepts
1371
+ filtered_args = filter_hook_args(hook, bg_args)
1372
+
1373
+ # Add to background tasks (both sync and async hooks supported)
1374
+ background_tasks.add_task(hook, **filtered_args)
1375
+ return
1376
+
1262
1377
  all_args.update(kwargs)
1263
1378
 
1264
1379
  for i, hook in enumerate(hooks):
1380
+ # Check if this specific hook should run in background (via @hook decorator)
1381
+ if should_run_hook_in_background(hook) and background_tasks is not None:
1382
+ # Copy args to prevent race conditions
1383
+ bg_args = copy_args_for_background(all_args)
1384
+ filtered_args = filter_hook_args(hook, bg_args)
1385
+ background_tasks.add_task(hook, **filtered_args)
1386
+ continue
1387
+
1265
1388
  if stream_events:
1266
1389
  yield handle_event( # type: ignore
1267
1390
  run_response=run_output,
@@ -1310,6 +1433,7 @@ class Team:
1310
1433
  add_session_state_to_context: Optional[bool] = None,
1311
1434
  response_format: Optional[Union[Dict, Type[BaseModel]]] = None,
1312
1435
  debug_mode: Optional[bool] = None,
1436
+ background_tasks: Optional[Any] = None,
1313
1437
  **kwargs: Any,
1314
1438
  ) -> TeamRunOutput:
1315
1439
  """Run the Team and return the response.
@@ -1330,9 +1454,6 @@ class Team:
1330
1454
  13. Cleanup and store (scrub, stop timer, add to session, calculate metrics, save session)
1331
1455
  """
1332
1456
 
1333
- # Register run for cancellation tracking
1334
- register_run(run_response.run_id) # type: ignore
1335
-
1336
1457
  # 1. Execute pre-hooks
1337
1458
  run_input = cast(TeamRunInput, run_response.input)
1338
1459
  self.model = cast(Model, self.model)
@@ -1346,6 +1467,7 @@ class Team:
1346
1467
  session=session,
1347
1468
  user_id=user_id,
1348
1469
  debug_mode=debug_mode,
1470
+ background_tasks=background_tasks,
1349
1471
  **kwargs,
1350
1472
  )
1351
1473
  # Consume the generator without yielding
@@ -1458,13 +1580,14 @@ class Team:
1458
1580
  session=session,
1459
1581
  user_id=user_id,
1460
1582
  debug_mode=debug_mode,
1583
+ background_tasks=background_tasks,
1461
1584
  **kwargs,
1462
1585
  )
1463
1586
  deque(iterator, maxlen=0)
1464
1587
  raise_if_cancelled(run_response.run_id) # type: ignore
1465
1588
 
1466
1589
  # 11. Wait for background memory creation
1467
- wait_for_background_tasks(memory_future=memory_future)
1590
+ wait_for_open_threads(memory_future=memory_future)
1468
1591
 
1469
1592
  raise_if_cancelled(run_response.run_id) # type: ignore
1470
1593
 
@@ -1517,6 +1640,7 @@ class Team:
1517
1640
  stream_events: bool = False,
1518
1641
  yield_run_output: bool = False,
1519
1642
  debug_mode: Optional[bool] = None,
1643
+ background_tasks: Optional[Any] = None,
1520
1644
  **kwargs: Any,
1521
1645
  ) -> Iterator[Union[TeamRunOutputEvent, RunOutputEvent, TeamRunOutput]]:
1522
1646
  """Run the Team and return the response iterator.
@@ -1533,8 +1657,6 @@ class Team:
1533
1657
  9. Create session summary
1534
1658
  10. Cleanup and store (scrub, add to session, calculate metrics, save session)
1535
1659
  """
1536
- # Register run for cancellation tracking
1537
- register_run(run_response.run_id) # type: ignore
1538
1660
 
1539
1661
  # 1. Execute pre-hooks
1540
1662
  run_input = cast(TeamRunInput, run_response.input)
@@ -1550,6 +1672,7 @@ class Team:
1550
1672
  user_id=user_id,
1551
1673
  debug_mode=debug_mode,
1552
1674
  stream_events=stream_events,
1675
+ background_tasks=background_tasks,
1553
1676
  **kwargs,
1554
1677
  )
1555
1678
  for pre_hook_event in pre_hook_iterator:
@@ -1702,12 +1825,13 @@ class Team:
1702
1825
  user_id=user_id,
1703
1826
  debug_mode=debug_mode,
1704
1827
  stream_events=stream_events,
1828
+ background_tasks=background_tasks,
1705
1829
  **kwargs,
1706
1830
  )
1707
1831
  raise_if_cancelled(run_response.run_id) # type: ignore
1708
1832
 
1709
1833
  # 8. Wait for background memory creation
1710
- yield from wait_for_background_tasks_stream(
1834
+ yield from wait_for_thread_tasks_stream(
1711
1835
  run_response=run_response,
1712
1836
  memory_future=memory_future,
1713
1837
  stream_events=stream_events,
@@ -1879,6 +2003,10 @@ class Team:
1879
2003
  if self._has_async_db():
1880
2004
  raise Exception("run() is not supported with an async DB. Please use arun() instead.")
1881
2005
 
2006
+ # Create a run_id for this specific run and register immediately for cancellation tracking
2007
+ run_id = str(uuid4())
2008
+ register_run(run_id)
2009
+
1882
2010
  # Initialize Team
1883
2011
  self.initialize_team(debug_mode=debug_mode)
1884
2012
 
@@ -1894,8 +2022,11 @@ class Team:
1894
2022
  stacklevel=2,
1895
2023
  )
1896
2024
 
1897
- # Create a run_id for this specific run
1898
- run_id = str(uuid4())
2025
+ background_tasks = kwargs.pop("background_tasks", None)
2026
+ if background_tasks is not None:
2027
+ from fastapi import BackgroundTasks
2028
+
2029
+ background_tasks: BackgroundTasks = background_tasks # type: ignore
1899
2030
 
1900
2031
  # Validate input against input_schema if provided
1901
2032
  validated_input = self._validate_input(input)
@@ -2051,6 +2182,7 @@ class Team:
2051
2182
  stream_events=stream_events,
2052
2183
  yield_run_output=yield_run_output,
2053
2184
  debug_mode=debug_mode,
2185
+ background_tasks=background_tasks,
2054
2186
  **kwargs,
2055
2187
  )
2056
2188
 
@@ -2066,6 +2198,7 @@ class Team:
2066
2198
  add_session_state_to_context=add_session_state,
2067
2199
  response_format=response_format,
2068
2200
  debug_mode=debug_mode,
2201
+ background_tasks=background_tasks,
2069
2202
  **kwargs,
2070
2203
  )
2071
2204
 
@@ -2123,6 +2256,7 @@ class Team:
2123
2256
  add_session_state_to_context: Optional[bool] = None,
2124
2257
  add_history_to_context: Optional[bool] = None,
2125
2258
  debug_mode: Optional[bool] = None,
2259
+ background_tasks: Optional[Any] = None,
2126
2260
  **kwargs: Any,
2127
2261
  ) -> TeamRunOutput:
2128
2262
  """Run the Team and return the response.
@@ -2146,8 +2280,6 @@ class Team:
2146
2280
  """
2147
2281
  log_debug(f"Team Run Start: {run_response.run_id}", center=True)
2148
2282
 
2149
- register_run(run_response.run_id) # type: ignore
2150
-
2151
2283
  if run_context.dependencies is not None:
2152
2284
  await self._aresolve_run_dependencies(run_context=run_context)
2153
2285
 
@@ -2184,6 +2316,7 @@ class Team:
2184
2316
  session=team_session,
2185
2317
  user_id=user_id,
2186
2318
  debug_mode=debug_mode,
2319
+ background_tasks=background_tasks,
2187
2320
  **kwargs,
2188
2321
  )
2189
2322
 
@@ -2241,9 +2374,6 @@ class Team:
2241
2374
  log_debug("Starting memory creation in background task.")
2242
2375
  memory_task = asyncio.create_task(self._amake_memories(run_messages=run_messages, user_id=user_id))
2243
2376
 
2244
- # Register run for cancellation tracking
2245
- register_run(run_response.run_id) # type: ignore
2246
-
2247
2377
  try:
2248
2378
  raise_if_cancelled(run_response.run_id) # type: ignore
2249
2379
  # 7. Reason about the task if reasoning is enabled
@@ -2299,6 +2429,7 @@ class Team:
2299
2429
  session=team_session,
2300
2430
  user_id=user_id,
2301
2431
  debug_mode=debug_mode,
2432
+ background_tasks=background_tasks,
2302
2433
  **kwargs,
2303
2434
  ):
2304
2435
  pass
@@ -2306,7 +2437,7 @@ class Team:
2306
2437
  raise_if_cancelled(run_response.run_id) # type: ignore
2307
2438
 
2308
2439
  # 13. Wait for background memory creation
2309
- await await_for_background_tasks(memory_task=memory_task)
2440
+ await await_for_open_threads(memory_task=memory_task)
2310
2441
 
2311
2442
  raise_if_cancelled(run_response.run_id) # type: ignore
2312
2443
  # 14. Create session summary
@@ -2367,6 +2498,7 @@ class Team:
2367
2498
  add_session_state_to_context: Optional[bool] = None,
2368
2499
  add_history_to_context: Optional[bool] = None,
2369
2500
  debug_mode: Optional[bool] = None,
2501
+ background_tasks: Optional[Any] = None,
2370
2502
  **kwargs: Any,
2371
2503
  ) -> AsyncIterator[Union[TeamRunOutputEvent, RunOutputEvent, TeamRunOutput]]:
2372
2504
  """Run the Team and return the response.
@@ -2425,6 +2557,7 @@ class Team:
2425
2557
  user_id=user_id,
2426
2558
  debug_mode=debug_mode,
2427
2559
  stream_events=stream_events,
2560
+ background_tasks=background_tasks,
2428
2561
  **kwargs,
2429
2562
  )
2430
2563
  async for pre_hook_event in pre_hook_iterator:
@@ -2477,9 +2610,6 @@ class Team:
2477
2610
  log_debug("Starting memory creation in background task.")
2478
2611
  memory_task = asyncio.create_task(self._amake_memories(run_messages=run_messages, user_id=user_id))
2479
2612
 
2480
- # Register run for cancellation tracking
2481
- register_run(run_response.run_id) # type: ignore
2482
-
2483
2613
  try:
2484
2614
  # Considering both stream_events and stream_intermediate_steps (deprecated)
2485
2615
  stream_events = stream_events or stream_intermediate_steps
@@ -2579,13 +2709,14 @@ class Team:
2579
2709
  user_id=user_id,
2580
2710
  debug_mode=debug_mode,
2581
2711
  stream_events=stream_events,
2712
+ background_tasks=background_tasks,
2582
2713
  **kwargs,
2583
2714
  ):
2584
2715
  yield event
2585
2716
 
2586
2717
  raise_if_cancelled(run_response.run_id) # type: ignore
2587
2718
  # 11. Wait for background memory creation
2588
- async for event in await_for_background_tasks_stream(
2719
+ async for event in await_for_thread_tasks_stream(
2589
2720
  run_response=run_response,
2590
2721
  memory_task=memory_task,
2591
2722
  stream_events=stream_events,
@@ -2767,6 +2898,10 @@ class Team:
2767
2898
  ) -> Union[TeamRunOutput, AsyncIterator[Union[RunOutputEvent, TeamRunOutputEvent]]]:
2768
2899
  """Run the Team asynchronously and return the response."""
2769
2900
 
2901
+ # Create a run_id for this specific run and register immediately for cancellation tracking
2902
+ run_id = str(uuid4())
2903
+ register_run(run_id)
2904
+
2770
2905
  if (add_history_to_context or self.add_history_to_context) and not self.db and not self.parent_team_id:
2771
2906
  log_warning(
2772
2907
  "add_history_to_context is True, but no database has been assigned to the team. History will not be added to the context."
@@ -2779,8 +2914,11 @@ class Team:
2779
2914
  stacklevel=2,
2780
2915
  )
2781
2916
 
2782
- # Create a run_id for this specific run
2783
- run_id = str(uuid4())
2917
+ background_tasks = kwargs.pop("background_tasks", None)
2918
+ if background_tasks is not None:
2919
+ from fastapi import BackgroundTasks
2920
+
2921
+ background_tasks: BackgroundTasks = background_tasks # type: ignore
2784
2922
 
2785
2923
  # Validate input against input_schema if provided
2786
2924
  validated_input = self._validate_input(input)
@@ -2927,6 +3065,7 @@ class Team:
2927
3065
  stream_events=stream_events,
2928
3066
  yield_run_output=yield_run_output,
2929
3067
  debug_mode=debug_mode,
3068
+ background_tasks=background_tasks,
2930
3069
  **kwargs,
2931
3070
  )
2932
3071
  return response_iterator # type: ignore
@@ -2942,6 +3081,7 @@ class Team:
2942
3081
  add_session_state_to_context=add_session_state,
2943
3082
  response_format=response_format,
2944
3083
  debug_mode=debug_mode,
3084
+ background_tasks=background_tasks,
2945
3085
  **kwargs,
2946
3086
  )
2947
3087
 
@@ -3047,7 +3187,9 @@ class Team:
3047
3187
  run_response.messages = messages_for_run_response
3048
3188
 
3049
3189
  # Update the TeamRunOutput metrics
3050
- run_response.metrics = self._calculate_metrics(messages_for_run_response)
3190
+ run_response.metrics = self._calculate_metrics(
3191
+ messages_for_run_response, current_run_metrics=run_response.metrics
3192
+ )
3051
3193
 
3052
3194
  if model_response.tool_executions:
3053
3195
  for tool_call in model_response.tool_executions:
@@ -3141,7 +3283,9 @@ class Team:
3141
3283
  # Update the TeamRunOutput messages
3142
3284
  run_response.messages = messages_for_run_response
3143
3285
  # Update the TeamRunOutput metrics
3144
- run_response.metrics = self._calculate_metrics(messages_for_run_response)
3286
+ run_response.metrics = self._calculate_metrics(
3287
+ messages_for_run_response, current_run_metrics=run_response.metrics
3288
+ )
3145
3289
 
3146
3290
  # Update the run_response audio if streaming
3147
3291
  if full_model_response.audio is not None:
@@ -3225,7 +3369,9 @@ class Team:
3225
3369
  # Update the TeamRunOutput messages
3226
3370
  run_response.messages = messages_for_run_response
3227
3371
  # Update the TeamRunOutput metrics
3228
- run_response.metrics = self._calculate_metrics(messages_for_run_response)
3372
+ run_response.metrics = self._calculate_metrics(
3373
+ messages_for_run_response, current_run_metrics=run_response.metrics
3374
+ )
3229
3375
 
3230
3376
  if stream_events and reasoning_state["reasoning_started"]:
3231
3377
  all_reasoning_steps: List[ReasoningStep] = []
@@ -3600,7 +3746,7 @@ class Team:
3600
3746
  session.upsert_run(run_response=run_response)
3601
3747
 
3602
3748
  # Calculate session metrics
3603
- self._update_session_metrics(session=session)
3749
+ self._update_session_metrics(session=session, run_response=run_response)
3604
3750
 
3605
3751
  # Save session to memory
3606
3752
  self.save_session(session=session)
@@ -3617,7 +3763,7 @@ class Team:
3617
3763
  session.upsert_run(run_response=run_response)
3618
3764
 
3619
3765
  # Calculate session metrics
3620
- self._update_session_metrics(session=session)
3766
+ self._update_session_metrics(session=session, run_response=run_response)
3621
3767
 
3622
3768
  # Save session to memory
3623
3769
  await self.asave_session(session=session)
@@ -3957,7 +4103,9 @@ class Team:
3957
4103
  # Update the RunResponse messages
3958
4104
  run_response.messages = messages_for_run_response
3959
4105
  # Update the RunResponse metrics
3960
- run_response.metrics = self._calculate_metrics(messages_for_run_response)
4106
+ run_response.metrics = self._calculate_metrics(
4107
+ messages_for_run_response, current_run_metrics=run_response.metrics
4108
+ )
3961
4109
 
3962
4110
  async def _agenerate_response_with_output_model(
3963
4111
  self, model_response: ModelResponse, run_messages: RunMessages
@@ -4022,7 +4170,9 @@ class Team:
4022
4170
  # Update the RunResponse messages
4023
4171
  run_response.messages = messages_for_run_response
4024
4172
  # Update the RunResponse metrics
4025
- run_response.metrics = self._calculate_metrics(messages_for_run_response)
4173
+ run_response.metrics = self._calculate_metrics(
4174
+ messages_for_run_response, current_run_metrics=run_response.metrics
4175
+ )
4026
4176
 
4027
4177
  def _handle_event(
4028
4178
  self,
@@ -4458,42 +4608,41 @@ class Team:
4458
4608
  async for item in reason_generator:
4459
4609
  yield item
4460
4610
 
4461
- def _calculate_session_metrics(self, messages: List[Message]) -> Metrics:
4462
- """Sum the metrics of the given messages into a Metrics object"""
4463
- session_metrics = Metrics()
4464
- assistant_message_role = self.model.assistant_message_role if self.model is not None else "assistant"
4465
-
4466
- # Get metrics of the team leader's messages
4467
- for m in messages:
4468
- if m.role == assistant_message_role and m.metrics is not None:
4469
- session_metrics += m.metrics
4470
-
4471
- return session_metrics
4472
-
4473
- def _calculate_metrics(self, messages: List[Message]) -> Metrics:
4474
- metrics = Metrics()
4611
+ def _calculate_metrics(self, messages: List[Message], current_run_metrics: Optional[Metrics] = None) -> Metrics:
4612
+ metrics = current_run_metrics or Metrics()
4475
4613
  assistant_message_role = self.model.assistant_message_role if self.model is not None else "assistant"
4476
4614
 
4477
4615
  for m in messages:
4478
4616
  if m.role == assistant_message_role and m.metrics is not None and m.from_history is False:
4479
4617
  metrics += m.metrics
4480
4618
 
4481
- return metrics
4619
+ # If the run metrics were already initialized, keep the time related metrics
4620
+ if current_run_metrics is not None:
4621
+ metrics.timer = current_run_metrics.timer
4622
+ metrics.duration = current_run_metrics.duration
4623
+ metrics.time_to_first_token = current_run_metrics.time_to_first_token
4482
4624
 
4483
- def _update_session_metrics(self, session: TeamSession):
4484
- """Calculate session metrics"""
4625
+ return metrics
4485
4626
 
4486
- session_messages: List[Message] = []
4487
- for run in session.runs or []:
4488
- if run.messages is not None:
4489
- for m in run.messages:
4490
- # Skipping messages from history to avoid duplicates
4491
- if not m.from_history:
4492
- session_messages.append(m)
4627
+ def _get_session_metrics(self, session: TeamSession) -> Metrics:
4628
+ # Get the session_metrics from the database
4629
+ if session.session_data is not None and "session_metrics" in session.session_data:
4630
+ session_metrics_from_db = session.session_data.get("session_metrics")
4631
+ if session_metrics_from_db is not None:
4632
+ if isinstance(session_metrics_from_db, dict):
4633
+ return Metrics(**session_metrics_from_db)
4634
+ elif isinstance(session_metrics_from_db, Metrics):
4635
+ return session_metrics_from_db
4493
4636
 
4494
- # Calculate initial metrics
4495
- session_metrics = self._calculate_session_metrics(session_messages)
4637
+ return Metrics()
4496
4638
 
4639
+ def _update_session_metrics(self, session: TeamSession, run_response: TeamRunOutput):
4640
+ """Calculate session metrics"""
4641
+ session_metrics = self._get_session_metrics(session=session)
4642
+ # Add the metrics for the current run to the session metrics
4643
+ if run_response.metrics is not None:
4644
+ session_metrics += run_response.metrics
4645
+ session_metrics.time_to_first_token = None
4497
4646
  if session.session_data is not None:
4498
4647
  session.session_data["session_metrics"] = session_metrics
4499
4648