agno 2.3.10__py3-none-any.whl → 2.3.11__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.
@@ -2434,86 +2434,110 @@ class AsyncMySQLDb(AsyncBaseDb):
2434
2434
  # Fallback if spans table doesn't exist
2435
2435
  return select(table, literal(0).label("total_spans"), literal(0).label("error_count"))
2436
2436
 
2437
- async def create_trace(self, trace: "Trace") -> None:
2438
- """Create a single trace record in the database.
2437
+ def _get_trace_component_level_expr(self, workflow_id_col, team_id_col, agent_id_col, name_col):
2438
+ """Build a SQL CASE expression that returns the component level for a trace.
2439
+
2440
+ Component levels (higher = more important):
2441
+ - 3: Workflow root (.run or .arun with workflow_id)
2442
+ - 2: Team root (.run or .arun with team_id)
2443
+ - 1: Agent root (.run or .arun with agent_id)
2444
+ - 0: Child span (not a root)
2445
+
2446
+ Args:
2447
+ workflow_id_col: SQL column/expression for workflow_id
2448
+ team_id_col: SQL column/expression for team_id
2449
+ agent_id_col: SQL column/expression for agent_id
2450
+ name_col: SQL column/expression for name
2451
+
2452
+ Returns:
2453
+ SQLAlchemy CASE expression returning the component level as an integer.
2454
+ """
2455
+ from sqlalchemy import and_, case, or_
2456
+
2457
+ is_root_name = or_(name_col.like("%.run%"), name_col.like("%.arun%"))
2458
+
2459
+ return case(
2460
+ # Workflow root (level 3)
2461
+ (and_(workflow_id_col.isnot(None), is_root_name), 3),
2462
+ # Team root (level 2)
2463
+ (and_(team_id_col.isnot(None), is_root_name), 2),
2464
+ # Agent root (level 1)
2465
+ (and_(agent_id_col.isnot(None), is_root_name), 1),
2466
+ # Child span or unknown (level 0)
2467
+ else_=0,
2468
+ )
2469
+
2470
+ async def upsert_trace(self, trace: "Trace") -> None:
2471
+ """Create or update a single trace record in the database.
2472
+
2473
+ Uses INSERT ... ON DUPLICATE KEY UPDATE (upsert) to handle concurrent inserts
2474
+ atomically and avoid race conditions.
2439
2475
 
2440
2476
  Args:
2441
2477
  trace: The Trace object to store (one per trace_id).
