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
agno/db/mongo/mongo.py ADDED
@@ -0,0 +1,1416 @@
1
+ import time
2
+ from datetime import date, datetime, timedelta, timezone
3
+ from typing import Any, Dict, List, Optional, Tuple, Union
4
+ from uuid import uuid4
5
+
6
+ from agno.db.base import BaseDb, SessionType
7
+ from agno.db.mongo.utils import (
8
+ apply_pagination,
9
+ apply_sorting,
10
+ bulk_upsert_metrics,
11
+ calculate_date_metrics,
12
+ create_collection_indexes,
13
+ fetch_all_sessions_data,
14
+ get_dates_to_calculate_metrics_for,
15
+ )
16
+ from agno.db.schemas.evals import EvalFilterType, EvalRunRecord, EvalType
17
+ from agno.db.schemas.knowledge import KnowledgeRow
18
+ from agno.db.schemas.memory import UserMemory
19
+ from agno.db.utils import deserialize_session_json_fields, serialize_session_json_fields
20
+ from agno.session import AgentSession, Session, TeamSession, WorkflowSession
21
+ from agno.utils.log import log_debug, log_error, log_info
22
+
23
+ try:
24
+ from pymongo import MongoClient, ReturnDocument
25
+ from pymongo.collection import Collection
26
+ from pymongo.database import Database
27
+ from pymongo.errors import OperationFailure
28
+ except ImportError:
29
+ raise ImportError("`pymongo` not installed. Please install it using `pip install pymongo`")
30
+
31
+
32
+ class MongoDb(BaseDb):
33
+ def __init__(
34
+ self,
35
+ db_client: Optional[MongoClient] = None,
36
+ db_name: Optional[str] = None,
37
+ db_url: Optional[str] = None,
38
+ session_collection: Optional[str] = None,
39
+ memory_collection: Optional[str] = None,
40
+ metrics_collection: Optional[str] = None,
41
+ eval_collection: Optional[str] = None,
42
+ knowledge_collection: Optional[str] = None,
43
+ ):
44
+ """
45
+ Interface for interacting with a MongoDB database.
46
+
47
+ Args:
48
+ db_client (Optional[MongoClient]): The MongoDB client to use.
49
+ db_name (Optional[str]): The name of the database to use.
50
+ db_url (Optional[str]): The database URL to connect to.
51
+ session_collection (Optional[str]): Name of the collection to store sessions.
52
+ memory_collection (Optional[str]): Name of the collection to store memories.
53
+ metrics_collection (Optional[str]): Name of the collection to store metrics.
54
+ eval_collection (Optional[str]): Name of the collection to store evaluation runs.
55
+ knowledge_collection (Optional[str]): Name of the collection to store knowledge documents.
56
+
57
+ Raises:
58
+ ValueError: If neither db_url nor db_client is provided.
59
+ """
60
+ super().__init__(
61
+ session_table=session_collection,
62
+ memory_table=memory_collection,
63
+ metrics_table=metrics_collection,
64
+ eval_table=eval_collection,
65
+ knowledge_table=knowledge_collection,
66
+ )
67
+
68
+ _client: Optional[MongoClient] = db_client
69
+ if _client is None and db_url is not None:
70
+ _client = MongoClient(db_url)
71
+ if _client is None:
72
+ raise ValueError("One of db_url or db_client must be provided")
73
+
74
+ self.db_url: Optional[str] = db_url
75
+ self.db_client: MongoClient = _client
76
+ self.db_name: str = db_name if db_name is not None else "agno"
77
+
78
+ self._database: Optional[Database] = None
79
+
80
+ @property
81
+ def database(self) -> Database:
82
+ if self._database is None:
83
+ self._database = self.db_client[self.db_name]
84
+ return self._database
85
+
86
+ # -- DB methods --
87
+
88
+ def _get_collection(
89
+ self, table_type: str, create_collection_if_not_found: Optional[bool] = True
90
+ ) -> Optional[Collection]:
91
+ """Get or create a collection based on table type.
92
+
93
+ Args:
94
+ table_type (str): The type of table to get or create.
95
+
96
+ Returns:
97
+ Collection: The collection object.
98
+ """
99
+ if table_type == "sessions":
100
+ if not hasattr(self, "session_collection"):
101
+ if self.session_table_name is None:
102
+ raise ValueError("Session collection was not provided on initialization")
103
+ self.session_collection = self._get_or_create_collection(
104
+ collection_name=self.session_table_name,
105
+ collection_type="sessions",
106
+ create_collection_if_not_found=create_collection_if_not_found,
107
+ )
108
+ return self.session_collection
109
+
110
+ if table_type == "memories":
111
+ if not hasattr(self, "memory_collection"):
112
+ if self.memory_table_name is None:
113
+ raise ValueError("Memory collection was not provided on initialization")
114
+ self.memory_collection = self._get_or_create_collection(
115
+ collection_name=self.memory_table_name,
116
+ collection_type="memories",
117
+ create_collection_if_not_found=create_collection_if_not_found,
118
+ )
119
+ return self.memory_collection
120
+
121
+ if table_type == "metrics":
122
+ if not hasattr(self, "metrics_collection"):
123
+ if self.metrics_table_name is None:
124
+ raise ValueError("Metrics collection was not provided on initialization")
125
+ self.metrics_collection = self._get_or_create_collection(
126
+ collection_name=self.metrics_table_name,
127
+ collection_type="metrics",
128
+ create_collection_if_not_found=create_collection_if_not_found,
129
+ )
130
+ return self.metrics_collection
131
+
132
+ if table_type == "evals":
133
+ if not hasattr(self, "eval_collection"):
134
+ if self.eval_table_name is None:
135
+ raise ValueError("Eval collection was not provided on initialization")
136
+ self.eval_collection = self._get_or_create_collection(
137
+ collection_name=self.eval_table_name,
138
+ collection_type="evals",
139
+ create_collection_if_not_found=create_collection_if_not_found,
140
+ )
141
+ return self.eval_collection
142
+
143
+ if table_type == "knowledge":
144
+ if not hasattr(self, "knowledge_collection"):
145
+ if self.knowledge_table_name is None:
146
+ raise ValueError("Knowledge collection was not provided on initialization")
147
+ self.knowledge_collection = self._get_or_create_collection(
148
+ collection_name=self.knowledge_table_name,
149
+ collection_type="knowledge",
150
+ create_collection_if_not_found=create_collection_if_not_found,
151
+ )
152
+ return self.knowledge_collection
153
+
154
+ raise ValueError(f"Unknown table type: {table_type}")
155
+
156
+ def _get_or_create_collection(
157
+ self, collection_name: str, collection_type: str, create_collection_if_not_found: Optional[bool] = True
158
+ ) -> Optional[Collection]:
159
+ """Get or create a collection with proper indexes.
160
+
161
+ Args:
162
+ collection_name (str): The name of the collection to get or create.
163
+ collection_type (str): The type of collection to get or create.
164
+ create_collection_if_not_found (Optional[bool]): Whether to create the collection if it doesn't exist.
165
+
166
+ Returns:
167
+ Optional[Collection]: The collection object.
168
+ """
169
+ try:
170
+ collection = self.database[collection_name]
171
+
172
+ if not hasattr(self, f"_{collection_name}_initialized"):
173
+ if not create_collection_if_not_found:
174
+ return None
175
+ create_collection_indexes(collection, collection_type)
176
+ setattr(self, f"_{collection_name}_initialized", True)
177
+ log_debug(f"Initialized collection '{collection_name}'")
178
+ else:
179
+ log_debug(f"Collection '{collection_name}' already initialized")
180
+
181
+ return collection
182
+
183
+ except Exception as e:
184
+ log_error(f"Error getting collection {collection_name}: {e}")
185
+ raise
186
+
187
+ # -- Session methods --
188
+
189
+ def delete_session(self, session_id: str) -> bool:
190
+ """Delete a session from the database.
191
+
192
+ Args:
193
+ session_id (str): The ID of the session to delete.
194
+
195
+ Returns:
196
+ bool: True if the session was deleted, False otherwise.
197
+
198
+ Raises:
199
+ Exception: If there is an error deleting the session.
200
+ """
201
+ try:
202
+ collection = self._get_collection(table_type="sessions")
203
+ if collection is None:
204
+ return False
205
+
206
+ result = collection.delete_one({"session_id": session_id})
207
+ if result.deleted_count == 0:
208
+ log_debug(f"No session found to delete with session_id: {session_id}")
209
+ return False
210
+ else:
211
+ log_debug(f"Successfully deleted session with session_id: {session_id}")
212
+ return True
213
+
214
+ except Exception as e:
215
+ log_error(f"Error deleting session: {e}")
216
+ return False
217
+
218
+ def delete_sessions(self, session_ids: List[str]) -> None:
219
+ """Delete multiple sessions from the database.
220
+
221
+ Args:
222
+ session_ids (List[str]): The IDs of the sessions to delete.
223
+ """
224
+ try:
225
+ collection = self._get_collection(table_type="sessions")
226
+ if collection is None:
227
+ return
228
+
229
+ result = collection.delete_many({"session_id": {"$in": session_ids}})
230
+ log_debug(f"Successfully deleted {result.deleted_count} sessions")
231
+
232
+ except Exception as e:
233
+ log_error(f"Error deleting sessions: {e}")
234
+
235
+ def get_session(
236
+ self,
237
+ session_id: str,
238
+ session_type: SessionType,
239
+ user_id: Optional[str] = None,
240
+ deserialize: Optional[bool] = True,
241
+ ) -> Optional[Union[Session, Dict[str, Any]]]:
242
+ """Read a session from the database.
243
+
244
+ Args:
245
+ session_id (str): The ID of the session to get.
246
+ user_id (Optional[str]): The ID of the user to get the session for.
247
+ session_type (Optional[SessionType]): The type of session to get.
248
+ deserialize (Optional[bool]): Whether to serialize the session. Defaults to True.
249
+
250
+ Returns:
251
+ Union[Session, Dict[str, Any], None]:
252
+ - When deserialize=True: Session object
253
+ - When deserialize=False: Session dictionary
254
+
255
+ Raises:
256
+ Exception: If there is an error reading the session.
257
+ """
258
+ try:
259
+ collection = self._get_collection(table_type="sessions")
260
+ if collection is None:
261
+ return None
262
+
263
+ query = {"session_id": session_id}
264
+ if user_id is not None:
265
+ query["user_id"] = user_id
266
+ if session_type is not None:
267
+ query["session_type"] = session_type
268
+
269
+ result = collection.find_one(query)
270
+ if result is None:
271
+ return None
272
+
273
+ session = deserialize_session_json_fields(result)
274
+
275
+ if not deserialize:
276
+ return session
277
+
278
+ if session_type == SessionType.AGENT.value:
279
+ return AgentSession.from_dict(session)
280
+ elif session_type == SessionType.TEAM.value:
281
+ return TeamSession.from_dict(session)
282
+ else:
283
+ return WorkflowSession.from_dict(session)
284
+
285
+ except Exception as e:
286
+ log_error(f"Exception reading session: {e}")
287
+ return None
288
+
289
+ def get_sessions(
290
+ self,
291
+ session_type: Optional[SessionType] = None,
292
+ user_id: Optional[str] = None,
293
+ component_id: Optional[str] = None,
294
+ session_name: Optional[str] = None,
295
+ start_timestamp: Optional[int] = None,
296
+ end_timestamp: Optional[int] = None,
297
+ limit: Optional[int] = None,
298
+ page: Optional[int] = None,
299
+ sort_by: Optional[str] = None,
300
+ sort_order: Optional[str] = None,
301
+ deserialize: Optional[bool] = True,
302
+ ) -> Union[List[Session], Tuple[List[Dict[str, Any]], int]]:
303
+ """Get all sessions.
304
+
305
+ Args:
306
+ session_type (Optional[SessionType]): The type of session to get.
307
+ user_id (Optional[str]): The ID of the user to get the session for.
308
+ component_id (Optional[str]): The ID of the component to get the session for.
309
+ session_name (Optional[str]): The name of the session to filter by.
310
+ start_timestamp (Optional[int]): The start timestamp to filter sessions by.
311
+ end_timestamp (Optional[int]): The end timestamp to filter sessions by.
312
+ limit (Optional[int]): The limit of the sessions to get.
313
+ page (Optional[int]): The page number to get.
314
+ sort_by (Optional[str]): The field to sort the sessions by.
315
+ sort_order (Optional[str]): The order to sort the sessions by.
316
+ deserialize (Optional[bool]): Whether to serialize the sessions. Defaults to True.
317
+ create_table_if_not_found (Optional[bool]): Whether to create the collection if it doesn't exist.
318
+
319
+ Returns:
320
+ Union[List[AgentSession], List[TeamSession], List[WorkflowSession], Tuple[List[Dict[str, Any]], int]]:
321
+ - When deserialize=True: List of Session objects
322
+ - When deserialize=False: List of session dictionaries and the total count
323
+
324
+ Raises:
325
+ Exception: If there is an error reading the sessions.
326
+ """
327
+ try:
328
+ collection = self._get_collection(table_type="sessions")
329
+ if collection is None:
330
+ return [] if deserialize else ([], 0)
331
+
332
+ # Filtering
333
+ query: Dict[str, Any] = {}
334
+ if user_id is not None:
335
+ query["user_id"] = user_id
336
+ if session_type is not None:
337
+ query["session_type"] = session_type
338
+ if component_id is not None:
339
+ if session_type == SessionType.AGENT:
340
+ query["agent_id"] = component_id
341
+ elif session_type == SessionType.TEAM:
342
+ query["team_id"] = component_id
343
+ elif session_type == SessionType.WORKFLOW:
344
+ query["workflow_id"] = component_id
345
+ if start_timestamp is not None:
346
+ query["created_at"] = {"$gte": start_timestamp}
347
+ if end_timestamp is not None:
348
+ if "created_at" in query:
349
+ query["created_at"]["$lte"] = end_timestamp
350
+ else:
351
+ query["created_at"] = {"$lte": end_timestamp}
352
+ if session_name is not None:
353
+ query["session_data.session_name"] = {"$regex": session_name, "$options": "i"}
354
+
355
+ # Get total count
356
+ total_count = collection.count_documents(query)
357
+
358
+ cursor = collection.find(query)
359
+
360
+ # Sorting
361
+ sort_criteria = apply_sorting({}, sort_by, sort_order)
362
+ if sort_criteria:
363
+ cursor = cursor.sort(sort_criteria)
364
+
365
+ # Pagination
366
+ query_args = apply_pagination({}, limit, page)
367
+ if query_args.get("skip"):
368
+ cursor = cursor.skip(query_args["skip"])
369
+ if query_args.get("limit"):
370
+ cursor = cursor.limit(query_args["limit"])
371
+
372
+ records = list(cursor)
373
+ if records is None:
374
+ return [] if deserialize else ([], 0)
375
+
376
+ sessions_raw = [deserialize_session_json_fields(record) for record in records]
377
+
378
+ if not deserialize:
379
+ return sessions_raw, total_count
380
+
381
+ sessions: List[Union[AgentSession, TeamSession, WorkflowSession]] = []
382
+ for record in sessions_raw:
383
+ if session_type == SessionType.AGENT.value:
384
+ agent_session = AgentSession.from_dict(record)
385
+ if agent_session is not None:
386
+ sessions.append(agent_session)
387
+ elif session_type == SessionType.TEAM.value:
388
+ team_session = TeamSession.from_dict(record)
389
+ if team_session is not None:
390
+ sessions.append(team_session)
391
+ elif session_type == SessionType.WORKFLOW.value:
392
+ workflow_session = WorkflowSession.from_dict(record)
393
+ if workflow_session is not None:
394
+ sessions.append(workflow_session)
395
+
396
+ return sessions
397
+
398
+ except Exception as e:
399
+ log_error(f"Exception reading sessions: {e}")
400
+ return [] if deserialize else ([], 0)
401
+
402
+ def rename_session(
403
+ self, session_id: str, session_type: SessionType, session_name: str, deserialize: Optional[bool] = True
404
+ ) -> Optional[Union[Session, Dict[str, Any]]]:
405
+ """Rename a session in the database.
406
+
407
+ Args:
408
+ session_id (str): The ID of the session to rename.
409
+ session_type (SessionType): The type of session to rename.
410
+ session_name (str): The new name of the session.
411
+ deserialize (Optional[bool]): Whether to serialize the session. Defaults to True.
412
+
413
+ Returns:
414
+ Optional[Union[Session, Dict[str, Any]]]:
415
+ - When deserialize=True: Session object
416
+ - When deserialize=False: Session dictionary
417
+
418
+ Raises:
419
+ Exception: If there is an error renaming the session.
420
+ """
421
+ try:
422
+ collection = self._get_collection(table_type="sessions")
423
+ if collection is None:
424
+ return None
425
+
426
+ try:
427
+ result = collection.find_one_and_update(
428
+ {"session_id": session_id},
429
+ {"$set": {"session_data.session_name": session_name, "updated_at": int(time.time())}},
430
+ return_document=ReturnDocument.AFTER,
431
+ upsert=False,
432
+ )
433
+ except OperationFailure:
434
+ # If the update fails because session_data doesn't contain a session_name yet, we initialize session_data
435
+ result = collection.find_one_and_update(
436
+ {"session_id": session_id},
437
+ {"$set": {"session_data": {"session_name": session_name}, "updated_at": int(time.time())}},
438
+ return_document=ReturnDocument.AFTER,
439
+ upsert=False,
440
+ )
441
+ if not result:
442
+ return None
443
+
444
+ deserialized_session = deserialize_session_json_fields(result)
445
+
446
+ if not deserialize:
447
+ return deserialized_session
448
+
449
+ if session_type == SessionType.AGENT.value:
450
+ return AgentSession.from_dict(deserialized_session)
451
+ elif session_type == SessionType.TEAM.value:
452
+ return TeamSession.from_dict(deserialized_session)
453
+ else:
454
+ return WorkflowSession.from_dict(deserialized_session)
455
+
456
+ except Exception as e:
457
+ log_error(f"Exception renaming session: {e}")
458
+ return None
459
+
460
+ def upsert_session(
461
+ self, session: Session, deserialize: Optional[bool] = True
462
+ ) -> Optional[Union[Session, Dict[str, Any]]]:
463
+ """Insert or update a session in the database.
464
+
465
+ Args:
466
+ session (Session): The session to upsert.
467
+
468
+ Returns:
469
+ Optional[Session]: The upserted session.
470
+
471
+ Raises:
472
+ Exception: If there is an error upserting the session.
473
+ """
474
+ try:
475
+ collection = self._get_collection(table_type="sessions", create_collection_if_not_found=True)
476
+ if collection is None:
477
+ return None
478
+
479
+ serialized_session_dict = serialize_session_json_fields(session.to_dict())
480
+
481
+ if isinstance(session, AgentSession):
482
+ record = {
483
+ "session_id": serialized_session_dict.get("session_id"),
484
+ "session_type": SessionType.AGENT.value,
485
+ "agent_id": serialized_session_dict.get("agent_id"),
486
+ "user_id": serialized_session_dict.get("user_id"),
487
+ "runs": serialized_session_dict.get("runs"),
488
+ "agent_data": serialized_session_dict.get("agent_data"),
489
+ "session_data": serialized_session_dict.get("session_data"),
490
+ "summary": serialized_session_dict.get("summary"),
491
+ "metadata": serialized_session_dict.get("metadata"),
492
+ "created_at": serialized_session_dict.get("created_at"),
493
+ "updated_at": int(time.time()),
494
+ }
495
+
496
+ result = collection.find_one_and_replace(
497
+ filter={"session_id": serialized_session_dict.get("session_id")},
498
+ replacement=record,
499
+ upsert=True,
500
+ return_document=ReturnDocument.AFTER,
501
+ )
502
+ if not result:
503
+ return None
504
+
505
+ session = deserialize_session_json_fields(result) # type: ignore
506
+
507
+ if not deserialize:
508
+ return session
509
+
510
+ return AgentSession.from_dict(session) # type: ignore
511
+
512
+ elif isinstance(session, TeamSession):
513
+ record = {
514
+ "session_id": serialized_session_dict.get("session_id"),
515
+ "session_type": SessionType.TEAM.value,
516
+ "team_id": serialized_session_dict.get("team_id"),
517
+ "user_id": serialized_session_dict.get("user_id"),
518
+ "runs": serialized_session_dict.get("runs"),
519
+ "team_data": serialized_session_dict.get("team_data"),
520
+ "session_data": serialized_session_dict.get("session_data"),
521
+ "summary": serialized_session_dict.get("summary"),
522
+ "metadata": serialized_session_dict.get("metadata"),
523
+ "created_at": serialized_session_dict.get("created_at"),
524
+ "updated_at": int(time.time()),
525
+ }
526
+
527
+ result = collection.find_one_and_replace(
528
+ filter={"session_id": serialized_session_dict.get("session_id")},
529
+ replacement=record,
530
+ upsert=True,
531
+ return_document=ReturnDocument.AFTER,
532
+ )
533
+ if not result:
534
+ return None
535
+
536
+ session = deserialize_session_json_fields(result) # type: ignore
537
+
538
+ if not deserialize:
539
+ return session
540
+
541
+ return TeamSession.from_dict(session) # type: ignore
542
+
543
+ else:
544
+ record = {
545
+ "session_id": serialized_session_dict.get("session_id"),
546
+ "session_type": SessionType.WORKFLOW.value,
547
+ "workflow_id": serialized_session_dict.get("workflow_id"),
548
+ "user_id": serialized_session_dict.get("user_id"),
549
+ "runs": serialized_session_dict.get("runs"),
550
+ "workflow_data": serialized_session_dict.get("workflow_data"),
551
+ "session_data": serialized_session_dict.get("session_data"),
552
+ "summary": serialized_session_dict.get("summary"),
553
+ "metadata": serialized_session_dict.get("metadata"),
554
+ "created_at": serialized_session_dict.get("created_at"),
555
+ "updated_at": int(time.time()),
556
+ }
557
+
558
+ result = collection.find_one_and_replace(
559
+ filter={"session_id": serialized_session_dict.get("session_id")},
560
+ replacement=record,
561
+ upsert=True,
562
+ return_document=ReturnDocument.AFTER,
563
+ )
564
+ if not result:
565
+ return None
566
+
567
+ session = deserialize_session_json_fields(result) # type: ignore
568
+
569
+ if not deserialize:
570
+ return session
571
+
572
+ return WorkflowSession.from_dict(session) # type: ignore
573
+
574
+ except Exception as e:
575
+ log_error(f"Exception upserting session: {e}")
576
+ return None
577
+
578
+ # -- Memory methods --
579
+
580
+ def delete_user_memory(self, memory_id: str):
581
+ """Delete a user memory from the database.
582
+
583
+ Args:
584
+ memory_id (str): The ID of the memory to delete.
585
+
586
+ Returns:
587
+ bool: True if the memory was deleted, False otherwise.
588
+
589
+ Raises:
590
+ Exception: If there is an error deleting the memory.
591
+ """
592
+ try:
593
+ collection = self._get_collection(table_type="memories")
594
+ if collection is None:
595
+ return
596
+
597
+ result = collection.delete_one({"memory_id": memory_id})
598
+
599
+ success = result.deleted_count > 0
600
+ if success:
601
+ log_debug(f"Successfully deleted memory id: {memory_id}")
602
+ else:
603
+ log_debug(f"No memory found with id: {memory_id}")
604
+
605
+ except Exception as e:
606
+ log_error(f"Error deleting memory: {e}")
607
+
608
+ def delete_user_memories(self, memory_ids: List[str]) -> None:
609
+ """Delete user memories from the database.
610
+
611
+ Args:
612
+ memory_ids (List[str]): The IDs of the memories to delete.
613
+
614
+ Raises:
615
+ Exception: If there is an error deleting the memories.
616
+ """
617
+ try:
618
+ collection = self._get_collection(table_type="memories")
619
+ if collection is None:
620
+ return
621
+
622
+ result = collection.delete_many({"memory_id": {"$in": memory_ids}})
623
+
624
+ if result.deleted_count == 0:
625
+ log_debug(f"No memories found with ids: {memory_ids}")
626
+
627
+ except Exception as e:
628
+ log_error(f"Error deleting memories: {e}")
629
+
630
+ def get_all_memory_topics(self) -> List[str]:
631
+ """Get all memory topics from the database.
632
+
633
+ Returns:
634
+ List[str]: The topics.
635
+
636
+ Raises:
637
+ Exception: If there is an error getting the topics.
638
+ """
639
+ try:
640
+ collection = self._get_collection(table_type="memories")
641
+ if collection is None:
642
+ return []
643
+
644
+ topics = collection.distinct("topics")
645
+ return [topic for topic in topics if topic]
646
+
647
+ except Exception as e:
648
+ log_error(f"Exception reading from collection: {e}")
649
+ return []
650
+
651
+ def get_user_memory(self, memory_id: str, deserialize: Optional[bool] = True) -> Optional[UserMemory]:
652
+ """Get a memory from the database.
653
+
654
+ Args:
655
+ memory_id (str): The ID of the memory to get.
656
+ deserialize (Optional[bool]): Whether to serialize the memory. Defaults to True.
657
+
658
+ Returns:
659
+ Optional[UserMemory]:
660
+ - When deserialize=True: UserMemory object
661
+ - When deserialize=False: Memory dictionary
662
+
663
+ Raises:
664
+ Exception: If there is an error getting the memory.
665
+ """
666
+ try:
667
+ collection = self._get_collection(table_type="memories")
668
+ if collection is None:
669
+ return None
670
+
671
+ result = collection.find_one({"memory_id": memory_id})
672
+ if result is None or not deserialize:
673
+ return result
674
+
675
+ # Remove MongoDB's _id field before creating UserMemory object
676
+ result_filtered = {k: v for k, v in result.items() if k != "_id"}
677
+ return UserMemory.from_dict(result_filtered)
678
+
679
+ except Exception as e:
680
+ log_error(f"Exception reading from collection: {e}")
681
+ return None
682
+
683
+ def get_user_memories(
684
+ self,
685
+ user_id: Optional[str] = None,
686
+ agent_id: Optional[str] = None,
687
+ team_id: Optional[str] = None,
688
+ topics: Optional[List[str]] = None,
689
+ search_content: Optional[str] = None,
690
+ limit: Optional[int] = None,
691
+ page: Optional[int] = None,
692
+ sort_by: Optional[str] = None,
693
+ sort_order: Optional[str] = None,
694
+ deserialize: Optional[bool] = True,
695
+ ) -> Union[List[UserMemory], Tuple[List[Dict[str, Any]], int]]:
696
+ """Get all memories from the database as UserMemory objects.
697
+
698
+ Args:
699
+ user_id (Optional[str]): The ID of the user to get the memories for.
700
+ agent_id (Optional[str]): The ID of the agent to get the memories for.
701
+ team_id (Optional[str]): The ID of the team to get the memories for.
702
+ topics (Optional[List[str]]): The topics to filter the memories by.
703
+ search_content (Optional[str]): The content to filter the memories by.
704
+ limit (Optional[int]): The limit of the memories to get.
705
+ page (Optional[int]): The page number to get.
706
+ sort_by (Optional[str]): The field to sort the memories by.
707
+ sort_order (Optional[str]): The order to sort the memories by.
708
+ deserialize (Optional[bool]): Whether to serialize the memories. Defaults to True.
709
+ create_table_if_not_found: Whether to create the collection if it doesn't exist.
710
+
711
+ Returns:
712
+ Tuple[List[Dict[str, Any]], int]: A tuple containing the memories and the total count.
713
+
714
+ Raises:
715
+ Exception: If there is an error getting the memories.
716
+ """
717
+ try:
718
+ collection = self._get_collection(table_type="memories")
719
+ if collection is None:
720
+ return [] if deserialize else ([], 0)
721
+
722
+ query: Dict[str, Any] = {}
723
+ if user_id is not None:
724
+ query["user_id"] = user_id
725
+ if agent_id is not None:
726
+ query["agent_id"] = agent_id
727
+ if team_id is not None:
728
+ query["team_id"] = team_id
729
+ if topics is not None:
730
+ query["topics"] = {"$in": topics}
731
+ if search_content is not None:
732
+ query["memory"] = {"$regex": search_content, "$options": "i"}
733
+
734
+ # Get total count
735
+ total_count = collection.count_documents(query)
736
+
737
+ # Apply sorting
738
+ sort_criteria = apply_sorting({}, sort_by, sort_order)
739
+
740
+ # Apply pagination
741
+ query_args = apply_pagination({}, limit, page)
742
+
743
+ cursor = collection.find(query)
744
+ if sort_criteria:
745
+ cursor = cursor.sort(sort_criteria)
746
+ if query_args.get("skip"):
747
+ cursor = cursor.skip(query_args["skip"])
748
+ if query_args.get("limit"):
749
+ cursor = cursor.limit(query_args["limit"])
750
+
751
+ records = list(cursor)
752
+ if not deserialize:
753
+ return records, total_count
754
+
755
+ # Remove MongoDB's _id field before creating UserMemory objects
756
+ return [UserMemory.from_dict({k: v for k, v in record.items() if k != "_id"}) for record in records]
757
+
758
+ except Exception as e:
759
+ log_error(f"Exception reading from collection: {e}")
760
+ return []
761
+
762
+ def get_user_memory_stats(
763
+ self,
764
+ limit: Optional[int] = None,
765
+ page: Optional[int] = None,
766
+ ) -> Tuple[List[Dict[str, Any]], int]:
767
+ """Get user memories stats.
768
+
769
+ Args:
770
+ limit (Optional[int]): The limit of the memories to get.
771
+ page (Optional[int]): The page number to get.
772
+
773
+ Returns:
774
+ Tuple[List[Dict[str, Any]], int]: A tuple containing the memories stats and the total count.
775
+
776
+ Raises:
777
+ Exception: If there is an error getting the memories stats.
778
+ """
779
+ try:
780
+ collection = self._get_collection(table_type="memories")
781
+ if collection is None:
782
+ return [], 0
783
+
784
+ pipeline = [
785
+ {"$match": {"user_id": {"$ne": None}}},
786
+ {
787
+ "$group": {
788
+ "_id": "$user_id",
789
+ "total_memories": {"$sum": 1},
790
+ "last_memory_updated_at": {"$max": "$updated_at"},
791
+ }
792
+ },
793
+ {"$sort": {"last_memory_updated_at": -1}},
794
+ ]
795
+
796
+ # Get total count
797
+ count_pipeline = pipeline + [{"$count": "total"}]
798
+ count_result = list(collection.aggregate(count_pipeline)) # type: ignore
799
+ total_count = count_result[0]["total"] if count_result else 0
800
+
801
+ # Apply pagination
802
+ if limit is not None:
803
+ if page is not None:
804
+ pipeline.append({"$skip": (page - 1) * limit})
805
+ pipeline.append({"$limit": limit})
806
+
807
+ results = list(collection.aggregate(pipeline)) # type: ignore
808
+
809
+ formatted_results = [
810
+ {
811
+ "user_id": result["_id"],
812
+ "total_memories": result["total_memories"],
813
+ "last_memory_updated_at": result["last_memory_updated_at"],
814
+ }
815
+ for result in results
816
+ ]
817
+
818
+ return formatted_results, total_count
819
+
820
+ except Exception as e:
821
+ log_error(f"Exception getting user memory stats: {e}")
822
+ return [], 0
823
+
824
+ def upsert_user_memory(
825
+ self, memory: UserMemory, deserialize: Optional[bool] = True
826
+ ) -> Optional[Union[UserMemory, Dict[str, Any]]]:
827
+ """Upsert a user memory in the database.
828
+
829
+ Args:
830
+ memory (UserMemory): The memory to upsert.
831
+ deserialize (Optional[bool]): Whether to serialize the memory. Defaults to True.
832
+
833
+ Returns:
834
+ Optional[Union[UserMemory, Dict[str, Any]]]:
835
+ - When deserialize=True: UserMemory object
836
+ - When deserialize=False: Memory dictionary
837
+
838
+ Raises:
839
+ Exception: If there is an error upserting the memory.
840
+ """
841
+ try:
842
+ collection = self._get_collection(table_type="memories", create_collection_if_not_found=True)
843
+ if collection is None:
844
+ return None
845
+
846
+ if memory.memory_id is None:
847
+ memory.memory_id = str(uuid4())
848
+
849
+ update_doc = {
850
+ "user_id": memory.user_id,
851
+ "agent_id": memory.agent_id,
852
+ "team_id": memory.team_id,
853
+ "memory_id": memory.memory_id,
854
+ "memory": memory.memory,
855
+ "topics": memory.topics,
856
+ "updated_at": int(time.time()),
857
+ }
858
+
859
+ result = collection.replace_one({"memory_id": memory.memory_id}, update_doc, upsert=True)
860
+
861
+ if result.upserted_id:
862
+ update_doc["_id"] = result.upserted_id
863
+
864
+ if not deserialize:
865
+ return update_doc
866
+
867
+ # Remove MongoDB's _id field before creating UserMemory object
868
+ update_doc_filtered = {k: v for k, v in update_doc.items() if k != "_id"}
869
+ return UserMemory.from_dict(update_doc_filtered)
870
+
871
+ except Exception as e:
872
+ log_error(f"Exception upserting user memory: {e}")
873
+ return None
874
+
875
+ def clear_memories(self) -> None:
876
+ """Delete all memories from the database.
877
+
878
+ Raises:
879
+ Exception: If an error occurs during deletion.
880
+ """
881
+ try:
882
+ collection = self._get_collection(table_type="memories")
883
+ if collection is None:
884
+ return
885
+
886
+ collection.delete_many({})
887
+
888
+ except Exception as e:
889
+ from agno.utils.log import log_warning
890
+
891
+ log_warning(f"Exception deleting all memories: {e}")
892
+
893
+ # -- Metrics methods --
894
+
895
+ def _get_all_sessions_for_metrics_calculation(
896
+ self, start_timestamp: Optional[int] = None, end_timestamp: Optional[int] = None
897
+ ) -> List[Dict[str, Any]]:
898
+ """Get all sessions of all types for metrics calculation."""
899
+ try:
900
+ collection = self._get_collection(table_type="sessions")
901
+ if collection is None:
902
+ return []
903
+
904
+ query = {}
905
+ if start_timestamp is not None:
906
+ query["created_at"] = {"$gte": start_timestamp}
907
+ if end_timestamp is not None:
908
+ if "created_at" in query:
909
+ query["created_at"]["$lte"] = end_timestamp
910
+ else:
911
+ query["created_at"] = {"$lte": end_timestamp}
912
+
913
+ projection = {
914
+ "user_id": 1,
915
+ "session_data": 1,
916
+ "runs": 1,
917
+ "created_at": 1,
918
+ "session_type": 1,
919
+ }
920
+
921
+ results = list(collection.find(query, projection))
922
+ return results
923
+
924
+ except Exception as e:
925
+ log_error(f"Exception reading from sessions collection: {e}")
926
+ return []
927
+
928
+ def _get_metrics_calculation_starting_date(self, collection: Collection) -> Optional[date]:
929
+ """Get the first date for which metrics calculation is needed."""
930
+ try:
931
+ result = collection.find_one({}, sort=[("date", -1)], limit=1)
932
+
933
+ if result is not None:
934
+ result_date = datetime.strptime(result["date"], "%Y-%m-%d").date()
935
+ if result.get("completed"):
936
+ return result_date + timedelta(days=1)
937
+ else:
938
+ return result_date
939
+
940
+ # No metrics records. Return the date of the first recorded session.
941
+ first_session_result = self.get_sessions(sort_by="created_at", sort_order="asc", limit=1, deserialize=False)
942
+ first_session_date = first_session_result[0][0]["created_at"] if first_session_result[0] else None # type: ignore
943
+
944
+ if first_session_date is None:
945
+ return None
946
+
947
+ return datetime.fromtimestamp(first_session_date, tz=timezone.utc).date()
948
+
949
+ except Exception as e:
950
+ log_error(f"Exception getting metrics calculation starting date: {e}")
951
+ return None
952
+
953
+ def calculate_metrics(self) -> Optional[list[dict]]:
954
+ """Calculate metrics for all dates without complete metrics."""
955
+ try:
956
+ collection = self._get_collection(table_type="metrics", create_collection_if_not_found=True)
957
+ if collection is None:
958
+ return None
959
+
960
+ starting_date = self._get_metrics_calculation_starting_date(collection)
961
+ if starting_date is None:
962
+ log_info("No session data found. Won't calculate metrics.")
963
+ return None
964
+
965
+ dates_to_process = get_dates_to_calculate_metrics_for(starting_date)
966
+ if not dates_to_process:
967
+ log_info("Metrics already calculated for all relevant dates.")
968
+ return None
969
+
970
+ start_timestamp = int(
971
+ datetime.combine(dates_to_process[0], datetime.min.time()).replace(tzinfo=timezone.utc).timestamp()
972
+ )
973
+ end_timestamp = int(
974
+ datetime.combine(dates_to_process[-1] + timedelta(days=1), datetime.min.time())
975
+ .replace(tzinfo=timezone.utc)
976
+ .timestamp()
977
+ )
978
+
979
+ sessions = self._get_all_sessions_for_metrics_calculation(
980
+ start_timestamp=start_timestamp, end_timestamp=end_timestamp
981
+ )
982
+ all_sessions_data = fetch_all_sessions_data(
983
+ sessions=sessions, dates_to_process=dates_to_process, start_timestamp=start_timestamp
984
+ )
985
+ if not all_sessions_data:
986
+ log_info("No new session data found. Won't calculate metrics.")
987
+ return None
988
+
989
+ results = []
990
+ metrics_records = []
991
+
992
+ for date_to_process in dates_to_process:
993
+ date_key = date_to_process.isoformat()
994
+ sessions_for_date = all_sessions_data.get(date_key, {})
995
+
996
+ # Skip dates with no sessions
997
+ if not any(len(sessions) > 0 for sessions in sessions_for_date.values()):
998
+ continue
999
+
1000
+ metrics_record = calculate_date_metrics(date_to_process, sessions_for_date)
1001
+ metrics_records.append(metrics_record)
1002
+
1003
+ if metrics_records:
1004
+ results = bulk_upsert_metrics(collection, metrics_records)
1005
+
1006
+ return results
1007
+
1008
+ except Exception as e:
1009
+ log_error(f"Error calculating metrics: {e}")
1010
+ raise e
1011
+
1012
+ def get_metrics(
1013
+ self,
1014
+ starting_date: Optional[date] = None,
1015
+ ending_date: Optional[date] = None,
1016
+ ) -> Tuple[List[dict], Optional[int]]:
1017
+ """Get all metrics matching the given date range."""
1018
+ try:
1019
+ collection = self._get_collection(table_type="metrics")
1020
+ if collection is None:
1021
+ return [], None
1022
+
1023
+ query = {}
1024
+ if starting_date:
1025
+ query["date"] = {"$gte": starting_date.isoformat()}
1026
+ if ending_date:
1027
+ if "date" in query:
1028
+ query["date"]["$lte"] = ending_date.isoformat()
1029
+ else:
1030
+ query["date"] = {"$lte": ending_date.isoformat()}
1031
+
1032
+ records = list(collection.find(query))
1033
+ if not records:
1034
+ return [], None
1035
+
1036
+ # Get the latest updated_at
1037
+ latest_updated_at = max(record.get("updated_at", 0) for record in records)
1038
+
1039
+ return records, latest_updated_at
1040
+
1041
+ except Exception as e:
1042
+ log_error(f"Error getting metrics: {e}")
1043
+ return [], None
1044
+
1045
+ # -- Knowledge methods --
1046
+
1047
+ def delete_knowledge_content(self, id: str):
1048
+ """Delete a knowledge row from the database.
1049
+
1050
+ Args:
1051
+ id (str): The ID of the knowledge row to delete.
1052
+
1053
+ Raises:
1054
+ Exception: If an error occurs during deletion.
1055
+ """
1056
+ try:
1057
+ collection = self._get_collection(table_type="knowledge")
1058
+ if collection is None:
1059
+ return
1060
+
1061
+ collection.delete_one({"id": id})
1062
+
1063
+ log_debug(f"Deleted knowledge content with id '{id}'")
1064
+
1065
+ except Exception as e:
1066
+ log_error(f"Error deleting knowledge content: {e}")
1067
+ raise
1068
+
1069
+ def get_knowledge_content(self, id: str) -> Optional[KnowledgeRow]:
1070
+ """Get a knowledge row from the database.
1071
+
1072
+ Args:
1073
+ id (str): The ID of the knowledge row to get.
1074
+
1075
+ Returns:
1076
+ Optional[KnowledgeRow]: The knowledge row, or None if it doesn't exist.
1077
+
1078
+ Raises:
1079
+ Exception: If an error occurs during retrieval.
1080
+ """
1081
+ try:
1082
+ collection = self._get_collection(table_type="knowledge")
1083
+ if collection is None:
1084
+ return None
1085
+
1086
+ result = collection.find_one({"id": id})
1087
+ if result is None:
1088
+ return None
1089
+
1090
+ return KnowledgeRow.model_validate(result)
1091
+
1092
+ except Exception as e:
1093
+ log_error(f"Error getting knowledge content: {e}")
1094
+ return None
1095
+
1096
+ def get_knowledge_contents(
1097
+ self,
1098
+ limit: Optional[int] = None,
1099
+ page: Optional[int] = None,
1100
+ sort_by: Optional[str] = None,
1101
+ sort_order: Optional[str] = None,
1102
+ ) -> Tuple[List[KnowledgeRow], int]:
1103
+ """Get all knowledge contents from the database.
1104
+
1105
+ Args:
1106
+ limit (Optional[int]): The maximum number of knowledge contents to return.
1107
+ page (Optional[int]): The page number.
1108
+ sort_by (Optional[str]): The column to sort by.
1109
+ sort_order (Optional[str]): The order to sort by.
1110
+ create_table_if_not_found (Optional[bool]): Whether to create the collection if it doesn't exist.
1111
+
1112
+ Returns:
1113
+ Tuple[List[KnowledgeRow], int]: The knowledge contents and total count.
1114
+
1115
+ Raises:
1116
+ Exception: If an error occurs during retrieval.
1117
+ """
1118
+ try:
1119
+ collection = self._get_collection(table_type="knowledge")
1120
+ if collection is None:
1121
+ return [], 0
1122
+
1123
+ query: Dict[str, Any] = {}
1124
+
1125
+ # Get total count
1126
+ total_count = collection.count_documents(query)
1127
+
1128
+ # Apply sorting
1129
+ sort_criteria = apply_sorting({}, sort_by, sort_order)
1130
+
1131
+ # Apply pagination
1132
+ query_args = apply_pagination({}, limit, page)
1133
+
1134
+ cursor = collection.find(query)
1135
+ if sort_criteria:
1136
+ cursor = cursor.sort(sort_criteria)
1137
+ if query_args.get("skip"):
1138
+ cursor = cursor.skip(query_args["skip"])
1139
+ if query_args.get("limit"):
1140
+ cursor = cursor.limit(query_args["limit"])
1141
+
1142
+ records = list(cursor)
1143
+ knowledge_rows = [KnowledgeRow.model_validate(record) for record in records]
1144
+
1145
+ return knowledge_rows, total_count
1146
+
1147
+ except Exception as e:
1148
+ log_error(f"Error getting knowledge contents: {e}")
1149
+ return [], 0
1150
+
1151
+ def upsert_knowledge_content(self, knowledge_row: KnowledgeRow):
1152
+ """Upsert knowledge content in the database.
1153
+
1154
+ Args:
1155
+ knowledge_row (KnowledgeRow): The knowledge row to upsert.
1156
+
1157
+ Returns:
1158
+ Optional[KnowledgeRow]: The upserted knowledge row, or None if the operation fails.
1159
+
1160
+ Raises:
1161
+ Exception: If an error occurs during upsert.
1162
+ """
1163
+ try:
1164
+ collection = self._get_collection(table_type="knowledge", create_collection_if_not_found=True)
1165
+ if collection is None:
1166
+ return None
1167
+
1168
+ update_doc = knowledge_row.model_dump()
1169
+ collection.replace_one({"id": knowledge_row.id}, update_doc, upsert=True)
1170
+
1171
+ return knowledge_row
1172
+
1173
+ except Exception as e:
1174
+ log_error(f"Error upserting knowledge content: {e}")
1175
+ return None
1176
+
1177
+ # -- Eval methods --
1178
+
1179
+ def create_eval_run(self, eval_run: EvalRunRecord) -> Optional[EvalRunRecord]:
1180
+ """Create an EvalRunRecord in the database."""
1181
+ try:
1182
+ collection = self._get_collection(table_type="evals", create_collection_if_not_found=True)
1183
+ if collection is None:
1184
+ return None
1185
+
1186
+ current_time = int(time.time())
1187
+ eval_dict = eval_run.model_dump()
1188
+ eval_dict["created_at"] = current_time
1189
+ eval_dict["updated_at"] = current_time
1190
+
1191
+ collection.insert_one(eval_dict)
1192
+
1193
+ log_debug(f"Created eval run with id '{eval_run.run_id}'")
1194
+
1195
+ return eval_run
1196
+
1197
+ except Exception as e:
1198
+ log_error(f"Error creating eval run: {e}")
1199
+ return None
1200
+
1201
+ def delete_eval_run(self, eval_run_id: str) -> None:
1202
+ """Delete an eval run from the database."""
1203
+ try:
1204
+ collection = self._get_collection(table_type="evals")
1205
+ if collection is None:
1206
+ return
1207
+
1208
+ result = collection.delete_one({"run_id": eval_run_id})
1209
+
1210
+ if result.deleted_count == 0:
1211
+ log_debug(f"No eval run found with ID: {eval_run_id}")
1212
+ else:
1213
+ log_debug(f"Deleted eval run with ID: {eval_run_id}")
1214
+
1215
+ except Exception as e:
1216
+ log_error(f"Error deleting eval run {eval_run_id}: {e}")
1217
+ raise
1218
+
1219
+ def delete_eval_runs(self, eval_run_ids: List[str]) -> None:
1220
+ """Delete multiple eval runs from the database."""
1221
+ try:
1222
+ collection = self._get_collection(table_type="evals")
1223
+ if collection is None:
1224
+ return
1225
+
1226
+ result = collection.delete_many({"run_id": {"$in": eval_run_ids}})
1227
+
1228
+ if result.deleted_count == 0:
1229
+ log_debug(f"No eval runs found with IDs: {eval_run_ids}")
1230
+ else:
1231
+ log_debug(f"Deleted {result.deleted_count} eval runs")
1232
+
1233
+ except Exception as e:
1234
+ log_error(f"Error deleting eval runs {eval_run_ids}: {e}")
1235
+ raise
1236
+
1237
+ def get_eval_run_raw(self, eval_run_id: str) -> Optional[Dict[str, Any]]:
1238
+ """Get an eval run from the database as a raw dictionary."""
1239
+ try:
1240
+ collection = self._get_collection(table_type="evals")
1241
+ if collection is None:
1242
+ return None
1243
+
1244
+ result = collection.find_one({"run_id": eval_run_id})
1245
+ return result
1246
+
1247
+ except Exception as e:
1248
+ log_error(f"Exception getting eval run {eval_run_id}: {e}")
1249
+ return None
1250
+
1251
+ def get_eval_run(self, eval_run_id: str, deserialize: Optional[bool] = True) -> Optional[EvalRunRecord]:
1252
+ """Get an eval run from the database.
1253
+
1254
+ Args:
1255
+ eval_run_id (str): The ID of the eval run to get.
1256
+ deserialize (Optional[bool]): Whether to serialize the eval run. Defaults to True.
1257
+
1258
+ Returns:
1259
+ Optional[EvalRunRecord]:
1260
+ - When deserialize=True: EvalRunRecord object
1261
+ - When deserialize=False: EvalRun dictionary
1262
+
1263
+ Raises:
1264
+ Exception: If there is an error getting the eval run.
1265
+ """
1266
+ try:
1267
+ collection = self._get_collection(table_type="evals")
1268
+ if collection is None:
1269
+ return None
1270
+
1271
+ eval_run_raw = collection.find_one({"run_id": eval_run_id})
1272
+
1273
+ if not eval_run_raw:
1274
+ return None
1275
+
1276
+ if not deserialize:
1277
+ return eval_run_raw
1278
+
1279
+ return EvalRunRecord.model_validate(eval_run_raw)
1280
+
1281
+ except Exception as e:
1282
+ log_error(f"Exception getting eval run {eval_run_id}: {e}")
1283
+ return None
1284
+
1285
+ def get_eval_runs(
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
+ agent_id: Optional[str] = None,
1292
+ team_id: Optional[str] = None,
1293
+ workflow_id: Optional[str] = None,
1294
+ model_id: Optional[str] = None,
1295
+ filter_type: Optional[EvalFilterType] = None,
1296
+ eval_type: Optional[List[EvalType]] = None,
1297
+ deserialize: Optional[bool] = True,
1298
+ ) -> Union[List[EvalRunRecord], Tuple[List[Dict[str, Any]], int]]:
1299
+ """Get all eval runs from the database.
1300
+
1301
+ Args:
1302
+ limit (Optional[int]): The maximum number of eval runs to return.
1303
+ page (Optional[int]): The page number to return.
1304
+ sort_by (Optional[str]): The field to sort by.
1305
+ sort_order (Optional[str]): The order to sort by.
1306
+ agent_id (Optional[str]): The ID of the agent to filter by.
1307
+ team_id (Optional[str]): The ID of the team to filter by.
1308
+ workflow_id (Optional[str]): The ID of the workflow to filter by.
1309
+ model_id (Optional[str]): The ID of the model to filter by.
1310
+ eval_type (Optional[List[EvalType]]): The type of eval to filter by.
1311
+ filter_type (Optional[EvalFilterType]): The type of filter to apply.
1312
+ deserialize (Optional[bool]): Whether to serialize the eval runs. Defaults to True.
1313
+ create_table_if_not_found (Optional[bool]): Whether to create the collection if it doesn't exist.
1314
+
1315
+ Returns:
1316
+ Union[List[EvalRunRecord], Tuple[List[Dict[str, Any]], int]]:
1317
+ - When deserialize=True: List of EvalRunRecord objects
1318
+ - When deserialize=False: List of eval run dictionaries and the total count
1319
+
1320
+ Raises:
1321
+ Exception: If there is an error getting the eval runs.
1322
+ """
1323
+ try:
1324
+ collection = self._get_collection(table_type="evals")
1325
+ if collection is None:
1326
+ return [] if deserialize else ([], 0)
1327
+
1328
+ query: Dict[str, Any] = {}
1329
+ if agent_id is not None:
1330
+ query["agent_id"] = agent_id
1331
+ if team_id is not None:
1332
+ query["team_id"] = team_id
1333
+ if workflow_id is not None:
1334
+ query["workflow_id"] = workflow_id
1335
+ if model_id is not None:
1336
+ query["model_id"] = model_id
1337
+ if eval_type is not None and len(eval_type) > 0:
1338
+ query["eval_type"] = {"$in": eval_type}
1339
+ if filter_type is not None:
1340
+ if filter_type == EvalFilterType.AGENT:
1341
+ query["agent_id"] = {"$ne": None}
1342
+ elif filter_type == EvalFilterType.TEAM:
1343
+ query["team_id"] = {"$ne": None}
1344
+ elif filter_type == EvalFilterType.WORKFLOW:
1345
+ query["workflow_id"] = {"$ne": None}
1346
+
1347
+ # Get total count
1348
+ total_count = collection.count_documents(query)
1349
+
1350
+ # Apply default sorting by created_at desc if no sort parameters provided
1351
+ if sort_by is None:
1352
+ sort_criteria = [("created_at", -1)]
1353
+ else:
1354
+ sort_criteria = apply_sorting({}, sort_by, sort_order)
1355
+
1356
+ # Apply pagination
1357
+ query_args = apply_pagination({}, limit, page)
1358
+
1359
+ cursor = collection.find(query)
1360
+ if sort_criteria:
1361
+ cursor = cursor.sort(sort_criteria)
1362
+ if query_args.get("skip"):
1363
+ cursor = cursor.skip(query_args["skip"])
1364
+ if query_args.get("limit"):
1365
+ cursor = cursor.limit(query_args["limit"])
1366
+
1367
+ records = list(cursor)
1368
+ if not records:
1369
+ return [] if deserialize else ([], 0)
1370
+
1371
+ if not deserialize:
1372
+ return records, total_count
1373
+
1374
+ return [EvalRunRecord.model_validate(row) for row in records]
1375
+
1376
+ except Exception as e:
1377
+ log_debug(f"Exception getting eval runs: {e}")
1378
+ return [] if deserialize else ([], 0)
1379
+
1380
+ def rename_eval_run(
1381
+ self, eval_run_id: str, name: str, deserialize: Optional[bool] = True
1382
+ ) -> Optional[Union[EvalRunRecord, Dict[str, Any]]]:
1383
+ """Update the name of an eval run in the database.
1384
+
1385
+ Args:
1386
+ eval_run_id (str): The ID of the eval run to update.
1387
+ name (str): The new name of the eval run.
1388
+ deserialize (Optional[bool]): Whether to serialize the eval run. Defaults to True.
1389
+
1390
+ Returns:
1391
+ Optional[Union[EvalRunRecord, Dict[str, Any]]]:
1392
+ - When deserialize=True: EvalRunRecord object
1393
+ - When deserialize=False: EvalRun dictionary
1394
+
1395
+ Raises:
1396
+ Exception: If there is an error updating the eval run.
1397
+ """
1398
+ try:
1399
+ collection = self._get_collection(table_type="evals")
1400
+ if collection is None:
1401
+ return None
1402
+
1403
+ result = collection.find_one_and_update(
1404
+ {"run_id": eval_run_id}, {"$set": {"name": name, "updated_at": int(time.time())}}
1405
+ )
1406
+
1407
+ log_debug(f"Renamed eval run with id '{eval_run_id}' to '{name}'")
1408
+
1409
+ if not result or not deserialize:
1410
+ return result
1411
+
1412
+ return EvalRunRecord.model_validate(result)
1413
+
1414
+ except Exception as e:
1415
+ log_error(f"Error updating eval run name {eval_run_id}: {e}")
1416
+ raise