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.
- agno/db/base.py +5 -5
- agno/db/dynamo/dynamo.py +2 -2
- agno/db/firestore/firestore.py +2 -2
- agno/db/gcs_json/gcs_json_db.py +2 -2
- agno/db/in_memory/in_memory_db.py +2 -2
- agno/db/json/json_db.py +2 -2
- agno/db/mongo/async_mongo.py +171 -69
- agno/db/mongo/mongo.py +171 -77
- agno/db/mysql/async_mysql.py +93 -69
- agno/db/mysql/mysql.py +93 -68
- agno/db/postgres/async_postgres.py +104 -78
- agno/db/postgres/postgres.py +97 -69
- agno/db/redis/redis.py +2 -2
- agno/db/singlestore/singlestore.py +91 -66
- agno/db/sqlite/async_sqlite.py +101 -78
- agno/db/sqlite/sqlite.py +97 -69
- agno/db/surrealdb/surrealdb.py +2 -2
- agno/knowledge/chunking/fixed.py +4 -1
- agno/knowledge/knowledge.py +22 -4
- agno/knowledge/utils.py +52 -7
- agno/models/openai/chat.py +21 -0
- agno/os/routers/knowledge/knowledge.py +21 -9
- agno/tracing/exporter.py +2 -2
- agno/vectordb/base.py +15 -2
- agno/vectordb/pgvector/pgvector.py +7 -7
- {agno-2.3.10.dist-info → agno-2.3.11.dist-info}/METADATA +1 -1
- {agno-2.3.10.dist-info → agno-2.3.11.dist-info}/RECORD +30 -30
- {agno-2.3.10.dist-info → agno-2.3.11.dist-info}/WHEEL +0 -0
- {agno-2.3.10.dist-info → agno-2.3.11.dist-info}/licenses/LICENSE +0 -0
- {agno-2.3.10.dist-info → agno-2.3.11.dist-info}/top_level.txt +0 -0
agno/db/mysql/async_mysql.py
CHANGED
|
@@ -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
|
-
|
|
2438
|
-
"""
|
|
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
|
-
|
|
2449
|
-
|
|
2450
|
-
|
|
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
|
-
|
|
2489
|
-
|
|
2490
|
-
|
|
2491
|
-
|
|
2492
|
-
|
|
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
|
-
|
|
2496
|
-
|
|
2497
|
-
|
|
2498
|
-
|
|
2499
|
-
|
|
2500
|
-
|
|
2501
|
-
|
|
2502
|
-
|
|
2503
|
-
|
|
2504
|
-
|
|
2505
|
-
|
|
2506
|
-
|
|
2507
|
-
|
|
2508
|
-
|
|
2509
|
-
|
|
2510
|
-
|
|
2511
|
-
|
|
2512
|
-
|
|
2513
|
-
|
|
2514
|
-
|
|
2515
|
-
|
|
2516
|
-
|
|
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
|
|
2456
|
-
"""
|
|
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
|
-
|
|
2467
|
-
|
|
2468
|
-
|
|
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
|
-
|
|
2497
|
-
|
|
2498
|
-
|
|
2499
|
-
|
|
2500
|
-
|
|
2501
|
-
|
|
2502
|
-
|
|
2503
|
-
|
|
2504
|
-
|
|
2505
|
-
|
|
2506
|
-
|
|
2507
|
-
|
|
2508
|
-
|
|
2509
|
-
|
|
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
|
-
|
|
2513
|
-
|
|
2514
|
-
|
|
2515
|
-
|
|
2516
|
-
|
|
2517
|
-
|
|
2518
|
-
|
|
2519
|
-
|
|
2520
|
-
|
|
2521
|
-
|
|
2522
|
-
|
|
2523
|
-
|
|
2524
|
-
|
|
2525
|
-
|
|
2526
|
-
|
|
2527
|
-
|
|
2528
|
-
|
|
2529
|
-
|
|
2530
|
-
|
|
2531
|
-
|
|
2532
|
-
|
|
2533
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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
|
-
|
|
1154
|
-
|
|
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
|
-
|
|
1227
|
-
|
|
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
|
-
|
|
2125
|
-
"""
|
|
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
|
-
|
|
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
|
-
|
|
2131
|
-
|
|
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
|
-
|
|
2134
|
-
|
|
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
|
-
|
|
2162
|
-
|
|
2158
|
+
Uses INSERT ... ON CONFLICT DO UPDATE (upsert) to handle concurrent inserts
|
|
2159
|
+
atomically and avoid race conditions.
|
|
2163
2160
|
|
|
2164
|
-
|
|
2165
|
-
|
|
2166
|
-
|
|
2167
|
-
|
|
2168
|
-
|
|
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
|
-
|
|
2167
|
+
trace_dict = trace.to_dict()
|
|
2168
|
+
trace_dict.pop("total_spans", None)
|
|
2169
|
+
trace_dict.pop("error_count", None)
|
|
2172
2170
|
|
|
2173
|
-
|
|
2174
|
-
|
|
2175
|
-
|
|
2176
|
-
|
|
2177
|
-
|
|
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
|
-
|
|
2181
|
-
|
|
2182
|
-
|
|
2183
|
-
|
|
2184
|
-
|
|
2185
|
-
|
|
2186
|
-
|
|
2187
|
-
|
|
2188
|
-
|
|
2189
|
-
|
|
2190
|
-
|
|
2191
|
-
|
|
2192
|
-
|
|
2193
|
-
|
|
2194
|
-
|
|
2195
|
-
|
|
2196
|
-
|
|
2197
|
-
|
|
2198
|
-
|
|
2199
|
-
|
|
2200
|
-
|
|
2201
|
-
|
|
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}")
|