2442
2478
  """
2479
+ from sqlalchemy import case
2480
+
2443
2481
  try:
2444
2482
  table = await self._get_table(table_type="traces", create_table_if_not_found=True)
2445
2483
  if table is None:
2446
2484
  return
2447
2485
 
2448
- async with self.async_session_factory() as sess, sess.begin():
2449
- # Check if trace exists
2450
- result = await sess.execute(select(table).where(table.c.trace_id == trace.trace_id))
2451
- existing = result.fetchone()
2452
-
2453
- if existing:
2454
- # workflow (level 3) > team (level 2) > agent (level 1) > child/unknown (level 0)
2455
-
2456
- def get_component_level(workflow_id, team_id, agent_id, name):
2457
- # Check if name indicates a root span
2458
- is_root_name = ".run" in name or ".arun" in name
2459
-
2460
- if not is_root_name:
2461
- return 0 # Child span (not a root)
2462
- elif workflow_id:
2463
- return 3 # Workflow root
2464
- elif team_id:
2465
- return 2 # Team root
2466
- elif agent_id:
2467
- return 1 # Agent root
2468
- else:
2469
- return 0 # Unknown
2470
-
2471
- existing_level = get_component_level(
2472
- existing.workflow_id, existing.team_id, existing.agent_id, existing.name
2473
- )
2474
- new_level = get_component_level(trace.workflow_id, trace.team_id, trace.agent_id, trace.name)
2475
-
2476
- # Only update name if new trace is from a higher or equal level
2477
- should_update_name = new_level > existing_level
2478
-
2479
- # Parse existing start_time to calculate correct duration
2480
- existing_start_time_str = existing.start_time
2481
- if isinstance(existing_start_time_str, str):
2482
- existing_start_time = datetime.fromisoformat(existing_start_time_str.replace("Z", "+00:00"))
2483
- else:
2484
- existing_start_time = trace.start_time
2485
-
2486
- recalculated_duration_ms = int((trace.end_time - existing_start_time).total_seconds() * 1000)
2486
+ trace_dict = trace.to_dict()
2487
+ trace_dict.pop("total_spans", None)
2488
+ trace_dict.pop("error_count", None)
2487
2489
 
2488
- update_values = {
2489
- "end_time": trace.end_time.isoformat(),
2490
- "duration_ms": recalculated_duration_ms,
2491
- "status": trace.status,
2492
- "name": trace.name if should_update_name else existing.name,
2493
- }
2490
+ async with self.async_session_factory() as sess, sess.begin():
2491
+ # Use upsert to handle concurrent inserts atomically
2492
+ # On conflict, update fields while preserving existing non-null context values
2493
+ # and keeping the earliest start_time
2494
+ insert_stmt = mysql.insert(table).values(trace_dict)
2495
+
2496
+ # Build component level expressions for comparing trace priority
2497
+ new_level = self._get_trace_component_level_expr(
2498
+ insert_stmt.inserted.workflow_id,
2499
+ insert_stmt.inserted.team_id,
2500
+ insert_stmt.inserted.agent_id,
2501
+ insert_stmt.inserted.name,
2502
+ )
2503
+ existing_level = self._get_trace_component_level_expr(
2504
+ table.c.workflow_id,
2505
+ table.c.team_id,
2506
+ table.c.agent_id,
2507
+ table.c.name,
2508
+ )
2494
2509
 
2495
- # Update context fields ONLY if new value is not None (preserve non-null values)
2496
- if trace.run_id is not None:
2497
- update_values["run_id"] = trace.run_id
2498
- if trace.session_id is not None:
2499
- update_values["session_id"] = trace.session_id
2500
- if trace.user_id is not None:
2501
- update_values["user_id"] = trace.user_id
2502
- if trace.agent_id is not None:
2503
- update_values["agent_id"] = trace.agent_id
2504
- if trace.team_id is not None:
2505
- update_values["team_id"] = trace.team_id
2506
- if trace.workflow_id is not None:
2507
- update_values["workflow_id"] = trace.workflow_id
2508
-
2509
- stmt = update(table).where(table.c.trace_id == trace.trace_id).values(**update_values)
2510
- await sess.execute(stmt)
2511
- else:
2512
- trace_dict = trace.to_dict()
2513
- trace_dict.pop("total_spans", None)
2514
- trace_dict.pop("error_count", None)
2515
- stmt = mysql.insert(table).values(trace_dict)
2516
- await sess.execute(stmt)
2510
+ # Build the ON DUPLICATE KEY UPDATE clause
2511
+ # Use LEAST for start_time, GREATEST for end_time to capture full trace duration
2512
+ # MySQL stores timestamps as ISO strings, so string comparison works for ISO format
2513
+ # Duration is calculated using TIMESTAMPDIFF in microseconds then converted to ms
2514
+ upsert_stmt = insert_stmt.on_duplicate_key_update(
2515
+ end_time=func.greatest(table.c.end_time, insert_stmt.inserted.end_time),
2516
+ start_time=func.least(table.c.start_time, insert_stmt.inserted.start_time),
2517
+ # Calculate duration in milliseconds using TIMESTAMPDIFF
2518
+ # TIMESTAMPDIFF(MICROSECOND, start, end) / 1000 gives milliseconds
2519
+ duration_ms=func.timestampdiff(
2520
+ text("MICROSECOND"),
2521
+ func.least(table.c.start_time, insert_stmt.inserted.start_time),
2522
+ func.greatest(table.c.end_time, insert_stmt.inserted.end_time),
2523
+ )
2524
+ / 1000,
2525
+ status=insert_stmt.inserted.status,
2526
+ # Update name only if new trace is from a higher-level component
2527
+ # Priority: workflow (3) > team (2) > agent (1) > child spans (0)
2528
+ name=case(
2529
+ (new_level > existing_level, insert_stmt.inserted.name),
2530
+ else_=table.c.name,
2531
+ ),
2532
+ # Preserve existing non-null context values using COALESCE
2533
+ run_id=func.coalesce(insert_stmt.inserted.run_id, table.c.run_id),
2534
+ session_id=func.coalesce(insert_stmt.inserted.session_id, table.c.session_id),
2535
+ user_id=func.coalesce(insert_stmt.inserted.user_id, table.c.user_id),
2536
+ agent_id=func.coalesce(insert_stmt.inserted.agent_id, table.c.agent_id),
2537
+ team_id=func.coalesce(insert_stmt.inserted.team_id, table.c.team_id),
2538
+ workflow_id=func.coalesce(insert_stmt.inserted.workflow_id, table.c.workflow_id),
2539
+ )
2540
+ await sess.execute(upsert_stmt)
2517
2541
 
2518
2542
  except Exception as e:
2519
2543
  log_error(f"Error creating trace: {e}")
agno/db/mysql/mysql.py CHANGED
@@ -2452,85 +2452,110 @@ class MySQLDb(BaseDb):
2452
2452
  # Fallback if spans table doesn't exist
2453
2453
  return select(table, literal(0).label("total_spans"), literal(0).label("error_count"))
2454
2454
 
2455
- def create_trace(self, trace: "Trace") -> None:
2456
- """Create a single trace record in the database.
2455
+ def _get_trace_component_level_expr(self, workflow_id_col, team_id_col, agent_id_col, name_col):
2456
+ """Build a SQL CASE expression that returns the component level for a trace.
2457
+
2458
+ Component levels (higher = more important):
2459
+ - 3: Workflow root (.run or .arun with workflow_id)
2460
+ - 2: Team root (.run or .arun with team_id)
2461
+ - 1: Agent root (.run or .arun with agent_id)
2462
+ - 0: Child span (not a root)
2463
+
2464
+ Args:
2465
+ workflow_id_col: SQL column/expression for workflow_id
2466
+ team_id_col: SQL column/expression for team_id
2467
+ agent_id_col: SQL column/expression for agent_id
2468
+ name_col: SQL column/expression for name
2469
+
2470
+ Returns:
2471
+ SQLAlchemy CASE expression returning the component level as an integer.
2472
+ """
2473
+ from sqlalchemy import and_, case, or_
2474
+
2475
+ is_root_name = or_(name_col.like("%.run%"), name_col.like("%.arun%"))
2476
+
2477
+ return case(
2478
+ # Workflow root (level 3)
2479
+ (and_(workflow_id_col.isnot(None), is_root_name), 3),
2480
+ # Team root (level 2)
2481
+ (and_(team_id_col.isnot(None), is_root_name), 2),
2482
+ # Agent root (level 1)
2483
+ (and_(agent_id_col.isnot(None), is_root_name), 1),
2484
+ # Child span or unknown (level 0)
2485
+ else_=0,
2486
+ )
2487
+
2488
+ def upsert_trace(self, trace: "Trace") -> None:
2489
+ """Create or update a single trace record in the database.
2490
+
2491
+ Uses INSERT ... ON DUPLICATE KEY UPDATE (upsert) to handle concurrent inserts
2492
+ atomically and avoid race conditions.
2457
2493
 
2458
2494
  Args:
2459
2495
  trace: The Trace object to store (one per trace_id).
2460
2496
  """
