agno 2.1.2__py3-none-any.whl → 2.3.13__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/agent/agent.py +5540 -2273
- agno/api/api.py +2 -0
- agno/api/os.py +1 -1
- agno/compression/__init__.py +3 -0
- agno/compression/manager.py +247 -0
- agno/culture/__init__.py +3 -0
- agno/culture/manager.py +956 -0
- agno/db/async_postgres/__init__.py +3 -0
- agno/db/base.py +689 -6
- agno/db/dynamo/dynamo.py +933 -37
- agno/db/dynamo/schemas.py +174 -10
- agno/db/dynamo/utils.py +63 -4
- agno/db/firestore/firestore.py +831 -9
- agno/db/firestore/schemas.py +51 -0
- agno/db/firestore/utils.py +102 -4
- agno/db/gcs_json/gcs_json_db.py +660 -12
- agno/db/gcs_json/utils.py +60 -26
- agno/db/in_memory/in_memory_db.py +287 -14
- agno/db/in_memory/utils.py +60 -2
- agno/db/json/json_db.py +590 -14
- agno/db/json/utils.py +60 -26
- agno/db/migrations/manager.py +199 -0
- agno/db/migrations/v1_to_v2.py +43 -13
- agno/db/migrations/versions/__init__.py +0 -0
- agno/db/migrations/versions/v2_3_0.py +938 -0
- agno/db/mongo/__init__.py +15 -1
- agno/db/mongo/async_mongo.py +2760 -0
- agno/db/mongo/mongo.py +879 -11
- agno/db/mongo/schemas.py +42 -0
- agno/db/mongo/utils.py +80 -8
- agno/db/mysql/__init__.py +2 -1
- agno/db/mysql/async_mysql.py +2912 -0
- agno/db/mysql/mysql.py +946 -68
- agno/db/mysql/schemas.py +72 -10
- agno/db/mysql/utils.py +198 -7
- agno/db/postgres/__init__.py +2 -1
- agno/db/postgres/async_postgres.py +2579 -0
- agno/db/postgres/postgres.py +942 -57
- agno/db/postgres/schemas.py +81 -18
- agno/db/postgres/utils.py +164 -2
- agno/db/redis/redis.py +671 -7
- agno/db/redis/schemas.py +50 -0
- agno/db/redis/utils.py +65 -7
- agno/db/schemas/__init__.py +2 -1
- agno/db/schemas/culture.py +120 -0
- agno/db/schemas/evals.py +1 -0
- agno/db/schemas/memory.py +17 -2
- agno/db/singlestore/schemas.py +63 -0
- agno/db/singlestore/singlestore.py +949 -83
- agno/db/singlestore/utils.py +60 -2
- agno/db/sqlite/__init__.py +2 -1
- agno/db/sqlite/async_sqlite.py +2911 -0
- agno/db/sqlite/schemas.py +62 -0
- agno/db/sqlite/sqlite.py +965 -46
- agno/db/sqlite/utils.py +169 -8
- agno/db/surrealdb/__init__.py +3 -0
- agno/db/surrealdb/metrics.py +292 -0
- agno/db/surrealdb/models.py +334 -0
- agno/db/surrealdb/queries.py +71 -0
- agno/db/surrealdb/surrealdb.py +1908 -0
- agno/db/surrealdb/utils.py +147 -0
- agno/db/utils.py +2 -0
- agno/eval/__init__.py +10 -0
- agno/eval/accuracy.py +75 -55
- agno/eval/agent_as_judge.py +861 -0
- agno/eval/base.py +29 -0
- agno/eval/performance.py +16 -7
- agno/eval/reliability.py +28 -16
- agno/eval/utils.py +35 -17
- agno/exceptions.py +27 -2
- agno/filters.py +354 -0
- agno/guardrails/prompt_injection.py +1 -0
- agno/hooks/__init__.py +3 -0
- agno/hooks/decorator.py +164 -0
- agno/integrations/discord/client.py +1 -1
- agno/knowledge/chunking/agentic.py +13 -10
- agno/knowledge/chunking/fixed.py +4 -1
- agno/knowledge/chunking/semantic.py +9 -4
- agno/knowledge/chunking/strategy.py +59 -15
- agno/knowledge/embedder/fastembed.py +1 -1
- agno/knowledge/embedder/nebius.py +1 -1
- agno/knowledge/embedder/ollama.py +8 -0
- agno/knowledge/embedder/openai.py +8 -8
- agno/knowledge/embedder/sentence_transformer.py +6 -2
- agno/knowledge/embedder/vllm.py +262 -0
- agno/knowledge/knowledge.py +1618 -318
- agno/knowledge/reader/base.py +6 -2
- agno/knowledge/reader/csv_reader.py +8 -10
- agno/knowledge/reader/docx_reader.py +5 -6
- agno/knowledge/reader/field_labeled_csv_reader.py +16 -20
- agno/knowledge/reader/json_reader.py +5 -4
- agno/knowledge/reader/markdown_reader.py +8 -8
- agno/knowledge/reader/pdf_reader.py +17 -19
- agno/knowledge/reader/pptx_reader.py +101 -0
- agno/knowledge/reader/reader_factory.py +32 -3
- agno/knowledge/reader/s3_reader.py +3 -3
- agno/knowledge/reader/tavily_reader.py +193 -0
- agno/knowledge/reader/text_reader.py +22 -10
- agno/knowledge/reader/web_search_reader.py +1 -48
- agno/knowledge/reader/website_reader.py +10 -10
- agno/knowledge/reader/wikipedia_reader.py +33 -1
- agno/knowledge/types.py +1 -0
- agno/knowledge/utils.py +72 -7
- agno/media.py +22 -6
- agno/memory/__init__.py +14 -1
- agno/memory/manager.py +544 -83
- agno/memory/strategies/__init__.py +15 -0
- agno/memory/strategies/base.py +66 -0
- agno/memory/strategies/summarize.py +196 -0
- agno/memory/strategies/types.py +37 -0
- agno/models/aimlapi/aimlapi.py +17 -0
- agno/models/anthropic/claude.py +515 -40
- agno/models/aws/bedrock.py +102 -21
- agno/models/aws/claude.py +131 -274
- agno/models/azure/ai_foundry.py +41 -19
- agno/models/azure/openai_chat.py +39 -8
- agno/models/base.py +1249 -525
- agno/models/cerebras/cerebras.py +91 -21
- agno/models/cerebras/cerebras_openai.py +21 -2
- agno/models/cohere/chat.py +40 -6
- agno/models/cometapi/cometapi.py +18 -1
- agno/models/dashscope/dashscope.py +2 -3
- agno/models/deepinfra/deepinfra.py +18 -1
- agno/models/deepseek/deepseek.py +69 -3
- agno/models/fireworks/fireworks.py +18 -1
- agno/models/google/gemini.py +877 -80
- agno/models/google/utils.py +22 -0
- agno/models/groq/groq.py +51 -18
- agno/models/huggingface/huggingface.py +17 -6
- agno/models/ibm/watsonx.py +16 -6
- agno/models/internlm/internlm.py +18 -1
- agno/models/langdb/langdb.py +13 -1
- agno/models/litellm/chat.py +44 -9
- agno/models/litellm/litellm_openai.py +18 -1
- agno/models/message.py +28 -5
- agno/models/meta/llama.py +47 -14
- agno/models/meta/llama_openai.py +22 -17
- agno/models/mistral/mistral.py +8 -4
- agno/models/nebius/nebius.py +6 -7
- agno/models/nvidia/nvidia.py +20 -3
- agno/models/ollama/chat.py +24 -8
- agno/models/openai/chat.py +104 -29
- agno/models/openai/responses.py +101 -81
- agno/models/openrouter/openrouter.py +60 -3
- agno/models/perplexity/perplexity.py +17 -1
- agno/models/portkey/portkey.py +7 -6
- agno/models/requesty/requesty.py +24 -4
- agno/models/response.py +73 -2
- agno/models/sambanova/sambanova.py +20 -3
- agno/models/siliconflow/siliconflow.py +19 -2
- agno/models/together/together.py +20 -3
- agno/models/utils.py +254 -8
- agno/models/vercel/v0.py +20 -3
- agno/models/vertexai/__init__.py +0 -0
- agno/models/vertexai/claude.py +190 -0
- agno/models/vllm/vllm.py +19 -14
- agno/models/xai/xai.py +19 -2
- agno/os/app.py +549 -152
- agno/os/auth.py +190 -3
- agno/os/config.py +23 -0
- agno/os/interfaces/a2a/router.py +8 -11
- agno/os/interfaces/a2a/utils.py +1 -1
- agno/os/interfaces/agui/router.py +18 -3
- agno/os/interfaces/agui/utils.py +152 -39
- agno/os/interfaces/slack/router.py +55 -37
- agno/os/interfaces/slack/slack.py +9 -1
- agno/os/interfaces/whatsapp/router.py +0 -1
- agno/os/interfaces/whatsapp/security.py +3 -1
- agno/os/mcp.py +110 -52
- agno/os/middleware/__init__.py +2 -0
- agno/os/middleware/jwt.py +676 -112
- agno/os/router.py +40 -1478
- agno/os/routers/agents/__init__.py +3 -0
- agno/os/routers/agents/router.py +599 -0
- agno/os/routers/agents/schema.py +261 -0
- agno/os/routers/evals/evals.py +96 -39
- agno/os/routers/evals/schemas.py +65 -33
- agno/os/routers/evals/utils.py +80 -10
- agno/os/routers/health.py +10 -4
- agno/os/routers/knowledge/knowledge.py +196 -38
- agno/os/routers/knowledge/schemas.py +82 -22
- agno/os/routers/memory/memory.py +279 -52
- agno/os/routers/memory/schemas.py +46 -17
- agno/os/routers/metrics/metrics.py +20 -8
- agno/os/routers/metrics/schemas.py +16 -16
- agno/os/routers/session/session.py +462 -34
- agno/os/routers/teams/__init__.py +3 -0
- agno/os/routers/teams/router.py +512 -0
- agno/os/routers/teams/schema.py +257 -0
- agno/os/routers/traces/__init__.py +3 -0
- agno/os/routers/traces/schemas.py +414 -0
- agno/os/routers/traces/traces.py +499 -0
- agno/os/routers/workflows/__init__.py +3 -0
- agno/os/routers/workflows/router.py +624 -0
- agno/os/routers/workflows/schema.py +75 -0
- agno/os/schema.py +256 -693
- agno/os/scopes.py +469 -0
- agno/os/utils.py +514 -36
- agno/reasoning/anthropic.py +80 -0
- agno/reasoning/gemini.py +73 -0
- agno/reasoning/openai.py +5 -0
- agno/reasoning/vertexai.py +76 -0
- agno/run/__init__.py +6 -0
- agno/run/agent.py +155 -32
- agno/run/base.py +55 -3
- agno/run/requirement.py +181 -0
- agno/run/team.py +125 -38
- agno/run/workflow.py +72 -18
- agno/session/agent.py +102 -89
- agno/session/summary.py +56 -15
- agno/session/team.py +164 -90
- agno/session/workflow.py +405 -40
- agno/table.py +10 -0
- agno/team/team.py +3974 -1903
- agno/tools/dalle.py +2 -4
- agno/tools/eleven_labs.py +23 -25
- agno/tools/exa.py +21 -16
- agno/tools/file.py +153 -23
- agno/tools/file_generation.py +16 -10
- agno/tools/firecrawl.py +15 -7
- agno/tools/function.py +193 -38
- agno/tools/gmail.py +238 -14
- agno/tools/google_drive.py +271 -0
- agno/tools/googlecalendar.py +36 -8
- agno/tools/googlesheets.py +20 -5
- agno/tools/jira.py +20 -0
- agno/tools/mcp/__init__.py +10 -0
- agno/tools/mcp/mcp.py +331 -0
- agno/tools/mcp/multi_mcp.py +347 -0
- agno/tools/mcp/params.py +24 -0
- agno/tools/mcp_toolbox.py +3 -3
- agno/tools/models/nebius.py +5 -5
- agno/tools/models_labs.py +20 -10
- agno/tools/nano_banana.py +151 -0
- agno/tools/notion.py +204 -0
- agno/tools/parallel.py +314 -0
- agno/tools/postgres.py +76 -36
- agno/tools/redshift.py +406 -0
- agno/tools/scrapegraph.py +1 -1
- agno/tools/shopify.py +1519 -0
- agno/tools/slack.py +18 -3
- agno/tools/spotify.py +919 -0
- agno/tools/tavily.py +146 -0
- agno/tools/toolkit.py +25 -0
- agno/tools/workflow.py +8 -1
- agno/tools/yfinance.py +12 -11
- agno/tracing/__init__.py +12 -0
- agno/tracing/exporter.py +157 -0
- agno/tracing/schemas.py +276 -0
- agno/tracing/setup.py +111 -0
- agno/utils/agent.py +938 -0
- agno/utils/cryptography.py +22 -0
- agno/utils/dttm.py +33 -0
- agno/utils/events.py +151 -3
- agno/utils/gemini.py +15 -5
- agno/utils/hooks.py +118 -4
- agno/utils/http.py +113 -2
- agno/utils/knowledge.py +12 -5
- agno/utils/log.py +1 -0
- agno/utils/mcp.py +92 -2
- agno/utils/media.py +187 -1
- agno/utils/merge_dict.py +3 -3
- agno/utils/message.py +60 -0
- agno/utils/models/ai_foundry.py +9 -2
- agno/utils/models/claude.py +49 -14
- agno/utils/models/cohere.py +9 -2
- agno/utils/models/llama.py +9 -2
- agno/utils/models/mistral.py +4 -2
- agno/utils/print_response/agent.py +109 -16
- agno/utils/print_response/team.py +223 -30
- agno/utils/print_response/workflow.py +251 -34
- agno/utils/streamlit.py +1 -1
- agno/utils/team.py +98 -9
- agno/utils/tokens.py +657 -0
- agno/vectordb/base.py +39 -7
- agno/vectordb/cassandra/cassandra.py +21 -5
- agno/vectordb/chroma/chromadb.py +43 -12
- agno/vectordb/clickhouse/clickhousedb.py +21 -5
- agno/vectordb/couchbase/couchbase.py +29 -5
- agno/vectordb/lancedb/lance_db.py +92 -181
- agno/vectordb/langchaindb/langchaindb.py +24 -4
- agno/vectordb/lightrag/lightrag.py +17 -3
- agno/vectordb/llamaindex/llamaindexdb.py +25 -5
- agno/vectordb/milvus/milvus.py +50 -37
- agno/vectordb/mongodb/__init__.py +7 -1
- agno/vectordb/mongodb/mongodb.py +36 -30
- agno/vectordb/pgvector/pgvector.py +201 -77
- agno/vectordb/pineconedb/pineconedb.py +41 -23
- agno/vectordb/qdrant/qdrant.py +67 -54
- agno/vectordb/redis/__init__.py +9 -0
- agno/vectordb/redis/redisdb.py +682 -0
- agno/vectordb/singlestore/singlestore.py +50 -29
- agno/vectordb/surrealdb/surrealdb.py +31 -41
- agno/vectordb/upstashdb/upstashdb.py +34 -6
- agno/vectordb/weaviate/weaviate.py +53 -14
- agno/workflow/__init__.py +2 -0
- agno/workflow/agent.py +299 -0
- agno/workflow/condition.py +120 -18
- agno/workflow/loop.py +77 -10
- agno/workflow/parallel.py +231 -143
- agno/workflow/router.py +118 -17
- agno/workflow/step.py +609 -170
- agno/workflow/steps.py +73 -6
- agno/workflow/types.py +96 -21
- agno/workflow/workflow.py +2039 -262
- {agno-2.1.2.dist-info → agno-2.3.13.dist-info}/METADATA +201 -66
- agno-2.3.13.dist-info/RECORD +613 -0
- agno/tools/googlesearch.py +0 -98
- agno/tools/mcp.py +0 -679
- agno/tools/memori.py +0 -339
- agno-2.1.2.dist-info/RECORD +0 -543
- {agno-2.1.2.dist-info → agno-2.3.13.dist-info}/WHEEL +0 -0
- {agno-2.1.2.dist-info → agno-2.3.13.dist-info}/licenses/LICENSE +0 -0
- {agno-2.1.2.dist-info → agno-2.3.13.dist-info}/top_level.txt +0 -0
agno/db/dynamo/dynamo.py
CHANGED
|
@@ -2,7 +2,10 @@ import json
|
|
|
2
2
|
import time
|
|
3
3
|
from datetime import date, datetime, timedelta, timezone
|
|
4
4
|
from os import getenv
|
|
5
|
-
from typing import Any, Dict, List, Optional, Tuple, Union
|
|
5
|
+
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
|
|
6
|
+
|
|
7
|
+
if TYPE_CHECKING:
|
|
8
|
+
from agno.tracing.schemas import Span, Trace
|
|
6
9
|
|
|
7
10
|
from agno.db.base import BaseDb, SessionType
|
|
8
11
|
from agno.db.dynamo.schemas import get_table_schema_definition
|
|
@@ -13,6 +16,7 @@ from agno.db.dynamo.utils import (
|
|
|
13
16
|
build_topic_filter_expression,
|
|
14
17
|
calculate_date_metrics,
|
|
15
18
|
create_table_if_not_exists,
|
|
19
|
+
deserialize_cultural_knowledge_from_db,
|
|
16
20
|
deserialize_eval_record,
|
|
17
21
|
deserialize_from_dynamodb_item,
|
|
18
22
|
deserialize_knowledge_row,
|
|
@@ -23,10 +27,12 @@ from agno.db.dynamo.utils import (
|
|
|
23
27
|
get_dates_to_calculate_metrics_for,
|
|
24
28
|
merge_with_existing_session,
|
|
25
29
|
prepare_session_data,
|
|
30
|
+
serialize_cultural_knowledge_for_db,
|
|
26
31
|
serialize_eval_record,
|
|
27
32
|
serialize_knowledge_row,
|
|
28
33
|
serialize_to_dynamo_item,
|
|
29
34
|
)
|
|
35
|
+
from agno.db.schemas.culture import CulturalKnowledge
|
|
30
36
|
from agno.db.schemas.evals import EvalFilterType, EvalRunRecord, EvalType
|
|
31
37
|
from agno.db.schemas.knowledge import KnowledgeRow
|
|
32
38
|
from agno.db.schemas.memory import UserMemory
|
|
@@ -52,10 +58,13 @@ class DynamoDb(BaseDb):
|
|
|
52
58
|
aws_access_key_id: Optional[str] = None,
|
|
53
59
|
aws_secret_access_key: Optional[str] = None,
|
|
54
60
|
session_table: Optional[str] = None,
|
|
61
|
+
culture_table: Optional[str] = None,
|
|
55
62
|
memory_table: Optional[str] = None,
|
|
56
63
|
metrics_table: Optional[str] = None,
|
|
57
64
|
eval_table: Optional[str] = None,
|
|
58
65
|
knowledge_table: Optional[str] = None,
|
|
66
|
+
traces_table: Optional[str] = None,
|
|
67
|
+
spans_table: Optional[str] = None,
|
|
59
68
|
id: Optional[str] = None,
|
|
60
69
|
):
|
|
61
70
|
"""
|
|
@@ -67,10 +76,13 @@ class DynamoDb(BaseDb):
|
|
|
67
76
|
aws_access_key_id: AWS access key ID.
|
|
68
77
|
aws_secret_access_key: AWS secret access key.
|
|
69
78
|
session_table: The name of the session table.
|
|
79
|
+
culture_table: The name of the culture table.
|
|
70
80
|
memory_table: The name of the memory table.
|
|
71
81
|
metrics_table: The name of the metrics table.
|
|
72
82
|
eval_table: The name of the eval table.
|
|
73
83
|
knowledge_table: The name of the knowledge table.
|
|
84
|
+
traces_table: The name of the traces table.
|
|
85
|
+
spans_table: The name of the spans table.
|
|
74
86
|
id: ID of the database.
|
|
75
87
|
"""
|
|
76
88
|
if id is None:
|
|
@@ -80,10 +92,13 @@ class DynamoDb(BaseDb):
|
|
|
80
92
|
super().__init__(
|
|
81
93
|
id=id,
|
|
82
94
|
session_table=session_table,
|
|
95
|
+
culture_table=culture_table,
|
|
83
96
|
memory_table=memory_table,
|
|
84
97
|
metrics_table=metrics_table,
|
|
85
98
|
eval_table=eval_table,
|
|
86
99
|
knowledge_table=knowledge_table,
|
|
100
|
+
traces_table=traces_table,
|
|
101
|
+
spans_table=spans_table,
|
|
87
102
|
)
|
|
88
103
|
|
|
89
104
|
if db_client is not None:
|
|
@@ -106,27 +121,8 @@ class DynamoDb(BaseDb):
|
|
|
106
121
|
session = boto3.Session(**session_kwargs)
|
|
107
122
|
self.client = session.client("dynamodb")
|
|
108
123
|
|
|
109
|
-
def
|
|
110
|
-
|
|
111
|
-
(self.session_table_name, "sessions"),
|
|
112
|
-
(self.memory_table_name, "memories"),
|
|
113
|
-
(self.metrics_table_name, "metrics"),
|
|
114
|
-
(self.eval_table_name, "evals"),
|
|
115
|
-
(self.knowledge_table_name, "knowledge_sources"),
|
|
116
|
-
]
|
|
117
|
-
|
|
118
|
-
for table_name, table_type in tables_to_create:
|
|
119
|
-
if table_name:
|
|
120
|
-
try:
|
|
121
|
-
schema = get_table_schema_definition(table_type)
|
|
122
|
-
schema["TableName"] = table_name
|
|
123
|
-
create_table_if_not_exists(self.client, table_name, schema)
|
|
124
|
-
|
|
125
|
-
except Exception as e:
|
|
126
|
-
log_error(f"Failed to create table {table_name}: {e}")
|
|
127
|
-
|
|
128
|
-
def _table_exists(self, table_name: str) -> bool:
|
|
129
|
-
"""Check if a DynamoDB table with the given name exists.
|
|
124
|
+
def table_exists(self, table_name: str) -> bool:
|
|
125
|
+
"""Check if a DynamoDB table exists.
|
|
130
126
|
|
|
131
127
|
Args:
|
|
132
128
|
table_name: The name of the table to check
|
|
@@ -139,16 +135,30 @@ class DynamoDb(BaseDb):
|
|
|
139
135
|
return True
|
|
140
136
|
except self.client.exceptions.ResourceNotFoundException:
|
|
141
137
|
return False
|
|
142
|
-
|
|
143
|
-
|
|
144
|
-
|
|
138
|
+
|
|
139
|
+
def _create_all_tables(self):
|
|
140
|
+
"""Create all configured DynamoDB tables if they don't exist."""
|
|
141
|
+
tables_to_create = [
|
|
142
|
+
("sessions", self.session_table_name),
|
|
143
|
+
("memories", self.memory_table_name),
|
|
144
|
+
("metrics", self.metrics_table_name),
|
|
145
|
+
("evals", self.eval_table_name),
|
|
146
|
+
("knowledge", self.knowledge_table_name),
|
|
147
|
+
("culture", self.culture_table_name),
|
|
148
|
+
]
|
|
149
|
+
|
|
150
|
+
for table_type, table_name in tables_to_create:
|
|
151
|
+
if not self.table_exists(table_name):
|
|
152
|
+
schema = get_table_schema_definition(table_type)
|
|
153
|
+
schema["TableName"] = table_name
|
|
154
|
+
create_table_if_not_exists(self.client, table_name, schema)
|
|
145
155
|
|
|
146
156
|
def _get_table(self, table_type: str, create_table_if_not_found: Optional[bool] = True) -> Optional[str]:
|
|
147
157
|
"""
|
|
148
158
|
Get table name and ensure the table exists, creating it if needed.
|
|
149
159
|
|
|
150
160
|
Args:
|
|
151
|
-
table_type: Type of table ("sessions", "memories", "metrics", "evals", "
|
|
161
|
+
table_type: Type of table ("sessions", "memories", "metrics", "evals", "knowledge", "culture", "traces", "spans")
|
|
152
162
|
|
|
153
163
|
Returns:
|
|
154
164
|
str: The table name
|
|
@@ -168,17 +178,33 @@ class DynamoDb(BaseDb):
|
|
|
168
178
|
table_name = self.eval_table_name
|
|
169
179
|
elif table_type == "knowledge":
|
|
170
180
|
table_name = self.knowledge_table_name
|
|
181
|
+
elif table_type == "culture":
|
|
182
|
+
table_name = self.culture_table_name
|
|
183
|
+
elif table_type == "traces":
|
|
184
|
+
table_name = self.trace_table_name
|
|
185
|
+
elif table_type == "spans":
|
|
186
|
+
# Ensure traces table exists first (spans reference traces)
|
|
187
|
+
self._get_table("traces", create_table_if_not_found=True)
|
|
188
|
+
table_name = self.span_table_name
|
|
171
189
|
else:
|
|
172
190
|
raise ValueError(f"Unknown table type: {table_type}")
|
|
173
191
|
|
|
174
192
|
# Check if table exists, create if it doesn't
|
|
175
|
-
if not self.
|
|
193
|
+
if not self.table_exists(table_name) and create_table_if_not_found:
|
|
176
194
|
schema = get_table_schema_definition(table_type)
|
|
177
195
|
schema["TableName"] = table_name
|
|
178
196
|
create_table_if_not_exists(self.client, table_name, schema)
|
|
179
197
|
|
|
180
198
|
return table_name
|
|
181
199
|
|
|
200
|
+
def get_latest_schema_version(self):
|
|
201
|
+
"""Get the latest version of the database schema."""
|
|
202
|
+
pass
|
|
203
|
+
|
|
204
|
+
def upsert_schema_version(self, version: str) -> None:
|
|
205
|
+
"""Upsert the schema version into the database."""
|
|
206
|
+
pass
|
|
207
|
+
|
|
182
208
|
# --- Sessions ---
|
|
183
209
|
|
|
184
210
|
def delete_session(self, session_id: Optional[str] = None) -> bool:
|
|
@@ -269,8 +295,6 @@ class DynamoDb(BaseDb):
|
|
|
269
295
|
|
|
270
296
|
session = deserialize_from_dynamodb_item(item)
|
|
271
297
|
|
|
272
|
-
if session.get("session_type") != session_type.value:
|
|
273
|
-
return None
|
|
274
298
|
if user_id and session.get("user_id") != user_id:
|
|
275
299
|
return None
|
|
276
300
|
|
|
@@ -524,7 +548,7 @@ class DynamoDb(BaseDb):
|
|
|
524
548
|
raise e
|
|
525
549
|
|
|
526
550
|
def upsert_sessions(
|
|
527
|
-
self, sessions: List[Session], deserialize: Optional[bool] = True
|
|
551
|
+
self, sessions: List[Session], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
|
|
528
552
|
) -> List[Union[Session, Dict[str, Any]]]:
|
|
529
553
|
"""
|
|
530
554
|
Bulk upsert multiple sessions for improved performance on large datasets.
|
|
@@ -845,6 +869,7 @@ class DynamoDb(BaseDb):
|
|
|
845
869
|
self,
|
|
846
870
|
limit: Optional[int] = None,
|
|
847
871
|
page: Optional[int] = None,
|
|
872
|
+
user_id: Optional[str] = None,
|
|
848
873
|
) -> Tuple[List[Dict[str, Any]], int]:
|
|
849
874
|
"""Get user memories stats.
|
|
850
875
|
|
|
@@ -872,7 +897,17 @@ class DynamoDb(BaseDb):
|
|
|
872
897
|
table_name = self._get_table("memories")
|
|
873
898
|
|
|
874
899
|
# Build filter expression for user_id if provided
|
|
900
|
+
filter_expression = None
|
|
901
|
+
expression_attribute_values = {}
|
|
902
|
+
if user_id:
|
|
903
|
+
filter_expression = "user_id = :user_id"
|
|
904
|
+
expression_attribute_values[":user_id"] = {"S": user_id}
|
|
905
|
+
|
|
875
906
|
scan_kwargs = {"TableName": table_name}
|
|
907
|
+
if filter_expression:
|
|
908
|
+
scan_kwargs["FilterExpression"] = filter_expression
|
|
909
|
+
if expression_attribute_values:
|
|
910
|
+
scan_kwargs["ExpressionAttributeValues"] = expression_attribute_values # type: ignore
|
|
876
911
|
|
|
877
912
|
response = self.client.scan(**scan_kwargs)
|
|
878
913
|
items = response.get("Items", [])
|
|
@@ -962,7 +997,7 @@ class DynamoDb(BaseDb):
|
|
|
962
997
|
raise e
|
|
963
998
|
|
|
964
999
|
def upsert_memories(
|
|
965
|
-
self, memories: List[UserMemory], deserialize: Optional[bool] = True
|
|
1000
|
+
self, memories: List[UserMemory], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
|
|
966
1001
|
) -> List[Union[UserMemory, Dict[str, Any]]]:
|
|
967
1002
|
"""
|
|
968
1003
|
Bulk upsert multiple user memories for improved performance on large datasets.
|
|
@@ -1454,17 +1489,17 @@ class DynamoDb(BaseDb):
|
|
|
1454
1489
|
"""
|
|
1455
1490
|
import json
|
|
1456
1491
|
|
|
1457
|
-
item = {}
|
|
1492
|
+
item: Dict[str, Any] = {}
|
|
1458
1493
|
for key, value in data.items():
|
|
1459
1494
|
if value is not None:
|
|
1460
1495
|
if isinstance(value, bool):
|
|
1461
|
-
item[key] = {"BOOL":
|
|
1496
|
+
item[key] = {"BOOL": value}
|
|
1462
1497
|
elif isinstance(value, (int, float)):
|
|
1463
1498
|
item[key] = {"N": str(value)}
|
|
1464
1499
|
elif isinstance(value, str):
|
|
1465
1500
|
item[key] = {"S": str(value)}
|
|
1466
1501
|
elif isinstance(value, (dict, list)):
|
|
1467
|
-
item[key] = {"S": json.dumps(
|
|
1502
|
+
item[key] = {"S": json.dumps(value)}
|
|
1468
1503
|
else:
|
|
1469
1504
|
item[key] = {"S": str(value)}
|
|
1470
1505
|
return item
|
|
@@ -1803,14 +1838,16 @@ class DynamoDb(BaseDb):
|
|
|
1803
1838
|
|
|
1804
1839
|
if filter_type is not None:
|
|
1805
1840
|
if filter_type == EvalFilterType.AGENT:
|
|
1806
|
-
filter_expressions.append("agent_id
|
|
1841
|
+
filter_expressions.append("attribute_exists(agent_id)")
|
|
1807
1842
|
elif filter_type == EvalFilterType.TEAM:
|
|
1808
|
-
filter_expressions.append("team_id
|
|
1843
|
+
filter_expressions.append("attribute_exists(team_id)")
|
|
1809
1844
|
elif filter_type == EvalFilterType.WORKFLOW:
|
|
1810
|
-
filter_expressions.append("workflow_id
|
|
1845
|
+
filter_expressions.append("attribute_exists(workflow_id)")
|
|
1811
1846
|
|
|
1812
1847
|
if filter_expressions:
|
|
1813
1848
|
scan_kwargs["FilterExpression"] = " AND ".join(filter_expressions)
|
|
1849
|
+
|
|
1850
|
+
if expression_values:
|
|
1814
1851
|
scan_kwargs["ExpressionAttributeValues"] = expression_values # type: ignore
|
|
1815
1852
|
|
|
1816
1853
|
# Execute scan
|
|
@@ -1883,3 +1920,862 @@ class DynamoDb(BaseDb):
|
|
|
1883
1920
|
except Exception as e:
|
|
1884
1921
|
log_error(f"Failed to rename eval run {eval_run_id}: {e}")
|
|
1885
1922
|
raise e
|
|
1923
|
+
|
|
1924
|
+
# -- Culture methods --
|
|
1925
|
+
|
|
1926
|
+
def clear_cultural_knowledge(self) -> None:
|
|
1927
|
+
"""Delete all cultural knowledge from the database."""
|
|
1928
|
+
try:
|
|
1929
|
+
table_name = self._get_table("culture")
|
|
1930
|
+
response = self.client.scan(TableName=table_name, ProjectionExpression="id")
|
|
1931
|
+
|
|
1932
|
+
with self.client.batch_writer(table_name) as batch:
|
|
1933
|
+
for item in response.get("Items", []):
|
|
1934
|
+
batch.delete_item(Key={"id": item["id"]})
|
|
1935
|
+
except Exception as e:
|
|
1936
|
+
log_error(f"Failed to clear cultural knowledge: {e}")
|
|
1937
|
+
raise e
|
|
1938
|
+
|
|
1939
|
+
def delete_cultural_knowledge(self, id: str) -> None:
|
|
1940
|
+
"""Delete a cultural knowledge entry from the database."""
|
|
1941
|
+
try:
|
|
1942
|
+
table_name = self._get_table("culture")
|
|
1943
|
+
self.client.delete_item(TableName=table_name, Key={"id": {"S": id}})
|
|
1944
|
+
except Exception as e:
|
|
1945
|
+
log_error(f"Failed to delete cultural knowledge {id}: {e}")
|
|
1946
|
+
raise e
|
|
1947
|
+
|
|
1948
|
+
def get_cultural_knowledge(
|
|
1949
|
+
self, id: str, deserialize: Optional[bool] = True
|
|
1950
|
+
) -> Optional[Union[CulturalKnowledge, Dict[str, Any]]]:
|
|
1951
|
+
"""Get a cultural knowledge entry from the database."""
|
|
1952
|
+
try:
|
|
1953
|
+
table_name = self._get_table("culture")
|
|
1954
|
+
response = self.client.get_item(TableName=table_name, Key={"id": {"S": id}})
|
|
1955
|
+
|
|
1956
|
+
item = response.get("Item")
|
|
1957
|
+
if not item:
|
|
1958
|
+
return None
|
|
1959
|
+
|
|
1960
|
+
db_row = deserialize_from_dynamodb_item(item)
|
|
1961
|
+
if not deserialize:
|
|
1962
|
+
return db_row
|
|
1963
|
+
|
|
1964
|
+
return deserialize_cultural_knowledge_from_db(db_row)
|
|
1965
|
+
except Exception as e:
|
|
1966
|
+
log_error(f"Failed to get cultural knowledge {id}: {e}")
|
|
1967
|
+
raise e
|
|
1968
|
+
|
|
1969
|
+
def get_all_cultural_knowledge(
|
|
1970
|
+
self,
|
|
1971
|
+
name: Optional[str] = None,
|
|
1972
|
+
agent_id: Optional[str] = None,
|
|
1973
|
+
team_id: Optional[str] = None,
|
|
1974
|
+
limit: Optional[int] = None,
|
|
1975
|
+
page: Optional[int] = None,
|
|
1976
|
+
sort_by: Optional[str] = None,
|
|
1977
|
+
sort_order: Optional[str] = None,
|
|
1978
|
+
deserialize: Optional[bool] = True,
|
|
1979
|
+
) -> Union[List[CulturalKnowledge], Tuple[List[Dict[str, Any]], int]]:
|
|
1980
|
+
"""Get all cultural knowledge from the database."""
|
|
1981
|
+
try:
|
|
1982
|
+
table_name = self._get_table("culture")
|
|
1983
|
+
|
|
1984
|
+
# Build filter expression
|
|
1985
|
+
filter_expressions = []
|
|
1986
|
+
expression_values = {}
|
|
1987
|
+
|
|
1988
|
+
if name:
|
|
1989
|
+
filter_expressions.append("#name = :name")
|
|
1990
|
+
expression_values[":name"] = {"S": name}
|
|
1991
|
+
if agent_id:
|
|
1992
|
+
filter_expressions.append("agent_id = :agent_id")
|
|
1993
|
+
expression_values[":agent_id"] = {"S": agent_id}
|
|
1994
|
+
if team_id:
|
|
1995
|
+
filter_expressions.append("team_id = :team_id")
|
|
1996
|
+
expression_values[":team_id"] = {"S": team_id}
|
|
1997
|
+
|
|
1998
|
+
scan_kwargs: Dict[str, Any] = {"TableName": table_name}
|
|
1999
|
+
if filter_expressions:
|
|
2000
|
+
scan_kwargs["FilterExpression"] = " AND ".join(filter_expressions)
|
|
2001
|
+
scan_kwargs["ExpressionAttributeValues"] = expression_values
|
|
2002
|
+
if name:
|
|
2003
|
+
scan_kwargs["ExpressionAttributeNames"] = {"#name": "name"}
|
|
2004
|
+
|
|
2005
|
+
# Execute scan
|
|
2006
|
+
response = self.client.scan(**scan_kwargs)
|
|
2007
|
+
items = response.get("Items", [])
|
|
2008
|
+
|
|
2009
|
+
# Continue scanning if there's more data
|
|
2010
|
+
while "LastEvaluatedKey" in response:
|
|
2011
|
+
scan_kwargs["ExclusiveStartKey"] = response["LastEvaluatedKey"]
|
|
2012
|
+
response = self.client.scan(**scan_kwargs)
|
|
2013
|
+
items.extend(response.get("Items", []))
|
|
2014
|
+
|
|
2015
|
+
# Deserialize items from DynamoDB format
|
|
2016
|
+
db_rows = [deserialize_from_dynamodb_item(item) for item in items]
|
|
2017
|
+
|
|
2018
|
+
# Apply sorting
|
|
2019
|
+
if sort_by:
|
|
2020
|
+
reverse = sort_order == "desc" if sort_order else False
|
|
2021
|
+
db_rows.sort(key=lambda x: x.get(sort_by, ""), reverse=reverse)
|
|
2022
|
+
|
|
2023
|
+
# Apply pagination
|
|
2024
|
+
total_count = len(db_rows)
|
|
2025
|
+
if limit and page:
|
|
2026
|
+
start = (page - 1) * limit
|
|
2027
|
+
db_rows = db_rows[start : start + limit]
|
|
2028
|
+
elif limit:
|
|
2029
|
+
db_rows = db_rows[:limit]
|
|
2030
|
+
|
|
2031
|
+
if not deserialize:
|
|
2032
|
+
return db_rows, total_count
|
|
2033
|
+
|
|
2034
|
+
return [deserialize_cultural_knowledge_from_db(row) for row in db_rows]
|
|
2035
|
+
except Exception as e:
|
|
2036
|
+
log_error(f"Failed to get all cultural knowledge: {e}")
|
|
2037
|
+
raise e
|
|
2038
|
+
|
|
2039
|
+
def upsert_cultural_knowledge(
|
|
2040
|
+
self, cultural_knowledge: CulturalKnowledge, deserialize: Optional[bool] = True
|
|
2041
|
+
) -> Optional[Union[CulturalKnowledge, Dict[str, Any]]]:
|
|
2042
|
+
"""Upsert a cultural knowledge entry into the database."""
|
|
2043
|
+
try:
|
|
2044
|
+
from uuid import uuid4
|
|
2045
|
+
|
|
2046
|
+
table_name = self._get_table("culture", create_table_if_not_found=True)
|
|
2047
|
+
|
|
2048
|
+
if not cultural_knowledge.id:
|
|
2049
|
+
cultural_knowledge.id = str(uuid4())
|
|
2050
|
+
|
|
2051
|
+
# Serialize content, categories, and notes into a dict for DB storage
|
|
2052
|
+
content_dict = serialize_cultural_knowledge_for_db(cultural_knowledge)
|
|
2053
|
+
|
|
2054
|
+
# Create the item dict with serialized content
|
|
2055
|
+
item_dict = {
|
|
2056
|
+
"id": cultural_knowledge.id,
|
|
2057
|
+
"name": cultural_knowledge.name,
|
|
2058
|
+
"summary": cultural_knowledge.summary,
|
|
2059
|
+
"content": content_dict if content_dict else None,
|
|
2060
|
+
"metadata": cultural_knowledge.metadata,
|
|
2061
|
+
"input": cultural_knowledge.input,
|
|
2062
|
+
"created_at": cultural_knowledge.created_at,
|
|
2063
|
+
"updated_at": int(time.time()),
|
|
2064
|
+
"agent_id": cultural_knowledge.agent_id,
|
|
2065
|
+
"team_id": cultural_knowledge.team_id,
|
|
2066
|
+
}
|
|
2067
|
+
|
|
2068
|
+
# Convert to DynamoDB format
|
|
2069
|
+
item = serialize_to_dynamo_item(item_dict)
|
|
2070
|
+
self.client.put_item(TableName=table_name, Item=item)
|
|
2071
|
+
|
|
2072
|
+
return self.get_cultural_knowledge(cultural_knowledge.id, deserialize=deserialize)
|
|
2073
|
+
|
|
2074
|
+
except Exception as e:
|
|
2075
|
+
log_error(f"Failed to upsert cultural knowledge: {e}")
|
|
2076
|
+
raise e
|
|
2077
|
+
|
|
2078
|
+
# --- Traces ---
|
|
2079
|
+
def upsert_trace(self, trace: "Trace") -> None:
|
|
2080
|
+
"""Create or update a single trace record in the database.
|
|
2081
|
+
|
|
2082
|
+
Args:
|
|
2083
|
+
trace: The Trace object to store (one per trace_id).
|
|
2084
|
+
"""
|
|
2085
|
+
try:
|
|
2086
|
+
table_name = self._get_table("traces", create_table_if_not_found=True)
|
|
2087
|
+
if table_name is None:
|
|
2088
|
+
return
|
|
2089
|
+
|
|
2090
|
+
# Check if trace already exists
|
|
2091
|
+
response = self.client.get_item(
|
|
2092
|
+
TableName=table_name,
|
|
2093
|
+
Key={"trace_id": {"S": trace.trace_id}},
|
|
2094
|
+
)
|
|
2095
|
+
|
|
2096
|
+
existing_item = response.get("Item")
|
|
2097
|
+
if existing_item:
|
|
2098
|
+
# Update existing trace
|
|
2099
|
+
existing = deserialize_from_dynamodb_item(existing_item)
|
|
2100
|
+
|
|
2101
|
+
# Determine component level for name update priority
|
|
2102
|
+
def get_component_level(workflow_id, team_id, agent_id, name):
|
|
2103
|
+
is_root_name = ".run" in name or ".arun" in name
|
|
2104
|
+
if not is_root_name:
|
|
2105
|
+
return 0
|
|
2106
|
+
elif workflow_id:
|
|
2107
|
+
return 3
|
|
2108
|
+
elif team_id:
|
|
2109
|
+
return 2
|
|
2110
|
+
elif agent_id:
|
|
2111
|
+
return 1
|
|
2112
|
+
else:
|
|
2113
|
+
return 0
|
|
2114
|
+
|
|
2115
|
+
existing_level = get_component_level(
|
|
2116
|
+
existing.get("workflow_id"),
|
|
2117
|
+
existing.get("team_id"),
|
|
2118
|
+
existing.get("agent_id"),
|
|
2119
|
+
existing.get("name", ""),
|
|
2120
|
+
)
|
|
2121
|
+
new_level = get_component_level(trace.workflow_id, trace.team_id, trace.agent_id, trace.name)
|
|
2122
|
+
should_update_name = new_level > existing_level
|
|
2123
|
+
|
|
2124
|
+
# Parse existing start_time to calculate correct duration
|
|
2125
|
+
existing_start_time_str = existing.get("start_time")
|
|
2126
|
+
if isinstance(existing_start_time_str, str):
|
|
2127
|
+
existing_start_time = datetime.fromisoformat(existing_start_time_str.replace("Z", "+00:00"))
|
|
2128
|
+
else:
|
|
2129
|
+
existing_start_time = trace.start_time
|
|
2130
|
+
|
|
2131
|
+
recalculated_duration_ms = int((trace.end_time - existing_start_time).total_seconds() * 1000)
|
|
2132
|
+
|
|
2133
|
+
# Build update expression
|
|
2134
|
+
update_parts = [
|
|
2135
|
+
"end_time = :end_time",
|
|
2136
|
+
"duration_ms = :duration_ms",
|
|
2137
|
+
"#status = :status",
|
|
2138
|
+
]
|
|
2139
|
+
expression_attr_names = {"#status": "status"}
|
|
2140
|
+
expression_attr_values: Dict[str, Any] = {
|
|
2141
|
+
":end_time": {"S": trace.end_time.isoformat()},
|
|
2142
|
+
":duration_ms": {"N": str(recalculated_duration_ms)},
|
|
2143
|
+
":status": {"S": trace.status},
|
|
2144
|
+
}
|
|
2145
|
+
|
|
2146
|
+
if should_update_name:
|
|
2147
|
+
update_parts.append("#name = :name")
|
|
2148
|
+
expression_attr_names["#name"] = "name"
|
|
2149
|
+
expression_attr_values[":name"] = {"S": trace.name}
|
|
2150
|
+
|
|
2151
|
+
if trace.run_id is not None:
|
|
2152
|
+
update_parts.append("run_id = :run_id")
|
|
2153
|
+
expression_attr_values[":run_id"] = {"S": trace.run_id}
|
|
2154
|
+
if trace.session_id is not None:
|
|
2155
|
+
update_parts.append("session_id = :session_id")
|
|
2156
|
+
expression_attr_values[":session_id"] = {"S": trace.session_id}
|
|
2157
|
+
if trace.user_id is not None:
|
|
2158
|
+
update_parts.append("user_id = :user_id")
|
|
2159
|
+
expression_attr_values[":user_id"] = {"S": trace.user_id}
|
|
2160
|
+
if trace.agent_id is not None:
|
|
2161
|
+
update_parts.append("agent_id = :agent_id")
|
|
2162
|
+
expression_attr_values[":agent_id"] = {"S": trace.agent_id}
|
|
2163
|
+
if trace.team_id is not None:
|
|
2164
|
+
update_parts.append("team_id = :team_id")
|
|
2165
|
+
expression_attr_values[":team_id"] = {"S": trace.team_id}
|
|
2166
|
+
if trace.workflow_id is not None:
|
|
2167
|
+
update_parts.append("workflow_id = :workflow_id")
|
|
2168
|
+
expression_attr_values[":workflow_id"] = {"S": trace.workflow_id}
|
|
2169
|
+
|
|
2170
|
+
self.client.update_item(
|
|
2171
|
+
TableName=table_name,
|
|
2172
|
+
Key={"trace_id": {"S": trace.trace_id}},
|
|
2173
|
+
UpdateExpression="SET " + ", ".join(update_parts),
|
|
2174
|
+
ExpressionAttributeNames=expression_attr_names,
|
|
2175
|
+
ExpressionAttributeValues=expression_attr_values,
|
|
2176
|
+
)
|
|
2177
|
+
else:
|
|
2178
|
+
# Create new trace with initialized counters
|
|
2179
|
+
trace_dict = trace.to_dict()
|
|
2180
|
+
trace_dict["total_spans"] = 0
|
|
2181
|
+
trace_dict["error_count"] = 0
|
|
2182
|
+
item = serialize_to_dynamo_item(trace_dict)
|
|
2183
|
+
self.client.put_item(TableName=table_name, Item=item)
|
|
2184
|
+
|
|
2185
|
+
except Exception as e:
|
|
2186
|
+
log_error(f"Error creating trace: {e}")
|
|
2187
|
+
|
|
2188
|
+
def get_trace(
|
|
2189
|
+
self,
|
|
2190
|
+
trace_id: Optional[str] = None,
|
|
2191
|
+
run_id: Optional[str] = None,
|
|
2192
|
+
):
|
|
2193
|
+
"""Get a single trace by trace_id or other filters.
|
|
2194
|
+
|
|
2195
|
+
Args:
|
|
2196
|
+
trace_id: The unique trace identifier.
|
|
2197
|
+
run_id: Filter by run ID (returns first match).
|
|
2198
|
+
|
|
2199
|
+
Returns:
|
|
2200
|
+
Optional[Trace]: The trace if found, None otherwise.
|
|
2201
|
+
|
|
2202
|
+
Note:
|
|
2203
|
+
If multiple filters are provided, trace_id takes precedence.
|
|
2204
|
+
For other filters, the most recent trace is returned.
|
|
2205
|
+
"""
|
|
2206
|
+
try:
|
|
2207
|
+
from agno.tracing.schemas import Trace
|
|
2208
|
+
|
|
2209
|
+
table_name = self._get_table("traces")
|
|
2210
|
+
if table_name is None:
|
|
2211
|
+
return None
|
|
2212
|
+
|
|
2213
|
+
if trace_id:
|
|
2214
|
+
# Direct lookup by primary key
|
|
2215
|
+
response = self.client.get_item(
|
|
2216
|
+
TableName=table_name,
|
|
2217
|
+
Key={"trace_id": {"S": trace_id}},
|
|
2218
|
+
)
|
|
2219
|
+
item = response.get("Item")
|
|
2220
|
+
if item:
|
|
2221
|
+
trace_data = deserialize_from_dynamodb_item(item)
|
|
2222
|
+
trace_data.setdefault("total_spans", 0)
|
|
2223
|
+
trace_data.setdefault("error_count", 0)
|
|
2224
|
+
return Trace.from_dict(trace_data)
|
|
2225
|
+
return None
|
|
2226
|
+
|
|
2227
|
+
elif run_id:
|
|
2228
|
+
# Query using GSI
|
|
2229
|
+
response = self.client.query(
|
|
2230
|
+
TableName=table_name,
|
|
2231
|
+
IndexName="run_id-start_time-index",
|
|
2232
|
+
KeyConditionExpression="run_id = :run_id",
|
|
2233
|
+
ExpressionAttributeValues={":run_id": {"S": run_id}},
|
|
2234
|
+
ScanIndexForward=False, # Descending order
|
|
2235
|
+
Limit=1,
|
|
2236
|
+
)
|
|
2237
|
+
items = response.get("Items", [])
|
|
2238
|
+
if items:
|
|
2239
|
+
trace_data = deserialize_from_dynamodb_item(items[0])
|
|
2240
|
+
# Use stored values (default to 0 if not present)
|
|
2241
|
+
trace_data.setdefault("total_spans", 0)
|
|
2242
|
+
trace_data.setdefault("error_count", 0)
|
|
2243
|
+
return Trace.from_dict(trace_data)
|
|
2244
|
+
return None
|
|
2245
|
+
|
|
2246
|
+
else:
|
|
2247
|
+
log_debug("get_trace called without any filter parameters")
|
|
2248
|
+
return None
|
|
2249
|
+
|
|
2250
|
+
except Exception as e:
|
|
2251
|
+
log_error(f"Error getting trace: {e}")
|
|
2252
|
+
return None
|
|
2253
|
+
|
|
2254
|
+
def get_traces(
|
|
2255
|
+
self,
|
|
2256
|
+
run_id: Optional[str] = None,
|
|
2257
|
+
session_id: Optional[str] = None,
|
|
2258
|
+
user_id: Optional[str] = None,
|
|
2259
|
+
agent_id: Optional[str] = None,
|
|
2260
|
+
team_id: Optional[str] = None,
|
|
2261
|
+
workflow_id: Optional[str] = None,
|
|
2262
|
+
status: Optional[str] = None,
|
|
2263
|
+
start_time: Optional[datetime] = None,
|
|
2264
|
+
end_time: Optional[datetime] = None,
|
|
2265
|
+
limit: Optional[int] = 20,
|
|
2266
|
+
page: Optional[int] = 1,
|
|
2267
|
+
) -> tuple[List, int]:
|
|
2268
|
+
"""Get traces matching the provided filters.
|
|
2269
|
+
|
|
2270
|
+
Args:
|
|
2271
|
+
run_id: Filter by run ID.
|
|
2272
|
+
session_id: Filter by session ID.
|
|
2273
|
+
user_id: Filter by user ID.
|
|
2274
|
+
agent_id: Filter by agent ID.
|
|
2275
|
+
team_id: Filter by team ID.
|
|
2276
|
+
workflow_id: Filter by workflow ID.
|
|
2277
|
+
status: Filter by status (OK, ERROR, UNSET).
|
|
2278
|
+
start_time: Filter traces starting after this datetime.
|
|
2279
|
+
end_time: Filter traces ending before this datetime.
|
|
2280
|
+
limit: Maximum number of traces to return per page.
|
|
2281
|
+
page: Page number (1-indexed).
|
|
2282
|
+
|
|
2283
|
+
Returns:
|
|
2284
|
+
tuple[List[Trace], int]: Tuple of (list of matching traces, total count).
|
|
2285
|
+
"""
|
|
2286
|
+
try:
|
|
2287
|
+
from agno.tracing.schemas import Trace
|
|
2288
|
+
|
|
2289
|
+
table_name = self._get_table("traces")
|
|
2290
|
+
if table_name is None:
|
|
2291
|
+
return [], 0
|
|
2292
|
+
|
|
2293
|
+
# Determine if we can use a GSI query or need to scan
|
|
2294
|
+
use_gsi = False
|
|
2295
|
+
gsi_name = None
|
|
2296
|
+
key_condition = None
|
|
2297
|
+
key_values: Dict[str, Any] = {}
|
|
2298
|
+
|
|
2299
|
+
# Check for GSI-compatible filters (only one can be used as key condition)
|
|
2300
|
+
if session_id:
|
|
2301
|
+
use_gsi = True
|
|
2302
|
+
gsi_name = "session_id-start_time-index"
|
|
2303
|
+
key_condition = "session_id = :session_id"
|
|
2304
|
+
key_values[":session_id"] = {"S": session_id}
|
|
2305
|
+
elif user_id:
|
|
2306
|
+
use_gsi = True
|
|
2307
|
+
gsi_name = "user_id-start_time-index"
|
|
2308
|
+
key_condition = "user_id = :user_id"
|
|
2309
|
+
key_values[":user_id"] = {"S": user_id}
|
|
2310
|
+
elif agent_id:
|
|
2311
|
+
use_gsi = True
|
|
2312
|
+
gsi_name = "agent_id-start_time-index"
|
|
2313
|
+
key_condition = "agent_id = :agent_id"
|
|
2314
|
+
key_values[":agent_id"] = {"S": agent_id}
|
|
2315
|
+
elif team_id:
|
|
2316
|
+
use_gsi = True
|
|
2317
|
+
gsi_name = "team_id-start_time-index"
|
|
2318
|
+
key_condition = "team_id = :team_id"
|
|
2319
|
+
key_values[":team_id"] = {"S": team_id}
|
|
2320
|
+
elif workflow_id:
|
|
2321
|
+
use_gsi = True
|
|
2322
|
+
gsi_name = "workflow_id-start_time-index"
|
|
2323
|
+
key_condition = "workflow_id = :workflow_id"
|
|
2324
|
+
key_values[":workflow_id"] = {"S": workflow_id}
|
|
2325
|
+
elif run_id:
|
|
2326
|
+
use_gsi = True
|
|
2327
|
+
gsi_name = "run_id-start_time-index"
|
|
2328
|
+
key_condition = "run_id = :run_id"
|
|
2329
|
+
key_values[":run_id"] = {"S": run_id}
|
|
2330
|
+
elif status:
|
|
2331
|
+
use_gsi = True
|
|
2332
|
+
gsi_name = "status-start_time-index"
|
|
2333
|
+
key_condition = "#status = :status"
|
|
2334
|
+
key_values[":status"] = {"S": status}
|
|
2335
|
+
|
|
2336
|
+
# Build filter expression for additional filters
|
|
2337
|
+
filter_parts = []
|
|
2338
|
+
filter_values: Dict[str, Any] = {}
|
|
2339
|
+
expression_attr_names: Dict[str, str] = {}
|
|
2340
|
+
|
|
2341
|
+
if start_time:
|
|
2342
|
+
filter_parts.append("start_time >= :start_time")
|
|
2343
|
+
filter_values[":start_time"] = {"S": start_time.isoformat()}
|
|
2344
|
+
if end_time:
|
|
2345
|
+
filter_parts.append("end_time <= :end_time")
|
|
2346
|
+
filter_values[":end_time"] = {"S": end_time.isoformat()}
|
|
2347
|
+
|
|
2348
|
+
if status and gsi_name != "status-start_time-index":
|
|
2349
|
+
filter_parts.append("#status = :filter_status")
|
|
2350
|
+
filter_values[":filter_status"] = {"S": status}
|
|
2351
|
+
expression_attr_names["#status"] = "status"
|
|
2352
|
+
|
|
2353
|
+
items = []
|
|
2354
|
+
if use_gsi and gsi_name and key_condition:
|
|
2355
|
+
# Use GSI query
|
|
2356
|
+
query_kwargs: Dict[str, Any] = {
|
|
2357
|
+
"TableName": table_name,
|
|
2358
|
+
"IndexName": gsi_name,
|
|
2359
|
+
"KeyConditionExpression": key_condition,
|
|
2360
|
+
"ExpressionAttributeValues": {**key_values, **filter_values},
|
|
2361
|
+
"ScanIndexForward": False, # Descending order by start_time
|
|
2362
|
+
}
|
|
2363
|
+
if gsi_name == "status-start_time-index":
|
|
2364
|
+
expression_attr_names["#status"] = "status"
|
|
2365
|
+
if expression_attr_names:
|
|
2366
|
+
query_kwargs["ExpressionAttributeNames"] = expression_attr_names
|
|
2367
|
+
if filter_parts:
|
|
2368
|
+
query_kwargs["FilterExpression"] = " AND ".join(filter_parts)
|
|
2369
|
+
|
|
2370
|
+
response = self.client.query(**query_kwargs)
|
|
2371
|
+
items.extend(response.get("Items", []))
|
|
2372
|
+
|
|
2373
|
+
while "LastEvaluatedKey" in response:
|
|
2374
|
+
query_kwargs["ExclusiveStartKey"] = response["LastEvaluatedKey"]
|
|
2375
|
+
response = self.client.query(**query_kwargs)
|
|
2376
|
+
items.extend(response.get("Items", []))
|
|
2377
|
+
else:
|
|
2378
|
+
# Use scan
|
|
2379
|
+
scan_kwargs: Dict[str, Any] = {"TableName": table_name}
|
|
2380
|
+
if filter_parts:
|
|
2381
|
+
scan_kwargs["FilterExpression"] = " AND ".join(filter_parts)
|
|
2382
|
+
scan_kwargs["ExpressionAttributeValues"] = filter_values
|
|
2383
|
+
if expression_attr_names:
|
|
2384
|
+
scan_kwargs["ExpressionAttributeNames"] = expression_attr_names
|
|
2385
|
+
|
|
2386
|
+
response = self.client.scan(**scan_kwargs)
|
|
2387
|
+
items.extend(response.get("Items", []))
|
|
2388
|
+
|
|
2389
|
+
while "LastEvaluatedKey" in response:
|
|
2390
|
+
scan_kwargs["ExclusiveStartKey"] = response["LastEvaluatedKey"]
|
|
2391
|
+
response = self.client.scan(**scan_kwargs)
|
|
2392
|
+
items.extend(response.get("Items", []))
|
|
2393
|
+
|
|
2394
|
+
# Deserialize items
|
|
2395
|
+
traces_data = [deserialize_from_dynamodb_item(item) for item in items]
|
|
2396
|
+
|
|
2397
|
+
# Sort by start_time descending
|
|
2398
|
+
traces_data.sort(key=lambda x: x.get("start_time", ""), reverse=True)
|
|
2399
|
+
|
|
2400
|
+
# Get total count
|
|
2401
|
+
total_count = len(traces_data)
|
|
2402
|
+
|
|
2403
|
+
# Apply pagination
|
|
2404
|
+
offset = (page - 1) * limit if page and limit else 0
|
|
2405
|
+
paginated_data = traces_data[offset : offset + limit] if limit else traces_data
|
|
2406
|
+
|
|
2407
|
+
# Use stored total_spans and error_count (default to 0 if not present)
|
|
2408
|
+
traces = []
|
|
2409
|
+
for trace_data in paginated_data:
|
|
2410
|
+
# Use stored values - these are updated by create_spans
|
|
2411
|
+
trace_data.setdefault("total_spans", 0)
|
|
2412
|
+
trace_data.setdefault("error_count", 0)
|
|
2413
|
+
traces.append(Trace.from_dict(trace_data))
|
|
2414
|
+
|
|
2415
|
+
return traces, total_count
|
|
2416
|
+
|
|
2417
|
+
except Exception as e:
|
|
2418
|
+
log_error(f"Error getting traces: {e}")
|
|
2419
|
+
return [], 0
|
|
2420
|
+
|
|
2421
|
+
def get_trace_stats(
|
|
2422
|
+
self,
|
|
2423
|
+
user_id: Optional[str] = None,
|
|
2424
|
+
agent_id: Optional[str] = None,
|
|
2425
|
+
team_id: Optional[str] = None,
|
|
2426
|
+
workflow_id: Optional[str] = None,
|
|
2427
|
+
start_time: Optional[datetime] = None,
|
|
2428
|
+
end_time: Optional[datetime] = None,
|
|
2429
|
+
limit: Optional[int] = 20,
|
|
2430
|
+
page: Optional[int] = 1,
|
|
2431
|
+
) -> tuple[List[Dict[str, Any]], int]:
|
|
2432
|
+
"""Get trace statistics grouped by session.
|
|
2433
|
+
|
|
2434
|
+
Args:
|
|
2435
|
+
user_id: Filter by user ID.
|
|
2436
|
+
agent_id: Filter by agent ID.
|
|
2437
|
+
team_id: Filter by team ID.
|
|
2438
|
+
workflow_id: Filter by workflow ID.
|
|
2439
|
+
start_time: Filter sessions with traces created after this datetime.
|
|
2440
|
+
end_time: Filter sessions with traces created before this datetime.
|
|
2441
|
+
limit: Maximum number of sessions to return per page.
|
|
2442
|
+
page: Page number (1-indexed).
|
|
2443
|
+
|
|
2444
|
+
Returns:
|
|
2445
|
+
tuple[List[Dict], int]: Tuple of (list of session stats dicts, total count).
|
|
2446
|
+
Each dict contains: session_id, user_id, agent_id, team_id, workflow_id, total_traces,
|
|
2447
|
+
first_trace_at, last_trace_at.
|
|
2448
|
+
"""
|
|
2449
|
+
try:
|
|
2450
|
+
table_name = self._get_table("traces")
|
|
2451
|
+
if table_name is None:
|
|
2452
|
+
return [], 0
|
|
2453
|
+
|
|
2454
|
+
# Fetch all traces and aggregate in memory (DynamoDB doesn't support GROUP BY)
|
|
2455
|
+
scan_kwargs: Dict[str, Any] = {"TableName": table_name}
|
|
2456
|
+
|
|
2457
|
+
# Build filter expression
|
|
2458
|
+
filter_parts = []
|
|
2459
|
+
filter_values: Dict[str, Any] = {}
|
|
2460
|
+
|
|
2461
|
+
if user_id:
|
|
2462
|
+
filter_parts.append("user_id = :user_id")
|
|
2463
|
+
filter_values[":user_id"] = {"S": user_id}
|
|
2464
|
+
if agent_id:
|
|
2465
|
+
filter_parts.append("agent_id = :agent_id")
|
|
2466
|
+
filter_values[":agent_id"] = {"S": agent_id}
|
|
2467
|
+
if team_id:
|
|
2468
|
+
filter_parts.append("team_id = :team_id")
|
|
2469
|
+
filter_values[":team_id"] = {"S": team_id}
|
|
2470
|
+
if workflow_id:
|
|
2471
|
+
filter_parts.append("workflow_id = :workflow_id")
|
|
2472
|
+
filter_values[":workflow_id"] = {"S": workflow_id}
|
|
2473
|
+
if start_time:
|
|
2474
|
+
filter_parts.append("created_at >= :start_time")
|
|
2475
|
+
filter_values[":start_time"] = {"S": start_time.isoformat()}
|
|
2476
|
+
if end_time:
|
|
2477
|
+
filter_parts.append("created_at <= :end_time")
|
|
2478
|
+
filter_values[":end_time"] = {"S": end_time.isoformat()}
|
|
2479
|
+
|
|
2480
|
+
# Filter for records with session_id
|
|
2481
|
+
filter_parts.append("attribute_exists(session_id)")
|
|
2482
|
+
|
|
2483
|
+
if filter_parts:
|
|
2484
|
+
scan_kwargs["FilterExpression"] = " AND ".join(filter_parts)
|
|
2485
|
+
if filter_values:
|
|
2486
|
+
scan_kwargs["ExpressionAttributeValues"] = filter_values
|
|
2487
|
+
|
|
2488
|
+
# Scan all matching traces
|
|
2489
|
+
items = []
|
|
2490
|
+
response = self.client.scan(**scan_kwargs)
|
|
2491
|
+
items.extend(response.get("Items", []))
|
|
2492
|
+
|
|
2493
|
+
while "LastEvaluatedKey" in response:
|
|
2494
|
+
scan_kwargs["ExclusiveStartKey"] = response["LastEvaluatedKey"]
|
|
2495
|
+
response = self.client.scan(**scan_kwargs)
|
|
2496
|
+
items.extend(response.get("Items", []))
|
|
2497
|
+
|
|
2498
|
+
# Aggregate by session_id
|
|
2499
|
+
session_stats: Dict[str, Dict[str, Any]] = {}
|
|
2500
|
+
for item in items:
|
|
2501
|
+
trace_data = deserialize_from_dynamodb_item(item)
|
|
2502
|
+
session_id = trace_data.get("session_id")
|
|
2503
|
+
if not session_id:
|
|
2504
|
+
continue
|
|
2505
|
+
|
|
2506
|
+
if session_id not in session_stats:
|
|
2507
|
+
session_stats[session_id] = {
|
|
2508
|
+
"session_id": session_id,
|
|
2509
|
+
"user_id": trace_data.get("user_id"),
|
|
2510
|
+
"agent_id": trace_data.get("agent_id"),
|
|
2511
|
+
"team_id": trace_data.get("team_id"),
|
|
2512
|
+
"workflow_id": trace_data.get("workflow_id"),
|
|
2513
|
+
"total_traces": 0,
|
|
2514
|
+
"first_trace_at": trace_data.get("created_at"),
|
|
2515
|
+
"last_trace_at": trace_data.get("created_at"),
|
|
2516
|
+
}
|
|
2517
|
+
|
|
2518
|
+
session_stats[session_id]["total_traces"] += 1
|
|
2519
|
+
|
|
2520
|
+
created_at = trace_data.get("created_at")
|
|
2521
|
+
if (
|
|
2522
|
+
created_at
|
|
2523
|
+
and session_stats[session_id]["first_trace_at"]
|
|
2524
|
+
and session_stats[session_id]["last_trace_at"]
|
|
2525
|
+
):
|
|
2526
|
+
if created_at < session_stats[session_id]["first_trace_at"]:
|
|
2527
|
+
session_stats[session_id]["first_trace_at"] = created_at
|
|
2528
|
+
if created_at > session_stats[session_id]["last_trace_at"]:
|
|
2529
|
+
session_stats[session_id]["last_trace_at"] = created_at
|
|
2530
|
+
|
|
2531
|
+
# Convert to list and sort by last_trace_at descending
|
|
2532
|
+
stats_list = list(session_stats.values())
|
|
2533
|
+
stats_list.sort(key=lambda x: x.get("last_trace_at", ""), reverse=True)
|
|
2534
|
+
|
|
2535
|
+
# Convert datetime strings to datetime objects
|
|
2536
|
+
for stat in stats_list:
|
|
2537
|
+
first_trace_at = stat["first_trace_at"]
|
|
2538
|
+
last_trace_at = stat["last_trace_at"]
|
|
2539
|
+
if isinstance(first_trace_at, str):
|
|
2540
|
+
stat["first_trace_at"] = datetime.fromisoformat(first_trace_at.replace("Z", "+00:00"))
|
|
2541
|
+
if isinstance(last_trace_at, str):
|
|
2542
|
+
stat["last_trace_at"] = datetime.fromisoformat(last_trace_at.replace("Z", "+00:00"))
|
|
2543
|
+
|
|
2544
|
+
# Get total count
|
|
2545
|
+
total_count = len(stats_list)
|
|
2546
|
+
|
|
2547
|
+
# Apply pagination
|
|
2548
|
+
offset = (page - 1) * limit if page and limit else 0
|
|
2549
|
+
paginated_stats = stats_list[offset : offset + limit] if limit else stats_list
|
|
2550
|
+
|
|
2551
|
+
return paginated_stats, total_count
|
|
2552
|
+
|
|
2553
|
+
except Exception as e:
|
|
2554
|
+
log_error(f"Error getting trace stats: {e}")
|
|
2555
|
+
return [], 0
|
|
2556
|
+
|
|
2557
|
+
# --- Spans ---
|
|
2558
|
+
def create_span(self, span: "Span") -> None:
|
|
2559
|
+
"""Create a single span in the database.
|
|
2560
|
+
|
|
2561
|
+
Args:
|
|
2562
|
+
span: The Span object to store.
|
|
2563
|
+
"""
|
|
2564
|
+
try:
|
|
2565
|
+
table_name = self._get_table("spans", create_table_if_not_found=True)
|
|
2566
|
+
if table_name is None:
|
|
2567
|
+
return
|
|
2568
|
+
|
|
2569
|
+
span_dict = span.to_dict()
|
|
2570
|
+
# Serialize attributes as JSON string
|
|
2571
|
+
if "attributes" in span_dict and isinstance(span_dict["attributes"], dict):
|
|
2572
|
+
span_dict["attributes"] = json.dumps(span_dict["attributes"])
|
|
2573
|
+
|
|
2574
|
+
item = serialize_to_dynamo_item(span_dict)
|
|
2575
|
+
self.client.put_item(TableName=table_name, Item=item)
|
|
2576
|
+
|
|
2577
|
+
# Increment total_spans and error_count on trace
|
|
2578
|
+
traces_table_name = self._get_table("traces")
|
|
2579
|
+
if traces_table_name:
|
|
2580
|
+
try:
|
|
2581
|
+
update_expr = "ADD total_spans :inc"
|
|
2582
|
+
expr_values: Dict[str, Any] = {":inc": {"N": "1"}}
|
|
2583
|
+
|
|
2584
|
+
if span.status_code == "ERROR":
|
|
2585
|
+
update_expr += ", error_count :inc"
|
|
2586
|
+
|
|
2587
|
+
self.client.update_item(
|
|
2588
|
+
TableName=traces_table_name,
|
|
2589
|
+
Key={"trace_id": {"S": span.trace_id}},
|
|
2590
|
+
UpdateExpression=update_expr,
|
|
2591
|
+
ExpressionAttributeValues=expr_values,
|
|
2592
|
+
)
|
|
2593
|
+
except Exception as update_error:
|
|
2594
|
+
log_debug(f"Could not update trace span counts: {update_error}")
|
|
2595
|
+
|
|
2596
|
+
except Exception as e:
|
|
2597
|
+
log_error(f"Error creating span: {e}")
|
|
2598
|
+
|
|
2599
|
+
def create_spans(self, spans: List) -> None:
|
|
2600
|
+
"""Create multiple spans in the database as a batch.
|
|
2601
|
+
|
|
2602
|
+
Args:
|
|
2603
|
+
spans: List of Span objects to store.
|
|
2604
|
+
"""
|
|
2605
|
+
if not spans:
|
|
2606
|
+
return
|
|
2607
|
+
|
|
2608
|
+
try:
|
|
2609
|
+
table_name = self._get_table("spans", create_table_if_not_found=True)
|
|
2610
|
+
if table_name is None:
|
|
2611
|
+
return
|
|
2612
|
+
|
|
2613
|
+
for i in range(0, len(spans), DYNAMO_BATCH_SIZE_LIMIT):
|
|
2614
|
+
batch = spans[i : i + DYNAMO_BATCH_SIZE_LIMIT]
|
|
2615
|
+
put_requests = []
|
|
2616
|
+
|
|
2617
|
+
for span in batch:
|
|
2618
|
+
span_dict = span.to_dict()
|
|
2619
|
+
# Serialize attributes as JSON string
|
|
2620
|
+
if "attributes" in span_dict and isinstance(span_dict["attributes"], dict):
|
|
2621
|
+
span_dict["attributes"] = json.dumps(span_dict["attributes"])
|
|
2622
|
+
|
|
2623
|
+
item = serialize_to_dynamo_item(span_dict)
|
|
2624
|
+
put_requests.append({"PutRequest": {"Item": item}})
|
|
2625
|
+
|
|
2626
|
+
if put_requests:
|
|
2627
|
+
self.client.batch_write_item(RequestItems={table_name: put_requests})
|
|
2628
|
+
|
|
2629
|
+
# Update trace with total_spans and error_count using ADD (atomic increment)
|
|
2630
|
+
trace_id = spans[0].trace_id
|
|
2631
|
+
spans_count = len(spans)
|
|
2632
|
+
error_count = sum(1 for s in spans if s.status_code == "ERROR")
|
|
2633
|
+
|
|
2634
|
+
traces_table_name = self._get_table("traces")
|
|
2635
|
+
if traces_table_name:
|
|
2636
|
+
try:
|
|
2637
|
+
# Use ADD for atomic increment - works even if attributes don't exist yet
|
|
2638
|
+
update_expr = "ADD total_spans :spans_inc"
|
|
2639
|
+
expr_values: Dict[str, Any] = {":spans_inc": {"N": str(spans_count)}}
|
|
2640
|
+
|
|
2641
|
+
if error_count > 0:
|
|
2642
|
+
update_expr += ", error_count :error_inc"
|
|
2643
|
+
expr_values[":error_inc"] = {"N": str(error_count)}
|
|
2644
|
+
|
|
2645
|
+
self.client.update_item(
|
|
2646
|
+
TableName=traces_table_name,
|
|
2647
|
+
Key={"trace_id": {"S": trace_id}},
|
|
2648
|
+
UpdateExpression=update_expr,
|
|
2649
|
+
ExpressionAttributeValues=expr_values,
|
|
2650
|
+
)
|
|
2651
|
+
except Exception as update_error:
|
|
2652
|
+
log_debug(f"Could not update trace span counts: {update_error}")
|
|
2653
|
+
|
|
2654
|
+
except Exception as e:
|
|
2655
|
+
log_error(f"Error creating spans batch: {e}")
|
|
2656
|
+
|
|
2657
|
+
def get_span(self, span_id: str):
|
|
2658
|
+
"""Get a single span by its span_id.
|
|
2659
|
+
|
|
2660
|
+
Args:
|
|
2661
|
+
span_id: The unique span identifier.
|
|
2662
|
+
|
|
2663
|
+
Returns:
|
|
2664
|
+
Optional[Span]: The span if found, None otherwise.
|
|
2665
|
+
"""
|
|
2666
|
+
try:
|
|
2667
|
+
from agno.tracing.schemas import Span
|
|
2668
|
+
|
|
2669
|
+
table_name = self._get_table("spans")
|
|
2670
|
+
if table_name is None:
|
|
2671
|
+
return None
|
|
2672
|
+
|
|
2673
|
+
response = self.client.get_item(
|
|
2674
|
+
TableName=table_name,
|
|
2675
|
+
Key={"span_id": {"S": span_id}},
|
|
2676
|
+
)
|
|
2677
|
+
|
|
2678
|
+
item = response.get("Item")
|
|
2679
|
+
if item:
|
|
2680
|
+
span_data = deserialize_from_dynamodb_item(item)
|
|
2681
|
+
# Deserialize attributes from JSON string
|
|
2682
|
+
if "attributes" in span_data and isinstance(span_data["attributes"], str):
|
|
2683
|
+
span_data["attributes"] = json.loads(span_data["attributes"])
|
|
2684
|
+
return Span.from_dict(span_data)
|
|
2685
|
+
return None
|
|
2686
|
+
|
|
2687
|
+
except Exception as e:
|
|
2688
|
+
log_error(f"Error getting span: {e}")
|
|
2689
|
+
return None
|
|
2690
|
+
|
|
2691
|
+
def get_spans(
|
|
2692
|
+
self,
|
|
2693
|
+
trace_id: Optional[str] = None,
|
|
2694
|
+
parent_span_id: Optional[str] = None,
|
|
2695
|
+
limit: Optional[int] = 1000,
|
|
2696
|
+
) -> List:
|
|
2697
|
+
"""Get spans matching the provided filters.
|
|
2698
|
+
|
|
2699
|
+
Args:
|
|
2700
|
+
trace_id: Filter by trace ID.
|
|
2701
|
+
parent_span_id: Filter by parent span ID.
|
|
2702
|
+
limit: Maximum number of spans to return.
|
|
2703
|
+
|
|
2704
|
+
Returns:
|
|
2705
|
+
List[Span]: List of matching spans.
|
|
2706
|
+
"""
|
|
2707
|
+
try:
|
|
2708
|
+
from agno.tracing.schemas import Span
|
|
2709
|
+
|
|
2710
|
+
table_name = self._get_table("spans")
|
|
2711
|
+
if table_name is None:
|
|
2712
|
+
return []
|
|
2713
|
+
|
|
2714
|
+
items = []
|
|
2715
|
+
|
|
2716
|
+
if trace_id:
|
|
2717
|
+
# Use GSI query
|
|
2718
|
+
query_kwargs: Dict[str, Any] = {
|
|
2719
|
+
"TableName": table_name,
|
|
2720
|
+
"IndexName": "trace_id-start_time-index",
|
|
2721
|
+
"KeyConditionExpression": "trace_id = :trace_id",
|
|
2722
|
+
"ExpressionAttributeValues": {":trace_id": {"S": trace_id}},
|
|
2723
|
+
}
|
|
2724
|
+
if limit:
|
|
2725
|
+
query_kwargs["Limit"] = limit
|
|
2726
|
+
|
|
2727
|
+
response = self.client.query(**query_kwargs)
|
|
2728
|
+
items.extend(response.get("Items", []))
|
|
2729
|
+
|
|
2730
|
+
while "LastEvaluatedKey" in response and (limit is None or len(items) < limit):
|
|
2731
|
+
query_kwargs["ExclusiveStartKey"] = response["LastEvaluatedKey"]
|
|
2732
|
+
response = self.client.query(**query_kwargs)
|
|
2733
|
+
items.extend(response.get("Items", []))
|
|
2734
|
+
|
|
2735
|
+
elif parent_span_id:
|
|
2736
|
+
# Use GSI query
|
|
2737
|
+
query_kwargs = {
|
|
2738
|
+
"TableName": table_name,
|
|
2739
|
+
"IndexName": "parent_span_id-start_time-index",
|
|
2740
|
+
"KeyConditionExpression": "parent_span_id = :parent_span_id",
|
|
2741
|
+
"ExpressionAttributeValues": {":parent_span_id": {"S": parent_span_id}},
|
|
2742
|
+
}
|
|
2743
|
+
if limit:
|
|
2744
|
+
query_kwargs["Limit"] = limit
|
|
2745
|
+
|
|
2746
|
+
response = self.client.query(**query_kwargs)
|
|
2747
|
+
items.extend(response.get("Items", []))
|
|
2748
|
+
|
|
2749
|
+
while "LastEvaluatedKey" in response and (limit is None or len(items) < limit):
|
|
2750
|
+
query_kwargs["ExclusiveStartKey"] = response["LastEvaluatedKey"]
|
|
2751
|
+
response = self.client.query(**query_kwargs)
|
|
2752
|
+
items.extend(response.get("Items", []))
|
|
2753
|
+
|
|
2754
|
+
else:
|
|
2755
|
+
# Scan all spans
|
|
2756
|
+
scan_kwargs: Dict[str, Any] = {"TableName": table_name}
|
|
2757
|
+
if limit:
|
|
2758
|
+
scan_kwargs["Limit"] = limit
|
|
2759
|
+
|
|
2760
|
+
response = self.client.scan(**scan_kwargs)
|
|
2761
|
+
items.extend(response.get("Items", []))
|
|
2762
|
+
|
|
2763
|
+
while "LastEvaluatedKey" in response and (limit is None or len(items) < limit):
|
|
2764
|
+
scan_kwargs["ExclusiveStartKey"] = response["LastEvaluatedKey"]
|
|
2765
|
+
response = self.client.scan(**scan_kwargs)
|
|
2766
|
+
items.extend(response.get("Items", []))
|
|
2767
|
+
|
|
2768
|
+
# Deserialize items
|
|
2769
|
+
spans = []
|
|
2770
|
+
for item in items[:limit] if limit else items:
|
|
2771
|
+
span_data = deserialize_from_dynamodb_item(item)
|
|
2772
|
+
# Deserialize attributes from JSON string
|
|
2773
|
+
if "attributes" in span_data and isinstance(span_data["attributes"], str):
|
|
2774
|
+
span_data["attributes"] = json.loads(span_data["attributes"])
|
|
2775
|
+
spans.append(Span.from_dict(span_data))
|
|
2776
|
+
|
|
2777
|
+
return spans
|
|
2778
|
+
|
|
2779
|
+
except Exception as e:
|
|
2780
|
+
log_error(f"Error getting spans: {e}")
|
|
2781
|
+
return []
|