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/postgres/postgres.py
CHANGED
|
@@ -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
|
|
2404
|
-
"""
|
|
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
|
-
|
|
2415
|
-
|
|
2416
|
-
|
|
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
|
-
|
|
2437
|
-
|
|
2438
|
-
|
|
2439
|
-
|
|
2440
|
-
|
|
2441
|
-
|
|
2442
|
-
|
|
2443
|
-
|
|
2444
|
-
|
|
2445
|
-
|
|
2446
|
-
|
|
2447
|
-
|
|
2448
|
-
|
|
2449
|
-
|
|
2450
|
-
|
|
2451
|
-
|
|
2452
|
-
|
|
2453
|
-
|
|
2454
|
-
|
|
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
|
-
|
|
2461
|
-
|
|
2462
|
-
|
|
2463
|
-
|
|
2464
|
-
|
|
2465
|
-
|
|
2466
|
-
|
|
2467
|
-
|
|
2468
|
-
|
|
2469
|
-
|
|
2470
|
-
|
|
2471
|
-
|
|
2472
|
-
|
|
2473
|
-
|
|
2474
|
-
|
|
2475
|
-
|
|
2476
|
-
|
|
2477
|
-
|
|
2478
|
-
|
|
2479
|
-
|
|
2480
|
-
|
|
2481
|
-
|
|
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
|
|
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
|
|
2399
|
-
"""
|
|
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
|
-
|
|
2411
|
-
|
|
2412
|
-
|
|
2413
|
-
|
|
2414
|
-
|
|
2415
|
-
|
|
2416
|
-
|
|
2417
|
-
|
|
2418
|
-
|
|
2419
|
-
|
|
2420
|
-
|
|
2421
|
-
|
|
2422
|
-
|
|
2423
|
-
|
|
2424
|
-
|
|
2425
|
-
|
|
2426
|
-
|
|
2427
|
-
|
|
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
|
-
|
|
2431
|
-
|
|
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
|
-
|
|
2434
|
-
|
|
2435
|
-
#
|
|
2436
|
-
|
|
2437
|
-
|
|
2438
|
-
|
|
2439
|
-
|
|
2440
|
-
|
|
2441
|
-
|
|
2442
|
-
|
|
2443
|
-
|
|
2444
|
-
|
|
2445
|
-
|
|
2446
|
-
|
|
2447
|
-
|
|
2448
|
-
|
|
2449
|
-
|
|
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}")
|
agno/db/sqlite/async_sqlite.py
CHANGED
|
@@ -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
|
|
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
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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
|
-
|
|
2444
|
-
"""
|
|
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
|
-
|
|
2455
|
-
|
|
2456
|
-
|
|
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
|
-
|
|
2483
|
-
|
|
2484
|
-
|
|
2485
|
-
|
|
2486
|
-
|
|
2487
|
-
|
|
2488
|
-
|
|
2489
|
-
|
|
2490
|
-
|
|
2491
|
-
|
|
2492
|
-
|
|
2493
|
-
|
|
2494
|
-
|
|
2495
|
-
|
|
2496
|
-
|
|
2497
|
-
|
|
2498
|
-
|
|
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
|
-
|
|
2502
|
-
|
|
2503
|
-
|
|
2504
|
-
|
|
2505
|
-
|
|
2506
|
-
|
|
2507
|
-
|
|
2508
|
-
|
|
2509
|
-
|
|
2510
|
-
|
|
2511
|
-
|
|
2512
|
-
|
|
2513
|
-
|
|
2514
|
-
|
|
2515
|
-
|
|
2516
|
-
|
|
2517
|
-
|
|
2518
|
-
|
|
2519
|
-
|
|
2520
|
-
|
|
2521
|
-
|
|
2522
|
-
|
|
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}")
|