agno 2.2.13__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 (575) hide show
  1. agno/__init__.py +8 -0
  2. agno/agent/__init__.py +51 -0
  3. agno/agent/agent.py +10405 -0
  4. agno/api/__init__.py +0 -0
  5. agno/api/agent.py +28 -0
  6. agno/api/api.py +40 -0
  7. agno/api/evals.py +22 -0
  8. agno/api/os.py +17 -0
  9. agno/api/routes.py +13 -0
  10. agno/api/schemas/__init__.py +9 -0
  11. agno/api/schemas/agent.py +16 -0
  12. agno/api/schemas/evals.py +16 -0
  13. agno/api/schemas/os.py +14 -0
  14. agno/api/schemas/response.py +6 -0
  15. agno/api/schemas/team.py +16 -0
  16. agno/api/schemas/utils.py +21 -0
  17. agno/api/schemas/workflows.py +16 -0
  18. agno/api/settings.py +53 -0
  19. agno/api/team.py +30 -0
  20. agno/api/workflow.py +28 -0
  21. agno/cloud/aws/base.py +214 -0
  22. agno/cloud/aws/s3/__init__.py +2 -0
  23. agno/cloud/aws/s3/api_client.py +43 -0
  24. agno/cloud/aws/s3/bucket.py +195 -0
  25. agno/cloud/aws/s3/object.py +57 -0
  26. agno/culture/__init__.py +3 -0
  27. agno/culture/manager.py +956 -0
  28. agno/db/__init__.py +24 -0
  29. agno/db/async_postgres/__init__.py +3 -0
  30. agno/db/base.py +598 -0
  31. agno/db/dynamo/__init__.py +3 -0
  32. agno/db/dynamo/dynamo.py +2042 -0
  33. agno/db/dynamo/schemas.py +314 -0
  34. agno/db/dynamo/utils.py +743 -0
  35. agno/db/firestore/__init__.py +3 -0
  36. agno/db/firestore/firestore.py +1795 -0
  37. agno/db/firestore/schemas.py +140 -0
  38. agno/db/firestore/utils.py +376 -0
  39. agno/db/gcs_json/__init__.py +3 -0
  40. agno/db/gcs_json/gcs_json_db.py +1335 -0
  41. agno/db/gcs_json/utils.py +228 -0
  42. agno/db/in_memory/__init__.py +3 -0
  43. agno/db/in_memory/in_memory_db.py +1160 -0
  44. agno/db/in_memory/utils.py +230 -0
  45. agno/db/json/__init__.py +3 -0
  46. agno/db/json/json_db.py +1328 -0
  47. agno/db/json/utils.py +230 -0
  48. agno/db/migrations/__init__.py +0 -0
  49. agno/db/migrations/v1_to_v2.py +635 -0
  50. agno/db/mongo/__init__.py +17 -0
  51. agno/db/mongo/async_mongo.py +2026 -0
  52. agno/db/mongo/mongo.py +1982 -0
  53. agno/db/mongo/schemas.py +87 -0
  54. agno/db/mongo/utils.py +259 -0
  55. agno/db/mysql/__init__.py +3 -0
  56. agno/db/mysql/mysql.py +2308 -0
  57. agno/db/mysql/schemas.py +138 -0
  58. agno/db/mysql/utils.py +355 -0
  59. agno/db/postgres/__init__.py +4 -0
  60. agno/db/postgres/async_postgres.py +1927 -0
  61. agno/db/postgres/postgres.py +2260 -0
  62. agno/db/postgres/schemas.py +139 -0
  63. agno/db/postgres/utils.py +442 -0
  64. agno/db/redis/__init__.py +3 -0
  65. agno/db/redis/redis.py +1660 -0
  66. agno/db/redis/schemas.py +123 -0
  67. agno/db/redis/utils.py +346 -0
  68. agno/db/schemas/__init__.py +4 -0
  69. agno/db/schemas/culture.py +120 -0
  70. agno/db/schemas/evals.py +33 -0
  71. agno/db/schemas/knowledge.py +40 -0
  72. agno/db/schemas/memory.py +46 -0
  73. agno/db/schemas/metrics.py +0 -0
  74. agno/db/singlestore/__init__.py +3 -0
  75. agno/db/singlestore/schemas.py +130 -0
  76. agno/db/singlestore/singlestore.py +2272 -0
  77. agno/db/singlestore/utils.py +384 -0
  78. agno/db/sqlite/__init__.py +4 -0
  79. agno/db/sqlite/async_sqlite.py +2293 -0
  80. agno/db/sqlite/schemas.py +133 -0
  81. agno/db/sqlite/sqlite.py +2288 -0
  82. agno/db/sqlite/utils.py +431 -0
  83. agno/db/surrealdb/__init__.py +3 -0
  84. agno/db/surrealdb/metrics.py +292 -0
  85. agno/db/surrealdb/models.py +309 -0
  86. agno/db/surrealdb/queries.py +71 -0
  87. agno/db/surrealdb/surrealdb.py +1353 -0
  88. agno/db/surrealdb/utils.py +147 -0
  89. agno/db/utils.py +116 -0
  90. agno/debug.py +18 -0
  91. agno/eval/__init__.py +14 -0
  92. agno/eval/accuracy.py +834 -0
  93. agno/eval/performance.py +773 -0
  94. agno/eval/reliability.py +306 -0
  95. agno/eval/utils.py +119 -0
  96. agno/exceptions.py +161 -0
  97. agno/filters.py +354 -0
  98. agno/guardrails/__init__.py +6 -0
  99. agno/guardrails/base.py +19 -0
  100. agno/guardrails/openai.py +144 -0
  101. agno/guardrails/pii.py +94 -0
  102. agno/guardrails/prompt_injection.py +52 -0
  103. agno/integrations/__init__.py +0 -0
  104. agno/integrations/discord/__init__.py +3 -0
  105. agno/integrations/discord/client.py +203 -0
  106. agno/knowledge/__init__.py +5 -0
  107. agno/knowledge/chunking/__init__.py +0 -0
  108. agno/knowledge/chunking/agentic.py +79 -0
  109. agno/knowledge/chunking/document.py +91 -0
  110. agno/knowledge/chunking/fixed.py +57 -0
  111. agno/knowledge/chunking/markdown.py +151 -0
  112. agno/knowledge/chunking/recursive.py +63 -0
  113. agno/knowledge/chunking/row.py +39 -0
  114. agno/knowledge/chunking/semantic.py +86 -0
  115. agno/knowledge/chunking/strategy.py +165 -0
  116. agno/knowledge/content.py +74 -0
  117. agno/knowledge/document/__init__.py +5 -0
  118. agno/knowledge/document/base.py +58 -0
  119. agno/knowledge/embedder/__init__.py +5 -0
  120. agno/knowledge/embedder/aws_bedrock.py +343 -0
  121. agno/knowledge/embedder/azure_openai.py +210 -0
  122. agno/knowledge/embedder/base.py +23 -0
  123. agno/knowledge/embedder/cohere.py +323 -0
  124. agno/knowledge/embedder/fastembed.py +62 -0
  125. agno/knowledge/embedder/fireworks.py +13 -0
  126. agno/knowledge/embedder/google.py +258 -0
  127. agno/knowledge/embedder/huggingface.py +94 -0
  128. agno/knowledge/embedder/jina.py +182 -0
  129. agno/knowledge/embedder/langdb.py +22 -0
  130. agno/knowledge/embedder/mistral.py +206 -0
  131. agno/knowledge/embedder/nebius.py +13 -0
  132. agno/knowledge/embedder/ollama.py +154 -0
  133. agno/knowledge/embedder/openai.py +195 -0
  134. agno/knowledge/embedder/sentence_transformer.py +63 -0
  135. agno/knowledge/embedder/together.py +13 -0
  136. agno/knowledge/embedder/vllm.py +262 -0
  137. agno/knowledge/embedder/voyageai.py +165 -0
  138. agno/knowledge/knowledge.py +1988 -0
  139. agno/knowledge/reader/__init__.py +7 -0
  140. agno/knowledge/reader/arxiv_reader.py +81 -0
  141. agno/knowledge/reader/base.py +95 -0
  142. agno/knowledge/reader/csv_reader.py +166 -0
  143. agno/knowledge/reader/docx_reader.py +82 -0
  144. agno/knowledge/reader/field_labeled_csv_reader.py +292 -0
  145. agno/knowledge/reader/firecrawl_reader.py +201 -0
  146. agno/knowledge/reader/json_reader.py +87 -0
  147. agno/knowledge/reader/markdown_reader.py +137 -0
  148. agno/knowledge/reader/pdf_reader.py +431 -0
  149. agno/knowledge/reader/pptx_reader.py +101 -0
  150. agno/knowledge/reader/reader_factory.py +313 -0
  151. agno/knowledge/reader/s3_reader.py +89 -0
  152. agno/knowledge/reader/tavily_reader.py +194 -0
  153. agno/knowledge/reader/text_reader.py +115 -0
  154. agno/knowledge/reader/web_search_reader.py +372 -0
  155. agno/knowledge/reader/website_reader.py +455 -0
  156. agno/knowledge/reader/wikipedia_reader.py +59 -0
  157. agno/knowledge/reader/youtube_reader.py +78 -0
  158. agno/knowledge/remote_content/__init__.py +0 -0
  159. agno/knowledge/remote_content/remote_content.py +88 -0
  160. agno/knowledge/reranker/__init__.py +3 -0
  161. agno/knowledge/reranker/base.py +14 -0
  162. agno/knowledge/reranker/cohere.py +64 -0
  163. agno/knowledge/reranker/infinity.py +195 -0
  164. agno/knowledge/reranker/sentence_transformer.py +54 -0
  165. agno/knowledge/types.py +39 -0
  166. agno/knowledge/utils.py +189 -0
  167. agno/media.py +462 -0
  168. agno/memory/__init__.py +3 -0
  169. agno/memory/manager.py +1327 -0
  170. agno/models/__init__.py +0 -0
  171. agno/models/aimlapi/__init__.py +5 -0
  172. agno/models/aimlapi/aimlapi.py +45 -0
  173. agno/models/anthropic/__init__.py +5 -0
  174. agno/models/anthropic/claude.py +757 -0
  175. agno/models/aws/__init__.py +15 -0
  176. agno/models/aws/bedrock.py +701 -0
  177. agno/models/aws/claude.py +378 -0
  178. agno/models/azure/__init__.py +18 -0
  179. agno/models/azure/ai_foundry.py +485 -0
  180. agno/models/azure/openai_chat.py +131 -0
  181. agno/models/base.py +2175 -0
  182. agno/models/cerebras/__init__.py +12 -0
  183. agno/models/cerebras/cerebras.py +501 -0
  184. agno/models/cerebras/cerebras_openai.py +112 -0
  185. agno/models/cohere/__init__.py +5 -0
  186. agno/models/cohere/chat.py +389 -0
  187. agno/models/cometapi/__init__.py +5 -0
  188. agno/models/cometapi/cometapi.py +57 -0
  189. agno/models/dashscope/__init__.py +5 -0
  190. agno/models/dashscope/dashscope.py +91 -0
  191. agno/models/deepinfra/__init__.py +5 -0
  192. agno/models/deepinfra/deepinfra.py +28 -0
  193. agno/models/deepseek/__init__.py +5 -0
  194. agno/models/deepseek/deepseek.py +61 -0
  195. agno/models/defaults.py +1 -0
  196. agno/models/fireworks/__init__.py +5 -0
  197. agno/models/fireworks/fireworks.py +26 -0
  198. agno/models/google/__init__.py +5 -0
  199. agno/models/google/gemini.py +1085 -0
  200. agno/models/groq/__init__.py +5 -0
  201. agno/models/groq/groq.py +556 -0
  202. agno/models/huggingface/__init__.py +5 -0
  203. agno/models/huggingface/huggingface.py +491 -0
  204. agno/models/ibm/__init__.py +5 -0
  205. agno/models/ibm/watsonx.py +422 -0
  206. agno/models/internlm/__init__.py +3 -0
  207. agno/models/internlm/internlm.py +26 -0
  208. agno/models/langdb/__init__.py +1 -0
  209. agno/models/langdb/langdb.py +48 -0
  210. agno/models/litellm/__init__.py +14 -0
  211. agno/models/litellm/chat.py +468 -0
  212. agno/models/litellm/litellm_openai.py +25 -0
  213. agno/models/llama_cpp/__init__.py +5 -0
  214. agno/models/llama_cpp/llama_cpp.py +22 -0
  215. agno/models/lmstudio/__init__.py +5 -0
  216. agno/models/lmstudio/lmstudio.py +25 -0
  217. agno/models/message.py +434 -0
  218. agno/models/meta/__init__.py +12 -0
  219. agno/models/meta/llama.py +475 -0
  220. agno/models/meta/llama_openai.py +78 -0
  221. agno/models/metrics.py +120 -0
  222. agno/models/mistral/__init__.py +5 -0
  223. agno/models/mistral/mistral.py +432 -0
  224. agno/models/nebius/__init__.py +3 -0
  225. agno/models/nebius/nebius.py +54 -0
  226. agno/models/nexus/__init__.py +3 -0
  227. agno/models/nexus/nexus.py +22 -0
  228. agno/models/nvidia/__init__.py +5 -0
  229. agno/models/nvidia/nvidia.py +28 -0
  230. agno/models/ollama/__init__.py +5 -0
  231. agno/models/ollama/chat.py +441 -0
  232. agno/models/openai/__init__.py +9 -0
  233. agno/models/openai/chat.py +883 -0
  234. agno/models/openai/like.py +27 -0
  235. agno/models/openai/responses.py +1050 -0
  236. agno/models/openrouter/__init__.py +5 -0
  237. agno/models/openrouter/openrouter.py +66 -0
  238. agno/models/perplexity/__init__.py +5 -0
  239. agno/models/perplexity/perplexity.py +187 -0
  240. agno/models/portkey/__init__.py +3 -0
  241. agno/models/portkey/portkey.py +81 -0
  242. agno/models/requesty/__init__.py +5 -0
  243. agno/models/requesty/requesty.py +52 -0
  244. agno/models/response.py +199 -0
  245. agno/models/sambanova/__init__.py +5 -0
  246. agno/models/sambanova/sambanova.py +28 -0
  247. agno/models/siliconflow/__init__.py +5 -0
  248. agno/models/siliconflow/siliconflow.py +25 -0
  249. agno/models/together/__init__.py +5 -0
  250. agno/models/together/together.py +25 -0
  251. agno/models/utils.py +266 -0
  252. agno/models/vercel/__init__.py +3 -0
  253. agno/models/vercel/v0.py +26 -0
  254. agno/models/vertexai/__init__.py +0 -0
  255. agno/models/vertexai/claude.py +70 -0
  256. agno/models/vllm/__init__.py +3 -0
  257. agno/models/vllm/vllm.py +78 -0
  258. agno/models/xai/__init__.py +3 -0
  259. agno/models/xai/xai.py +113 -0
  260. agno/os/__init__.py +3 -0
  261. agno/os/app.py +876 -0
  262. agno/os/auth.py +57 -0
  263. agno/os/config.py +104 -0
  264. agno/os/interfaces/__init__.py +1 -0
  265. agno/os/interfaces/a2a/__init__.py +3 -0
  266. agno/os/interfaces/a2a/a2a.py +42 -0
  267. agno/os/interfaces/a2a/router.py +250 -0
  268. agno/os/interfaces/a2a/utils.py +924 -0
  269. agno/os/interfaces/agui/__init__.py +3 -0
  270. agno/os/interfaces/agui/agui.py +47 -0
  271. agno/os/interfaces/agui/router.py +144 -0
  272. agno/os/interfaces/agui/utils.py +534 -0
  273. agno/os/interfaces/base.py +25 -0
  274. agno/os/interfaces/slack/__init__.py +3 -0
  275. agno/os/interfaces/slack/router.py +148 -0
  276. agno/os/interfaces/slack/security.py +30 -0
  277. agno/os/interfaces/slack/slack.py +47 -0
  278. agno/os/interfaces/whatsapp/__init__.py +3 -0
  279. agno/os/interfaces/whatsapp/router.py +211 -0
  280. agno/os/interfaces/whatsapp/security.py +53 -0
  281. agno/os/interfaces/whatsapp/whatsapp.py +36 -0
  282. agno/os/mcp.py +292 -0
  283. agno/os/middleware/__init__.py +7 -0
  284. agno/os/middleware/jwt.py +233 -0
  285. agno/os/router.py +1763 -0
  286. agno/os/routers/__init__.py +3 -0
  287. agno/os/routers/evals/__init__.py +3 -0
  288. agno/os/routers/evals/evals.py +430 -0
  289. agno/os/routers/evals/schemas.py +142 -0
  290. agno/os/routers/evals/utils.py +162 -0
  291. agno/os/routers/health.py +31 -0
  292. agno/os/routers/home.py +52 -0
  293. agno/os/routers/knowledge/__init__.py +3 -0
  294. agno/os/routers/knowledge/knowledge.py +997 -0
  295. agno/os/routers/knowledge/schemas.py +178 -0
  296. agno/os/routers/memory/__init__.py +3 -0
  297. agno/os/routers/memory/memory.py +515 -0
  298. agno/os/routers/memory/schemas.py +62 -0
  299. agno/os/routers/metrics/__init__.py +3 -0
  300. agno/os/routers/metrics/metrics.py +190 -0
  301. agno/os/routers/metrics/schemas.py +47 -0
  302. agno/os/routers/session/__init__.py +3 -0
  303. agno/os/routers/session/session.py +997 -0
  304. agno/os/schema.py +1055 -0
  305. agno/os/settings.py +43 -0
  306. agno/os/utils.py +630 -0
  307. agno/py.typed +0 -0
  308. agno/reasoning/__init__.py +0 -0
  309. agno/reasoning/anthropic.py +80 -0
  310. agno/reasoning/azure_ai_foundry.py +67 -0
  311. agno/reasoning/deepseek.py +63 -0
  312. agno/reasoning/default.py +97 -0
  313. agno/reasoning/gemini.py +73 -0
  314. agno/reasoning/groq.py +71 -0
  315. agno/reasoning/helpers.py +63 -0
  316. agno/reasoning/ollama.py +67 -0
  317. agno/reasoning/openai.py +86 -0
  318. agno/reasoning/step.py +31 -0
  319. agno/reasoning/vertexai.py +76 -0
  320. agno/run/__init__.py +6 -0
  321. agno/run/agent.py +787 -0
  322. agno/run/base.py +229 -0
  323. agno/run/cancel.py +81 -0
  324. agno/run/messages.py +32 -0
  325. agno/run/team.py +753 -0
  326. agno/run/workflow.py +708 -0
  327. agno/session/__init__.py +10 -0
  328. agno/session/agent.py +295 -0
  329. agno/session/summary.py +265 -0
  330. agno/session/team.py +392 -0
  331. agno/session/workflow.py +205 -0
  332. agno/team/__init__.py +37 -0
  333. agno/team/team.py +8793 -0
  334. agno/tools/__init__.py +10 -0
  335. agno/tools/agentql.py +120 -0
  336. agno/tools/airflow.py +69 -0
  337. agno/tools/api.py +122 -0
  338. agno/tools/apify.py +314 -0
  339. agno/tools/arxiv.py +127 -0
  340. agno/tools/aws_lambda.py +53 -0
  341. agno/tools/aws_ses.py +66 -0
  342. agno/tools/baidusearch.py +89 -0
  343. agno/tools/bitbucket.py +292 -0
  344. agno/tools/brandfetch.py +213 -0
  345. agno/tools/bravesearch.py +106 -0
  346. agno/tools/brightdata.py +367 -0
  347. agno/tools/browserbase.py +209 -0
  348. agno/tools/calcom.py +255 -0
  349. agno/tools/calculator.py +151 -0
  350. agno/tools/cartesia.py +187 -0
  351. agno/tools/clickup.py +244 -0
  352. agno/tools/confluence.py +240 -0
  353. agno/tools/crawl4ai.py +158 -0
  354. agno/tools/csv_toolkit.py +185 -0
  355. agno/tools/dalle.py +110 -0
  356. agno/tools/daytona.py +475 -0
  357. agno/tools/decorator.py +262 -0
  358. agno/tools/desi_vocal.py +108 -0
  359. agno/tools/discord.py +161 -0
  360. agno/tools/docker.py +716 -0
  361. agno/tools/duckdb.py +379 -0
  362. agno/tools/duckduckgo.py +91 -0
  363. agno/tools/e2b.py +703 -0
  364. agno/tools/eleven_labs.py +196 -0
  365. agno/tools/email.py +67 -0
  366. agno/tools/evm.py +129 -0
  367. agno/tools/exa.py +396 -0
  368. agno/tools/fal.py +127 -0
  369. agno/tools/file.py +240 -0
  370. agno/tools/file_generation.py +350 -0
  371. agno/tools/financial_datasets.py +288 -0
  372. agno/tools/firecrawl.py +143 -0
  373. agno/tools/function.py +1187 -0
  374. agno/tools/giphy.py +93 -0
  375. agno/tools/github.py +1760 -0
  376. agno/tools/gmail.py +922 -0
  377. agno/tools/google_bigquery.py +117 -0
  378. agno/tools/google_drive.py +270 -0
  379. agno/tools/google_maps.py +253 -0
  380. agno/tools/googlecalendar.py +674 -0
  381. agno/tools/googlesearch.py +98 -0
  382. agno/tools/googlesheets.py +377 -0
  383. agno/tools/hackernews.py +77 -0
  384. agno/tools/jina.py +101 -0
  385. agno/tools/jira.py +170 -0
  386. agno/tools/knowledge.py +218 -0
  387. agno/tools/linear.py +426 -0
  388. agno/tools/linkup.py +58 -0
  389. agno/tools/local_file_system.py +90 -0
  390. agno/tools/lumalab.py +183 -0
  391. agno/tools/mcp/__init__.py +10 -0
  392. agno/tools/mcp/mcp.py +331 -0
  393. agno/tools/mcp/multi_mcp.py +347 -0
  394. agno/tools/mcp/params.py +24 -0
  395. agno/tools/mcp_toolbox.py +284 -0
  396. agno/tools/mem0.py +193 -0
  397. agno/tools/memori.py +339 -0
  398. agno/tools/memory.py +419 -0
  399. agno/tools/mlx_transcribe.py +139 -0
  400. agno/tools/models/__init__.py +0 -0
  401. agno/tools/models/azure_openai.py +190 -0
  402. agno/tools/models/gemini.py +203 -0
  403. agno/tools/models/groq.py +158 -0
  404. agno/tools/models/morph.py +186 -0
  405. agno/tools/models/nebius.py +124 -0
  406. agno/tools/models_labs.py +195 -0
  407. agno/tools/moviepy_video.py +349 -0
  408. agno/tools/neo4j.py +134 -0
  409. agno/tools/newspaper.py +46 -0
  410. agno/tools/newspaper4k.py +93 -0
  411. agno/tools/notion.py +204 -0
  412. agno/tools/openai.py +202 -0
  413. agno/tools/openbb.py +160 -0
  414. agno/tools/opencv.py +321 -0
  415. agno/tools/openweather.py +233 -0
  416. agno/tools/oxylabs.py +385 -0
  417. agno/tools/pandas.py +102 -0
  418. agno/tools/parallel.py +314 -0
  419. agno/tools/postgres.py +257 -0
  420. agno/tools/pubmed.py +188 -0
  421. agno/tools/python.py +205 -0
  422. agno/tools/reasoning.py +283 -0
  423. agno/tools/reddit.py +467 -0
  424. agno/tools/replicate.py +117 -0
  425. agno/tools/resend.py +62 -0
  426. agno/tools/scrapegraph.py +222 -0
  427. agno/tools/searxng.py +152 -0
  428. agno/tools/serpapi.py +116 -0
  429. agno/tools/serper.py +255 -0
  430. agno/tools/shell.py +53 -0
  431. agno/tools/slack.py +136 -0
  432. agno/tools/sleep.py +20 -0
  433. agno/tools/spider.py +116 -0
  434. agno/tools/sql.py +154 -0
  435. agno/tools/streamlit/__init__.py +0 -0
  436. agno/tools/streamlit/components.py +113 -0
  437. agno/tools/tavily.py +254 -0
  438. agno/tools/telegram.py +48 -0
  439. agno/tools/todoist.py +218 -0
  440. agno/tools/tool_registry.py +1 -0
  441. agno/tools/toolkit.py +146 -0
  442. agno/tools/trafilatura.py +388 -0
  443. agno/tools/trello.py +274 -0
  444. agno/tools/twilio.py +186 -0
  445. agno/tools/user_control_flow.py +78 -0
  446. agno/tools/valyu.py +228 -0
  447. agno/tools/visualization.py +467 -0
  448. agno/tools/webbrowser.py +28 -0
  449. agno/tools/webex.py +76 -0
  450. agno/tools/website.py +54 -0
  451. agno/tools/webtools.py +45 -0
  452. agno/tools/whatsapp.py +286 -0
  453. agno/tools/wikipedia.py +63 -0
  454. agno/tools/workflow.py +278 -0
  455. agno/tools/x.py +335 -0
  456. agno/tools/yfinance.py +257 -0
  457. agno/tools/youtube.py +184 -0
  458. agno/tools/zendesk.py +82 -0
  459. agno/tools/zep.py +454 -0
  460. agno/tools/zoom.py +382 -0
  461. agno/utils/__init__.py +0 -0
  462. agno/utils/agent.py +820 -0
  463. agno/utils/audio.py +49 -0
  464. agno/utils/certs.py +27 -0
  465. agno/utils/code_execution.py +11 -0
  466. agno/utils/common.py +132 -0
  467. agno/utils/dttm.py +13 -0
  468. agno/utils/enum.py +22 -0
  469. agno/utils/env.py +11 -0
  470. agno/utils/events.py +696 -0
  471. agno/utils/format_str.py +16 -0
  472. agno/utils/functions.py +166 -0
  473. agno/utils/gemini.py +426 -0
  474. agno/utils/hooks.py +57 -0
  475. agno/utils/http.py +74 -0
  476. agno/utils/json_schema.py +234 -0
  477. agno/utils/knowledge.py +36 -0
  478. agno/utils/location.py +19 -0
  479. agno/utils/log.py +255 -0
  480. agno/utils/mcp.py +214 -0
  481. agno/utils/media.py +352 -0
  482. agno/utils/merge_dict.py +41 -0
  483. agno/utils/message.py +118 -0
  484. agno/utils/models/__init__.py +0 -0
  485. agno/utils/models/ai_foundry.py +43 -0
  486. agno/utils/models/claude.py +358 -0
  487. agno/utils/models/cohere.py +87 -0
  488. agno/utils/models/llama.py +78 -0
  489. agno/utils/models/mistral.py +98 -0
  490. agno/utils/models/openai_responses.py +140 -0
  491. agno/utils/models/schema_utils.py +153 -0
  492. agno/utils/models/watsonx.py +41 -0
  493. agno/utils/openai.py +257 -0
  494. agno/utils/pickle.py +32 -0
  495. agno/utils/pprint.py +178 -0
  496. agno/utils/print_response/__init__.py +0 -0
  497. agno/utils/print_response/agent.py +842 -0
  498. agno/utils/print_response/team.py +1724 -0
  499. agno/utils/print_response/workflow.py +1668 -0
  500. agno/utils/prompts.py +111 -0
  501. agno/utils/reasoning.py +108 -0
  502. agno/utils/response.py +163 -0
  503. agno/utils/response_iterator.py +17 -0
  504. agno/utils/safe_formatter.py +24 -0
  505. agno/utils/serialize.py +32 -0
  506. agno/utils/shell.py +22 -0
  507. agno/utils/streamlit.py +487 -0
  508. agno/utils/string.py +231 -0
  509. agno/utils/team.py +139 -0
  510. agno/utils/timer.py +41 -0
  511. agno/utils/tools.py +102 -0
  512. agno/utils/web.py +23 -0
  513. agno/utils/whatsapp.py +305 -0
  514. agno/utils/yaml_io.py +25 -0
  515. agno/vectordb/__init__.py +3 -0
  516. agno/vectordb/base.py +127 -0
  517. agno/vectordb/cassandra/__init__.py +5 -0
  518. agno/vectordb/cassandra/cassandra.py +501 -0
  519. agno/vectordb/cassandra/extra_param_mixin.py +11 -0
  520. agno/vectordb/cassandra/index.py +13 -0
  521. agno/vectordb/chroma/__init__.py +5 -0
  522. agno/vectordb/chroma/chromadb.py +929 -0
  523. agno/vectordb/clickhouse/__init__.py +9 -0
  524. agno/vectordb/clickhouse/clickhousedb.py +835 -0
  525. agno/vectordb/clickhouse/index.py +9 -0
  526. agno/vectordb/couchbase/__init__.py +3 -0
  527. agno/vectordb/couchbase/couchbase.py +1442 -0
  528. agno/vectordb/distance.py +7 -0
  529. agno/vectordb/lancedb/__init__.py +6 -0
  530. agno/vectordb/lancedb/lance_db.py +995 -0
  531. agno/vectordb/langchaindb/__init__.py +5 -0
  532. agno/vectordb/langchaindb/langchaindb.py +163 -0
  533. agno/vectordb/lightrag/__init__.py +5 -0
  534. agno/vectordb/lightrag/lightrag.py +388 -0
  535. agno/vectordb/llamaindex/__init__.py +3 -0
  536. agno/vectordb/llamaindex/llamaindexdb.py +166 -0
  537. agno/vectordb/milvus/__init__.py +4 -0
  538. agno/vectordb/milvus/milvus.py +1182 -0
  539. agno/vectordb/mongodb/__init__.py +9 -0
  540. agno/vectordb/mongodb/mongodb.py +1417 -0
  541. agno/vectordb/pgvector/__init__.py +12 -0
  542. agno/vectordb/pgvector/index.py +23 -0
  543. agno/vectordb/pgvector/pgvector.py +1462 -0
  544. agno/vectordb/pineconedb/__init__.py +5 -0
  545. agno/vectordb/pineconedb/pineconedb.py +747 -0
  546. agno/vectordb/qdrant/__init__.py +5 -0
  547. agno/vectordb/qdrant/qdrant.py +1134 -0
  548. agno/vectordb/redis/__init__.py +9 -0
  549. agno/vectordb/redis/redisdb.py +694 -0
  550. agno/vectordb/search.py +7 -0
  551. agno/vectordb/singlestore/__init__.py +10 -0
  552. agno/vectordb/singlestore/index.py +41 -0
  553. agno/vectordb/singlestore/singlestore.py +763 -0
  554. agno/vectordb/surrealdb/__init__.py +3 -0
  555. agno/vectordb/surrealdb/surrealdb.py +699 -0
  556. agno/vectordb/upstashdb/__init__.py +5 -0
  557. agno/vectordb/upstashdb/upstashdb.py +718 -0
  558. agno/vectordb/weaviate/__init__.py +8 -0
  559. agno/vectordb/weaviate/index.py +15 -0
  560. agno/vectordb/weaviate/weaviate.py +1005 -0
  561. agno/workflow/__init__.py +23 -0
  562. agno/workflow/agent.py +299 -0
  563. agno/workflow/condition.py +738 -0
  564. agno/workflow/loop.py +735 -0
  565. agno/workflow/parallel.py +824 -0
  566. agno/workflow/router.py +702 -0
  567. agno/workflow/step.py +1432 -0
  568. agno/workflow/steps.py +592 -0
  569. agno/workflow/types.py +520 -0
  570. agno/workflow/workflow.py +4321 -0
  571. agno-2.2.13.dist-info/METADATA +614 -0
  572. agno-2.2.13.dist-info/RECORD +575 -0
  573. agno-2.2.13.dist-info/WHEEL +5 -0
  574. agno-2.2.13.dist-info/licenses/LICENSE +201 -0
  575. agno-2.2.13.dist-info/top_level.txt +1 -0
