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.
@@ -30,8 +30,9 @@ from agno.utils.log import log_debug, log_error, log_info, log_warning
30
30
  from agno.utils.string import generate_id
31
31
 
32
32
  try:
33
- from sqlalchemy import ForeignKey, Index, String, UniqueConstraint, func, select, update
33
+ from sqlalchemy import ForeignKey, Index, String, UniqueConstraint, and_, case, func, or_, select, update
34
34
  from sqlalchemy.dialects import postgresql
35
+ from sqlalchemy.dialects.postgresql import TIMESTAMP
35
36
  from sqlalchemy.engine import Engine, create_engine
36
37
  from sqlalchemy.exc import ProgrammingError
37
38
  from sqlalchemy.orm import scoped_session, sessionmaker
@@ -2400,8 +2401,42 @@ class PostgresDb(BaseDb):
2400
2401
  # Fallback if spans table doesn't exist
2401
2402
  return select(table, literal(0).label("total_spans"), literal(0).label("error_count"))
2402
2403
 
2403
- def create_trace(self, trace: "Trace") -> None:
2404
- """Create a single trace record in the database.
2404
+ def _get_trace_component_level_expr(self, workflow_id_col, team_id_col, agent_id_col, name_col):
2405
+ """Build a SQL CASE expression that returns the component level for a trace.
2406
+
2407
+ Component levels (higher = more important):
2408
+ - 3: Workflow root (.run or .arun with workflow_id)
2409
+ - 2: Team root (.run or .arun with team_id)
2410
+ - 1: Agent root (.run or .arun with agent_id)
2411
+ - 0: Child span (not a root)
2412
+
2413
+ Args:
2414
+ workflow_id_col: SQL column/expression for workflow_id
2415
+ team_id_col: SQL column/expression for team_id
2416
+ agent_id_col: SQL column/expression for agent_id
2417
+ name_col: SQL column/expression for name
2418
+
2419
+ Returns:
2420
+ SQLAlchemy CASE expression returning the component level as an integer.
2421
+ """
2422
+ is_root_name = or_(name_col.contains(".run"), name_col.contains(".arun"))
2423
+
2424
+ return case(
2425
+ # Workflow root (level 3)
2426
+ (and_(workflow_id_col.isnot(None), is_root_name), 3),
2427
+ # Team root (level 2)
2428
+ (and_(team_id_col.isnot(None), is_root_name), 2),
2429
+ # Agent root (level 1)
2430
+ (and_(agent_id_col.isnot(None), is_root_name), 1),
2431
+ # Child span or unknown (level 0)
2432
+ else_=0,
2433
+ )
2434
+
2435
+ def upsert_trace(self, trace: "Trace") -> None:
2436
+ """Create or update a single trace record in the database.
2437
+
2438
+ Uses INSERT ... ON CONFLICT DO UPDATE (upsert) to handle concurrent inserts
2439
+ atomically and avoid race conditions.
2405
2440
 
2406
2441
  Args:
2407
2442
  trace: The Trace object to store (one per trace_id).
@@ -2411,74 +2446,67 @@ class PostgresDb(BaseDb):
2411
2446
  if table is None:
2412
2447
  return
2413
2448
 
2414
- with self.Session() as sess, sess.begin():
2415
- # Check if trace exists
2416
- existing = sess.execute(select(table).where(table.c.trace_id == trace.trace_id)).fetchone()
2417
-
2418
- if existing:
2419
- # workflow (level 3) > team (level 2) > agent (level 1) > child/unknown (level 0)
2420
-
2421
- def get_component_level(workflow_id, team_id, agent_id, name):
2422
- # Check if name indicates a root span
2423
- is_root_name = ".run" in name or ".arun" in name
2424
-
2425
- if not is_root_name:
2426
- return 0 # Child span (not a root)
2427
- elif workflow_id:
2428
- return 3 # Workflow root
2429
- elif team_id:
2430
- return 2 # Team root
2431
- elif agent_id:
2432
- return 1 # Agent root
2433
- else:
2434
- return 0 # Unknown
2449
+ trace_dict = trace.to_dict()
2450
+ trace_dict.pop("total_spans", None)
2451
+ trace_dict.pop("error_count", None)
2435
2452
 
