agno 1.8.1__py3-none-any.whl → 2.0.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 (590) hide show
  1. agno/__init__.py +8 -0
  2. agno/agent/__init__.py +19 -27
  3. agno/agent/agent.py +3143 -4170
  4. agno/api/agent.py +11 -67
  5. agno/api/api.py +5 -46
  6. agno/api/evals.py +8 -19
  7. agno/api/os.py +17 -0
  8. agno/api/routes.py +6 -41
  9. agno/api/schemas/__init__.py +9 -0
  10. agno/api/schemas/agent.py +5 -21
  11. agno/api/schemas/evals.py +7 -16
  12. agno/api/schemas/os.py +14 -0
  13. agno/api/schemas/team.py +5 -21
  14. agno/api/schemas/utils.py +21 -0
  15. agno/api/schemas/workflows.py +11 -7
  16. agno/api/settings.py +53 -0
  17. agno/api/team.py +11 -66
  18. agno/api/workflow.py +28 -0
  19. agno/cloud/aws/base.py +214 -0
  20. agno/cloud/aws/s3/__init__.py +2 -0
  21. agno/cloud/aws/s3/api_client.py +43 -0
  22. agno/cloud/aws/s3/bucket.py +195 -0
  23. agno/cloud/aws/s3/object.py +57 -0
  24. agno/db/__init__.py +24 -0
  25. agno/db/base.py +245 -0
  26. agno/db/dynamo/__init__.py +3 -0
  27. agno/db/dynamo/dynamo.py +1743 -0
  28. agno/db/dynamo/schemas.py +278 -0
  29. agno/db/dynamo/utils.py +684 -0
  30. agno/db/firestore/__init__.py +3 -0
  31. agno/db/firestore/firestore.py +1432 -0
  32. agno/db/firestore/schemas.py +130 -0
  33. agno/db/firestore/utils.py +278 -0
  34. agno/db/gcs_json/__init__.py +3 -0
  35. agno/db/gcs_json/gcs_json_db.py +1001 -0
  36. agno/db/gcs_json/utils.py +194 -0
  37. agno/db/in_memory/__init__.py +3 -0
  38. agno/db/in_memory/in_memory_db.py +882 -0
  39. agno/db/in_memory/utils.py +172 -0
  40. agno/db/json/__init__.py +3 -0
  41. agno/db/json/json_db.py +1045 -0
  42. agno/db/json/utils.py +196 -0
  43. agno/db/migrations/v1_to_v2.py +162 -0
  44. agno/db/mongo/__init__.py +3 -0
  45. agno/db/mongo/mongo.py +1416 -0
  46. agno/db/mongo/schemas.py +77 -0
  47. agno/db/mongo/utils.py +204 -0
  48. agno/db/mysql/__init__.py +3 -0
  49. agno/db/mysql/mysql.py +1719 -0
  50. agno/db/mysql/schemas.py +124 -0
  51. agno/db/mysql/utils.py +297 -0
  52. agno/db/postgres/__init__.py +3 -0
  53. agno/db/postgres/postgres.py +1710 -0
  54. agno/db/postgres/schemas.py +124 -0
  55. agno/db/postgres/utils.py +280 -0
  56. agno/db/redis/__init__.py +3 -0
  57. agno/db/redis/redis.py +1367 -0
  58. agno/db/redis/schemas.py +109 -0
  59. agno/db/redis/utils.py +288 -0
  60. agno/db/schemas/__init__.py +3 -0
  61. agno/db/schemas/evals.py +33 -0
  62. agno/db/schemas/knowledge.py +40 -0
  63. agno/db/schemas/memory.py +46 -0
  64. agno/db/singlestore/__init__.py +3 -0
  65. agno/db/singlestore/schemas.py +116 -0
  66. agno/db/singlestore/singlestore.py +1712 -0
  67. agno/db/singlestore/utils.py +326 -0
  68. agno/db/sqlite/__init__.py +3 -0
  69. agno/db/sqlite/schemas.py +119 -0
  70. agno/db/sqlite/sqlite.py +1676 -0
  71. agno/db/sqlite/utils.py +268 -0
  72. agno/db/utils.py +88 -0
  73. agno/eval/__init__.py +14 -0
  74. agno/eval/accuracy.py +154 -48
  75. agno/eval/performance.py +88 -23
  76. agno/eval/reliability.py +73 -20
  77. agno/eval/utils.py +23 -13
  78. agno/integrations/discord/__init__.py +3 -0
  79. agno/{app → integrations}/discord/client.py +15 -11
  80. agno/knowledge/__init__.py +2 -2
  81. agno/{document → knowledge}/chunking/agentic.py +2 -2
  82. agno/{document → knowledge}/chunking/document.py +2 -2
  83. agno/{document → knowledge}/chunking/fixed.py +3 -3
  84. agno/{document → knowledge}/chunking/markdown.py +2 -2
  85. agno/{document → knowledge}/chunking/recursive.py +2 -2
  86. agno/{document → knowledge}/chunking/row.py +2 -2
  87. agno/knowledge/chunking/semantic.py +59 -0
  88. agno/knowledge/chunking/strategy.py +121 -0
  89. agno/knowledge/content.py +74 -0
  90. agno/knowledge/document/__init__.py +5 -0
  91. agno/{document → knowledge/document}/base.py +12 -2
  92. agno/knowledge/embedder/__init__.py +5 -0
  93. agno/{embedder → knowledge/embedder}/aws_bedrock.py +127 -1
  94. agno/{embedder → knowledge/embedder}/azure_openai.py +65 -1
  95. agno/{embedder → knowledge/embedder}/base.py +6 -0
  96. agno/{embedder → knowledge/embedder}/cohere.py +72 -1
  97. agno/{embedder → knowledge/embedder}/fastembed.py +17 -1
  98. agno/{embedder → knowledge/embedder}/fireworks.py +1 -1
  99. agno/{embedder → knowledge/embedder}/google.py +74 -1
  100. agno/{embedder → knowledge/embedder}/huggingface.py +36 -2
  101. agno/{embedder → knowledge/embedder}/jina.py +48 -2
  102. agno/knowledge/embedder/langdb.py +22 -0
  103. agno/knowledge/embedder/mistral.py +139 -0
  104. agno/{embedder → knowledge/embedder}/nebius.py +1 -1
  105. agno/{embedder → knowledge/embedder}/ollama.py +54 -3
  106. agno/knowledge/embedder/openai.py +223 -0
  107. agno/{embedder → knowledge/embedder}/sentence_transformer.py +16 -1
  108. agno/{embedder → knowledge/embedder}/together.py +1 -1
  109. agno/{embedder → knowledge/embedder}/voyageai.py +49 -1
  110. agno/knowledge/knowledge.py +1551 -0
  111. agno/knowledge/reader/__init__.py +7 -0
  112. agno/{document → knowledge}/reader/arxiv_reader.py +32 -4
  113. agno/knowledge/reader/base.py +88 -0
  114. agno/{document → knowledge}/reader/csv_reader.py +47 -65
  115. agno/knowledge/reader/docx_reader.py +83 -0
  116. agno/{document → knowledge}/reader/firecrawl_reader.py +42 -21
  117. agno/{document → knowledge}/reader/json_reader.py +30 -9
  118. agno/{document → knowledge}/reader/markdown_reader.py +58 -9
  119. agno/{document → knowledge}/reader/pdf_reader.py +71 -126
  120. agno/knowledge/reader/reader_factory.py +268 -0
  121. agno/knowledge/reader/s3_reader.py +101 -0
  122. agno/{document → knowledge}/reader/text_reader.py +31 -10
  123. agno/knowledge/reader/url_reader.py +128 -0
  124. agno/knowledge/reader/web_search_reader.py +366 -0
  125. agno/{document → knowledge}/reader/website_reader.py +37 -10
  126. agno/knowledge/reader/wikipedia_reader.py +59 -0
  127. agno/knowledge/reader/youtube_reader.py +78 -0
  128. agno/knowledge/remote_content/remote_content.py +88 -0
  129. agno/{reranker → knowledge/reranker}/base.py +1 -1
  130. agno/{reranker → knowledge/reranker}/cohere.py +2 -2
  131. agno/{reranker → knowledge/reranker}/infinity.py +2 -2
  132. agno/{reranker → knowledge/reranker}/sentence_transformer.py +2 -2
  133. agno/knowledge/types.py +30 -0
  134. agno/knowledge/utils.py +169 -0
  135. agno/media.py +269 -268
  136. agno/memory/__init__.py +2 -10
  137. agno/memory/manager.py +1003 -148
  138. agno/models/aimlapi/__init__.py +2 -2
  139. agno/models/aimlapi/aimlapi.py +6 -6
  140. agno/models/anthropic/claude.py +131 -131
  141. agno/models/aws/bedrock.py +110 -182
  142. agno/models/aws/claude.py +64 -18
  143. agno/models/azure/ai_foundry.py +73 -23
  144. agno/models/base.py +346 -290
  145. agno/models/cerebras/cerebras.py +84 -27
  146. agno/models/cohere/chat.py +106 -98
  147. agno/models/google/gemini.py +105 -46
  148. agno/models/groq/groq.py +97 -35
  149. agno/models/huggingface/huggingface.py +92 -27
  150. agno/models/ibm/watsonx.py +72 -13
  151. agno/models/litellm/chat.py +85 -13
  152. agno/models/message.py +46 -151
  153. agno/models/meta/llama.py +85 -49
  154. agno/models/metrics.py +120 -0
  155. agno/models/mistral/mistral.py +90 -21
  156. agno/models/ollama/__init__.py +0 -2
  157. agno/models/ollama/chat.py +85 -47
  158. agno/models/openai/chat.py +154 -37
  159. agno/models/openai/responses.py +178 -105
  160. agno/models/perplexity/perplexity.py +26 -2
  161. agno/models/portkey/portkey.py +0 -7
  162. agno/models/response.py +15 -9
  163. agno/models/utils.py +20 -0
  164. agno/models/vercel/__init__.py +2 -2
  165. agno/models/vercel/v0.py +1 -1
  166. agno/models/vllm/__init__.py +2 -2
  167. agno/models/vllm/vllm.py +3 -3
  168. agno/models/xai/xai.py +10 -10
  169. agno/os/__init__.py +3 -0
  170. agno/os/app.py +497 -0
  171. agno/os/auth.py +47 -0
  172. agno/os/config.py +103 -0
  173. agno/os/interfaces/agui/__init__.py +3 -0
  174. agno/os/interfaces/agui/agui.py +31 -0
  175. agno/{app/agui/async_router.py → os/interfaces/agui/router.py} +16 -16
  176. agno/{app → os/interfaces}/agui/utils.py +77 -33
  177. agno/os/interfaces/base.py +21 -0
  178. agno/os/interfaces/slack/__init__.py +3 -0
  179. agno/{app/slack/async_router.py → os/interfaces/slack/router.py} +3 -5
  180. agno/os/interfaces/slack/slack.py +32 -0
  181. agno/os/interfaces/whatsapp/__init__.py +3 -0
  182. agno/{app/whatsapp/async_router.py → os/interfaces/whatsapp/router.py} +4 -7
  183. agno/os/interfaces/whatsapp/whatsapp.py +29 -0
  184. agno/os/mcp.py +235 -0
  185. agno/os/router.py +1400 -0
  186. agno/os/routers/__init__.py +3 -0
  187. agno/os/routers/evals/__init__.py +3 -0
  188. agno/os/routers/evals/evals.py +393 -0
  189. agno/os/routers/evals/schemas.py +142 -0
  190. agno/os/routers/evals/utils.py +161 -0
  191. agno/os/routers/knowledge/__init__.py +3 -0
  192. agno/os/routers/knowledge/knowledge.py +850 -0
  193. agno/os/routers/knowledge/schemas.py +118 -0
  194. agno/os/routers/memory/__init__.py +3 -0
  195. agno/os/routers/memory/memory.py +410 -0
  196. agno/os/routers/memory/schemas.py +58 -0
  197. agno/os/routers/metrics/__init__.py +3 -0
  198. agno/os/routers/metrics/metrics.py +178 -0
  199. agno/os/routers/metrics/schemas.py +47 -0
  200. agno/os/routers/session/__init__.py +3 -0
  201. agno/os/routers/session/session.py +536 -0
  202. agno/os/schema.py +945 -0
  203. agno/{app/playground → os}/settings.py +7 -15
  204. agno/os/utils.py +270 -0
  205. agno/reasoning/azure_ai_foundry.py +4 -4
  206. agno/reasoning/deepseek.py +4 -4
  207. agno/reasoning/default.py +6 -11
  208. agno/reasoning/groq.py +4 -4
  209. agno/reasoning/helpers.py +4 -6
  210. agno/reasoning/ollama.py +4 -4
  211. agno/reasoning/openai.py +4 -4
  212. agno/run/agent.py +633 -0
  213. agno/run/base.py +53 -77
  214. agno/run/cancel.py +81 -0
  215. agno/run/team.py +243 -96
  216. agno/run/workflow.py +550 -12
  217. agno/session/__init__.py +10 -0
  218. agno/session/agent.py +244 -0
  219. agno/session/summary.py +225 -0
  220. agno/session/team.py +262 -0
  221. agno/{storage/session/v2 → session}/workflow.py +47 -24
  222. agno/team/__init__.py +15 -16
  223. agno/team/team.py +3260 -4824
  224. agno/tools/agentql.py +14 -5
  225. agno/tools/airflow.py +9 -4
  226. agno/tools/api.py +7 -3
  227. agno/tools/apify.py +2 -46
  228. agno/tools/arxiv.py +8 -3
  229. agno/tools/aws_lambda.py +7 -5
  230. agno/tools/aws_ses.py +7 -1
  231. agno/tools/baidusearch.py +4 -1
  232. agno/tools/bitbucket.py +4 -4
  233. agno/tools/brandfetch.py +14 -11
  234. agno/tools/bravesearch.py +4 -1
  235. agno/tools/brightdata.py +43 -23
  236. agno/tools/browserbase.py +13 -4
  237. agno/tools/calcom.py +12 -10
  238. agno/tools/calculator.py +10 -27
  239. agno/tools/cartesia.py +20 -17
  240. agno/tools/{clickup_tool.py → clickup.py} +12 -25
  241. agno/tools/confluence.py +8 -8
  242. agno/tools/crawl4ai.py +7 -1
  243. agno/tools/csv_toolkit.py +9 -8
  244. agno/tools/dalle.py +22 -12
  245. agno/tools/daytona.py +13 -16
  246. agno/tools/decorator.py +6 -3
  247. agno/tools/desi_vocal.py +17 -8
  248. agno/tools/discord.py +11 -8
  249. agno/tools/docker.py +30 -42
  250. agno/tools/duckdb.py +34 -53
  251. agno/tools/duckduckgo.py +8 -7
  252. agno/tools/e2b.py +62 -62
  253. agno/tools/eleven_labs.py +36 -29
  254. agno/tools/email.py +4 -1
  255. agno/tools/evm.py +7 -1
  256. agno/tools/exa.py +19 -14
  257. agno/tools/fal.py +30 -30
  258. agno/tools/file.py +9 -8
  259. agno/tools/financial_datasets.py +25 -44
  260. agno/tools/firecrawl.py +22 -22
  261. agno/tools/function.py +127 -18
  262. agno/tools/giphy.py +23 -11
  263. agno/tools/github.py +48 -126
  264. agno/tools/gmail.py +45 -61
  265. agno/tools/google_bigquery.py +7 -6
  266. agno/tools/google_maps.py +11 -26
  267. agno/tools/googlesearch.py +7 -2
  268. agno/tools/googlesheets.py +21 -17
  269. agno/tools/hackernews.py +9 -5
  270. agno/tools/jina.py +5 -4
  271. agno/tools/jira.py +18 -9
  272. agno/tools/knowledge.py +31 -32
  273. agno/tools/linear.py +19 -34
  274. agno/tools/linkup.py +5 -1
  275. agno/tools/local_file_system.py +8 -5
  276. agno/tools/lumalab.py +32 -20
  277. agno/tools/mcp.py +1 -2
  278. agno/tools/mem0.py +18 -12
  279. agno/tools/memori.py +14 -10
  280. agno/tools/mlx_transcribe.py +3 -2
  281. agno/tools/models/azure_openai.py +33 -15
  282. agno/tools/models/gemini.py +59 -32
  283. agno/tools/models/groq.py +30 -23
  284. agno/tools/models/nebius.py +28 -12
  285. agno/tools/models_labs.py +40 -16
  286. agno/tools/moviepy_video.py +7 -6
  287. agno/tools/neo4j.py +10 -8
  288. agno/tools/newspaper.py +7 -2
  289. agno/tools/newspaper4k.py +8 -3
  290. agno/tools/openai.py +58 -32
  291. agno/tools/openbb.py +12 -11
  292. agno/tools/opencv.py +63 -47
  293. agno/tools/openweather.py +14 -12
  294. agno/tools/pandas.py +11 -3
  295. agno/tools/postgres.py +4 -12
  296. agno/tools/pubmed.py +4 -1
  297. agno/tools/python.py +9 -22
  298. agno/tools/reasoning.py +35 -27
  299. agno/tools/reddit.py +11 -26
  300. agno/tools/replicate.py +55 -42
  301. agno/tools/resend.py +4 -1
  302. agno/tools/scrapegraph.py +15 -14
  303. agno/tools/searxng.py +10 -23
  304. agno/tools/serpapi.py +6 -3
  305. agno/tools/serper.py +13 -4
  306. agno/tools/shell.py +9 -2
  307. agno/tools/slack.py +12 -11
  308. agno/tools/sleep.py +3 -2
  309. agno/tools/spider.py +24 -4
  310. agno/tools/sql.py +7 -6
  311. agno/tools/tavily.py +6 -4
  312. agno/tools/telegram.py +12 -4
  313. agno/tools/todoist.py +11 -31
  314. agno/tools/toolkit.py +1 -1
  315. agno/tools/trafilatura.py +22 -6
  316. agno/tools/trello.py +9 -22
  317. agno/tools/twilio.py +10 -3
  318. agno/tools/user_control_flow.py +6 -1
  319. agno/tools/valyu.py +34 -5
  320. agno/tools/visualization.py +19 -28
  321. agno/tools/webbrowser.py +4 -3
  322. agno/tools/webex.py +11 -7
  323. agno/tools/website.py +15 -46
  324. agno/tools/webtools.py +12 -4
  325. agno/tools/whatsapp.py +5 -9
  326. agno/tools/wikipedia.py +20 -13
  327. agno/tools/x.py +14 -13
  328. agno/tools/yfinance.py +13 -40
  329. agno/tools/youtube.py +26 -20
  330. agno/tools/zendesk.py +7 -2
  331. agno/tools/zep.py +10 -7
  332. agno/tools/zoom.py +10 -9
  333. agno/utils/common.py +1 -19
  334. agno/utils/events.py +100 -123
  335. agno/utils/gemini.py +32 -2
  336. agno/utils/knowledge.py +29 -0
  337. agno/utils/log.py +54 -4
  338. agno/utils/mcp.py +68 -10
  339. agno/utils/media.py +39 -0
  340. agno/utils/message.py +12 -1
  341. agno/utils/models/aws_claude.py +1 -1
  342. agno/utils/models/claude.py +47 -4
  343. agno/utils/models/cohere.py +1 -1
  344. agno/utils/models/mistral.py +8 -7
  345. agno/utils/models/schema_utils.py +3 -3
  346. agno/utils/models/watsonx.py +1 -1
  347. agno/utils/openai.py +1 -1
  348. agno/utils/pprint.py +33 -32
  349. agno/utils/print_response/agent.py +779 -0
  350. agno/utils/print_response/team.py +1669 -0
  351. agno/utils/print_response/workflow.py +1451 -0
  352. agno/utils/prompts.py +14 -14
  353. agno/utils/reasoning.py +87 -0
  354. agno/utils/response.py +42 -42
  355. agno/utils/streamlit.py +481 -0
  356. agno/utils/string.py +8 -22
  357. agno/utils/team.py +50 -0
  358. agno/utils/timer.py +2 -2
  359. agno/vectordb/base.py +33 -21
  360. agno/vectordb/cassandra/cassandra.py +287 -23
  361. agno/vectordb/chroma/chromadb.py +482 -59
  362. agno/vectordb/clickhouse/clickhousedb.py +270 -63
  363. agno/vectordb/couchbase/couchbase.py +309 -29
  364. agno/vectordb/lancedb/lance_db.py +360 -21
  365. agno/vectordb/langchaindb/__init__.py +5 -0
  366. agno/vectordb/langchaindb/langchaindb.py +145 -0
  367. agno/vectordb/lightrag/__init__.py +5 -0
  368. agno/vectordb/lightrag/lightrag.py +374 -0
  369. agno/vectordb/llamaindex/llamaindexdb.py +127 -0
  370. agno/vectordb/milvus/milvus.py +242 -32
  371. agno/vectordb/mongodb/mongodb.py +200 -24
  372. agno/vectordb/pgvector/pgvector.py +319 -37
  373. agno/vectordb/pineconedb/pineconedb.py +221 -27
  374. agno/vectordb/qdrant/qdrant.py +334 -14
  375. agno/vectordb/singlestore/singlestore.py +286 -29
  376. agno/vectordb/surrealdb/surrealdb.py +187 -7
  377. agno/vectordb/upstashdb/upstashdb.py +342 -26
  378. agno/vectordb/weaviate/weaviate.py +227 -165
  379. agno/workflow/__init__.py +17 -13
  380. agno/workflow/{v2/condition.py → condition.py} +135 -32
  381. agno/workflow/{v2/loop.py → loop.py} +115 -28
  382. agno/workflow/{v2/parallel.py → parallel.py} +138 -108
  383. agno/workflow/{v2/router.py → router.py} +133 -32
  384. agno/workflow/{v2/step.py → step.py} +207 -49
  385. agno/workflow/{v2/steps.py → steps.py} +147 -66
  386. agno/workflow/types.py +482 -0
  387. agno/workflow/workflow.py +2410 -696
  388. agno-2.0.0.dist-info/METADATA +494 -0
  389. agno-2.0.0.dist-info/RECORD +515 -0
  390. agno-2.0.0.dist-info/licenses/LICENSE +201 -0
  391. agno/agent/metrics.py +0 -107
  392. agno/api/app.py +0 -35
  393. agno/api/playground.py +0 -92
  394. agno/api/schemas/app.py +0 -12
  395. agno/api/schemas/playground.py +0 -22
  396. agno/api/schemas/user.py +0 -35
  397. agno/api/schemas/workspace.py +0 -46
  398. agno/api/user.py +0 -160
  399. agno/api/workflows.py +0 -33
  400. agno/api/workspace.py +0 -175
  401. agno/app/agui/__init__.py +0 -3
  402. agno/app/agui/app.py +0 -17
  403. agno/app/agui/sync_router.py +0 -120
  404. agno/app/base.py +0 -186
  405. agno/app/discord/__init__.py +0 -3
  406. agno/app/fastapi/__init__.py +0 -3
  407. agno/app/fastapi/app.py +0 -107
  408. agno/app/fastapi/async_router.py +0 -457
  409. agno/app/fastapi/sync_router.py +0 -448
  410. agno/app/playground/app.py +0 -228
  411. agno/app/playground/async_router.py +0 -1050
  412. agno/app/playground/deploy.py +0 -249
  413. agno/app/playground/operator.py +0 -183
  414. agno/app/playground/schemas.py +0 -220
  415. agno/app/playground/serve.py +0 -55
  416. agno/app/playground/sync_router.py +0 -1042
  417. agno/app/playground/utils.py +0 -46
  418. agno/app/settings.py +0 -15
  419. agno/app/slack/__init__.py +0 -3
  420. agno/app/slack/app.py +0 -19
  421. agno/app/slack/sync_router.py +0 -92
  422. agno/app/utils.py +0 -54
  423. agno/app/whatsapp/__init__.py +0 -3
  424. agno/app/whatsapp/app.py +0 -15
  425. agno/app/whatsapp/sync_router.py +0 -197
  426. agno/cli/auth_server.py +0 -249
  427. agno/cli/config.py +0 -274
  428. agno/cli/console.py +0 -88
  429. agno/cli/credentials.py +0 -23
  430. agno/cli/entrypoint.py +0 -571
  431. agno/cli/operator.py +0 -357
  432. agno/cli/settings.py +0 -96
  433. agno/cli/ws/ws_cli.py +0 -817
  434. agno/constants.py +0 -13
  435. agno/document/__init__.py +0 -5
  436. agno/document/chunking/semantic.py +0 -45
  437. agno/document/chunking/strategy.py +0 -31
  438. agno/document/reader/__init__.py +0 -5
  439. agno/document/reader/base.py +0 -47
  440. agno/document/reader/docx_reader.py +0 -60
  441. agno/document/reader/gcs/pdf_reader.py +0 -44
  442. agno/document/reader/s3/pdf_reader.py +0 -59
  443. agno/document/reader/s3/text_reader.py +0 -63
  444. agno/document/reader/url_reader.py +0 -59
  445. agno/document/reader/youtube_reader.py +0 -58
  446. agno/embedder/__init__.py +0 -5
  447. agno/embedder/langdb.py +0 -80
  448. agno/embedder/mistral.py +0 -82
  449. agno/embedder/openai.py +0 -78
  450. agno/file/__init__.py +0 -5
  451. agno/file/file.py +0 -16
  452. agno/file/local/csv.py +0 -32
  453. agno/file/local/txt.py +0 -19
  454. agno/infra/app.py +0 -240
  455. agno/infra/base.py +0 -144
  456. agno/infra/context.py +0 -20
  457. agno/infra/db_app.py +0 -52
  458. agno/infra/resource.py +0 -205
  459. agno/infra/resources.py +0 -55
  460. agno/knowledge/agent.py +0 -702
  461. agno/knowledge/arxiv.py +0 -33
  462. agno/knowledge/combined.py +0 -36
  463. agno/knowledge/csv.py +0 -144
  464. agno/knowledge/csv_url.py +0 -124
  465. agno/knowledge/document.py +0 -223
  466. agno/knowledge/docx.py +0 -137
  467. agno/knowledge/firecrawl.py +0 -34
  468. agno/knowledge/gcs/__init__.py +0 -0
  469. agno/knowledge/gcs/base.py +0 -39
  470. agno/knowledge/gcs/pdf.py +0 -125
  471. agno/knowledge/json.py +0 -137
  472. agno/knowledge/langchain.py +0 -71
  473. agno/knowledge/light_rag.py +0 -273
  474. agno/knowledge/llamaindex.py +0 -66
  475. agno/knowledge/markdown.py +0 -154
  476. agno/knowledge/pdf.py +0 -164
  477. agno/knowledge/pdf_bytes.py +0 -42
  478. agno/knowledge/pdf_url.py +0 -148
  479. agno/knowledge/s3/__init__.py +0 -0
  480. agno/knowledge/s3/base.py +0 -64
  481. agno/knowledge/s3/pdf.py +0 -33
  482. agno/knowledge/s3/text.py +0 -34
  483. agno/knowledge/text.py +0 -141
  484. agno/knowledge/url.py +0 -46
  485. agno/knowledge/website.py +0 -179
  486. agno/knowledge/wikipedia.py +0 -32
  487. agno/knowledge/youtube.py +0 -35
  488. agno/memory/agent.py +0 -423
  489. agno/memory/classifier.py +0 -104
  490. agno/memory/db/__init__.py +0 -5
  491. agno/memory/db/base.py +0 -42
  492. agno/memory/db/mongodb.py +0 -189
  493. agno/memory/db/postgres.py +0 -203
  494. agno/memory/db/sqlite.py +0 -193
  495. agno/memory/memory.py +0 -22
  496. agno/memory/row.py +0 -36
  497. agno/memory/summarizer.py +0 -201
  498. agno/memory/summary.py +0 -19
  499. agno/memory/team.py +0 -415
  500. agno/memory/v2/__init__.py +0 -2
  501. agno/memory/v2/db/__init__.py +0 -1
  502. agno/memory/v2/db/base.py +0 -42
  503. agno/memory/v2/db/firestore.py +0 -339
  504. agno/memory/v2/db/mongodb.py +0 -196
  505. agno/memory/v2/db/postgres.py +0 -214
  506. agno/memory/v2/db/redis.py +0 -187
  507. agno/memory/v2/db/schema.py +0 -54
  508. agno/memory/v2/db/sqlite.py +0 -209
  509. agno/memory/v2/manager.py +0 -437
  510. agno/memory/v2/memory.py +0 -1097
  511. agno/memory/v2/schema.py +0 -55
  512. agno/memory/v2/summarizer.py +0 -215
  513. agno/memory/workflow.py +0 -38
  514. agno/models/ollama/tools.py +0 -430
  515. agno/models/qwen/__init__.py +0 -5
  516. agno/playground/__init__.py +0 -10
  517. agno/playground/deploy.py +0 -3
  518. agno/playground/playground.py +0 -3
  519. agno/playground/serve.py +0 -3
  520. agno/playground/settings.py +0 -3
  521. agno/reranker/__init__.py +0 -0
  522. agno/run/response.py +0 -467
  523. agno/run/v2/__init__.py +0 -0
  524. agno/run/v2/workflow.py +0 -567
  525. agno/storage/__init__.py +0 -0
  526. agno/storage/agent/__init__.py +0 -0
  527. agno/storage/agent/dynamodb.py +0 -1
  528. agno/storage/agent/json.py +0 -1
  529. agno/storage/agent/mongodb.py +0 -1
  530. agno/storage/agent/postgres.py +0 -1
  531. agno/storage/agent/singlestore.py +0 -1
  532. agno/storage/agent/sqlite.py +0 -1
  533. agno/storage/agent/yaml.py +0 -1
  534. agno/storage/base.py +0 -60
  535. agno/storage/dynamodb.py +0 -673
  536. agno/storage/firestore.py +0 -297
  537. agno/storage/gcs_json.py +0 -261
  538. agno/storage/in_memory.py +0 -234
  539. agno/storage/json.py +0 -237
  540. agno/storage/mongodb.py +0 -328
  541. agno/storage/mysql.py +0 -685
  542. agno/storage/postgres.py +0 -682
  543. agno/storage/redis.py +0 -336
  544. agno/storage/session/__init__.py +0 -16
  545. agno/storage/session/agent.py +0 -64
  546. agno/storage/session/team.py +0 -63
  547. agno/storage/session/v2/__init__.py +0 -5
  548. agno/storage/session/workflow.py +0 -61
  549. agno/storage/singlestore.py +0 -606
  550. agno/storage/sqlite.py +0 -646
  551. agno/storage/workflow/__init__.py +0 -0
  552. agno/storage/workflow/mongodb.py +0 -1
  553. agno/storage/workflow/postgres.py +0 -1
  554. agno/storage/workflow/sqlite.py +0 -1
  555. agno/storage/yaml.py +0 -241
  556. agno/tools/thinking.py +0 -73
  557. agno/utils/defaults.py +0 -57
  558. agno/utils/filesystem.py +0 -39
  559. agno/utils/git.py +0 -52
  560. agno/utils/json_io.py +0 -30
  561. agno/utils/load_env.py +0 -19
  562. agno/utils/py_io.py +0 -19
  563. agno/utils/pyproject.py +0 -18
  564. agno/utils/resource_filter.py +0 -31
  565. agno/workflow/v2/__init__.py +0 -21
  566. agno/workflow/v2/types.py +0 -357
  567. agno/workflow/v2/workflow.py +0 -3312
  568. agno/workspace/__init__.py +0 -0
  569. agno/workspace/config.py +0 -325
  570. agno/workspace/enums.py +0 -6
  571. agno/workspace/helpers.py +0 -52
  572. agno/workspace/operator.py +0 -757
  573. agno/workspace/settings.py +0 -158
  574. agno-1.8.1.dist-info/METADATA +0 -982
  575. agno-1.8.1.dist-info/RECORD +0 -566
  576. agno-1.8.1.dist-info/entry_points.txt +0 -3
  577. agno-1.8.1.dist-info/licenses/LICENSE +0 -375
  578. /agno/{app → db/migrations}/__init__.py +0 -0
  579. /agno/{app/playground/__init__.py → db/schemas/metrics.py} +0 -0
  580. /agno/{cli → integrations}/__init__.py +0 -0
  581. /agno/{cli/ws → knowledge/chunking}/__init__.py +0 -0
  582. /agno/{document/chunking → knowledge/remote_content}/__init__.py +0 -0
  583. /agno/{document/reader/gcs → knowledge/reranker}/__init__.py +0 -0
  584. /agno/{document/reader/s3 → os/interfaces}/__init__.py +0 -0
  585. /agno/{app → os/interfaces}/slack/security.py +0 -0
  586. /agno/{app → os/interfaces}/whatsapp/security.py +0 -0
  587. /agno/{file/local → utils/print_response}/__init__.py +0 -0
  588. /agno/{infra → vectordb/llamaindex}/__init__.py +0 -0
  589. {agno-1.8.1.dist-info → agno-2.0.0.dist-info}/WHEEL +0 -0
  590. {agno-1.8.1.dist-info → agno-2.0.0.dist-info}/top_level.txt +0 -0
