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
@@ -0,0 +1,1795 @@
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.firestore.utils import (
8
+ apply_pagination,
9
+ apply_pagination_to_records,
10
+ apply_sorting,
11
+ apply_sorting_to_records,
12
+ bulk_upsert_metrics,
13
+ calculate_date_metrics,
14
+ create_collection_indexes,
15
+ deserialize_cultural_knowledge_from_db,
16
+ fetch_all_sessions_data,
17
+ get_dates_to_calculate_metrics_for,
18
+ serialize_cultural_knowledge_for_db,
19
+ )
20
+ from agno.db.schemas.culture import CulturalKnowledge
21
+ from agno.db.schemas.evals import EvalFilterType, EvalRunRecord, EvalType
22
+ from agno.db.schemas.knowledge import KnowledgeRow
23
+ from agno.db.schemas.memory import UserMemory
24
+ from agno.db.utils import deserialize_session_json_fields, serialize_session_json_fields
25
+ from agno.session import AgentSession, Session, TeamSession, WorkflowSession
26
+ from agno.utils.log import log_debug, log_error, log_info
27
+ from agno.utils.string import generate_id
28
+
29
+ try:
30
+ from google.cloud.firestore import Client, FieldFilter # type: ignore[import-untyped]
31
+ except ImportError:
32
+ raise ImportError(
33
+ "`google-cloud-firestore` not installed. Please install it using `pip install google-cloud-firestore`"
34
+ )
35
+
36
+
37
+ class FirestoreDb(BaseDb):
38
+ def __init__(
39
+ self,
40
+ db_client: Optional[Client] = None,
41
+ project_id: 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 Firestore database.
52
+
53
+ Args:
54
+ db_client (Optional[Client]): The Firestore client to use.
55
+ project_id (Optional[str]): The GCP project ID for Firestore.
56
+ session_collection (Optional[str]): Name of the collection to store sessions.
57
+ memory_collection (Optional[str]): Name of the collection to store memories.
58
+ metrics_collection (Optional[str]): Name of the collection to store metrics.
59
+ eval_collection (Optional[str]): Name of the collection to store evaluation runs.
60
+ knowledge_collection (Optional[str]): Name of the collection to store knowledge documents.
61
+ culture_collection (Optional[str]): Name of the collection to store cultural knowledge.
62
+ id (Optional[str]): ID of the database.
63
+
64
+ Raises:
65
+ ValueError: If neither project_id nor db_client is provided.
66
+ """
67
+ if id is None:
68
+ seed = project_id or str(db_client)
69
+ id = generate_id(seed)
70
+
71
+ super().__init__(
72
+ id=id,
73
+ session_table=session_collection,
74
+ memory_table=memory_collection,
75
+ metrics_table=metrics_collection,
76
+ eval_table=eval_collection,
77
+ knowledge_table=knowledge_collection,
78
+ culture_table=culture_collection,
79
+ )
80
+
81
+ _client: Optional[Client] = db_client
82
+ if _client is None and project_id is not None:
83
+ _client = Client(project=project_id)
84
+ if _client is None:
85
+ raise ValueError("One of project_id or db_client must be provided")
86
+
87
+ self.project_id: Optional[str] = project_id
88
+ self.db_client: Client = _client
89
+
90
+ # -- DB methods --
91
+
92
+ def table_exists(self, table_name: str) -> bool:
93
+ """Check if a collection with the given name exists in the Firestore database.
94
+
95
+ Args:
96
+ table_name: Name of the collection to check
97
+
98
+ Returns:
99
+ bool: True if the collection exists in the database, False otherwise
100
+ """
101
+ return table_name in self.db_client.list_collections()
102
+
103
+ def _get_collection(self, table_type: str, create_collection_if_not_found: Optional[bool] = True):
104
+ """Get or create a collection based on table type.
105
+
106
+ Args:
107
+ table_type (str): The type of table to get or create.
108
+ create_collection_if_not_found (Optional[bool]): Whether to create the collection if it doesn't exist.
109
+
110
+ Returns:
111
+ CollectionReference: The collection reference.
112
+ """
113
+ if table_type == "sessions":
114
+ if not hasattr(self, "session_collection"):
115
+ if self.session_table_name is None:
116
+ raise ValueError("Session collection was not provided on initialization")
117
+ self.session_collection = self._get_or_create_collection(
118
+ collection_name=self.session_table_name,
119
+ collection_type="sessions",
120
+ create_collection_if_not_found=create_collection_if_not_found,
121
+ )
122
+ return self.session_collection
123
+
124
+ if table_type == "memories":
125
+ if not hasattr(self, "memory_collection"):
126
+ if self.memory_table_name is None:
127
+ raise ValueError("Memory collection was not provided on initialization")
128
+ self.memory_collection = self._get_or_create_collection(
129
+ collection_name=self.memory_table_name,
130
+ collection_type="memories",
131
+ create_collection_if_not_found=create_collection_if_not_found,
132
+ )
133
+ return self.memory_collection
134
+
135
+ if table_type == "metrics":
136
+ if not hasattr(self, "metrics_collection"):
137
+ if self.metrics_table_name is None:
138
+ raise ValueError("Metrics collection was not provided on initialization")
139
+ self.metrics_collection = self._get_or_create_collection(
140
+ collection_name=self.metrics_table_name,
141
+ collection_type="metrics",
142
+ create_collection_if_not_found=create_collection_if_not_found,
143
+ )
144
+ return self.metrics_collection
145
+
146
+ if table_type == "evals":
147
+ if not hasattr(self, "eval_collection"):
148
+ if self.eval_table_name is None:
149
+ raise ValueError("Eval collection was not provided on initialization")
150
+ self.eval_collection = self._get_or_create_collection(
151
+ collection_name=self.eval_table_name,
152
+ collection_type="evals",
153
+ create_collection_if_not_found=create_collection_if_not_found,
154
+ )
155
+ return self.eval_collection
156
+
157
+ if table_type == "knowledge":
158
+ if not hasattr(self, "knowledge_collection"):
159
+ if self.knowledge_table_name is None:
160
+ raise ValueError("Knowledge collection was not provided on initialization")
161
+ self.knowledge_collection = self._get_or_create_collection(
162
+ collection_name=self.knowledge_table_name,
163
+ collection_type="knowledge",
164
+ create_collection_if_not_found=create_collection_if_not_found,
165
+ )
166
+ return self.knowledge_collection
167
+
168
+ if table_type == "culture":
169
+ if not hasattr(self, "culture_collection"):
170
+ if self.culture_table_name is None:
171
+ raise ValueError("Culture collection was not provided on initialization")
172
+ self.culture_collection = self._get_or_create_collection(
173
+ collection_name=self.culture_table_name,
174
+ collection_type="culture",
175
+ create_collection_if_not_found=create_collection_if_not_found,
176
+ )
177
+ return self.culture_collection
178
+
179
+ raise ValueError(f"Unknown table type: {table_type}")
180
+
181
+ def _get_or_create_collection(
182
+ self, collection_name: str, collection_type: str, create_collection_if_not_found: Optional[bool] = True
183
+ ):
184
+ """Get or create a collection with proper indexes.
185
+
186
+ Args:
187
+ collection_name (str): The name of the collection to get or create.
188
+ collection_type (str): The type of collection to get or create.
189
+ create_collection_if_not_found (Optional[bool]): Whether to create the collection if it doesn't exist.
190
+
191
+ Returns:
192
+ Optional[CollectionReference]: The collection reference.
193
+ """
194
+ try:
195
+ collection_ref = self.db_client.collection(collection_name)
196
+
197
+ if not hasattr(self, f"_{collection_name}_initialized"):
198
+ if not create_collection_if_not_found:
199
+ return None
200
+ create_collection_indexes(self.db_client, collection_name, collection_type)
201
+ setattr(self, f"_{collection_name}_initialized", True)
202
+
203
+ return collection_ref
204
+
205
+ except Exception as e:
206
+ log_error(f"Error getting collection {collection_name}: {e}")
207
+ raise
208
+
209
+ # -- Session methods --
210
+
211
+ def delete_session(self, session_id: str) -> bool:
212
+ """Delete a session from the database.
213
+
214
+ Args:
215
+ session_id (str): The ID of the session to delete.
216
+ session_type (SessionType): The type of session to delete. Defaults to SessionType.AGENT.
217
+
218
+ Returns:
219
+ bool: True if the session was deleted, False otherwise.
220
+
221
+ Raises:
222
+ Exception: If there is an error deleting the session.
223
+ """
224
+ try:
225
+ collection_ref = self._get_collection(table_type="sessions")
226
+ docs = collection_ref.where(filter=FieldFilter("session_id", "==", session_id)).stream()
227
+
228
+ for doc in docs:
229
+ doc.reference.delete()
230
+ log_debug(f"Successfully deleted session with session_id: {session_id}")
231
+ return True
232
+
233
+ log_debug(f"No session found to delete with session_id: {session_id}")
234
+ return False
235
+
236
+ except Exception as e:
237
+ log_error(f"Error deleting session: {e}")
238
+ raise e
239
+
240
+ def delete_sessions(self, session_ids: List[str]) -> None:
241
+ """Delete multiple sessions from the database.
242
+
243
+ Args:
244
+ session_ids (List[str]): The IDs of the sessions to delete.
245
+ """
246
+ try:
247
+ collection_ref = self._get_collection(table_type="sessions")
248
+ batch = self.db_client.batch()
249
+
250
+ deleted_count = 0
251
+ for session_id in session_ids:
252
+ docs = collection_ref.where(filter=FieldFilter("session_id", "==", session_id)).stream()
253
+ for doc in docs:
254
+ batch.delete(doc.reference)
255
+ deleted_count += 1
256
+
257
+ batch.commit()
258
+
259
+ log_debug(f"Successfully deleted {deleted_count} sessions")
260
+
261
+ except Exception as e:
262
+ log_error(f"Error deleting sessions: {e}")
263
+ raise e
264
+
265
+ def get_session(
266
+ self,
267
+ session_id: str,
268
+ session_type: SessionType,
269
+ user_id: Optional[str] = None,
270
+ deserialize: Optional[bool] = True,
271
+ ) -> Optional[Union[Session, Dict[str, Any]]]:
272
+ """Read a session from the database.
273
+
274
+ Args:
275
+ session_id (str): The ID of the session to get.
276
+ session_type (SessionType): The type of session to get.
277
+ user_id (Optional[str]): The ID of the user to get the session for.
278
+ deserialize (Optional[bool]): Whether to serialize the session. Defaults to True.
279
+
280
+ Returns:
281
+ Union[Session, Dict[str, Any], None]:
282
+ - When deserialize=True: Session object
283
+ - When deserialize=False: Session dictionary
284
+
285
+ Raises:
286
+ Exception: If there is an error reading the session.
287
+ """
288
+ try:
289
+ collection_ref = self._get_collection(table_type="sessions")
290
+ query = collection_ref.where(filter=FieldFilter("session_id", "==", session_id))
291
+
292
+ if user_id is not None:
293
+ query = query.where(filter=FieldFilter("user_id", "==", user_id))
294
+
295
+ docs = query.stream()
296
+ result = None
297
+ for doc in docs:
298
+ result = doc.to_dict()
299
+ break
300
+
301
+ if result is None:
302
+ return None
303
+
304
+ session = deserialize_session_json_fields(result)
305
+
306
+ if not deserialize:
307
+ return session
308
+
309
+ if session_type == SessionType.AGENT:
310
+ return AgentSession.from_dict(session)
311
+ elif session_type == SessionType.TEAM:
312
+ return TeamSession.from_dict(session)
313
+ elif session_type == SessionType.WORKFLOW:
314
+ return WorkflowSession.from_dict(session)
315
+ else:
316
+ raise ValueError(f"Invalid session type: {session_type}")
317
+
318
+ except Exception as e:
319
+ log_error(f"Exception reading session: {e}")
320
+ raise e
321
+
322
+ def get_sessions(
323
+ self,
324
+ session_type: Optional[SessionType] = None,
325
+ user_id: Optional[str] = None,
326
+ component_id: Optional[str] = None,
327
+ session_name: Optional[str] = None,
328
+ start_timestamp: Optional[int] = None,
329
+ end_timestamp: Optional[int] = None,
330
+ limit: Optional[int] = None,
331
+ page: Optional[int] = None,
332
+ sort_by: Optional[str] = None,
333
+ sort_order: Optional[str] = None,
334
+ deserialize: Optional[bool] = True,
335
+ ) -> Union[List[Session], Tuple[List[Dict[str, Any]], int]]:
336
+ """Get all sessions.
337
+
338
+ Args:
339
+ session_type (Optional[SessionType]): The type of session to get.
340
+ user_id (Optional[str]): The ID of the user to get the session for.
341
+ component_id (Optional[str]): The ID of the component to get the session for.
342
+ session_name (Optional[str]): The name of the session to filter by.
343
+ start_timestamp (Optional[int]): The start timestamp to filter sessions by.
344
+ end_timestamp (Optional[int]): The end timestamp to filter sessions by.
345
+ limit (Optional[int]): The limit of the sessions to get.
346
+ page (Optional[int]): The page number to get.
347
+ sort_by (Optional[str]): The field to sort the sessions by.
348
+ sort_order (Optional[str]): The order to sort the sessions by.
349
+ deserialize (Optional[bool]): Whether to serialize the sessions. Defaults to True.
350
+
351
+ Returns:
352
+ Union[List[AgentSession], List[TeamSession], List[WorkflowSession], Tuple[List[Dict[str, Any]], int]]:
353
+ - When deserialize=True: List of Session objects
354
+ - When deserialize=False: List of session dictionaries and the total count
355
+
356
+ Raises:
357
+ Exception: If there is an error reading the sessions.
358
+ """
359
+ try:
360
+ collection_ref = self._get_collection(table_type="sessions")
361
+ if collection_ref is None:
362
+ return [] if deserialize else ([], 0)
363
+
364
+ query = collection_ref
365
+
366
+ if user_id is not None:
367
+ query = query.where(filter=FieldFilter("user_id", "==", user_id))
368
+ if session_type is not None:
369
+ query = query.where(filter=FieldFilter("session_type", "==", session_type.value))
370
+ if component_id is not None:
371
+ if session_type == SessionType.AGENT:
372
+ query = query.where(filter=FieldFilter("agent_id", "==", component_id))
373
+ elif session_type == SessionType.TEAM:
374
+ query = query.where(filter=FieldFilter("team_id", "==", component_id))
375
+ elif session_type == SessionType.WORKFLOW:
376
+ query = query.where(filter=FieldFilter("workflow_id", "==", component_id))
377
+ if start_timestamp is not None:
378
+ query = query.where(filter=FieldFilter("created_at", ">=", start_timestamp))
379
+ if end_timestamp is not None:
380
+ query = query.where(filter=FieldFilter("created_at", "<=", end_timestamp))
381
+ if session_name is not None:
382
+ query = query.where(filter=FieldFilter("session_data.session_name", "==", session_name))
383
+
384
+ # Apply sorting
385
+ query = apply_sorting(query, sort_by, sort_order)
386
+
387
+ # Get all documents for counting before pagination
388
+ all_docs = query.stream()
389
+ all_records = [doc.to_dict() for doc in all_docs]
390
+
391
+ if not all_records:
392
+ return [] if deserialize else ([], 0)
393
+
394
+ all_sessions_raw = [deserialize_session_json_fields(record) for record in all_records]
395
+
396
+ # Get total count before pagination
397
+ total_count = len(all_sessions_raw)
398
+
399
+ # Apply pagination to the results
400
+ if limit is not None and page is not None:
401
+ start_index = (page - 1) * limit
402
+ end_index = start_index + limit
403
+ sessions_raw = all_sessions_raw[start_index:end_index]
404
+ elif limit is not None:
405
+ sessions_raw = all_sessions_raw[:limit]
406
+ else:
407
+ sessions_raw = all_sessions_raw
408
+
409
+ if not deserialize:
410
+ return sessions_raw, total_count
411
+
412
+ sessions: List[Union[AgentSession, TeamSession, WorkflowSession]] = []
413
+ for session in sessions_raw:
414
+ if session["session_type"] == SessionType.AGENT.value:
415
+ agent_session = AgentSession.from_dict(session)
416
+ if agent_session is not None:
417
+ sessions.append(agent_session)
418
+ elif session["session_type"] == SessionType.TEAM.value:
419
+ team_session = TeamSession.from_dict(session)
420
+ if team_session is not None:
421
+ sessions.append(team_session)
422
+ elif session["session_type"] == SessionType.WORKFLOW.value:
423
+ workflow_session = WorkflowSession.from_dict(session)
424
+ if workflow_session is not None:
425
+ sessions.append(workflow_session)
426
+
427
+ if not sessions:
428
+ return [] if deserialize else ([], 0)
429
+
430
+ return sessions
431
+
432
+ except Exception as e:
433
+ log_error(f"Exception reading sessions: {e}")
434
+ raise e
435
+
436
+ def rename_session(
437
+ self, session_id: str, session_type: SessionType, session_name: str, deserialize: Optional[bool] = True
438
+ ) -> Optional[Union[Session, Dict[str, Any]]]:
439
+ """Rename a session in the database.
440
+
441
+ Args:
442
+ session_id (str): The ID of the session to rename.
443
+ session_type (SessionType): The type of session to rename.
444
+ session_name (str): The new name of the session.
445
+ deserialize (Optional[bool]): Whether to serialize the session. Defaults to True.
446
+
447
+ Returns:
448
+ Optional[Union[Session, Dict[str, Any]]]:
449
+ - When deserialize=True: Session object
450
+ - When deserialize=False: Session dictionary
451
+
452
+ Raises:
453
+ Exception: If there is an error renaming the session.
454
+ """
455
+ try:
456
+ collection_ref = self._get_collection(table_type="sessions")
457
+
458
+ docs = collection_ref.where(filter=FieldFilter("session_id", "==", session_id)).stream()
459
+ doc_ref = next((doc.reference for doc in docs), None)
460
+
461
+ if doc_ref is None:
462
+ return None
463
+
464
+ doc_ref.update({"session_data.session_name": session_name, "updated_at": int(time.time())})
465
+
466
+ updated_doc = doc_ref.get()
467
+ if not updated_doc.exists:
468
+ return None
469
+
470
+ result = updated_doc.to_dict()
471
+ if result is None:
472
+ return None
473
+ deserialized_session = deserialize_session_json_fields(result)
474
+
475
+ log_debug(f"Renamed session with id '{session_id}' to '{session_name}'")
476
+
477
+ if not deserialize:
478
+ return deserialized_session
479
+
480
+ if session_type == SessionType.AGENT:
481
+ return AgentSession.from_dict(deserialized_session)
482
+ elif session_type == SessionType.TEAM:
483
+ return TeamSession.from_dict(deserialized_session)
484
+ else:
485
+ return WorkflowSession.from_dict(deserialized_session)
486
+
487
+ except Exception as e:
488
+ log_error(f"Exception renaming session: {e}")
489
+ raise e
490
+
491
+ def upsert_session(
492
+ self, session: Session, deserialize: Optional[bool] = True
493
+ ) -> Optional[Union[Session, Dict[str, Any]]]:
494
+ """Insert or update a session in the database.
495
+
496
+ Args:
497
+ session (Session): The session to upsert.
498
+
499
+ Returns:
500
+ Optional[Session]: The upserted session.
501
+
502
+ Raises:
503
+ Exception: If there is an error upserting the session.
504
+ """
505
+ try:
506
+ collection_ref = self._get_collection(table_type="sessions", create_collection_if_not_found=True)
507
+ serialized_session_dict = serialize_session_json_fields(session.to_dict())
508
+
509
+ if isinstance(session, AgentSession):
510
+ record = {
511
+ "session_id": serialized_session_dict.get("session_id"),
512
+ "session_type": SessionType.AGENT.value,
513
+ "agent_id": serialized_session_dict.get("agent_id"),
514
+ "user_id": serialized_session_dict.get("user_id"),
515
+ "runs": serialized_session_dict.get("runs"),
516
+ "agent_data": serialized_session_dict.get("agent_data"),
517
+ "session_data": serialized_session_dict.get("session_data"),
518
+ "summary": serialized_session_dict.get("summary"),
519
+ "metadata": serialized_session_dict.get("metadata"),
520
+ "created_at": serialized_session_dict.get("created_at"),
521
+ "updated_at": int(time.time()),
522
+ }
523
+
524
+ elif isinstance(session, TeamSession):
525
+ record = {
526
+ "session_id": serialized_session_dict.get("session_id"),
527
+ "session_type": SessionType.TEAM.value,
528
+ "team_id": serialized_session_dict.get("team_id"),
529
+ "user_id": serialized_session_dict.get("user_id"),
530
+ "runs": serialized_session_dict.get("runs"),
531
+ "team_data": serialized_session_dict.get("team_data"),
532
+ "session_data": serialized_session_dict.get("session_data"),
533
+ "summary": serialized_session_dict.get("summary"),
534
+ "metadata": serialized_session_dict.get("metadata"),
535
+ "created_at": serialized_session_dict.get("created_at"),
536
+ "updated_at": int(time.time()),
537
+ }
538
+
539
+ elif isinstance(session, WorkflowSession):
540
+ record = {
541
+ "session_id": serialized_session_dict.get("session_id"),
542
+ "session_type": SessionType.WORKFLOW.value,
543
+ "workflow_id": serialized_session_dict.get("workflow_id"),
544
+ "user_id": serialized_session_dict.get("user_id"),
545
+ "runs": serialized_session_dict.get("runs"),
546
+ "workflow_data": serialized_session_dict.get("workflow_data"),
547
+ "session_data": serialized_session_dict.get("session_data"),
548
+ "summary": serialized_session_dict.get("summary"),
549
+ "metadata": serialized_session_dict.get("metadata"),
550
+ "created_at": serialized_session_dict.get("created_at"),
551
+ "updated_at": int(time.time()),
552
+ }
553
+
554
+ # Find existing document or create new one
555
+ docs = collection_ref.where(filter=FieldFilter("session_id", "==", record["session_id"])).stream()
556
+ doc_ref = next((doc.reference for doc in docs), None)
557
+
558
+ if doc_ref is None:
559
+ # Create new document
560
+ doc_ref = collection_ref.document()
561
+
562
+ doc_ref.set(record, merge=True)
563
+
564
+ # Get the updated document
565
+ updated_doc = doc_ref.get()
566
+ if not updated_doc.exists:
567
+ return None
568
+
569
+ result = updated_doc.to_dict()
570
+ if result is None:
571
+ return None
572
+ deserialized_session = deserialize_session_json_fields(result)
573
+
574
+ if not deserialize:
575
+ return deserialized_session
576
+
577
+ if isinstance(session, AgentSession):
578
+ return AgentSession.from_dict(deserialized_session)
579
+ elif isinstance(session, TeamSession):
580
+ return TeamSession.from_dict(deserialized_session)
581
+ else:
582
+ return WorkflowSession.from_dict(deserialized_session)
583
+
584
+ except Exception as e:
585
+ log_error(f"Exception upserting session: {e}")
586
+ raise e
587
+
588
+ def upsert_sessions(
589
+ self, sessions: List[Session], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
590
+ ) -> List[Union[Session, Dict[str, Any]]]:
591
+ """
592
+ Bulk upsert multiple sessions for improved performance on large datasets.
593
+
594
+ Args:
595
+ sessions (List[Session]): List of sessions to upsert.
596
+ deserialize (Optional[bool]): Whether to deserialize the sessions. Defaults to True.
597
+
598
+ Returns:
599
+ List[Union[Session, Dict[str, Any]]]: List of upserted sessions.
600
+
601
+ Raises:
602
+ Exception: If an error occurs during bulk upsert.
603
+ """
604
+ if not sessions:
605
+ return []
606
+
607
+ try:
608
+ log_info(
609
+ f"FirestoreDb doesn't support efficient bulk operations, falling back to individual upserts for {len(sessions)} sessions"
610
+ )
611
+
612
+ # Fall back to individual upserts
613
+ results = []
614
+ for session in sessions:
615
+ if session is not None:
616
+ result = self.upsert_session(session, deserialize=deserialize)
617
+ if result is not None:
618
+ results.append(result)
619
+ return results
620
+
621
+ except Exception as e:
622
+ log_error(f"Exception during bulk session upsert: {e}")
623
+ return []
624
+
625
+ # -- Memory methods --
626
+
627
+ def delete_user_memory(self, memory_id: str, user_id: Optional[str] = None):
628
+ """Delete a user memory from the database.
629
+
630
+ Args:
631
+ memory_id (str): The ID of the memory to delete.
632
+ user_id (Optional[str]): The ID of the user (optional, for filtering).
633
+
634
+ Returns:
635
+ bool: True if the memory was deleted, False otherwise.
636
+
637
+ Raises:
638
+ Exception: If there is an error deleting the memory.
639
+ """
640
+ try:
641
+ collection_ref = self._get_collection(table_type="memories")
642
+
643
+ # If user_id is provided, verify the memory belongs to the user before deleting
644
+ if user_id:
645
+ docs = collection_ref.where(filter=FieldFilter("memory_id", "==", memory_id)).stream()
646
+ for doc in docs:
647
+ data = doc.to_dict()
648
+ if data.get("user_id") != user_id:
649
+ log_debug(f"Memory {memory_id} does not belong to user {user_id}")
650
+ return
651
+ doc.reference.delete()
652
+ log_debug(f"Successfully deleted user memory id: {memory_id}")
653
+ return
654
+ else:
655
+ docs = collection_ref.where(filter=FieldFilter("memory_id", "==", memory_id)).stream()
656
+ deleted_count = 0
657
+ for doc in docs:
658
+ doc.reference.delete()
659
+ deleted_count += 1
660
+
661
+ success = deleted_count > 0
662
+ if success:
663
+ log_debug(f"Successfully deleted user memory id: {memory_id}")
664
+ else:
665
+ log_debug(f"No user memory found with id: {memory_id}")
666
+
667
+ except Exception as e:
668
+ log_error(f"Error deleting user memory: {e}")
669
+ raise e
670
+
671
+ def delete_user_memories(self, memory_ids: List[str], user_id: Optional[str] = None) -> None:
672
+ """Delete user memories from the database.
673
+
674
+ Args:
675
+ memory_ids (List[str]): The IDs of the memories to delete.
676
+ user_id (Optional[str]): The ID of the user (optional, for filtering).
677
+
678
+ Raises:
679
+ Exception: If there is an error deleting the memories.
680
+ """
681
+ try:
682
+ collection_ref = self._get_collection(table_type="memories")
683
+ batch = self.db_client.batch()
684
+ deleted_count = 0
685
+
686
+ # If user_id is provided, filter memory_ids to only those belonging to the user
687
+ if user_id:
688
+ for memory_id in memory_ids:
689
+ docs = collection_ref.where(filter=FieldFilter("memory_id", "==", memory_id)).stream()
690
+ for doc in docs:
691
+ data = doc.to_dict()
692
+ if data.get("user_id") == user_id:
693
+ batch.delete(doc.reference)
694
+ deleted_count += 1
695
+ else:
696
+ for memory_id in memory_ids:
697
+ docs = collection_ref.where(filter=FieldFilter("memory_id", "==", memory_id)).stream()
698
+ for doc in docs:
699
+ batch.delete(doc.reference)
700
+ deleted_count += 1
701
+
702
+ batch.commit()
703
+
704
+ if deleted_count == 0:
705
+ log_info(f"No memories found with ids: {memory_ids}")
706
+ else:
707
+ log_info(f"Successfully deleted {deleted_count} memories")
708
+
709
+ except Exception as e:
710
+ log_error(f"Error deleting memories: {e}")
711
+ raise e
712
+
713
+ def get_all_memory_topics(self, create_collection_if_not_found: Optional[bool] = True) -> List[str]:
714
+ """Get all memory topics from the database.
715
+
716
+ Returns:
717
+ List[str]: The topics.
718
+
719
+ Raises:
720
+ Exception: If there is an error getting the topics.
721
+ """
722
+ try:
723
+ collection_ref = self._get_collection(table_type="memories")
724
+ if collection_ref is None:
725
+ return []
726
+
727
+ docs = collection_ref.stream()
728
+
729
+ all_topics = set()
730
+ for doc in docs:
731
+ data = doc.to_dict()
732
+ topics = data.get("topics", [])
733
+ if topics:
734
+ all_topics.update(topics)
735
+
736
+ return [topic for topic in all_topics if topic]
737
+
738
+ except Exception as e:
739
+ log_error(f"Exception getting all memory topics: {e}")
740
+ raise e
741
+
742
+ def get_user_memory(
743
+ self, memory_id: str, deserialize: Optional[bool] = True, user_id: Optional[str] = None
744
+ ) -> Optional[UserMemory]:
745
+ """Get a memory from the database.
746
+
747
+ Args:
748
+ memory_id (str): The ID of the memory to get.
749
+ deserialize (Optional[bool]): Whether to serialize the memory. Defaults to True.
750
+ user_id (Optional[str]): The ID of the user (optional, for filtering).
751
+
752
+ Returns:
753
+ Optional[UserMemory]:
754
+ - When deserialize=True: UserMemory object
755
+ - When deserialize=False: Memory dictionary
756
+
757
+ Raises:
758
+ Exception: If there is an error getting the memory.
759
+ """
760
+ try:
761
+ collection_ref = self._get_collection(table_type="memories")
762
+ docs = collection_ref.where(filter=FieldFilter("memory_id", "==", memory_id)).stream()
763
+
764
+ result = None
765
+ for doc in docs:
766
+ result = doc.to_dict()
767
+ break
768
+
769
+ if result is None:
770
+ return None
771
+
772
+ # Filter by user_id if provided
773
+ if user_id and result.get("user_id") != user_id:
774
+ return None
775
+
776
+ if not deserialize:
777
+ return result
778
+
779
+ return UserMemory.from_dict(result)
780
+
781
+ except Exception as e:
782
+ log_error(f"Exception getting user memory: {e}")
783
+ raise e
784
+
785
+ def get_user_memories(
786
+ self,
787
+ user_id: Optional[str] = None,
788
+ agent_id: Optional[str] = None,
789
+ team_id: Optional[str] = None,
790
+ topics: Optional[List[str]] = None,
791
+ search_content: Optional[str] = None,
792
+ limit: Optional[int] = None,
793
+ page: Optional[int] = None,
794
+ sort_by: Optional[str] = None,
795
+ sort_order: Optional[str] = None,
796
+ deserialize: Optional[bool] = True,
797
+ ) -> Union[List[UserMemory], Tuple[List[Dict[str, Any]], int]]:
798
+ """Get all memories from the database as UserMemory objects.
799
+
800
+ Args:
801
+ user_id (Optional[str]): The ID of the user to get the memories for.
802
+ agent_id (Optional[str]): The ID of the agent to get the memories for.
803
+ team_id (Optional[str]): The ID of the team to get the memories for.
804
+ topics (Optional[List[str]]): The topics to filter the memories by.
805
+ search_content (Optional[str]): The content to filter the memories by.
806
+ limit (Optional[int]): The limit of the memories to get.
807
+ page (Optional[int]): The page number to get.
808
+ sort_by (Optional[str]): The field to sort the memories by.
809
+ sort_order (Optional[str]): The order to sort the memories by.
810
+ deserialize (Optional[bool]): Whether to serialize the memories. Defaults to True.
811
+ create_table_if_not_found: Whether to create the index if it doesn't exist.
812
+
813
+ Returns:
814
+ Tuple[List[Dict[str, Any]], int]: A tuple containing the memories and the total count.
815
+
816
+ Raises:
817
+ Exception: If there is an error getting the memories.
818
+ """
819
+ try:
820
+ collection_ref = self._get_collection(table_type="memories")
821
+ if collection_ref is None:
822
+ return [] if deserialize else ([], 0)
823
+
824
+ query = collection_ref
825
+
826
+ if user_id is not None:
827
+ query = query.where(filter=FieldFilter("user_id", "==", user_id))
828
+ if agent_id is not None:
829
+ query = query.where(filter=FieldFilter("agent_id", "==", agent_id))
830
+ if team_id is not None:
831
+ query = query.where(filter=FieldFilter("team_id", "==", team_id))
832
+ if topics is not None and len(topics) > 0:
833
+ query = query.where(filter=FieldFilter("topics", "array_contains_any", topics))
834
+ if search_content is not None:
835
+ query = query.where(filter=FieldFilter("memory", "==", search_content))
836
+
837
+ # Apply sorting
838
+ query = apply_sorting(query, sort_by, sort_order)
839
+
840
+ # Get all documents
841
+ docs = query.stream()
842
+ all_records = [doc.to_dict() for doc in docs]
843
+
844
+ total_count = len(all_records)
845
+
846
+ # Apply pagination to the filtered results
847
+ if limit is not None and page is not None:
848
+ start_index = (page - 1) * limit
849
+ end_index = start_index + limit
850
+ records = all_records[start_index:end_index]
851
+ elif limit is not None:
852
+ records = all_records[:limit]
853
+ else:
854
+ records = all_records
855
+ if not deserialize:
856
+ return records, total_count
857
+
858
+ return [UserMemory.from_dict(record) for record in records]
859
+
860
+ except Exception as e:
861
+ log_error(f"Exception getting user memories: {e}")
862
+ raise e
863
+
864
+ def get_user_memory_stats(
865
+ self,
866
+ limit: Optional[int] = None,
867
+ page: Optional[int] = None,
868
+ ) -> Tuple[List[Dict[str, Any]], int]:
869
+ """Get user memories stats.
870
+
871
+ Args:
872
+ limit (Optional[int]): The limit of the memories to get.
873
+ page (Optional[int]): The page number to get.
874
+
875
+ Returns:
876
+ Tuple[List[Dict[str, Any]], int]: A tuple containing the memories stats and the total count.
877
+
878
+ Raises:
879
+ Exception: If there is an error getting the memories stats.
880
+ """
881
+ try:
882
+ collection_ref = self._get_collection(table_type="memories")
883
+
884
+ query = collection_ref.where(filter=FieldFilter("user_id", "!=", None))
885
+
886
+ docs = query.stream()
887
+
888
+ user_stats = {}
889
+ for doc in docs:
890
+ data = doc.to_dict()
891
+ current_user_id = data.get("user_id")
892
+ if current_user_id:
893
+ if current_user_id not in user_stats:
894
+ user_stats[current_user_id] = {
895
+ "user_id": current_user_id,
896
+ "total_memories": 0,
897
+ "last_memory_updated_at": 0,
898
+ }
899
+ user_stats[current_user_id]["total_memories"] += 1
900
+ updated_at = data.get("updated_at", 0)
901
+ if updated_at > user_stats[current_user_id]["last_memory_updated_at"]:
902
+ user_stats[current_user_id]["last_memory_updated_at"] = updated_at
903
+
904
+ # Convert to list and sort
905
+ formatted_results = list(user_stats.values())
906
+ formatted_results.sort(key=lambda x: x["last_memory_updated_at"], reverse=True)
907
+
908
+ total_count = len(formatted_results)
909
+
910
+ # Apply pagination
911
+ if limit is not None:
912
+ start_idx = 0
913
+ if page is not None:
914
+ start_idx = (page - 1) * limit
915
+ formatted_results = formatted_results[start_idx : start_idx + limit]
916
+
917
+ return formatted_results, total_count
918
+
919
+ except Exception as e:
920
+ log_error(f"Exception getting user memory stats: {e}")
921
+ raise e
922
+
923
+ def upsert_user_memory(
924
+ self, memory: UserMemory, deserialize: Optional[bool] = True
925
+ ) -> Optional[Union[UserMemory, Dict[str, Any]]]:
926
+ """Upsert a user memory in the database.
927
+
928
+ Args:
929
+ memory (UserMemory): The memory to upsert.
930
+ deserialize (Optional[bool]): Whether to serialize the memory. Defaults to True.
931
+
932
+ Returns:
933
+ Optional[Union[UserMemory, Dict[str, Any]]]:
934
+ - When deserialize=True: UserMemory object
935
+ - When deserialize=False: Memory dictionary
936
+
937
+ Raises:
938
+ Exception: If there is an error upserting the memory.
939
+ """
940
+ try:
941
+ collection_ref = self._get_collection(table_type="memories", create_collection_if_not_found=True)
942
+ if collection_ref is None:
943
+ return None
944
+
945
+ if memory.memory_id is None:
946
+ memory.memory_id = str(uuid4())
947
+
948
+ update_doc = memory.to_dict()
949
+ update_doc["updated_at"] = int(time.time())
950
+
951
+ # Find existing document or create new one
952
+ docs = collection_ref.where("memory_id", "==", memory.memory_id).stream()
953
+ doc_ref = next((doc.reference for doc in docs), None)
954
+
955
+ if doc_ref is None:
956
+ doc_ref = collection_ref.document()
957
+
958
+ doc_ref.set(update_doc, merge=True)
959
+
960
+ if not deserialize:
961
+ return update_doc
962
+
963
+ return UserMemory.from_dict(update_doc)
964
+
965
+ except Exception as e:
966
+ log_error(f"Exception upserting user memory: {e}")
967
+ raise e
968
+
969
+ def upsert_memories(
970
+ self, memories: List[UserMemory], deserialize: Optional[bool] = True, preserve_updated_at: bool = False
971
+ ) -> List[Union[UserMemory, Dict[str, Any]]]:
972
+ """
973
+ Bulk upsert multiple user memories for improved performance on large datasets.
974
+
975
+ Args:
976
+ memories (List[UserMemory]): List of memories to upsert.
977
+ deserialize (Optional[bool]): Whether to deserialize the memories. Defaults to True.
978
+
979
+ Returns:
980
+ List[Union[UserMemory, Dict[str, Any]]]: List of upserted memories.
981
+
982
+ Raises:
983
+ Exception: If an error occurs during bulk upsert.
984
+ """
985
+ if not memories:
986
+ return []
987
+
988
+ try:
989
+ log_info(
990
+ f"FirestoreDb doesn't support efficient bulk operations, falling back to individual upserts for {len(memories)} memories"
991
+ )
992
+ # Fall back to individual upserts
993
+ results = []
994
+ for memory in memories:
995
+ if memory is not None:
996
+ result = self.upsert_user_memory(memory, deserialize=deserialize)
997
+ if result is not None:
998
+ results.append(result)
999
+ return results
1000
+
1001
+ except Exception as e:
1002
+ log_error(f"Exception during bulk memory upsert: {e}")
1003
+ return []
1004
+
1005
+ def clear_memories(self) -> None:
1006
+ """Delete all memories from the database.
1007
+
1008
+ Raises:
1009
+ Exception: If an error occurs during deletion.
1010
+ """
1011
+ try:
1012
+ collection_ref = self._get_collection(table_type="memories")
1013
+
1014
+ # Get all documents in the collection
1015
+ docs = collection_ref.stream()
1016
+
1017
+ # Delete all documents in batches
1018
+ batch = self.db_client.batch()
1019
+ batch_count = 0
1020
+
1021
+ for doc in docs:
1022
+ batch.delete(doc.reference)
1023
+ batch_count += 1
1024
+
1025
+ # Firestore batch has a limit of 500 operations
1026
+ if batch_count >= 500:
1027
+ batch.commit()
1028
+ batch = self.db_client.batch()
1029
+ batch_count = 0
1030
+
1031
+ # Commit remaining operations
1032
+ if batch_count > 0:
1033
+ batch.commit()
1034
+
1035
+ except Exception as e:
1036
+ log_error(f"Exception deleting all memories: {e}")
1037
+ raise e
1038
+
1039
+ # -- Cultural Knowledge methods --
1040
+ def clear_cultural_knowledge(self) -> None:
1041
+ """Delete all cultural knowledge from the database.
1042
+
1043
+ Raises:
1044
+ Exception: If an error occurs during deletion.
1045
+ """
1046
+ try:
1047
+ collection_ref = self._get_collection(table_type="culture")
1048
+
1049
+ # Get all documents in the collection
1050
+ docs = collection_ref.stream()
1051
+
1052
+ # Delete all documents in batches
1053
+ batch = self.db_client.batch()
1054
+ batch_count = 0
1055
+
1056
+ for doc in docs:
1057
+ batch.delete(doc.reference)
1058
+ batch_count += 1
1059
+
1060
+ # Firestore batch has a limit of 500 operations
1061
+ if batch_count >= 500:
1062
+ batch.commit()
1063
+ batch = self.db_client.batch()
1064
+ batch_count = 0
1065
+
1066
+ # Commit remaining operations
1067
+ if batch_count > 0:
1068
+ batch.commit()
1069
+
1070
+ except Exception as e:
1071
+ log_error(f"Exception deleting all cultural knowledge: {e}")
1072
+ raise e
1073
+
1074
+ def delete_cultural_knowledge(self, id: str) -> None:
1075
+ """Delete cultural knowledge by ID.
1076
+
1077
+ Args:
1078
+ id (str): The ID of the cultural knowledge to delete.
1079
+
1080
+ Raises:
1081
+ Exception: If an error occurs during deletion.
1082
+ """
1083
+ try:
1084
+ collection_ref = self._get_collection(table_type="culture")
1085
+ docs = collection_ref.where(filter=FieldFilter("id", "==", id)).stream()
1086
+
1087
+ for doc in docs:
1088
+ doc.reference.delete()
1089
+ log_debug(f"Deleted cultural knowledge with ID: {id}")
1090
+
1091
+ except Exception as e:
1092
+ log_error(f"Error deleting cultural knowledge: {e}")
1093
+ raise e
1094
+
1095
+ def get_cultural_knowledge(
1096
+ self, id: str, deserialize: Optional[bool] = True
1097
+ ) -> Optional[Union[CulturalKnowledge, Dict[str, Any]]]:
1098
+ """Get cultural knowledge by ID.
1099
+
1100
+ Args:
1101
+ id (str): The ID of the cultural knowledge to retrieve.
1102
+ deserialize (Optional[bool]): Whether to deserialize to CulturalKnowledge object. Defaults to True.
1103
+
1104
+ Returns:
1105
+ Optional[Union[CulturalKnowledge, Dict[str, Any]]]: The cultural knowledge if found, None otherwise.
1106
+
1107
+ Raises:
1108
+ Exception: If an error occurs during retrieval.
1109
+ """
1110
+ try:
1111
+ collection_ref = self._get_collection(table_type="culture")
1112
+ docs = collection_ref.where(filter=FieldFilter("id", "==", id)).limit(1).stream()
1113
+
1114
+ for doc in docs:
1115
+ result = doc.to_dict()
1116
+ if not deserialize:
1117
+ return result
1118
+ return deserialize_cultural_knowledge_from_db(result)
1119
+
1120
+ return None
1121
+
1122
+ except Exception as e:
1123
+ log_error(f"Error getting cultural knowledge: {e}")
1124
+ raise e
1125
+
1126
+ def get_all_cultural_knowledge(
1127
+ self,
1128
+ agent_id: Optional[str] = None,
1129
+ team_id: Optional[str] = None,
1130
+ name: Optional[str] = None,
1131
+ limit: Optional[int] = None,
1132
+ page: Optional[int] = None,
1133
+ sort_by: Optional[str] = None,
1134
+ sort_order: Optional[str] = None,
1135
+ deserialize: Optional[bool] = True,
1136
+ ) -> Union[List[CulturalKnowledge], Tuple[List[Dict[str, Any]], int]]:
1137
+ """Get all cultural knowledge with filtering and pagination.
1138
+
1139
+ Args:
1140
+ agent_id (Optional[str]): Filter by agent ID.
1141
+ team_id (Optional[str]): Filter by team ID.
1142
+ name (Optional[str]): Filter by name (case-insensitive partial match).
1143
+ limit (Optional[int]): Maximum number of results to return.
1144
+ page (Optional[int]): Page number for pagination.
1145
+ sort_by (Optional[str]): Field to sort by.
1146
+ sort_order (Optional[str]): Sort order ('asc' or 'desc').
1147
+ deserialize (Optional[bool]): Whether to deserialize to CulturalKnowledge objects. Defaults to True.
1148
+
1149
+ Returns:
1150
+ Union[List[CulturalKnowledge], Tuple[List[Dict[str, Any]], int]]:
1151
+ - When deserialize=True: List of CulturalKnowledge objects
1152
+ - When deserialize=False: Tuple with list of dictionaries and total count
1153
+
1154
+ Raises:
1155
+ Exception: If an error occurs during retrieval.
1156
+ """
1157
+ try:
1158
+ collection_ref = self._get_collection(table_type="culture")
1159
+
1160
+ # Build query with filters
1161
+ query = collection_ref
1162
+ if agent_id is not None:
1163
+ query = query.where(filter=FieldFilter("agent_id", "==", agent_id))
1164
+ if team_id is not None:
1165
+ query = query.where(filter=FieldFilter("team_id", "==", team_id))
1166
+
1167
+ # Get all matching documents
1168
+ docs = query.stream()
1169
+ results = [doc.to_dict() for doc in docs]
1170
+
1171
+ # Apply name filter (Firestore doesn't support regex in queries)
1172
+ if name is not None:
1173
+ results = [r for r in results if name.lower() in r.get("name", "").lower()]
1174
+
1175
+ total_count = len(results)
1176
+
1177
+ # Apply sorting and pagination to in-memory results
1178
+ sorted_results = apply_sorting_to_records(records=results, sort_by=sort_by, sort_order=sort_order)
1179
+ paginated_results = apply_pagination_to_records(records=sorted_results, limit=limit, page=page)
1180
+
1181
+ if not deserialize:
1182
+ return paginated_results, total_count
1183
+
1184
+ return [deserialize_cultural_knowledge_from_db(item) for item in paginated_results]
1185
+
1186
+ except Exception as e:
1187
+ log_error(f"Error getting all cultural knowledge: {e}")
1188
+ raise e
1189
+
1190
+ def upsert_cultural_knowledge(
1191
+ self, cultural_knowledge: CulturalKnowledge, deserialize: Optional[bool] = True
1192
+ ) -> Optional[Union[CulturalKnowledge, Dict[str, Any]]]:
1193
+ """Upsert cultural knowledge in Firestore.
1194
+
1195
+ Args:
1196
+ cultural_knowledge (CulturalKnowledge): The cultural knowledge to upsert.
1197
+ deserialize (Optional[bool]): Whether to deserialize the result. Defaults to True.
1198
+
1199
+ Returns:
1200
+ Optional[Union[CulturalKnowledge, Dict[str, Any]]]: The upserted cultural knowledge.
1201
+
1202
+ Raises:
1203
+ Exception: If an error occurs during upsert.
1204
+ """
1205
+ try:
1206
+ collection_ref = self._get_collection(table_type="culture", create_collection_if_not_found=True)
1207
+
1208
+ # Serialize content, categories, and notes into a dict for DB storage
1209
+ content_dict = serialize_cultural_knowledge_for_db(cultural_knowledge)
1210
+
1211
+ # Create the update document with serialized content
1212
+ update_doc = {
1213
+ "id": cultural_knowledge.id,
1214
+ "name": cultural_knowledge.name,
1215
+ "summary": cultural_knowledge.summary,
1216
+ "content": content_dict if content_dict else None,
1217
+ "metadata": cultural_knowledge.metadata,
1218
+ "input": cultural_knowledge.input,
1219
+ "created_at": cultural_knowledge.created_at,
1220
+ "updated_at": int(time.time()),
1221
+ "agent_id": cultural_knowledge.agent_id,
1222
+ "team_id": cultural_knowledge.team_id,
1223
+ }
1224
+
1225
+ # Find and update or create new document
1226
+ docs = collection_ref.where(filter=FieldFilter("id", "==", cultural_knowledge.id)).limit(1).stream()
1227
+
1228
+ doc_found = False
1229
+ for doc in docs:
1230
+ doc.reference.set(update_doc)
1231
+ doc_found = True
1232
+ break
1233
+
1234
+ if not doc_found:
1235
+ collection_ref.add(update_doc)
1236
+
1237
+ if not deserialize:
1238
+ return update_doc
1239
+
1240
+ return deserialize_cultural_knowledge_from_db(update_doc)
1241
+
1242
+ except Exception as e:
1243
+ log_error(f"Error upserting cultural knowledge: {e}")
1244
+ raise e
1245
+
1246
+ # -- Metrics methods --
1247
+
1248
+ def _get_all_sessions_for_metrics_calculation(
1249
+ self, start_timestamp: Optional[int] = None, end_timestamp: Optional[int] = None
1250
+ ) -> List[Dict[str, Any]]:
1251
+ """Get all sessions of all types for metrics calculation."""
1252
+ try:
1253
+ collection_ref = self._get_collection(table_type="sessions")
1254
+
1255
+ query = collection_ref
1256
+ if start_timestamp is not None:
1257
+ query = query.where(filter=FieldFilter("created_at", ">=", start_timestamp))
1258
+ if end_timestamp is not None:
1259
+ query = query.where(filter=FieldFilter("created_at", "<=", end_timestamp))
1260
+
1261
+ docs = query.stream()
1262
+ results = []
1263
+ for doc in docs:
1264
+ data = doc.to_dict()
1265
+ # Only include required fields for metrics
1266
+ result = {
1267
+ "user_id": data.get("user_id"),
1268
+ "session_data": data.get("session_data"),
1269
+ "runs": data.get("runs"),
1270
+ "created_at": data.get("created_at"),
1271
+ "session_type": data.get("session_type"),
1272
+ }
1273
+ results.append(result)
1274
+
1275
+ return results
1276
+
1277
+ except Exception as e:
1278
+ log_error(f"Exception getting all sessions for metrics calculation: {e}")
1279
+ raise e
1280
+
1281
+ def _get_metrics_calculation_starting_date(self, collection_ref) -> Optional[date]:
1282
+ """Get the first date for which metrics calculation is needed."""
1283
+ try:
1284
+ query = collection_ref.order_by("date", direction="DESCENDING").limit(1)
1285
+ docs = query.stream()
1286
+
1287
+ for doc in docs:
1288
+ data = doc.to_dict()
1289
+ result_date = datetime.strptime(data["date"], "%Y-%m-%d").date()
1290
+ if data.get("completed"):
1291
+ return result_date + timedelta(days=1)
1292
+ else:
1293
+ return result_date
1294
+
1295
+ # No metrics records. Return the date of the first recorded session.
1296
+ first_session_result = self.get_sessions(sort_by="created_at", sort_order="asc", limit=1, deserialize=False)
1297
+ first_session_date = None
1298
+
1299
+ if isinstance(first_session_result, list) and len(first_session_result) > 0:
1300
+ first_session_date = first_session_result[0].created_at # type: ignore
1301
+ elif isinstance(first_session_result, tuple) and len(first_session_result[0]) > 0:
1302
+ first_session_date = first_session_result[0][0].get("created_at")
1303
+
1304
+ if first_session_date is None:
1305
+ return None
1306
+
1307
+ return datetime.fromtimestamp(first_session_date, tz=timezone.utc).date()
1308
+
1309
+ except Exception as e:
1310
+ log_error(f"Exception getting metrics calculation starting date: {e}")
1311
+ raise e
1312
+
1313
+ def calculate_metrics(self) -> Optional[list[dict]]:
1314
+ """Calculate metrics for all dates without complete metrics."""
1315
+ try:
1316
+ collection_ref = self._get_collection(table_type="metrics", create_collection_if_not_found=True)
1317
+
1318
+ starting_date = self._get_metrics_calculation_starting_date(collection_ref)
1319
+ if starting_date is None:
1320
+ log_info("No session data found. Won't calculate metrics.")
1321
+ return None
1322
+
1323
+ dates_to_process = get_dates_to_calculate_metrics_for(starting_date)
1324
+ if not dates_to_process:
1325
+ log_info("Metrics already calculated for all relevant dates.")
1326
+ return None
1327
+
1328
+ start_timestamp = int(datetime.combine(dates_to_process[0], datetime.min.time()).timestamp())
1329
+ end_timestamp = int(
1330
+ datetime.combine(dates_to_process[-1] + timedelta(days=1), datetime.min.time()).timestamp()
1331
+ )
1332
+
1333
+ sessions = self._get_all_sessions_for_metrics_calculation(
1334
+ start_timestamp=start_timestamp, end_timestamp=end_timestamp
1335
+ )
1336
+ all_sessions_data = fetch_all_sessions_data(
1337
+ sessions=sessions, dates_to_process=dates_to_process, start_timestamp=start_timestamp
1338
+ )
1339
+ if not all_sessions_data:
1340
+ log_info("No new session data found. Won't calculate metrics.")
1341
+ return None
1342
+
1343
+ results = []
1344
+ metrics_records = []
1345
+
1346
+ for date_to_process in dates_to_process:
1347
+ date_key = date_to_process.isoformat()
1348
+ sessions_for_date = all_sessions_data.get(date_key, {})
1349
+
1350
+ # Skip dates with no sessions
1351
+ if not any(len(sessions) > 0 for sessions in sessions_for_date.values()):
1352
+ continue
1353
+
1354
+ metrics_record = calculate_date_metrics(date_to_process, sessions_for_date)
1355
+ metrics_records.append(metrics_record)
1356
+
1357
+ if metrics_records:
1358
+ results = bulk_upsert_metrics(collection_ref, metrics_records)
1359
+
1360
+ log_debug("Updated metrics calculations")
1361
+
1362
+ return results
1363
+
1364
+ except Exception as e:
1365
+ log_error(f"Exception calculating metrics: {e}")
1366
+ raise e
1367
+
1368
+ def get_metrics(
1369
+ self,
1370
+ starting_date: Optional[date] = None,
1371
+ ending_date: Optional[date] = None,
1372
+ ) -> Tuple[List[dict], Optional[int]]:
1373
+ """Get all metrics matching the given date range."""
1374
+ try:
1375
+ collection_ref = self._get_collection(table_type="metrics")
1376
+ if collection_ref is None:
1377
+ return [], None
1378
+
1379
+ query = collection_ref
1380
+ if starting_date:
1381
+ query = query.where(filter=FieldFilter("date", ">=", starting_date.isoformat()))
1382
+ if ending_date:
1383
+ query = query.where(filter=FieldFilter("date", "<=", ending_date.isoformat()))
1384
+
1385
+ docs = query.stream()
1386
+ records = []
1387
+ latest_updated_at = 0
1388
+
1389
+ for doc in docs:
1390
+ data = doc.to_dict()
1391
+ records.append(data)
1392
+ updated_at = data.get("updated_at", 0)
1393
+ if updated_at > latest_updated_at:
1394
+ latest_updated_at = updated_at
1395
+
1396
+ if not records:
1397
+ return [], None
1398
+
1399
+ return records, latest_updated_at
1400
+
1401
+ except Exception as e:
1402
+ log_error(f"Exception getting metrics: {e}")
1403
+ raise e
1404
+
1405
+ # -- Knowledge methods --
1406
+
1407
+ def delete_knowledge_content(self, id: str):
1408
+ """Delete a knowledge row from the database.
1409
+
1410
+ Args:
1411
+ id (str): The ID of the knowledge row to delete.
1412
+
1413
+ Raises:
1414
+ Exception: If an error occurs during deletion.
1415
+ """
1416
+ try:
1417
+ collection_ref = self._get_collection(table_type="knowledge")
1418
+ docs = collection_ref.where(filter=FieldFilter("id", "==", id)).stream()
1419
+
1420
+ for doc in docs:
1421
+ doc.reference.delete()
1422
+
1423
+ except Exception as e:
1424
+ log_error(f"Error deleting knowledge content: {e}")
1425
+ raise e
1426
+
1427
+ def get_knowledge_content(self, id: str) -> Optional[KnowledgeRow]:
1428
+ """Get a knowledge row from the database.
1429
+
1430
+ Args:
1431
+ id (str): The ID of the knowledge row to get.
1432
+
1433
+ Returns:
1434
+ Optional[KnowledgeRow]: The knowledge row, or None if it doesn't exist.
1435
+
1436
+ Raises:
1437
+ Exception: If an error occurs during retrieval.
1438
+ """
1439
+ try:
1440
+ collection_ref = self._get_collection(table_type="knowledge")
1441
+ docs = collection_ref.where(filter=FieldFilter("id", "==", id)).stream()
1442
+
1443
+ for doc in docs:
1444
+ data = doc.to_dict()
1445
+ return KnowledgeRow.model_validate(data)
1446
+
1447
+ return None
1448
+
1449
+ except Exception as e:
1450
+ log_error(f"Error getting knowledge content: {e}")
1451
+ raise e
1452
+
1453
+ def get_knowledge_contents(
1454
+ self,
1455
+ limit: Optional[int] = None,
1456
+ page: Optional[int] = None,
1457
+ sort_by: Optional[str] = None,
1458
+ sort_order: Optional[str] = None,
1459
+ ) -> Tuple[List[KnowledgeRow], int]:
1460
+ """Get all knowledge contents from the database.
1461
+
1462
+ Args:
1463
+ limit (Optional[int]): The maximum number of knowledge contents to return.
1464
+ page (Optional[int]): The page number.
1465
+ sort_by (Optional[str]): The column to sort by.
1466
+ sort_order (Optional[str]): The order to sort by.
1467
+ create_table_if_not_found (Optional[bool]): Whether to create the table if it doesn't exist.
1468
+
1469
+ Returns:
1470
+ Tuple[List[KnowledgeRow], int]: The knowledge contents and total count.
1471
+
1472
+ Raises:
1473
+ Exception: If an error occurs during retrieval.
1474
+ """
1475
+ try:
1476
+ collection_ref = self._get_collection(table_type="knowledge")
1477
+ if collection_ref is None:
1478
+ return [], 0
1479
+
1480
+ query = collection_ref
1481
+
1482
+ # Apply sorting
1483
+ query = apply_sorting(query, sort_by, sort_order)
1484
+
1485
+ # Apply pagination
1486
+ query = apply_pagination(query, limit, page)
1487
+
1488
+ docs = query.stream()
1489
+ records = []
1490
+ for doc in docs:
1491
+ records.append(doc.to_dict())
1492
+
1493
+ knowledge_rows = [KnowledgeRow.model_validate(record) for record in records]
1494
+ total_count = len(knowledge_rows) # Simplified count
1495
+
1496
+ return knowledge_rows, total_count
1497
+
1498
+ except Exception as e:
1499
+ log_error(f"Error getting knowledge contents: {e}")
1500
+ raise e
1501
+
1502
+ def upsert_knowledge_content(self, knowledge_row: KnowledgeRow):
1503
+ """Upsert knowledge content in the database.
1504
+
1505
+ Args:
1506
+ knowledge_row (KnowledgeRow): The knowledge row to upsert.
1507
+
1508
+ Returns:
1509
+ Optional[KnowledgeRow]: The upserted knowledge row, or None if the operation fails.
1510
+ """
1511
+ try:
1512
+ collection_ref = self._get_collection(table_type="knowledge", create_collection_if_not_found=True)
1513
+ if collection_ref is None:
1514
+ return None
1515
+
1516
+ update_doc = knowledge_row.model_dump()
1517
+
1518
+ # Find existing document or create new one
1519
+ docs = collection_ref.where(filter=FieldFilter("id", "==", knowledge_row.id)).stream()
1520
+ doc_ref = next((doc.reference for doc in docs), None)
1521
+
1522
+ if doc_ref is None:
1523
+ doc_ref = collection_ref.document()
1524
+
1525
+ doc_ref.set(update_doc, merge=True)
1526
+
1527
+ return knowledge_row
1528
+
1529
+ except Exception as e:
1530
+ log_error(f"Error upserting knowledge content: {e}")
1531
+ raise e
1532
+
1533
+ # -- Eval methods --
1534
+
1535
+ def create_eval_run(self, eval_run: EvalRunRecord) -> Optional[EvalRunRecord]:
1536
+ """Create an EvalRunRecord in the database."""
1537
+ try:
1538
+ collection_ref = self._get_collection(table_type="evals", create_collection_if_not_found=True)
1539
+
1540
+ current_time = int(time.time())
1541
+ eval_dict = eval_run.model_dump()
1542
+ eval_dict["created_at"] = current_time
1543
+ eval_dict["updated_at"] = current_time
1544
+
1545
+ doc_ref = collection_ref.document()
1546
+ doc_ref.set(eval_dict)
1547
+
1548
+ log_debug(f"Created eval run with id '{eval_run.run_id}'")
1549
+
1550
+ return eval_run
1551
+
1552
+ except Exception as e:
1553
+ log_error(f"Error creating eval run: {e}")
1554
+ raise e
1555
+
1556
+ def delete_eval_run(self, eval_run_id: str) -> None:
1557
+ """Delete an eval run from the database."""
1558
+ try:
1559
+ collection_ref = self._get_collection(table_type="evals")
1560
+ docs = collection_ref.where(filter=FieldFilter("run_id", "==", eval_run_id)).stream()
1561
+
1562
+ deleted_count = 0
1563
+ for doc in docs:
1564
+ doc.reference.delete()
1565
+ deleted_count += 1
1566
+
1567
+ if deleted_count == 0:
1568
+ log_info(f"No eval run found with ID: {eval_run_id}")
1569
+ else:
1570
+ log_info(f"Deleted eval run with ID: {eval_run_id}")
1571
+
1572
+ except Exception as e:
1573
+ log_error(f"Error deleting eval run {eval_run_id}: {e}")
1574
+ raise e
1575
+
1576
+ def delete_eval_runs(self, eval_run_ids: List[str]) -> None:
1577
+ """Delete multiple eval runs from the database.
1578
+
1579
+ Args:
1580
+ eval_run_ids (List[str]): The IDs of the eval runs to delete.
1581
+
1582
+ Raises:
1583
+ Exception: If there is an error deleting the eval runs.
1584
+ """
1585
+ try:
1586
+ collection_ref = self._get_collection(table_type="evals")
1587
+ batch = self.db_client.batch()
1588
+ deleted_count = 0
1589
+
1590
+ for eval_run_id in eval_run_ids:
1591
+ docs = collection_ref.where(filter=FieldFilter("run_id", "==", eval_run_id)).stream()
1592
+ for doc in docs:
1593
+ batch.delete(doc.reference)
1594
+ deleted_count += 1
1595
+
1596
+ batch.commit()
1597
+
1598
+ if deleted_count == 0:
1599
+ log_info(f"No eval runs found with IDs: {eval_run_ids}")
1600
+ else:
1601
+ log_info(f"Deleted {deleted_count} eval runs")
1602
+
1603
+ except Exception as e:
1604
+ log_error(f"Error deleting eval runs {eval_run_ids}: {e}")
1605
+ raise e
1606
+
1607
+ def get_eval_run(
1608
+ self, eval_run_id: str, deserialize: Optional[bool] = True
1609
+ ) -> Optional[Union[EvalRunRecord, Dict[str, Any]]]:
1610
+ """Get an eval run from the database.
1611
+
1612
+ Args:
1613
+ eval_run_id (str): The ID of the eval run to get.
1614
+ deserialize (Optional[bool]): Whether to serialize the eval run. Defaults to True.
1615
+
1616
+ Returns:
1617
+ Optional[Union[EvalRunRecord, Dict[str, Any]]]:
1618
+ - When deserialize=True: EvalRunRecord object
1619
+ - When deserialize=False: EvalRun dictionary
1620
+
1621
+ Raises:
1622
+ Exception: If there is an error getting the eval run.
1623
+ """
1624
+ try:
1625
+ collection_ref = self._get_collection(table_type="evals")
1626
+ if not collection_ref:
1627
+ return None
1628
+
1629
+ docs = collection_ref.where(filter=FieldFilter("run_id", "==", eval_run_id)).stream()
1630
+
1631
+ eval_run_raw = None
1632
+ for doc in docs:
1633
+ eval_run_raw = doc.to_dict()
1634
+ break
1635
+
1636
+ if not eval_run_raw:
1637
+ return None
1638
+
1639
+ if not deserialize:
1640
+ return eval_run_raw
1641
+
1642
+ return EvalRunRecord.model_validate(eval_run_raw)
1643
+
1644
+ except Exception as e:
1645
+ log_error(f"Exception getting eval run {eval_run_id}: {e}")
1646
+ raise e
1647
+
1648
+ def get_eval_runs(
1649
+ self,
1650
+ limit: Optional[int] = None,
1651
+ page: Optional[int] = None,
1652
+ sort_by: Optional[str] = None,
1653
+ sort_order: Optional[str] = None,
1654
+ agent_id: Optional[str] = None,
1655
+ team_id: Optional[str] = None,
1656
+ workflow_id: Optional[str] = None,
1657
+ model_id: Optional[str] = None,
1658
+ filter_type: Optional[EvalFilterType] = None,
1659
+ eval_type: Optional[List[EvalType]] = None,
1660
+ deserialize: Optional[bool] = True,
1661
+ ) -> Union[List[EvalRunRecord], Tuple[List[Dict[str, Any]], int]]:
1662
+ """Get all eval runs from the database.
1663
+
1664
+ Args:
1665
+ limit (Optional[int]): The maximum number of eval runs to return.
1666
+ page (Optional[int]): The page number to return.
1667
+ sort_by (Optional[str]): The field to sort by.
1668
+ sort_order (Optional[str]): The order to sort by.
1669
+ agent_id (Optional[str]): The ID of the agent to filter by.
1670
+ team_id (Optional[str]): The ID of the team to filter by.
1671
+ workflow_id (Optional[str]): The ID of the workflow to filter by.
1672
+ model_id (Optional[str]): The ID of the model to filter by.
1673
+ eval_type (Optional[List[EvalType]]): The type of eval to filter by.
1674
+ filter_type (Optional[EvalFilterType]): The type of filter to apply.
1675
+ deserialize (Optional[bool]): Whether to serialize the eval runs. Defaults to True.
1676
+ create_table_if_not_found (Optional[bool]): Whether to create the table if it doesn't exist.
1677
+
1678
+ Returns:
1679
+ Union[List[EvalRunRecord], Tuple[List[Dict[str, Any]], int]]:
1680
+ - When deserialize=True: List of EvalRunRecord objects
1681
+ - When deserialize=False: List of eval run dictionaries and the total count
1682
+
1683
+ Raises:
1684
+ Exception: If there is an error getting the eval runs.
1685
+ """
1686
+ try:
1687
+ collection_ref = self._get_collection(table_type="evals")
1688
+ if collection_ref is None:
1689
+ return [] if deserialize else ([], 0)
1690
+
1691
+ query = collection_ref
1692
+
1693
+ if agent_id is not None:
1694
+ query = query.where(filter=FieldFilter("agent_id", "==", agent_id))
1695
+ if team_id is not None:
1696
+ query = query.where(filter=FieldFilter("team_id", "==", team_id))
1697
+ if workflow_id is not None:
1698
+ query = query.where(filter=FieldFilter("workflow_id", "==", workflow_id))
1699
+ if model_id is not None:
1700
+ query = query.where(filter=FieldFilter("model_id", "==", model_id))
1701
+ if eval_type is not None and len(eval_type) > 0:
1702
+ eval_values = [et.value for et in eval_type]
1703
+ query = query.where(filter=FieldFilter("eval_type", "in", eval_values))
1704
+ if filter_type is not None:
1705
+ if filter_type == EvalFilterType.AGENT:
1706
+ query = query.where(filter=FieldFilter("agent_id", "!=", None))
1707
+ elif filter_type == EvalFilterType.TEAM:
1708
+ query = query.where(filter=FieldFilter("team_id", "!=", None))
1709
+ elif filter_type == EvalFilterType.WORKFLOW:
1710
+ query = query.where(filter=FieldFilter("workflow_id", "!=", None))
1711
+
1712
+ # Apply default sorting by created_at desc if no sort parameters provided
1713
+ if sort_by is None:
1714
+ from google.cloud.firestore import Query
1715
+
1716
+ query = query.order_by("created_at", direction=Query.DESCENDING)
1717
+ else:
1718
+ query = apply_sorting(query, sort_by, sort_order)
1719
+
1720
+ # Get all documents for counting before pagination
1721
+ all_docs = query.stream()
1722
+ all_records = [doc.to_dict() for doc in all_docs]
1723
+
1724
+ if not all_records:
1725
+ return [] if deserialize else ([], 0)
1726
+
1727
+ # Get total count before pagination
1728
+ total_count = len(all_records)
1729
+
1730
+ # Apply pagination to the results
1731
+ if limit is not None and page is not None:
1732
+ start_index = (page - 1) * limit
1733
+ end_index = start_index + limit
1734
+ records = all_records[start_index:end_index]
1735
+ elif limit is not None:
1736
+ records = all_records[:limit]
1737
+ else:
1738
+ records = all_records
1739
+
1740
+ if not deserialize:
1741
+ return records, total_count
1742
+
1743
+ return [EvalRunRecord.model_validate(row) for row in records]
1744
+
1745
+ except Exception as e:
1746
+ log_error(f"Exception getting eval runs: {e}")
1747
+ raise e
1748
+
1749
+ def rename_eval_run(
1750
+ self, eval_run_id: str, name: str, deserialize: Optional[bool] = True
1751
+ ) -> Optional[Union[EvalRunRecord, Dict[str, Any]]]:
1752
+ """Update the name of an eval run in the database.
1753
+
1754
+ Args:
1755
+ eval_run_id (str): The ID of the eval run to update.
1756
+ name (str): The new name of the eval run.
1757
+ deserialize (Optional[bool]): Whether to serialize the eval run. Defaults to True.
1758
+
1759
+ Returns:
1760
+ Optional[Union[EvalRunRecord, Dict[str, Any]]]:
1761
+ - When deserialize=True: EvalRunRecord object
1762
+ - When deserialize=False: EvalRun dictionary
1763
+
1764
+ Raises:
1765
+ Exception: If there is an error updating the eval run.
1766
+ """
1767
+ try:
1768
+ collection_ref = self._get_collection(table_type="evals")
1769
+ if not collection_ref:
1770
+ return None
1771
+
1772
+ docs = collection_ref.where(filter=FieldFilter("run_id", "==", eval_run_id)).stream()
1773
+ doc_ref = next((doc.reference for doc in docs), None)
1774
+
1775
+ if doc_ref is None:
1776
+ return None
1777
+
1778
+ doc_ref.update({"name": name, "updated_at": int(time.time())})
1779
+
1780
+ updated_doc = doc_ref.get()
1781
+ if not updated_doc.exists:
1782
+ return None
1783
+
1784
+ result = updated_doc.to_dict()
1785
+
1786
+ log_debug(f"Renamed eval run with id '{eval_run_id}' to '{name}'")
1787
+
1788
+ if not result or not deserialize:
1789
+ return result
1790
+
1791
+ return EvalRunRecord.model_validate(result)
1792
+
1793
+ except Exception as e:
1794
+ log_error(f"Error updating eval run name {eval_run_id}: {e}")
1795
+ raise e