2436
- existing_level = get_component_level(
2437
- existing.workflow_id, existing.team_id, existing.agent_id, existing.name
2438
- )
2439
- new_level = get_component_level(trace.workflow_id, trace.team_id, trace.agent_id, trace.name)
2440
-
2441
- # Only update name if new trace is from a higher or equal level
2442
- should_update_name = new_level > existing_level
2443
-
2444
- # Parse existing start_time to calculate correct duration
2445
- existing_start_time_str = existing.start_time
2446
- if isinstance(existing_start_time_str, str):
2447
- existing_start_time = datetime.fromisoformat(existing_start_time_str.replace("Z", "+00:00"))
2448
- else:
2449
- existing_start_time = trace.start_time
2450
-
2451
- recalculated_duration_ms = int((trace.end_time - existing_start_time).total_seconds() * 1000)
2452
-
2453
- update_values = {
2454
- "end_time": trace.end_time.isoformat(),
2455
- "duration_ms": recalculated_duration_ms,
2456
- "status": trace.status,
2457
- "name": trace.name if should_update_name else existing.name,
2458
- }
2453
+ with self.Session() as sess, sess.begin():
2454
+ # Use upsert to handle concurrent inserts atomically
2455
+ # On conflict, update fields while preserving existing non-null context values
2456
+ # and keeping the earliest start_time
2457
+ insert_stmt = postgresql.insert(table).values(trace_dict)
2458
+
2459
+ # Build component level expressions for comparing trace priority
2460
+ new_level = self._get_trace_component_level_expr(
2461
+ insert_stmt.excluded.workflow_id,
2462
+ insert_stmt.excluded.team_id,
2463
+ insert_stmt.excluded.agent_id,
2464
+ insert_stmt.excluded.name,
2465
+ )
2466
+ existing_level = self._get_trace_component_level_expr(
2467
+ table.c.workflow_id,
2468
+ table.c.team_id,
2469
+ table.c.agent_id,
2470
+ table.c.name,
2471
+ )
2459
2472
 
