agno 2.0.0rc2__py3-none-any.whl → 2.3.0__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 +6009 -2874
- agno/api/api.py +2 -0
- agno/api/os.py +1 -1
- agno/culture/__init__.py +3 -0
- agno/culture/manager.py +956 -0
- agno/db/async_postgres/__init__.py +3 -0
- agno/db/base.py +385 -6
- agno/db/dynamo/dynamo.py +388 -81
- agno/db/dynamo/schemas.py +47 -10
- agno/db/dynamo/utils.py +63 -4
- agno/db/firestore/firestore.py +435 -64
- agno/db/firestore/schemas.py +11 -0
- agno/db/firestore/utils.py +102 -4
- agno/db/gcs_json/gcs_json_db.py +384 -42
- agno/db/gcs_json/utils.py +60 -26
- agno/db/in_memory/in_memory_db.py +351 -66
- agno/db/in_memory/utils.py +60 -2
- agno/db/json/json_db.py +339 -48
- agno/db/json/utils.py +60 -26
- agno/db/migrations/manager.py +199 -0
- agno/db/migrations/v1_to_v2.py +510 -37
- 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 +2036 -0
- agno/db/mongo/mongo.py +653 -76
- agno/db/mongo/schemas.py +13 -0
- agno/db/mongo/utils.py +80 -8
- agno/db/mysql/mysql.py +687 -25
- agno/db/mysql/schemas.py +61 -37
- agno/db/mysql/utils.py +60 -2
- agno/db/postgres/__init__.py +2 -1
- agno/db/postgres/async_postgres.py +2001 -0
- agno/db/postgres/postgres.py +676 -57
- agno/db/postgres/schemas.py +43 -18
- agno/db/postgres/utils.py +164 -2
- agno/db/redis/redis.py +344 -38
- agno/db/redis/schemas.py +18 -0
- agno/db/redis/utils.py +60 -2
- agno/db/schemas/__init__.py +2 -1
- agno/db/schemas/culture.py +120 -0
- agno/db/schemas/memory.py +13 -0
- agno/db/singlestore/schemas.py +26 -1
- agno/db/singlestore/singlestore.py +687 -53
- agno/db/singlestore/utils.py +60 -2
- agno/db/sqlite/__init__.py +2 -1
- agno/db/sqlite/async_sqlite.py +2371 -0
- agno/db/sqlite/schemas.py +24 -0
- agno/db/sqlite/sqlite.py +774 -85
- agno/db/sqlite/utils.py +168 -5
- agno/db/surrealdb/__init__.py +3 -0
- agno/db/surrealdb/metrics.py +292 -0
- agno/db/surrealdb/models.py +309 -0
- agno/db/surrealdb/queries.py +71 -0
- agno/db/surrealdb/surrealdb.py +1361 -0
- agno/db/surrealdb/utils.py +147 -0
- agno/db/utils.py +50 -22
- agno/eval/accuracy.py +50 -43
- agno/eval/performance.py +6 -3
- agno/eval/reliability.py +6 -3
- agno/eval/utils.py +33 -16
- agno/exceptions.py +68 -1
- agno/filters.py +354 -0
- agno/guardrails/__init__.py +6 -0
- agno/guardrails/base.py +19 -0
- agno/guardrails/openai.py +144 -0
- agno/guardrails/pii.py +94 -0
- agno/guardrails/prompt_injection.py +52 -0
- agno/integrations/discord/client.py +1 -0
- agno/knowledge/chunking/agentic.py +13 -10
- agno/knowledge/chunking/fixed.py +1 -1
- agno/knowledge/chunking/semantic.py +40 -8
- agno/knowledge/chunking/strategy.py +59 -15
- agno/knowledge/embedder/aws_bedrock.py +9 -4
- agno/knowledge/embedder/azure_openai.py +54 -0
- agno/knowledge/embedder/base.py +2 -0
- agno/knowledge/embedder/cohere.py +184 -5
- agno/knowledge/embedder/fastembed.py +1 -1
- agno/knowledge/embedder/google.py +79 -1
- agno/knowledge/embedder/huggingface.py +9 -4
- agno/knowledge/embedder/jina.py +63 -0
- agno/knowledge/embedder/mistral.py +78 -11
- agno/knowledge/embedder/nebius.py +1 -1
- agno/knowledge/embedder/ollama.py +13 -0
- agno/knowledge/embedder/openai.py +37 -65
- agno/knowledge/embedder/sentence_transformer.py +8 -4
- agno/knowledge/embedder/vllm.py +262 -0
- agno/knowledge/embedder/voyageai.py +69 -16
- agno/knowledge/knowledge.py +595 -187
- agno/knowledge/reader/base.py +9 -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 +290 -0
- agno/knowledge/reader/json_reader.py +6 -5
- agno/knowledge/reader/markdown_reader.py +13 -13
- agno/knowledge/reader/pdf_reader.py +43 -68
- agno/knowledge/reader/pptx_reader.py +101 -0
- agno/knowledge/reader/reader_factory.py +51 -6
- agno/knowledge/reader/s3_reader.py +3 -15
- agno/knowledge/reader/tavily_reader.py +194 -0
- agno/knowledge/reader/text_reader.py +13 -13
- agno/knowledge/reader/web_search_reader.py +2 -43
- agno/knowledge/reader/website_reader.py +43 -25
- agno/knowledge/reranker/__init__.py +3 -0
- agno/knowledge/types.py +9 -0
- agno/knowledge/utils.py +20 -0
- agno/media.py +339 -266
- agno/memory/manager.py +336 -82
- agno/models/aimlapi/aimlapi.py +2 -2
- agno/models/anthropic/claude.py +183 -37
- agno/models/aws/bedrock.py +52 -112
- agno/models/aws/claude.py +33 -1
- agno/models/azure/ai_foundry.py +33 -15
- agno/models/azure/openai_chat.py +25 -8
- agno/models/base.py +1011 -566
- agno/models/cerebras/cerebras.py +19 -13
- agno/models/cerebras/cerebras_openai.py +8 -5
- agno/models/cohere/chat.py +27 -1
- agno/models/cometapi/__init__.py +5 -0
- agno/models/cometapi/cometapi.py +57 -0
- agno/models/dashscope/dashscope.py +1 -0
- agno/models/deepinfra/deepinfra.py +2 -2
- agno/models/deepseek/deepseek.py +2 -2
- agno/models/fireworks/fireworks.py +2 -2
- agno/models/google/gemini.py +110 -37
- agno/models/groq/groq.py +28 -11
- agno/models/huggingface/huggingface.py +2 -1
- agno/models/internlm/internlm.py +2 -2
- agno/models/langdb/langdb.py +4 -4
- agno/models/litellm/chat.py +18 -1
- agno/models/litellm/litellm_openai.py +2 -2
- agno/models/llama_cpp/__init__.py +5 -0
- agno/models/llama_cpp/llama_cpp.py +22 -0
- agno/models/message.py +143 -4
- agno/models/meta/llama.py +27 -10
- agno/models/meta/llama_openai.py +5 -17
- agno/models/nebius/nebius.py +6 -6
- agno/models/nexus/__init__.py +3 -0
- agno/models/nexus/nexus.py +22 -0
- agno/models/nvidia/nvidia.py +2 -2
- agno/models/ollama/chat.py +60 -6
- agno/models/openai/chat.py +102 -43
- agno/models/openai/responses.py +103 -106
- agno/models/openrouter/openrouter.py +41 -3
- agno/models/perplexity/perplexity.py +4 -5
- agno/models/portkey/portkey.py +3 -3
- agno/models/requesty/__init__.py +5 -0
- agno/models/requesty/requesty.py +52 -0
- agno/models/response.py +81 -5
- agno/models/sambanova/sambanova.py +2 -2
- agno/models/siliconflow/__init__.py +5 -0
- agno/models/siliconflow/siliconflow.py +25 -0
- agno/models/together/together.py +2 -2
- agno/models/utils.py +254 -8
- agno/models/vercel/v0.py +2 -2
- agno/models/vertexai/__init__.py +0 -0
- agno/models/vertexai/claude.py +96 -0
- agno/models/vllm/vllm.py +1 -0
- agno/models/xai/xai.py +3 -2
- agno/os/app.py +543 -175
- agno/os/auth.py +24 -14
- agno/os/config.py +1 -0
- agno/os/interfaces/__init__.py +1 -0
- agno/os/interfaces/a2a/__init__.py +3 -0
- agno/os/interfaces/a2a/a2a.py +42 -0
- agno/os/interfaces/a2a/router.py +250 -0
- agno/os/interfaces/a2a/utils.py +924 -0
- agno/os/interfaces/agui/agui.py +23 -7
- agno/os/interfaces/agui/router.py +27 -3
- agno/os/interfaces/agui/utils.py +242 -142
- agno/os/interfaces/base.py +6 -2
- agno/os/interfaces/slack/router.py +81 -23
- agno/os/interfaces/slack/slack.py +29 -14
- agno/os/interfaces/whatsapp/router.py +11 -4
- agno/os/interfaces/whatsapp/whatsapp.py +14 -7
- agno/os/mcp.py +111 -54
- agno/os/middleware/__init__.py +7 -0
- agno/os/middleware/jwt.py +233 -0
- agno/os/router.py +556 -139
- agno/os/routers/evals/evals.py +71 -34
- agno/os/routers/evals/schemas.py +31 -31
- agno/os/routers/evals/utils.py +6 -5
- agno/os/routers/health.py +31 -0
- agno/os/routers/home.py +52 -0
- agno/os/routers/knowledge/knowledge.py +185 -38
- agno/os/routers/knowledge/schemas.py +82 -22
- agno/os/routers/memory/memory.py +158 -53
- agno/os/routers/memory/schemas.py +20 -16
- agno/os/routers/metrics/metrics.py +20 -8
- agno/os/routers/metrics/schemas.py +16 -16
- agno/os/routers/session/session.py +499 -38
- agno/os/schema.py +308 -198
- agno/os/utils.py +401 -41
- agno/reasoning/anthropic.py +80 -0
- agno/reasoning/azure_ai_foundry.py +2 -2
- agno/reasoning/deepseek.py +2 -2
- agno/reasoning/default.py +3 -1
- agno/reasoning/gemini.py +73 -0
- agno/reasoning/groq.py +2 -2
- agno/reasoning/ollama.py +2 -2
- agno/reasoning/openai.py +7 -2
- agno/reasoning/vertexai.py +76 -0
- agno/run/__init__.py +6 -0
- agno/run/agent.py +266 -112
- agno/run/base.py +53 -24
- agno/run/team.py +252 -111
- agno/run/workflow.py +156 -45
- agno/session/agent.py +105 -89
- agno/session/summary.py +65 -25
- agno/session/team.py +176 -96
- agno/session/workflow.py +406 -40
- agno/team/team.py +3854 -1692
- agno/tools/brightdata.py +3 -3
- agno/tools/cartesia.py +3 -5
- agno/tools/dalle.py +9 -8
- agno/tools/decorator.py +4 -2
- agno/tools/desi_vocal.py +2 -2
- agno/tools/duckduckgo.py +15 -11
- agno/tools/e2b.py +20 -13
- agno/tools/eleven_labs.py +26 -28
- agno/tools/exa.py +21 -16
- agno/tools/fal.py +4 -4
- agno/tools/file.py +153 -23
- agno/tools/file_generation.py +350 -0
- agno/tools/firecrawl.py +4 -4
- agno/tools/function.py +257 -37
- agno/tools/giphy.py +2 -2
- agno/tools/gmail.py +238 -14
- agno/tools/google_drive.py +270 -0
- agno/tools/googlecalendar.py +36 -8
- agno/tools/googlesheets.py +20 -5
- agno/tools/jira.py +20 -0
- agno/tools/knowledge.py +3 -3
- agno/tools/lumalab.py +3 -3
- 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 +284 -0
- agno/tools/mem0.py +11 -17
- agno/tools/memori.py +1 -53
- agno/tools/memory.py +419 -0
- agno/tools/models/azure_openai.py +2 -2
- agno/tools/models/gemini.py +3 -3
- agno/tools/models/groq.py +3 -5
- agno/tools/models/nebius.py +7 -7
- agno/tools/models_labs.py +25 -15
- agno/tools/notion.py +204 -0
- agno/tools/openai.py +4 -9
- agno/tools/opencv.py +3 -3
- agno/tools/parallel.py +314 -0
- agno/tools/replicate.py +7 -7
- agno/tools/scrapegraph.py +58 -31
- agno/tools/searxng.py +2 -2
- agno/tools/serper.py +2 -2
- agno/tools/slack.py +18 -3
- agno/tools/spider.py +2 -2
- agno/tools/tavily.py +146 -0
- agno/tools/whatsapp.py +1 -1
- agno/tools/workflow.py +278 -0
- agno/tools/yfinance.py +12 -11
- agno/utils/agent.py +820 -0
- agno/utils/audio.py +27 -0
- agno/utils/common.py +90 -1
- agno/utils/events.py +222 -7
- agno/utils/gemini.py +181 -23
- agno/utils/hooks.py +57 -0
- agno/utils/http.py +111 -0
- agno/utils/knowledge.py +12 -5
- agno/utils/log.py +1 -0
- agno/utils/mcp.py +95 -5
- agno/utils/media.py +188 -10
- agno/utils/merge_dict.py +22 -1
- agno/utils/message.py +60 -0
- agno/utils/models/claude.py +40 -11
- agno/utils/models/cohere.py +1 -1
- agno/utils/models/watsonx.py +1 -1
- agno/utils/openai.py +1 -1
- agno/utils/print_response/agent.py +105 -21
- agno/utils/print_response/team.py +103 -38
- agno/utils/print_response/workflow.py +251 -34
- agno/utils/reasoning.py +22 -1
- agno/utils/serialize.py +32 -0
- agno/utils/streamlit.py +16 -10
- agno/utils/string.py +41 -0
- agno/utils/team.py +98 -9
- agno/utils/tools.py +1 -1
- agno/vectordb/base.py +23 -4
- agno/vectordb/cassandra/cassandra.py +65 -9
- agno/vectordb/chroma/chromadb.py +182 -38
- agno/vectordb/clickhouse/clickhousedb.py +64 -11
- agno/vectordb/couchbase/couchbase.py +105 -10
- agno/vectordb/lancedb/lance_db.py +183 -135
- agno/vectordb/langchaindb/langchaindb.py +25 -7
- agno/vectordb/lightrag/lightrag.py +17 -3
- agno/vectordb/llamaindex/__init__.py +3 -0
- agno/vectordb/llamaindex/llamaindexdb.py +46 -7
- agno/vectordb/milvus/milvus.py +126 -9
- agno/vectordb/mongodb/__init__.py +7 -1
- agno/vectordb/mongodb/mongodb.py +112 -7
- agno/vectordb/pgvector/pgvector.py +142 -21
- agno/vectordb/pineconedb/pineconedb.py +80 -8
- agno/vectordb/qdrant/qdrant.py +125 -39
- agno/vectordb/redis/__init__.py +9 -0
- agno/vectordb/redis/redisdb.py +694 -0
- agno/vectordb/singlestore/singlestore.py +111 -25
- agno/vectordb/surrealdb/surrealdb.py +31 -5
- agno/vectordb/upstashdb/upstashdb.py +76 -8
- agno/vectordb/weaviate/weaviate.py +86 -15
- agno/workflow/__init__.py +2 -0
- agno/workflow/agent.py +299 -0
- agno/workflow/condition.py +112 -18
- agno/workflow/loop.py +69 -10
- agno/workflow/parallel.py +266 -118
- agno/workflow/router.py +110 -17
- agno/workflow/step.py +645 -136
- agno/workflow/steps.py +65 -6
- agno/workflow/types.py +71 -33
- agno/workflow/workflow.py +2113 -300
- agno-2.3.0.dist-info/METADATA +618 -0
- agno-2.3.0.dist-info/RECORD +577 -0
- agno-2.3.0.dist-info/licenses/LICENSE +201 -0
- agno/knowledge/reader/url_reader.py +0 -128
- agno/tools/googlesearch.py +0 -98
- agno/tools/mcp.py +0 -610
- agno/utils/models/aws_claude.py +0 -170
- agno-2.0.0rc2.dist-info/METADATA +0 -355
- agno-2.0.0rc2.dist-info/RECORD +0 -515
- agno-2.0.0rc2.dist-info/licenses/LICENSE +0 -375
- {agno-2.0.0rc2.dist-info → agno-2.3.0.dist-info}/WHEEL +0 -0
- {agno-2.0.0rc2.dist-info → agno-2.3.0.dist-info}/top_level.txt +0 -0
agno/db/dynamo/dynamo.py
CHANGED
|
@@ -13,6 +13,7 @@ from agno.db.dynamo.utils import (
|
|
|
13
13
|
build_topic_filter_expression,
|
|
14
14
|
calculate_date_metrics,
|
|
15
15
|
create_table_if_not_exists,
|
|
16
|
+
deserialize_cultural_knowledge_from_db,
|
|
16
17
|
deserialize_eval_record,
|
|
17
18
|
deserialize_from_dynamodb_item,
|
|
18
19
|
deserialize_knowledge_row,
|
|
@@ -23,15 +24,18 @@ from agno.db.dynamo.utils import (
|
|
|
23
24
|
get_dates_to_calculate_metrics_for,
|
|
24
25
|
merge_with_existing_session,
|
|
25
26
|
prepare_session_data,
|
|
27
|
+
serialize_cultural_knowledge_for_db,
|
|
26
28
|
serialize_eval_record,
|
|
27
29
|
serialize_knowledge_row,
|
|
28
30
|
serialize_to_dynamo_item,
|
|
29
31
|
)
|
|
32
|
+
from agno.db.schemas.culture import CulturalKnowledge
|
|
30
33
|
from agno.db.schemas.evals import EvalFilterType, EvalRunRecord, EvalType
|
|
31
34
|
from agno.db.schemas.knowledge import KnowledgeRow
|
|
32
35
|
from agno.db.schemas.memory import UserMemory
|
|
33
36
|
from agno.session import AgentSession, Session, TeamSession, WorkflowSession
|
|
34
|
-
from agno.utils.log import log_debug, log_error
|
|
37
|
+
from agno.utils.log import log_debug, log_error, log_info
|
|
38
|
+
from agno.utils.string import generate_id
|
|
35
39
|
|
|
36
40
|
try:
|
|
37
41
|
import boto3 # type: ignore[import-untyped]
|
|
@@ -51,10 +55,12 @@ class DynamoDb(BaseDb):
|
|
|
51
55
|
aws_access_key_id: Optional[str] = None,
|
|
52
56
|
aws_secret_access_key: Optional[str] = None,
|
|
53
57
|
session_table: Optional[str] = None,
|
|
58
|
+
culture_table: Optional[str] = None,
|
|
54
59
|
memory_table: Optional[str] = None,
|
|
55
60
|
metrics_table: Optional[str] = None,
|
|
56
61
|
eval_table: Optional[str] = None,
|
|
57
62
|
knowledge_table: Optional[str] = None,
|
|
63
|
+
id: Optional[str] = None,
|
|
58
64
|
):
|
|
59
65
|
"""
|
|
60
66
|
Interface for interacting with a DynamoDB database.
|
|
@@ -65,13 +71,21 @@ class DynamoDb(BaseDb):
|
|
|
65
71
|
aws_access_key_id: AWS access key ID.
|
|
66
72
|
aws_secret_access_key: AWS secret access key.
|
|
67
73
|
session_table: The name of the session table.
|
|
74
|
+
culture_table: The name of the culture table.
|
|
68
75
|
memory_table: The name of the memory table.
|
|
69
76
|
metrics_table: The name of the metrics table.
|
|
70
77
|
eval_table: The name of the eval table.
|
|
71
78
|
knowledge_table: The name of the knowledge table.
|
|
79
|
+
id: ID of the database.
|
|
72
80
|
"""
|
|
81
|
+
if id is None:
|
|
82
|
+
seed = str(db_client) if db_client else f"{region_name}_{aws_access_key_id}"
|
|
83
|
+
id = generate_id(seed)
|
|
84
|
+
|
|
73
85
|
super().__init__(
|
|
86
|
+
id=id,
|
|
74
87
|
session_table=session_table,
|
|
88
|
+
culture_table=culture_table,
|
|
75
89
|
memory_table=memory_table,
|
|
76
90
|
metrics_table=metrics_table,
|
|
77
91
|
eval_table=eval_table,
|
|
@@ -98,27 +112,8 @@ class DynamoDb(BaseDb):
|
|
|
98
112
|
session = boto3.Session(**session_kwargs)
|
|
99
113
|
self.client = session.client("dynamodb")
|
|
100
114
|
|
|
101
|
-
def
|
|
102
|
-
|
|
103
|
-
(self.session_table_name, "sessions"),
|
|
104
|
-
(self.memory_table_name, "memories"),
|
|
105
|
-
(self.metrics_table_name, "metrics"),
|
|
106
|
-
(self.eval_table_name, "evals"),
|
|
107
|
-
(self.knowledge_table_name, "knowledge_sources"),
|
|
108
|
-
]
|
|
109
|
-
|
|
110
|
-
for table_name, table_type in tables_to_create:
|
|
111
|
-
if table_name:
|
|
112
|
-
try:
|
|
113
|
-
schema = get_table_schema_definition(table_type)
|
|
114
|
-
schema["TableName"] = table_name
|
|
115
|
-
create_table_if_not_exists(self.client, table_name, schema)
|
|
116
|
-
|
|
117
|
-
except Exception as e:
|
|
118
|
-
log_error(f"Failed to create table {table_name}: {e}")
|
|
119
|
-
|
|
120
|
-
def _table_exists(self, table_name: str) -> bool:
|
|
121
|
-
"""Check if a DynamoDB table with the given name exists.
|
|
115
|
+
def table_exists(self, table_name: str) -> bool:
|
|
116
|
+
"""Check if a DynamoDB table exists.
|
|
122
117
|
|
|
123
118
|
Args:
|
|
124
119
|
table_name: The name of the table to check
|
|
@@ -131,9 +126,23 @@ class DynamoDb(BaseDb):
|
|
|
131
126
|
return True
|
|
132
127
|
except self.client.exceptions.ResourceNotFoundException:
|
|
133
128
|
return False
|
|
134
|
-
|
|
135
|
-
|
|
136
|
-
|
|
129
|
+
|
|
130
|
+
def _create_all_tables(self):
|
|
131
|
+
"""Create all configured DynamoDB tables if they don't exist."""
|
|
132
|
+
tables_to_create = [
|
|
133
|
+
("sessions", self.session_table_name),
|
|
134
|
+
("memories", self.memory_table_name),
|
|
135
|
+
("metrics", self.metrics_table_name),
|
|
136
|
+
("evals", self.eval_table_name),
|
|
137
|
+
("knowledge", self.knowledge_table_name),
|
|
138
|
+
("culture", self.culture_table_name),
|
|
139
|
+
]
|
|
140
|
+
|
|
141
|
+
for table_type, table_name in tables_to_create:
|
|
142
|
+
if not self.table_exists(table_name):
|
|
143
|
+
schema = get_table_schema_definition(table_type)
|
|
144
|
+
schema["TableName"] = table_name
|
|
145
|
+
create_table_if_not_exists(self.client, table_name, schema)
|
|
137
146
|
|
|
138
147
|
def _get_table(self, table_type: str, create_table_if_not_found: Optional[bool] = True) -> Optional[str]:
|
|
139
148
|
"""
|
|
@@ -160,20 +169,30 @@ class DynamoDb(BaseDb):
|
|
|
160
169
|
table_name = self.eval_table_name
|
|
161
170
|
elif table_type == "knowledge":
|
|
162
171
|
table_name = self.knowledge_table_name
|
|
172
|
+
elif table_type == "culture":
|
|
173
|
+
table_name = self.culture_table_name
|
|
163
174
|
else:
|
|
164
175
|
raise ValueError(f"Unknown table type: {table_type}")
|
|
165
176
|
|
|
166
177
|
# Check if table exists, create if it doesn't
|
|
167
|
-
if not self.
|
|
178
|
+
if not self.table_exists(table_name) and create_table_if_not_found:
|
|
168
179
|
schema = get_table_schema_definition(table_type)
|
|
169
180
|
schema["TableName"] = table_name
|
|
170
181
|
create_table_if_not_exists(self.client, table_name, schema)
|
|
171
182
|
|
|
172
183
|
return table_name
|
|
173
184
|
|
|
185
|
+
def get_latest_schema_version(self):
|
|
186
|
+
"""Get the latest version of the database schema."""
|
|
187
|
+
pass
|
|
188
|
+
|
|
189
|
+
def upsert_schema_version(self, version: str) -> None:
|
|
190
|
+
"""Upsert the schema version into the database."""
|
|
191
|
+
pass
|
|
192
|
+
|
|
174
193
|
# --- Sessions ---
|
|
175
194
|
|
|
176
|
-
def delete_session(self, session_id: Optional[str] = None
|
|
195
|
+
def delete_session(self, session_id: Optional[str] = None) -> bool:
|
|
177
196
|
"""
|
|
178
197
|
Delete a session from the database.
|
|
179
198
|
|
|
@@ -224,11 +243,12 @@ class DynamoDb(BaseDb):
|
|
|
224
243
|
|
|
225
244
|
except Exception as e:
|
|
226
245
|
log_error(f"Failed to delete sessions: {e}")
|
|
246
|
+
raise e
|
|
227
247
|
|
|
228
248
|
def get_session(
|
|
229
249
|
self,
|
|
230
250
|
session_id: str,
|
|
231
|
-
session_type:
|
|
251
|
+
session_type: SessionType,
|
|
232
252
|
user_id: Optional[str] = None,
|
|
233
253
|
deserialize: Optional[bool] = True,
|
|
234
254
|
) -> Optional[Union[Session, Dict[str, Any]]]:
|
|
@@ -237,7 +257,7 @@ class DynamoDb(BaseDb):
|
|
|
237
257
|
|
|
238
258
|
Args:
|
|
239
259
|
session_id (str): The ID of the session to get.
|
|
240
|
-
session_type (
|
|
260
|
+
session_type (SessionType): The type of session to get.
|
|
241
261
|
user_id (Optional[str]): The ID of the user to get the session for.
|
|
242
262
|
deserialize (Optional[bool]): Whether to deserialize the session.
|
|
243
263
|
|
|
@@ -260,8 +280,6 @@ class DynamoDb(BaseDb):
|
|
|
260
280
|
|
|
261
281
|
session = deserialize_from_dynamodb_item(item)
|
|
262
282
|
|
|
263
|
-
if session_type and session.get("session_type") != session_type.value:
|
|
264
|
-
return None
|
|
265
283
|
if user_id and session.get("user_id") != user_id:
|
|
266
284
|
return None
|
|
267
285
|
|
|
@@ -275,12 +293,14 @@ class DynamoDb(BaseDb):
|
|
|
275
293
|
return AgentSession.from_dict(session)
|
|
276
294
|
elif session_type == SessionType.TEAM:
|
|
277
295
|
return TeamSession.from_dict(session)
|
|
278
|
-
|
|
296
|
+
elif session_type == SessionType.WORKFLOW:
|
|
279
297
|
return WorkflowSession.from_dict(session)
|
|
298
|
+
else:
|
|
299
|
+
raise ValueError(f"Invalid session type: {session_type}")
|
|
280
300
|
|
|
281
301
|
except Exception as e:
|
|
282
302
|
log_error(f"Failed to get session {session_id}: {e}")
|
|
283
|
-
|
|
303
|
+
raise e
|
|
284
304
|
|
|
285
305
|
def get_sessions(
|
|
286
306
|
self,
|
|
@@ -400,7 +420,7 @@ class DynamoDb(BaseDb):
|
|
|
400
420
|
|
|
401
421
|
except Exception as e:
|
|
402
422
|
log_error(f"Failed to get sessions: {e}")
|
|
403
|
-
|
|
423
|
+
raise e
|
|
404
424
|
|
|
405
425
|
def rename_session(
|
|
406
426
|
self,
|
|
@@ -468,7 +488,7 @@ class DynamoDb(BaseDb):
|
|
|
468
488
|
|
|
469
489
|
except Exception as e:
|
|
470
490
|
log_error(f"Failed to rename session {session_id}: {e}")
|
|
471
|
-
|
|
491
|
+
raise e
|
|
472
492
|
|
|
473
493
|
def upsert_session(
|
|
474
494
|
self, session: Session, deserialize: Optional[bool] = True
|
|
@@ -510,21 +530,72 @@ class DynamoDb(BaseDb):
|
|
|
510
530
|
|
|
511
531
|
except Exception as e:
|
|
512
532
|
log_error(f"Failed to upsert session {session.session_id}: {e}")
|
|
513
|
-
|
|
533
|
+
raise e
|
|
534
|
+
|
|
535
|
+
def upsert_sessions(
|
|
536
|
+
self, sessions: List[Session], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
|
|
537
|
+
) -> List[Union[Session, Dict[str, Any]]]:
|
|
538
|
+
"""
|
|
539
|
+
Bulk upsert multiple sessions for improved performance on large datasets.
|
|
540
|
+
|
|
541
|
+
Args:
|
|
542
|
+
sessions (List[Session]): List of sessions to upsert.
|
|
543
|
+
deserialize (Optional[bool]): Whether to deserialize the sessions. Defaults to True.
|
|
544
|
+
|
|
545
|
+
Returns:
|
|
546
|
+
List[Union[Session, Dict[str, Any]]]: List of upserted sessions.
|
|
547
|
+
|
|
548
|
+
Raises:
|
|
549
|
+
Exception: If an error occurs during bulk upsert.
|
|
550
|
+
"""
|
|
551
|
+
if not sessions:
|
|
552
|
+
return []
|
|
553
|
+
|
|
554
|
+
try:
|
|
555
|
+
log_info(
|
|
556
|
+
f"DynamoDb doesn't support efficient bulk operations, falling back to individual upserts for {len(sessions)} sessions"
|
|
557
|
+
)
|
|
558
|
+
|
|
559
|
+
# Fall back to individual upserts
|
|
560
|
+
results = []
|
|
561
|
+
for session in sessions:
|
|
562
|
+
if session is not None:
|
|
563
|
+
result = self.upsert_session(session, deserialize=deserialize)
|
|
564
|
+
if result is not None:
|
|
565
|
+
results.append(result)
|
|
566
|
+
return results
|
|
567
|
+
|
|
568
|
+
except Exception as e:
|
|
569
|
+
log_error(f"Exception during bulk session upsert: {e}")
|
|
570
|
+
return []
|
|
514
571
|
|
|
515
572
|
# --- User Memory ---
|
|
516
573
|
|
|
517
|
-
def delete_user_memory(self, memory_id: str) -> None:
|
|
574
|
+
def delete_user_memory(self, memory_id: str, user_id: Optional[str] = None) -> None:
|
|
518
575
|
"""
|
|
519
576
|
Delete a user memory from the database.
|
|
520
577
|
|
|
521
578
|
Args:
|
|
522
579
|
memory_id: The ID of the memory to delete.
|
|
580
|
+
user_id: The ID of the user (optional, for filtering).
|
|
523
581
|
|
|
524
582
|
Raises:
|
|
525
583
|
Exception: If any error occurs while deleting the user memory.
|
|
526
584
|
"""
|
|
527
585
|
try:
|
|
586
|
+
# If user_id is provided, verify the memory belongs to the user before deleting
|
|
587
|
+
if user_id:
|
|
588
|
+
response = self.client.get_item(
|
|
589
|
+
TableName=self.memory_table_name,
|
|
590
|
+
Key={"memory_id": {"S": memory_id}},
|
|
591
|
+
)
|
|
592
|
+
item = response.get("Item")
|
|
593
|
+
if item:
|
|
594
|
+
memory_data = deserialize_from_dynamodb_item(item)
|
|
595
|
+
if memory_data.get("user_id") != user_id:
|
|
596
|
+
log_debug(f"Memory {memory_id} does not belong to user {user_id}")
|
|
597
|
+
return
|
|
598
|
+
|
|
528
599
|
self.client.delete_item(
|
|
529
600
|
TableName=self.memory_table_name,
|
|
530
601
|
Key={"memory_id": {"S": memory_id}},
|
|
@@ -533,19 +604,36 @@ class DynamoDb(BaseDb):
|
|
|
533
604
|
|
|
534
605
|
except Exception as e:
|
|
535
606
|
log_error(f"Failed to delete user memory {memory_id}: {e}")
|
|
607
|
+
raise e
|
|
536
608
|
|
|
537
|
-
def delete_user_memories(self, memory_ids: List[str]) -> None:
|
|
609
|
+
def delete_user_memories(self, memory_ids: List[str], user_id: Optional[str] = None) -> None:
|
|
538
610
|
"""
|
|
539
611
|
Delete user memories from the database in batches.
|
|
540
612
|
|
|
541
613
|
Args:
|
|
542
614
|
memory_ids: List of memory IDs to delete
|
|
615
|
+
user_id: The ID of the user (optional, for filtering).
|
|
543
616
|
|
|
544
617
|
Raises:
|
|
545
618
|
Exception: If any error occurs while deleting the user memories.
|
|
546
619
|
"""
|
|
547
620
|
|
|
548
621
|
try:
|
|
622
|
+
# If user_id is provided, filter memory_ids to only those belonging to the user
|
|
623
|
+
if user_id:
|
|
624
|
+
filtered_memory_ids = []
|
|
625
|
+
for memory_id in memory_ids:
|
|
626
|
+
response = self.client.get_item(
|
|
627
|
+
TableName=self.memory_table_name,
|
|
628
|
+
Key={"memory_id": {"S": memory_id}},
|
|
629
|
+
)
|
|
630
|
+
item = response.get("Item")
|
|
631
|
+
if item:
|
|
632
|
+
memory_data = deserialize_from_dynamodb_item(item)
|
|
633
|
+
if memory_data.get("user_id") == user_id:
|
|
634
|
+
filtered_memory_ids.append(memory_id)
|
|
635
|
+
memory_ids = filtered_memory_ids
|
|
636
|
+
|
|
549
637
|
for i in range(0, len(memory_ids), DYNAMO_BATCH_SIZE_LIMIT):
|
|
550
638
|
batch = memory_ids[i : i + DYNAMO_BATCH_SIZE_LIMIT]
|
|
551
639
|
|
|
@@ -557,10 +645,14 @@ class DynamoDb(BaseDb):
|
|
|
557
645
|
|
|
558
646
|
except Exception as e:
|
|
559
647
|
log_error(f"Failed to delete user memories: {e}")
|
|
648
|
+
raise e
|
|
560
649
|
|
|
561
650
|
def get_all_memory_topics(self) -> List[str]:
|
|
562
651
|
"""Get all memory topics from the database.
|
|
563
652
|
|
|
653
|
+
Args:
|
|
654
|
+
user_id: The ID of the user (optional, for filtering).
|
|
655
|
+
|
|
564
656
|
Returns:
|
|
565
657
|
List[str]: List of unique memory topics.
|
|
566
658
|
"""
|
|
@@ -569,13 +661,17 @@ class DynamoDb(BaseDb):
|
|
|
569
661
|
if table_name is None:
|
|
570
662
|
return []
|
|
571
663
|
|
|
572
|
-
#
|
|
573
|
-
|
|
664
|
+
# Build filter expression for user_id if provided
|
|
665
|
+
scan_kwargs = {"TableName": table_name}
|
|
666
|
+
|
|
667
|
+
# Scan the table to get memories
|
|
668
|
+
response = self.client.scan(**scan_kwargs)
|
|
574
669
|
items = response.get("Items", [])
|
|
575
670
|
|
|
576
671
|
# Handle pagination
|
|
577
672
|
while "LastEvaluatedKey" in response:
|
|
578
|
-
|
|
673
|
+
scan_kwargs["ExclusiveStartKey"] = response["LastEvaluatedKey"]
|
|
674
|
+
response = self.client.scan(**scan_kwargs)
|
|
579
675
|
items.extend(response.get("Items", []))
|
|
580
676
|
|
|
581
677
|
# Extract topics from all memories
|
|
@@ -589,16 +685,21 @@ class DynamoDb(BaseDb):
|
|
|
589
685
|
|
|
590
686
|
except Exception as e:
|
|
591
687
|
log_error(f"Exception reading from memory table: {e}")
|
|
592
|
-
|
|
688
|
+
raise e
|
|
593
689
|
|
|
594
690
|
def get_user_memory(
|
|
595
|
-
self,
|
|
691
|
+
self,
|
|
692
|
+
memory_id: str,
|
|
693
|
+
deserialize: Optional[bool] = True,
|
|
694
|
+
user_id: Optional[str] = None,
|
|
596
695
|
) -> Optional[Union[UserMemory, Dict[str, Any]]]:
|
|
597
696
|
"""
|
|
598
697
|
Get a user memory from the database as a UserMemory object.
|
|
599
698
|
|
|
600
699
|
Args:
|
|
601
700
|
memory_id: The ID of the memory to get.
|
|
701
|
+
deserialize: Whether to deserialize the memory.
|
|
702
|
+
user_id: The ID of the user (optional, for filtering).
|
|
602
703
|
|
|
603
704
|
Returns:
|
|
604
705
|
Optional[UserMemory]: The user memory data if found, None otherwise.
|
|
@@ -615,6 +716,11 @@ class DynamoDb(BaseDb):
|
|
|
615
716
|
return None
|
|
616
717
|
|
|
617
718
|
item = deserialize_from_dynamodb_item(item)
|
|
719
|
+
|
|
720
|
+
# Filter by user_id if provided
|
|
721
|
+
if user_id and item.get("user_id") != user_id:
|
|
722
|
+
return None
|
|
723
|
+
|
|
618
724
|
if not deserialize:
|
|
619
725
|
return item
|
|
620
726
|
|
|
@@ -622,7 +728,7 @@ class DynamoDb(BaseDb):
|
|
|
622
728
|
|
|
623
729
|
except Exception as e:
|
|
624
730
|
log_error(f"Failed to get user memory {memory_id}: {e}")
|
|
625
|
-
|
|
731
|
+
raise e
|
|
626
732
|
|
|
627
733
|
def get_user_memories(
|
|
628
734
|
self,
|
|
@@ -742,7 +848,7 @@ class DynamoDb(BaseDb):
|
|
|
742
848
|
|
|
743
849
|
except Exception as e:
|
|
744
850
|
log_error(f"Failed to get user memories: {e}")
|
|
745
|
-
|
|
851
|
+
raise e
|
|
746
852
|
|
|
747
853
|
def get_user_memory_stats(
|
|
748
854
|
self,
|
|
@@ -754,6 +860,7 @@ class DynamoDb(BaseDb):
|
|
|
754
860
|
Args:
|
|
755
861
|
limit (Optional[int]): The maximum number of user stats to return.
|
|
756
862
|
page (Optional[int]): The page number.
|
|
863
|
+
user_id (Optional[str]): The ID of the user (optional, for filtering).
|
|
757
864
|
|
|
758
865
|
Returns:
|
|
759
866
|
Tuple[List[Dict[str, Any]], int]: A list of dictionaries containing user stats and total count.
|
|
@@ -773,29 +880,33 @@ class DynamoDb(BaseDb):
|
|
|
773
880
|
try:
|
|
774
881
|
table_name = self._get_table("memories")
|
|
775
882
|
|
|
776
|
-
|
|
883
|
+
# Build filter expression for user_id if provided
|
|
884
|
+
scan_kwargs = {"TableName": table_name}
|
|
885
|
+
|
|
886
|
+
response = self.client.scan(**scan_kwargs)
|
|
777
887
|
items = response.get("Items", [])
|
|
778
888
|
|
|
779
889
|
# Handle pagination
|
|
780
890
|
while "LastEvaluatedKey" in response:
|
|
781
|
-
|
|
891
|
+
scan_kwargs["ExclusiveStartKey"] = response["LastEvaluatedKey"]
|
|
892
|
+
response = self.client.scan(**scan_kwargs)
|
|
782
893
|
items.extend(response.get("Items", []))
|
|
783
894
|
|
|
784
895
|
# Aggregate stats by user_id
|
|
785
896
|
user_stats = {}
|
|
786
897
|
for item in items:
|
|
787
898
|
memory_data = deserialize_from_dynamodb_item(item)
|
|
788
|
-
|
|
899
|
+
current_user_id = memory_data.get("user_id")
|
|
789
900
|
|
|
790
|
-
if
|
|
791
|
-
if
|
|
792
|
-
user_stats[
|
|
793
|
-
"user_id":
|
|
901
|
+
if current_user_id:
|
|
902
|
+
if current_user_id not in user_stats:
|
|
903
|
+
user_stats[current_user_id] = {
|
|
904
|
+
"user_id": current_user_id,
|
|
794
905
|
"total_memories": 0,
|
|
795
906
|
"last_memory_updated_at": None,
|
|
796
907
|
}
|
|
797
908
|
|
|
798
|
-
user_stats[
|
|
909
|
+
user_stats[current_user_id]["total_memories"] += 1
|
|
799
910
|
|
|
800
911
|
updated_at = memory_data.get("updated_at")
|
|
801
912
|
if updated_at:
|
|
@@ -803,10 +914,10 @@ class DynamoDb(BaseDb):
|
|
|
803
914
|
updated_at_timestamp = int(updated_at_dt.timestamp())
|
|
804
915
|
|
|
805
916
|
if updated_at_timestamp and (
|
|
806
|
-
user_stats[
|
|
807
|
-
or updated_at_timestamp > user_stats[
|
|
917
|
+
user_stats[current_user_id]["last_memory_updated_at"] is None
|
|
918
|
+
or updated_at_timestamp > user_stats[current_user_id]["last_memory_updated_at"]
|
|
808
919
|
):
|
|
809
|
-
user_stats[
|
|
920
|
+
user_stats[current_user_id]["last_memory_updated_at"] = updated_at_timestamp
|
|
810
921
|
|
|
811
922
|
# Convert to list and apply sorting
|
|
812
923
|
stats_list = list(user_stats.values())
|
|
@@ -828,7 +939,7 @@ class DynamoDb(BaseDb):
|
|
|
828
939
|
|
|
829
940
|
except Exception as e:
|
|
830
941
|
log_error(f"Failed to get user memory stats: {e}")
|
|
831
|
-
|
|
942
|
+
raise e
|
|
832
943
|
|
|
833
944
|
def upsert_user_memory(
|
|
834
945
|
self, memory: UserMemory, deserialize: Optional[bool] = True
|
|
@@ -857,7 +968,44 @@ class DynamoDb(BaseDb):
|
|
|
857
968
|
|
|
858
969
|
except Exception as e:
|
|
859
970
|
log_error(f"Failed to upsert user memory: {e}")
|
|
860
|
-
|
|
971
|
+
raise e
|
|
972
|
+
|
|
973
|
+
def upsert_memories(
|
|
974
|
+
self, memories: List[UserMemory], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
|
|
975
|
+
) -> List[Union[UserMemory, Dict[str, Any]]]:
|
|
976
|
+
"""
|
|
977
|
+
Bulk upsert multiple user memories for improved performance on large datasets.
|
|
978
|
+
|
|
979
|
+
Args:
|
|
980
|
+
memories (List[UserMemory]): List of memories to upsert.
|
|
981
|
+
deserialize (Optional[bool]): Whether to deserialize the memories. Defaults to True.
|
|
982
|
+
|
|
983
|
+
Returns:
|
|
984
|
+
List[Union[UserMemory, Dict[str, Any]]]: List of upserted memories.
|
|
985
|
+
|
|
986
|
+
Raises:
|
|
987
|
+
Exception: If an error occurs during bulk upsert.
|
|
988
|
+
"""
|
|
989
|
+
if not memories:
|
|
990
|
+
return []
|
|
991
|
+
|
|
992
|
+
try:
|
|
993
|
+
log_info(
|
|
994
|
+
f"DynamoDb doesn't support efficient bulk operations, falling back to individual upserts for {len(memories)} memories"
|
|
995
|
+
)
|
|
996
|
+
|
|
997
|
+
# Fall back to individual upserts
|
|
998
|
+
results = []
|
|
999
|
+
for memory in memories:
|
|
1000
|
+
if memory is not None:
|
|
1001
|
+
result = self.upsert_user_memory(memory, deserialize=deserialize)
|
|
1002
|
+
if result is not None:
|
|
1003
|
+
results.append(result)
|
|
1004
|
+
return results
|
|
1005
|
+
|
|
1006
|
+
except Exception as e:
|
|
1007
|
+
log_error(f"Exception during bulk memory upsert: {e}")
|
|
1008
|
+
return []
|
|
861
1009
|
|
|
862
1010
|
def clear_memories(self) -> None:
|
|
863
1011
|
"""Delete all memories from the database.
|
|
@@ -898,6 +1046,7 @@ class DynamoDb(BaseDb):
|
|
|
898
1046
|
from agno.utils.log import log_warning
|
|
899
1047
|
|
|
900
1048
|
log_warning(f"Exception deleting all memories: {e}")
|
|
1049
|
+
raise e
|
|
901
1050
|
|
|
902
1051
|
# --- Metrics ---
|
|
903
1052
|
|
|
@@ -975,7 +1124,7 @@ class DynamoDb(BaseDb):
|
|
|
975
1124
|
|
|
976
1125
|
except Exception as e:
|
|
977
1126
|
log_error(f"Failed to calculate metrics: {e}")
|
|
978
|
-
|
|
1127
|
+
raise e
|
|
979
1128
|
|
|
980
1129
|
def _get_metrics_calculation_starting_date(self) -> Optional[date]:
|
|
981
1130
|
"""Get the first date for which metrics calculation is needed:
|
|
@@ -1061,7 +1210,7 @@ class DynamoDb(BaseDb):
|
|
|
1061
1210
|
|
|
1062
1211
|
except Exception as e:
|
|
1063
1212
|
log_error(f"Failed to get metrics calculation starting date: {e}")
|
|
1064
|
-
|
|
1213
|
+
raise e
|
|
1065
1214
|
|
|
1066
1215
|
def _get_all_sessions_for_metrics_calculation(
|
|
1067
1216
|
self, start_timestamp: int, end_timestamp: int
|
|
@@ -1119,7 +1268,7 @@ class DynamoDb(BaseDb):
|
|
|
1119
1268
|
|
|
1120
1269
|
except Exception as e:
|
|
1121
1270
|
log_error(f"Failed to get sessions for metrics calculation: {e}")
|
|
1122
|
-
|
|
1271
|
+
raise e
|
|
1123
1272
|
|
|
1124
1273
|
def _bulk_upsert_metrics(self, metrics_records: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
|
1125
1274
|
"""Bulk upsert metrics records into DynamoDB with proper deduplication.
|
|
@@ -1147,7 +1296,7 @@ class DynamoDb(BaseDb):
|
|
|
1147
1296
|
|
|
1148
1297
|
except Exception as e:
|
|
1149
1298
|
log_error(f"Failed to bulk upsert metrics: {e}")
|
|
1150
|
-
|
|
1299
|
+
raise e
|
|
1151
1300
|
|
|
1152
1301
|
def _upsert_single_metrics_record(self, table_name: str, record: Dict[str, Any]) -> Optional[Dict[str, Any]]:
|
|
1153
1302
|
"""Upsert a single metrics record, checking for existing records with the same date.
|
|
@@ -1179,7 +1328,7 @@ class DynamoDb(BaseDb):
|
|
|
1179
1328
|
|
|
1180
1329
|
except Exception as e:
|
|
1181
1330
|
log_error(f"Failed to upsert single metrics record: {e}")
|
|
1182
|
-
|
|
1331
|
+
raise e
|
|
1183
1332
|
|
|
1184
1333
|
def _get_existing_metrics_record(self, table_name: str, date_str: str) -> Optional[Dict[str, Any]]:
|
|
1185
1334
|
"""Get existing metrics record for a given date.
|
|
@@ -1212,7 +1361,7 @@ class DynamoDb(BaseDb):
|
|
|
1212
1361
|
|
|
1213
1362
|
except Exception as e:
|
|
1214
1363
|
log_error(f"Failed to get existing metrics record for date {date_str}: {e}")
|
|
1215
|
-
|
|
1364
|
+
raise e
|
|
1216
1365
|
|
|
1217
1366
|
def _update_existing_metrics_record(
|
|
1218
1367
|
self,
|
|
@@ -1246,7 +1395,7 @@ class DynamoDb(BaseDb):
|
|
|
1246
1395
|
|
|
1247
1396
|
except Exception as e:
|
|
1248
1397
|
log_error(f"Failed to update existing metrics record: {e}")
|
|
1249
|
-
|
|
1398
|
+
raise e
|
|
1250
1399
|
|
|
1251
1400
|
def _create_new_metrics_record(self, table_name: str, record: Dict[str, Any]) -> Optional[Dict[str, Any]]:
|
|
1252
1401
|
"""Create a new metrics record.
|
|
@@ -1270,7 +1419,7 @@ class DynamoDb(BaseDb):
|
|
|
1270
1419
|
|
|
1271
1420
|
except Exception as e:
|
|
1272
1421
|
log_error(f"Failed to create new metrics record: {e}")
|
|
1273
|
-
|
|
1422
|
+
raise e
|
|
1274
1423
|
|
|
1275
1424
|
def _prepare_metrics_record_for_dynamo(self, record: Dict[str, Any]) -> Dict[str, Any]:
|
|
1276
1425
|
"""Prepare a metrics record for DynamoDB serialization by converting all data types properly.
|
|
@@ -1314,17 +1463,17 @@ class DynamoDb(BaseDb):
|
|
|
1314
1463
|
"""
|
|
1315
1464
|
import json
|
|
1316
1465
|
|
|
1317
|
-
item = {}
|
|
1466
|
+
item: Dict[str, Any] = {}
|
|
1318
1467
|
for key, value in data.items():
|
|
1319
1468
|
if value is not None:
|
|
1320
1469
|
if isinstance(value, bool):
|
|
1321
|
-
item[key] = {"BOOL":
|
|
1470
|
+
item[key] = {"BOOL": value}
|
|
1322
1471
|
elif isinstance(value, (int, float)):
|
|
1323
1472
|
item[key] = {"N": str(value)}
|
|
1324
1473
|
elif isinstance(value, str):
|
|
1325
1474
|
item[key] = {"S": str(value)}
|
|
1326
1475
|
elif isinstance(value, (dict, list)):
|
|
1327
|
-
item[key] = {"S": json.dumps(
|
|
1476
|
+
item[key] = {"S": json.dumps(value)}
|
|
1328
1477
|
else:
|
|
1329
1478
|
item[key] = {"S": str(value)}
|
|
1330
1479
|
return item
|
|
@@ -1393,7 +1542,7 @@ class DynamoDb(BaseDb):
|
|
|
1393
1542
|
|
|
1394
1543
|
except Exception as e:
|
|
1395
1544
|
log_error(f"Failed to get metrics: {e}")
|
|
1396
|
-
|
|
1545
|
+
raise e
|
|
1397
1546
|
|
|
1398
1547
|
# --- Knowledge methods ---
|
|
1399
1548
|
|
|
@@ -1415,6 +1564,7 @@ class DynamoDb(BaseDb):
|
|
|
1415
1564
|
|
|
1416
1565
|
except Exception as e:
|
|
1417
1566
|
log_error(f"Failed to delete knowledge content {id}: {e}")
|
|
1567
|
+
raise e
|
|
1418
1568
|
|
|
1419
1569
|
def get_knowledge_content(self, id: str) -> Optional[KnowledgeRow]:
|
|
1420
1570
|
"""Get a knowledge row from the database.
|
|
@@ -1437,7 +1587,7 @@ class DynamoDb(BaseDb):
|
|
|
1437
1587
|
|
|
1438
1588
|
except Exception as e:
|
|
1439
1589
|
log_error(f"Failed to get knowledge content {id}: {e}")
|
|
1440
|
-
|
|
1590
|
+
raise e
|
|
1441
1591
|
|
|
1442
1592
|
def get_knowledge_contents(
|
|
1443
1593
|
self,
|
|
@@ -1509,7 +1659,7 @@ class DynamoDb(BaseDb):
|
|
|
1509
1659
|
|
|
1510
1660
|
except Exception as e:
|
|
1511
1661
|
log_error(f"Failed to get knowledge contents: {e}")
|
|
1512
|
-
|
|
1662
|
+
raise e
|
|
1513
1663
|
|
|
1514
1664
|
def upsert_knowledge_content(self, knowledge_row: KnowledgeRow):
|
|
1515
1665
|
"""Upsert knowledge content in the database.
|
|
@@ -1530,7 +1680,7 @@ class DynamoDb(BaseDb):
|
|
|
1530
1680
|
|
|
1531
1681
|
except Exception as e:
|
|
1532
1682
|
log_error(f"Failed to upsert knowledge content {knowledge_row.id}: {e}")
|
|
1533
|
-
|
|
1683
|
+
raise e
|
|
1534
1684
|
|
|
1535
1685
|
# --- Eval ---
|
|
1536
1686
|
|
|
@@ -1560,7 +1710,7 @@ class DynamoDb(BaseDb):
|
|
|
1560
1710
|
|
|
1561
1711
|
except Exception as e:
|
|
1562
1712
|
log_error(f"Failed to create eval run: {e}")
|
|
1563
|
-
|
|
1713
|
+
raise e
|
|
1564
1714
|
|
|
1565
1715
|
def delete_eval_runs(self, eval_run_ids: List[str]) -> None:
|
|
1566
1716
|
if not eval_run_ids or not self.eval_table_name:
|
|
@@ -1578,6 +1728,7 @@ class DynamoDb(BaseDb):
|
|
|
1578
1728
|
|
|
1579
1729
|
except Exception as e:
|
|
1580
1730
|
log_error(f"Failed to delete eval runs: {e}")
|
|
1731
|
+
raise e
|
|
1581
1732
|
|
|
1582
1733
|
def get_eval_run_raw(self, eval_run_id: str, table: Optional[Any] = None) -> Optional[Dict[str, Any]]:
|
|
1583
1734
|
if not self.eval_table_name:
|
|
@@ -1593,7 +1744,7 @@ class DynamoDb(BaseDb):
|
|
|
1593
1744
|
|
|
1594
1745
|
except Exception as e:
|
|
1595
1746
|
log_error(f"Failed to get eval run {eval_run_id}: {e}")
|
|
1596
|
-
|
|
1747
|
+
raise e
|
|
1597
1748
|
|
|
1598
1749
|
def get_eval_run(self, eval_run_id: str, table: Optional[Any] = None) -> Optional[EvalRunRecord]:
|
|
1599
1750
|
if not self.eval_table_name:
|
|
@@ -1609,7 +1760,7 @@ class DynamoDb(BaseDb):
|
|
|
1609
1760
|
|
|
1610
1761
|
except Exception as e:
|
|
1611
1762
|
log_error(f"Failed to get eval run {eval_run_id}: {e}")
|
|
1612
|
-
|
|
1763
|
+
raise e
|
|
1613
1764
|
|
|
1614
1765
|
def get_eval_runs(
|
|
1615
1766
|
self,
|
|
@@ -1661,14 +1812,16 @@ class DynamoDb(BaseDb):
|
|
|
1661
1812
|
|
|
1662
1813
|
if filter_type is not None:
|
|
1663
1814
|
if filter_type == EvalFilterType.AGENT:
|
|
1664
|
-
filter_expressions.append("agent_id
|
|
1815
|
+
filter_expressions.append("attribute_exists(agent_id)")
|
|
1665
1816
|
elif filter_type == EvalFilterType.TEAM:
|
|
1666
|
-
filter_expressions.append("team_id
|
|
1817
|
+
filter_expressions.append("attribute_exists(team_id)")
|
|
1667
1818
|
elif filter_type == EvalFilterType.WORKFLOW:
|
|
1668
|
-
filter_expressions.append("workflow_id
|
|
1819
|
+
filter_expressions.append("attribute_exists(workflow_id)")
|
|
1669
1820
|
|
|
1670
1821
|
if filter_expressions:
|
|
1671
1822
|
scan_kwargs["FilterExpression"] = " AND ".join(filter_expressions)
|
|
1823
|
+
|
|
1824
|
+
if expression_values:
|
|
1672
1825
|
scan_kwargs["ExpressionAttributeValues"] = expression_values # type: ignore
|
|
1673
1826
|
|
|
1674
1827
|
# Execute scan
|
|
@@ -1708,7 +1861,7 @@ class DynamoDb(BaseDb):
|
|
|
1708
1861
|
|
|
1709
1862
|
except Exception as e:
|
|
1710
1863
|
log_error(f"Failed to get eval runs: {e}")
|
|
1711
|
-
|
|
1864
|
+
raise e
|
|
1712
1865
|
|
|
1713
1866
|
def rename_eval_run(
|
|
1714
1867
|
self, eval_run_id: str, name: str, deserialize: Optional[bool] = True
|
|
@@ -1740,4 +1893,158 @@ class DynamoDb(BaseDb):
|
|
|
1740
1893
|
|
|
1741
1894
|
except Exception as e:
|
|
1742
1895
|
log_error(f"Failed to rename eval run {eval_run_id}: {e}")
|
|
1743
|
-
|
|
1896
|
+
raise e
|
|
1897
|
+
|
|
1898
|
+
# -- Culture methods --
|
|
1899
|
+
|
|
1900
|
+
def clear_cultural_knowledge(self) -> None:
|
|
1901
|
+
"""Delete all cultural knowledge from the database."""
|
|
1902
|
+
try:
|
|
1903
|
+
table_name = self._get_table("culture")
|
|
1904
|
+
response = self.client.scan(TableName=table_name, ProjectionExpression="id")
|
|
1905
|
+
|
|
1906
|
+
with self.client.batch_writer(table_name) as batch:
|
|
1907
|
+
for item in response.get("Items", []):
|
|
1908
|
+
batch.delete_item(Key={"id": item["id"]})
|
|
1909
|
+
except Exception as e:
|
|
1910
|
+
log_error(f"Failed to clear cultural knowledge: {e}")
|
|
1911
|
+
raise e
|
|
1912
|
+
|
|
1913
|
+
def delete_cultural_knowledge(self, id: str) -> None:
|
|
1914
|
+
"""Delete a cultural knowledge entry from the database."""
|
|
1915
|
+
try:
|
|
1916
|
+
table_name = self._get_table("culture")
|
|
1917
|
+
self.client.delete_item(TableName=table_name, Key={"id": {"S": id}})
|
|
1918
|
+
except Exception as e:
|
|
1919
|
+
log_error(f"Failed to delete cultural knowledge {id}: {e}")
|
|
1920
|
+
raise e
|
|
1921
|
+
|
|
1922
|
+
def get_cultural_knowledge(
|
|
1923
|
+
self, id: str, deserialize: Optional[bool] = True
|
|
1924
|
+
) -> Optional[Union[CulturalKnowledge, Dict[str, Any]]]:
|
|
1925
|
+
"""Get a cultural knowledge entry from the database."""
|
|
1926
|
+
try:
|
|
1927
|
+
table_name = self._get_table("culture")
|
|
1928
|
+
response = self.client.get_item(TableName=table_name, Key={"id": {"S": id}})
|
|
1929
|
+
|
|
1930
|
+
item = response.get("Item")
|
|
1931
|
+
if not item:
|
|
1932
|
+
return None
|
|
1933
|
+
|
|
1934
|
+
db_row = deserialize_from_dynamodb_item(item)
|
|
1935
|
+
if not deserialize:
|
|
1936
|
+
return db_row
|
|
1937
|
+
|
|
1938
|
+
return deserialize_cultural_knowledge_from_db(db_row)
|
|
1939
|
+
except Exception as e:
|
|
1940
|
+
log_error(f"Failed to get cultural knowledge {id}: {e}")
|
|
1941
|
+
raise e
|
|
1942
|
+
|
|
1943
|
+
def get_all_cultural_knowledge(
|
|
1944
|
+
self,
|
|
1945
|
+
name: Optional[str] = None,
|
|
1946
|
+
agent_id: Optional[str] = None,
|
|
1947
|
+
team_id: Optional[str] = None,
|
|
1948
|
+
limit: Optional[int] = None,
|
|
1949
|
+
page: Optional[int] = None,
|
|
1950
|
+
sort_by: Optional[str] = None,
|
|
1951
|
+
sort_order: Optional[str] = None,
|
|
1952
|
+
deserialize: Optional[bool] = True,
|
|
1953
|
+
) -> Union[List[CulturalKnowledge], Tuple[List[Dict[str, Any]], int]]:
|
|
1954
|
+
"""Get all cultural knowledge from the database."""
|
|
1955
|
+
try:
|
|
1956
|
+
table_name = self._get_table("culture")
|
|
1957
|
+
|
|
1958
|
+
# Build filter expression
|
|
1959
|
+
filter_expressions = []
|
|
1960
|
+
expression_values = {}
|
|
1961
|
+
|
|
1962
|
+
if name:
|
|
1963
|
+
filter_expressions.append("#name = :name")
|
|
1964
|
+
expression_values[":name"] = {"S": name}
|
|
1965
|
+
if agent_id:
|
|
1966
|
+
filter_expressions.append("agent_id = :agent_id")
|
|
1967
|
+
expression_values[":agent_id"] = {"S": agent_id}
|
|
1968
|
+
if team_id:
|
|
1969
|
+
filter_expressions.append("team_id = :team_id")
|
|
1970
|
+
expression_values[":team_id"] = {"S": team_id}
|
|
1971
|
+
|
|
1972
|
+
scan_kwargs: Dict[str, Any] = {"TableName": table_name}
|
|
1973
|
+
if filter_expressions:
|
|
1974
|
+
scan_kwargs["FilterExpression"] = " AND ".join(filter_expressions)
|
|
1975
|
+
scan_kwargs["ExpressionAttributeValues"] = expression_values
|
|
1976
|
+
if name:
|
|
1977
|
+
scan_kwargs["ExpressionAttributeNames"] = {"#name": "name"}
|
|
1978
|
+
|
|
1979
|
+
# Execute scan
|
|
1980
|
+
response = self.client.scan(**scan_kwargs)
|
|
1981
|
+
items = response.get("Items", [])
|
|
1982
|
+
|
|
1983
|
+
# Continue scanning if there's more data
|
|
1984
|
+
while "LastEvaluatedKey" in response:
|
|
1985
|
+
scan_kwargs["ExclusiveStartKey"] = response["LastEvaluatedKey"]
|
|
1986
|
+
response = self.client.scan(**scan_kwargs)
|
|
1987
|
+
items.extend(response.get("Items", []))
|
|
1988
|
+
|
|
1989
|
+
# Deserialize items from DynamoDB format
|
|
1990
|
+
db_rows = [deserialize_from_dynamodb_item(item) for item in items]
|
|
1991
|
+
|
|
1992
|
+
# Apply sorting
|
|
1993
|
+
if sort_by:
|
|
1994
|
+
reverse = sort_order == "desc" if sort_order else False
|
|
1995
|
+
db_rows.sort(key=lambda x: x.get(sort_by, ""), reverse=reverse)
|
|
1996
|
+
|
|
1997
|
+
# Apply pagination
|
|
1998
|
+
total_count = len(db_rows)
|
|
1999
|
+
if limit and page:
|
|
2000
|
+
start = (page - 1) * limit
|
|
2001
|
+
db_rows = db_rows[start : start + limit]
|
|
2002
|
+
elif limit:
|
|
2003
|
+
db_rows = db_rows[:limit]
|
|
2004
|
+
|
|
2005
|
+
if not deserialize:
|
|
2006
|
+
return db_rows, total_count
|
|
2007
|
+
|
|
2008
|
+
return [deserialize_cultural_knowledge_from_db(row) for row in db_rows]
|
|
2009
|
+
except Exception as e:
|
|
2010
|
+
log_error(f"Failed to get all cultural knowledge: {e}")
|
|
2011
|
+
raise e
|
|
2012
|
+
|
|
2013
|
+
def upsert_cultural_knowledge(
|
|
2014
|
+
self, cultural_knowledge: CulturalKnowledge, deserialize: Optional[bool] = True
|
|
2015
|
+
) -> Optional[Union[CulturalKnowledge, Dict[str, Any]]]:
|
|
2016
|
+
"""Upsert a cultural knowledge entry into the database."""
|
|
2017
|
+
try:
|
|
2018
|
+
from uuid import uuid4
|
|
2019
|
+
|
|
2020
|
+
table_name = self._get_table("culture", create_table_if_not_found=True)
|
|
2021
|
+
|
|
2022
|
+
if not cultural_knowledge.id:
|
|
2023
|
+
cultural_knowledge.id = str(uuid4())
|
|
2024
|
+
|
|
2025
|
+
# Serialize content, categories, and notes into a dict for DB storage
|
|
2026
|
+
content_dict = serialize_cultural_knowledge_for_db(cultural_knowledge)
|
|
2027
|
+
|
|
2028
|
+
# Create the item dict with serialized content
|
|
2029
|
+
item_dict = {
|
|
2030
|
+
"id": cultural_knowledge.id,
|
|
2031
|
+
"name": cultural_knowledge.name,
|
|
2032
|
+
"summary": cultural_knowledge.summary,
|
|
2033
|
+
"content": content_dict if content_dict else None,
|
|
2034
|
+
"metadata": cultural_knowledge.metadata,
|
|
2035
|
+
"input": cultural_knowledge.input,
|
|
2036
|
+
"created_at": cultural_knowledge.created_at,
|
|
2037
|
+
"updated_at": int(time.time()),
|
|
2038
|
+
"agent_id": cultural_knowledge.agent_id,
|
|
2039
|
+
"team_id": cultural_knowledge.team_id,
|
|
2040
|
+
}
|
|
2041
|
+
|
|
2042
|
+
# Convert to DynamoDB format
|
|
2043
|
+
item = serialize_to_dynamo_item(item_dict)
|
|
2044
|
+
self.client.put_item(TableName=table_name, Item=item)
|
|
2045
|
+
|
|
2046
|
+
return self.get_cultural_knowledge(cultural_knowledge.id, deserialize=deserialize)
|
|
2047
|
+
|
|
2048
|
+
except Exception as e:
|
|
2049
|
+
log_error(f"Failed to upsert cultural knowledge: {e}")
|
|
2050
|
+
raise e
|