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