2497
+ from sqlalchemy import case
2498
+
2461
2499
  try:
2462
2500
  table = self._get_table(table_type="traces", create_table_if_not_found=True)
2463
2501
  if table is None:
2464
2502
  return
2465
2503
 
2466
- with self.Session() as sess, sess.begin():
2467
- # Check if trace exists
2468
- existing = sess.execute(select(table).where(table.c.trace_id == trace.trace_id)).fetchone()
2469
-
2470
- if existing:
2471
- # workflow (level 3) > team (level 2) > agent (level 1) > child/unknown (level 0)
2472
-
2473
- def get_component_level(workflow_id, team_id, agent_id, name):
2474
- # Check if name indicates a root span
2475
- is_root_name = ".run" in name or ".arun" in name
2476
-
2477
- if not is_root_name:
2478
- return 0 # Child span (not a root)
2479
- elif workflow_id:
2480
- return 3 # Workflow root
2481
- elif team_id:
2482
- return 2 # Team root
2483
- elif agent_id:
2484
- return 1 # Agent root
2485
- else:
2486
- return 0 # Unknown
2487
-
2488
- existing_level = get_component_level(
2489
- existing.workflow_id, existing.team_id, existing.agent_id, existing.name
2490
- )
2491
- new_level = get_component_level(trace.workflow_id, trace.team_id, trace.agent_id, trace.name)
2492
-
2493
- # Only update name if new trace is from a higher or equal level
2494
- should_update_name = new_level > existing_level
2504
+ trace_dict = trace.to_dict()
2505
+ trace_dict.pop("total_spans", None)
2506
+ trace_dict.pop("error_count", None)
2495
2507
 
