agno 2.0.1__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.
Files changed (314) hide show
  1. agno/agent/agent.py +6015 -2823
  2. agno/api/api.py +2 -0
  3. agno/api/os.py +1 -1
  4. agno/culture/__init__.py +3 -0
  5. agno/culture/manager.py +956 -0
  6. agno/db/async_postgres/__init__.py +3 -0
  7. agno/db/base.py +385 -6
  8. agno/db/dynamo/dynamo.py +388 -81
  9. agno/db/dynamo/schemas.py +47 -10
  10. agno/db/dynamo/utils.py +63 -4
  11. agno/db/firestore/firestore.py +435 -64
  12. agno/db/firestore/schemas.py +11 -0
  13. agno/db/firestore/utils.py +102 -4
  14. agno/db/gcs_json/gcs_json_db.py +384 -42
  15. agno/db/gcs_json/utils.py +60 -26
  16. agno/db/in_memory/in_memory_db.py +351 -66
  17. agno/db/in_memory/utils.py +60 -2
  18. agno/db/json/json_db.py +339 -48
  19. agno/db/json/utils.py +60 -26
  20. agno/db/migrations/manager.py +199 -0
  21. agno/db/migrations/v1_to_v2.py +510 -37
  22. agno/db/migrations/versions/__init__.py +0 -0
  23. agno/db/migrations/versions/v2_3_0.py +938 -0
  24. agno/db/mongo/__init__.py +15 -1
  25. agno/db/mongo/async_mongo.py +2036 -0
  26. agno/db/mongo/mongo.py +653 -76
  27. agno/db/mongo/schemas.py +13 -0
  28. agno/db/mongo/utils.py +80 -8
  29. agno/db/mysql/mysql.py +687 -25
  30. agno/db/mysql/schemas.py +61 -37
  31. agno/db/mysql/utils.py +60 -2
  32. agno/db/postgres/__init__.py +2 -1
  33. agno/db/postgres/async_postgres.py +2001 -0
  34. agno/db/postgres/postgres.py +676 -57
  35. agno/db/postgres/schemas.py +43 -18
  36. agno/db/postgres/utils.py +164 -2
  37. agno/db/redis/redis.py +344 -38
  38. agno/db/redis/schemas.py +18 -0
  39. agno/db/redis/utils.py +60 -2
  40. agno/db/schemas/__init__.py +2 -1
  41. agno/db/schemas/culture.py +120 -0
  42. agno/db/schemas/memory.py +13 -0
  43. agno/db/singlestore/schemas.py +26 -1
  44. agno/db/singlestore/singlestore.py +687 -53
  45. agno/db/singlestore/utils.py +60 -2
  46. agno/db/sqlite/__init__.py +2 -1
  47. agno/db/sqlite/async_sqlite.py +2371 -0
  48. agno/db/sqlite/schemas.py +24 -0
  49. agno/db/sqlite/sqlite.py +774 -85
  50. agno/db/sqlite/utils.py +168 -5
  51. agno/db/surrealdb/__init__.py +3 -0
  52. agno/db/surrealdb/metrics.py +292 -0
  53. agno/db/surrealdb/models.py +309 -0
  54. agno/db/surrealdb/queries.py +71 -0
  55. agno/db/surrealdb/surrealdb.py +1361 -0
  56. agno/db/surrealdb/utils.py +147 -0
  57. agno/db/utils.py +50 -22
  58. agno/eval/accuracy.py +50 -43
  59. agno/eval/performance.py +6 -3
  60. agno/eval/reliability.py +6 -3
  61. agno/eval/utils.py +33 -16
  62. agno/exceptions.py +68 -1
  63. agno/filters.py +354 -0
  64. agno/guardrails/__init__.py +6 -0
  65. agno/guardrails/base.py +19 -0
  66. agno/guardrails/openai.py +144 -0
  67. agno/guardrails/pii.py +94 -0
  68. agno/guardrails/prompt_injection.py +52 -0
  69. agno/integrations/discord/client.py +1 -0
  70. agno/knowledge/chunking/agentic.py +13 -10
  71. agno/knowledge/chunking/fixed.py +1 -1
  72. agno/knowledge/chunking/semantic.py +40 -8
  73. agno/knowledge/chunking/strategy.py +59 -15
  74. agno/knowledge/embedder/aws_bedrock.py +9 -4
  75. agno/knowledge/embedder/azure_openai.py +54 -0
  76. agno/knowledge/embedder/base.py +2 -0
  77. agno/knowledge/embedder/cohere.py +184 -5
  78. agno/knowledge/embedder/fastembed.py +1 -1
  79. agno/knowledge/embedder/google.py +79 -1
  80. agno/knowledge/embedder/huggingface.py +9 -4
  81. agno/knowledge/embedder/jina.py +63 -0
  82. agno/knowledge/embedder/mistral.py +78 -11
  83. agno/knowledge/embedder/nebius.py +1 -1
  84. agno/knowledge/embedder/ollama.py +13 -0
  85. agno/knowledge/embedder/openai.py +37 -65
  86. agno/knowledge/embedder/sentence_transformer.py +8 -4
  87. agno/knowledge/embedder/vllm.py +262 -0
  88. agno/knowledge/embedder/voyageai.py +69 -16
  89. agno/knowledge/knowledge.py +594 -186
  90. agno/knowledge/reader/base.py +9 -2
  91. agno/knowledge/reader/csv_reader.py +8 -10
  92. agno/knowledge/reader/docx_reader.py +5 -6
  93. agno/knowledge/reader/field_labeled_csv_reader.py +290 -0
  94. agno/knowledge/reader/json_reader.py +6 -5
  95. agno/knowledge/reader/markdown_reader.py +13 -13
  96. agno/knowledge/reader/pdf_reader.py +43 -68
  97. agno/knowledge/reader/pptx_reader.py +101 -0
  98. agno/knowledge/reader/reader_factory.py +51 -6
  99. agno/knowledge/reader/s3_reader.py +3 -15
  100. agno/knowledge/reader/tavily_reader.py +194 -0
  101. agno/knowledge/reader/text_reader.py +13 -13
  102. agno/knowledge/reader/web_search_reader.py +2 -43
  103. agno/knowledge/reader/website_reader.py +43 -25
  104. agno/knowledge/reranker/__init__.py +2 -8
  105. agno/knowledge/types.py +9 -0
  106. agno/knowledge/utils.py +20 -0
  107. agno/media.py +72 -0
  108. agno/memory/manager.py +336 -82
  109. agno/models/aimlapi/aimlapi.py +2 -2
  110. agno/models/anthropic/claude.py +183 -37
  111. agno/models/aws/bedrock.py +52 -112
  112. agno/models/aws/claude.py +33 -1
  113. agno/models/azure/ai_foundry.py +33 -15
  114. agno/models/azure/openai_chat.py +25 -8
  115. agno/models/base.py +999 -519
  116. agno/models/cerebras/cerebras.py +19 -13
  117. agno/models/cerebras/cerebras_openai.py +8 -5
  118. agno/models/cohere/chat.py +27 -1
  119. agno/models/cometapi/__init__.py +5 -0
  120. agno/models/cometapi/cometapi.py +57 -0
  121. agno/models/dashscope/dashscope.py +1 -0
  122. agno/models/deepinfra/deepinfra.py +2 -2
  123. agno/models/deepseek/deepseek.py +2 -2
  124. agno/models/fireworks/fireworks.py +2 -2
  125. agno/models/google/gemini.py +103 -31
  126. agno/models/groq/groq.py +28 -11
  127. agno/models/huggingface/huggingface.py +2 -1
  128. agno/models/internlm/internlm.py +2 -2
  129. agno/models/langdb/langdb.py +4 -4
  130. agno/models/litellm/chat.py +18 -1
  131. agno/models/litellm/litellm_openai.py +2 -2
  132. agno/models/llama_cpp/__init__.py +5 -0
  133. agno/models/llama_cpp/llama_cpp.py +22 -0
  134. agno/models/message.py +139 -0
  135. agno/models/meta/llama.py +27 -10
  136. agno/models/meta/llama_openai.py +5 -17
  137. agno/models/nebius/nebius.py +6 -6
  138. agno/models/nexus/__init__.py +3 -0
  139. agno/models/nexus/nexus.py +22 -0
  140. agno/models/nvidia/nvidia.py +2 -2
  141. agno/models/ollama/chat.py +59 -5
  142. agno/models/openai/chat.py +69 -29
  143. agno/models/openai/responses.py +103 -106
  144. agno/models/openrouter/openrouter.py +41 -3
  145. agno/models/perplexity/perplexity.py +4 -5
  146. agno/models/portkey/portkey.py +3 -3
  147. agno/models/requesty/__init__.py +5 -0
  148. agno/models/requesty/requesty.py +52 -0
  149. agno/models/response.py +77 -1
  150. agno/models/sambanova/sambanova.py +2 -2
  151. agno/models/siliconflow/__init__.py +5 -0
  152. agno/models/siliconflow/siliconflow.py +25 -0
  153. agno/models/together/together.py +2 -2
  154. agno/models/utils.py +254 -8
  155. agno/models/vercel/v0.py +2 -2
  156. agno/models/vertexai/__init__.py +0 -0
  157. agno/models/vertexai/claude.py +96 -0
  158. agno/models/vllm/vllm.py +1 -0
  159. agno/models/xai/xai.py +3 -2
  160. agno/os/app.py +543 -178
  161. agno/os/auth.py +24 -14
  162. agno/os/config.py +1 -0
  163. agno/os/interfaces/__init__.py +1 -0
  164. agno/os/interfaces/a2a/__init__.py +3 -0
  165. agno/os/interfaces/a2a/a2a.py +42 -0
  166. agno/os/interfaces/a2a/router.py +250 -0
  167. agno/os/interfaces/a2a/utils.py +924 -0
  168. agno/os/interfaces/agui/agui.py +23 -7
  169. agno/os/interfaces/agui/router.py +27 -3
  170. agno/os/interfaces/agui/utils.py +242 -142
  171. agno/os/interfaces/base.py +6 -2
  172. agno/os/interfaces/slack/router.py +81 -23
  173. agno/os/interfaces/slack/slack.py +29 -14
  174. agno/os/interfaces/whatsapp/router.py +11 -4
  175. agno/os/interfaces/whatsapp/whatsapp.py +14 -7
  176. agno/os/mcp.py +111 -54
  177. agno/os/middleware/__init__.py +7 -0
  178. agno/os/middleware/jwt.py +233 -0
  179. agno/os/router.py +556 -139
  180. agno/os/routers/evals/evals.py +71 -34
  181. agno/os/routers/evals/schemas.py +31 -31
  182. agno/os/routers/evals/utils.py +6 -5
  183. agno/os/routers/health.py +31 -0
  184. agno/os/routers/home.py +52 -0
  185. agno/os/routers/knowledge/knowledge.py +185 -38
  186. agno/os/routers/knowledge/schemas.py +82 -22
  187. agno/os/routers/memory/memory.py +158 -53
  188. agno/os/routers/memory/schemas.py +20 -16
  189. agno/os/routers/metrics/metrics.py +20 -8
  190. agno/os/routers/metrics/schemas.py +16 -16
  191. agno/os/routers/session/session.py +499 -38
  192. agno/os/schema.py +308 -198
  193. agno/os/utils.py +401 -41
  194. agno/reasoning/anthropic.py +80 -0
  195. agno/reasoning/azure_ai_foundry.py +2 -2
  196. agno/reasoning/deepseek.py +2 -2
  197. agno/reasoning/default.py +3 -1
  198. agno/reasoning/gemini.py +73 -0
  199. agno/reasoning/groq.py +2 -2
  200. agno/reasoning/ollama.py +2 -2
  201. agno/reasoning/openai.py +7 -2
  202. agno/reasoning/vertexai.py +76 -0
  203. agno/run/__init__.py +6 -0
  204. agno/run/agent.py +248 -94
  205. agno/run/base.py +44 -5
  206. agno/run/team.py +238 -97
  207. agno/run/workflow.py +144 -33
  208. agno/session/agent.py +105 -89
  209. agno/session/summary.py +65 -25
  210. agno/session/team.py +176 -96
  211. agno/session/workflow.py +406 -40
  212. agno/team/team.py +3854 -1610
  213. agno/tools/dalle.py +2 -4
  214. agno/tools/decorator.py +4 -2
  215. agno/tools/duckduckgo.py +15 -11
  216. agno/tools/e2b.py +14 -7
  217. agno/tools/eleven_labs.py +23 -25
  218. agno/tools/exa.py +21 -16
  219. agno/tools/file.py +153 -23
  220. agno/tools/file_generation.py +350 -0
  221. agno/tools/firecrawl.py +4 -4
  222. agno/tools/function.py +250 -30
  223. agno/tools/gmail.py +238 -14
  224. agno/tools/google_drive.py +270 -0
  225. agno/tools/googlecalendar.py +36 -8
  226. agno/tools/googlesheets.py +20 -5
  227. agno/tools/jira.py +20 -0
  228. agno/tools/knowledge.py +3 -3
  229. agno/tools/mcp/__init__.py +10 -0
  230. agno/tools/mcp/mcp.py +331 -0
  231. agno/tools/mcp/multi_mcp.py +347 -0
  232. agno/tools/mcp/params.py +24 -0
  233. agno/tools/mcp_toolbox.py +284 -0
  234. agno/tools/mem0.py +11 -17
  235. agno/tools/memori.py +1 -53
  236. agno/tools/memory.py +419 -0
  237. agno/tools/models/nebius.py +5 -5
  238. agno/tools/models_labs.py +20 -10
  239. agno/tools/notion.py +204 -0
  240. agno/tools/parallel.py +314 -0
  241. agno/tools/scrapegraph.py +58 -31
  242. agno/tools/searxng.py +2 -2
  243. agno/tools/serper.py +2 -2
  244. agno/tools/slack.py +18 -3
  245. agno/tools/spider.py +2 -2
  246. agno/tools/tavily.py +146 -0
  247. agno/tools/whatsapp.py +1 -1
  248. agno/tools/workflow.py +278 -0
  249. agno/tools/yfinance.py +12 -11
  250. agno/utils/agent.py +820 -0
  251. agno/utils/audio.py +27 -0
  252. agno/utils/common.py +90 -1
  253. agno/utils/events.py +217 -2
  254. agno/utils/gemini.py +180 -22
  255. agno/utils/hooks.py +57 -0
  256. agno/utils/http.py +111 -0
  257. agno/utils/knowledge.py +12 -5
  258. agno/utils/log.py +1 -0
  259. agno/utils/mcp.py +92 -2
  260. agno/utils/media.py +188 -10
  261. agno/utils/merge_dict.py +22 -1
  262. agno/utils/message.py +60 -0
  263. agno/utils/models/claude.py +40 -11
  264. agno/utils/print_response/agent.py +105 -21
  265. agno/utils/print_response/team.py +103 -38
  266. agno/utils/print_response/workflow.py +251 -34
  267. agno/utils/reasoning.py +22 -1
  268. agno/utils/serialize.py +32 -0
  269. agno/utils/streamlit.py +16 -10
  270. agno/utils/string.py +41 -0
  271. agno/utils/team.py +98 -9
  272. agno/utils/tools.py +1 -1
  273. agno/vectordb/base.py +23 -4
  274. agno/vectordb/cassandra/cassandra.py +65 -9
  275. agno/vectordb/chroma/chromadb.py +182 -38
  276. agno/vectordb/clickhouse/clickhousedb.py +64 -11
  277. agno/vectordb/couchbase/couchbase.py +105 -10
  278. agno/vectordb/lancedb/lance_db.py +124 -133
  279. agno/vectordb/langchaindb/langchaindb.py +25 -7
  280. agno/vectordb/lightrag/lightrag.py +17 -3
  281. agno/vectordb/llamaindex/__init__.py +3 -0
  282. agno/vectordb/llamaindex/llamaindexdb.py +46 -7
  283. agno/vectordb/milvus/milvus.py +126 -9
  284. agno/vectordb/mongodb/__init__.py +7 -1
  285. agno/vectordb/mongodb/mongodb.py +112 -7
  286. agno/vectordb/pgvector/pgvector.py +142 -21
  287. agno/vectordb/pineconedb/pineconedb.py +80 -8
  288. agno/vectordb/qdrant/qdrant.py +125 -39
  289. agno/vectordb/redis/__init__.py +9 -0
  290. agno/vectordb/redis/redisdb.py +694 -0
  291. agno/vectordb/singlestore/singlestore.py +111 -25
  292. agno/vectordb/surrealdb/surrealdb.py +31 -5
  293. agno/vectordb/upstashdb/upstashdb.py +76 -8
  294. agno/vectordb/weaviate/weaviate.py +86 -15
  295. agno/workflow/__init__.py +2 -0
  296. agno/workflow/agent.py +299 -0
  297. agno/workflow/condition.py +112 -18
  298. agno/workflow/loop.py +69 -10
  299. agno/workflow/parallel.py +266 -118
  300. agno/workflow/router.py +110 -17
  301. agno/workflow/step.py +638 -129
  302. agno/workflow/steps.py +65 -6
  303. agno/workflow/types.py +61 -23
  304. agno/workflow/workflow.py +2085 -272
  305. {agno-2.0.1.dist-info → agno-2.3.0.dist-info}/METADATA +182 -58
  306. agno-2.3.0.dist-info/RECORD +577 -0
  307. agno/knowledge/reader/url_reader.py +0 -128
  308. agno/tools/googlesearch.py +0 -98
  309. agno/tools/mcp.py +0 -610
  310. agno/utils/models/aws_claude.py +0 -170
  311. agno-2.0.1.dist-info/RECORD +0 -515
  312. {agno-2.0.1.dist-info → agno-2.3.0.dist-info}/WHEEL +0 -0
  313. {agno-2.0.1.dist-info → agno-2.3.0.dist-info}/licenses/LICENSE +0 -0
  314. {agno-2.0.1.dist-info → agno-2.3.0.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,2371 @@
1
+ import time
2
+ from datetime import date, datetime, timedelta, timezone
3
+ from pathlib import Path
4
+ from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, cast
5
+ from uuid import uuid4
6
+
7
+ from agno.db.base import AsyncBaseDb, SessionType
8
+ from agno.db.migrations.manager import MigrationManager
9
+ from agno.db.schemas.culture import CulturalKnowledge
10
+ from agno.db.schemas.evals import EvalFilterType, EvalRunRecord, EvalType
11
+ from agno.db.schemas.knowledge import KnowledgeRow
12
+ from agno.db.schemas.memory import UserMemory
13
+ from agno.db.sqlite.schemas import get_table_schema_definition
14
+ from agno.db.sqlite.utils import (
15
+ abulk_upsert_metrics,
16
+ ais_table_available,
17
+ ais_valid_table,
18
+ apply_sorting,
19
+ calculate_date_metrics,
20
+ deserialize_cultural_knowledge_from_db,
21
+ fetch_all_sessions_data,
22
+ get_dates_to_calculate_metrics_for,
23
+ serialize_cultural_knowledge_for_db,
24
+ )
25
+ from agno.db.utils import deserialize_session_json_fields, serialize_session_json_fields
26
+ from agno.session import AgentSession, Session, TeamSession, WorkflowSession
27
+ from agno.utils.log import log_debug, log_error, log_info, log_warning
28
+ from agno.utils.string import generate_id
29
+
30
+ try:
31
+ from sqlalchemy import Column, MetaData, Table, and_, func, select, text
32
+ from sqlalchemy.dialects import sqlite
33
+ from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker, create_async_engine
34
+ from sqlalchemy.schema import Index, UniqueConstraint
35
+ except ImportError:
36
+ raise ImportError("`sqlalchemy` not installed. Please install it using `pip install sqlalchemy`")
37
+
38
+
39
+ class AsyncSqliteDb(AsyncBaseDb):
40
+ def __init__(
41
+ self,
42
+ db_file: Optional[str] = None,
43
+ db_engine: Optional[AsyncEngine] = None,
44
+ db_url: Optional[str] = None,
45
+ session_table: Optional[str] = None,
46
+ culture_table: Optional[str] = None,
47
+ memory_table: Optional[str] = None,
48
+ metrics_table: Optional[str] = None,
49
+ eval_table: Optional[str] = None,
50
+ knowledge_table: Optional[str] = None,
51
+ versions_table: Optional[str] = None,
52
+ id: Optional[str] = None,
53
+ ):
54
+ """
55
+ Async interface for interacting with a SQLite database.
56
+
57
+ The following order is used to determine the database connection:
58
+ 1. Use the db_engine
59
+ 2. Use the db_url
60
+ 3. Use the db_file
61
+ 4. Create a new database in the current directory
62
+
63
+ Args:
64
+ db_file (Optional[str]): The database file to connect to.
65
+ db_engine (Optional[AsyncEngine]): The SQLAlchemy async database engine to use.
66
+ db_url (Optional[str]): The database URL to connect to.
67
+ session_table (Optional[str]): Name of the table to store Agent, Team and Workflow sessions.
68
+ culture_table (Optional[str]): Name of the table to store cultural notions.
69
+ memory_table (Optional[str]): Name of the table to store user memories.
70
+ metrics_table (Optional[str]): Name of the table to store metrics.
71
+ eval_table (Optional[str]): Name of the table to store evaluation runs data.
72
+ knowledge_table (Optional[str]): Name of the table to store knowledge documents data.
73
+ versions_table (Optional[str]): Name of the table to store schema versions.
74
+ id (Optional[str]): ID of the database.
75
+
76
+ Raises:
77
+ ValueError: If none of the tables are provided.
78
+ """
79
+ if id is None:
80
+ seed = db_url or db_file or str(db_engine.url) if db_engine else "sqlite+aiosqlite:///agno.db"
81
+ id = generate_id(seed)
82
+
83
+ super().__init__(
84
+ id=id,
85
+ session_table=session_table,
86
+ culture_table=culture_table,
87
+ memory_table=memory_table,
88
+ metrics_table=metrics_table,
89
+ eval_table=eval_table,
90
+ knowledge_table=knowledge_table,
91
+ versions_table=versions_table,
92
+ )
93
+
94
+ _engine: Optional[AsyncEngine] = db_engine
95
+ if _engine is None:
96
+ if db_url is not None:
97
+ _engine = create_async_engine(db_url)
98
+ elif db_file is not None:
99
+ db_path = Path(db_file).resolve()
100
+ db_path.parent.mkdir(parents=True, exist_ok=True)
101
+ db_file = str(db_path)
102
+ _engine = create_async_engine(f"sqlite+aiosqlite:///{db_path}")
103
+ else:
104
+ # If none of db_engine, db_url, or db_file are provided, create a db in the current directory
105
+ default_db_path = Path("./agno.db").resolve()
106
+ _engine = create_async_engine(f"sqlite+aiosqlite:///{default_db_path}")
107
+ db_file = str(default_db_path)
108
+ log_debug(f"Created SQLite database: {default_db_path}")
109
+
110
+ self.db_engine: AsyncEngine = _engine
111
+ self.db_url: Optional[str] = db_url
112
+ self.db_file: Optional[str] = db_file
113
+ self.metadata: MetaData = MetaData()
114
+
115
+ # Initialize database session factory
116
+ self.async_session_factory = async_sessionmaker(bind=self.db_engine, expire_on_commit=False)
117
+
118
+ # -- DB methods --
119
+ async def table_exists(self, table_name: str) -> bool:
120
+ """Check if a table with the given name exists in the SQLite database.
121
+
122
+ Args:
123
+ table_name: Name of the table to check
124
+
125
+ Returns:
126
+ bool: True if the table exists in the database, False otherwise
127
+ """
128
+ async with self.async_session_factory() as sess:
129
+ return await ais_table_available(session=sess, table_name=table_name)
130
+
131
+ async def _create_all_tables(self):
132
+ """Create all tables for the database."""
133
+ tables_to_create = [
134
+ (self.session_table_name, "sessions"),
135
+ (self.memory_table_name, "memories"),
136
+ (self.metrics_table_name, "metrics"),
137
+ (self.eval_table_name, "evals"),
138
+ (self.knowledge_table_name, "knowledge"),
139
+ (self.versions_table_name, "versions"),
140
+ ]
141
+
142
+ for table_name, table_type in tables_to_create:
143
+ if table_name != self.versions_table_name:
144
+ # Also store the schema version for the created table
145
+ latest_schema_version = MigrationManager(self).latest_schema_version
146
+ await self.upsert_schema_version(table_name=table_name, version=latest_schema_version.public)
147
+
148
+ await self._create_table(table_name=table_name, table_type=table_type)
149
+
150
+ async def _create_table(self, table_name: str, table_type: str) -> Table:
151
+ """
152
+ Create a table with the appropriate schema based on the table type.
153
+
154
+ Args:
155
+ table_name (str): Name of the table to create
156
+ table_type (str): Type of table (used to get schema definition)
157
+
158
+ Returns:
159
+ Table: SQLAlchemy Table object
160
+ """
161
+ try:
162
+ table_schema = get_table_schema_definition(table_type)
163
+ log_debug(f"Creating table {table_name}")
164
+
165
+ columns: List[Column] = []
166
+ indexes: List[str] = []
167
+ unique_constraints: List[str] = []
168
+ schema_unique_constraints = table_schema.pop("_unique_constraints", [])
169
+
170
+ # Get the columns, indexes, and unique constraints from the table schema
171
+ for col_name, col_config in table_schema.items():
172
+ column_args = [col_name, col_config["type"]()]
173
+ column_kwargs = {}
174
+
175
+ if col_config.get("primary_key", False):
176
+ column_kwargs["primary_key"] = True
177
+ if "nullable" in col_config:
178
+ column_kwargs["nullable"] = col_config["nullable"]
179
+ if col_config.get("index", False):
180
+ indexes.append(col_name)
181
+ if col_config.get("unique", False):
182
+ column_kwargs["unique"] = True
183
+ unique_constraints.append(col_name)
184
+
185
+ columns.append(Column(*column_args, **column_kwargs)) # type: ignore
186
+
187
+ # Create the table object
188
+ table_metadata = MetaData()
189
+ table = Table(table_name, table_metadata, *columns)
190
+
191
+ # Add multi-column unique constraints with table-specific names
192
+ for constraint in schema_unique_constraints:
193
+ constraint_name = f"{table_name}_{constraint['name']}"
194
+ constraint_columns = constraint["columns"]
195
+ table.append_constraint(UniqueConstraint(*constraint_columns, name=constraint_name))
196
+
197
+ # Add indexes to the table definition
198
+ for idx_col in indexes:
199
+ idx_name = f"idx_{table_name}_{idx_col}"
200
+ table.append_constraint(Index(idx_name, idx_col))
201
+
202
+ # Create table
203
+ async with self.db_engine.begin() as conn:
204
+ await conn.run_sync(table.create, checkfirst=True)
205
+
206
+ # Create indexes
207
+ for idx in table.indexes:
208
+ try:
209
+ log_debug(f"Creating index: {idx.name}")
210
+ # Check if index already exists
211
+ async with self.async_session_factory() as sess:
212
+ exists_query = text("SELECT 1 FROM sqlite_master WHERE type = 'index' AND name = :index_name")
213
+ result = await sess.execute(exists_query, {"index_name": idx.name})
214
+ exists = result.scalar() is not None
215
+ if exists:
216
+ log_debug(f"Index {idx.name} already exists in table {table_name}, skipping creation")
217
+ continue
218
+
219
+ async with self.db_engine.begin() as conn:
220
+ await conn.run_sync(idx.create)
221
+
222
+ except Exception as e:
223
+ log_warning(f"Error creating index {idx.name}: {e}")
224
+
225
+ log_debug(f"Successfully created table '{table_name}'")
226
+ return table
227
+
228
+ except Exception as e:
229
+ log_error(f"Could not create table '{table_name}': {e}")
230
+ raise e
231
+
232
+ async def _get_table(self, table_type: str) -> Optional[Table]:
233
+ if table_type == "sessions":
234
+ if not hasattr(self, "session_table"):
235
+ self.session_table = await self._get_or_create_table(
236
+ table_name=self.session_table_name,
237
+ table_type=table_type,
238
+ )
239
+ return self.session_table
240
+
241
+ elif table_type == "memories":
242
+ if not hasattr(self, "memory_table"):
243
+ self.memory_table = await self._get_or_create_table(
244
+ table_name=self.memory_table_name,
245
+ table_type="memories",
246
+ )
247
+ return self.memory_table
248
+
249
+ elif table_type == "metrics":
250
+ if not hasattr(self, "metrics_table"):
251
+ self.metrics_table = await self._get_or_create_table(
252
+ table_name=self.metrics_table_name,
253
+ table_type="metrics",
254
+ )
255
+ return self.metrics_table
256
+
257
+ elif table_type == "evals":
258
+ if not hasattr(self, "eval_table"):
259
+ self.eval_table = await self._get_or_create_table(
260
+ table_name=self.eval_table_name,
261
+ table_type="evals",
262
+ )
263
+ return self.eval_table
264
+
265
+ elif table_type == "knowledge":
266
+ if not hasattr(self, "knowledge_table"):
267
+ self.knowledge_table = await self._get_or_create_table(
268
+ table_name=self.knowledge_table_name,
269
+ table_type="knowledge",
270
+ )
271
+ return self.knowledge_table
272
+
273
+ elif table_type == "culture":
274
+ if not hasattr(self, "culture_table"):
275
+ self.culture_table = await self._get_or_create_table(
276
+ table_name=self.culture_table_name,
277
+ table_type="culture",
278
+ )
279
+ return self.culture_table
280
+
281
+ elif table_type == "versions":
282
+ if not hasattr(self, "versions_table"):
283
+ self.versions_table = await self._get_or_create_table(
284
+ table_name=self.versions_table_name,
285
+ table_type="versions",
286
+ )
287
+ return self.versions_table
288
+
289
+ else:
290
+ raise ValueError(f"Unknown table type: '{table_type}'")
291
+
292
+ async def _get_or_create_table(
293
+ self,
294
+ table_name: str,
295
+ table_type: str,
296
+ ) -> Table:
297
+ """
298
+ Check if the table exists and is valid, else create it.
299
+
300
+ Args:
301
+ table_name (str): Name of the table to get or create
302
+ table_type (str): Type of table (used to get schema definition)
303
+
304
+ Returns:
305
+ Table: SQLAlchemy Table object
306
+ """
307
+ async with self.async_session_factory() as sess, sess.begin():
308
+ table_is_available = await ais_table_available(session=sess, table_name=table_name)
309
+
310
+ if not table_is_available:
311
+ if table_name != self.versions_table_name:
312
+ # Also store the schema version for the created table
313
+ latest_schema_version = MigrationManager(self).latest_schema_version
314
+ await self.upsert_schema_version(table_name=table_name, version=latest_schema_version.public)
315
+
316
+ return await self._create_table(table_name=table_name, table_type=table_type)
317
+
318
+ # SQLite version of table validation (no schema)
319
+ if not await ais_valid_table(db_engine=self.db_engine, table_name=table_name, table_type=table_type):
320
+ raise ValueError(f"Table {table_name} has an invalid schema")
321
+
322
+ try:
323
+ async with self.db_engine.connect() as conn:
324
+
325
+ def load_table(connection):
326
+ return Table(table_name, self.metadata, autoload_with=connection)
327
+
328
+ table = await conn.run_sync(load_table)
329
+ log_debug(f"Loaded existing table {table_name}")
330
+ return table
331
+
332
+ except Exception as e:
333
+ log_error(f"Error loading existing table {table_name}: {e}")
334
+ raise e
335
+
336
+ async def get_latest_schema_version(self, table_name: str) -> str:
337
+ """Get the latest version of the database schema."""
338
+ table = await self._get_table(table_type="versions")
339
+ if table is None:
340
+ return "2.0.0"
341
+ async with self.async_session_factory() as sess:
342
+ stmt = select(table)
343
+ # Latest version for the given table
344
+ stmt = stmt.where(table.c.table_name == table_name)
345
+ stmt = stmt.order_by(table.c.version.desc()).limit(1)
346
+ result = await sess.execute(stmt)
347
+ row = result.fetchone()
348
+ if row is None:
349
+ return "2.0.0"
350
+ version_dict = dict(row._mapping)
351
+ return version_dict.get("version") or "2.0.0"
352
+
353
+ async def upsert_schema_version(self, table_name: str, version: str) -> None:
354
+ """Upsert the schema version into the database."""
355
+ table = await self._get_table(table_type="versions")
356
+ if table is None:
357
+ return
358
+ current_datetime = datetime.now().isoformat()
359
+ async with self.async_session_factory() as sess, sess.begin():
360
+ stmt = sqlite.insert(table).values(
361
+ table_name=table_name,
362
+ version=version,
363
+ created_at=current_datetime, # Store as ISO format string
364
+ updated_at=current_datetime,
365
+ )
366
+ # Update version if table_name already exists
367
+ stmt = stmt.on_conflict_do_update(
368
+ index_elements=["table_name"],
369
+ set_=dict(version=version, updated_at=current_datetime),
370
+ )
371
+ await sess.execute(stmt)
372
+
373
+ # -- Session methods --
374
+
375
+ async def delete_session(self, session_id: str) -> bool:
376
+ """
377
+ Delete a session from the database.
378
+
379
+ Args:
380
+ session_id (str): ID of the session to delete
381
+
382
+ Returns:
383
+ bool: True if the session was deleted, False otherwise.
384
+
385
+ Raises:
386
+ Exception: If an error occurs during deletion.
387
+ """
388
+ try:
389
+ table = await self._get_table(table_type="sessions")
390
+ if table is None:
391
+ return False
392
+
393
+ async with self.async_session_factory() as sess, sess.begin():
394
+ delete_stmt = table.delete().where(table.c.session_id == session_id)
395
+ result = await sess.execute(delete_stmt)
396
+ if result.rowcount == 0: # type: ignore
397
+ log_debug(f"No session found to delete with session_id: {session_id}")
398
+ return False
399
+ else:
400
+ log_debug(f"Successfully deleted session with session_id: {session_id}")
401
+ return True
402
+
403
+ except Exception as e:
404
+ log_error(f"Error deleting session: {e}")
405
+ return False
406
+
407
+ async def delete_sessions(self, session_ids: List[str]) -> None:
408
+ """Delete all given sessions from the database.
409
+ Can handle multiple session types in the same run.
410
+
411
+ Args:
412
+ session_ids (List[str]): The IDs of the sessions to delete.
413
+
414
+ Raises:
415
+ Exception: If an error occurs during deletion.
416
+ """
417
+ try:
418
+ table = await self._get_table(table_type="sessions")
419
+ if table is None:
420
+ return
421
+
422
+ async with self.async_session_factory() as sess, sess.begin():
423
+ delete_stmt = table.delete().where(table.c.session_id.in_(session_ids))
424
+ result = await sess.execute(delete_stmt)
425
+
426
+ log_debug(f"Successfully deleted {result.rowcount} sessions") # type: ignore
427
+
428
+ except Exception as e:
429
+ log_error(f"Error deleting sessions: {e}")
430
+
431
+ async def get_session(
432
+ self,
433
+ session_id: str,
434
+ session_type: SessionType,
435
+ user_id: Optional[str] = None,
436
+ deserialize: Optional[bool] = True,
437
+ ) -> Optional[Union[Session, Dict[str, Any]]]:
438
+ """
439
+ Read a session from the database.
440
+
441
+ Args:
442
+ session_id (str): ID of the session to read.
443
+ session_type (SessionType): Type of session to get.
444
+ user_id (Optional[str]): User ID to filter by. Defaults to None.
445
+ deserialize (Optional[bool]): Whether to serialize the session. Defaults to True.
446
+
447
+ Returns:
448
+ Optional[Union[Session, Dict[str, Any]]]:
449
+ - When deserialize=True: Session object
450
+ - When deserialize=False: Session dictionary
451
+
452
+ Raises:
453
+ Exception: If an error occurs during retrieval.
454
+ """
455
+ try:
456
+ table = await self._get_table(table_type="sessions")
457
+ if table is None:
458
+ return None
459
+
460
+ async with self.async_session_factory() as sess, sess.begin():
461
+ stmt = select(table).where(table.c.session_id == session_id)
462
+
463
+ # Filtering
464
+ if user_id is not None:
465
+ stmt = stmt.where(table.c.user_id == user_id)
466
+
467
+ result = await sess.execute(stmt)
468
+ row = result.fetchone()
469
+ if row is None:
470
+ return None
471
+
472
+ session_raw = deserialize_session_json_fields(dict(row._mapping))
473
+ if not session_raw or not deserialize:
474
+ return session_raw
475
+
476
+ if session_type == SessionType.AGENT:
477
+ return AgentSession.from_dict(session_raw)
478
+ elif session_type == SessionType.TEAM:
479
+ return TeamSession.from_dict(session_raw)
480
+ elif session_type == SessionType.WORKFLOW:
481
+ return WorkflowSession.from_dict(session_raw)
482
+ else:
483
+ raise ValueError(f"Invalid session type: {session_type}")
484
+
485
+ except Exception as e:
486
+ log_debug(f"Exception reading from sessions table: {e}")
487
+ raise e
488
+
489
+ async def get_sessions(
490
+ self,
491
+ session_type: Optional[SessionType] = None,
492
+ user_id: Optional[str] = None,
493
+ component_id: Optional[str] = None,
494
+ session_name: Optional[str] = None,
495
+ start_timestamp: Optional[int] = None,
496
+ end_timestamp: Optional[int] = None,
497
+ limit: Optional[int] = None,
498
+ page: Optional[int] = None,
499
+ sort_by: Optional[str] = None,
500
+ sort_order: Optional[str] = None,
501
+ deserialize: Optional[bool] = True,
502
+ ) -> Union[List[Session], Tuple[List[Dict[str, Any]], int]]:
503
+ """
504
+ Get all sessions in the given table. Can filter by user_id and entity_id.
505
+ Args:
506
+ session_type (Optional[SessionType]): The type of session to get.
507
+ user_id (Optional[str]): The ID of the user to filter by.
508
+ component_id (Optional[str]): The ID of the agent / workflow to filter by.
509
+ session_name (Optional[str]): The name of the session to filter by.
510
+ start_timestamp (Optional[int]): The start timestamp to filter by.
511
+ end_timestamp (Optional[int]): The end timestamp to filter by.
512
+ limit (Optional[int]): The maximum number of sessions to return. Defaults to None.
513
+ page (Optional[int]): The page number to return. Defaults to None.
514
+ sort_by (Optional[str]): The field to sort by. Defaults to None.
515
+ sort_order (Optional[str]): The sort order. Defaults to None.
516
+ deserialize (Optional[bool]): Whether to serialize the sessions. Defaults to True.
517
+
518
+ Returns:
519
+ List[Session]:
520
+ - When deserialize=True: List of Session objects matching the criteria.
521
+ - When deserialize=False: List of Session dictionaries matching the criteria.
522
+
523
+ Raises:
524
+ Exception: If an error occurs during retrieval.
525
+ """
526
+ try:
527
+ table = await self._get_table(table_type="sessions")
528
+ if table is None:
529
+ return [] if deserialize else ([], 0)
530
+
531
+ async with self.async_session_factory() as sess, sess.begin():
532
+ stmt = select(table)
533
+
534
+ # Filtering
535
+ if user_id is not None:
536
+ stmt = stmt.where(table.c.user_id == user_id)
537
+ if component_id is not None:
538
+ if session_type == SessionType.AGENT:
539
+ stmt = stmt.where(table.c.agent_id == component_id)
540
+ elif session_type == SessionType.TEAM:
541
+ stmt = stmt.where(table.c.team_id == component_id)
542
+ elif session_type == SessionType.WORKFLOW:
543
+ stmt = stmt.where(table.c.workflow_id == component_id)
544
+ if start_timestamp is not None:
545
+ stmt = stmt.where(table.c.created_at >= start_timestamp)
546
+ if end_timestamp is not None:
547
+ stmt = stmt.where(table.c.created_at <= end_timestamp)
548
+ if session_name is not None:
549
+ stmt = stmt.where(table.c.session_data.like(f"%{session_name}%"))
550
+ if session_type is not None:
551
+ stmt = stmt.where(table.c.session_type == session_type.value)
552
+
553
+ # Getting total count
554
+ count_stmt = select(func.count()).select_from(stmt.alias())
555
+ count_result = await sess.execute(count_stmt)
556
+ total_count = count_result.scalar() or 0
557
+
558
+ # Sorting
559
+ stmt = apply_sorting(stmt, table, sort_by, sort_order)
560
+
561
+ # Paginating
562
+ if limit is not None:
563
+ stmt = stmt.limit(limit)
564
+ if page is not None:
565
+ stmt = stmt.offset((page - 1) * limit)
566
+
567
+ result = await sess.execute(stmt)
568
+ records = result.fetchall()
569
+ if records is None:
570
+ return [] if deserialize else ([], 0)
571
+
572
+ sessions_raw = [deserialize_session_json_fields(dict(record._mapping)) for record in records]
573
+ if not deserialize:
574
+ return sessions_raw, total_count
575
+ if not sessions_raw:
576
+ return []
577
+
578
+ if session_type == SessionType.AGENT:
579
+ return [AgentSession.from_dict(record) for record in sessions_raw] # type: ignore
580
+ elif session_type == SessionType.TEAM:
581
+ return [TeamSession.from_dict(record) for record in sessions_raw] # type: ignore
582
+ elif session_type == SessionType.WORKFLOW:
583
+ return [WorkflowSession.from_dict(record) for record in sessions_raw] # type: ignore
584
+ else:
585
+ raise ValueError(f"Invalid session type: {session_type}")
586
+
587
+ except Exception as e:
588
+ log_debug(f"Exception reading from sessions table: {e}")
589
+ raise e
590
+
591
+ async def rename_session(
592
+ self,
593
+ session_id: str,
594
+ session_type: SessionType,
595
+ session_name: str,
596
+ deserialize: Optional[bool] = True,
597
+ ) -> Optional[Union[Session, Dict[str, Any]]]:
598
+ """
599
+ Rename a session in the database.
600
+
601
+ Args:
602
+ session_id (str): The ID of the session to rename.
603
+ session_type (SessionType): The type of session to rename.
604
+ session_name (str): The new name for the session.
605
+ deserialize (Optional[bool]): Whether to serialize the session. Defaults to True.
606
+
607
+ Returns:
608
+ Optional[Union[Session, Dict[str, Any]]]:
609
+ - When deserialize=True: Session object
610
+ - When deserialize=False: Session dictionary
611
+
612
+ Raises:
613
+ Exception: If an error occurs during renaming.
614
+ """
615
+ try:
616
+ # Get the current session as a deserialized object
617
+ session = await self.get_session(session_id, session_type, deserialize=True)
618
+ if session is None:
619
+ return None
620
+
621
+ session = cast(Session, session)
622
+ # Update the session name
623
+ if session.session_data is None:
624
+ session.session_data = {}
625
+ session.session_data["session_name"] = session_name
626
+
627
+ # Upsert the updated session back to the database
628
+ return await self.upsert_session(session, deserialize=deserialize)
629
+
630
+ except Exception as e:
631
+ log_error(f"Exception renaming session: {e}")
632
+ raise e
633
+
634
+ async def upsert_session(
635
+ self, session: Session, deserialize: Optional[bool] = True
636
+ ) -> Optional[Union[Session, Dict[str, Any]]]:
637
+ """
638
+ Insert or update a session in the database.
639
+
640
+ Args:
641
+ session (Session): The session data to upsert.
642
+ deserialize (Optional[bool]): Whether to serialize the session. Defaults to True.
643
+
644
+ Returns:
645
+ Optional[Session]:
646
+ - When deserialize=True: Session object
647
+ - When deserialize=False: Session dictionary
648
+
649
+ Raises:
650
+ Exception: If an error occurs during upserting.
651
+ """
652
+ try:
653
+ table = await self._get_table(table_type="sessions")
654
+ if table is None:
655
+ return None
656
+
657
+ serialized_session = serialize_session_json_fields(session.to_dict())
658
+
659
+ if isinstance(session, AgentSession):
660
+ async with self.async_session_factory() as sess, sess.begin():
661
+ stmt = sqlite.insert(table).values(
662
+ session_id=serialized_session.get("session_id"),
663
+ session_type=SessionType.AGENT.value,
664
+ agent_id=serialized_session.get("agent_id"),
665
+ user_id=serialized_session.get("user_id"),
666
+ agent_data=serialized_session.get("agent_data"),
667
+ session_data=serialized_session.get("session_data"),
668
+ metadata=serialized_session.get("metadata"),
669
+ runs=serialized_session.get("runs"),
670
+ summary=serialized_session.get("summary"),
671
+ created_at=serialized_session.get("created_at"),
672
+ updated_at=serialized_session.get("created_at"),
673
+ )
674
+ stmt = stmt.on_conflict_do_update(
675
+ index_elements=["session_id"],
676
+ set_=dict(
677
+ agent_id=serialized_session.get("agent_id"),
678
+ user_id=serialized_session.get("user_id"),
679
+ runs=serialized_session.get("runs"),
680
+ summary=serialized_session.get("summary"),
681
+ agent_data=serialized_session.get("agent_data"),
682
+ session_data=serialized_session.get("session_data"),
683
+ metadata=serialized_session.get("metadata"),
684
+ updated_at=int(time.time()),
685
+ ),
686
+ )
687
+ stmt = stmt.returning(*table.columns) # type: ignore
688
+ result = await sess.execute(stmt)
689
+ row = result.fetchone()
690
+
691
+ session_raw = deserialize_session_json_fields(dict(row._mapping)) if row else None
692
+ if session_raw is None or not deserialize:
693
+ return session_raw
694
+ return AgentSession.from_dict(session_raw)
695
+
696
+ elif isinstance(session, TeamSession):
697
+ async with self.async_session_factory() as sess, sess.begin():
698
+ stmt = sqlite.insert(table).values(
699
+ session_id=serialized_session.get("session_id"),
700
+ session_type=SessionType.TEAM.value,
701
+ team_id=serialized_session.get("team_id"),
702
+ user_id=serialized_session.get("user_id"),
703
+ runs=serialized_session.get("runs"),
704
+ summary=serialized_session.get("summary"),
705
+ created_at=serialized_session.get("created_at"),
706
+ updated_at=serialized_session.get("created_at"),
707
+ team_data=serialized_session.get("team_data"),
708
+ session_data=serialized_session.get("session_data"),
709
+ metadata=serialized_session.get("metadata"),
710
+ )
711
+
712
+ stmt = stmt.on_conflict_do_update(
713
+ index_elements=["session_id"],
714
+ set_=dict(
715
+ team_id=serialized_session.get("team_id"),
716
+ user_id=serialized_session.get("user_id"),
717
+ summary=serialized_session.get("summary"),
718
+ runs=serialized_session.get("runs"),
719
+ team_data=serialized_session.get("team_data"),
720
+ session_data=serialized_session.get("session_data"),
721
+ metadata=serialized_session.get("metadata"),
722
+ updated_at=int(time.time()),
723
+ ),
724
+ )
725
+ stmt = stmt.returning(*table.columns) # type: ignore
726
+ result = await sess.execute(stmt)
727
+ row = result.fetchone()
728
+
729
+ session_raw = deserialize_session_json_fields(dict(row._mapping)) if row else None
730
+ if session_raw is None or not deserialize:
731
+ return session_raw
732
+ return TeamSession.from_dict(session_raw)
733
+
734
+ else:
735
+ async with self.async_session_factory() as sess, sess.begin():
736
+ stmt = sqlite.insert(table).values(
737
+ session_id=serialized_session.get("session_id"),
738
+ session_type=SessionType.WORKFLOW.value,
739
+ workflow_id=serialized_session.get("workflow_id"),
740
+ user_id=serialized_session.get("user_id"),
741
+ runs=serialized_session.get("runs"),
742
+ summary=serialized_session.get("summary"),
743
+ created_at=serialized_session.get("created_at") or int(time.time()),
744
+ updated_at=serialized_session.get("updated_at") or int(time.time()),
745
+ workflow_data=serialized_session.get("workflow_data"),
746
+ session_data=serialized_session.get("session_data"),
747
+ metadata=serialized_session.get("metadata"),
748
+ )
749
+ stmt = stmt.on_conflict_do_update(
750
+ index_elements=["session_id"],
751
+ set_=dict(
752
+ workflow_id=serialized_session.get("workflow_id"),
753
+ user_id=serialized_session.get("user_id"),
754
+ summary=serialized_session.get("summary"),
755
+ runs=serialized_session.get("runs"),
756
+ workflow_data=serialized_session.get("workflow_data"),
757
+ session_data=serialized_session.get("session_data"),
758
+ metadata=serialized_session.get("metadata"),
759
+ updated_at=int(time.time()),
760
+ ),
761
+ )
762
+ stmt = stmt.returning(*table.columns) # type: ignore
763
+ result = await sess.execute(stmt)
764
+ row = result.fetchone()
765
+
766
+ session_raw = deserialize_session_json_fields(dict(row._mapping)) if row else None
767
+ if session_raw is None or not deserialize:
768
+ return session_raw
769
+ return WorkflowSession.from_dict(session_raw)
770
+
771
+ except Exception as e:
772
+ log_warning(f"Exception upserting into table: {e}")
773
+ raise e
774
+
775
+ async def upsert_sessions(
776
+ self,
777
+ sessions: List[Session],
778
+ deserialize: Optional[bool] = True,
779
+ preserve_updated_at: bool = False,
780
+ ) -> List[Union[Session, Dict[str, Any]]]:
781
+ """
782
+ Bulk upsert multiple sessions for improved performance on large datasets.
783
+
784
+ Args:
785
+ sessions (List[Session]): List of sessions to upsert.
786
+ deserialize (Optional[bool]): Whether to deserialize the sessions. Defaults to True.
787
+ preserve_updated_at (bool): If True, preserve the updated_at from the session object.
788
+
789
+ Returns:
790
+ List[Union[Session, Dict[str, Any]]]: List of upserted sessions.
791
+
792
+ Raises:
793
+ Exception: If an error occurs during bulk upsert.
794
+ """
795
+ if not sessions:
796
+ return []
797
+
798
+ try:
799
+ table = await self._get_table(table_type="sessions")
800
+ if table is None:
801
+ log_info("Sessions table not available, falling back to individual upserts")
802
+ return [
803
+ result
804
+ for session in sessions
805
+ if session is not None
806
+ for result in [await self.upsert_session(session, deserialize=deserialize)]
807
+ if result is not None
808
+ ]
809
+
810
+ # Group sessions by type for batch processing
811
+ agent_sessions = []
812
+ team_sessions = []
813
+ workflow_sessions = []
814
+
815
+ for session in sessions:
816
+ if isinstance(session, AgentSession):
817
+ agent_sessions.append(session)
818
+ elif isinstance(session, TeamSession):
819
+ team_sessions.append(session)
820
+ elif isinstance(session, WorkflowSession):
821
+ workflow_sessions.append(session)
822
+
823
+ results: List[Union[Session, Dict[str, Any]]] = []
824
+
825
+ async with self.async_session_factory() as sess, sess.begin():
826
+ # Bulk upsert agent sessions
827
+ if agent_sessions:
828
+ agent_data = []
829
+ for session in agent_sessions:
830
+ serialized_session = serialize_session_json_fields(session.to_dict())
831
+ # Use preserved updated_at if flag is set and value exists, otherwise use current time
832
+ updated_at = serialized_session.get("updated_at") if preserve_updated_at else int(time.time())
833
+ agent_data.append(
834
+ {
835
+ "session_id": serialized_session.get("session_id"),
836
+ "session_type": SessionType.AGENT.value,
837
+ "agent_id": serialized_session.get("agent_id"),
838
+ "user_id": serialized_session.get("user_id"),
839
+ "agent_data": serialized_session.get("agent_data"),
840
+ "session_data": serialized_session.get("session_data"),
841
+ "metadata": serialized_session.get("metadata"),
842
+ "runs": serialized_session.get("runs"),
843
+ "summary": serialized_session.get("summary"),
844
+ "created_at": serialized_session.get("created_at"),
845
+ "updated_at": updated_at,
846
+ }
847
+ )
848
+
849
+ if agent_data:
850
+ stmt = sqlite.insert(table)
851
+ stmt = stmt.on_conflict_do_update(
852
+ index_elements=["session_id"],
853
+ set_=dict(
854
+ agent_id=stmt.excluded.agent_id,
855
+ user_id=stmt.excluded.user_id,
856
+ agent_data=stmt.excluded.agent_data,
857
+ session_data=stmt.excluded.session_data,
858
+ metadata=stmt.excluded.metadata,
859
+ runs=stmt.excluded.runs,
860
+ summary=stmt.excluded.summary,
861
+ updated_at=stmt.excluded.updated_at,
862
+ ),
863
+ )
864
+ await sess.execute(stmt, agent_data)
865
+
866
+ # Fetch the results for agent sessions
867
+ agent_ids = [session.session_id for session in agent_sessions]
868
+ select_stmt = select(table).where(table.c.session_id.in_(agent_ids))
869
+ result = (await sess.execute(select_stmt)).fetchall()
870
+
871
+ for row in result:
872
+ session_dict = deserialize_session_json_fields(dict(row._mapping))
873
+ if deserialize:
874
+ deserialized_agent_session = AgentSession.from_dict(session_dict)
875
+ if deserialized_agent_session is None:
876
+ continue
877
+ results.append(deserialized_agent_session)
878
+ else:
879
+ results.append(session_dict)
880
+
881
+ # Bulk upsert team sessions
882
+ if team_sessions:
883
+ team_data = []
884
+ for session in team_sessions:
885
+ serialized_session = serialize_session_json_fields(session.to_dict())
886
+ # Use preserved updated_at if flag is set and value exists, otherwise use current time
887
+ updated_at = serialized_session.get("updated_at") if preserve_updated_at else int(time.time())
888
+ team_data.append(
889
+ {
890
+ "session_id": serialized_session.get("session_id"),
891
+ "session_type": SessionType.TEAM.value,
892
+ "team_id": serialized_session.get("team_id"),
893
+ "user_id": serialized_session.get("user_id"),
894
+ "runs": serialized_session.get("runs"),
895
+ "summary": serialized_session.get("summary"),
896
+ "created_at": serialized_session.get("created_at"),
897
+ "updated_at": updated_at,
898
+ "team_data": serialized_session.get("team_data"),
899
+ "session_data": serialized_session.get("session_data"),
900
+ "metadata": serialized_session.get("metadata"),
901
+ }
902
+ )
903
+
904
+ if team_data:
905
+ stmt = sqlite.insert(table)
906
+ stmt = stmt.on_conflict_do_update(
907
+ index_elements=["session_id"],
908
+ set_=dict(
909
+ team_id=stmt.excluded.team_id,
910
+ user_id=stmt.excluded.user_id,
911
+ team_data=stmt.excluded.team_data,
912
+ session_data=stmt.excluded.session_data,
913
+ metadata=stmt.excluded.metadata,
914
+ runs=stmt.excluded.runs,
915
+ summary=stmt.excluded.summary,
916
+ updated_at=stmt.excluded.updated_at,
917
+ ),
918
+ )
919
+ await sess.execute(stmt, team_data)
920
+
921
+ # Fetch the results for team sessions
922
+ team_ids = [session.session_id for session in team_sessions]
923
+ select_stmt = select(table).where(table.c.session_id.in_(team_ids))
924
+ result = (await sess.execute(select_stmt)).fetchall()
925
+
926
+ for row in result:
927
+ session_dict = deserialize_session_json_fields(dict(row._mapping))
928
+ if deserialize:
929
+ deserialized_team_session = TeamSession.from_dict(session_dict)
930
+ if deserialized_team_session is None:
931
+ continue
932
+ results.append(deserialized_team_session)
933
+ else:
934
+ results.append(session_dict)
935
+
936
+ # Bulk upsert workflow sessions
937
+ if workflow_sessions:
938
+ workflow_data = []
939
+ for session in workflow_sessions:
940
+ serialized_session = serialize_session_json_fields(session.to_dict())
941
+ # Use preserved updated_at if flag is set and value exists, otherwise use current time
942
+ updated_at = serialized_session.get("updated_at") if preserve_updated_at else int(time.time())
943
+ workflow_data.append(
944
+ {
945
+ "session_id": serialized_session.get("session_id"),
946
+ "session_type": SessionType.WORKFLOW.value,
947
+ "workflow_id": serialized_session.get("workflow_id"),
948
+ "user_id": serialized_session.get("user_id"),
949
+ "runs": serialized_session.get("runs"),
950
+ "summary": serialized_session.get("summary"),
951
+ "created_at": serialized_session.get("created_at"),
952
+ "updated_at": updated_at,
953
+ "workflow_data": serialized_session.get("workflow_data"),
954
+ "session_data": serialized_session.get("session_data"),
955
+ "metadata": serialized_session.get("metadata"),
956
+ }
957
+ )
958
+
959
+ if workflow_data:
960
+ stmt = sqlite.insert(table)
961
+ stmt = stmt.on_conflict_do_update(
962
+ index_elements=["session_id"],
963
+ set_=dict(
964
+ workflow_id=stmt.excluded.workflow_id,
965
+ user_id=stmt.excluded.user_id,
966
+ workflow_data=stmt.excluded.workflow_data,
967
+ session_data=stmt.excluded.session_data,
968
+ metadata=stmt.excluded.metadata,
969
+ runs=stmt.excluded.runs,
970
+ summary=stmt.excluded.summary,
971
+ updated_at=stmt.excluded.updated_at,
972
+ ),
973
+ )
974
+ await sess.execute(stmt, workflow_data)
975
+
976
+ # Fetch the results for workflow sessions
977
+ workflow_ids = [session.session_id for session in workflow_sessions]
978
+ select_stmt = select(table).where(table.c.session_id.in_(workflow_ids))
979
+ result = (await sess.execute(select_stmt)).fetchall()
980
+
981
+ for row in result:
982
+ session_dict = deserialize_session_json_fields(dict(row._mapping))
983
+ if deserialize:
984
+ deserialized_workflow_session = WorkflowSession.from_dict(session_dict)
985
+ if deserialized_workflow_session is None:
986
+ continue
987
+ results.append(deserialized_workflow_session)
988
+ else:
989
+ results.append(session_dict)
990
+
991
+ return results
992
+
993
+ except Exception as e:
994
+ log_error(f"Exception during bulk session upsert, falling back to individual upserts: {e}")
995
+ # Fallback to individual upserts
996
+ return [
997
+ result
998
+ for session in sessions
999
+ if session is not None
1000
+ for result in [await self.upsert_session(session, deserialize=deserialize)]
1001
+ if result is not None
1002
+ ]
1003
+
1004
+ # -- Memory methods --
1005
+
1006
+ async def delete_user_memory(self, memory_id: str, user_id: Optional[str] = None):
1007
+ """Delete a user memory from the database.
1008
+
1009
+ Args:
1010
+ memory_id (str): The ID of the memory to delete.
1011
+ user_id (Optional[str]): The user ID to filter by. Defaults to None.
1012
+
1013
+ Returns:
1014
+ bool: True if deletion was successful, False otherwise.
1015
+
1016
+ Raises:
1017
+ Exception: If an error occurs during deletion.
1018
+ """
1019
+ try:
1020
+ table = await self._get_table(table_type="memories")
1021
+ if table is None:
1022
+ return
1023
+
1024
+ async with self.async_session_factory() as sess, sess.begin():
1025
+ delete_stmt = table.delete().where(table.c.memory_id == memory_id)
1026
+ if user_id is not None:
1027
+ delete_stmt = delete_stmt.where(table.c.user_id == user_id)
1028
+ result = await sess.execute(delete_stmt)
1029
+
1030
+ success = result.rowcount > 0 # type: ignore
1031
+ if success:
1032
+ log_debug(f"Successfully deleted user memory id: {memory_id}")
1033
+ else:
1034
+ log_debug(f"No user memory found with id: {memory_id}")
1035
+
1036
+ except Exception as e:
1037
+ log_error(f"Error deleting user memory: {e}")
1038
+ raise e
1039
+
1040
+ async def delete_user_memories(self, memory_ids: List[str], user_id: Optional[str] = None) -> None:
1041
+ """Delete user memories from the database.
1042
+
1043
+ Args:
1044
+ memory_ids (List[str]): The IDs of the memories to delete.
1045
+ user_id (Optional[str]): The user ID to filter by. Defaults to None.
1046
+
1047
+ Raises:
1048
+ Exception: If an error occurs during deletion.
1049
+ """
1050
+ try:
1051
+ table = await self._get_table(table_type="memories")
1052
+ if table is None:
1053
+ return
1054
+
1055
+ async with self.async_session_factory() as sess, sess.begin():
1056
+ delete_stmt = table.delete().where(table.c.memory_id.in_(memory_ids))
1057
+ if user_id is not None:
1058
+ delete_stmt = delete_stmt.where(table.c.user_id == user_id)
1059
+ result = await sess.execute(delete_stmt)
1060
+ if result.rowcount == 0: # type: ignore
1061
+ log_debug(f"No user memories found with ids: {memory_ids}")
1062
+
1063
+ except Exception as e:
1064
+ log_error(f"Error deleting user memories: {e}")
1065
+ raise e
1066
+
1067
+ async def get_all_memory_topics(self) -> List[str]:
1068
+ """Get all memory topics from the database.
1069
+
1070
+ Returns:
1071
+ List[str]: List of memory topics.
1072
+ """
1073
+ try:
1074
+ table = await self._get_table(table_type="memories")
1075
+ if table is None:
1076
+ return []
1077
+
1078
+ async with self.async_session_factory() as sess, sess.begin():
1079
+ # Select topics from all results
1080
+ stmt = select(func.json_array_elements_text(table.c.topics)).select_from(table)
1081
+ result = (await sess.execute(stmt)).fetchall()
1082
+
1083
+ return list(set([record[0] for record in result]))
1084
+
1085
+ except Exception as e:
1086
+ log_debug(f"Exception reading from memory table: {e}")
1087
+ raise e
1088
+
1089
+ async def get_user_memory(
1090
+ self,
1091
+ memory_id: str,
1092
+ deserialize: Optional[bool] = True,
1093
+ user_id: Optional[str] = None,
1094
+ ) -> Optional[Union[UserMemory, Dict[str, Any]]]:
1095
+ """Get a memory from the database.
1096
+
1097
+ Args:
1098
+ memory_id (str): The ID of the memory to get.
1099
+ deserialize (Optional[bool]): Whether to serialize the memory. Defaults to True.
1100
+ user_id (Optional[str]): The user ID to filter by. Defaults to None.
1101
+
1102
+ Returns:
1103
+ Optional[Union[UserMemory, Dict[str, Any]]]:
1104
+ - When deserialize=True: UserMemory object
1105
+ - When deserialize=False: Memory dictionary
1106
+
1107
+ Raises:
1108
+ Exception: If an error occurs during retrieval.
1109
+ """
1110
+ try:
1111
+ table = await self._get_table(table_type="memories")
1112
+ if table is None:
1113
+ return None
1114
+
1115
+ async with self.async_session_factory() as sess, sess.begin():
1116
+ stmt = select(table).where(table.c.memory_id == memory_id)
1117
+ if user_id is not None:
1118
+ stmt = stmt.where(table.c.user_id == user_id)
1119
+ result = (await sess.execute(stmt)).fetchone()
1120
+ if result is None:
1121
+ return None
1122
+
1123
+ memory_raw = dict(result._mapping)
1124
+ if not memory_raw or not deserialize:
1125
+ return memory_raw
1126
+
1127
+ return UserMemory.from_dict(memory_raw)
1128
+
1129
+ except Exception as e:
1130
+ log_debug(f"Exception reading from memorytable: {e}")
1131
+ raise e
1132
+
1133
+ async def get_user_memories(
1134
+ self,
1135
+ user_id: Optional[str] = None,
1136
+ agent_id: Optional[str] = None,
1137
+ team_id: Optional[str] = None,
1138
+ topics: Optional[List[str]] = None,
1139
+ search_content: Optional[str] = None,
1140
+ limit: Optional[int] = None,
1141
+ page: Optional[int] = None,
1142
+ sort_by: Optional[str] = None,
1143
+ sort_order: Optional[str] = None,
1144
+ deserialize: Optional[bool] = True,
1145
+ ) -> Union[List[UserMemory], Tuple[List[Dict[str, Any]], int]]:
1146
+ """Get all memories from the database as UserMemory objects.
1147
+
1148
+ Args:
1149
+ user_id (Optional[str]): The ID of the user to filter by.
1150
+ agent_id (Optional[str]): The ID of the agent to filter by.
1151
+ team_id (Optional[str]): The ID of the team to filter by.
1152
+ topics (Optional[List[str]]): The topics to filter by.
1153
+ search_content (Optional[str]): The content to search for.
1154
+ limit (Optional[int]): The maximum number of memories to return.
1155
+ page (Optional[int]): The page number.
1156
+ sort_by (Optional[str]): The column to sort by.
1157
+ sort_order (Optional[str]): The order to sort by.
1158
+ deserialize (Optional[bool]): Whether to serialize the memories. Defaults to True.
1159
+
1160
+
1161
+ Returns:
1162
+ Union[List[UserMemory], Tuple[List[Dict[str, Any]], int]]:
1163
+ - When deserialize=True: List of UserMemory objects
1164
+ - When deserialize=False: List of UserMemory dictionaries and total count
1165
+
1166
+ Raises:
1167
+ Exception: If an error occurs during retrieval.
1168
+ """
1169
+ try:
1170
+ table = await self._get_table(table_type="memories")
1171
+ if table is None:
1172
+ return [] if deserialize else ([], 0)
1173
+
1174
+ async with self.async_session_factory() as sess, sess.begin():
1175
+ stmt = select(table)
1176
+
1177
+ # Filtering
1178
+ if user_id is not None:
1179
+ stmt = stmt.where(table.c.user_id == user_id)
1180
+ if agent_id is not None:
1181
+ stmt = stmt.where(table.c.agent_id == agent_id)
1182
+ if team_id is not None:
1183
+ stmt = stmt.where(table.c.team_id == team_id)
1184
+ if topics is not None:
1185
+ topic_conditions = [text(f"topics::text LIKE '%\"{topic}\"%'") for topic in topics]
1186
+ stmt = stmt.where(and_(*topic_conditions))
1187
+ if search_content is not None:
1188
+ stmt = stmt.where(table.c.memory.ilike(f"%{search_content}%"))
1189
+
1190
+ # Get total count after applying filtering
1191
+ count_stmt = select(func.count()).select_from(stmt.alias())
1192
+ total_count = (await sess.execute(count_stmt)).scalar() or 0
1193
+
1194
+ # Sorting
1195
+ stmt = apply_sorting(stmt, table, sort_by, sort_order)
1196
+ # Paginating
1197
+ if limit is not None:
1198
+ stmt = stmt.limit(limit)
1199
+ if page is not None:
1200
+ stmt = stmt.offset((page - 1) * limit)
1201
+
1202
+ result = (await sess.execute(stmt)).fetchall()
1203
+ if not result:
1204
+ return [] if deserialize else ([], 0)
1205
+
1206
+ memories_raw = [dict(record._mapping) for record in result]
1207
+
1208
+ if not deserialize:
1209
+ return memories_raw, total_count
1210
+
1211
+ return [UserMemory.from_dict(record) for record in memories_raw]
1212
+
1213
+ except Exception as e:
1214
+ log_error(f"Error reading from memory table: {e}")
1215
+ raise e
1216
+
1217
+ async def get_user_memory_stats(
1218
+ self,
1219
+ limit: Optional[int] = None,
1220
+ page: Optional[int] = None,
1221
+ ) -> Tuple[List[Dict[str, Any]], int]:
1222
+ """Get user memories stats.
1223
+
1224
+ Args:
1225
+ limit (Optional[int]): The maximum number of user stats to return.
1226
+ page (Optional[int]): The page number.
1227
+
1228
+ Returns:
1229
+ Tuple[List[Dict[str, Any]], int]: A list of dictionaries containing user stats and total count.
1230
+
1231
+ Example:
1232
+ (
1233
+ [
1234
+ {
1235
+ "user_id": "123",
1236
+ "total_memories": 10,
1237
+ "last_memory_updated_at": 1714560000,
1238
+ },
1239
+ ],
1240
+ total_count: 1,
1241
+ )
1242
+ """
1243
+ try:
1244
+ table = await self._get_table(table_type="memories")
1245
+ if table is None:
1246
+ return [], 0
1247
+
1248
+ async with self.async_session_factory() as sess, sess.begin():
1249
+ stmt = (
1250
+ select(
1251
+ table.c.user_id,
1252
+ func.count(table.c.memory_id).label("total_memories"),
1253
+ func.max(table.c.updated_at).label("last_memory_updated_at"),
1254
+ )
1255
+ .where(table.c.user_id.is_not(None))
1256
+ .group_by(table.c.user_id)
1257
+ .order_by(func.max(table.c.updated_at).desc())
1258
+ )
1259
+
1260
+ count_stmt = select(func.count()).select_from(stmt.alias())
1261
+ total_count = (await sess.execute(count_stmt)).scalar() or 0
1262
+
1263
+ # Pagination
1264
+ if limit is not None:
1265
+ stmt = stmt.limit(limit)
1266
+ if page is not None:
1267
+ stmt = stmt.offset((page - 1) * limit)
1268
+
1269
+ result = (await sess.execute(stmt)).fetchall()
1270
+ if not result:
1271
+ return [], 0
1272
+
1273
+ return [
1274
+ {
1275
+ "user_id": record.user_id, # type: ignore
1276
+ "total_memories": record.total_memories,
1277
+ "last_memory_updated_at": record.last_memory_updated_at,
1278
+ }
1279
+ for record in result
1280
+ ], total_count
1281
+
1282
+ except Exception as e:
1283
+ log_error(f"Error getting user memory stats: {e}")
1284
+ raise e
1285
+
1286
+ async def upsert_user_memory(
1287
+ self, memory: UserMemory, deserialize: Optional[bool] = True
1288
+ ) -> Optional[Union[UserMemory, Dict[str, Any]]]:
1289
+ """Upsert a user memory in the database.
1290
+
1291
+ Args:
1292
+ memory (UserMemory): The user memory to upsert.
1293
+ deserialize (Optional[bool]): Whether to serialize the memory. Defaults to True.
1294
+
1295
+ Returns:
1296
+ Optional[Union[UserMemory, Dict[str, Any]]]:
1297
+ - When deserialize=True: UserMemory object
1298
+ - When deserialize=False: UserMemory dictionary
1299
+
1300
+ Raises:
1301
+ Exception: If an error occurs during upsert.
1302
+ """
1303
+ try:
1304
+ table = await self._get_table(table_type="memories")
1305
+ if table is None:
1306
+ return None
1307
+
1308
+ if memory.memory_id is None:
1309
+ memory.memory_id = str(uuid4())
1310
+
1311
+ current_time = int(time.time())
1312
+
1313
+ async with self.async_session_factory() as sess:
1314
+ async with sess.begin():
1315
+ stmt = sqlite.insert(table).values(
1316
+ user_id=memory.user_id,
1317
+ agent_id=memory.agent_id,
1318
+ team_id=memory.team_id,
1319
+ memory_id=memory.memory_id,
1320
+ memory=memory.memory,
1321
+ topics=memory.topics,
1322
+ input=memory.input,
1323
+ feedback=memory.feedback,
1324
+ created_at=memory.created_at,
1325
+ updated_at=memory.created_at,
1326
+ )
1327
+ stmt = stmt.on_conflict_do_update( # type: ignore
1328
+ index_elements=["memory_id"],
1329
+ set_=dict(
1330
+ memory=memory.memory,
1331
+ topics=memory.topics,
1332
+ input=memory.input,
1333
+ agent_id=memory.agent_id,
1334
+ team_id=memory.team_id,
1335
+ feedback=memory.feedback,
1336
+ updated_at=current_time,
1337
+ # Preserve created_at on update - don't overwrite existing value
1338
+ created_at=table.c.created_at,
1339
+ ),
1340
+ ).returning(table)
1341
+
1342
+ result = await sess.execute(stmt)
1343
+ row = result.fetchone()
1344
+
1345
+ if row is None:
1346
+ return None
1347
+
1348
+ memory_raw = dict(row._mapping)
1349
+ if not memory_raw or not deserialize:
1350
+ return memory_raw
1351
+
1352
+ return UserMemory.from_dict(memory_raw)
1353
+
1354
+ except Exception as e:
1355
+ log_error(f"Error upserting user memory: {e}")
1356
+ raise e
1357
+
1358
+ async def upsert_memories(
1359
+ self,
1360
+ memories: List[UserMemory],
1361
+ deserialize: Optional[bool] = True,
1362
+ preserve_updated_at: bool = False,
1363
+ ) -> List[Union[UserMemory, Dict[str, Any]]]:
1364
+ """
1365
+ Bulk upsert multiple user memories for improved performance on large datasets.
1366
+
1367
+ Args:
1368
+ memories (List[UserMemory]): List of memories to upsert.
1369
+ deserialize (Optional[bool]): Whether to deserialize the memories. Defaults to True.
1370
+
1371
+ Returns:
1372
+ List[Union[UserMemory, Dict[str, Any]]]: List of upserted memories.
1373
+
1374
+ Raises:
1375
+ Exception: If an error occurs during bulk upsert.
1376
+ """
1377
+ if not memories:
1378
+ return []
1379
+
1380
+ try:
1381
+ table = await self._get_table(table_type="memories")
1382
+ if table is None:
1383
+ log_info("Memories table not available, falling back to individual upserts")
1384
+ return [
1385
+ result
1386
+ for memory in memories
1387
+ if memory is not None
1388
+ for result in [await self.upsert_user_memory(memory, deserialize=deserialize)]
1389
+ if result is not None
1390
+ ]
1391
+ # Prepare bulk data
1392
+ bulk_data = []
1393
+ current_time = int(time.time())
1394
+
1395
+ for memory in memories:
1396
+ if memory.memory_id is None:
1397
+ memory.memory_id = str(uuid4())
1398
+
1399
+ # Use preserved updated_at if flag is set and value exists, otherwise use current time
1400
+ updated_at = memory.updated_at if preserve_updated_at else current_time
1401
+
1402
+ bulk_data.append(
1403
+ {
1404
+ "user_id": memory.user_id,
1405
+ "agent_id": memory.agent_id,
1406
+ "team_id": memory.team_id,
1407
+ "memory_id": memory.memory_id,
1408
+ "memory": memory.memory,
1409
+ "topics": memory.topics,
1410
+ "input": memory.input,
1411
+ "feedback": memory.feedback,
1412
+ "created_at": memory.created_at,
1413
+ "updated_at": updated_at,
1414
+ }
1415
+ )
1416
+
1417
+ results: List[Union[UserMemory, Dict[str, Any]]] = []
1418
+
1419
+ async with self.async_session_factory() as sess, sess.begin():
1420
+ # Bulk upsert memories using SQLite ON CONFLICT DO UPDATE
1421
+ stmt = sqlite.insert(table)
1422
+ stmt = stmt.on_conflict_do_update(
1423
+ index_elements=["memory_id"],
1424
+ set_=dict(
1425
+ memory=stmt.excluded.memory,
1426
+ topics=stmt.excluded.topics,
1427
+ input=stmt.excluded.input,
1428
+ agent_id=stmt.excluded.agent_id,
1429
+ team_id=stmt.excluded.team_id,
1430
+ feedback=stmt.excluded.feedback,
1431
+ updated_at=stmt.excluded.updated_at,
1432
+ # Preserve created_at on update
1433
+ created_at=table.c.created_at,
1434
+ ),
1435
+ )
1436
+ await sess.execute(stmt, bulk_data)
1437
+
1438
+ # Fetch results
1439
+ memory_ids = [memory.memory_id for memory in memories if memory.memory_id]
1440
+ select_stmt = select(table).where(table.c.memory_id.in_(memory_ids))
1441
+ result = (await sess.execute(select_stmt)).fetchall()
1442
+
1443
+ for row in result:
1444
+ memory_dict = dict(row._mapping)
1445
+ if deserialize:
1446
+ results.append(UserMemory.from_dict(memory_dict))
1447
+ else:
1448
+ results.append(memory_dict)
1449
+
1450
+ return results
1451
+
1452
+ except Exception as e:
1453
+ log_error(f"Exception during bulk memory upsert, falling back to individual upserts: {e}")
1454
+
1455
+ # Fallback to individual upserts
1456
+ return [
1457
+ result
1458
+ for memory in memories
1459
+ if memory is not None
1460
+ for result in [await self.upsert_user_memory(memory, deserialize=deserialize)]
1461
+ if result is not None
1462
+ ]
1463
+
1464
+ async def clear_memories(self) -> None:
1465
+ """Delete all memories from the database.
1466
+
1467
+ Raises:
1468
+ Exception: If an error occurs during deletion.
1469
+ """
1470
+ try:
1471
+ table = await self._get_table(table_type="memories")
1472
+ if table is None:
1473
+ return
1474
+
1475
+ async with self.async_session_factory() as sess, sess.begin():
1476
+ await sess.execute(table.delete())
1477
+
1478
+ except Exception as e:
1479
+ from agno.utils.log import log_warning
1480
+
1481
+ log_warning(f"Exception deleting all memories: {e}")
1482
+ raise e
1483
+
1484
+ # -- Metrics methods --
1485
+
1486
+ async def _get_all_sessions_for_metrics_calculation(
1487
+ self, start_timestamp: Optional[int] = None, end_timestamp: Optional[int] = None
1488
+ ) -> List[Dict[str, Any]]:
1489
+ """
1490
+ Get all sessions of all types (agent, team, workflow) as raw dictionaries.
1491
+
1492
+ Args:
1493
+ start_timestamp (Optional[int]): The start timestamp to filter by. Defaults to None.
1494
+ end_timestamp (Optional[int]): The end timestamp to filter by. Defaults to None.
1495
+
1496
+ Returns:
1497
+ List[Dict[str, Any]]: List of session dictionaries with session_type field.
1498
+
1499
+ Raises:
1500
+ Exception: If an error occurs during retrieval.
1501
+ """
1502
+ try:
1503
+ table = await self._get_table(table_type="sessions")
1504
+ if table is None:
1505
+ return []
1506
+
1507
+ stmt = select(
1508
+ table.c.user_id,
1509
+ table.c.session_data,
1510
+ table.c.runs,
1511
+ table.c.created_at,
1512
+ table.c.session_type,
1513
+ )
1514
+
1515
+ if start_timestamp is not None:
1516
+ stmt = stmt.where(table.c.created_at >= start_timestamp)
1517
+ if end_timestamp is not None:
1518
+ stmt = stmt.where(table.c.created_at <= end_timestamp)
1519
+
1520
+ async with self.async_session_factory() as sess:
1521
+ result = (await sess.execute(stmt)).fetchall()
1522
+ return [dict(record._mapping) for record in result]
1523
+
1524
+ except Exception as e:
1525
+ log_error(f"Error reading from sessions table: {e}")
1526
+ raise e
1527
+
1528
+ async def _get_metrics_calculation_starting_date(self, table: Table) -> Optional[date]:
1529
+ """Get the first date for which metrics calculation is needed:
1530
+
1531
+ 1. If there are metrics records, return the date of the first day without a complete metrics record.
1532
+ 2. If there are no metrics records, return the date of the first recorded session.
1533
+ 3. If there are no metrics records and no sessions records, return None.
1534
+
1535
+ Args:
1536
+ table (Table): The table to get the starting date for.
1537
+
1538
+ Returns:
1539
+ Optional[date]: The starting date for which metrics calculation is needed.
1540
+ """
1541
+ async with self.async_session_factory() as sess:
1542
+ stmt = select(table).order_by(table.c.date.desc()).limit(1)
1543
+ result = (await sess.execute(stmt)).fetchone()
1544
+
1545
+ # 1. Return the date of the first day without a complete metrics record.
1546
+ if result is not None:
1547
+ if result.completed:
1548
+ return result._mapping["date"] + timedelta(days=1)
1549
+ else:
1550
+ return result._mapping["date"]
1551
+
1552
+ # 2. No metrics records. Return the date of the first recorded session.
1553
+ first_session, _ = await self.get_sessions(sort_by="created_at", sort_order="asc", limit=1, deserialize=False)
1554
+ first_session_date = first_session[0]["created_at"] if first_session else None # type: ignore
1555
+
1556
+ # 3. No metrics records and no sessions records. Return None.
1557
+ if not first_session_date:
1558
+ return None
1559
+
1560
+ return datetime.fromtimestamp(first_session_date, tz=timezone.utc).date()
1561
+
1562
+ async def calculate_metrics(self) -> Optional[list[dict]]:
1563
+ """Calculate metrics for all dates without complete metrics.
1564
+
1565
+ Returns:
1566
+ Optional[list[dict]]: The calculated metrics.
1567
+
1568
+ Raises:
1569
+ Exception: If an error occurs during metrics calculation.
1570
+ """
1571
+ try:
1572
+ table = await self._get_table(table_type="metrics")
1573
+ if table is None:
1574
+ return None
1575
+
1576
+ starting_date = await self._get_metrics_calculation_starting_date(table)
1577
+ if starting_date is None:
1578
+ log_info("No session data found. Won't calculate metrics.")
1579
+ return None
1580
+
1581
+ dates_to_process = get_dates_to_calculate_metrics_for(starting_date)
1582
+ if not dates_to_process:
1583
+ log_info("Metrics already calculated for all relevant dates.")
1584
+ return None
1585
+
1586
+ start_timestamp = int(
1587
+ datetime.combine(dates_to_process[0], datetime.min.time()).replace(tzinfo=timezone.utc).timestamp()
1588
+ )
1589
+ end_timestamp = int(
1590
+ datetime.combine(dates_to_process[-1] + timedelta(days=1), datetime.min.time())
1591
+ .replace(tzinfo=timezone.utc)
1592
+ .timestamp()
1593
+ )
1594
+
1595
+ sessions = await self._get_all_sessions_for_metrics_calculation(
1596
+ start_timestamp=start_timestamp, end_timestamp=end_timestamp
1597
+ )
1598
+ all_sessions_data = fetch_all_sessions_data(
1599
+ sessions=sessions,
1600
+ dates_to_process=dates_to_process,
1601
+ start_timestamp=start_timestamp,
1602
+ )
1603
+ if not all_sessions_data:
1604
+ log_info("No new session data found. Won't calculate metrics.")
1605
+ return None
1606
+
1607
+ results = []
1608
+ metrics_records = []
1609
+
1610
+ for date_to_process in dates_to_process:
1611
+ date_key = date_to_process.isoformat()
1612
+ sessions_for_date = all_sessions_data.get(date_key, {})
1613
+
1614
+ # Skip dates with no sessions
1615
+ if not any(len(sessions) > 0 for sessions in sessions_for_date.values()):
1616
+ continue
1617
+
1618
+ metrics_record = calculate_date_metrics(date_to_process, sessions_for_date)
1619
+ metrics_records.append(metrics_record)
1620
+
1621
+ if metrics_records:
1622
+ async with self.async_session_factory() as sess, sess.begin():
1623
+ results = await abulk_upsert_metrics(session=sess, table=table, metrics_records=metrics_records)
1624
+
1625
+ log_debug("Updated metrics calculations")
1626
+
1627
+ return results
1628
+
1629
+ except Exception as e:
1630
+ log_error(f"Error refreshing metrics: {e}")
1631
+ raise e
1632
+
1633
+ async def get_metrics(
1634
+ self,
1635
+ starting_date: Optional[date] = None,
1636
+ ending_date: Optional[date] = None,
1637
+ ) -> Tuple[List[dict], Optional[int]]:
1638
+ """Get all metrics matching the given date range.
1639
+
1640
+ Args:
1641
+ starting_date (Optional[date]): The starting date to filter metrics by.
1642
+ ending_date (Optional[date]): The ending date to filter metrics by.
1643
+
1644
+ Returns:
1645
+ Tuple[List[dict], Optional[int]]: A tuple containing the metrics and the timestamp of the latest update.
1646
+
1647
+ Raises:
1648
+ Exception: If an error occurs during retrieval.
1649
+ """
1650
+ try:
1651
+ table = await self._get_table(table_type="metrics")
1652
+ if table is None:
1653
+ return [], None
1654
+
1655
+ async with self.async_session_factory() as sess, sess.begin():
1656
+ stmt = select(table)
1657
+ if starting_date:
1658
+ stmt = stmt.where(table.c.date >= starting_date)
1659
+ if ending_date:
1660
+ stmt = stmt.where(table.c.date <= ending_date)
1661
+ result = (await sess.execute(stmt)).fetchall()
1662
+ if not result:
1663
+ return [], None
1664
+
1665
+ # Get the latest updated_at
1666
+ latest_stmt = select(func.max(table.c.updated_at))
1667
+ latest_updated_at = (await sess.execute(latest_stmt)).scalar()
1668
+
1669
+ return [dict(row._mapping) for row in result], latest_updated_at
1670
+
1671
+ except Exception as e:
1672
+ log_error(f"Error getting metrics: {e}")
1673
+ raise e
1674
+
1675
+ # -- Knowledge methods --
1676
+
1677
+ async def delete_knowledge_content(self, id: str):
1678
+ """Delete a knowledge row from the database.
1679
+
1680
+ Args:
1681
+ id (str): The ID of the knowledge row to delete.
1682
+
1683
+ Raises:
1684
+ Exception: If an error occurs during deletion.
1685
+ """
1686
+ table = await self._get_table(table_type="knowledge")
1687
+ if table is None:
1688
+ return
1689
+
1690
+ try:
1691
+ async with self.async_session_factory() as sess, sess.begin():
1692
+ stmt = table.delete().where(table.c.id == id)
1693
+ await sess.execute(stmt)
1694
+
1695
+ except Exception as e:
1696
+ log_error(f"Error deleting knowledge content: {e}")
1697
+ raise e
1698
+
1699
+ async def get_knowledge_content(self, id: str) -> Optional[KnowledgeRow]:
1700
+ """Get a knowledge row from the database.
1701
+
1702
+ Args:
1703
+ id (str): The ID of the knowledge row to get.
1704
+
1705
+ Returns:
1706
+ Optional[KnowledgeRow]: The knowledge row, or None if it doesn't exist.
1707
+
1708
+ Raises:
1709
+ Exception: If an error occurs during retrieval.
1710
+ """
1711
+ table = await self._get_table(table_type="knowledge")
1712
+ if table is None:
1713
+ return None
1714
+
1715
+ try:
1716
+ async with self.async_session_factory() as sess, sess.begin():
1717
+ stmt = select(table).where(table.c.id == id)
1718
+ result = (await sess.execute(stmt)).fetchone()
1719
+ if result is None:
1720
+ return None
1721
+
1722
+ return KnowledgeRow.model_validate(result._mapping)
1723
+
1724
+ except Exception as e:
1725
+ log_error(f"Error getting knowledge content: {e}")
1726
+ raise e
1727
+
1728
+ async def get_knowledge_contents(
1729
+ self,
1730
+ limit: Optional[int] = None,
1731
+ page: Optional[int] = None,
1732
+ sort_by: Optional[str] = None,
1733
+ sort_order: Optional[str] = None,
1734
+ ) -> Tuple[List[KnowledgeRow], int]:
1735
+ """Get all knowledge contents from the database.
1736
+
1737
+ Args:
1738
+ limit (Optional[int]): The maximum number of knowledge contents to return.
1739
+ page (Optional[int]): The page number.
1740
+ sort_by (Optional[str]): The column to sort by.
1741
+ sort_order (Optional[str]): The order to sort by.
1742
+
1743
+ Returns:
1744
+ Tuple[List[KnowledgeRow], int]: The knowledge contents and total count.
1745
+
1746
+ Raises:
1747
+ Exception: If an error occurs during retrieval.
1748
+ """
1749
+ table = await self._get_table(table_type="knowledge")
1750
+ if table is None:
1751
+ return [], 0
1752
+
1753
+ try:
1754
+ async with self.async_session_factory() as sess, sess.begin():
1755
+ stmt = select(table)
1756
+
1757
+ # Apply sorting
1758
+ if sort_by is not None:
1759
+ stmt = stmt.order_by(getattr(table.c, sort_by) * (1 if sort_order == "asc" else -1))
1760
+
1761
+ # Get total count before applying limit and pagination
1762
+ count_stmt = select(func.count()).select_from(stmt.alias())
1763
+ total_count = (await sess.execute(count_stmt)).scalar() or 0
1764
+
1765
+ # Apply pagination after count
1766
+ if limit is not None:
1767
+ stmt = stmt.limit(limit)
1768
+ if page is not None:
1769
+ stmt = stmt.offset((page - 1) * limit)
1770
+
1771
+ result = (await sess.execute(stmt)).fetchall()
1772
+ return [KnowledgeRow.model_validate(record._mapping) for record in result], total_count
1773
+
1774
+ except Exception as e:
1775
+ log_error(f"Error getting knowledge contents: {e}")
1776
+ raise e
1777
+
1778
+ async def upsert_knowledge_content(self, knowledge_row: KnowledgeRow):
1779
+ """Upsert knowledge content in the database.
1780
+
1781
+ Args:
1782
+ knowledge_row (KnowledgeRow): The knowledge row to upsert.
1783
+
1784
+ Returns:
1785
+ Optional[KnowledgeRow]: The upserted knowledge row, or None if the operation fails.
1786
+ """
1787
+ try:
1788
+ table = await self._get_table(table_type="knowledge")
1789
+ if table is None:
1790
+ return None
1791
+
1792
+ async with self.async_session_factory() as sess, sess.begin():
1793
+ update_fields = {
1794
+ k: v
1795
+ for k, v in {
1796
+ "name": knowledge_row.name,
1797
+ "description": knowledge_row.description,
1798
+ "metadata": knowledge_row.metadata,
1799
+ "type": knowledge_row.type,
1800
+ "size": knowledge_row.size,
1801
+ "linked_to": knowledge_row.linked_to,
1802
+ "access_count": knowledge_row.access_count,
1803
+ "status": knowledge_row.status,
1804
+ "status_message": knowledge_row.status_message,
1805
+ "created_at": knowledge_row.created_at,
1806
+ "updated_at": knowledge_row.updated_at,
1807
+ "external_id": knowledge_row.external_id,
1808
+ }.items()
1809
+ # Filtering out None fields if updating
1810
+ if v is not None
1811
+ }
1812
+
1813
+ stmt = (
1814
+ sqlite.insert(table)
1815
+ .values(knowledge_row.model_dump())
1816
+ .on_conflict_do_update(index_elements=["id"], set_=update_fields)
1817
+ )
1818
+ await sess.execute(stmt)
1819
+
1820
+ return knowledge_row
1821
+
1822
+ except Exception as e:
1823
+ log_error(f"Error upserting knowledge content: {e}")
1824
+ raise e
1825
+
1826
+ # -- Eval methods --
1827
+
1828
+ async def create_eval_run(self, eval_run: EvalRunRecord) -> Optional[EvalRunRecord]:
1829
+ """Create an EvalRunRecord in the database.
1830
+
1831
+ Args:
1832
+ eval_run (EvalRunRecord): The eval run to create.
1833
+
1834
+ Returns:
1835
+ Optional[EvalRunRecord]: The created eval run, or None if the operation fails.
1836
+
1837
+ Raises:
1838
+ Exception: If an error occurs during creation.
1839
+ """
1840
+ try:
1841
+ table = await self._get_table(table_type="evals")
1842
+ if table is None:
1843
+ return None
1844
+
1845
+ async with self.async_session_factory() as sess, sess.begin():
1846
+ current_time = int(time.time())
1847
+ stmt = sqlite.insert(table).values(
1848
+ {
1849
+ "created_at": current_time,
1850
+ "updated_at": current_time,
1851
+ **eval_run.model_dump(),
1852
+ }
1853
+ )
1854
+ await sess.execute(stmt)
1855
+
1856
+ log_debug(f"Created eval run with id '{eval_run.run_id}'")
1857
+
1858
+ return eval_run
1859
+
1860
+ except Exception as e:
1861
+ log_error(f"Error creating eval run: {e}")
1862
+ raise e
1863
+
1864
+ async def delete_eval_run(self, eval_run_id: str) -> None:
1865
+ """Delete an eval run from the database.
1866
+
1867
+ Args:
1868
+ eval_run_id (str): The ID of the eval run to delete.
1869
+ """
1870
+ try:
1871
+ table = await self._get_table(table_type="evals")
1872
+ if table is None:
1873
+ return
1874
+
1875
+ async with self.async_session_factory() as sess, sess.begin():
1876
+ stmt = table.delete().where(table.c.run_id == eval_run_id)
1877
+ result = await sess.execute(stmt)
1878
+ if result.rowcount == 0: # type: ignore
1879
+ log_warning(f"No eval run found with ID: {eval_run_id}")
1880
+ else:
1881
+ log_debug(f"Deleted eval run with ID: {eval_run_id}")
1882
+
1883
+ except Exception as e:
1884
+ log_error(f"Error deleting eval run {eval_run_id}: {e}")
1885
+ raise e
1886
+
1887
+ async def delete_eval_runs(self, eval_run_ids: List[str]) -> None:
1888
+ """Delete multiple eval runs from the database.
1889
+
1890
+ Args:
1891
+ eval_run_ids (List[str]): List of eval run IDs to delete.
1892
+ """
1893
+ try:
1894
+ table = await self._get_table(table_type="evals")
1895
+ if table is None:
1896
+ return
1897
+
1898
+ async with self.async_session_factory() as sess, sess.begin():
1899
+ stmt = table.delete().where(table.c.run_id.in_(eval_run_ids))
1900
+ result = await sess.execute(stmt)
1901
+ if result.rowcount == 0: # type: ignore
1902
+ log_debug(f"No eval runs found with IDs: {eval_run_ids}")
1903
+ else:
1904
+ log_debug(f"Deleted {result.rowcount} eval runs") # type: ignore
1905
+
1906
+ except Exception as e:
1907
+ log_error(f"Error deleting eval runs {eval_run_ids}: {e}")
1908
+ raise e
1909
+
1910
+ async def get_eval_run(
1911
+ self, eval_run_id: str, deserialize: Optional[bool] = True
1912
+ ) -> Optional[Union[EvalRunRecord, Dict[str, Any]]]:
1913
+ """Get an eval run from the database.
1914
+
1915
+ Args:
1916
+ eval_run_id (str): The ID of the eval run to get.
1917
+ deserialize (Optional[bool]): Whether to serialize the eval run. Defaults to True.
1918
+
1919
+ Returns:
1920
+ Optional[Union[EvalRunRecord, Dict[str, Any]]]:
1921
+ - When deserialize=True: EvalRunRecord object
1922
+ - When deserialize=False: EvalRun dictionary
1923
+
1924
+ Raises:
1925
+ Exception: If an error occurs during retrieval.
1926
+ """
1927
+ try:
1928
+ table = await self._get_table(table_type="evals")
1929
+ if table is None:
1930
+ return None
1931
+
1932
+ async with self.async_session_factory() as sess, sess.begin():
1933
+ stmt = select(table).where(table.c.run_id == eval_run_id)
1934
+ result = (await sess.execute(stmt)).fetchone()
1935
+ if result is None:
1936
+ return None
1937
+
1938
+ eval_run_raw = dict(result._mapping)
1939
+ if not eval_run_raw or not deserialize:
1940
+ return eval_run_raw
1941
+
1942
+ return EvalRunRecord.model_validate(eval_run_raw)
1943
+
1944
+ except Exception as e:
1945
+ log_error(f"Exception getting eval run {eval_run_id}: {e}")
1946
+ raise e
1947
+
1948
+ async def get_eval_runs(
1949
+ self,
1950
+ limit: Optional[int] = None,
1951
+ page: Optional[int] = None,
1952
+ sort_by: Optional[str] = None,
1953
+ sort_order: Optional[str] = None,
1954
+ agent_id: Optional[str] = None,
1955
+ team_id: Optional[str] = None,
1956
+ workflow_id: Optional[str] = None,
1957
+ model_id: Optional[str] = None,
1958
+ filter_type: Optional[EvalFilterType] = None,
1959
+ eval_type: Optional[List[EvalType]] = None,
1960
+ deserialize: Optional[bool] = True,
1961
+ ) -> Union[List[EvalRunRecord], Tuple[List[Dict[str, Any]], int]]:
1962
+ """Get all eval runs from the database.
1963
+
1964
+ Args:
1965
+ limit (Optional[int]): The maximum number of eval runs to return.
1966
+ page (Optional[int]): The page number.
1967
+ sort_by (Optional[str]): The column to sort by.
1968
+ sort_order (Optional[str]): The order to sort by.
1969
+ agent_id (Optional[str]): The ID of the agent to filter by.
1970
+ team_id (Optional[str]): The ID of the team to filter by.
1971
+ workflow_id (Optional[str]): The ID of the workflow to filter by.
1972
+ model_id (Optional[str]): The ID of the model to filter by.
1973
+ eval_type (Optional[List[EvalType]]): The type(s) of eval to filter by.
1974
+ filter_type (Optional[EvalFilterType]): Filter by component type (agent, team, workflow).
1975
+ deserialize (Optional[bool]): Whether to serialize the eval runs. Defaults to True.
1976
+ create_table_if_not_found (Optional[bool]): Whether to create the table if it doesn't exist.
1977
+
1978
+ Returns:
1979
+ Union[List[EvalRunRecord], Tuple[List[Dict[str, Any]], int]]:
1980
+ - When deserialize=True: List of EvalRunRecord objects
1981
+ - When deserialize=False: List of EvalRun dictionaries and total count
1982
+
1983
+ Raises:
1984
+ Exception: If an error occurs during retrieval.
1985
+ """
1986
+ try:
1987
+ table = await self._get_table(table_type="evals")
1988
+ if table is None:
1989
+ return [] if deserialize else ([], 0)
1990
+
1991
+ async with self.async_session_factory() as sess, sess.begin():
1992
+ stmt = select(table)
1993
+
1994
+ # Filtering
1995
+ if agent_id is not None:
1996
+ stmt = stmt.where(table.c.agent_id == agent_id)
1997
+ if team_id is not None:
1998
+ stmt = stmt.where(table.c.team_id == team_id)
1999
+ if workflow_id is not None:
2000
+ stmt = stmt.where(table.c.workflow_id == workflow_id)
2001
+ if model_id is not None:
2002
+ stmt = stmt.where(table.c.model_id == model_id)
2003
+ if eval_type is not None and len(eval_type) > 0:
2004
+ stmt = stmt.where(table.c.eval_type.in_(eval_type))
2005
+ if filter_type is not None:
2006
+ if filter_type == EvalFilterType.AGENT:
2007
+ stmt = stmt.where(table.c.agent_id.is_not(None))
2008
+ elif filter_type == EvalFilterType.TEAM:
2009
+ stmt = stmt.where(table.c.team_id.is_not(None))
2010
+ elif filter_type == EvalFilterType.WORKFLOW:
2011
+ stmt = stmt.where(table.c.workflow_id.is_not(None))
2012
+
2013
+ # Get total count after applying filtering
2014
+ count_stmt = select(func.count()).select_from(stmt.alias())
2015
+ total_count = (await sess.execute(count_stmt)).scalar() or 0
2016
+
2017
+ # Sorting - apply default sort by created_at desc if no sort parameters provided
2018
+ if sort_by is None:
2019
+ stmt = stmt.order_by(table.c.created_at.desc())
2020
+ else:
2021
+ stmt = apply_sorting(stmt, table, sort_by, sort_order)
2022
+ # Paginating
2023
+ if limit is not None:
2024
+ stmt = stmt.limit(limit)
2025
+ if page is not None:
2026
+ stmt = stmt.offset((page - 1) * limit)
2027
+
2028
+ result = (await sess.execute(stmt)).fetchall()
2029
+ if not result:
2030
+ return [] if deserialize else ([], 0)
2031
+
2032
+ eval_runs_raw = [dict(row._mapping) for row in result]
2033
+ if not deserialize:
2034
+ return eval_runs_raw, total_count
2035
+
2036
+ return [EvalRunRecord.model_validate(row) for row in eval_runs_raw]
2037
+
2038
+ except Exception as e:
2039
+ log_error(f"Exception getting eval runs: {e}")
2040
+ raise e
2041
+
2042
+ async def rename_eval_run(
2043
+ self, eval_run_id: str, name: str, deserialize: Optional[bool] = True
2044
+ ) -> Optional[Union[EvalRunRecord, Dict[str, Any]]]:
2045
+ """Upsert the name of an eval run in the database, returning raw dictionary.
2046
+
2047
+ Args:
2048
+ eval_run_id (str): The ID of the eval run to update.
2049
+ name (str): The new name of the eval run.
2050
+ deserialize (Optional[bool]): Whether to serialize the eval run. Defaults to True.
2051
+
2052
+ Returns:
2053
+ Optional[Union[EvalRunRecord, Dict[str, Any]]]:
2054
+ - When deserialize=True: EvalRunRecord object
2055
+ - When deserialize=False: EvalRun dictionary
2056
+
2057
+ Raises:
2058
+ Exception: If an error occurs during update.
2059
+ """
2060
+ try:
2061
+ table = await self._get_table(table_type="evals")
2062
+ if table is None:
2063
+ return None
2064
+
2065
+ async with self.async_session_factory() as sess, sess.begin():
2066
+ stmt = (
2067
+ table.update().where(table.c.run_id == eval_run_id).values(name=name, updated_at=int(time.time()))
2068
+ )
2069
+ await sess.execute(stmt)
2070
+
2071
+ eval_run_raw = await self.get_eval_run(eval_run_id=eval_run_id, deserialize=deserialize)
2072
+
2073
+ log_debug(f"Renamed eval run with id '{eval_run_id}' to '{name}'")
2074
+
2075
+ if not eval_run_raw or not deserialize:
2076
+ return eval_run_raw
2077
+
2078
+ return EvalRunRecord.model_validate(eval_run_raw)
2079
+
2080
+ except Exception as e:
2081
+ log_error(f"Error renaming eval run {eval_run_id}: {e}")
2082
+ raise e
2083
+
2084
+ # -- Migrations --
2085
+
2086
+ async def migrate_table_from_v1_to_v2(self, v1_db_schema: str, v1_table_name: str, v1_table_type: str):
2087
+ """Migrate all content in the given table to the right v2 table"""
2088
+
2089
+ from agno.db.migrations.v1_to_v2 import (
2090
+ get_all_table_content,
2091
+ parse_agent_sessions,
2092
+ parse_memories,
2093
+ parse_team_sessions,
2094
+ parse_workflow_sessions,
2095
+ )
2096
+
2097
+ # Get all content from the old table
2098
+ old_content: list[dict[str, Any]] = get_all_table_content(
2099
+ db=self,
2100
+ db_schema=v1_db_schema,
2101
+ table_name=v1_table_name,
2102
+ )
2103
+ if not old_content:
2104
+ log_info(f"No content to migrate from table {v1_table_name}")
2105
+ return
2106
+
2107
+ # Parse the content into the new format
2108
+ memories: List[UserMemory] = []
2109
+ sessions: Sequence[Union[AgentSession, TeamSession, WorkflowSession]] = []
2110
+ if v1_table_type == "agent_sessions":
2111
+ sessions = parse_agent_sessions(old_content)
2112
+ elif v1_table_type == "team_sessions":
2113
+ sessions = parse_team_sessions(old_content)
2114
+ elif v1_table_type == "workflow_sessions":
2115
+ sessions = parse_workflow_sessions(old_content)
2116
+ elif v1_table_type == "memories":
2117
+ memories = parse_memories(old_content)
2118
+ else:
2119
+ raise ValueError(f"Invalid table type: {v1_table_type}")
2120
+
2121
+ # Insert the new content into the new table
2122
+ if v1_table_type == "agent_sessions":
2123
+ for session in sessions:
2124
+ await self.upsert_session(session)
2125
+ log_info(f"Migrated {len(sessions)} Agent sessions to table: {self.session_table_name}")
2126
+
2127
+ elif v1_table_type == "team_sessions":
2128
+ for session in sessions:
2129
+ await self.upsert_session(session)
2130
+ log_info(f"Migrated {len(sessions)} Team sessions to table: {self.session_table_name}")
2131
+
2132
+ elif v1_table_type == "workflow_sessions":
2133
+ for session in sessions:
2134
+ await self.upsert_session(session)
2135
+ log_info(f"Migrated {len(sessions)} Workflow sessions to table: {self.session_table_name}")
2136
+
2137
+ elif v1_table_type == "memories":
2138
+ for memory in memories:
2139
+ await self.upsert_user_memory(memory)
2140
+ log_info(f"Migrated {len(memories)} memories to table: {self.memory_table}")
2141
+
2142
+ # -- Culture methods --
2143
+
2144
+ async def clear_cultural_knowledge(self) -> None:
2145
+ """Delete all cultural artifacts from the database.
2146
+
2147
+ Raises:
2148
+ Exception: If an error occurs during deletion.
2149
+ """
2150
+ try:
2151
+ table = await self._get_table(table_type="culture")
2152
+ if table is None:
2153
+ return
2154
+
2155
+ async with self.async_session_factory() as sess, sess.begin():
2156
+ await sess.execute(table.delete())
2157
+
2158
+ except Exception as e:
2159
+ from agno.utils.log import log_warning
2160
+
2161
+ log_warning(f"Exception deleting all cultural artifacts: {e}")
2162
+ raise e
2163
+
2164
+ async def delete_cultural_knowledge(self, id: str) -> None:
2165
+ """Delete a cultural artifact from the database.
2166
+
2167
+ Args:
2168
+ id (str): The ID of the cultural artifact to delete.
2169
+
2170
+ Raises:
2171
+ Exception: If an error occurs during deletion.
2172
+ """
2173
+ try:
2174
+ table = await self._get_table(table_type="culture")
2175
+ if table is None:
2176
+ return
2177
+
2178
+ async with self.async_session_factory() as sess, sess.begin():
2179
+ delete_stmt = table.delete().where(table.c.id == id)
2180
+ result = await sess.execute(delete_stmt)
2181
+
2182
+ success = result.rowcount > 0 # type: ignore
2183
+ if success:
2184
+ log_debug(f"Successfully deleted cultural artifact id: {id}")
2185
+ else:
2186
+ log_debug(f"No cultural artifact found with id: {id}")
2187
+
2188
+ except Exception as e:
2189
+ log_error(f"Error deleting cultural artifact: {e}")
2190
+ raise e
2191
+
2192
+ async def get_cultural_knowledge(
2193
+ self, id: str, deserialize: Optional[bool] = True
2194
+ ) -> Optional[Union[CulturalKnowledge, Dict[str, Any]]]:
2195
+ """Get a cultural artifact from the database.
2196
+
2197
+ Args:
2198
+ id (str): The ID of the cultural artifact to get.
2199
+ deserialize (Optional[bool]): Whether to serialize the cultural artifact. Defaults to True.
2200
+
2201
+ Returns:
2202
+ Optional[CulturalKnowledge]: The cultural artifact, or None if it doesn't exist.
2203
+
2204
+ Raises:
2205
+ Exception: If an error occurs during retrieval.
2206
+ """
2207
+ try:
2208
+ table = await self._get_table(table_type="culture")
2209
+ if table is None:
2210
+ return None
2211
+
2212
+ async with self.async_session_factory() as sess, sess.begin():
2213
+ stmt = select(table).where(table.c.id == id)
2214
+ result = (await sess.execute(stmt)).fetchone()
2215
+ if result is None:
2216
+ return None
2217
+
2218
+ db_row = dict(result._mapping)
2219
+ if not db_row or not deserialize:
2220
+ return db_row
2221
+
2222
+ return deserialize_cultural_knowledge_from_db(db_row)
2223
+
2224
+ except Exception as e:
2225
+ log_error(f"Exception reading from cultural artifacts table: {e}")
2226
+ raise e
2227
+
2228
+ async def get_all_cultural_knowledge(
2229
+ self,
2230
+ name: Optional[str] = None,
2231
+ agent_id: Optional[str] = None,
2232
+ team_id: Optional[str] = None,
2233
+ limit: Optional[int] = None,
2234
+ page: Optional[int] = None,
2235
+ sort_by: Optional[str] = None,
2236
+ sort_order: Optional[str] = None,
2237
+ deserialize: Optional[bool] = True,
2238
+ ) -> Union[List[CulturalKnowledge], Tuple[List[Dict[str, Any]], int]]:
2239
+ """Get all cultural artifacts from the database as CulturalNotion objects.
2240
+
2241
+ Args:
2242
+ name (Optional[str]): The name of the cultural artifact to filter by.
2243
+ agent_id (Optional[str]): The ID of the agent to filter by.
2244
+ team_id (Optional[str]): The ID of the team to filter by.
2245
+ limit (Optional[int]): The maximum number of cultural artifacts to return.
2246
+ page (Optional[int]): The page number.
2247
+ sort_by (Optional[str]): The column to sort by.
2248
+ sort_order (Optional[str]): The order to sort by.
2249
+ deserialize (Optional[bool]): Whether to serialize the cultural artifacts. Defaults to True.
2250
+
2251
+ Returns:
2252
+ Union[List[CulturalKnowledge], Tuple[List[Dict[str, Any]], int]]:
2253
+ - When deserialize=True: List of CulturalNotion objects
2254
+ - When deserialize=False: List of CulturalNotion dictionaries and total count
2255
+
2256
+ Raises:
2257
+ Exception: If an error occurs during retrieval.
2258
+ """
2259
+ try:
2260
+ table = await self._get_table(table_type="culture")
2261
+ if table is None:
2262
+ return [] if deserialize else ([], 0)
2263
+
2264
+ async with self.async_session_factory() as sess, sess.begin():
2265
+ stmt = select(table)
2266
+
2267
+ # Filtering
2268
+ if name is not None:
2269
+ stmt = stmt.where(table.c.name == name)
2270
+ if agent_id is not None:
2271
+ stmt = stmt.where(table.c.agent_id == agent_id)
2272
+ if team_id is not None:
2273
+ stmt = stmt.where(table.c.team_id == team_id)
2274
+
2275
+ # Get total count after applying filtering
2276
+ count_stmt = select(func.count()).select_from(stmt.alias())
2277
+ total_count = (await sess.execute(count_stmt)).scalar() or 0
2278
+
2279
+ # Sorting
2280
+ stmt = apply_sorting(stmt, table, sort_by, sort_order)
2281
+ # Paginating
2282
+ if limit is not None:
2283
+ stmt = stmt.limit(limit)
2284
+ if page is not None:
2285
+ stmt = stmt.offset((page - 1) * limit)
2286
+
2287
+ result = (await sess.execute(stmt)).fetchall()
2288
+ if not result:
2289
+ return [] if deserialize else ([], 0)
2290
+
2291
+ db_rows = [dict(record._mapping) for record in result]
2292
+
2293
+ if not deserialize:
2294
+ return db_rows, total_count
2295
+
2296
+ return [deserialize_cultural_knowledge_from_db(row) for row in db_rows]
2297
+
2298
+ except Exception as e:
2299
+ log_error(f"Error reading from cultural artifacts table: {e}")
2300
+ raise e
2301
+
2302
+ async def upsert_cultural_knowledge(
2303
+ self, cultural_knowledge: CulturalKnowledge, deserialize: Optional[bool] = True
2304
+ ) -> Optional[Union[CulturalKnowledge, Dict[str, Any]]]:
2305
+ """Upsert a cultural artifact into the database.
2306
+
2307
+ Args:
2308
+ cultural_knowledge (CulturalKnowledge): The cultural artifact to upsert.
2309
+ deserialize (Optional[bool]): Whether to serialize the cultural artifact. Defaults to True.
2310
+
2311
+ Returns:
2312
+ Optional[Union[CulturalNotion, Dict[str, Any]]]:
2313
+ - When deserialize=True: CulturalNotion object
2314
+ - When deserialize=False: CulturalNotion dictionary
2315
+
2316
+ Raises:
2317
+ Exception: If an error occurs during upsert.
2318
+ """
2319
+ try:
2320
+ table = await self._get_table(table_type="culture")
2321
+ if table is None:
2322
+ return None
2323
+
2324
+ if cultural_knowledge.id is None:
2325
+ cultural_knowledge.id = str(uuid4())
2326
+
2327
+ # Serialize content, categories, and notes into a JSON string for DB storage (SQLite requires strings)
2328
+ content_json_str = serialize_cultural_knowledge_for_db(cultural_knowledge)
2329
+
2330
+ async with self.async_session_factory() as sess, sess.begin():
2331
+ stmt = sqlite.insert(table).values(
2332
+ id=cultural_knowledge.id,
2333
+ name=cultural_knowledge.name,
2334
+ summary=cultural_knowledge.summary,
2335
+ content=content_json_str,
2336
+ metadata=cultural_knowledge.metadata,
2337
+ input=cultural_knowledge.input,
2338
+ created_at=cultural_knowledge.created_at,
2339
+ updated_at=int(time.time()),
2340
+ agent_id=cultural_knowledge.agent_id,
2341
+ team_id=cultural_knowledge.team_id,
2342
+ )
2343
+ stmt = stmt.on_conflict_do_update( # type: ignore
2344
+ index_elements=["id"],
2345
+ set_=dict(
2346
+ name=cultural_knowledge.name,
2347
+ summary=cultural_knowledge.summary,
2348
+ content=content_json_str,
2349
+ metadata=cultural_knowledge.metadata,
2350
+ input=cultural_knowledge.input,
2351
+ updated_at=int(time.time()),
2352
+ agent_id=cultural_knowledge.agent_id,
2353
+ team_id=cultural_knowledge.team_id,
2354
+ ),
2355
+ ).returning(table)
2356
+
2357
+ result = await sess.execute(stmt)
2358
+ row = result.fetchone()
2359
+
2360
+ if row is None:
2361
+ return None
2362
+
2363
+ db_row: Dict[str, Any] = dict(row._mapping)
2364
+ if not db_row or not deserialize:
2365
+ return db_row
2366
+
2367
+ return deserialize_cultural_knowledge_from_db(db_row)
2368
+
2369
+ except Exception as e:
2370
+ log_error(f"Error upserting cultural knowledge: {e}")
2371
+ raise e