2460
- # Update context fields ONLY if new value is not None (preserve non-null values)
2461
- if trace.run_id is not None:
2462
- update_values["run_id"] = trace.run_id
2463
- if trace.session_id is not None:
2464
- update_values["session_id"] = trace.session_id
2465
- if trace.user_id is not None:
2466
- update_values["user_id"] = trace.user_id
2467
- if trace.agent_id is not None:
2468
- update_values["agent_id"] = trace.agent_id
2469
- if trace.team_id is not None:
2470
- update_values["team_id"] = trace.team_id
2471
- if trace.workflow_id is not None:
2472
- update_values["workflow_id"] = trace.workflow_id
2473
-
2474
- stmt = update(table).where(table.c.trace_id == trace.trace_id).values(**update_values)
2475
- sess.execute(stmt)
2476
- else:
2477
- trace_dict = trace.to_dict()
2478
- trace_dict.pop("total_spans", None)
2479
- trace_dict.pop("error_count", None)
2480
- stmt = postgresql.insert(table).values(trace_dict)
2481
- sess.execute(stmt)
2473
+ # Build the ON CONFLICT DO UPDATE clause
2474
+ # Use LEAST for start_time, GREATEST for end_time to capture full trace duration
2475
+ # Use COALESCE to preserve existing non-null context values
2476
+ upsert_stmt = insert_stmt.on_conflict_do_update(
2477
+ index_elements=["trace_id"],
2478
+ set_={
2479
+ "end_time": func.greatest(table.c.end_time, insert_stmt.excluded.end_time),
2480
+ "start_time": func.least(table.c.start_time, insert_stmt.excluded.start_time),
2481
+ "duration_ms": func.extract(
2482
+ "epoch",
2483
+ func.cast(
2484
+ func.greatest(table.c.end_time, insert_stmt.excluded.end_time),
2485
+ TIMESTAMP(timezone=True),
2486
+ )
2487
+ - func.cast(
2488
+ func.least(table.c.start_time, insert_stmt.excluded.start_time),
2489
+ TIMESTAMP(timezone=True),
2490
+ ),
2491
+ )
2492
+ * 1000,
2493
+ "status": insert_stmt.excluded.status,
2494
+ # Update name only if new trace is from a higher-level component
2495
+ # Priority: workflow (3) > team (2) > agent (1) > child spans (0)
2496
+ "name": case(
2497
+ (new_level > existing_level, insert_stmt.excluded.name),
2498
+ else_=table.c.name,
2499
+ ),
2500
+ # Preserve existing non-null context values using COALESCE
2501
+ "run_id": func.coalesce(insert_stmt.excluded.run_id, table.c.run_id),
2502
+ "session_id": func.coalesce(insert_stmt.excluded.session_id, table.c.session_id),
2503
+ "user_id": func.coalesce(insert_stmt.excluded.user_id, table.c.user_id),
2504
+ "agent_id": func.coalesce(insert_stmt.excluded.agent_id, table.c.agent_id),
2505
+ "team_id": func.coalesce(insert_stmt.excluded.team_id, table.c.team_id),
2506
+ "workflow_id": func.coalesce(insert_stmt.excluded.workflow_id, table.c.workflow_id),
2507
+ },
2508
+ )
2509
+ sess.execute(upsert_stmt)
2482
2510
 
2483
2511
  except Exception as e:
2484
2512
  log_error(f"Error creating trace: {e}")
agno/db/redis/redis.py CHANGED
@@ -1693,8 +1693,8 @@ class RedisDb(BaseDb):
1693
1693
  raise e
1694
1694
 
1695
1695
  # --- Traces ---
1696
- def create_trace(self, trace: "Trace") -> None:
1697
- """Create a single trace record in the database.
1696
+ def upsert_trace(self, trace: "Trace") -> None:
1697
+ """Create or update a single trace record in the database.
1698
1698
 
1699
1699
  Args:
1700
1700
  trace: The Trace object to store (one per trace_id).
@@ -2395,84 +2395,109 @@ class SingleStoreDb(BaseDb):
2395
2395
  # Fallback if spans table doesn't exist
2396
2396
  return select(table, literal(0).label("total_spans"), literal(0).label("error_count"))
2397
2397
 
2398
- def create_trace(self, trace: "Trace") -> None:
2399
- """Create a single trace record in the database.
2398
+ def _get_trace_component_level_expr(self, workflow_id_col, team_id_col, agent_id_col, name_col):
2399
+ """Build a SQL CASE expression that returns the component level for a trace.
2400
+
2401
+ Component levels (higher = more important):
2402
+ - 3: Workflow root (.run or .arun with workflow_id)
2403
+ - 2: Team root (.run or .arun with team_id)
2404
+ - 1: Agent root (.run or .arun with agent_id)
2405
+ - 0: Child span (not a root)
2406
+
2407
+ Args:
2408
+ workflow_id_col: SQL column/expression for workflow_id
2409
+ team_id_col: SQL column/expression for team_id
2410
+ agent_id_col: SQL column/expression for agent_id
2411
+ name_col: SQL column/expression for name
2412
+
2413
+ Returns:
2414
+ SQLAlchemy CASE expression returning the component level as an integer.
2415
+ """
2416
+ from sqlalchemy import case, or_
2417
+
2418
+ is_root_name = or_(name_col.like("%.run%"), name_col.like("%.arun%"))
2419
+
2420
+ return case(
2421
+ # Workflow root (level 3)
2422
+ (and_(workflow_id_col.isnot(None), is_root_name), 3),
2423
+ # Team root (level 2)
2424
+ (and_(team_id_col.isnot(None), is_root_name), 2),
2425
+ # Agent root (level 1)
2426
+ (and_(agent_id_col.isnot(None), is_root_name), 1),
2427
+ # Child span or unknown (level 0)
2428
+ else_=0,
2429
+ )
2430
+
2431
+ def upsert_trace(self, trace: "Trace") -> None:
2432
+ """Create or update a single trace record in the database.
2433
+
2434
+ Uses INSERT ... ON DUPLICATE KEY UPDATE (upsert) to handle concurrent inserts
2435
+ atomically and avoid race conditions.
2400
2436
 