2496
- # Parse existing start_time to calculate correct duration
2497
- existing_start_time_str = existing.start_time
2498
- if isinstance(existing_start_time_str, str):
2499
- existing_start_time = datetime.fromisoformat(existing_start_time_str.replace("Z", "+00:00"))
2500
- else:
2501
- existing_start_time = trace.start_time
2502
-
2503
- recalculated_duration_ms = int((trace.end_time - existing_start_time).total_seconds() * 1000)
2504
-
2505
- update_values = {
2506
- "end_time": trace.end_time.isoformat(),
2507
- "duration_ms": recalculated_duration_ms,
2508
- "status": trace.status,
2509
- "name": trace.name if should_update_name else existing.name,
2510
- }
2508
+ with self.Session() as sess, sess.begin():
2509
+ # Use upsert to handle concurrent inserts atomically
2510
+ # On conflict, update fields while preserving existing non-null context values
2511
+ # and keeping the earliest start_time
2512
+ insert_stmt = mysql.insert(table).values(trace_dict)
2513
+
2514
+ # Build component level expressions for comparing trace priority
2515
+ new_level = self._get_trace_component_level_expr(
2516
+ insert_stmt.inserted.workflow_id,
2517
+ insert_stmt.inserted.team_id,
2518
+ insert_stmt.inserted.agent_id,
2519
+ insert_stmt.inserted.name,
2520
+ )
2521
+ existing_level = self._get_trace_component_level_expr(
2522
+ table.c.workflow_id,
2523
+ table.c.team_id,
2524
+ table.c.agent_id,
2525
+ table.c.name,
2526
+ )
2511
2527
 
2512
- # Update context fields ONLY if new value is not None (preserve non-null values)
2513
- if trace.run_id is not None:
2514
- update_values["run_id"] = trace.run_id
2515
- if trace.session_id is not None:
2516
- update_values["session_id"] = trace.session_id
2517
- if trace.user_id is not None:
2518
- update_values["user_id"] = trace.user_id
2519
- if trace.agent_id is not None:
2520
- update_values["agent_id"] = trace.agent_id
2521
- if trace.team_id is not None:
2522
- update_values["team_id"] = trace.team_id
2523
- if trace.workflow_id is not None:
2524
- update_values["workflow_id"] = trace.workflow_id
2525
-
2526
- stmt = update(table).where(table.c.trace_id == trace.trace_id).values(**update_values)
2527
- sess.execute(stmt)
2528
- else:
2529
- trace_dict = trace.to_dict()
2530
- trace_dict.pop("total_spans", None)
2531
- trace_dict.pop("error_count", None)
2532
- stmt = mysql.insert(table).values(trace_dict)
2533
- sess.execute(stmt)
2528
+ # Build the ON DUPLICATE KEY UPDATE clause
2529
+ # Use LEAST for start_time, GREATEST for end_time to capture full trace duration
2530
+ # MySQL stores timestamps as ISO strings, so string comparison works for ISO format
2531
+ # Duration is calculated using TIMESTAMPDIFF in microseconds then converted to ms
2532
+ upsert_stmt = insert_stmt.on_duplicate_key_update(
2533
+ end_time=func.greatest(table.c.end_time, insert_stmt.inserted.end_time),
2534
+ start_time=func.least(table.c.start_time, insert_stmt.inserted.start_time),
2535
+ # Calculate duration in milliseconds using TIMESTAMPDIFF
2536
+ # TIMESTAMPDIFF(MICROSECOND, start, end) / 1000 gives milliseconds
2537
+ duration_ms=func.timestampdiff(
2538
+ text("MICROSECOND"),
2539
+ func.least(table.c.start_time, insert_stmt.inserted.start_time),
2540
+ func.greatest(table.c.end_time, insert_stmt.inserted.end_time),
2541
+ )
2542
+ / 1000,
2543
+ status=insert_stmt.inserted.status,
2544
+ # Update name only if new trace is from a higher-level component
2545
+ # Priority: workflow (3) > team (2) > agent (1) > child spans (0)
2546
+ name=case(
2547
+ (new_level > existing_level, insert_stmt.inserted.name),
2548
+ else_=table.c.name,
2549
+ ),
2550
+ # Preserve existing non-null context values using COALESCE
2551
+ run_id=func.coalesce(insert_stmt.inserted.run_id, table.c.run_id),
2552
+ session_id=func.coalesce(insert_stmt.inserted.session_id, table.c.session_id),
2553
+ user_id=func.coalesce(insert_stmt.inserted.user_id, table.c.user_id),
2554
+ agent_id=func.coalesce(insert_stmt.inserted.agent_id, table.c.agent_id),
2555
+ team_id=func.coalesce(insert_stmt.inserted.team_id, table.c.team_id),
2556
+ workflow_id=func.coalesce(insert_stmt.inserted.workflow_id, table.c.workflow_id),
2557
+ )
2558
+ sess.execute(upsert_stmt)
2534
2559
 
