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