agno/db/mysql/mysql.py ADDED
@@ -0,0 +1,1719 @@
1
+ import time
2
+ from datetime import date, datetime, timedelta, timezone
3
+ from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
4
+ from uuid import uuid4
5
+
6
+ from sqlalchemy import Index, UniqueConstraint
7
+
8
+ from agno.db.base import BaseDb, SessionType
9
+ from agno.db.mysql.schemas import get_table_schema_definition
10
+ from agno.db.mysql.utils import (
11
+ apply_sorting,
12
+ bulk_upsert_metrics,
13
+ calculate_date_metrics,
14
+ create_schema,
15
+ fetch_all_sessions_data,
16
+ get_dates_to_calculate_metrics_for,
17
+ is_table_available,
18
+ is_valid_table,
19
+ )
20
+ from agno.db.schemas.evals import EvalFilterType, EvalRunRecord, EvalType
21
+ from agno.db.schemas.knowledge import KnowledgeRow
22
+ from agno.db.schemas.memory import UserMemory
23
+ from agno.session import AgentSession, Session, TeamSession, WorkflowSession
24
+ from agno.utils.log import log_debug, log_error, log_info
25
+
26
+ try:
27
+ from sqlalchemy import TEXT, and_, cast, func, update
28
+ from sqlalchemy.dialects import mysql
29
+ from sqlalchemy.engine import Engine, create_engine
30
+ from sqlalchemy.orm import scoped_session, sessionmaker
31
+ from sqlalchemy.schema import Column, MetaData, Table
32
+ from sqlalchemy.sql.expression import select, text
33
+ except ImportError:
34
+ raise ImportError("`sqlalchemy` not installed. Please install it using `pip install sqlalchemy`")
35
+
36
+
37
+ class MySQLDb(BaseDb):
38
+ def __init__(
39
+ self,
40
+ db_engine: Optional[Engine] = None,
41
+ db_schema: Optional[str] = None,
42
+ db_url: Optional[str] = None,
43
+ session_table: Optional[str] = None,
44
+ memory_table: Optional[str] = None,
45
+ metrics_table: Optional[str] = None,
46
+ eval_table: Optional[str] = None,
47
+ knowledge_table: Optional[str] = None,
48
+ ):
49
+ """
50
+ Interface for interacting with a MySQL database.
51
+
52
+ The following order is used to determine the database connection:
53
+ 1. Use the db_engine if provided
54
+ 2. Use the db_url
55
+ 3. Raise an error if neither is provided
56
+
57
+ Args:
58
+ db_url (Optional[str]): The database URL to connect to.
59
+ db_engine (Optional[Engine]): The SQLAlchemy database engine to use.
60
+ db_schema (Optional[str]): The database schema to use.
61
+ session_table (Optional[str]): Name of the table to store Agent, Team and Workflow sessions.
62
+ memory_table (Optional[str]): Name of the table to store memories.
63
+ metrics_table (Optional[str]): Name of the table to store metrics.
64
+ eval_table (Optional[str]): Name of the table to store evaluation runs data.
65
+ knowledge_table (Optional[str]): Name of the table to store knowledge content.
66
+
67
+ Raises:
68
+ ValueError: If neither db_url nor db_engine is provided.
69
+ ValueError: If none of the tables are provided.
70
+ """
71
+ super().__init__(
72
+ session_table=session_table,
73
+ memory_table=memory_table,
74
+ metrics_table=metrics_table,
75
+ eval_table=eval_table,
76
+ knowledge_table=knowledge_table,
77
+ )
78
+
79
+ _engine: Optional[Engine] = db_engine
80
+ if _engine is None and db_url is not None:
81
+ _engine = create_engine(db_url)
82
+ if _engine is None:
83
+ raise ValueError("One of db_url or db_engine must be provided")
84
+
85
+ self.db_url: Optional[str] = db_url
86
+ self.db_engine: Engine = _engine
87
+ self.db_schema: str = db_schema if db_schema is not None else "ai"
88
+ self.metadata: MetaData = MetaData()
89
+
90
+ # Initialize database session
91
+ self.Session: scoped_session = scoped_session(sessionmaker(bind=self.db_engine))
92
+
93
+ # -- DB methods --
94
+ def _create_table(self, table_name: str, table_type: str, db_schema: str) -> Table:
95
+ """
96
+ Create a table with the appropriate schema based on the table type.
97
+
98
+ Args:
99
+ table_name (str): Name of the table to create
100
+ table_type (str): Type of table (used to get schema definition)
101
+ db_schema (str): Database schema name
102
+
103
+ Returns:
104
+ Table: SQLAlchemy Table object
105
+ """
106
+ try:
107
+ table_schema = get_table_schema_definition(table_type)
108
+
109
+ log_debug(f"Creating table {db_schema}.{table_name} with schema: {table_schema}")
110
+
111
+ columns: List[Column] = []
112
+ indexes: List[str] = []
113
+ unique_constraints: List[str] = []
114
+ schema_unique_constraints = table_schema.pop("_unique_constraints", [])
115
+
116
+ # Get the columns, indexes, and unique constraints from the table schema
117
+ for col_name, col_config in table_schema.items():
118
+ column_args = [col_name, col_config["type"]()]
119
+ column_kwargs = {}
120
+ if col_config.get("primary_key", False):
121
+ column_kwargs["primary_key"] = True
122
+ if "nullable" in col_config:
123
+ column_kwargs["nullable"] = col_config["nullable"]
124
+ if col_config.get("index", False):
125
+ indexes.append(col_name)
126
+ if col_config.get("unique", False):
127
+ column_kwargs["unique"] = True
128
+ unique_constraints.append(col_name)
129
+ columns.append(Column(*column_args, **column_kwargs)) # type: ignore
130
+
131
+ # Create the table object
132
+ table_metadata = MetaData(schema=db_schema)
133
+ table = Table(table_name, table_metadata, *columns, schema=db_schema)
134
+
135
+ # Add multi-column unique constraints with table-specific names
136
+ for constraint in schema_unique_constraints:
137
+ constraint_name = f"{table_name}_{constraint['name']}"
138
+ constraint_columns = constraint["columns"]
139
+ table.append_constraint(UniqueConstraint(*constraint_columns, name=constraint_name))
140
+
141
+ # Add indexes to the table definition
142
+ for idx_col in indexes:
143
+ idx_name = f"idx_{table_name}_{idx_col}"
144
+ table.append_constraint(Index(idx_name, idx_col))
145
+
146
+ with self.Session() as sess, sess.begin():
147
+ create_schema(session=sess, db_schema=db_schema)
148
+
149
+ # Create table
150
+ table.create(self.db_engine, checkfirst=True)
151
+
152
+ # Create indexes
153
+ for idx in table.indexes:
154
+ try:
155
+ log_debug(f"Creating index: {idx.name}")
156
+
157
+ # Check if index already exists
158
+ with self.Session() as sess:
159
+ exists_query = text(
160
+ "SELECT 1 FROM information_schema.statistics WHERE table_schema = :schema "
161
+ "AND table_name = :table_name AND index_name = :index_name"
162
+ )
163
+ exists = (
164
+ sess.execute(
165
+ exists_query, {"schema": db_schema, "table_name": table_name, "index_name": idx.name}
166
+ ).scalar()
167
+ is not None
168
+ )
169
+ if exists:
170
+ log_debug(f"Index {idx.name} already exists in {db_schema}.{table_name}, skipping creation")
171
+ continue
172
+
173
+ idx.create(self.db_engine)
174
+
175
+ except Exception as e:
176
+ log_error(f"Error creating index {idx.name}: {e}")
177
+
178
+ log_info(f"Successfully created table {db_schema}.{table_name}")
179
+ return table
180
+
181
+ except Exception as e:
182
+ log_error(f"Could not create table {db_schema}.{table_name}: {e}")
183
+ raise
184
+
185
+ def _get_table(self, table_type: str, create_table_if_not_found: Optional[bool] = False) -> Optional[Table]:
186
+ if table_type == "sessions":
187
+ self.session_table = self._get_or_create_table(
188
+ table_name=self.session_table_name,
189
+ table_type="sessions",
190
+ db_schema=self.db_schema,
191
+ create_table_if_not_found=create_table_if_not_found,
192
+ )
193
+ return self.session_table
194
+
195
+ if table_type == "memories":
196
+ self.memory_table = self._get_or_create_table(
197
+ table_name=self.memory_table_name,
198
+ table_type="memories",
199
+ db_schema=self.db_schema,
200
+ create_table_if_not_found=create_table_if_not_found,
201
+ )
202
+ return self.memory_table
203
+
204
+ if table_type == "metrics":
205
+ self.metrics_table = self._get_or_create_table(
206
+ table_name=self.metrics_table_name,
207
+ table_type="metrics",
208
+ db_schema=self.db_schema,
209
+ create_table_if_not_found=create_table_if_not_found,
210
+ )
211
+ return self.metrics_table
212
+
213
+ if table_type == "evals":
214
+ self.eval_table = self._get_or_create_table(
215
+ table_name=self.eval_table_name,
216
+ table_type="evals",
217
+ db_schema=self.db_schema,
218
+ create_table_if_not_found=create_table_if_not_found,
219
+ )
220
+ return self.eval_table
221
+
222
+ if table_type == "knowledge":
223
+ self.knowledge_table = self._get_or_create_table(
224
+ table_name=self.knowledge_table_name,
225
+ table_type="knowledge",
226
+ db_schema=self.db_schema,
227
+ create_table_if_not_found=create_table_if_not_found,
228
+ )
229
+ return self.knowledge_table
230
+
231
+ raise ValueError(f"Unknown table type: {table_type}")
232
+
233
+ def _get_or_create_table(
234
+ self, table_name: str, table_type: str, db_schema: str, create_table_if_not_found: Optional[bool] = False
235
+ ) -> Optional[Table]:
236
+ """
237
+ Check if the table exists and is valid, else create it.
238
+
239
+ Args:
240
+ table_name (str): Name of the table to get or create
241
+ table_type (str): Type of table (used to get schema definition)
242
+ db_schema (str): Database schema name
243
+
244
+ Returns:
245
+ Table: SQLAlchemy Table object representing the schema.
246
+ """
247
+
248
+ with self.Session() as sess, sess.begin():
249
+ table_is_available = is_table_available(session=sess, table_name=table_name, db_schema=db_schema)
250
+
251
+ if not table_is_available:
252
+ if not create_table_if_not_found:
253
+ return None
254
+
255
+ return self._create_table(table_name=table_name, table_type=table_type, db_schema=db_schema)
256
+
257
+ if not is_valid_table(
258
+ db_engine=self.db_engine,
259
+ table_name=table_name,
260
+ table_type=table_type,
261
+ db_schema=db_schema,
262
+ ):
263
+ raise ValueError(f"Table {db_schema}.{table_name} has an invalid schema")
264
+
265
+ try:
266
+ table = Table(table_name, self.metadata, schema=db_schema, autoload_with=self.db_engine)
267
+ log_debug(f"Loaded existing table {db_schema}.{table_name}")
268
+ return table
269
+
270
+ except Exception as e:
271
+ log_error(f"Error loading existing table {db_schema}.{table_name}: {e}")
272
+ raise
273
+
274
+ # -- Session methods --
275
+ def delete_session(self, session_id: str) -> bool:
276
+ """
277
+ Delete a session from the database.
278
+
279
+ Args:
280
+ session_id (str): ID of the session to delete
281
+
282
+ Returns:
283
+ bool: True if the session was deleted, False otherwise.
284
+
285
+ Raises:
286
+ Exception: If an error occurs during deletion.
287
+ """
288
+ try:
289
+ table = self._get_table(table_type="sessions")
290
+ if table is None:
291
+ return False
292
+
293
+ with self.Session() as sess, sess.begin():
294
+ delete_stmt = table.delete().where(table.c.session_id == session_id)
295
+ result = sess.execute(delete_stmt)
296
+ if result.rowcount == 0:
297
+ log_debug(f"No session found to delete with session_id: {session_id} in table {table.name}")
298
+ return False
299
+ else:
300
+ log_debug(f"Successfully deleted session with session_id: {session_id} in table {table.name}")
301
+ return True
302
+
303
+ except Exception as e:
304
+ log_error(f"Error deleting session: {e}")
305
+ return False
306
+
307
+ def delete_sessions(self, session_ids: List[str]) -> None:
308
+ """Delete all given sessions from the database.
309
+ Can handle multiple session types in the same run.
310
+
311
+ Args:
312
+ session_ids (List[str]): The IDs of the sessions to delete.
313
+
314
+ Raises:
315
+ Exception: If an error occurs during deletion.
316
+ """
317
+ try:
318
+ table = self._get_table(table_type="sessions")
319
+ if table is None:
320
+ return
321
+
322
+ with self.Session() as sess, sess.begin():
323
+ delete_stmt = table.delete().where(table.c.session_id.in_(session_ids))
324
+ result = sess.execute(delete_stmt)
325
+
326
+ log_debug(f"Successfully deleted {result.rowcount} sessions")
327
+
328
+ except Exception as e:
329
+ log_error(f"Error deleting sessions: {e}")
330
+
331
+ def get_session(
332
+ self,
333
+ session_id: str,
334
+ session_type: SessionType,
335
+ user_id: Optional[str] = None,
336
+ deserialize: Optional[bool] = True,
337
+ ) -> Optional[Union[Session, Dict[str, Any]]]:
338
+ """
339
+ Read a session from the database.
340
+
341
+ Args:
342
+ session_id (str): ID of the session to read.
343
+ user_id (Optional[str]): User ID to filter by. Defaults to None.
344
+ session_type (Optional[SessionType]): Type of session to read. Defaults to None.
345
+ deserialize (Optional[bool]): Whether to serialize the session. Defaults to True.
346
+
347
+ Returns:
348
+ Union[Session, Dict[str, Any], None]:
349
+ - When deserialize=True: Session object
350
+ - When deserialize=False: Session dictionary
351
+
352
+ Raises:
353
+ Exception: If an error occurs during retrieval.
354
+ """
355
+ try:
356
+ table = self._get_table(table_type="sessions")
357
+ if table is None:
358
+ return None
359
+
360
+ with self.Session() as sess:
361
+ stmt = select(table).where(table.c.session_id == session_id)
362
+
363
+ if user_id is not None:
364
+ stmt = stmt.where(table.c.user_id == user_id)
365
+ if session_type is not None:
366
+ session_type_value = session_type.value if isinstance(session_type, SessionType) else session_type
367
+ stmt = stmt.where(table.c.session_type == session_type_value)
368
+ result = sess.execute(stmt).fetchone()
369
+ if result is None:
370
+ return None
371
+
372
+ session = dict(result._mapping)
373
+
374
+ if not deserialize:
375
+ return session
376
+
377
+ if session_type == SessionType.AGENT:
378
+ return AgentSession.from_dict(session)
379
+ elif session_type == SessionType.TEAM:
380
+ return TeamSession.from_dict(session)
381
+ elif session_type == SessionType.WORKFLOW:
382
+ return WorkflowSession.from_dict(session)
383
+ else:
384
+ raise ValueError(f"Invalid session type: {session_type}")
385
+
386
+ except Exception as e:
387
+ log_error(f"Exception reading from session table: {e}")
388
+ return None
389
+
390
+ def get_sessions(
391
+ self,
392
+ session_type: Optional[SessionType] = None,
393
+ user_id: Optional[str] = None,
394
+ component_id: Optional[str] = None,
395
+ session_name: Optional[str] = None,
396
+ start_timestamp: Optional[int] = None,
397
+ end_timestamp: Optional[int] = None,
398
+ limit: Optional[int] = None,
399
+ page: Optional[int] = None,
400
+ sort_by: Optional[str] = None,
401
+ sort_order: Optional[str] = None,
402
+ deserialize: Optional[bool] = True,
403
+ ) -> Union[List[Session], Tuple[List[Dict[str, Any]], int]]:
404
+ """
405
+ Get all sessions in the given table. Can filter by user_id and entity_id.
406
+
407
+ Args:
408
+ user_id (Optional[str]): The ID of the user to filter by.
409
+ entity_id (Optional[str]): The ID of the agent / workflow to filter by.
410
+ start_timestamp (Optional[int]): The start timestamp to filter by.
411
+ end_timestamp (Optional[int]): The end timestamp to filter by.
412
+ session_name (Optional[str]): The name of the session to filter by.
413
+ limit (Optional[int]): The maximum number of sessions to return. Defaults to None.
414
+ page (Optional[int]): The page number to return. Defaults to None.
415
+ sort_by (Optional[str]): The field to sort by. Defaults to None.
416
+ sort_order (Optional[str]): The sort order. Defaults to None.
417
+ deserialize (Optional[bool]): Whether to serialize the sessions. Defaults to True.
418
+ create_table_if_not_found (Optional[bool]): Whether to create the table if it doesn't exist.
419
+
420
+ Returns:
421
+ Union[List[Session], Tuple[List[Dict], int]]:
422
+ - When deserialize=True: List of Session objects
423
+ - When deserialize=False: Tuple of (session dictionaries, total count)
424
+
425
+ Raises:
426
+ Exception: If an error occurs during retrieval.
427
+ """
428
+ try:
429
+ table = self._get_table(table_type="sessions")
430
+ if table is None:
431
+ return [] if deserialize else ([], 0)
432
+
433
+ with self.Session() as sess, sess.begin():
434
+ stmt = select(table)
435
+
436
+ # Filtering
437
+ if user_id is not None:
438
+ stmt = stmt.where(table.c.user_id == user_id)
439
+ if component_id is not None:
440
+ if session_type == SessionType.AGENT:
441
+ stmt = stmt.where(table.c.agent_id == component_id)
442
+ elif session_type == SessionType.TEAM:
443
+ stmt = stmt.where(table.c.team_id == component_id)
444
+ elif session_type == SessionType.WORKFLOW:
445
+ stmt = stmt.where(table.c.workflow_id == component_id)
446
+ if start_timestamp is not None:
447
+ stmt = stmt.where(table.c.created_at >= start_timestamp)
448
+ if end_timestamp is not None:
449
+ stmt = stmt.where(table.c.created_at <= end_timestamp)
450
+ if session_name is not None:
451
+ # MySQL JSON extraction syntax
452
+ stmt = stmt.where(
453
+ func.coalesce(
454
+ func.json_unquote(func.json_extract(table.c.session_data, "$.session_name")), ""
455
+ ).ilike(f"%{session_name}%")
456
+ )
457
+ if session_type is not None:
458
+ session_type_value = session_type.value if isinstance(session_type, SessionType) else session_type
459
+ stmt = stmt.where(table.c.session_type == session_type_value)
460
+
461
+ count_stmt = select(func.count()).select_from(stmt.alias())
462
+ total_count = sess.execute(count_stmt).scalar()
463
+
464
+ # Sorting
465
+ stmt = apply_sorting(stmt, table, sort_by, sort_order)
466
+
467
+ # Paginating
468
+ if limit is not None:
469
+ stmt = stmt.limit(limit)
470
+ if page is not None:
471
+ stmt = stmt.offset((page - 1) * limit)
472
+
473
+ result = sess.execute(stmt).fetchall()
474
+ if not result:
475
+ return [] if deserialize else ([], 0)
476
+
477
+ session_dicts = [dict(row._mapping) for row in result]
478
+ if not deserialize:
479
+ return session_dicts, total_count
480
+
481
+ if session_type == SessionType.AGENT:
482
+ return [AgentSession.from_dict(record) for record in session_dicts] # type: ignore
483
+ elif session_type == SessionType.TEAM:
484
+ return [TeamSession.from_dict(record) for record in session_dicts] # type: ignore
485
+ elif session_type == SessionType.WORKFLOW:
486
+ return [WorkflowSession.from_dict(record) for record in session_dicts] # type: ignore
487
+ else:
488
+ raise ValueError(f"Invalid session type: {session_type}")
489
+
490
+ except Exception as e:
491
+ log_error(f"Exception getting eval runs: {e}")
492
+ return [] if deserialize else ([], 0)
493
+
494
+ def rename_session(
495
+ self, session_id: str, session_type: SessionType, session_name: str, deserialize: Optional[bool] = True
496
+ ) -> Optional[Union[Session, Dict[str, Any]]]:
497
+ """
498
+ Rename a session in the database.
499
+
500
+ Args:
501
+ session_id (str): The ID of the session to rename.
502
+ session_type (SessionType): The type of session to rename.
503
+ session_name (str): The new name for the session.
504
+ deserialize (Optional[bool]): Whether to serialize the session. Defaults to True.
505
+
506
+ Returns:
507
+ Optional[Union[Session, Dict[str, Any]]]:
508
+ - When deserialize=True: Session object
509
+ - When deserialize=False: Session dictionary
510
+
511
+ Raises:
512
+ Exception: If an error occurs during renaming.
513
+ """
514
+ try:
515
+ table = self._get_table(table_type="sessions")
516
+ if table is None:
517
+ return None
518
+
519
+ with self.Session() as sess, sess.begin():
520
+ # MySQL JSON_SET syntax
521
+ stmt = (
522
+ update(table)
523
+ .where(table.c.session_id == session_id)
524
+ .where(table.c.session_type == session_type.value)
525
+ .values(session_data=func.json_set(table.c.session_data, "$.session_name", session_name))
526
+ )
527
+ sess.execute(stmt)
528
+
529
+ # Fetch the updated row
530
+ select_stmt = select(table).where(table.c.session_id == session_id)
531
+ result = sess.execute(select_stmt)
532
+ row = result.fetchone()
533
+ if not row:
534
+ return None
535
+
536
+ session = dict(row._mapping)
537
+ if not deserialize:
538
+ return session
539
+
540
+ # Return the appropriate session type
541
+ if session_type == SessionType.AGENT:
542
+ return AgentSession.from_dict(session)
543
+ elif session_type == SessionType.TEAM:
544
+ return TeamSession.from_dict(session)
545
+ elif session_type == SessionType.WORKFLOW:
546
+ return WorkflowSession.from_dict(session)
547
+ else:
548
+ raise ValueError(f"Invalid session type: {session_type}")
549
+
550
+ except Exception as e:
551
+ log_error(f"Exception renaming session: {e}")
552
+ return None
553
+
554
+ def upsert_session(
555
+ self, session: Session, deserialize: Optional[bool] = True
556
+ ) -> Optional[Union[Session, Dict[str, Any]]]:
557
+ """
558
+ Insert or update a session in the database.
559
+
560
+ Args:
561
+ session (Session): The session data to upsert.
562
+ deserialize (Optional[bool]): Whether to deserialize the session. Defaults to True.
563
+
564
+ Returns:
565
+ Optional[Union[Session, Dict[str, Any]]]:
566
+ - When deserialize=True: Session object
567
+ - When deserialize=False: Session dictionary
568
+
569
+ Raises:
570
+ Exception: If an error occurs during upsert.
571
+ """
572
+ try:
573
+ table = self._get_table(table_type="sessions", create_table_if_not_found=True)
574
+ if table is None:
575
+ return None
576
+
577
+ session_dict = session.to_dict()
578
+
579
+ if isinstance(session, AgentSession):
580
+ with self.Session() as sess, sess.begin():
581
+ stmt = mysql.insert(table).values(
582
+ session_id=session_dict.get("session_id"),
583
+ session_type=SessionType.AGENT.value,
584
+ agent_id=session_dict.get("agent_id"),
585
+ user_id=session_dict.get("user_id"),
586
+ runs=session_dict.get("runs"),
587
+ agent_data=session_dict.get("agent_data"),
588
+ session_data=session_dict.get("session_data"),
589
+ summary=session_dict.get("summary"),
590
+ metadata=session_dict.get("metadata"),
591
+ created_at=session_dict.get("created_at"),
592
+ updated_at=session_dict.get("created_at"),
593
+ )
594
+ stmt = stmt.on_duplicate_key_update(
595
+ agent_id=session_dict.get("agent_id"),
596
+ user_id=session_dict.get("user_id"),
597
+ agent_data=session_dict.get("agent_data"),
598
+ session_data=session_dict.get("session_data"),
599
+ summary=session_dict.get("summary"),
600
+ metadata=session_dict.get("metadata"),
601
+ runs=session_dict.get("runs"),
602
+ updated_at=int(time.time()),
603
+ )
604
+ sess.execute(stmt)
605
+
606
+ # Fetch the row
607
+ select_stmt = select(table).where(table.c.session_id == session_dict.get("session_id"))
608
+ result = sess.execute(select_stmt)
609
+ row = result.fetchone()
610
+ if not row:
611
+ return None
612
+ session_dict = dict(row._mapping)
613
+ if session_dict is None or not deserialize:
614
+ return session_dict
615
+ return AgentSession.from_dict(session_dict)
616
+
617
+ elif isinstance(session, TeamSession):
618
+ with self.Session() as sess, sess.begin():
619
+ stmt = mysql.insert(table).values(
620
+ session_id=session_dict.get("session_id"),
621
+ session_type=SessionType.TEAM.value,
622
+ team_id=session_dict.get("team_id"),
623
+ user_id=session_dict.get("user_id"),
624
+ runs=session_dict.get("runs"),
625
+ team_data=session_dict.get("team_data"),
626
+ session_data=session_dict.get("session_data"),
627
+ summary=session_dict.get("summary"),
628
+ metadata=session_dict.get("metadata"),
629
+ created_at=session_dict.get("created_at"),
630
+ updated_at=session_dict.get("created_at"),
631
+ )
632
+ stmt = stmt.on_duplicate_key_update(
633
+ team_id=session_dict.get("team_id"),
634
+ user_id=session_dict.get("user_id"),
635
+ team_data=session_dict.get("team_data"),
636
+ session_data=session_dict.get("session_data"),
637
+ summary=session_dict.get("summary"),
638
+ metadata=session_dict.get("metadata"),
639
+ runs=session_dict.get("runs"),
640
+ updated_at=int(time.time()),
641
+ )
642
+ sess.execute(stmt)
643
+
644
+ # Fetch the row
645
+ select_stmt = select(table).where(table.c.session_id == session_dict.get("session_id"))
646
+ result = sess.execute(select_stmt)
647
+ row = result.fetchone()
648
+ if not row:
649
+ return None
650
+ session_dict = dict(row._mapping)
651
+ if session_dict is None or not deserialize:
652
+ return session_dict
653
+ return TeamSession.from_dict(session_dict)
654
+
655
+ else:
656
+ with self.Session() as sess, sess.begin():
657
+ stmt = mysql.insert(table).values(
658
+ session_id=session_dict.get("session_id"),
659
+ session_type=SessionType.WORKFLOW.value,
660
+ workflow_id=session_dict.get("workflow_id"),
661
+ user_id=session_dict.get("user_id"),
662
+ runs=session_dict.get("runs"),
663
+ workflow_data=session_dict.get("workflow_data"),
664
+ session_data=session_dict.get("session_data"),
665
+ summary=session_dict.get("summary"),
666
+ metadata=session_dict.get("metadata"),
667
+ created_at=session_dict.get("created_at"),
668
+ updated_at=session_dict.get("created_at"),
669
+ )
670
+ stmt = stmt.on_duplicate_key_update(
671
+ workflow_id=session_dict.get("workflow_id"),
672
+ user_id=session_dict.get("user_id"),
673
+ workflow_data=session_dict.get("workflow_data"),
674
+ session_data=session_dict.get("session_data"),
675
+ summary=session_dict.get("summary"),
676
+ metadata=session_dict.get("metadata"),
677
+ runs=session_dict.get("runs"),
678
+ updated_at=int(time.time()),
679
+ )
680
+ sess.execute(stmt)
681
+
682
+ # Fetch the row
683
+ select_stmt = select(table).where(table.c.session_id == session_dict.get("session_id"))
684
+ result = sess.execute(select_stmt)
685
+ row = result.fetchone()
686
+ if not row:
687
+ return None
688
+ session_dict = dict(row._mapping)
689
+ if session_dict is None or not deserialize:
690
+ return session_dict
691
+ return WorkflowSession.from_dict(session_dict)
692
+
693
+ except Exception as e:
694
+ log_error(f"Exception upserting into sessions table: {e}")
695
+ return None
696
+
697
+ # -- Memory methods --
698
+ def delete_user_memory(self, memory_id: str):
699
+ """Delete a user memory from the database.
700
+
701
+ Returns:
702
+ bool: True if deletion was successful, False otherwise.
703
+
704
+ Raises:
705
+ Exception: If an error occurs during deletion.
706
+ """
707
+ try:
708
+ table = self._get_table(table_type="memories")
709
+ if table is None:
710
+ return
711
+
712
+ with self.Session() as sess, sess.begin():
713
+ delete_stmt = table.delete().where(table.c.memory_id == memory_id)
714
+ result = sess.execute(delete_stmt)
715
+
716
+ success = result.rowcount > 0
717
+ if success:
718
+ log_debug(f"Successfully deleted user memory id: {memory_id}")
719
+ else:
720
+ log_debug(f"No user memory found with id: {memory_id}")
721
+
722
+ except Exception as e:
723
+ log_error(f"Error deleting user memory: {e}")
724
+
725
+ def delete_user_memories(self, memory_ids: List[str]) -> None:
726
+ """Delete user memories from the database.
727
+
728
+ Args:
729
+ memory_ids (List[str]): The IDs of the memories to delete.
730
+
731
+ Raises:
732
+ Exception: If an error occurs during deletion.
733
+ """
734
+ try:
735
+ table = self._get_table(table_type="memories")
736
+ if table is None:
737
+ return
738
+
739
+ with self.Session() as sess, sess.begin():
740
+ delete_stmt = table.delete().where(table.c.memory_id.in_(memory_ids))
741
+ result = sess.execute(delete_stmt)
742
+ if result.rowcount == 0:
743
+ log_debug(f"No user memories found with ids: {memory_ids}")
744
+
745
+ except Exception as e:
746
+ log_error(f"Error deleting user memories: {e}")
747
+
748
+ def get_all_memory_topics(self) -> List[str]:
749
+ """Get all memory topics from the database.
750
+
751
+ Returns:
752
+ List[str]: List of memory topics.
753
+ """
754
+ try:
755
+ table = self._get_table(table_type="memories")
756
+ if table is None:
757
+ return []
758
+
759
+ with self.Session() as sess, sess.begin():
760
+ # MySQL approach: extract JSON array elements differently
761
+ stmt = select(table.c.topics)
762
+ result = sess.execute(stmt).fetchall()
763
+
764
+ topics_set = set()
765
+ for row in result:
766
+ if row[0]:
767
+ # Parse JSON array and add topics to set
768
+ import json
769
+
770
+ try:
771
+ topics = json.loads(row[0]) if isinstance(row[0], str) else row[0]
772
+ if isinstance(topics, list):
773
+ topics_set.update(topics)
774
+ except Exception:
775
+ pass
776
+
777
+ return list(topics_set)
778
+
779
+ except Exception as e:
780
+ log_error(f"Exception reading from memory table: {e}")
781
+ return []
782
+
783
+ def get_user_memory(self, memory_id: str, deserialize: Optional[bool] = True) -> Optional[UserMemory]:
784
+ """Get a memory from the database.
785
+
786
+ Args:
787
+ memory_id (str): The ID of the memory to get.
788
+ deserialize (Optional[bool]): Whether to serialize the memory. Defaults to True.
789
+
790
+ Returns:
791
+ Union[UserMemory, Dict[str, Any], None]:
792
+ - When deserialize=True: UserMemory object
793
+ - When deserialize=False: UserMemory dictionary
794
+
795
+ Raises:
796
+ Exception: If an error occurs during retrieval.
797
+ """
798
+ try:
799
+ table = self._get_table(table_type="memories")
800
+ if table is None:
801
+ return None
802
+
803
+ with self.Session() as sess, sess.begin():
804
+ stmt = select(table).where(table.c.memory_id == memory_id)
805
+
806
+ result = sess.execute(stmt).fetchone()
807
+ if not result:
808
+ return None
809
+
810
+ memory_raw = result._mapping
811
+ if not deserialize:
812
+ return memory_raw
813
+ return UserMemory.from_dict(memory_raw)
814
+
815
+ except Exception as e:
816
+ log_error(f"Exception reading from memory table: {e}")
817
+ return None
818
+
819
+ def get_user_memories(
820
+ self,
821
+ user_id: Optional[str] = None,
822
+ agent_id: Optional[str] = None,
823
+ team_id: Optional[str] = None,
824
+ topics: Optional[List[str]] = None,
825
+ search_content: Optional[str] = None,
826
+ limit: Optional[int] = None,
827
+ page: Optional[int] = None,
828
+ sort_by: Optional[str] = None,
829
+ sort_order: Optional[str] = None,
830
+ deserialize: Optional[bool] = True,
831
+ ) -> Union[List[UserMemory], Tuple[List[Dict[str, Any]], int]]:
832
+ """Get all memories from the database as MemoryRow objects.
833
+
834
+ Args:
835
+ user_id (Optional[str]): The ID of the user to filter by.
836
+ agent_id (Optional[str]): The ID of the agent to filter by.
837
+ team_id (Optional[str]): The ID of the team to filter by.
838
+ topics (Optional[List[str]]): The topics to filter by.
839
+ search_content (Optional[str]): The content to search for.
840
+ limit (Optional[int]): The maximum number of memories to return.
841
+ page (Optional[int]): The page number.
842
+ sort_by (Optional[str]): The column to sort by.
843
+ sort_order (Optional[str]): The order to sort by.
844
+ deserialize (Optional[bool]): Whether to serialize the memories. Defaults to True.
845
+
846
+
847
+ Returns:
848
+ Union[List[UserMemory], Tuple[List[Dict[str, Any]], int]]:
849
+ - When deserialize=True: List of UserMemory objects
850
+ - When deserialize=False: Tuple of (memory dictionaries, total count)
851
+
852
+ Raises:
853
+ Exception: If an error occurs during retrieval.
854
+ """
855
+ try:
856
+ table = self._get_table(table_type="memories")
857
+ if table is None:
858
+ return [] if deserialize else ([], 0)
859
+
860
+ with self.Session() as sess, sess.begin():
861
+ stmt = select(table)
862
+ # Filtering
863
+ if user_id is not None:
864
+ stmt = stmt.where(table.c.user_id == user_id)
865
+ if agent_id is not None:
866
+ stmt = stmt.where(table.c.agent_id == agent_id)
867
+ if team_id is not None:
868
+ stmt = stmt.where(table.c.team_id == team_id)
869
+ if topics is not None:
870
+ # MySQL JSON contains syntax
871
+ topic_conditions = []
872
+ for topic in topics:
873
+ topic_conditions.append(func.json_contains(table.c.topics, f'"{topic}"'))
874
+ stmt = stmt.where(and_(*topic_conditions))
875
+ if search_content is not None:
876
+ stmt = stmt.where(cast(table.c.memory, TEXT).ilike(f"%{search_content}%"))
877
+
878
+ # Get total count after applying filtering
879
+ count_stmt = select(func.count()).select_from(stmt.alias())
880
+ total_count = sess.execute(count_stmt).scalar()
881
+
882
+ # Sorting
883
+ stmt = apply_sorting(stmt, table, sort_by, sort_order)
884
+
885
+ # Paginating
886
+ if limit is not None:
887
+ stmt = stmt.limit(limit)
888
+ if page is not None:
889
+ stmt = stmt.offset((page - 1) * limit)
890
+
891
+ result = sess.execute(stmt).fetchall()
892
+ if not result:
893
+ return [] if deserialize else ([], 0)
894
+
895
+ memories_raw = [record._mapping for record in result]
896
+ if not deserialize:
897
+ return memories_raw, total_count
898
+
899
+ return [UserMemory.from_dict(record) for record in memories_raw]
900
+
901
+ except Exception as e:
902
+ log_error(f"Exception reading from memory table: {e}")
903
+ return [] if deserialize else ([], 0)
904
+
905
+ def clear_memories(self) -> None:
906
+ """Clear all user memories from the database."""
907
+ try:
908
+ table = self._get_table(table_type="memories")
909
+ if table is None:
910
+ return
911
+
912
+ with self.Session() as sess, sess.begin():
913
+ sess.execute(table.delete())
914
+ except Exception as e:
915
+ log_error(f"Exception clearing user memories: {e}")
916
+
917
+ def get_user_memory_stats(
918
+ self, limit: Optional[int] = None, page: Optional[int] = None
919
+ ) -> Tuple[List[Dict[str, Any]], int]:
920
+ """Get user memories stats.
921
+
922
+ Args:
923
+ limit (Optional[int]): The maximum number of user stats to return.
924
+ page (Optional[int]): The page number.
925
+
926
+ Returns:
927
+ Tuple[List[Dict[str, Any]], int]: A list of dictionaries containing user stats and total count.
928
+
929
+ Example:
930
+ (
931
+ [
932
+ {
933
+ "user_id": "123",
934
+ "total_memories": 10,
935
+ "last_memory_updated_at": 1714560000,
936
+ },
937
+ ],
938
+ total_count: 1,
939
+ )
940
+ """
941
+ try:
942
+ table = self._get_table(table_type="memories")
943
+ if table is None:
944
+ return [], 0
945
+
946
+ with self.Session() as sess, sess.begin():
947
+ stmt = (
948
+ select(
949
+ table.c.user_id,
950
+ func.count(table.c.memory_id).label("total_memories"),
951
+ func.max(table.c.updated_at).label("last_memory_updated_at"),
952
+ )
953
+ .where(table.c.user_id.is_not(None))
954
+ .group_by(table.c.user_id)
955
+ .order_by(func.max(table.c.updated_at).desc())
956
+ )
957
+
958
+ count_stmt = select(func.count()).select_from(stmt.alias())
959
+ total_count = sess.execute(count_stmt).scalar()
960
+
961
+ # Pagination
962
+ if limit is not None:
963
+ stmt = stmt.limit(limit)
964
+ if page is not None:
965
+ stmt = stmt.offset((page - 1) * limit)
966
+
967
+ result = sess.execute(stmt).fetchall()
968
+ if not result:
969
+ return [], 0
970
+
971
+ return [
972
+ {
973
+ "user_id": record.user_id, # type: ignore
974
+ "total_memories": record.total_memories,
975
+ "last_memory_updated_at": record.last_memory_updated_at,
976
+ }
977
+ for record in result
978
+ ], total_count
979
+
980
+ except Exception as e:
981
+ log_error(f"Exception getting user memory stats: {e}")
982
+ return [], 0
983
+
984
+ def upsert_user_memory(
985
+ self, memory: UserMemory, deserialize: Optional[bool] = True
986
+ ) -> Optional[Union[UserMemory, Dict[str, Any]]]:
987
+ """Upsert a user memory in the database.
988
+
989
+ Args:
990
+ memory (UserMemory): The user memory to upsert.
991
+ deserialize (Optional[bool]): Whether to serialize the memory. Defaults to True.
992
+
993
+ Returns:
994
+ Optional[Union[UserMemory, Dict[str, Any]]]:
995
+ - When deserialize=True: UserMemory object
996
+ - When deserialize=False: UserMemory dictionary
997
+
998
+ Raises:
999
+ Exception: If an error occurs during upsert.
1000
+ """
1001
+ try:
1002
+ table = self._get_table(table_type="memories", create_table_if_not_found=True)
1003
+ if table is None:
1004
+ return None
1005
+
1006
+ with self.Session() as sess, sess.begin():
1007
+ if memory.memory_id is None:
1008
+ memory.memory_id = str(uuid4())
1009
+
1010
+ stmt = mysql.insert(table).values(
1011
+ memory_id=memory.memory_id,
1012
+ memory=memory.memory,
1013
+ input=memory.input,
1014
+ user_id=memory.user_id,
1015
+ agent_id=memory.agent_id,
1016
+ team_id=memory.team_id,
1017
+ topics=memory.topics,
1018
+ updated_at=int(time.time()),
1019
+ )
1020
+ stmt = stmt.on_duplicate_key_update(
1021
+ memory=memory.memory,
1022
+ topics=memory.topics,
1023
+ input=memory.input,
1024
+ agent_id=memory.agent_id,
1025
+ team_id=memory.team_id,
1026
+ updated_at=int(time.time()),
1027
+ )
1028
+ sess.execute(stmt)
1029
+
1030
+ # Fetch the row
1031
+ select_stmt = select(table).where(table.c.memory_id == memory.memory_id)
1032
+ result = sess.execute(select_stmt)
1033
+ row = result.fetchone()
1034
+ if not row:
1035
+ return None
1036
+
1037
+ memory_raw = row._mapping
1038
+ if not memory_raw or not deserialize:
1039
+ return memory_raw
1040
+
1041
+ return UserMemory.from_dict(memory_raw)
1042
+
1043
+ except Exception as e:
1044
+ log_error(f"Exception upserting user memory: {e}")
1045
+ return None
1046
+
1047
+ # -- Metrics methods --
1048
+ def _get_all_sessions_for_metrics_calculation(
1049
+ self, start_timestamp: Optional[int] = None, end_timestamp: Optional[int] = None
1050
+ ) -> List[Dict[str, Any]]:
1051
+ """
1052
+ Get all sessions of all types (agent, team, workflow) as raw dictionaries.
1053
+
1054
+ Args:
1055
+ start_timestamp (Optional[int]): The start timestamp to filter by. Defaults to None.
1056
+ end_timestamp (Optional[int]): The end timestamp to filter by. Defaults to None.
1057
+
1058
+ Returns:
1059
+ List[Dict[str, Any]]: List of session dictionaries with session_type field.
1060
+
1061
+ Raises:
1062
+ Exception: If an error occurs during retrieval.
1063
+ """
1064
+ try:
1065
+ table = self._get_table(table_type="sessions")
1066
+ if table is None:
1067
+ return []
1068
+
1069
+ stmt = select(
1070
+ table.c.user_id,
1071
+ table.c.session_data,
1072
+ table.c.runs,
1073
+ table.c.created_at,
1074
+ table.c.session_type,
1075
+ )
1076
+
1077
+ if start_timestamp is not None:
1078
+ stmt = stmt.where(table.c.created_at >= start_timestamp)
1079
+ if end_timestamp is not None:
1080
+ stmt = stmt.where(table.c.created_at <= end_timestamp)
1081
+
1082
+ with self.Session() as sess:
1083
+ result = sess.execute(stmt).fetchall()
1084
+ return [record._mapping for record in result]
1085
+
1086
+ except Exception as e:
1087
+ log_error(f"Exception reading from sessions table: {e}")
1088
+ return []
1089
+
1090
+ def _get_metrics_calculation_starting_date(self, table: Table) -> Optional[date]:
1091
+ """Get the first date for which metrics calculation is needed:
1092
+
1093
+ 1. If there are metrics records, return the date of the first day without a complete metrics record.
1094
+ 2. If there are no metrics records, return the date of the first recorded session.
1095
+ 3. If there are no metrics records and no sessions records, return None.
1096
+
1097
+ Args:
1098
+ table (Table): The table to get the starting date for.
1099
+
1100
+ Returns:
1101
+ Optional[date]: The starting date for which metrics calculation is needed.
1102
+ """
1103
+ with self.Session() as sess:
1104
+ stmt = select(table).order_by(table.c.date.desc()).limit(1)
1105
+ result = sess.execute(stmt).fetchone()
1106
+
1107
+ # 1. Return the date of the first day without a complete metrics record.
1108
+ if result is not None:
1109
+ if result.completed:
1110
+ return result._mapping["date"] + timedelta(days=1)
1111
+ else:
1112
+ return result._mapping["date"]
1113
+
1114
+ # 2. No metrics records. Return the date of the first recorded session.
1115
+ first_session, _ = self.get_sessions(sort_by="created_at", sort_order="asc", limit=1, deserialize=False)
1116
+ if not isinstance(first_session, list):
1117
+ raise ValueError("Error obtaining session list to calculate metrics")
1118
+
1119
+ first_session_date = first_session[0]["created_at"] if first_session else None
1120
+
1121
+ # 3. No metrics records and no sessions records. Return None.
1122
+ if first_session_date is None:
1123
+ return None
1124
+
1125
+ return datetime.fromtimestamp(first_session_date, tz=timezone.utc).date()
1126
+
1127
+ def calculate_metrics(self) -> Optional[list[dict]]:
1128
+ """Calculate metrics for all dates without complete metrics.
1129
+
1130
+ Returns:
1131
+ Optional[list[dict]]: The calculated metrics.
1132
+
1133
+ Raises:
1134
+ Exception: If an error occurs during metrics calculation.
1135
+ """
1136
+ try:
1137
+ table = self._get_table(table_type="metrics", create_table_if_not_found=True)
1138
+ if table is None:
1139
+ return None
1140
+
1141
+ starting_date = self._get_metrics_calculation_starting_date(table)
1142
+ if starting_date is None:
1143
+ log_info("No session data found. Won't calculate metrics.")
1144
+ return None
1145
+
1146
+ dates_to_process = get_dates_to_calculate_metrics_for(starting_date)
1147
+ if not dates_to_process:
1148
+ log_info("Metrics already calculated for all relevant dates.")
1149
+ return None
1150
+
1151
+ start_timestamp = int(
1152
+ datetime.combine(dates_to_process[0], datetime.min.time()).replace(tzinfo=timezone.utc).timestamp()
1153
+ )
1154
+ end_timestamp = int(
1155
+ datetime.combine(dates_to_process[-1] + timedelta(days=1), datetime.min.time())
1156
+ .replace(tzinfo=timezone.utc)
1157
+ .timestamp()
1158
+ )
1159
+
1160
+ sessions = self._get_all_sessions_for_metrics_calculation(
1161
+ start_timestamp=start_timestamp, end_timestamp=end_timestamp
1162
+ )
1163
+ all_sessions_data = fetch_all_sessions_data(
1164
+ sessions=sessions, dates_to_process=dates_to_process, start_timestamp=start_timestamp
1165
+ )
1166
+ if not all_sessions_data:
1167
+ log_info("No new session data found. Won't calculate metrics.")
1168
+ return None
1169
+
1170
+ results = []
1171
+ metrics_records = []
1172
+
1173
+ for date_to_process in dates_to_process:
1174
+ date_key = date_to_process.isoformat()
1175
+ sessions_for_date = all_sessions_data.get(date_key, {})
1176
+
1177
+ # Skip dates with no sessions
1178
+ if not any(len(sessions) > 0 for sessions in sessions_for_date.values()):
1179
+ continue
1180
+
1181
+ metrics_record = calculate_date_metrics(date_to_process, sessions_for_date)
1182
+ metrics_records.append(metrics_record)
1183
+
1184
+ if metrics_records:
1185
+ with self.Session() as sess, sess.begin():
1186
+ results = bulk_upsert_metrics(session=sess, table=table, metrics_records=metrics_records)
1187
+
1188
+ return results
1189
+
1190
+ except Exception as e:
1191
+ log_error(f"Exception refreshing metrics: {e}")
1192
+ return None
1193
+
1194
+ def get_metrics(
1195
+ self,
1196
+ starting_date: Optional[date] = None,
1197
+ ending_date: Optional[date] = None,
1198
+ ) -> Tuple[List[dict], Optional[int]]:
1199
+ """Get all metrics matching the given date range.
1200
+
1201
+ Args:
1202
+ starting_date (Optional[date]): The starting date to filter metrics by.
1203
+ ending_date (Optional[date]): The ending date to filter metrics by.
1204
+
1205
+ Returns:
1206
+ Tuple[List[dict], Optional[int]]: A tuple containing the metrics and the timestamp of the latest update.
1207
+
1208
+ Raises:
1209
+ Exception: If an error occurs during retrieval.
1210
+ """
1211
+ try:
1212
+ table = self._get_table(table_type="metrics", create_table_if_not_found=True)
1213
+ if table is None:
1214
+ return [], 0
1215
+
1216
+ with self.Session() as sess, sess.begin():
1217
+ stmt = select(table)
1218
+ if starting_date:
1219
+ stmt = stmt.where(table.c.date >= starting_date)
1220
+ if ending_date:
1221
+ stmt = stmt.where(table.c.date <= ending_date)
1222
+ result = sess.execute(stmt).fetchall()
1223
+ if not result:
1224
+ return [], None
1225
+
1226
+ # Get the latest updated_at
1227
+ latest_stmt = select(func.max(table.c.updated_at))
1228
+ latest_updated_at = sess.execute(latest_stmt).scalar()
1229
+
1230
+ return [row._mapping for row in result], latest_updated_at
1231
+
1232
+ except Exception as e:
1233
+ log_error(f"Exception getting metrics: {e}")
1234
+ return [], None
1235
+
1236
+ # -- Knowledge methods --
1237
+
1238
+ def delete_knowledge_content(self, id: str):
1239
+ """Delete a knowledge row from the database.
1240
+
1241
+ Args:
1242
+ id (str): The ID of the knowledge row to delete.
1243
+
1244
+ Raises:
1245
+ Exception: If an error occurs during deletion.
1246
+ """
1247
+ table = self._get_table(table_type="knowledge")
1248
+ if table is None:
1249
+ return None
1250
+
1251
+ try:
1252
+ with self.Session() as sess, sess.begin():
1253
+ stmt = table.delete().where(table.c.id == id)
1254
+ sess.execute(stmt)
1255
+
1256
+ except Exception as e:
1257
+ log_error(f"Exception deleting knowledge content: {e}")
1258
+
1259
+ def get_knowledge_content(self, id: str) -> Optional[KnowledgeRow]:
1260
+ """Get a knowledge row from the database.
1261
+
1262
+ Args:
1263
+ id (str): The ID of the knowledge row to get.
1264
+
1265
+ Returns:
1266
+ Optional[KnowledgeRow]: The knowledge row, or None if it doesn't exist.
1267
+
1268
+ Raises:
1269
+ Exception: If an error occurs during retrieval.
1270
+ """
1271
+ table = self._get_table(table_type="knowledge")
1272
+ if table is None:
1273
+ return None
1274
+
1275
+ try:
1276
+ with self.Session() as sess, sess.begin():
1277
+ stmt = select(table).where(table.c.id == id)
1278
+ result = sess.execute(stmt).fetchone()
1279
+ if result is None:
1280
+ return None
1281
+ return KnowledgeRow.model_validate(result._mapping)
1282
+
1283
+ except Exception as e:
1284
+ log_error(f"Exception getting knowledge content: {e}")
1285
+ return None
1286
+
1287
+ def get_knowledge_contents(
1288
+ self,
1289
+ limit: Optional[int] = None,
1290
+ page: Optional[int] = None,
1291
+ sort_by: Optional[str] = None,
1292
+ sort_order: Optional[str] = None,
1293
+ ) -> Tuple[List[KnowledgeRow], int]:
1294
+ """Get all knowledge contents from the database.
1295
+
1296
+ Args:
1297
+ limit (Optional[int]): The maximum number of knowledge contents to return.
1298
+ page (Optional[int]): The page number.
1299
+ sort_by (Optional[str]): The column to sort by.
1300
+ sort_order (Optional[str]): The order to sort by.
1301
+ create_table_if_not_found (Optional[bool]): Whether to create the table if it doesn't exist.
1302
+
1303
+ Returns:
1304
+ Tuple[List[KnowledgeRow], int]: The knowledge contents and total count.
1305
+
1306
+ Raises:
1307
+ Exception: If an error occurs during retrieval.
1308
+ """
1309
+ table = self._get_table(table_type="knowledge")
1310
+ if table is None:
1311
+ return [], 0
1312
+
1313
+ try:
1314
+ with self.Session() as sess, sess.begin():
1315
+ stmt = select(table)
1316
+
1317
+ # Apply sorting
1318
+ if sort_by is not None:
1319
+ stmt = stmt.order_by(getattr(table.c, sort_by) * (1 if sort_order == "asc" else -1))
1320
+
1321
+ # Get total count before applying limit and pagination
1322
+ count_stmt = select(func.count()).select_from(stmt.alias())
1323
+ total_count = sess.execute(count_stmt).scalar()
1324
+
1325
+ # Apply pagination after count
1326
+ if limit is not None:
1327
+ stmt = stmt.limit(limit)
1328
+ if page is not None:
1329
+ stmt = stmt.offset((page - 1) * limit)
1330
+
1331
+ result = sess.execute(stmt).fetchall()
1332
+ if not result:
1333
+ return [], 0
1334
+
1335
+ return [KnowledgeRow.model_validate(record._mapping) for record in result], total_count
1336
+
1337
+ except Exception as e:
1338
+ log_error(f"Exception getting knowledge contents: {e}")
1339
+ return [], 0
1340
+
1341
+ def upsert_knowledge_content(self, knowledge_row: KnowledgeRow):
1342
+ """Upsert knowledge content in the database.
1343
+
1344
+ Args:
1345
+ knowledge_row (KnowledgeRow): The knowledge row to upsert.
1346
+
1347
+ Returns:
1348
+ Optional[KnowledgeRow]: The upserted knowledge row, or None if the operation fails.
1349
+
1350
+ Raises:
1351
+ Exception: If an error occurs during upsert.
1352
+ """
1353
+ try:
1354
+ table = self._get_table(table_type="knowledge", create_table_if_not_found=True)
1355
+ if table is None:
1356
+ return None
1357
+
1358
+ with self.Session() as sess, sess.begin():
1359
+ # Get the actual table columns to avoid "unconsumed column names" error
1360
+ table_columns = set(table.columns.keys())
1361
+
1362
+ # Only include fields that exist in the table and are not None
1363
+ insert_data = {}
1364
+ update_fields = {}
1365
+
1366
+ # Map of KnowledgeRow fields to table columns
1367
+ field_mapping = {
1368
+ "id": "id",
1369
+ "name": "name",
1370
+ "description": "description",
1371
+ "metadata": "metadata",
1372
+ "type": "type",
1373
+ "size": "size",
1374
+ "linked_to": "linked_to",
1375
+ "access_count": "access_count",
1376
+ "status": "status",
1377
+ "status_message": "status_message",
1378
+ "created_at": "created_at",
1379
+ "updated_at": "updated_at",
1380
+ "external_id": "external_id",
1381
+ }
1382
+
1383
+ # Build insert and update data only for fields that exist in the table
1384
+ for model_field, table_column in field_mapping.items():
1385
+ if table_column in table_columns:
1386
+ value = getattr(knowledge_row, model_field, None)
1387
+ if value is not None:
1388
+ insert_data[table_column] = value
1389
+ # Don't include ID in update_fields since it's the primary key
1390
+ if table_column != "id":
1391
+ update_fields[table_column] = value
1392
+
1393
+ # Ensure id is always included for the insert
1394
+ if "id" in table_columns and knowledge_row.id:
1395
+ insert_data["id"] = knowledge_row.id
1396
+
1397
+ # Handle case where update_fields is empty (all fields are None or don't exist in table)
1398
+ if not update_fields:
1399
+ # If we have insert_data, just do an insert without conflict resolution
1400
+ if insert_data:
1401
+ stmt = mysql.insert(table).values(insert_data)
1402
+ sess.execute(stmt)
1403
+ else:
1404
+ # If we have no data at all, this is an error
1405
+ log_error("No valid fields found for knowledge row upsert")
1406
+ return None
1407
+ else:
1408
+ # Normal upsert with conflict resolution
1409
+ stmt = mysql.insert(table).values(insert_data).on_duplicate_key_update(**update_fields)
1410
+ sess.execute(stmt)
1411
+
1412
+ return knowledge_row
1413
+
1414
+ except Exception as e:
1415
+ log_error(f"Error upserting knowledge row: {e}")
1416
+ return None
1417
+
1418
+ # -- Eval methods --
1419
+
1420
+ def create_eval_run(self, eval_run: EvalRunRecord) -> Optional[EvalRunRecord]:
1421
+ """Create an EvalRunRecord in the database.
1422
+
1423
+ Args:
1424
+ eval_run (EvalRunRecord): The eval run to create.
1425
+
1426
+ Returns:
1427
+ Optional[EvalRunRecord]: The created eval run, or None if the operation fails.
1428
+
1429
+ Raises:
1430
+ Exception: If an error occurs during creation.
1431
+ """
1432
+ try:
1433
+ table = self._get_table(table_type="evals", create_table_if_not_found=True)
1434
+ if table is None:
1435
+ return None
1436
+
1437
+ with self.Session() as sess, sess.begin():
1438
+ current_time = int(time.time())
1439
+ stmt = mysql.insert(table).values(
1440
+ {"created_at": current_time, "updated_at": current_time, **eval_run.model_dump()}
1441
+ )
1442
+ sess.execute(stmt)
1443
+
1444
+ return eval_run
1445
+
1446
+ except Exception as e:
1447
+ log_error(f"Error creating eval run: {e}")
1448
+ return None
1449
+
1450
+ def delete_eval_run(self, eval_run_id: str) -> None:
1451
+ """Delete an eval run from the database.
1452
+
1453
+ Args:
1454
+ eval_run_id (str): The ID of the eval run to delete.
1455
+ """
1456
+ try:
1457
+ table = self._get_table(table_type="evals")
1458
+ if table is None:
1459
+ return
1460
+
1461
+ with self.Session() as sess, sess.begin():
1462
+ stmt = table.delete().where(table.c.run_id == eval_run_id)
1463
+ result = sess.execute(stmt)
1464
+ if result.rowcount == 0:
1465
+ log_error(f"No eval run found with ID: {eval_run_id}")
1466
+ else:
1467
+ log_debug(f"Deleted eval run with ID: {eval_run_id}")
1468
+
1469
+ except Exception as e:
1470
+ log_error(f"Error deleting eval run {eval_run_id}: {e}")
1471
+
1472
+ def delete_eval_runs(self, eval_run_ids: List[str]) -> None:
1473
+ """Delete multiple eval runs from the database.
1474
+
1475
+ Args:
1476
+ eval_run_ids (List[str]): List of eval run IDs to delete.
1477
+ """
1478
+ try:
1479
+ table = self._get_table(table_type="evals")
1480
+ if table is None:
1481
+ return
1482
+
1483
+ with self.Session() as sess, sess.begin():
1484
+ stmt = table.delete().where(table.c.run_id.in_(eval_run_ids))
1485
+ result = sess.execute(stmt)
1486
+ if result.rowcount == 0:
1487
+ log_error(f"No eval runs found with IDs: {eval_run_ids}")
1488
+ else:
1489
+ log_debug(f"Deleted {result.rowcount} eval runs")
1490
+
1491
+ except Exception as e:
1492
+ log_error(f"Error deleting eval runs {eval_run_ids}: {e}")
1493
+
1494
+ def get_eval_run(
1495
+ self, eval_run_id: str, deserialize: Optional[bool] = True
1496
+ ) -> Optional[Union[EvalRunRecord, Dict[str, Any]]]:
1497
+ """Get an eval run from the database.
1498
+
1499
+ Args:
1500
+ eval_run_id (str): The ID of the eval run to get.
1501
+ deserialize (Optional[bool]): Whether to serialize the eval run. Defaults to True.
1502
+
1503
+ Returns:
1504
+ Optional[Union[EvalRunRecord, Dict[str, Any]]]:
1505
+ - When deserialize=True: EvalRunRecord object
1506
+ - When deserialize=False: EvalRun dictionary
1507
+
1508
+ Raises:
1509
+ Exception: If an error occurs during retrieval.
1510
+ """
1511
+ try:
1512
+ table = self._get_table(table_type="evals")
1513
+ if table is None:
1514
+ return None
1515
+
1516
+ with self.Session() as sess, sess.begin():
1517
+ stmt = select(table).where(table.c.run_id == eval_run_id)
1518
+ result = sess.execute(stmt).fetchone()
1519
+ if result is None:
1520
+ return None
1521
+
1522
+ eval_run_raw = result._mapping
1523
+ if not deserialize:
1524
+ return eval_run_raw
1525
+
1526
+ return EvalRunRecord.model_validate(eval_run_raw)
1527
+
1528
+ except Exception as e:
1529
+ log_error(f"Exception getting eval run {eval_run_id}: {e}")
1530
+ return None
1531
+
1532
+ def get_eval_runs(
1533
+ self,
1534
+ limit: Optional[int] = None,
1535
+ page: Optional[int] = None,
1536
+ sort_by: Optional[str] = None,
1537
+ sort_order: Optional[str] = None,
1538
+ agent_id: Optional[str] = None,
1539
+ team_id: Optional[str] = None,
1540
+ workflow_id: Optional[str] = None,
1541
+ model_id: Optional[str] = None,
1542
+ filter_type: Optional[EvalFilterType] = None,
1543
+ eval_type: Optional[List[EvalType]] = None,
1544
+ deserialize: Optional[bool] = True,
1545
+ ) -> Union[List[EvalRunRecord], Tuple[List[Dict[str, Any]], int]]:
1546
+ """Get all eval runs from the database.
1547
+
1548
+ Args:
1549
+ limit (Optional[int]): The maximum number of eval runs to return.
1550
+ page (Optional[int]): The page number.
1551
+ sort_by (Optional[str]): The column to sort by.
1552
+ sort_order (Optional[str]): The order to sort by.
1553
+ agent_id (Optional[str]): The ID of the agent to filter by.
1554
+ team_id (Optional[str]): The ID of the team to filter by.
1555
+ workflow_id (Optional[str]): The ID of the workflow to filter by.
1556
+ model_id (Optional[str]): The ID of the model to filter by.
1557
+ eval_type (Optional[List[EvalType]]): The type(s) of eval to filter by.
1558
+ filter_type (Optional[EvalFilterType]): Filter by component type (agent, team, workflow).
1559
+ deserialize (Optional[bool]): Whether to serialize the eval runs. Defaults to True.
1560
+ create_table_if_not_found (Optional[bool]): Whether to create the table if it doesn't exist.
1561
+
1562
+ Returns:
1563
+ Union[List[EvalRunRecord], Tuple[List[Dict[str, Any]], int]]:
1564
+ - When deserialize=True: List of EvalRunRecord objects
1565
+ - When deserialize=False: List of dictionaries
1566
+
1567
+ Raises:
1568
+ Exception: If an error occurs during retrieval.
1569
+ """
1570
+ try:
1571
+ table = self._get_table(table_type="evals")
1572
+ if table is None:
1573
+ return [] if deserialize else ([], 0)
1574
+
1575
+ with self.Session() as sess, sess.begin():
1576
+ stmt = select(table)
1577
+
1578
+ # Filtering
1579
+ if agent_id is not None:
1580
+ stmt = stmt.where(table.c.agent_id == agent_id)
1581
+ if team_id is not None:
1582
+ stmt = stmt.where(table.c.team_id == team_id)
1583
+ if workflow_id is not None:
1584
+ stmt = stmt.where(table.c.workflow_id == workflow_id)
1585
+ if model_id is not None:
1586
+ stmt = stmt.where(table.c.model_id == model_id)
1587
+ if eval_type is not None and len(eval_type) > 0:
1588
+ stmt = stmt.where(table.c.eval_type.in_(eval_type))
1589
+ if filter_type is not None:
1590
+ if filter_type == EvalFilterType.AGENT:
1591
+ stmt = stmt.where(table.c.agent_id.is_not(None))
1592
+ elif filter_type == EvalFilterType.TEAM:
1593
+ stmt = stmt.where(table.c.team_id.is_not(None))
1594
+ elif filter_type == EvalFilterType.WORKFLOW:
1595
+ stmt = stmt.where(table.c.workflow_id.is_not(None))
1596
+
1597
+ # Get total count after applying filtering
1598
+ count_stmt = select(func.count()).select_from(stmt.alias())
1599
+ total_count = sess.execute(count_stmt).scalar()
1600
+
1601
+ # Sorting
1602
+ if sort_by is None:
1603
+ stmt = stmt.order_by(table.c.created_at.desc())
1604
+ else:
1605
+ stmt = apply_sorting(stmt, table, sort_by, sort_order)
1606
+
1607
+ # Paginating
1608
+ if limit is not None:
1609
+ stmt = stmt.limit(limit)
1610
+ if page is not None:
1611
+ stmt = stmt.offset((page - 1) * limit)
1612
+
1613
+ result = sess.execute(stmt).fetchall()
1614
+ if not result:
1615
+ return [] if deserialize else ([], 0)
1616
+
1617
+ eval_runs_raw = [row._mapping for row in result]
1618
+ if not deserialize:
1619
+ return eval_runs_raw, total_count
1620
+
1621
+ return [EvalRunRecord.model_validate(row) for row in eval_runs_raw]
1622
+
1623
+ except Exception as e:
1624
+ log_error(f"Exception getting eval runs: {e}")
1625
+ return [] if deserialize else ([], 0)
1626
+
1627
+ def rename_eval_run(
1628
+ self, eval_run_id: str, name: str, deserialize: Optional[bool] = True
1629
+ ) -> Optional[Union[EvalRunRecord, Dict[str, Any]]]:
1630
+ """Upsert the name of an eval run in the database, returning raw dictionary.
1631
+
1632
+ Args:
1633
+ eval_run_id (str): The ID of the eval run to update.
1634
+ name (str): The new name of the eval run.
1635
+
1636
+ Returns:
1637
+ Optional[Dict[str, Any]]: The updated eval run, or None if the operation fails.
1638
+
1639
+ Raises:
1640
+ Exception: If an error occurs during update.
1641
+ """
1642
+ try:
1643
+ table = self._get_table(table_type="evals")
1644
+ if table is None:
1645
+ return None
1646
+
1647
+ with self.Session() as sess, sess.begin():
1648
+ stmt = (
1649
+ table.update().where(table.c.run_id == eval_run_id).values(name=name, updated_at=int(time.time()))
1650
+ )
1651
+ sess.execute(stmt)
1652
+
1653
+ eval_run_raw = self.get_eval_run(eval_run_id=eval_run_id, deserialize=deserialize)
1654
+ if not eval_run_raw or not deserialize:
1655
+ return eval_run_raw
1656
+
1657
+ return EvalRunRecord.model_validate(eval_run_raw)
1658
+
1659
+ except Exception as e:
1660
+ log_error(f"Error upserting eval run name {eval_run_id}: {e}")
1661
+ return None
1662
+
1663
+ # -- Migrations --
1664
+
1665
+ def migrate_table_from_v1_to_v2(self, v1_db_schema: str, v1_table_name: str, v1_table_type: str):
1666
+ """Migrate all content in the given table to the right v2 table"""
1667
+
1668
+ from agno.db.migrations.v1_to_v2 import (
1669
+ get_all_table_content,
1670
+ parse_agent_sessions,
1671
+ parse_memories,
1672
+ parse_team_sessions,
1673
+ parse_workflow_sessions,
1674
+ )
1675
+
1676
+ # Get all content from the old table
1677
+ old_content: list[dict[str, Any]] = get_all_table_content(
1678
+ db=self,
1679
+ db_schema=v1_db_schema,
1680
+ table_name=v1_table_name,
1681
+ )
1682
+ if not old_content:
1683
+ log_info(f"No content to migrate from table {v1_table_name}")
1684
+ return
1685
+
1686
+ # Parse the content into the new format
1687
+ memories: List[UserMemory] = []
1688
+ sessions: Sequence[Union[AgentSession, TeamSession, WorkflowSession]] = []
1689
+ if v1_table_type == "agent_sessions":
1690
+ sessions = parse_agent_sessions(old_content)
1691
+ elif v1_table_type == "team_sessions":
1692
+ sessions = parse_team_sessions(old_content)
1693
+ elif v1_table_type == "workflow_sessions":
1694
+ sessions = parse_workflow_sessions(old_content)
1695
+ elif v1_table_type == "memories":
1696
+ memories = parse_memories(old_content)
1697
+ else:
1698
+ raise ValueError(f"Invalid table type: {v1_table_type}")
1699
+
1700
+ # Insert the new content into the new table
1701
+ if v1_table_type == "agent_sessions":
1702
+ for session in sessions:
1703
+ self.upsert_session(session)
1704
+ log_info(f"Migrated {len(sessions)} Agent sessions to table: {self.session_table}")
1705
+
1706
+ elif v1_table_type == "team_sessions":
1707
+ for session in sessions:
1708
+ self.upsert_session(session)
1709
+ log_info(f"Migrated {len(sessions)} Team sessions to table: {self.session_table}")
1710
+
1711
+ elif v1_table_type == "workflow_sessions":
1712
+ for session in sessions:
1713
+ self.upsert_session(session)
1714
+ log_info(f"Migrated {len(sessions)} Workflow sessions to table: {self.session_table}")
1715
+
1716
+ elif v1_table_type == "memories":
1717
+ for memory in memories:
1718
+ self.upsert_user_memory(memory)
1719
+ log_info(f"Migrated {len(memories)} memories to table: {self.memory_table}")