2535
2560
  except Exception as e:
2536
2561
  log_error(f"Error creating trace: {e}")
@@ -30,8 +30,9 @@ from agno.session import AgentSession, Session, TeamSession, WorkflowSession
30
30
  from agno.utils.log import log_debug, log_error, log_info, log_warning
31
31
 
32
32
  try:
33
- from sqlalchemy import Index, String, Table, UniqueConstraint, func, update
33
+ from sqlalchemy import Index, String, Table, UniqueConstraint, and_, case, func, or_, update
34
34
  from sqlalchemy.dialects import postgresql
35
+ from sqlalchemy.dialects.postgresql import TIMESTAMP
35
36
  from sqlalchemy.exc import ProgrammingError
36
37
  from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker, create_async_engine
37
38
  from sqlalchemy.schema import Column, MetaData
@@ -1094,7 +1095,7 @@ class AsyncPostgresDb(AsyncBaseDb):
1094
1095
  await sess.execute(table.delete())
1095
1096
 
1096
1097
  except Exception as e:
1097
- log_warning(f"Exception deleting all cultural knowledge: {e}")
1098
+ log_error(f"Exception deleting all cultural knowledge: {e}")
1098
1099
 
1099
1100
  async def delete_cultural_knowledge(self, id: str) -> None:
1100
1101
  """Delete cultural knowledge by ID.
@@ -1113,8 +1114,7 @@ class AsyncPostgresDb(AsyncBaseDb):
1113
1114
  await sess.execute(stmt)
1114
1115
 
1115
1116
  except Exception as e:
1116
- log_warning(f"Exception deleting cultural knowledge: {e}")
1117
- raise e
1117
+ log_error(f"Exception deleting cultural knowledge: {e}")
1118
1118
 
1119
1119
  async def get_cultural_knowledge(
1120
1120
  self, id: str, deserialize: Optional[bool] = True
@@ -1150,8 +1150,8 @@ class AsyncPostgresDb(AsyncBaseDb):
1150
1150
  return deserialize_cultural_knowledge(db_row)
1151
1151
 
1152
1152
  except Exception as e:
1153
- log_warning(f"Exception reading cultural knowledge: {e}")
1154
- raise e
1153
+ log_error(f"Exception reading cultural knowledge: {e}")
1154
+ return None
1155
1155
 
1156
1156
  async def get_all_cultural_knowledge(
1157
1157
  self,
@@ -1185,7 +1185,7 @@ class AsyncPostgresDb(AsyncBaseDb):
1185
1185
  Exception: If an error occurs during retrieval.
1186
1186
  """
1187
1187
  try:
1188
- table = await self._get_table(table_type="culture")
1188
+ table = await self._get_table(table_type="culture", create_table_if_not_found=True)
1189
1189
 
1190
1190
  async with self.async_session_factory() as sess:
1191
1191
  # Build query with filters
@@ -1223,8 +1223,8 @@ class AsyncPostgresDb(AsyncBaseDb):
1223
1223
  return [deserialize_cultural_knowledge(row) for row in db_rows]
1224
1224
 
1225
1225
  except Exception as e:
1226
- log_warning(f"Exception reading all cultural knowledge: {e}")
1227
- raise e
1226
+ log_error(f"Exception reading all cultural knowledge: {e}")
1227
+ return [] if deserialize else ([], 0)
1228
1228
 