2401
2437
  Args:
2402
2438
  trace: The Trace object to store (one per trace_id).
2403
2439
  """
2440
+ from sqlalchemy import case
2441
+
2404
2442
  try:
2405
2443
  table = self._get_table(table_type="traces", create_table_if_not_found=True)
2406
2444
  if table is None:
2407
2445
  return
2408
2446
 
2447
+ trace_dict = trace.to_dict()
2448
+ trace_dict.pop("total_spans", None)
2449
+ trace_dict.pop("error_count", None)
2450
+
2409
2451
  with self.Session() as sess, sess.begin():
2410
- existing = sess.execute(select(table).where(table.c.trace_id == trace.trace_id)).fetchone()
2411
-
2412
- if existing:
2413
- # workflow (level 3) > team (level 2) > agent (level 1) > child/unknown (level 0)
2414
-
2415
- def get_component_level(workflow_id, team_id, agent_id, name):
2416
- # Check if name indicates a root span
2417
- is_root_name = ".run" in name or ".arun" in name
2418
-
2419
- if not is_root_name:
2420
- return 0 # Child span (not a root)
2421
- elif workflow_id:
2422
- return 3 # Workflow root
2423
- elif team_id:
2424
- return 2 # Team root
2425
- elif agent_id:
2426
- return 1 # Agent root
2427
- else:
2428
- return 0 # Unknown
2452
+ # Use upsert to handle concurrent inserts atomically
2453
+ # On conflict, update fields while preserving existing non-null context values
2454
+ # and keeping the earliest start_time
2455
+ insert_stmt = mysql.insert(table).values(trace_dict)
2456
+
2457
+ # Build component level expressions for comparing trace priority
2458
+ new_level = self._get_trace_component_level_expr(
2459
+ insert_stmt.inserted.workflow_id,
2460
+ insert_stmt.inserted.team_id,
2461
+ insert_stmt.inserted.agent_id,
2462
+ insert_stmt.inserted.name,
2463
+ )
2464
+ existing_level = self._get_trace_component_level_expr(
2465
+ table.c.workflow_id,
2466
+ table.c.team_id,
2467
+ table.c.agent_id,
2468
+ table.c.name,
2469
+ )
2429
2470
 
2430
- existing_level = get_component_level(
2431
- existing.workflow_id, existing.team_id, existing.agent_id, existing.name
2471
+ # Build the ON DUPLICATE KEY UPDATE clause
2472
+ # Use LEAST for start_time, GREATEST for end_time to capture full trace duration
2473
+ # Duration is calculated using TIMESTAMPDIFF in microseconds then converted to ms
2474
+ upsert_stmt = insert_stmt.on_duplicate_key_update(
2475
+ end_time=func.greatest(table.c.end_time, insert_stmt.inserted.end_time),
2476
+ start_time=func.least(table.c.start_time, insert_stmt.inserted.start_time),
2477
+ # Calculate duration in milliseconds using TIMESTAMPDIFF
2478
+ # TIMESTAMPDIFF(MICROSECOND, start, end) / 1000 gives milliseconds
2479
+ duration_ms=func.timestampdiff(
2480
+ text("MICROSECOND"),
2481
+ func.least(table.c.start_time, insert_stmt.inserted.start_time),
2482
+ func.greatest(table.c.end_time, insert_stmt.inserted.end_time),
2432
2483
  )
2433
- new_level = get_component_level(trace.workflow_id, trace.team_id, trace.agent_id, trace.name)
2434
-
2435
- # Only update name if new trace is from a higher or equal level
2436
- should_update_name = new_level > existing_level
2437
-
2438
- # Parse existing start_time to calculate correct duration
2439
- existing_start_time_str = existing.start_time
2440
- if isinstance(existing_start_time_str, str):
2441
- existing_start_time = datetime.fromisoformat(existing_start_time_str.replace("Z", "+00:00"))
2442
- else:
2443
- existing_start_time = trace.start_time
2444
-
2445
- recalculated_duration_ms = int((trace.end_time - existing_start_time).total_seconds() * 1000)
2446
-
2447
- update_values = {
2448
- "end_time": trace.end_time.isoformat(),
2449
- "duration_ms": recalculated_duration_ms,
2450
- "status": trace.status,
2451
- "name": trace.name if should_update_name else existing.name,
2452
- }
2453
-
2454
- # Update context fields ONLY if new value is not None (preserve non-null values)
2455
- if trace.run_id is not None:
2456
- update_values["run_id"] = trace.run_id
2457
- if trace.session_id is not None:
2458
- update_values["session_id"] = trace.session_id
2459
- if trace.user_id is not None:
2460
- update_values["user_id"] = trace.user_id
2461
- if trace.agent_id is not None:
2462
- update_values["agent_id"] = trace.agent_id
2463
- if trace.team_id is not None:
2464
- update_values["team_id"] = trace.team_id
2465
- if trace.workflow_id is not None:
2466
- update_values["workflow_id"] = trace.workflow_id
2467
-
2468
- update_stmt = update(table).where(table.c.trace_id == trace.trace_id).values(**update_values)
2469
- sess.execute(update_stmt)
2470
- else:
2471
- trace_dict = trace.to_dict()
2472
- trace_dict.pop("total_spans", None)
2473
- trace_dict.pop("error_count", None)
2474
- insert_stmt = mysql.insert(table).values(trace_dict)
2475
- sess.execute(insert_stmt)
2484
+ / 1000,
2485
+ status=insert_stmt.inserted.status,
2486
+ # Update name only if new trace is from a higher-level component
2487
+ # Priority: workflow (3) > team (2) > agent (1) > child spans (0)
2488
+ name=case(
2489
+ (new_level > existing_level, insert_stmt.inserted.name),
2490
+ else_=table.c.name,
2491
+ ),
2492
+ # Preserve existing non-null context values using COALESCE
2493
+ run_id=func.coalesce(insert_stmt.inserted.run_id, table.c.run_id),
2494
+ session_id=func.coalesce(insert_stmt.inserted.session_id, table.c.session_id),
2495
+ user_id=func.coalesce(insert_stmt.inserted.user_id, table.c.user_id),
2496
+ agent_id=func.coalesce(insert_stmt.inserted.agent_id, table.c.agent_id),
2497
+ team_id=func.coalesce(insert_stmt.inserted.team_id, table.c.team_id),
2498
+ workflow_id=func.coalesce(insert_stmt.inserted.workflow_id, table.c.workflow_id),
2499
+ )
2500
+ sess.execute(upsert_stmt)
2476
2501
 
2477
2502
  except Exception as e:
2478
2503
  log_error(f"Error creating trace: {e}")
@@ -31,7 +31,7 @@ from agno.utils.log import log_debug, log_error, log_info, log_warning
31
31
  from agno.utils.string import generate_id
32
32
 
33
33
  try:
34
- from sqlalchemy import Column, MetaData, String, Table, func, select, text, update
34
+ from sqlalchemy import Column, MetaData, String, Table, func, select, text
35
35
  from sqlalchemy.dialects import sqlite
36
36
  from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker, create_async_engine
37
37
  from sqlalchemy.schema import Index, UniqueConstraint
@@ -2196,10 +2196,7 @@ class AsyncSqliteDb(AsyncBaseDb):
2196
2196
  await sess.execute(table.delete())
2197
2197
 
2198
2198
  except Exception as e:
2199
- from agno.utils.log import log_warning
2200
-
2201
- log_warning(f"Exception deleting all cultural artifacts: {e}")
2202
- raise e
2199
+ log_error(f"Exception deleting all cultural artifacts: {e}")
2203
2200
 
2204
2201
  async def delete_cultural_knowledge(self, id: str) -> None:
2205
2202
  """Delete a cultural artifact from the database.