agno/db/mongo/mongo.py ADDED
@@ -0,0 +1,1982 @@
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
+ deserialize_cultural_knowledge_from_db,
14
+ fetch_all_sessions_data,
15
+ get_dates_to_calculate_metrics_for,
16
+ serialize_cultural_knowledge_for_db,
17
+ )
18
+ from agno.db.schemas.culture import CulturalKnowledge
19
+ from agno.db.schemas.evals import EvalFilterType, EvalRunRecord, EvalType
20
+ from agno.db.schemas.knowledge import KnowledgeRow
21
+ from agno.db.schemas.memory import UserMemory
22
+ from agno.db.utils import deserialize_session_json_fields
23
+ from agno.session import AgentSession, Session, TeamSession, WorkflowSession
24
+ from agno.utils.log import log_debug, log_error, log_info
25
+ from agno.utils.string import generate_id
26
+
27
+ try:
28
+ from pymongo import MongoClient, ReturnDocument
29
+ from pymongo.collection import Collection
30
+ from pymongo.database import Database
31
+ from pymongo.errors import OperationFailure
32
+ except ImportError:
33
+ raise ImportError("`pymongo` not installed. Please install it using `pip install pymongo`")
34
+
35
+
36
+ class MongoDb(BaseDb):
37
+ def __init__(
38
+ self,
39
+ db_client: Optional[MongoClient] = None,
40
+ db_name: Optional[str] = None,
41
+ db_url: Optional[str] = None,
42
+ session_collection: Optional[str] = None,
43
+ memory_collection: Optional[str] = None,
44
+ metrics_collection: Optional[str] = None,
45
+ eval_collection: Optional[str] = None,
46
+ knowledge_collection: Optional[str] = None,
47
+ culture_collection: Optional[str] = None,
48
+ id: Optional[str] = None,
49
+ ):
50
+ """
51
+ Interface for interacting with a MongoDB database.
52
+
53
+ Args:
54
+ db_client (Optional[MongoClient]): The MongoDB client to use.
55
+ db_name (Optional[str]): The name of the database to use.
56
+ db_url (Optional[str]): The database URL to connect to.
57
+ session_collection (Optional[str]): Name of the collection to store sessions.
58
+ memory_collection (Optional[str]): Name of the collection to store memories.
59
+ metrics_collection (Optional[str]): Name of the collection to store metrics.
60
+ eval_collection (Optional[str]): Name of the collection to store evaluation runs.
61
+ knowledge_collection (Optional[str]): Name of the collection to store knowledge documents.
62
+ culture_collection (Optional[str]): Name of the collection to store cultural knowledge.
63
+ id (Optional[str]): ID of the database.
64
+
65
+ Raises:
66
+ ValueError: If neither db_url nor db_client is provided.
67
+ """
68
+ if id is None:
69
+ base_seed = db_url or str(db_client)
70
+ db_name_suffix = db_name if db_name is not None else "agno"
71
+ seed = f"{base_seed}#{db_name_suffix}"
72
+ id = generate_id(seed)
73
+
74
+ super().__init__(
75
+ id=id,
76
+ session_table=session_collection,
77
+ memory_table=memory_collection,
78
+ metrics_table=metrics_collection,
79
+ eval_table=eval_collection,
80
+ knowledge_table=knowledge_collection,
81
+ culture_table=culture_collection,
82
+ )
83
+
84
+ _client: Optional[MongoClient] = db_client
85
+ if _client is None and db_url is not None:
86
+ _client = MongoClient(db_url)
87
+ if _client is None:
88
+ raise ValueError("One of db_url or db_client must be provided")
89
+
90
+ self.db_url: Optional[str] = db_url
91
+ self.db_client: MongoClient = _client
92
+ self.db_name: str = db_name if db_name is not None else "agno"
93
+
94
+ self._database: Optional[Database] = None
95
+
96
+ @property
97
+ def database(self) -> Database:
98
+ if self._database is None:
99
+ self._database = self.db_client[self.db_name]
100
+ return self._database
101
+
102
+ # -- DB methods --
103
+ def table_exists(self, table_name: str) -> bool:
104
+ """Check if a collection with the given name exists in the MongoDB database.
105
+
106
+ Args:
107
+ table_name: Name of the collection to check
108
+
109
+ Returns:
110
+ bool: True if the collection exists in the database, False otherwise
111
+ """
112
+ return table_name in self.database.list_collection_names()
113
+
114
+ def _create_all_tables(self):
115
+ """Create all configured MongoDB collections if they don't exist."""
116
+ collections_to_create = [
117
+ ("sessions", self.session_table_name),
118
+ ("memories", self.memory_table_name),
119
+ ("metrics", self.metrics_table_name),
120
+ ("evals", self.eval_table_name),
121
+ ("knowledge", self.knowledge_table_name),
122
+ ("culture", self.culture_table_name),
123
+ ]
124
+
125
+ for collection_type, collection_name in collections_to_create:
126
+ if collection_name and not self.table_exists(collection_name):
127
+ self._get_collection(collection_type, create_collection_if_not_found=True)
128
+
129
+ def _get_collection(
130
+ self, table_type: str, create_collection_if_not_found: Optional[bool] = True
131
+ ) -> Optional[Collection]:
132
+ """Get or create a collection based on table type.
133
+
134
+ Args:
135
+ table_type (str): The type of table to get or create.
136
+
137
+ Returns:
138
+ Collection: The collection object.
139
+ """
140
+ if table_type == "sessions":
141
+ if not hasattr(self, "session_collection"):
142
+ if self.session_table_name is None:
143
+ raise ValueError("Session collection was not provided on initialization")
144
+ self.session_collection = self._get_or_create_collection(
145
+ collection_name=self.session_table_name,
146
+ collection_type="sessions",
147
+ create_collection_if_not_found=create_collection_if_not_found,
148
+ )
149
+ return self.session_collection
150
+
151
+ if table_type == "memories":
152
+ if not hasattr(self, "memory_collection"):
153
+ if self.memory_table_name is None:
154
+ raise ValueError("Memory collection was not provided on initialization")
155
+ self.memory_collection = self._get_or_create_collection(
156
+ collection_name=self.memory_table_name,
157
+ collection_type="memories",
158
+ create_collection_if_not_found=create_collection_if_not_found,
159
+ )
160
+ return self.memory_collection
161
+
162
+ if table_type == "metrics":
163
+ if not hasattr(self, "metrics_collection"):
164
+ if self.metrics_table_name is None:
165
+ raise ValueError("Metrics collection was not provided on initialization")
166
+ self.metrics_collection = self._get_or_create_collection(
167
+ collection_name=self.metrics_table_name,
168
+ collection_type="metrics",
169
+ create_collection_if_not_found=create_collection_if_not_found,
170
+ )
171
+ return self.metrics_collection
172
+
173
+ if table_type == "evals":
174
+ if not hasattr(self, "eval_collection"):
175
+ if self.eval_table_name is None:
176
+ raise ValueError("Eval collection was not provided on initialization")
177
+ self.eval_collection = self._get_or_create_collection(
178
+ collection_name=self.eval_table_name,
179
+ collection_type="evals",
180
+ create_collection_if_not_found=create_collection_if_not_found,
181
+ )
182
+ return self.eval_collection
183
+
184
+ if table_type == "knowledge":
185
+ if not hasattr(self, "knowledge_collection"):
186
+ if self.knowledge_table_name is None:
187
+ raise ValueError("Knowledge collection was not provided on initialization")
188
+ self.knowledge_collection = self._get_or_create_collection(
189
+ collection_name=self.knowledge_table_name,
190
+ collection_type="knowledge",
191
+ create_collection_if_not_found=create_collection_if_not_found,
192
+ )
193
+ return self.knowledge_collection
194
+
195
+ if table_type == "culture":
196
+ if not hasattr(self, "culture_collection"):
197
+ if self.culture_table_name is None:
198
+ raise ValueError("Culture collection was not provided on initialization")
199
+ self.culture_collection = self._get_or_create_collection(
200
+ collection_name=self.culture_table_name,
201
+ collection_type="culture",
202
+ create_collection_if_not_found=create_collection_if_not_found,
203
+ )
204
+ return self.culture_collection
205
+
206
+ raise ValueError(f"Unknown table type: {table_type}")
207
+
208
+ def _get_or_create_collection(
209
+ self, collection_name: str, collection_type: str, create_collection_if_not_found: Optional[bool] = True
210
+ ) -> Optional[Collection]:
211
+ """Get or create a collection with proper indexes.
212
+
213
+ Args:
214
+ collection_name (str): The name of the collection to get or create.
215
+ collection_type (str): The type of collection to get or create.
216
+ create_collection_if_not_found (Optional[bool]): Whether to create the collection if it doesn't exist.
217
+
218
+ Returns:
219
+ Optional[Collection]: The collection object.
220
+ """
221
+ try:
222
+ collection = self.database[collection_name]
223
+
224
+ if not hasattr(self, f"_{collection_name}_initialized"):
225
+ if not create_collection_if_not_found:
226
+ return None
227
+ create_collection_indexes(collection, collection_type)
228
+ setattr(self, f"_{collection_name}_initialized", True)
229
+ log_debug(f"Initialized collection '{collection_name}'")
230
+ else:
231
+ log_debug(f"Collection '{collection_name}' already initialized")
232
+
233
+ return collection
234
+
235
+ except Exception as e:
236
+ log_error(f"Error getting collection {collection_name}: {e}")
237
+ raise
238
+
239
+ # -- Session methods --
240
+
241
+ def delete_session(self, session_id: str) -> bool:
242
+ """Delete a session from the database.
243
+
244
+ Args:
245
+ session_id (str): The ID of the session to delete.
246
+
247
+ Returns:
248
+ bool: True if the session was deleted, False otherwise.
249
+
250
+ Raises:
251
+ Exception: If there is an error deleting the session.
252
+ """
253
+ try:
254
+ collection = self._get_collection(table_type="sessions")
255
+ if collection is None:
256
+ return False
257
+
258
+ result = collection.delete_one({"session_id": session_id})
259
+ if result.deleted_count == 0:
260
+ log_debug(f"No session found to delete with session_id: {session_id}")
261
+ return False
262
+ else:
263
+ log_debug(f"Successfully deleted session with session_id: {session_id}")
264
+ return True
265
+
266
+ except Exception as e:
267
+ log_error(f"Error deleting session: {e}")
268
+ raise e
269
+
270
+ def delete_sessions(self, session_ids: List[str]) -> None:
271
+ """Delete multiple sessions from the database.
272
+
273
+ Args:
274
+ session_ids (List[str]): The IDs of the sessions to delete.
275
+ """
276
+ try:
277
+ collection = self._get_collection(table_type="sessions")
278
+ if collection is None:
279
+ return
280
+
281
+ result = collection.delete_many({"session_id": {"$in": session_ids}})
282
+ log_debug(f"Successfully deleted {result.deleted_count} sessions")
283
+
284
+ except Exception as e:
285
+ log_error(f"Error deleting sessions: {e}")
286
+ raise e
287
+
288
+ def get_session(
289
+ self,
290
+ session_id: str,
291
+ session_type: SessionType,
292
+ user_id: Optional[str] = None,
293
+ deserialize: Optional[bool] = True,
294
+ ) -> Optional[Union[Session, Dict[str, Any]]]:
295
+ """Read a session from the database.
296
+
297
+ Args:
298
+ session_id (str): The ID of the session to get.
299
+ session_type (SessionType): The type of session to get.
300
+ user_id (Optional[str]): The ID of the user to get the session for.
301
+ deserialize (Optional[bool]): Whether to serialize the session. Defaults to True.
302
+
303
+ Returns:
304
+ Union[Session, Dict[str, Any], None]:
305
+ - When deserialize=True: Session object
306
+ - When deserialize=False: Session dictionary
307
+
308
+ Raises:
309
+ Exception: If there is an error reading the session.
310
+ """
311
+ try:
312
+ collection = self._get_collection(table_type="sessions")
313
+ if collection is None:
314
+ return None
315
+
316
+ query = {"session_id": session_id}
317
+ if user_id is not None:
318
+ query["user_id"] = user_id
319
+
320
+ result = collection.find_one(query)
321
+ if result is None:
322
+ return None
323
+
324
+ session = deserialize_session_json_fields(result)
325
+ if not deserialize:
326
+ return session
327
+
328
+ if session_type == SessionType.AGENT:
329
+ return AgentSession.from_dict(session)
330
+ elif session_type == SessionType.TEAM:
331
+ return TeamSession.from_dict(session)
332
+ elif session_type == SessionType.WORKFLOW:
333
+ return WorkflowSession.from_dict(session)
334
+ else:
335
+ raise ValueError(f"Invalid session type: {session_type}")
336
+
337
+ except Exception as e:
338
+ log_error(f"Exception reading session: {e}")
339
+ raise e
340
+
341
+ def get_sessions(
342
+ self,
343
+ session_type: Optional[SessionType] = None,
344
+ user_id: Optional[str] = None,
345
+ component_id: Optional[str] = None,
346
+ session_name: Optional[str] = None,
347
+ start_timestamp: Optional[int] = None,
348
+ end_timestamp: Optional[int] = None,
349
+ limit: Optional[int] = None,
350
+ page: Optional[int] = None,
351
+ sort_by: Optional[str] = None,
352
+ sort_order: Optional[str] = None,
353
+ deserialize: Optional[bool] = True,
354
+ ) -> Union[List[Session], Tuple[List[Dict[str, Any]], int]]:
355
+ """Get all sessions.
356
+
357
+ Args:
358
+ session_type (Optional[SessionType]): The type of session to get.
359
+ user_id (Optional[str]): The ID of the user to get the session for.
360
+ component_id (Optional[str]): The ID of the component to get the session for.
361
+ session_name (Optional[str]): The name of the session to filter by.
362
+ start_timestamp (Optional[int]): The start timestamp to filter sessions by.
363
+ end_timestamp (Optional[int]): The end timestamp to filter sessions by.
364
+ limit (Optional[int]): The limit of the sessions to get.
365
+ page (Optional[int]): The page number to get.
366
+ sort_by (Optional[str]): The field to sort the sessions by.
367
+ sort_order (Optional[str]): The order to sort the sessions by.
368
+ deserialize (Optional[bool]): Whether to serialize the sessions. Defaults to True.
369
+ create_table_if_not_found (Optional[bool]): Whether to create the collection if it doesn't exist.
370
+
371
+ Returns:
372
+ Union[List[AgentSession], List[TeamSession], List[WorkflowSession], Tuple[List[Dict[str, Any]], int]]:
373
+ - When deserialize=True: List of Session objects
374
+ - When deserialize=False: List of session dictionaries and the total count
375
+
376
+ Raises:
377
+ Exception: If there is an error reading the sessions.
378
+ """
379
+ try:
380
+ collection = self._get_collection(table_type="sessions")
381
+ if collection is None:
382
+ return [] if deserialize else ([], 0)
383
+
384
+ # Filtering
385
+ query: Dict[str, Any] = {}
386
+ if user_id is not None:
387
+ query["user_id"] = user_id
388
+ if session_type is not None:
389
+ query["session_type"] = session_type
390
+ if component_id is not None:
391
+ if session_type == SessionType.AGENT:
392
+ query["agent_id"] = component_id
393
+ elif session_type == SessionType.TEAM:
394
+ query["team_id"] = component_id
395
+ elif session_type == SessionType.WORKFLOW:
396
+ query["workflow_id"] = component_id
397
+ if start_timestamp is not None:
398
+ query["created_at"] = {"$gte": start_timestamp}
399
+ if end_timestamp is not None:
400
+ if "created_at" in query:
401
+ query["created_at"]["$lte"] = end_timestamp
402
+ else:
403
+ query["created_at"] = {"$lte": end_timestamp}
404
+ if session_name is not None:
405
+ query["session_data.session_name"] = {"$regex": session_name, "$options": "i"}
406
+
407
+ # Get total count
408
+ total_count = collection.count_documents(query)
409
+
410
+ cursor = collection.find(query)
411
+
412
+ # Sorting
413
+ sort_criteria = apply_sorting({}, sort_by, sort_order)
414
+ if sort_criteria:
415
+ cursor = cursor.sort(sort_criteria)
416
+
417
+ # Pagination
418
+ query_args = apply_pagination({}, limit, page)
419
+ if query_args.get("skip"):
420
+ cursor = cursor.skip(query_args["skip"])
421
+ if query_args.get("limit"):
422
+ cursor = cursor.limit(query_args["limit"])
423
+
424
+ records = list(cursor)
425
+ if records is None:
426
+ return [] if deserialize else ([], 0)
427
+ sessions_raw = [deserialize_session_json_fields(record) for record in records]
428
+
429
+ if not deserialize:
430
+ return sessions_raw, total_count
431
+
432
+ sessions: List[Union[AgentSession, TeamSession, WorkflowSession]] = []
433
+ for record in sessions_raw:
434
+ if session_type == SessionType.AGENT.value:
435
+ agent_session = AgentSession.from_dict(record)
436
+ if agent_session is not None:
437
+ sessions.append(agent_session)
438
+ elif session_type == SessionType.TEAM.value:
439
+ team_session = TeamSession.from_dict(record)
440
+ if team_session is not None:
441
+ sessions.append(team_session)
442
+ elif session_type == SessionType.WORKFLOW.value:
443
+ workflow_session = WorkflowSession.from_dict(record)
444
+ if workflow_session is not None:
445
+ sessions.append(workflow_session)
446
+
447
+ return sessions
448
+
449
+ except Exception as e:
450
+ log_error(f"Exception reading sessions: {e}")
451
+ raise e
452
+
453
+ def rename_session(
454
+ self, session_id: str, session_type: SessionType, session_name: str, deserialize: Optional[bool] = True
455
+ ) -> Optional[Union[Session, Dict[str, Any]]]:
456
+ """Rename a session in the database.
457
+
458
+ Args:
459
+ session_id (str): The ID of the session to rename.
460
+ session_type (SessionType): The type of session to rename.
461
+ session_name (str): The new name of the session.
462
+ deserialize (Optional[bool]): Whether to serialize the session. Defaults to True.
463
+
464
+ Returns:
465
+ Optional[Union[Session, Dict[str, Any]]]:
466
+ - When deserialize=True: Session object
467
+ - When deserialize=False: Session dictionary
468
+
469
+ Raises:
470
+ Exception: If there is an error renaming the session.
471
+ """
472
+ try:
473
+ collection = self._get_collection(table_type="sessions")
474
+ if collection is None:
475
+ return None
476
+
477
+ try:
478
+ result = collection.find_one_and_update(
479
+ {"session_id": session_id},
480
+ {"$set": {"session_data.session_name": session_name, "updated_at": int(time.time())}},
481
+ return_document=ReturnDocument.AFTER,
482
+ upsert=False,
483
+ )
484
+ except OperationFailure:
485
+ # If the update fails because session_data doesn't contain a session_name yet, we initialize session_data
486
+ result = collection.find_one_and_update(
487
+ {"session_id": session_id},
488
+ {"$set": {"session_data": {"session_name": session_name}, "updated_at": int(time.time())}},
489
+ return_document=ReturnDocument.AFTER,
490
+ upsert=False,
491
+ )
492
+ if not result:
493
+ return None
494
+
495
+ deserialized_session = deserialize_session_json_fields(result)
496
+
497
+ if not deserialize:
498
+ return deserialized_session
499
+
500
+ if session_type == SessionType.AGENT.value:
501
+ return AgentSession.from_dict(deserialized_session)
502
+ elif session_type == SessionType.TEAM.value:
503
+ return TeamSession.from_dict(deserialized_session)
504
+ else:
505
+ return WorkflowSession.from_dict(deserialized_session)
506
+
507
+ except Exception as e:
508
+ log_error(f"Exception renaming session: {e}")
509
+ raise e
510
+
511
+ def upsert_session(
512
+ self, session: Session, deserialize: Optional[bool] = True
513
+ ) -> Optional[Union[Session, Dict[str, Any]]]:
514
+ """Insert or update a session in the database.
515
+
516
+ Args:
517
+ session (Session): The session to upsert.
518
+
519
+ Returns:
520
+ Optional[Session]: The upserted session.
521
+
522
+ Raises:
523
+ Exception: If there is an error upserting the session.
524
+ """
525
+ try:
526
+ collection = self._get_collection(table_type="sessions", create_collection_if_not_found=True)
527
+ if collection is None:
528
+ return None
529
+
530
+ session_dict = session.to_dict()
531
+
532
+ if isinstance(session, AgentSession):
533
+ record = {
534
+ "session_id": session_dict.get("session_id"),
535
+ "session_type": SessionType.AGENT.value,
536
+ "agent_id": session_dict.get("agent_id"),
537
+ "user_id": session_dict.get("user_id"),
538
+ "runs": session_dict.get("runs"),
539
+ "agent_data": session_dict.get("agent_data"),
540
+ "session_data": session_dict.get("session_data"),
541
+ "summary": session_dict.get("summary"),
542
+ "metadata": session_dict.get("metadata"),
543
+ "created_at": session_dict.get("created_at"),
544
+ "updated_at": int(time.time()),
545
+ }
546
+
547
+ result = collection.find_one_and_replace(
548
+ filter={"session_id": session_dict.get("session_id")},
549
+ replacement=record,
550
+ upsert=True,
551
+ return_document=ReturnDocument.AFTER,
552
+ )
553
+ if not result:
554
+ return None
555
+
556
+ session = result # type: ignore
557
+
558
+ if not deserialize:
559
+ return session
560
+
561
+ return AgentSession.from_dict(session) # type: ignore
562
+
563
+ elif isinstance(session, TeamSession):
564
+ record = {
565
+ "session_id": session_dict.get("session_id"),
566
+ "session_type": SessionType.TEAM.value,
567
+ "team_id": session_dict.get("team_id"),
568
+ "user_id": session_dict.get("user_id"),
569
+ "runs": session_dict.get("runs"),
570
+ "team_data": session_dict.get("team_data"),
571
+ "session_data": session_dict.get("session_data"),
572
+ "summary": session_dict.get("summary"),
573
+ "metadata": session_dict.get("metadata"),
574
+ "created_at": session_dict.get("created_at"),
575
+ "updated_at": int(time.time()),
576
+ }
577
+
578
+ result = collection.find_one_and_replace(
579
+ filter={"session_id": session_dict.get("session_id")},
580
+ replacement=record,
581
+ upsert=True,
582
+ return_document=ReturnDocument.AFTER,
583
+ )
584
+ if not result:
585
+ return None
586
+
587
+ # MongoDB stores native objects, no deserialization needed for document fields
588
+ session = result # type: ignore
589
+
590
+ if not deserialize:
591
+ return session
592
+
593
+ return TeamSession.from_dict(session) # type: ignore
594
+
595
+ else:
596
+ record = {
597
+ "session_id": session_dict.get("session_id"),
598
+ "session_type": SessionType.WORKFLOW.value,
599
+ "workflow_id": session_dict.get("workflow_id"),
600
+ "user_id": session_dict.get("user_id"),
601
+ "runs": session_dict.get("runs"),
602
+ "workflow_data": session_dict.get("workflow_data"),
603
+ "session_data": session_dict.get("session_data"),
604
+ "summary": session_dict.get("summary"),
605
+ "metadata": session_dict.get("metadata"),
606
+ "created_at": session_dict.get("created_at"),
607
+ "updated_at": int(time.time()),
608
+ }
609
+
610
+ result = collection.find_one_and_replace(
611
+ filter={"session_id": session_dict.get("session_id")},
612
+ replacement=record,
613
+ upsert=True,
614
+ return_document=ReturnDocument.AFTER,
615
+ )
616
+ if not result:
617
+ return None
618
+
619
+ session = result # type: ignore
620
+
621
+ if not deserialize:
622
+ return session
623
+
624
+ return WorkflowSession.from_dict(session) # type: ignore
625
+
626
+ except Exception as e:
627
+ log_error(f"Exception upserting session: {e}")
628
+ raise e
629
+
630
+ def upsert_sessions(
631
+ self, sessions: List[Session], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
632
+ ) -> List[Union[Session, Dict[str, Any]]]:
633
+ """
634
+ Bulk upsert multiple sessions for improved performance on large datasets.
635
+
636
+ Args:
637
+ sessions (List[Session]): List of sessions to upsert.
638
+ deserialize (Optional[bool]): Whether to deserialize the sessions. Defaults to True.
639
+ preserve_updated_at (bool): If True, preserve the updated_at from the session object.
640
+
641
+ Returns:
642
+ List[Union[Session, Dict[str, Any]]]: List of upserted sessions.
643
+
644
+ Raises:
645
+ Exception: If an error occurs during bulk upsert.
646
+ """
647
+ if not sessions:
648
+ return []
649
+
650
+ try:
651
+ collection = self._get_collection(table_type="sessions", create_collection_if_not_found=True)
652
+ if collection is None:
653
+ log_info("Sessions collection not available, falling back to individual upserts")
654
+ return [
655
+ result
656
+ for session in sessions
657
+ if session is not None
658
+ for result in [self.upsert_session(session, deserialize=deserialize)]
659
+ if result is not None
660
+ ]
661
+
662
+ from pymongo import ReplaceOne
663
+
664
+ operations = []
665
+ results: List[Union[Session, Dict[str, Any]]] = []
666
+
667
+ for session in sessions:
668
+ if session is None:
669
+ continue
670
+
671
+ session_dict = session.to_dict()
672
+
673
+ # Use preserved updated_at if flag is set and value exists, otherwise use current time
674
+ updated_at = session_dict.get("updated_at") if preserve_updated_at else int(time.time())
675
+
676
+ if isinstance(session, AgentSession):
677
+ record = {
678
+ "session_id": session_dict.get("session_id"),
679
+ "session_type": SessionType.AGENT.value,
680
+ "agent_id": session_dict.get("agent_id"),
681
+ "user_id": session_dict.get("user_id"),
682
+ "runs": session_dict.get("runs"),
683
+ "agent_data": session_dict.get("agent_data"),
684
+ "session_data": session_dict.get("session_data"),
685
+ "summary": session_dict.get("summary"),
686
+ "metadata": session_dict.get("metadata"),
687
+ "created_at": session_dict.get("created_at"),
688
+ "updated_at": updated_at,
689
+ }
690
+ elif isinstance(session, TeamSession):
691
+ record = {
692
+ "session_id": session_dict.get("session_id"),
693
+ "session_type": SessionType.TEAM.value,
694
+ "team_id": session_dict.get("team_id"),
695
+ "user_id": session_dict.get("user_id"),
696
+ "runs": session_dict.get("runs"),
697
+ "team_data": session_dict.get("team_data"),
698
+ "session_data": session_dict.get("session_data"),
699
+ "summary": session_dict.get("summary"),
700
+ "metadata": session_dict.get("metadata"),
701
+ "created_at": session_dict.get("created_at"),
702
+ "updated_at": updated_at,
703
+ }
704
+ elif isinstance(session, WorkflowSession):
705
+ record = {
706
+ "session_id": session_dict.get("session_id"),
707
+ "session_type": SessionType.WORKFLOW.value,
708
+ "workflow_id": session_dict.get("workflow_id"),
709
+ "user_id": session_dict.get("user_id"),
710
+ "runs": session_dict.get("runs"),
711
+ "workflow_data": session_dict.get("workflow_data"),
712
+ "session_data": session_dict.get("session_data"),
713
+ "summary": session_dict.get("summary"),
714
+ "metadata": session_dict.get("metadata"),
715
+ "created_at": session_dict.get("created_at"),
716
+ "updated_at": updated_at,
717
+ }
718
+ else:
719
+ continue
720
+
721
+ operations.append(
722
+ ReplaceOne(filter={"session_id": record["session_id"]}, replacement=record, upsert=True)
723
+ )
724
+
725
+ if operations:
726
+ # Execute bulk write
727
+ collection.bulk_write(operations)
728
+
729
+ # Fetch the results
730
+ session_ids = [session.session_id for session in sessions if session and session.session_id]
731
+ cursor = collection.find({"session_id": {"$in": session_ids}})
732
+
733
+ for doc in cursor:
734
+ session_dict = doc
735
+
736
+ if deserialize:
737
+ session_type = doc.get("session_type")
738
+ if session_type == SessionType.AGENT.value:
739
+ deserialized_agent_session = AgentSession.from_dict(session_dict)
740
+ if deserialized_agent_session is None:
741
+ continue
742
+ results.append(deserialized_agent_session)
743
+
744
+ elif session_type == SessionType.TEAM.value:
745
+ deserialized_team_session = TeamSession.from_dict(session_dict)
746
+ if deserialized_team_session is None:
747
+ continue
748
+ results.append(deserialized_team_session)
749
+
750
+ elif session_type == SessionType.WORKFLOW.value:
751
+ deserialized_workflow_session = WorkflowSession.from_dict(session_dict)
752
+ if deserialized_workflow_session is None:
753
+ continue
754
+ results.append(deserialized_workflow_session)
755
+ else:
756
+ results.append(session_dict)
757
+
758
+ return results
759
+
760
+ except Exception as e:
761
+ log_error(f"Exception during bulk session upsert, falling back to individual upserts: {e}")
762
+
763
+ # Fallback to individual upserts
764
+ return [
765
+ result
766
+ for session in sessions
767
+ if session is not None
768
+ for result in [self.upsert_session(session, deserialize=deserialize)]
769
+ if result is not None
770
+ ]
771
+
772
+ # -- Memory methods --
773
+
774
+ def delete_user_memory(self, memory_id: str, user_id: Optional[str] = None):
775
+ """Delete a user memory from the database.
776
+
777
+ Args:
778
+ memory_id (str): The ID of the memory to delete.
779
+ user_id (Optional[str]): The ID of the user to verify ownership. If provided, only delete if the memory belongs to this user.
780
+
781
+ Returns:
782
+ bool: True if the memory was deleted, False otherwise.
783
+
784
+ Raises:
785
+ Exception: If there is an error deleting the memory.
786
+ """
787
+ try:
788
+ collection = self._get_collection(table_type="memories")
789
+ if collection is None:
790
+ return
791
+
792
+ query = {"memory_id": memory_id}
793
+ if user_id is not None:
794
+ query["user_id"] = user_id
795
+
796
+ result = collection.delete_one(query)
797
+
798
+ success = result.deleted_count > 0
799
+ if success:
800
+ log_debug(f"Successfully deleted memory id: {memory_id}")
801
+ else:
802
+ log_debug(f"No memory found with id: {memory_id}")
803
+
804
+ except Exception as e:
805
+ log_error(f"Error deleting memory: {e}")
806
+ raise e
807
+
808
+ def delete_user_memories(self, memory_ids: List[str], user_id: Optional[str] = None) -> None:
809
+ """Delete user memories from the database.
810
+
811
+ Args:
812
+ memory_ids (List[str]): The IDs of the memories to delete.
813
+ user_id (Optional[str]): The ID of the user to verify ownership. If provided, only delete memories that belong to this user.
814
+
815
+ Raises:
816
+ Exception: If there is an error deleting the memories.
817
+ """
818
+ try:
819
+ collection = self._get_collection(table_type="memories")
820
+ if collection is None:
821
+ return
822
+
823
+ query: Dict[str, Any] = {"memory_id": {"$in": memory_ids}}
824
+ if user_id is not None:
825
+ query["user_id"] = user_id
826
+
827
+ result = collection.delete_many(query)
828
+
829
+ if result.deleted_count == 0:
830
+ log_debug(f"No memories found with ids: {memory_ids}")
831
+
832
+ except Exception as e:
833
+ log_error(f"Error deleting memories: {e}")
834
+ raise e
835
+
836
+ def get_all_memory_topics(self) -> List[str]:
837
+ """Get all memory topics from the database.
838
+
839
+ Returns:
840
+ List[str]: The topics.
841
+
842
+ Raises:
843
+ Exception: If there is an error getting the topics.
844
+ """
845
+ try:
846
+ collection = self._get_collection(table_type="memories")
847
+ if collection is None:
848
+ return []
849
+
850
+ topics = collection.distinct("topics", {})
851
+ return [topic for topic in topics if topic]
852
+
853
+ except Exception as e:
854
+ log_error(f"Exception reading from collection: {e}")
855
+ raise e
856
+
857
+ def get_user_memory(
858
+ self, memory_id: str, deserialize: Optional[bool] = True, user_id: Optional[str] = None
859
+ ) -> Optional[UserMemory]:
860
+ """Get a memory from the database.
861
+
862
+ Args:
863
+ memory_id (str): The ID of the memory to get.
864
+ deserialize (Optional[bool]): Whether to serialize the memory. Defaults to True.
865
+ user_id (Optional[str]): The ID of the user to verify ownership. If provided, only return the memory if it belongs to this user.
866
+
867
+ Returns:
868
+ Optional[UserMemory]:
869
+ - When deserialize=True: UserMemory object
870
+ - When deserialize=False: Memory dictionary
871
+
872
+ Raises:
873
+ Exception: If there is an error getting the memory.
874
+ """
875
+ try:
876
+ collection = self._get_collection(table_type="memories")
877
+ if collection is None:
878
+ return None
879
+
880
+ query = {"memory_id": memory_id}
881
+ if user_id is not None:
882
+ query["user_id"] = user_id
883
+
884
+ result = collection.find_one(query)
885
+ if result is None or not deserialize:
886
+ return result
887
+
888
+ # Remove MongoDB's _id field before creating UserMemory object
889
+ result_filtered = {k: v for k, v in result.items() if k != "_id"}
890
+ return UserMemory.from_dict(result_filtered)
891
+
892
+ except Exception as e:
893
+ log_error(f"Exception reading from collection: {e}")
894
+ raise e
895
+
896
+ def get_user_memories(
897
+ self,
898
+ user_id: Optional[str] = None,
899
+ agent_id: Optional[str] = None,
900
+ team_id: Optional[str] = None,
901
+ topics: Optional[List[str]] = None,
902
+ search_content: Optional[str] = None,
903
+ limit: Optional[int] = None,
904
+ page: Optional[int] = None,
905
+ sort_by: Optional[str] = None,
906
+ sort_order: Optional[str] = None,
907
+ deserialize: Optional[bool] = True,
908
+ ) -> Union[List[UserMemory], Tuple[List[Dict[str, Any]], int]]:
909
+ """Get all memories from the database as UserMemory objects.
910
+
911
+ Args:
912
+ user_id (Optional[str]): The ID of the user to get the memories for.
913
+ agent_id (Optional[str]): The ID of the agent to get the memories for.
914
+ team_id (Optional[str]): The ID of the team to get the memories for.
915
+ topics (Optional[List[str]]): The topics to filter the memories by.
916
+ search_content (Optional[str]): The content to filter the memories by.
917
+ limit (Optional[int]): The limit of the memories to get.
918
+ page (Optional[int]): The page number to get.
919
+ sort_by (Optional[str]): The field to sort the memories by.
920
+ sort_order (Optional[str]): The order to sort the memories by.
921
+ deserialize (Optional[bool]): Whether to serialize the memories. Defaults to True.
922
+ create_table_if_not_found: Whether to create the collection if it doesn't exist.
923
+
924
+ Returns:
925
+ Tuple[List[Dict[str, Any]], int]: A tuple containing the memories and the total count.
926
+
927
+ Raises:
928
+ Exception: If there is an error getting the memories.
929
+ """
930
+ try:
931
+ collection = self._get_collection(table_type="memories")
932
+ if collection is None:
933
+ return [] if deserialize else ([], 0)
934
+
935
+ query: Dict[str, Any] = {}
936
+ if user_id is not None:
937
+ query["user_id"] = user_id
938
+ if agent_id is not None:
939
+ query["agent_id"] = agent_id
940
+ if team_id is not None:
941
+ query["team_id"] = team_id
942
+ if topics is not None:
943
+ query["topics"] = {"$in": topics}
944
+ if search_content is not None:
945
+ query["memory"] = {"$regex": search_content, "$options": "i"}
946
+
947
+ # Get total count
948
+ total_count = collection.count_documents(query)
949
+
950
+ # Apply sorting
951
+ sort_criteria = apply_sorting({}, sort_by, sort_order)
952
+
953
+ # Apply pagination
954
+ query_args = apply_pagination({}, limit, page)
955
+
956
+ cursor = collection.find(query)
957
+ if sort_criteria:
958
+ cursor = cursor.sort(sort_criteria)
959
+ if query_args.get("skip"):
960
+ cursor = cursor.skip(query_args["skip"])
961
+ if query_args.get("limit"):
962
+ cursor = cursor.limit(query_args["limit"])
963
+
964
+ records = list(cursor)
965
+ if not deserialize:
966
+ return records, total_count
967
+
968
+ # Remove MongoDB's _id field before creating UserMemory objects
969
+ return [UserMemory.from_dict({k: v for k, v in record.items() if k != "_id"}) for record in records]
970
+
971
+ except Exception as e:
972
+ log_error(f"Exception reading from collection: {e}")
973
+ raise e
974
+
975
+ def get_user_memory_stats(
976
+ self,
977
+ limit: Optional[int] = None,
978
+ page: Optional[int] = None,
979
+ ) -> Tuple[List[Dict[str, Any]], int]:
980
+ """Get user memories stats.
981
+
982
+ Args:
983
+ limit (Optional[int]): The limit of the memories to get.
984
+ page (Optional[int]): The page number to get.
985
+
986
+ Returns:
987
+ Tuple[List[Dict[str, Any]], int]: A tuple containing the memories stats and the total count.
988
+
989
+ Raises:
990
+ Exception: If there is an error getting the memories stats.
991
+ """
992
+ try:
993
+ collection = self._get_collection(table_type="memories")
994
+ if collection is None:
995
+ return [], 0
996
+
997
+ match_stage = {"user_id": {"$ne": None}}
998
+
999
+ pipeline = [
1000
+ {"$match": match_stage},
1001
+ {
1002
+ "$group": {
1003
+ "_id": "$user_id",
1004
+ "total_memories": {"$sum": 1},
1005
+ "last_memory_updated_at": {"$max": "$updated_at"},
1006
+ }
1007
+ },
1008
+ {"$sort": {"last_memory_updated_at": -1}},
1009
+ ]
1010
+
1011
+ # Get total count
1012
+ count_pipeline = pipeline + [{"$count": "total"}]
1013
+ count_result = list(collection.aggregate(count_pipeline)) # type: ignore
1014
+ total_count = count_result[0]["total"] if count_result else 0
1015
+
1016
+ # Apply pagination
1017
+ if limit is not None:
1018
+ if page is not None:
1019
+ pipeline.append({"$skip": (page - 1) * limit})
1020
+ pipeline.append({"$limit": limit})
1021
+
1022
+ results = list(collection.aggregate(pipeline)) # type: ignore
1023
+
1024
+ formatted_results = [
1025
+ {
1026
+ "user_id": result["_id"],
1027
+ "total_memories": result["total_memories"],
1028
+ "last_memory_updated_at": result["last_memory_updated_at"],
1029
+ }
1030
+ for result in results
1031
+ ]
1032
+
1033
+ return formatted_results, total_count
1034
+
1035
+ except Exception as e:
1036
+ log_error(f"Exception getting user memory stats: {e}")
1037
+ raise e
1038
+
1039
+ def upsert_user_memory(
1040
+ self, memory: UserMemory, deserialize: Optional[bool] = True
1041
+ ) -> Optional[Union[UserMemory, Dict[str, Any]]]:
1042
+ """Upsert a user memory in the database.
1043
+
1044
+ Args:
1045
+ memory (UserMemory): The memory to upsert.
1046
+ deserialize (Optional[bool]): Whether to serialize the memory. Defaults to True.
1047
+
1048
+ Returns:
1049
+ Optional[Union[UserMemory, Dict[str, Any]]]:
1050
+ - When deserialize=True: UserMemory object
1051
+ - When deserialize=False: Memory dictionary
1052
+
1053
+ Raises:
1054
+ Exception: If there is an error upserting the memory.
1055
+ """
1056
+ try:
1057
+ collection = self._get_collection(table_type="memories", create_collection_if_not_found=True)
1058
+ if collection is None:
1059
+ return None
1060
+
1061
+ if memory.memory_id is None:
1062
+ memory.memory_id = str(uuid4())
1063
+
1064
+ update_doc = {
1065
+ "user_id": memory.user_id,
1066
+ "agent_id": memory.agent_id,
1067
+ "team_id": memory.team_id,
1068
+ "memory_id": memory.memory_id,
1069
+ "memory": memory.memory,
1070
+ "topics": memory.topics,
1071
+ "updated_at": int(time.time()),
1072
+ }
1073
+
1074
+ result = collection.replace_one({"memory_id": memory.memory_id}, update_doc, upsert=True)
1075
+
1076
+ if result.upserted_id:
1077
+ update_doc["_id"] = result.upserted_id
1078
+
1079
+ if not deserialize:
1080
+ return update_doc
1081
+
1082
+ # Remove MongoDB's _id field before creating UserMemory object
1083
+ update_doc_filtered = {k: v for k, v in update_doc.items() if k != "_id"}
1084
+ return UserMemory.from_dict(update_doc_filtered)
1085
+
1086
+ except Exception as e:
1087
+ log_error(f"Exception upserting user memory: {e}")
1088
+ raise e
1089
+
1090
+ def upsert_memories(
1091
+ self, memories: List[UserMemory], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
1092
+ ) -> List[Union[UserMemory, Dict[str, Any]]]:
1093
+ """
1094
+ Bulk upsert multiple user memories for improved performance on large datasets.
1095
+
1096
+ Args:
1097
+ memories (List[UserMemory]): List of memories to upsert.
1098
+ deserialize (Optional[bool]): Whether to deserialize the memories. Defaults to True.
1099
+
1100
+ Returns:
1101
+ List[Union[UserMemory, Dict[str, Any]]]: List of upserted memories.
1102
+
1103
+ Raises:
1104
+ Exception: If an error occurs during bulk upsert.
1105
+ """
1106
+ if not memories:
1107
+ return []
1108
+
1109
+ try:
1110
+ collection = self._get_collection(table_type="memories", create_collection_if_not_found=True)
1111
+ if collection is None:
1112
+ log_info("Memories collection not available, falling back to individual upserts")
1113
+ return [
1114
+ result
1115
+ for memory in memories
1116
+ if memory is not None
1117
+ for result in [self.upsert_user_memory(memory, deserialize=deserialize)]
1118
+ if result is not None
1119
+ ]
1120
+
1121
+ from pymongo import ReplaceOne
1122
+
1123
+ operations = []
1124
+ results: List[Union[UserMemory, Dict[str, Any]]] = []
1125
+
1126
+ current_time = int(time.time())
1127
+ for memory in memories:
1128
+ if memory is None:
1129
+ continue
1130
+
1131
+ if memory.memory_id is None:
1132
+ memory.memory_id = str(uuid4())
1133
+
1134
+ # Use preserved updated_at if flag is set and value exists, otherwise use current time
1135
+ updated_at = memory.updated_at if preserve_updated_at else current_time
1136
+
1137
+ record = {
1138
+ "user_id": memory.user_id,
1139
+ "agent_id": memory.agent_id,
1140
+ "team_id": memory.team_id,
1141
+ "memory_id": memory.memory_id,
1142
+ "memory": memory.memory,
1143
+ "topics": memory.topics,
1144
+ "updated_at": updated_at,
1145
+ }
1146
+
1147
+ operations.append(ReplaceOne(filter={"memory_id": memory.memory_id}, replacement=record, upsert=True))
1148
+
1149
+ if operations:
1150
+ # Execute bulk write
1151
+ collection.bulk_write(operations)
1152
+
1153
+ # Fetch the results
1154
+ memory_ids = [memory.memory_id for memory in memories if memory and memory.memory_id]
1155
+ cursor = collection.find({"memory_id": {"$in": memory_ids}})
1156
+
1157
+ for doc in cursor:
1158
+ if deserialize:
1159
+ # Remove MongoDB's _id field before creating UserMemory object
1160
+ doc_filtered = {k: v for k, v in doc.items() if k != "_id"}
1161
+ results.append(UserMemory.from_dict(doc_filtered))
1162
+ else:
1163
+ results.append(doc)
1164
+
1165
+ return results
1166
+
1167
+ except Exception as e:
1168
+ log_error(f"Exception during bulk memory upsert, falling back to individual upserts: {e}")
1169
+
1170
+ # Fallback to individual upserts
1171
+ return [
1172
+ result
1173
+ for memory in memories
1174
+ if memory is not None
1175
+ for result in [self.upsert_user_memory(memory, deserialize=deserialize)]
1176
+ if result is not None
1177
+ ]
1178
+
1179
+ def clear_memories(self) -> None:
1180
+ """Delete all memories from the database.
1181
+
1182
+ Raises:
1183
+ Exception: If an error occurs during deletion.
1184
+ """
1185
+ try:
1186
+ collection = self._get_collection(table_type="memories")
1187
+ if collection is None:
1188
+ return
1189
+
1190
+ collection.delete_many({})
1191
+
1192
+ except Exception as e:
1193
+ log_error(f"Exception deleting all memories: {e}")
1194
+ raise e
1195
+
1196
+ # -- Cultural Knowledge methods --
1197
+ def clear_cultural_knowledge(self) -> None:
1198
+ """Delete all cultural knowledge from the database.
1199
+
1200
+ Raises:
1201
+ Exception: If an error occurs during deletion.
1202
+ """
1203
+ try:
1204
+ collection = self._get_collection(table_type="culture")
1205
+ if collection is None:
1206
+ return
1207
+
1208
+ collection.delete_many({})
1209
+
1210
+ except Exception as e:
1211
+ log_error(f"Exception deleting all cultural knowledge: {e}")
1212
+ raise e
1213
+
1214
+ def delete_cultural_knowledge(self, id: str) -> None:
1215
+ """Delete cultural knowledge by ID.
1216
+
1217
+ Args:
1218
+ id (str): The ID of the cultural knowledge to delete.
1219
+
1220
+ Raises:
1221
+ Exception: If an error occurs during deletion.
1222
+ """
1223
+ try:
1224
+ collection = self._get_collection(table_type="culture")
1225
+ if collection is None:
1226
+ return
1227
+
1228
+ collection.delete_one({"id": id})
1229
+ log_debug(f"Deleted cultural knowledge with ID: {id}")
1230
+
1231
+ except Exception as e:
1232
+ log_error(f"Error deleting cultural knowledge: {e}")
1233
+ raise e
1234
+
1235
+ def get_cultural_knowledge(
1236
+ self, id: str, deserialize: Optional[bool] = True
1237
+ ) -> Optional[Union[CulturalKnowledge, Dict[str, Any]]]:
1238
+ """Get cultural knowledge by ID.
1239
+
1240
+ Args:
1241
+ id (str): The ID of the cultural knowledge to retrieve.
1242
+ deserialize (Optional[bool]): Whether to deserialize to CulturalKnowledge object. Defaults to True.
1243
+
1244
+ Returns:
1245
+ Optional[Union[CulturalKnowledge, Dict[str, Any]]]: The cultural knowledge if found, None otherwise.
1246
+
1247
+ Raises:
1248
+ Exception: If an error occurs during retrieval.
1249
+ """
1250
+ try:
1251
+ collection = self._get_collection(table_type="culture")
1252
+ if collection is None:
1253
+ return None
1254
+
1255
+ result = collection.find_one({"id": id})
1256
+ if result is None:
1257
+ return None
1258
+
1259
+ # Remove MongoDB's _id field
1260
+ result_filtered = {k: v for k, v in result.items() if k != "_id"}
1261
+
1262
+ if not deserialize:
1263
+ return result_filtered
1264
+
1265
+ return deserialize_cultural_knowledge_from_db(result_filtered)
1266
+
1267
+ except Exception as e:
1268
+ log_error(f"Error getting cultural knowledge: {e}")
1269
+ raise e
1270
+
1271
+ def get_all_cultural_knowledge(
1272
+ self,
1273
+ agent_id: Optional[str] = None,
1274
+ team_id: Optional[str] = None,
1275
+ name: Optional[str] = None,
1276
+ limit: Optional[int] = None,
1277
+ page: Optional[int] = None,
1278
+ sort_by: Optional[str] = None,
1279
+ sort_order: Optional[str] = None,
1280
+ deserialize: Optional[bool] = True,
1281
+ ) -> Union[List[CulturalKnowledge], Tuple[List[Dict[str, Any]], int]]:
1282
+ """Get all cultural knowledge with filtering and pagination.
1283
+
1284
+ Args:
1285
+ agent_id (Optional[str]): Filter by agent ID.
1286
+ team_id (Optional[str]): Filter by team ID.
1287
+ name (Optional[str]): Filter by name (case-insensitive partial match).
1288
+ limit (Optional[int]): Maximum number of results to return.
1289
+ page (Optional[int]): Page number for pagination.
1290
+ sort_by (Optional[str]): Field to sort by.
1291
+ sort_order (Optional[str]): Sort order ('asc' or 'desc').
1292
+ deserialize (Optional[bool]): Whether to deserialize to CulturalKnowledge objects. Defaults to True.
1293
+
1294
+ Returns:
1295
+ Union[List[CulturalKnowledge], Tuple[List[Dict[str, Any]], int]]:
1296
+ - When deserialize=True: List of CulturalKnowledge objects
1297
+ - When deserialize=False: Tuple with list of dictionaries and total count
1298
+
1299
+ Raises:
1300
+ Exception: If an error occurs during retrieval.
1301
+ """
1302
+ try:
1303
+ collection = self._get_collection(table_type="culture")
1304
+ if collection is None:
1305
+ if not deserialize:
1306
+ return [], 0
1307
+ return []
1308
+
1309
+ # Build query
1310
+ query: Dict[str, Any] = {}
1311
+ if agent_id is not None:
1312
+ query["agent_id"] = agent_id
1313
+ if team_id is not None:
1314
+ query["team_id"] = team_id
1315
+ if name is not None:
1316
+ query["name"] = {"$regex": name, "$options": "i"}
1317
+
1318
+ # Get total count for pagination
1319
+ total_count = collection.count_documents(query)
1320
+
1321
+ # Apply sorting
1322
+ sort_criteria = apply_sorting({}, sort_by, sort_order)
1323
+
1324
+ # Apply pagination
1325
+ query_args = apply_pagination({}, limit, page)
1326
+
1327
+ cursor = collection.find(query)
1328
+ if sort_criteria:
1329
+ cursor = cursor.sort(sort_criteria)
1330
+ if query_args.get("skip"):
1331
+ cursor = cursor.skip(query_args["skip"])
1332
+ if query_args.get("limit"):
1333
+ cursor = cursor.limit(query_args["limit"])
1334
+
1335
+ # Remove MongoDB's _id field from all results
1336
+ results_filtered = [{k: v for k, v in item.items() if k != "_id"} for item in cursor]
1337
+
1338
+ if not deserialize:
1339
+ return results_filtered, total_count
1340
+
1341
+ return [deserialize_cultural_knowledge_from_db(item) for item in results_filtered]
1342
+
1343
+ except Exception as e:
1344
+ log_error(f"Error getting all cultural knowledge: {e}")
1345
+ raise e
1346
+
1347
+ def upsert_cultural_knowledge(
1348
+ self, cultural_knowledge: CulturalKnowledge, deserialize: Optional[bool] = True
1349
+ ) -> Optional[Union[CulturalKnowledge, Dict[str, Any]]]:
1350
+ """Upsert cultural knowledge in MongoDB.
1351
+
1352
+ Args:
1353
+ cultural_knowledge (CulturalKnowledge): The cultural knowledge to upsert.
1354
+ deserialize (Optional[bool]): Whether to deserialize the result. Defaults to True.
1355
+
1356
+ Returns:
1357
+ Optional[Union[CulturalKnowledge, Dict[str, Any]]]: The upserted cultural knowledge.
1358
+
1359
+ Raises:
1360
+ Exception: If an error occurs during upsert.
1361
+ """
1362
+ try:
1363
+ collection = self._get_collection(table_type="culture", create_collection_if_not_found=True)
1364
+ if collection is None:
1365
+ return None
1366
+
1367
+ # Serialize content, categories, and notes into a dict for DB storage
1368
+ content_dict = serialize_cultural_knowledge_for_db(cultural_knowledge)
1369
+
1370
+ # Create the document with serialized content
1371
+ update_doc = {
1372
+ "id": cultural_knowledge.id,
1373
+ "name": cultural_knowledge.name,
1374
+ "summary": cultural_knowledge.summary,
1375
+ "content": content_dict if content_dict else None,
1376
+ "metadata": cultural_knowledge.metadata,
1377
+ "input": cultural_knowledge.input,
1378
+ "created_at": cultural_knowledge.created_at,
1379
+ "updated_at": int(time.time()),
1380
+ "agent_id": cultural_knowledge.agent_id,
1381
+ "team_id": cultural_knowledge.team_id,
1382
+ }
1383
+
1384
+ result = collection.replace_one({"id": cultural_knowledge.id}, update_doc, upsert=True)
1385
+
1386
+ if result.upserted_id:
1387
+ update_doc["_id"] = result.upserted_id
1388
+
1389
+ # Remove MongoDB's _id field
1390
+ doc_filtered = {k: v for k, v in update_doc.items() if k != "_id"}
1391
+
1392
+ if not deserialize:
1393
+ return doc_filtered
1394
+
1395
+ return deserialize_cultural_knowledge_from_db(doc_filtered)
1396
+
1397
+ except Exception as e:
1398
+ log_error(f"Error upserting cultural knowledge: {e}")
1399
+ raise e
1400
+
1401
+ # -- Metrics methods --
1402
+
1403
+ def _get_all_sessions_for_metrics_calculation(
1404
+ self, start_timestamp: Optional[int] = None, end_timestamp: Optional[int] = None
1405
+ ) -> List[Dict[str, Any]]:
1406
+ """Get all sessions of all types for metrics calculation."""
1407
+ try:
1408
+ collection = self._get_collection(table_type="sessions")
1409
+ if collection is None:
1410
+ return []
1411
+
1412
+ query = {}
1413
+ if start_timestamp is not None:
1414
+ query["created_at"] = {"$gte": start_timestamp}
1415
+ if end_timestamp is not None:
1416
+ if "created_at" in query:
1417
+ query["created_at"]["$lte"] = end_timestamp
1418
+ else:
1419
+ query["created_at"] = {"$lte": end_timestamp}
1420
+
1421
+ projection = {
1422
+ "user_id": 1,
1423
+ "session_data": 1,
1424
+ "runs": 1,
1425
+ "created_at": 1,
1426
+ "session_type": 1,
1427
+ }
1428
+
1429
+ results = list(collection.find(query, projection))
1430
+ return results
1431
+
1432
+ except Exception as e:
1433
+ log_error(f"Exception reading from sessions collection: {e}")
1434
+ return []
1435
+
1436
+ def _get_metrics_calculation_starting_date(self, collection: Collection) -> Optional[date]:
1437
+ """Get the first date for which metrics calculation is needed."""
1438
+ try:
1439
+ result = collection.find_one({}, sort=[("date", -1)], limit=1)
1440
+
1441
+ if result is not None:
1442
+ result_date = datetime.strptime(result["date"], "%Y-%m-%d").date()
1443
+ if result.get("completed"):
1444
+ return result_date + timedelta(days=1)
1445
+ else:
1446
+ return result_date
1447
+
1448
+ # No metrics records. Return the date of the first recorded session.
1449
+ first_session_result = self.get_sessions(sort_by="created_at", sort_order="asc", limit=1, deserialize=False)
1450
+ first_session_date = first_session_result[0][0]["created_at"] if first_session_result[0] else None # type: ignore
1451
+
1452
+ if first_session_date is None:
1453
+ return None
1454
+
1455
+ return datetime.fromtimestamp(first_session_date, tz=timezone.utc).date()
1456
+
1457
+ except Exception as e:
1458
+ log_error(f"Exception getting metrics calculation starting date: {e}")
1459
+ return None
1460
+
1461
+ def calculate_metrics(self) -> Optional[list[dict]]:
1462
+ """Calculate metrics for all dates without complete metrics."""
1463
+ try:
1464
+ collection = self._get_collection(table_type="metrics", create_collection_if_not_found=True)
1465
+ if collection is None:
1466
+ return None
1467
+
1468
+ starting_date = self._get_metrics_calculation_starting_date(collection)
1469
+ if starting_date is None:
1470
+ log_info("No session data found. Won't calculate metrics.")
1471
+ return None
1472
+
1473
+ dates_to_process = get_dates_to_calculate_metrics_for(starting_date)
1474
+ if not dates_to_process:
1475
+ log_info("Metrics already calculated for all relevant dates.")
1476
+ return None
1477
+
1478
+ start_timestamp = int(
1479
+ datetime.combine(dates_to_process[0], datetime.min.time()).replace(tzinfo=timezone.utc).timestamp()
1480
+ )
1481
+ end_timestamp = int(
1482
+ datetime.combine(dates_to_process[-1] + timedelta(days=1), datetime.min.time())
1483
+ .replace(tzinfo=timezone.utc)
1484
+ .timestamp()
1485
+ )
1486
+
1487
+ sessions = self._get_all_sessions_for_metrics_calculation(
1488
+ start_timestamp=start_timestamp, end_timestamp=end_timestamp
1489
+ )
1490
+ all_sessions_data = fetch_all_sessions_data(
1491
+ sessions=sessions, dates_to_process=dates_to_process, start_timestamp=start_timestamp
1492
+ )
1493
+ if not all_sessions_data:
1494
+ log_info("No new session data found. Won't calculate metrics.")
1495
+ return None
1496
+
1497
+ results = []
1498
+ metrics_records = []
1499
+
1500
+ for date_to_process in dates_to_process:
1501
+ date_key = date_to_process.isoformat()
1502
+ sessions_for_date = all_sessions_data.get(date_key, {})
1503
+
1504
+ # Skip dates with no sessions
1505
+ if not any(len(sessions) > 0 for sessions in sessions_for_date.values()):
1506
+ continue
1507
+
1508
+ metrics_record = calculate_date_metrics(date_to_process, sessions_for_date)
1509
+ metrics_records.append(metrics_record)
1510
+
1511
+ if metrics_records:
1512
+ results = bulk_upsert_metrics(collection, metrics_records)
1513
+
1514
+ return results
1515
+
1516
+ except Exception as e:
1517
+ log_error(f"Error calculating metrics: {e}")
1518
+ raise e
1519
+
1520
+ def get_metrics(
1521
+ self,
1522
+ starting_date: Optional[date] = None,
1523
+ ending_date: Optional[date] = None,
1524
+ ) -> Tuple[List[dict], Optional[int]]:
1525
+ """Get all metrics matching the given date range."""
1526
+ try:
1527
+ collection = self._get_collection(table_type="metrics")
1528
+ if collection is None:
1529
+ return [], None
1530
+
1531
+ query = {}
1532
+ if starting_date:
1533
+ query["date"] = {"$gte": starting_date.isoformat()}
1534
+ if ending_date:
1535
+ if "date" in query:
1536
+ query["date"]["$lte"] = ending_date.isoformat()
1537
+ else:
1538
+ query["date"] = {"$lte": ending_date.isoformat()}
1539
+
1540
+ records = list(collection.find(query))
1541
+ if not records:
1542
+ return [], None
1543
+
1544
+ # Get the latest updated_at
1545
+ latest_updated_at = max(record.get("updated_at", 0) for record in records)
1546
+
1547
+ return records, latest_updated_at
1548
+
1549
+ except Exception as e:
1550
+ log_error(f"Error getting metrics: {e}")
1551
+ raise e
1552
+
1553
+ # -- Knowledge methods --
1554
+
1555
+ def delete_knowledge_content(self, id: str):
1556
+ """Delete a knowledge row from the database.
1557
+
1558
+ Args:
1559
+ id (str): The ID of the knowledge row to delete.
1560
+
1561
+ Raises:
1562
+ Exception: If an error occurs during deletion.
1563
+ """
1564
+ try:
1565
+ collection = self._get_collection(table_type="knowledge")
1566
+ if collection is None:
1567
+ return
1568
+
1569
+ collection.delete_one({"id": id})
1570
+
1571
+ log_debug(f"Deleted knowledge content with id '{id}'")
1572
+
1573
+ except Exception as e:
1574
+ log_error(f"Error deleting knowledge content: {e}")
1575
+ raise e
1576
+
1577
+ def get_knowledge_content(self, id: str) -> Optional[KnowledgeRow]:
1578
+ """Get a knowledge row from the database.
1579
+
1580
+ Args:
1581
+ id (str): The ID of the knowledge row to get.
1582
+
1583
+ Returns:
1584
+ Optional[KnowledgeRow]: The knowledge row, or None if it doesn't exist.
1585
+
1586
+ Raises:
1587
+ Exception: If an error occurs during retrieval.
1588
+ """
1589
+ try:
1590
+ collection = self._get_collection(table_type="knowledge")
1591
+ if collection is None:
1592
+ return None
1593
+
1594
+ result = collection.find_one({"id": id})
1595
+ if result is None:
1596
+ return None
1597
+
1598
+ return KnowledgeRow.model_validate(result)
1599
+
1600
+ except Exception as e:
1601
+ log_error(f"Error getting knowledge content: {e}")
1602
+ raise e
1603
+
1604
+ def get_knowledge_contents(
1605
+ self,
1606
+ limit: Optional[int] = None,
1607
+ page: Optional[int] = None,
1608
+ sort_by: Optional[str] = None,
1609
+ sort_order: Optional[str] = None,
1610
+ ) -> Tuple[List[KnowledgeRow], int]:
1611
+ """Get all knowledge contents from the database.
1612
+
1613
+ Args:
1614
+ limit (Optional[int]): The maximum number of knowledge contents to return.
1615
+ page (Optional[int]): The page number.
1616
+ sort_by (Optional[str]): The column to sort by.
1617
+ sort_order (Optional[str]): The order to sort by.
1618
+ create_table_if_not_found (Optional[bool]): Whether to create the collection if it doesn't exist.
1619
+
1620
+ Returns:
1621
+ Tuple[List[KnowledgeRow], int]: The knowledge contents and total count.
1622
+
1623
+ Raises:
1624
+ Exception: If an error occurs during retrieval.
1625
+ """
1626
+ try:
1627
+ collection = self._get_collection(table_type="knowledge")
1628
+ if collection is None:
1629
+ return [], 0
1630
+
1631
+ query: Dict[str, Any] = {}
1632
+
1633
+ # Get total count
1634
+ total_count = collection.count_documents(query)
1635
+
1636
+ # Apply sorting
1637
+ sort_criteria = apply_sorting({}, sort_by, sort_order)
1638
+
1639
+ # Apply pagination
1640
+ query_args = apply_pagination({}, limit, page)
1641
+
1642
+ cursor = collection.find(query)
1643
+ if sort_criteria:
1644
+ cursor = cursor.sort(sort_criteria)
1645
+ if query_args.get("skip"):
1646
+ cursor = cursor.skip(query_args["skip"])
1647
+ if query_args.get("limit"):
1648
+ cursor = cursor.limit(query_args["limit"])
1649
+
1650
+ records = list(cursor)
1651
+ knowledge_rows = [KnowledgeRow.model_validate(record) for record in records]
1652
+
1653
+ return knowledge_rows, total_count
1654
+
1655
+ except Exception as e:
1656
+ log_error(f"Error getting knowledge contents: {e}")
1657
+ raise e
1658
+
1659
+ def upsert_knowledge_content(self, knowledge_row: KnowledgeRow):
1660
+ """Upsert knowledge content in the database.
1661
+
1662
+ Args:
1663
+ knowledge_row (KnowledgeRow): The knowledge row to upsert.
1664
+
1665
+ Returns:
1666
+ Optional[KnowledgeRow]: The upserted knowledge row, or None if the operation fails.
1667
+
1668
+ Raises:
1669
+ Exception: If an error occurs during upsert.
1670
+ """
1671
+ try:
1672
+ collection = self._get_collection(table_type="knowledge", create_collection_if_not_found=True)
1673
+ if collection is None:
1674
+ return None
1675
+
1676
+ update_doc = knowledge_row.model_dump()
1677
+ collection.replace_one({"id": knowledge_row.id}, update_doc, upsert=True)
1678
+
1679
+ return knowledge_row
1680
+
1681
+ except Exception as e:
1682
+ log_error(f"Error upserting knowledge content: {e}")
1683
+ raise e
1684
+
1685
+ # -- Eval methods --
1686
+
1687
+ def create_eval_run(self, eval_run: EvalRunRecord) -> Optional[EvalRunRecord]:
1688
+ """Create an EvalRunRecord in the database."""
1689
+ try:
1690
+ collection = self._get_collection(table_type="evals", create_collection_if_not_found=True)
1691
+ if collection is None:
1692
+ return None
1693
+
1694
+ current_time = int(time.time())
1695
+ eval_dict = eval_run.model_dump()
1696
+ eval_dict["created_at"] = current_time
1697
+ eval_dict["updated_at"] = current_time
1698
+
1699
+ collection.insert_one(eval_dict)
1700
+
1701
+ log_debug(f"Created eval run with id '{eval_run.run_id}'")
1702
+
1703
+ return eval_run
1704
+
1705
+ except Exception as e:
1706
+ log_error(f"Error creating eval run: {e}")
1707
+ raise e
1708
+
1709
+ def delete_eval_run(self, eval_run_id: str) -> None:
1710
+ """Delete an eval run from the database."""
1711
+ try:
1712
+ collection = self._get_collection(table_type="evals")
1713
+ if collection is None:
1714
+ return
1715
+
1716
+ result = collection.delete_one({"run_id": eval_run_id})
1717
+
1718
+ if result.deleted_count == 0:
1719
+ log_debug(f"No eval run found with ID: {eval_run_id}")
1720
+ else:
1721
+ log_debug(f"Deleted eval run with ID: {eval_run_id}")
1722
+
1723
+ except Exception as e:
1724
+ log_error(f"Error deleting eval run {eval_run_id}: {e}")
1725
+ raise e
1726
+
1727
+ def delete_eval_runs(self, eval_run_ids: List[str]) -> None:
1728
+ """Delete multiple eval runs from the database."""
1729
+ try:
1730
+ collection = self._get_collection(table_type="evals")
1731
+ if collection is None:
1732
+ return
1733
+
1734
+ result = collection.delete_many({"run_id": {"$in": eval_run_ids}})
1735
+
1736
+ if result.deleted_count == 0:
1737
+ log_debug(f"No eval runs found with IDs: {eval_run_ids}")
1738
+ else:
1739
+ log_debug(f"Deleted {result.deleted_count} eval runs")
1740
+
1741
+ except Exception as e:
1742
+ log_error(f"Error deleting eval runs {eval_run_ids}: {e}")
1743
+ raise e
1744
+
1745
+ def get_eval_run_raw(self, eval_run_id: str) -> Optional[Dict[str, Any]]:
1746
+ """Get an eval run from the database as a raw dictionary."""
1747
+ try:
1748
+ collection = self._get_collection(table_type="evals")
1749
+ if collection is None:
1750
+ return None
1751
+
1752
+ result = collection.find_one({"run_id": eval_run_id})
1753
+ return result
1754
+
1755
+ except Exception as e:
1756
+ log_error(f"Exception getting eval run {eval_run_id}: {e}")
1757
+ raise e
1758
+
1759
+ def get_eval_run(self, eval_run_id: str, deserialize: Optional[bool] = True) -> Optional[EvalRunRecord]:
1760
+ """Get an eval run from the database.
1761
+
1762
+ Args:
1763
+ eval_run_id (str): The ID of the eval run to get.
1764
+ deserialize (Optional[bool]): Whether to serialize the eval run. Defaults to True.
1765
+
1766
+ Returns:
1767
+ Optional[EvalRunRecord]:
1768
+ - When deserialize=True: EvalRunRecord object
1769
+ - When deserialize=False: EvalRun dictionary
1770
+
1771
+ Raises:
1772
+ Exception: If there is an error getting the eval run.
1773
+ """
1774
+ try:
1775
+ collection = self._get_collection(table_type="evals")
1776
+ if collection is None:
1777
+ return None
1778
+
1779
+ eval_run_raw = collection.find_one({"run_id": eval_run_id})
1780
+
1781
+ if not eval_run_raw:
1782
+ return None
1783
+
1784
+ if not deserialize:
1785
+ return eval_run_raw
1786
+
1787
+ return EvalRunRecord.model_validate(eval_run_raw)
1788
+
1789
+ except Exception as e:
1790
+ log_error(f"Exception getting eval run {eval_run_id}: {e}")
1791
+ raise e
1792
+
1793
+ def get_eval_runs(
1794
+ self,
1795
+ limit: Optional[int] = None,
1796
+ page: Optional[int] = None,
1797
+ sort_by: Optional[str] = None,
1798
+ sort_order: Optional[str] = None,
1799
+ agent_id: Optional[str] = None,
1800
+ team_id: Optional[str] = None,
1801
+ workflow_id: Optional[str] = None,
1802
+ model_id: Optional[str] = None,
1803
+ filter_type: Optional[EvalFilterType] = None,
1804
+ eval_type: Optional[List[EvalType]] = None,
1805
+ deserialize: Optional[bool] = True,
1806
+ ) -> Union[List[EvalRunRecord], Tuple[List[Dict[str, Any]], int]]:
1807
+ """Get all eval runs from the database.
1808
+
1809
+ Args:
1810
+ limit (Optional[int]): The maximum number of eval runs to return.
1811
+ page (Optional[int]): The page number to return.
1812
+ sort_by (Optional[str]): The field to sort by.
1813
+ sort_order (Optional[str]): The order to sort by.
1814
+ agent_id (Optional[str]): The ID of the agent to filter by.
1815
+ team_id (Optional[str]): The ID of the team to filter by.
1816
+ workflow_id (Optional[str]): The ID of the workflow to filter by.
1817
+ model_id (Optional[str]): The ID of the model to filter by.
1818
+ eval_type (Optional[List[EvalType]]): The type of eval to filter by.
1819
+ filter_type (Optional[EvalFilterType]): The type of filter to apply.
1820
+ deserialize (Optional[bool]): Whether to serialize the eval runs. Defaults to True.
1821
+ create_table_if_not_found (Optional[bool]): Whether to create the collection if it doesn't exist.
1822
+
1823
+ Returns:
1824
+ Union[List[EvalRunRecord], Tuple[List[Dict[str, Any]], int]]:
1825
+ - When deserialize=True: List of EvalRunRecord objects
1826
+ - When deserialize=False: List of eval run dictionaries and the total count
1827
+
1828
+ Raises:
1829
+ Exception: If there is an error getting the eval runs.
1830
+ """
1831
+ try:
1832
+ collection = self._get_collection(table_type="evals")
1833
+ if collection is None:
1834
+ return [] if deserialize else ([], 0)
1835
+
1836
+ query: Dict[str, Any] = {}
1837
+ if agent_id is not None:
1838
+ query["agent_id"] = agent_id
1839
+ if team_id is not None:
1840
+ query["team_id"] = team_id
1841
+ if workflow_id is not None:
1842
+ query["workflow_id"] = workflow_id
1843
+ if model_id is not None:
1844
+ query["model_id"] = model_id
1845
+ if eval_type is not None and len(eval_type) > 0:
1846
+ query["eval_type"] = {"$in": eval_type}
1847
+ if filter_type is not None:
1848
+ if filter_type == EvalFilterType.AGENT:
1849
+ query["agent_id"] = {"$ne": None}
1850
+ elif filter_type == EvalFilterType.TEAM:
1851
+ query["team_id"] = {"$ne": None}
1852
+ elif filter_type == EvalFilterType.WORKFLOW:
1853
+ query["workflow_id"] = {"$ne": None}
1854
+
1855
+ # Get total count
1856
+ total_count = collection.count_documents(query)
1857
+
1858
+ # Apply default sorting by created_at desc if no sort parameters provided
1859
+ if sort_by is None:
1860
+ sort_criteria = [("created_at", -1)]
1861
+ else:
1862
+ sort_criteria = apply_sorting({}, sort_by, sort_order)
1863
+
1864
+ # Apply pagination
1865
+ query_args = apply_pagination({}, limit, page)
1866
+
1867
+ cursor = collection.find(query)
1868
+ if sort_criteria:
1869
+ cursor = cursor.sort(sort_criteria)
1870
+ if query_args.get("skip"):
1871
+ cursor = cursor.skip(query_args["skip"])
1872
+ if query_args.get("limit"):
1873
+ cursor = cursor.limit(query_args["limit"])
1874
+
1875
+ records = list(cursor)
1876
+ if not records:
1877
+ return [] if deserialize else ([], 0)
1878
+
1879
+ if not deserialize:
1880
+ return records, total_count
1881
+
1882
+ return [EvalRunRecord.model_validate(row) for row in records]
1883
+
1884
+ except Exception as e:
1885
+ log_error(f"Exception getting eval runs: {e}")
1886
+ raise e
1887
+
1888
+ def rename_eval_run(
1889
+ self, eval_run_id: str, name: str, deserialize: Optional[bool] = True
1890
+ ) -> Optional[Union[EvalRunRecord, Dict[str, Any]]]:
1891
+ """Update the name of an eval run in the database.
1892
+
1893
+ Args:
1894
+ eval_run_id (str): The ID of the eval run to update.
1895
+ name (str): The new name of the eval run.
1896
+ deserialize (Optional[bool]): Whether to serialize the eval run. Defaults to True.
1897
+
1898
+ Returns:
1899
+ Optional[Union[EvalRunRecord, Dict[str, Any]]]:
1900
+ - When deserialize=True: EvalRunRecord object
1901
+ - When deserialize=False: EvalRun dictionary
1902
+
1903
+ Raises:
1904
+ Exception: If there is an error updating the eval run.
1905
+ """
1906
+ try:
1907
+ collection = self._get_collection(table_type="evals")
1908
+ if collection is None:
1909
+ return None
1910
+
1911
+ result = collection.find_one_and_update(
1912
+ {"run_id": eval_run_id}, {"$set": {"name": name, "updated_at": int(time.time())}}
1913
+ )
1914
+
1915
+ log_debug(f"Renamed eval run with id '{eval_run_id}' to '{name}'")
1916
+
1917
+ if not result or not deserialize:
1918
+ return result
1919
+
1920
+ return EvalRunRecord.model_validate(result)
1921
+
1922
+ except Exception as e:
1923
+ log_error(f"Error updating eval run name {eval_run_id}: {e}")
1924
+ raise e
1925
+
1926
+ def migrate_table_from_v1_to_v2(self, v1_db_schema: str, v1_table_name: str, v1_table_type: str):
1927
+ """Migrate all content in the given collection to the right v2 collection"""
1928
+
1929
+ from typing import List, Sequence, Union
1930
+
1931
+ from agno.db.migrations.v1_to_v2 import (
1932
+ get_all_table_content,
1933
+ parse_agent_sessions,
1934
+ parse_memories,
1935
+ parse_team_sessions,
1936
+ parse_workflow_sessions,
1937
+ )
1938
+
1939
+ # Get all content from the old collection
1940
+ old_content: list[dict[str, Any]] = get_all_table_content(
1941
+ db=self,
1942
+ db_schema=v1_db_schema,
1943
+ table_name=v1_table_name,
1944
+ )
1945
+ if not old_content:
1946
+ log_info(f"No content to migrate from collection {v1_table_name}")
1947
+ return
1948
+
1949
+ # Parse the content into the new format
1950
+ memories: List[UserMemory] = []
1951
+ sessions: Sequence[Union[AgentSession, TeamSession, WorkflowSession]] = []
1952
+ if v1_table_type == "agent_sessions":
1953
+ sessions = parse_agent_sessions(old_content)
1954
+ elif v1_table_type == "team_sessions":
1955
+ sessions = parse_team_sessions(old_content)
1956
+ elif v1_table_type == "workflow_sessions":
1957
+ sessions = parse_workflow_sessions(old_content)
1958
+ elif v1_table_type == "memories":
1959
+ memories = parse_memories(old_content)
1960
+ else:
1961
+ raise ValueError(f"Invalid table type: {v1_table_type}")
1962
+
1963
+ # Insert the new content into the new collection
1964
+ if v1_table_type == "agent_sessions":
1965
+ for session in sessions:
1966
+ self.upsert_session(session)
1967
+ log_info(f"Migrated {len(sessions)} Agent sessions to collection: {self.session_table_name}")
1968
+
1969
+ elif v1_table_type == "team_sessions":
1970
+ for session in sessions:
1971
+ self.upsert_session(session)
1972
+ log_info(f"Migrated {len(sessions)} Team sessions to collection: {self.session_table_name}")
1973
+
1974
+ elif v1_table_type == "workflow_sessions":
1975
+ for session in sessions:
1976
+ self.upsert_session(session)
1977
+ log_info(f"Migrated {len(sessions)} Workflow sessions to collection: {self.session_table_name}")
1978
+
1979
+ elif v1_table_type == "memories":
1980
+ for memory in memories:
1981
+ self.upsert_user_memory(memory)
1982
+ log_info(f"Migrated {len(memories)} memories to collection: {self.memory_table_name}")