1229
1229
  async def upsert_cultural_knowledge(
1230
1230
  self, cultural_knowledge: CulturalKnowledge, deserialize: Optional[bool] = True
@@ -2121,84 +2121,110 @@ class AsyncPostgresDb(AsyncBaseDb):
2121
2121
  # Fallback if spans table doesn't exist
2122
2122
  return select(table, literal(0).label("total_spans"), literal(0).label("error_count"))
2123
2123
 
2124
- async def create_trace(self, trace: "Trace") -> None:
2125
- """Create a single trace record in the database.
2124
+ def _get_trace_component_level_expr(self, workflow_id_col, team_id_col, agent_id_col, name_col):
2125
+ """Build a SQL CASE expression that returns the component level for a trace.
2126
+
2127
+ Component levels (higher = more important):
2128
+ - 3: Workflow root (.run or .arun with workflow_id)
2129
+ - 2: Team root (.run or .arun with team_id)
2130
+ - 1: Agent root (.run or .arun with agent_id)
2131
+ - 0: Child span (not a root)
2126
2132
 
2127
2133
  Args:
2128
- trace: The Trace object to store (one per trace_id).
2134
+ workflow_id_col: SQL column/expression for workflow_id
2135
+ team_id_col: SQL column/expression for team_id
2136
+ agent_id_col: SQL column/expression for agent_id
2137
+ name_col: SQL column/expression for name
2138
+
2139
+ Returns:
2140
+ SQLAlchemy CASE expression returning the component level as an integer.
2129
2141
  """
2130
- try:
2131
- table = await self._get_table(table_type="traces", create_table_if_not_found=True)
2142
+ is_root_name = or_(name_col.contains(".run"), name_col.contains(".arun"))
2143
+
2144
+ return case(
2145
+ # Workflow root (level 3)
2146
+ (and_(workflow_id_col.isnot(None), is_root_name), 3),
2147
+ # Team root (level 2)
2148
+ (and_(team_id_col.isnot(None), is_root_name), 2),
2149
+ # Agent root (level 1)
2150
+ (and_(agent_id_col.isnot(None), is_root_name), 1),
2151
+ # Child span or unknown (level 0)
2152
+ else_=0,
2153
+ )
2132
2154
 
2133
- async with self.async_session_factory() as sess, sess.begin():
2134
- # Check if trace exists
2135
- result = await sess.execute(select(table).where(table.c.trace_id == trace.trace_id))
2136
- existing = result.fetchone()
2137
-
2138
- if existing:
2139
- # workflow (level 3) > team (level 2) > agent (level 1) > child/unknown (level 0)
2140
-
2141
- def get_component_level(workflow_id, team_id, agent_id, name):
2142
- # Check if name indicates a root span
2143
- is_root_name = ".run" in name or ".arun" in name
2144
-
2145
- if not is_root_name:
2146
- return 0 # Child span (not a root)
2147
- elif workflow_id:
2148
- return 3 # Workflow root
2149
- elif team_id:
2150
- return 2 # Team root
2151
- elif agent_id:
2152
- return 1 # Agent root
2153
- else:
2154
- return 0 # Unknown
2155
-
2156
- existing_level = get_component_level(
2157
- existing.workflow_id, existing.team_id, existing.agent_id, existing.name
2158
- )
2159
- new_level = get_component_level(trace.workflow_id, trace.team_id, trace.agent_id, trace.name)
2155
+ async def upsert_trace(self, trace: "Trace") -> None:
2156
+ """Create or update a single trace record in the database.
2160
2157
 
2161
- # Only update name if new trace is from a higher or equal level
2162
- should_update_name = new_level > existing_level
2158
+ Uses INSERT ... ON CONFLICT DO UPDATE (upsert) to handle concurrent inserts
2159
+ atomically and avoid race conditions.
2163
2160
 
2164
- # Parse existing start_time to calculate correct duration
2165
- existing_start_time_str = existing.start_time
2166
- if isinstance(existing_start_time_str, str):
2167
- existing_start_time = datetime.fromisoformat(existing_start_time_str.replace("Z", "+00:00"))
2168
- else:
2169
- existing_start_time = trace.start_time
2161
+ Args:
2162
+ trace: The Trace object to store (one per trace_id).
2163
+ """
2164
+ try:
2165
+ table = await self._get_table(table_type="traces", create_table_if_not_found=True)
2170
2166
 
2171
- recalculated_duration_ms = int((trace.end_time - existing_start_time).total_seconds() * 1000)
2167
+ trace_dict = trace.to_dict()
2168
+ trace_dict.pop("total_spans", None)
2169
+ trace_dict.pop("error_count", None)
2172
2170
 
2173
- update_values = {
2174
- "end_time": trace.end_time.isoformat(),
2175
- "duration_ms": recalculated_duration_ms,
2176
- "status": trace.status,
2177
- "name": trace.name if should_update_name else existing.name,
2178
- }
2171
+ async with self.async_session_factory() as sess, sess.begin():
2172
+ # Use upsert to handle concurrent inserts atomically
2173
+ # On conflict, update fields while preserving existing non-null context values
2174
+ # and keeping the earliest start_time
2175
+ insert_stmt = postgresql.insert(table).values(trace_dict)
2176
+
2177
+ # Build component level expressions for comparing trace priority
2178
+ new_level = self._get_trace_component_level_expr(
2179
+ insert_stmt.excluded.workflow_id,
2180
+ insert_stmt.excluded.team_id,
2181
+ insert_stmt.excluded.agent_id,
2182
+ insert_stmt.excluded.name,
2183
+ )
2184
+ existing_level = self._get_trace_component_level_expr(
2185
+ table.c.workflow_id,
2186
+ table.c.team_id,
2187
+ table.c.agent_id,
2188
+ table.c.name,
2189
+ )
2179
2190
 
2180
- # Update context fields ONLY if new value is not None (preserve non-null values)
2181
- if trace.run_id is not None:
2182
- update_values["run_id"] = trace.run_id
2183
- if trace.session_id is not None:
2184
- update_values["session_id"] = trace.session_id
2185
- if trace.user_id is not None:
2186
- update_values["user_id"] = trace.user_id
2187
- if trace.agent_id is not None:
2188
- update_values["agent_id"] = trace.agent_id
2189
- if trace.team_id is not None:
2190
- update_values["team_id"] = trace.team_id
2191
- if trace.workflow_id is not None:
2192
- update_values["workflow_id"] = trace.workflow_id
2193
-
2194
- stmt = update(table).where(table.c.trace_id == trace.trace_id).values(**update_values)
2195
- await sess.execute(stmt)
2196
- else:
2197
- trace_dict = trace.to_dict()
2198
- trace_dict.pop("total_spans", None)
2199
- trace_dict.pop("error_count", None)
2200
- stmt = postgresql.insert(table).values(trace_dict)
2201
- await sess.execute(stmt)
2191
+ # Build the ON CONFLICT DO UPDATE clause
2192
+ # Use LEAST for start_time, GREATEST for end_time to capture full trace duration
2193
+ # Use COALESCE to preserve existing non-null context values
2194
+ upsert_stmt = insert_stmt.on_conflict_do_update(
2195
+ index_elements=["trace_id"],
2196
+ set_={
2197
+ "end_time": func.greatest(table.c.end_time, insert_stmt.excluded.end_time),
2198
+ "start_time": func.least(table.c.start_time, insert_stmt.excluded.start_time),
2199
+ "duration_ms": func.extract(
2200
+ "epoch",
2201
+ func.cast(
2202
+ func.greatest(table.c.end_time, insert_stmt.excluded.end_time),
2203
+ TIMESTAMP(timezone=True),
2204
+ )
2205
+ - func.cast(
2206
+ func.least(table.c.start_time, insert_stmt.excluded.start_time),
2207
+ TIMESTAMP(timezone=True),
2208
+ ),
2209
+ )
2210
+ * 1000,
2211
+ "status": insert_stmt.excluded.status,
2212
+ # Update name only if new trace is from a higher-level component
2213
+ # Priority: workflow (3) > team (2) > agent (1) > child spans (0)
2214
+ "name": case(
2215
+ (new_level > existing_level, insert_stmt.excluded.name),
2216
+ else_=table.c.name,
2217
+ ),
2218
+ # Preserve existing non-null context values using COALESCE
2219
+ "run_id": func.coalesce(insert_stmt.excluded.run_id, table.c.run_id),
2220
+ "session_id": func.coalesce(insert_stmt.excluded.session_id, table.c.session_id),
2221
+ "user_id": func.coalesce(insert_stmt.excluded.user_id, table.c.user_id),
2222
+ "agent_id": func.coalesce(insert_stmt.excluded.agent_id, table.c.agent_id),
2223
+ "team_id": func.coalesce(insert_stmt.excluded.team_id, table.c.team_id),
2224
+ "workflow_id": func.coalesce(insert_stmt.excluded.workflow_id, table.c.workflow_id),
2225
+ },
2226
+ )
2227
+ await sess.execute(upsert_stmt)
2202
2228
 
2203
2229
  except Exception as e:
2204
2230
  log_error(f"Error creating trace: {e}")