@@ -2227,7 +2224,6 @@ class AsyncSqliteDb(AsyncBaseDb):
2227
2224
 
2228
2225
  except Exception as e:
2229
2226
  log_error(f"Error deleting cultural artifact: {e}")
2230
- raise e
2231
2227
 
2232
2228
  async def get_cultural_knowledge(
2233
2229
  self, id: str, deserialize: Optional[bool] = True
@@ -2263,7 +2259,7 @@ class AsyncSqliteDb(AsyncBaseDb):
2263
2259
 
2264
2260
  except Exception as e:
2265
2261
  log_error(f"Exception reading from cultural artifacts table: {e}")
2266
- raise e
2262
+ return None
2267
2263
 
2268
2264
  async def get_all_cultural_knowledge(
2269
2265
  self,
@@ -2337,7 +2333,7 @@ class AsyncSqliteDb(AsyncBaseDb):
2337
2333
 
2338
2334
  except Exception as e:
2339
2335
  log_error(f"Error reading from cultural artifacts table: {e}")
2340
- raise e
2336
+ return [] if deserialize else ([], 0)
2341
2337
 
2342
2338
  async def upsert_cultural_knowledge(
2343
2339
  self, cultural_knowledge: CulturalKnowledge, deserialize: Optional[bool] = True
@@ -2357,7 +2353,7 @@ class AsyncSqliteDb(AsyncBaseDb):
2357
2353
  Exception: If an error occurs during upsert.
2358
2354
  """
2359
2355
  try:
2360
- table = await self._get_table(table_type="culture")
2356
+ table = await self._get_table(table_type="culture", create_table_if_not_found=True)
2361
2357
  if table is None:
2362
2358
  return None
2363
2359
 
@@ -2440,86 +2436,113 @@ class AsyncSqliteDb(AsyncBaseDb):
2440
2436
  # Fallback if spans table doesn't exist
2441
2437
  return select(table, literal(0).label("total_spans"), literal(0).label("error_count"))
2442
2438
 
2443
- async def create_trace(self, trace: "Trace") -> None:
2444
- """Create a single trace record in the database.
2439
+ def _get_trace_component_level_expr(self, workflow_id_col, team_id_col, agent_id_col, name_col):
2440
+ """Build a SQL CASE expression that returns the component level for a trace.
2441
+
2442
+ Component levels (higher = more important):
2443
+ - 3: Workflow root (.run or .arun with workflow_id)
2444
+ - 2: Team root (.run or .arun with team_id)
2445
+ - 1: Agent root (.run or .arun with agent_id)
2446
+ - 0: Child span (not a root)
2447
+
2448
+ Args:
2449
+ workflow_id_col: SQL column/expression for workflow_id
2450
+ team_id_col: SQL column/expression for team_id
2451
+ agent_id_col: SQL column/expression for agent_id
2452
+ name_col: SQL column/expression for name
2453
+
2454
+ Returns:
2455
+ SQLAlchemy CASE expression returning the component level as an integer.
2456
+ """
2457
+ from sqlalchemy import and_, case, or_
2458
+
2459
+ is_root_name = or_(name_col.contains(".run"), name_col.contains(".arun"))
2460
+
2461
+ return case(
2462
+ # Workflow root (level 3)
2463
+ (and_(workflow_id_col.isnot(None), is_root_name), 3),
2464
+ # Team root (level 2)
2465
+ (and_(team_id_col.isnot(None), is_root_name), 2),
2466
+ # Agent root (level 1)
2467
+ (and_(agent_id_col.isnot(None), is_root_name), 1),
2468
+ # Child span or unknown (level 0)
2469
+ else_=0,
2470
+ )
2471
+
2472
+ async def upsert_trace(self, trace: "Trace") -> None:
2473
+ """Create or update a single trace record in the database.
2474
+
2475
+ Uses INSERT ... ON CONFLICT DO UPDATE (upsert) to handle concurrent inserts
2476
+ atomically and avoid race conditions.
2445
2477
 
2446
2478
  Args:
2447
2479
  trace: The Trace object to store (one per trace_id).
2448
2480
  """
2481
+ from sqlalchemy import case
2482
+
2449
2483
  try:
2450
2484
  table = await self._get_table(table_type="traces", create_table_if_not_found=True)
2451
2485
  if table is None:
2452
2486
  return
2453
2487
 
2454
- async with self.async_session_factory() as sess, sess.begin():
2455
- # Check if trace exists
2456
- result = await sess.execute(select(table).where(table.c.trace_id == trace.trace_id))
2457
- existing = result.fetchone()
2458
-
2459
- if existing:
2460
- # workflow (level 3) > team (level 2) > agent (level 1) > child/unknown (level 0)
2461
-
2462
- def get_component_level(workflow_id, team_id, agent_id, name):
2463
- # Check if name indicates a root span
2464
- is_root_name = ".run" in name or ".arun" in name
2465
-
2466
- if not is_root_name:
2467
- return 0 # Child span (not a root)
2468
- elif workflow_id:
2469
- return 3 # Workflow root
2470
- elif team_id:
2471
- return 2 # Team root
2472
- elif agent_id:
2473
- return 1 # Agent root
2474
- else:
2475
- return 0 # Unknown
2476
-
2477
- existing_level = get_component_level(
2478
- existing.workflow_id, existing.team_id, existing.agent_id, existing.name
2479
- )
2480
- new_level = get_component_level(trace.workflow_id, trace.team_id, trace.agent_id, trace.name)
2488
+ trace_dict = trace.to_dict()
2489
+ trace_dict.pop("total_spans", None)
2490
+ trace_dict.pop("error_count", None)
2481
2491
 
2482
- # Only update name if new trace is from a higher or equal level
2483
- should_update_name = new_level > existing_level
2484
-
2485
- # Parse existing start_time to calculate correct duration
2486
- existing_start_time_str = existing.start_time
2487
- if isinstance(existing_start_time_str, str):
2488
- existing_start_time = datetime.fromisoformat(existing_start_time_str.replace("Z", "+00:00"))
2489
- else:
2490
- existing_start_time = trace.start_time
2491
-
2492
- recalculated_duration_ms = int((trace.end_time - existing_start_time).total_seconds() * 1000)
2493
-
2494
- update_values = {
2495
- "end_time": trace.end_time.isoformat(),
2496
- "duration_ms": recalculated_duration_ms,
2497
- "status": trace.status,
2498
- "name": trace.name if should_update_name else existing.name,
2499
- }
2492
+ async with self.async_session_factory() as sess, sess.begin():
2493
+ # Use upsert to handle concurrent inserts atomically
2494
+ # On conflict, update fields while preserving existing non-null context values
2495
+ # and keeping the earliest start_time
2496
+ insert_stmt = sqlite.insert(table).values(trace_dict)
2497
+
2498
+ # Build component level expressions for comparing trace priority
2499
+ new_level = self._get_trace_component_level_expr(
2500
+ insert_stmt.excluded.workflow_id,
2501
+ insert_stmt.excluded.team_id,
2502
+ insert_stmt.excluded.agent_id,
2503
+ insert_stmt.excluded.name,
2504
+ )
2505
+ existing_level = self._get_trace_component_level_expr(
2506
+ table.c.workflow_id,
2507
+ table.c.team_id,
2508
+ table.c.agent_id,
2509
+ table.c.name,
2510
+ )
2500
2511
 
2501
- # Update context fields ONLY if new value is not None (preserve non-null values)
2502
- if trace.run_id is not None:
2503
- update_values["run_id"] = trace.run_id
2504
- if trace.session_id is not None:
2505
- update_values["session_id"] = trace.session_id
2506
- if trace.user_id is not None:
2507
- update_values["user_id"] = trace.user_id
2508
- if trace.agent_id is not None:
2509
- update_values["agent_id"] = trace.agent_id
2510
- if trace.team_id is not None:
2511
- update_values["team_id"] = trace.team_id
2512
- if trace.workflow_id is not None:
2513
- update_values["workflow_id"] = trace.workflow_id
2514
-
2515
- stmt = update(table).where(table.c.trace_id == trace.trace_id).values(**update_values)
2516
- await sess.execute(stmt)
2517
- else:
2518
- trace_dict = trace.to_dict()
2519
- trace_dict.pop("total_spans", None)
2520
- trace_dict.pop("error_count", None)
2521
- stmt = sqlite.insert(table).values(trace_dict)
2522
- await sess.execute(stmt)
2512
+ # Build the ON CONFLICT DO UPDATE clause
2513
+ # Use MIN for start_time, MAX for end_time to capture full trace duration
2514
+ # SQLite stores timestamps as ISO strings, so string comparison works for ISO format
2515
+ # Duration is calculated as: (MAX(end_time) - MIN(start_time)) in milliseconds
2516
+ # SQLite doesn't have epoch extraction, so we calculate duration using julianday
2517
+ upsert_stmt = insert_stmt.on_conflict_do_update(
2518
+ index_elements=["trace_id"],
2519
+ set_={
2520
+ "end_time": func.max(table.c.end_time, insert_stmt.excluded.end_time),
2521
+ "start_time": func.min(table.c.start_time, insert_stmt.excluded.start_time),
2522
+ # Calculate duration in milliseconds using julianday (SQLite-specific)
2523
+ # julianday returns days, so multiply by 86400000 to get milliseconds
2524
+ "duration_ms": (
2525
+ func.julianday(func.max(table.c.end_time, insert_stmt.excluded.end_time))
2526
+ - func.julianday(func.min(table.c.start_time, insert_stmt.excluded.start_time))
2527
+ )
2528
+ * 86400000,
2529
+ "status": insert_stmt.excluded.status,
2530
+ # Update name only if new trace is from a higher-level component
2531
+ # Priority: workflow (3) > team (2) > agent (1) > child spans (0)
2532
+ "name": case(
2533
+ (new_level > existing_level, insert_stmt.excluded.name),
2534
+ else_=table.c.name,
2535
+ ),
2536
+ # Preserve existing non-null context values using COALESCE
2537
+ "run_id": func.coalesce(insert_stmt.excluded.run_id, table.c.run_id),
2538
+ "session_id": func.coalesce(insert_stmt.excluded.session_id, table.c.session_id),
2539
+ "user_id": func.coalesce(insert_stmt.excluded.user_id, table.c.user_id),
2540
+ "agent_id": func.coalesce(insert_stmt.excluded.agent_id, table.c.agent_id),
2541
+ "team_id": func.coalesce(insert_stmt.excluded.team_id, table.c.team_id),
2542
+ "workflow_id": func.coalesce(insert_stmt.excluded.workflow_id, table.c.workflow_id),
2543
+ },
2544
+ )
2545
+ await sess.execute(upsert_stmt)
2523
2546
 
2524
2547
  except Exception as e:
2525
2548
  log_error(f"Error creating trace: {e}")