agno 1.8.1__py3-none-any.whl → 2.0.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (590) hide show
  1. agno/__init__.py +8 -0
  2. agno/agent/__init__.py +19 -27
  3. agno/agent/agent.py +3143 -4170
  4. agno/api/agent.py +11 -67
  5. agno/api/api.py +5 -46
  6. agno/api/evals.py +8 -19
  7. agno/api/os.py +17 -0
  8. agno/api/routes.py +6 -41
  9. agno/api/schemas/__init__.py +9 -0
  10. agno/api/schemas/agent.py +5 -21
  11. agno/api/schemas/evals.py +7 -16
  12. agno/api/schemas/os.py +14 -0
  13. agno/api/schemas/team.py +5 -21
  14. agno/api/schemas/utils.py +21 -0
  15. agno/api/schemas/workflows.py +11 -7
  16. agno/api/settings.py +53 -0
  17. agno/api/team.py +11 -66
  18. agno/api/workflow.py +28 -0
  19. agno/cloud/aws/base.py +214 -0
  20. agno/cloud/aws/s3/__init__.py +2 -0
  21. agno/cloud/aws/s3/api_client.py +43 -0
  22. agno/cloud/aws/s3/bucket.py +195 -0
  23. agno/cloud/aws/s3/object.py +57 -0
  24. agno/db/__init__.py +24 -0
  25. agno/db/base.py +245 -0
  26. agno/db/dynamo/__init__.py +3 -0
  27. agno/db/dynamo/dynamo.py +1743 -0
  28. agno/db/dynamo/schemas.py +278 -0
  29. agno/db/dynamo/utils.py +684 -0
  30. agno/db/firestore/__init__.py +3 -0
  31. agno/db/firestore/firestore.py +1432 -0
  32. agno/db/firestore/schemas.py +130 -0
  33. agno/db/firestore/utils.py +278 -0
  34. agno/db/gcs_json/__init__.py +3 -0
  35. agno/db/gcs_json/gcs_json_db.py +1001 -0
  36. agno/db/gcs_json/utils.py +194 -0
  37. agno/db/in_memory/__init__.py +3 -0
  38. agno/db/in_memory/in_memory_db.py +882 -0
  39. agno/db/in_memory/utils.py +172 -0
  40. agno/db/json/__init__.py +3 -0
  41. agno/db/json/json_db.py +1045 -0
  42. agno/db/json/utils.py +196 -0
  43. agno/db/migrations/v1_to_v2.py +162 -0
  44. agno/db/mongo/__init__.py +3 -0
  45. agno/db/mongo/mongo.py +1416 -0
  46. agno/db/mongo/schemas.py +77 -0
  47. agno/db/mongo/utils.py +204 -0
  48. agno/db/mysql/__init__.py +3 -0
  49. agno/db/mysql/mysql.py +1719 -0
  50. agno/db/mysql/schemas.py +124 -0
  51. agno/db/mysql/utils.py +297 -0
  52. agno/db/postgres/__init__.py +3 -0
  53. agno/db/postgres/postgres.py +1710 -0
  54. agno/db/postgres/schemas.py +124 -0
  55. agno/db/postgres/utils.py +280 -0
  56. agno/db/redis/__init__.py +3 -0
  57. agno/db/redis/redis.py +1367 -0
  58. agno/db/redis/schemas.py +109 -0
  59. agno/db/redis/utils.py +288 -0
  60. agno/db/schemas/__init__.py +3 -0
  61. agno/db/schemas/evals.py +33 -0
  62. agno/db/schemas/knowledge.py +40 -0
  63. agno/db/schemas/memory.py +46 -0
  64. agno/db/singlestore/__init__.py +3 -0
  65. agno/db/singlestore/schemas.py +116 -0
  66. agno/db/singlestore/singlestore.py +1712 -0
  67. agno/db/singlestore/utils.py +326 -0
  68. agno/db/sqlite/__init__.py +3 -0
  69. agno/db/sqlite/schemas.py +119 -0
  70. agno/db/sqlite/sqlite.py +1676 -0
  71. agno/db/sqlite/utils.py +268 -0
  72. agno/db/utils.py +88 -0
  73. agno/eval/__init__.py +14 -0
  74. agno/eval/accuracy.py +154 -48
  75. agno/eval/performance.py +88 -23
  76. agno/eval/reliability.py +73 -20
  77. agno/eval/utils.py +23 -13
  78. agno/integrations/discord/__init__.py +3 -0
  79. agno/{app → integrations}/discord/client.py +15 -11
  80. agno/knowledge/__init__.py +2 -2
  81. agno/{document → knowledge}/chunking/agentic.py +2 -2
  82. agno/{document → knowledge}/chunking/document.py +2 -2
  83. agno/{document → knowledge}/chunking/fixed.py +3 -3
  84. agno/{document → knowledge}/chunking/markdown.py +2 -2
  85. agno/{document → knowledge}/chunking/recursive.py +2 -2
  86. agno/{document → knowledge}/chunking/row.py +2 -2
  87. agno/knowledge/chunking/semantic.py +59 -0
  88. agno/knowledge/chunking/strategy.py +121 -0
  89. agno/knowledge/content.py +74 -0
  90. agno/knowledge/document/__init__.py +5 -0
  91. agno/{document → knowledge/document}/base.py +12 -2
  92. agno/knowledge/embedder/__init__.py +5 -0
  93. agno/{embedder → knowledge/embedder}/aws_bedrock.py +127 -1
  94. agno/{embedder → knowledge/embedder}/azure_openai.py +65 -1
  95. agno/{embedder → knowledge/embedder}/base.py +6 -0
  96. agno/{embedder → knowledge/embedder}/cohere.py +72 -1
  97. agno/{embedder → knowledge/embedder}/fastembed.py +17 -1
  98. agno/{embedder → knowledge/embedder}/fireworks.py +1 -1
  99. agno/{embedder → knowledge/embedder}/google.py +74 -1
  100. agno/{embedder → knowledge/embedder}/huggingface.py +36 -2
  101. agno/{embedder → knowledge/embedder}/jina.py +48 -2
  102. agno/knowledge/embedder/langdb.py +22 -0
  103. agno/knowledge/embedder/mistral.py +139 -0
  104. agno/{embedder → knowledge/embedder}/nebius.py +1 -1
  105. agno/{embedder → knowledge/embedder}/ollama.py +54 -3
  106. agno/knowledge/embedder/openai.py +223 -0
  107. agno/{embedder → knowledge/embedder}/sentence_transformer.py +16 -1
  108. agno/{embedder → knowledge/embedder}/together.py +1 -1
  109. agno/{embedder → knowledge/embedder}/voyageai.py +49 -1
  110. agno/knowledge/knowledge.py +1551 -0
  111. agno/knowledge/reader/__init__.py +7 -0
  112. agno/{document → knowledge}/reader/arxiv_reader.py +32 -4
  113. agno/knowledge/reader/base.py +88 -0
  114. agno/{document → knowledge}/reader/csv_reader.py +47 -65
  115. agno/knowledge/reader/docx_reader.py +83 -0
  116. agno/{document → knowledge}/reader/firecrawl_reader.py +42 -21
  117. agno/{document → knowledge}/reader/json_reader.py +30 -9
  118. agno/{document → knowledge}/reader/markdown_reader.py +58 -9
  119. agno/{document → knowledge}/reader/pdf_reader.py +71 -126
  120. agno/knowledge/reader/reader_factory.py +268 -0
  121. agno/knowledge/reader/s3_reader.py +101 -0
  122. agno/{document → knowledge}/reader/text_reader.py +31 -10
  123. agno/knowledge/reader/url_reader.py +128 -0
  124. agno/knowledge/reader/web_search_reader.py +366 -0
  125. agno/{document → knowledge}/reader/website_reader.py +37 -10
  126. agno/knowledge/reader/wikipedia_reader.py +59 -0
  127. agno/knowledge/reader/youtube_reader.py +78 -0
  128. agno/knowledge/remote_content/remote_content.py +88 -0
  129. agno/{reranker → knowledge/reranker}/base.py +1 -1
  130. agno/{reranker → knowledge/reranker}/cohere.py +2 -2
  131. agno/{reranker → knowledge/reranker}/infinity.py +2 -2
  132. agno/{reranker → knowledge/reranker}/sentence_transformer.py +2 -2
  133. agno/knowledge/types.py +30 -0
  134. agno/knowledge/utils.py +169 -0
  135. agno/media.py +269 -268
  136. agno/memory/__init__.py +2 -10
  137. agno/memory/manager.py +1003 -148
  138. agno/models/aimlapi/__init__.py +2 -2
  139. agno/models/aimlapi/aimlapi.py +6 -6
  140. agno/models/anthropic/claude.py +131 -131
  141. agno/models/aws/bedrock.py +110 -182
  142. agno/models/aws/claude.py +64 -18
  143. agno/models/azure/ai_foundry.py +73 -23
  144. agno/models/base.py +346 -290
  145. agno/models/cerebras/cerebras.py +84 -27
  146. agno/models/cohere/chat.py +106 -98
  147. agno/models/google/gemini.py +105 -46
  148. agno/models/groq/groq.py +97 -35
  149. agno/models/huggingface/huggingface.py +92 -27
  150. agno/models/ibm/watsonx.py +72 -13
  151. agno/models/litellm/chat.py +85 -13
  152. agno/models/message.py +46 -151
  153. agno/models/meta/llama.py +85 -49
  154. agno/models/metrics.py +120 -0
  155. agno/models/mistral/mistral.py +90 -21
  156. agno/models/ollama/__init__.py +0 -2
  157. agno/models/ollama/chat.py +85 -47
  158. agno/models/openai/chat.py +154 -37
  159. agno/models/openai/responses.py +178 -105
  160. agno/models/perplexity/perplexity.py +26 -2
  161. agno/models/portkey/portkey.py +0 -7
  162. agno/models/response.py +15 -9
  163. agno/models/utils.py +20 -0
  164. agno/models/vercel/__init__.py +2 -2
  165. agno/models/vercel/v0.py +1 -1
  166. agno/models/vllm/__init__.py +2 -2
  167. agno/models/vllm/vllm.py +3 -3
  168. agno/models/xai/xai.py +10 -10
  169. agno/os/__init__.py +3 -0
  170. agno/os/app.py +497 -0
  171. agno/os/auth.py +47 -0
  172. agno/os/config.py +103 -0
  173. agno/os/interfaces/agui/__init__.py +3 -0
  174. agno/os/interfaces/agui/agui.py +31 -0
  175. agno/{app/agui/async_router.py → os/interfaces/agui/router.py} +16 -16
  176. agno/{app → os/interfaces}/agui/utils.py +77 -33
  177. agno/os/interfaces/base.py +21 -0
  178. agno/os/interfaces/slack/__init__.py +3 -0
  179. agno/{app/slack/async_router.py → os/interfaces/slack/router.py} +3 -5
  180. agno/os/interfaces/slack/slack.py +32 -0
  181. agno/os/interfaces/whatsapp/__init__.py +3 -0
  182. agno/{app/whatsapp/async_router.py → os/interfaces/whatsapp/router.py} +4 -7
  183. agno/os/interfaces/whatsapp/whatsapp.py +29 -0
  184. agno/os/mcp.py +235 -0
  185. agno/os/router.py +1400 -0
  186. agno/os/routers/__init__.py +3 -0
  187. agno/os/routers/evals/__init__.py +3 -0
  188. agno/os/routers/evals/evals.py +393 -0
  189. agno/os/routers/evals/schemas.py +142 -0
  190. agno/os/routers/evals/utils.py +161 -0
  191. agno/os/routers/knowledge/__init__.py +3 -0
  192. agno/os/routers/knowledge/knowledge.py +850 -0
  193. agno/os/routers/knowledge/schemas.py +118 -0
  194. agno/os/routers/memory/__init__.py +3 -0
  195. agno/os/routers/memory/memory.py +410 -0
  196. agno/os/routers/memory/schemas.py +58 -0
  197. agno/os/routers/metrics/__init__.py +3 -0
  198. agno/os/routers/metrics/metrics.py +178 -0
  199. agno/os/routers/metrics/schemas.py +47 -0
  200. agno/os/routers/session/__init__.py +3 -0
  201. agno/os/routers/session/session.py +536 -0
  202. agno/os/schema.py +945 -0
  203. agno/{app/playground → os}/settings.py +7 -15
  204. agno/os/utils.py +270 -0
  205. agno/reasoning/azure_ai_foundry.py +4 -4
  206. agno/reasoning/deepseek.py +4 -4
  207. agno/reasoning/default.py +6 -11
  208. agno/reasoning/groq.py +4 -4
  209. agno/reasoning/helpers.py +4 -6
  210. agno/reasoning/ollama.py +4 -4
  211. agno/reasoning/openai.py +4 -4
  212. agno/run/agent.py +633 -0
  213. agno/run/base.py +53 -77
  214. agno/run/cancel.py +81 -0
  215. agno/run/team.py +243 -96
  216. agno/run/workflow.py +550 -12
  217. agno/session/__init__.py +10 -0
  218. agno/session/agent.py +244 -0
  219. agno/session/summary.py +225 -0
  220. agno/session/team.py +262 -0
  221. agno/{storage/session/v2 → session}/workflow.py +47 -24
  222. agno/team/__init__.py +15 -16
  223. agno/team/team.py +3260 -4824
  224. agno/tools/agentql.py +14 -5
  225. agno/tools/airflow.py +9 -4
  226. agno/tools/api.py +7 -3
  227. agno/tools/apify.py +2 -46
  228. agno/tools/arxiv.py +8 -3
  229. agno/tools/aws_lambda.py +7 -5
  230. agno/tools/aws_ses.py +7 -1
  231. agno/tools/baidusearch.py +4 -1
  232. agno/tools/bitbucket.py +4 -4
  233. agno/tools/brandfetch.py +14 -11
  234. agno/tools/bravesearch.py +4 -1
  235. agno/tools/brightdata.py +43 -23
  236. agno/tools/browserbase.py +13 -4
  237. agno/tools/calcom.py +12 -10
  238. agno/tools/calculator.py +10 -27
  239. agno/tools/cartesia.py +20 -17
  240. agno/tools/{clickup_tool.py → clickup.py} +12 -25
  241. agno/tools/confluence.py +8 -8
  242. agno/tools/crawl4ai.py +7 -1
  243. agno/tools/csv_toolkit.py +9 -8
  244. agno/tools/dalle.py +22 -12
  245. agno/tools/daytona.py +13 -16
  246. agno/tools/decorator.py +6 -3
  247. agno/tools/desi_vocal.py +17 -8
  248. agno/tools/discord.py +11 -8
  249. agno/tools/docker.py +30 -42
  250. agno/tools/duckdb.py +34 -53
  251. agno/tools/duckduckgo.py +8 -7
  252. agno/tools/e2b.py +62 -62
  253. agno/tools/eleven_labs.py +36 -29
  254. agno/tools/email.py +4 -1
  255. agno/tools/evm.py +7 -1
  256. agno/tools/exa.py +19 -14
  257. agno/tools/fal.py +30 -30
  258. agno/tools/file.py +9 -8
  259. agno/tools/financial_datasets.py +25 -44
  260. agno/tools/firecrawl.py +22 -22
  261. agno/tools/function.py +127 -18
  262. agno/tools/giphy.py +23 -11
  263. agno/tools/github.py +48 -126
  264. agno/tools/gmail.py +45 -61
  265. agno/tools/google_bigquery.py +7 -6
  266. agno/tools/google_maps.py +11 -26
  267. agno/tools/googlesearch.py +7 -2
  268. agno/tools/googlesheets.py +21 -17
  269. agno/tools/hackernews.py +9 -5
  270. agno/tools/jina.py +5 -4
  271. agno/tools/jira.py +18 -9
  272. agno/tools/knowledge.py +31 -32
  273. agno/tools/linear.py +19 -34
  274. agno/tools/linkup.py +5 -1
  275. agno/tools/local_file_system.py +8 -5
  276. agno/tools/lumalab.py +32 -20
  277. agno/tools/mcp.py +1 -2
  278. agno/tools/mem0.py +18 -12
  279. agno/tools/memori.py +14 -10
  280. agno/tools/mlx_transcribe.py +3 -2
  281. agno/tools/models/azure_openai.py +33 -15
  282. agno/tools/models/gemini.py +59 -32
  283. agno/tools/models/groq.py +30 -23
  284. agno/tools/models/nebius.py +28 -12
  285. agno/tools/models_labs.py +40 -16
  286. agno/tools/moviepy_video.py +7 -6
  287. agno/tools/neo4j.py +10 -8
  288. agno/tools/newspaper.py +7 -2
  289. agno/tools/newspaper4k.py +8 -3
  290. agno/tools/openai.py +58 -32
  291. agno/tools/openbb.py +12 -11
  292. agno/tools/opencv.py +63 -47
  293. agno/tools/openweather.py +14 -12
  294. agno/tools/pandas.py +11 -3
  295. agno/tools/postgres.py +4 -12
  296. agno/tools/pubmed.py +4 -1
  297. agno/tools/python.py +9 -22
  298. agno/tools/reasoning.py +35 -27
  299. agno/tools/reddit.py +11 -26
  300. agno/tools/replicate.py +55 -42
  301. agno/tools/resend.py +4 -1
  302. agno/tools/scrapegraph.py +15 -14
  303. agno/tools/searxng.py +10 -23
  304. agno/tools/serpapi.py +6 -3
  305. agno/tools/serper.py +13 -4
  306. agno/tools/shell.py +9 -2
  307. agno/tools/slack.py +12 -11
  308. agno/tools/sleep.py +3 -2
  309. agno/tools/spider.py +24 -4
  310. agno/tools/sql.py +7 -6
  311. agno/tools/tavily.py +6 -4
  312. agno/tools/telegram.py +12 -4
  313. agno/tools/todoist.py +11 -31
  314. agno/tools/toolkit.py +1 -1
  315. agno/tools/trafilatura.py +22 -6
  316. agno/tools/trello.py +9 -22
  317. agno/tools/twilio.py +10 -3
  318. agno/tools/user_control_flow.py +6 -1
  319. agno/tools/valyu.py +34 -5
  320. agno/tools/visualization.py +19 -28
  321. agno/tools/webbrowser.py +4 -3
  322. agno/tools/webex.py +11 -7
  323. agno/tools/website.py +15 -46
  324. agno/tools/webtools.py +12 -4
  325. agno/tools/whatsapp.py +5 -9
  326. agno/tools/wikipedia.py +20 -13
  327. agno/tools/x.py +14 -13
  328. agno/tools/yfinance.py +13 -40
  329. agno/tools/youtube.py +26 -20
  330. agno/tools/zendesk.py +7 -2
  331. agno/tools/zep.py +10 -7
  332. agno/tools/zoom.py +10 -9
  333. agno/utils/common.py +1 -19
  334. agno/utils/events.py +100 -123
  335. agno/utils/gemini.py +32 -2
  336. agno/utils/knowledge.py +29 -0
  337. agno/utils/log.py +54 -4
  338. agno/utils/mcp.py +68 -10
  339. agno/utils/media.py +39 -0
  340. agno/utils/message.py +12 -1
  341. agno/utils/models/aws_claude.py +1 -1
  342. agno/utils/models/claude.py +47 -4
  343. agno/utils/models/cohere.py +1 -1
  344. agno/utils/models/mistral.py +8 -7
  345. agno/utils/models/schema_utils.py +3 -3
  346. agno/utils/models/watsonx.py +1 -1
  347. agno/utils/openai.py +1 -1
  348. agno/utils/pprint.py +33 -32
  349. agno/utils/print_response/agent.py +779 -0
  350. agno/utils/print_response/team.py +1669 -0
  351. agno/utils/print_response/workflow.py +1451 -0
  352. agno/utils/prompts.py +14 -14
  353. agno/utils/reasoning.py +87 -0
  354. agno/utils/response.py +42 -42
  355. agno/utils/streamlit.py +481 -0
  356. agno/utils/string.py +8 -22
  357. agno/utils/team.py +50 -0
  358. agno/utils/timer.py +2 -2
  359. agno/vectordb/base.py +33 -21
  360. agno/vectordb/cassandra/cassandra.py +287 -23
  361. agno/vectordb/chroma/chromadb.py +482 -59
  362. agno/vectordb/clickhouse/clickhousedb.py +270 -63
  363. agno/vectordb/couchbase/couchbase.py +309 -29
  364. agno/vectordb/lancedb/lance_db.py +360 -21
  365. agno/vectordb/langchaindb/__init__.py +5 -0
  366. agno/vectordb/langchaindb/langchaindb.py +145 -0
  367. agno/vectordb/lightrag/__init__.py +5 -0
  368. agno/vectordb/lightrag/lightrag.py +374 -0
  369. agno/vectordb/llamaindex/llamaindexdb.py +127 -0
  370. agno/vectordb/milvus/milvus.py +242 -32
  371. agno/vectordb/mongodb/mongodb.py +200 -24
  372. agno/vectordb/pgvector/pgvector.py +319 -37
  373. agno/vectordb/pineconedb/pineconedb.py +221 -27
  374. agno/vectordb/qdrant/qdrant.py +334 -14
  375. agno/vectordb/singlestore/singlestore.py +286 -29
  376. agno/vectordb/surrealdb/surrealdb.py +187 -7
  377. agno/vectordb/upstashdb/upstashdb.py +342 -26
  378. agno/vectordb/weaviate/weaviate.py +227 -165
  379. agno/workflow/__init__.py +17 -13
  380. agno/workflow/{v2/condition.py → condition.py} +135 -32
  381. agno/workflow/{v2/loop.py → loop.py} +115 -28
  382. agno/workflow/{v2/parallel.py → parallel.py} +138 -108
  383. agno/workflow/{v2/router.py → router.py} +133 -32
  384. agno/workflow/{v2/step.py → step.py} +207 -49
  385. agno/workflow/{v2/steps.py → steps.py} +147 -66
  386. agno/workflow/types.py +482 -0
  387. agno/workflow/workflow.py +2410 -696
  388. agno-2.0.0.dist-info/METADATA +494 -0
  389. agno-2.0.0.dist-info/RECORD +515 -0
  390. agno-2.0.0.dist-info/licenses/LICENSE +201 -0
  391. agno/agent/metrics.py +0 -107
  392. agno/api/app.py +0 -35
  393. agno/api/playground.py +0 -92
  394. agno/api/schemas/app.py +0 -12
  395. agno/api/schemas/playground.py +0 -22
  396. agno/api/schemas/user.py +0 -35
  397. agno/api/schemas/workspace.py +0 -46
  398. agno/api/user.py +0 -160
  399. agno/api/workflows.py +0 -33
  400. agno/api/workspace.py +0 -175
  401. agno/app/agui/__init__.py +0 -3
  402. agno/app/agui/app.py +0 -17
  403. agno/app/agui/sync_router.py +0 -120
  404. agno/app/base.py +0 -186
  405. agno/app/discord/__init__.py +0 -3
  406. agno/app/fastapi/__init__.py +0 -3
  407. agno/app/fastapi/app.py +0 -107
  408. agno/app/fastapi/async_router.py +0 -457
  409. agno/app/fastapi/sync_router.py +0 -448
  410. agno/app/playground/app.py +0 -228
  411. agno/app/playground/async_router.py +0 -1050
  412. agno/app/playground/deploy.py +0 -249
  413. agno/app/playground/operator.py +0 -183
  414. agno/app/playground/schemas.py +0 -220
  415. agno/app/playground/serve.py +0 -55
  416. agno/app/playground/sync_router.py +0 -1042
  417. agno/app/playground/utils.py +0 -46
  418. agno/app/settings.py +0 -15
  419. agno/app/slack/__init__.py +0 -3
  420. agno/app/slack/app.py +0 -19
  421. agno/app/slack/sync_router.py +0 -92
  422. agno/app/utils.py +0 -54
  423. agno/app/whatsapp/__init__.py +0 -3
  424. agno/app/whatsapp/app.py +0 -15
  425. agno/app/whatsapp/sync_router.py +0 -197
  426. agno/cli/auth_server.py +0 -249
  427. agno/cli/config.py +0 -274
  428. agno/cli/console.py +0 -88
  429. agno/cli/credentials.py +0 -23
  430. agno/cli/entrypoint.py +0 -571
  431. agno/cli/operator.py +0 -357
  432. agno/cli/settings.py +0 -96
  433. agno/cli/ws/ws_cli.py +0 -817
  434. agno/constants.py +0 -13
  435. agno/document/__init__.py +0 -5
  436. agno/document/chunking/semantic.py +0 -45
  437. agno/document/chunking/strategy.py +0 -31
  438. agno/document/reader/__init__.py +0 -5
  439. agno/document/reader/base.py +0 -47
  440. agno/document/reader/docx_reader.py +0 -60
  441. agno/document/reader/gcs/pdf_reader.py +0 -44
  442. agno/document/reader/s3/pdf_reader.py +0 -59
  443. agno/document/reader/s3/text_reader.py +0 -63
  444. agno/document/reader/url_reader.py +0 -59
  445. agno/document/reader/youtube_reader.py +0 -58
  446. agno/embedder/__init__.py +0 -5
  447. agno/embedder/langdb.py +0 -80
  448. agno/embedder/mistral.py +0 -82
  449. agno/embedder/openai.py +0 -78
  450. agno/file/__init__.py +0 -5
  451. agno/file/file.py +0 -16
  452. agno/file/local/csv.py +0 -32
  453. agno/file/local/txt.py +0 -19
  454. agno/infra/app.py +0 -240
  455. agno/infra/base.py +0 -144
  456. agno/infra/context.py +0 -20
  457. agno/infra/db_app.py +0 -52
  458. agno/infra/resource.py +0 -205
  459. agno/infra/resources.py +0 -55
  460. agno/knowledge/agent.py +0 -702
  461. agno/knowledge/arxiv.py +0 -33
  462. agno/knowledge/combined.py +0 -36
  463. agno/knowledge/csv.py +0 -144
  464. agno/knowledge/csv_url.py +0 -124
  465. agno/knowledge/document.py +0 -223
  466. agno/knowledge/docx.py +0 -137
  467. agno/knowledge/firecrawl.py +0 -34
  468. agno/knowledge/gcs/__init__.py +0 -0
  469. agno/knowledge/gcs/base.py +0 -39
  470. agno/knowledge/gcs/pdf.py +0 -125
  471. agno/knowledge/json.py +0 -137
  472. agno/knowledge/langchain.py +0 -71
  473. agno/knowledge/light_rag.py +0 -273
  474. agno/knowledge/llamaindex.py +0 -66
  475. agno/knowledge/markdown.py +0 -154
  476. agno/knowledge/pdf.py +0 -164
  477. agno/knowledge/pdf_bytes.py +0 -42
  478. agno/knowledge/pdf_url.py +0 -148
  479. agno/knowledge/s3/__init__.py +0 -0
  480. agno/knowledge/s3/base.py +0 -64
  481. agno/knowledge/s3/pdf.py +0 -33
  482. agno/knowledge/s3/text.py +0 -34
  483. agno/knowledge/text.py +0 -141
  484. agno/knowledge/url.py +0 -46
  485. agno/knowledge/website.py +0 -179
  486. agno/knowledge/wikipedia.py +0 -32
  487. agno/knowledge/youtube.py +0 -35
  488. agno/memory/agent.py +0 -423
  489. agno/memory/classifier.py +0 -104
  490. agno/memory/db/__init__.py +0 -5
  491. agno/memory/db/base.py +0 -42
  492. agno/memory/db/mongodb.py +0 -189
  493. agno/memory/db/postgres.py +0 -203
  494. agno/memory/db/sqlite.py +0 -193
  495. agno/memory/memory.py +0 -22
  496. agno/memory/row.py +0 -36
  497. agno/memory/summarizer.py +0 -201
  498. agno/memory/summary.py +0 -19
  499. agno/memory/team.py +0 -415
  500. agno/memory/v2/__init__.py +0 -2
  501. agno/memory/v2/db/__init__.py +0 -1
  502. agno/memory/v2/db/base.py +0 -42
  503. agno/memory/v2/db/firestore.py +0 -339
  504. agno/memory/v2/db/mongodb.py +0 -196
  505. agno/memory/v2/db/postgres.py +0 -214
  506. agno/memory/v2/db/redis.py +0 -187
  507. agno/memory/v2/db/schema.py +0 -54
  508. agno/memory/v2/db/sqlite.py +0 -209
  509. agno/memory/v2/manager.py +0 -437
  510. agno/memory/v2/memory.py +0 -1097
  511. agno/memory/v2/schema.py +0 -55
  512. agno/memory/v2/summarizer.py +0 -215
  513. agno/memory/workflow.py +0 -38
  514. agno/models/ollama/tools.py +0 -430
  515. agno/models/qwen/__init__.py +0 -5
  516. agno/playground/__init__.py +0 -10
  517. agno/playground/deploy.py +0 -3
  518. agno/playground/playground.py +0 -3
  519. agno/playground/serve.py +0 -3
  520. agno/playground/settings.py +0 -3
  521. agno/reranker/__init__.py +0 -0
  522. agno/run/response.py +0 -467
  523. agno/run/v2/__init__.py +0 -0
  524. agno/run/v2/workflow.py +0 -567
  525. agno/storage/__init__.py +0 -0
  526. agno/storage/agent/__init__.py +0 -0
  527. agno/storage/agent/dynamodb.py +0 -1
  528. agno/storage/agent/json.py +0 -1
  529. agno/storage/agent/mongodb.py +0 -1
  530. agno/storage/agent/postgres.py +0 -1
  531. agno/storage/agent/singlestore.py +0 -1
  532. agno/storage/agent/sqlite.py +0 -1
  533. agno/storage/agent/yaml.py +0 -1
  534. agno/storage/base.py +0 -60
  535. agno/storage/dynamodb.py +0 -673
  536. agno/storage/firestore.py +0 -297
  537. agno/storage/gcs_json.py +0 -261
  538. agno/storage/in_memory.py +0 -234
  539. agno/storage/json.py +0 -237
  540. agno/storage/mongodb.py +0 -328
  541. agno/storage/mysql.py +0 -685
  542. agno/storage/postgres.py +0 -682
  543. agno/storage/redis.py +0 -336
  544. agno/storage/session/__init__.py +0 -16
  545. agno/storage/session/agent.py +0 -64
  546. agno/storage/session/team.py +0 -63
  547. agno/storage/session/v2/__init__.py +0 -5
  548. agno/storage/session/workflow.py +0 -61
  549. agno/storage/singlestore.py +0 -606
  550. agno/storage/sqlite.py +0 -646
  551. agno/storage/workflow/__init__.py +0 -0
  552. agno/storage/workflow/mongodb.py +0 -1
  553. agno/storage/workflow/postgres.py +0 -1
  554. agno/storage/workflow/sqlite.py +0 -1
  555. agno/storage/yaml.py +0 -241
  556. agno/tools/thinking.py +0 -73
  557. agno/utils/defaults.py +0 -57
  558. agno/utils/filesystem.py +0 -39
  559. agno/utils/git.py +0 -52
  560. agno/utils/json_io.py +0 -30
  561. agno/utils/load_env.py +0 -19
  562. agno/utils/py_io.py +0 -19
  563. agno/utils/pyproject.py +0 -18
  564. agno/utils/resource_filter.py +0 -31
  565. agno/workflow/v2/__init__.py +0 -21
  566. agno/workflow/v2/types.py +0 -357
  567. agno/workflow/v2/workflow.py +0 -3312
  568. agno/workspace/__init__.py +0 -0
  569. agno/workspace/config.py +0 -325
  570. agno/workspace/enums.py +0 -6
  571. agno/workspace/helpers.py +0 -52
  572. agno/workspace/operator.py +0 -757
  573. agno/workspace/settings.py +0 -158
  574. agno-1.8.1.dist-info/METADATA +0 -982
  575. agno-1.8.1.dist-info/RECORD +0 -566
  576. agno-1.8.1.dist-info/entry_points.txt +0 -3
  577. agno-1.8.1.dist-info/licenses/LICENSE +0 -375
  578. /agno/{app → db/migrations}/__init__.py +0 -0
  579. /agno/{app/playground/__init__.py → db/schemas/metrics.py} +0 -0
  580. /agno/{cli → integrations}/__init__.py +0 -0
  581. /agno/{cli/ws → knowledge/chunking}/__init__.py +0 -0
  582. /agno/{document/chunking → knowledge/remote_content}/__init__.py +0 -0
  583. /agno/{document/reader/gcs → knowledge/reranker}/__init__.py +0 -0
  584. /agno/{document/reader/s3 → os/interfaces}/__init__.py +0 -0
  585. /agno/{app → os/interfaces}/slack/security.py +0 -0
  586. /agno/{app → os/interfaces}/whatsapp/security.py +0 -0
  587. /agno/{file/local → utils/print_response}/__init__.py +0 -0
  588. /agno/{infra → vectordb/llamaindex}/__init__.py +0 -0
  589. {agno-1.8.1.dist-info → agno-2.0.0.dist-info}/WHEEL +0 -0
  590. {agno-1.8.1.dist-info → agno-2.0.0.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,1432 @@
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_sorting,
10
+ bulk_upsert_metrics,
11
+ calculate_date_metrics,
12
+ create_collection_indexes,
13
+ fetch_all_sessions_data,
14
+ get_dates_to_calculate_metrics_for,
15
+ )
16
+ from agno.db.schemas.evals import EvalFilterType, EvalRunRecord, EvalType
17
+ from agno.db.schemas.knowledge import KnowledgeRow
18
+ from agno.db.schemas.memory import UserMemory
19
+ from agno.db.utils import deserialize_session_json_fields, serialize_session_json_fields
20
+ from agno.session import AgentSession, Session, TeamSession, WorkflowSession
21
+ from agno.utils.log import log_debug, log_error, log_info
22
+
23
+ try:
24
+ from google.cloud.firestore import Client, FieldFilter # type: ignore[import-untyped]
25
+ except ImportError:
26
+ raise ImportError(
27
+ "`google-cloud-firestore` not installed. Please install it using `pip install google-cloud-firestore`"
28
+ )
29
+
30
+
31
+ class FirestoreDb(BaseDb):
32
+ def __init__(
33
+ self,
34
+ db_client: Optional[Client] = None,
35
+ project_id: Optional[str] = None,
36
+ session_collection: Optional[str] = None,
37
+ memory_collection: Optional[str] = None,
38
+ metrics_collection: Optional[str] = None,
39
+ eval_collection: Optional[str] = None,
40
+ knowledge_collection: Optional[str] = None,
41
+ ):
42
+ """
43
+ Interface for interacting with a Firestore database.
44
+
45
+ Args:
46
+ db_client (Optional[Client]): The Firestore client to use.
47
+ project_id (Optional[str]): The GCP project ID for Firestore.
48
+ session_collection (Optional[str]): Name of the collection to store sessions.
49
+ memory_collection (Optional[str]): Name of the collection to store memories.
50
+ metrics_collection (Optional[str]): Name of the collection to store metrics.
51
+ eval_collection (Optional[str]): Name of the collection to store evaluation runs.
52
+ knowledge_collection (Optional[str]): Name of the collection to store knowledge documents.
53
+
54
+ Raises:
55
+ ValueError: If neither project_id nor db_client is provided.
56
+ """
57
+ super().__init__(
58
+ session_table=session_collection,
59
+ memory_table=memory_collection,
60
+ metrics_table=metrics_collection,
61
+ eval_table=eval_collection,
62
+ knowledge_table=knowledge_collection,
63
+ )
64
+
65
+ _client: Optional[Client] = db_client
66
+ if _client is None and project_id is not None:
67
+ _client = Client(project=project_id)
68
+ if _client is None:
69
+ raise ValueError("One of project_id or db_client must be provided")
70
+
71
+ self.project_id: Optional[str] = project_id
72
+ self.db_client: Client = _client
73
+
74
+ # -- DB methods --
75
+
76
+ def _get_collection(self, table_type: str, create_collection_if_not_found: Optional[bool] = True):
77
+ """Get or create a collection based on table type.
78
+
79
+ Args:
80
+ table_type (str): The type of table to get or create.
81
+ create_collection_if_not_found (Optional[bool]): Whether to create the collection if it doesn't exist.
82
+
83
+ Returns:
84
+ CollectionReference: The collection reference.
85
+ """
86
+ if table_type == "sessions":
87
+ if not hasattr(self, "session_collection"):
88
+ if self.session_table_name is None:
89
+ raise ValueError("Session collection was not provided on initialization")
90
+ self.session_collection = self._get_or_create_collection(
91
+ collection_name=self.session_table_name,
92
+ collection_type="sessions",
93
+ create_collection_if_not_found=create_collection_if_not_found,
94
+ )
95
+ return self.session_collection
96
+
97
+ if table_type == "memories":
98
+ if not hasattr(self, "memory_collection"):
99
+ if self.memory_table_name is None:
100
+ raise ValueError("Memory collection was not provided on initialization")
101
+ self.memory_collection = self._get_or_create_collection(
102
+ collection_name=self.memory_table_name,
103
+ collection_type="memories",
104
+ create_collection_if_not_found=create_collection_if_not_found,
105
+ )
106
+ return self.memory_collection
107
+
108
+ if table_type == "metrics":
109
+ if not hasattr(self, "metrics_collection"):
110
+ if self.metrics_table_name is None:
111
+ raise ValueError("Metrics collection was not provided on initialization")
112
+ self.metrics_collection = self._get_or_create_collection(
113
+ collection_name=self.metrics_table_name,
114
+ collection_type="metrics",
115
+ create_collection_if_not_found=create_collection_if_not_found,
116
+ )
117
+ return self.metrics_collection
118
+
119
+ if table_type == "evals":
120
+ if not hasattr(self, "eval_collection"):
121
+ if self.eval_table_name is None:
122
+ raise ValueError("Eval collection was not provided on initialization")
123
+ self.eval_collection = self._get_or_create_collection(
124
+ collection_name=self.eval_table_name,
125
+ collection_type="evals",
126
+ create_collection_if_not_found=create_collection_if_not_found,
127
+ )
128
+ return self.eval_collection
129
+
130
+ if table_type == "knowledge":
131
+ if not hasattr(self, "knowledge_collection"):
132
+ if self.knowledge_table_name is None:
133
+ raise ValueError("Knowledge collection was not provided on initialization")
134
+ self.knowledge_collection = self._get_or_create_collection(
135
+ collection_name=self.knowledge_table_name,
136
+ collection_type="knowledge",
137
+ create_collection_if_not_found=create_collection_if_not_found,
138
+ )
139
+ return self.knowledge_collection
140
+
141
+ raise ValueError(f"Unknown table type: {table_type}")
142
+
143
+ def _get_or_create_collection(
144
+ self, collection_name: str, collection_type: str, create_collection_if_not_found: Optional[bool] = True
145
+ ):
146
+ """Get or create a collection with proper indexes.
147
+
148
+ Args:
149
+ collection_name (str): The name of the collection to get or create.
150
+ collection_type (str): The type of collection to get or create.
151
+ create_collection_if_not_found (Optional[bool]): Whether to create the collection if it doesn't exist.
152
+
153
+ Returns:
154
+ Optional[CollectionReference]: The collection reference.
155
+ """
156
+ try:
157
+ collection_ref = self.db_client.collection(collection_name)
158
+
159
+ if not hasattr(self, f"_{collection_name}_initialized"):
160
+ if not create_collection_if_not_found:
161
+ return None
162
+ create_collection_indexes(self.db_client, collection_name, collection_type)
163
+ setattr(self, f"_{collection_name}_initialized", True)
164
+
165
+ return collection_ref
166
+
167
+ except Exception as e:
168
+ log_error(f"Error getting collection {collection_name}: {e}")
169
+ raise
170
+
171
+ # -- Session methods --
172
+
173
+ def delete_session(self, session_id: str) -> bool:
174
+ """Delete a session from the database.
175
+
176
+ Args:
177
+ session_id (str): The ID of the session to delete.
178
+ session_type (SessionType): The type of session to delete. Defaults to SessionType.AGENT.
179
+
180
+ Returns:
181
+ bool: True if the session was deleted, False otherwise.
182
+
183
+ Raises:
184
+ Exception: If there is an error deleting the session.
185
+ """
186
+ try:
187
+ collection_ref = self._get_collection(table_type="sessions")
188
+ docs = collection_ref.where(filter=FieldFilter("session_id", "==", session_id)).stream()
189
+
190
+ for doc in docs:
191
+ doc.reference.delete()
192
+ log_debug(f"Successfully deleted session with session_id: {session_id}")
193
+ return True
194
+
195
+ log_debug(f"No session found to delete with session_id: {session_id}")
196
+ return False
197
+
198
+ except Exception as e:
199
+ log_error(f"Error deleting session: {e}")
200
+ return False
201
+
202
+ def delete_sessions(self, session_ids: List[str]) -> None:
203
+ """Delete multiple sessions from the database.
204
+
205
+ Args:
206
+ session_ids (List[str]): The IDs of the sessions to delete.
207
+ """
208
+ try:
209
+ collection_ref = self._get_collection(table_type="sessions")
210
+ batch = self.db_client.batch()
211
+
212
+ deleted_count = 0
213
+ for session_id in session_ids:
214
+ docs = collection_ref.where(filter=FieldFilter("session_id", "==", session_id)).stream()
215
+ for doc in docs:
216
+ batch.delete(doc.reference)
217
+ deleted_count += 1
218
+
219
+ batch.commit()
220
+
221
+ log_debug(f"Successfully deleted {deleted_count} sessions")
222
+
223
+ except Exception as e:
224
+ log_error(f"Error deleting sessions: {e}")
225
+
226
+ def get_session(
227
+ self,
228
+ session_id: str,
229
+ session_type: SessionType,
230
+ user_id: Optional[str] = None,
231
+ deserialize: Optional[bool] = True,
232
+ ) -> Optional[Union[Session, Dict[str, Any]]]:
233
+ """Read a session from the database.
234
+
235
+ Args:
236
+ session_id (str): The ID of the session to get.
237
+ user_id (Optional[str]): The ID of the user to get the session for.
238
+ session_type (Optional[SessionType]): The type of session to get.
239
+ deserialize (Optional[bool]): Whether to serialize the session. Defaults to True.
240
+
241
+ Returns:
242
+ Union[Session, Dict[str, Any], None]:
243
+ - When deserialize=True: Session object
244
+ - When deserialize=False: Session dictionary
245
+
246
+ Raises:
247
+ Exception: If there is an error reading the session.
248
+ """
249
+ try:
250
+ collection_ref = self._get_collection(table_type="sessions")
251
+ query = collection_ref.where(filter=FieldFilter("session_id", "==", session_id))
252
+
253
+ if user_id is not None:
254
+ query = query.where(filter=FieldFilter("user_id", "==", user_id))
255
+ if session_type is not None:
256
+ query = query.where(filter=FieldFilter("session_type", "==", session_type.value))
257
+
258
+ docs = query.stream()
259
+ result = None
260
+ for doc in docs:
261
+ result = doc.to_dict()
262
+ break
263
+
264
+ if result is None:
265
+ return None
266
+
267
+ session = deserialize_session_json_fields(result)
268
+
269
+ if not deserialize:
270
+ return session
271
+
272
+ if session_type == SessionType.AGENT:
273
+ return AgentSession.from_dict(session)
274
+ elif session_type == SessionType.TEAM:
275
+ return TeamSession.from_dict(session)
276
+ else:
277
+ return WorkflowSession.from_dict(session)
278
+
279
+ except Exception as e:
280
+ log_error(f"Exception reading session: {e}")
281
+ return None
282
+
283
+ def get_sessions(
284
+ self,
285
+ session_type: Optional[SessionType] = None,
286
+ user_id: Optional[str] = None,
287
+ component_id: Optional[str] = None,
288
+ session_name: Optional[str] = None,
289
+ start_timestamp: Optional[int] = None,
290
+ end_timestamp: Optional[int] = None,
291
+ limit: Optional[int] = None,
292
+ page: Optional[int] = None,
293
+ sort_by: Optional[str] = None,
294
+ sort_order: Optional[str] = None,
295
+ deserialize: Optional[bool] = True,
296
+ ) -> Union[List[Session], Tuple[List[Dict[str, Any]], int]]:
297
+ """Get all sessions.
298
+
299
+ Args:
300
+ session_type (Optional[SessionType]): The type of session to get.
301
+ user_id (Optional[str]): The ID of the user to get the session for.
302
+ component_id (Optional[str]): The ID of the component to get the session for.
303
+ session_name (Optional[str]): The name of the session to filter by.
304
+ start_timestamp (Optional[int]): The start timestamp to filter sessions by.
305
+ end_timestamp (Optional[int]): The end timestamp to filter sessions by.
306
+ limit (Optional[int]): The limit of the sessions to get.
307
+ page (Optional[int]): The page number to get.
308
+ sort_by (Optional[str]): The field to sort the sessions by.
309
+ sort_order (Optional[str]): The order to sort the sessions by.
310
+ deserialize (Optional[bool]): Whether to serialize the sessions. Defaults to True.
311
+
312
+ Returns:
313
+ Union[List[AgentSession], List[TeamSession], List[WorkflowSession], Tuple[List[Dict[str, Any]], int]]:
314
+ - When deserialize=True: List of Session objects
315
+ - When deserialize=False: List of session dictionaries and the total count
316
+
317
+ Raises:
318
+ Exception: If there is an error reading the sessions.
319
+ """
320
+ try:
321
+ collection_ref = self._get_collection(table_type="sessions")
322
+ if collection_ref is None:
323
+ return [] if deserialize else ([], 0)
324
+
325
+ query = collection_ref
326
+
327
+ if user_id is not None:
328
+ query = query.where(filter=FieldFilter("user_id", "==", user_id))
329
+ if session_type is not None:
330
+ query = query.where(filter=FieldFilter("session_type", "==", session_type.value))
331
+ if component_id is not None:
332
+ if session_type == SessionType.AGENT:
333
+ query = query.where(filter=FieldFilter("agent_id", "==", component_id))
334
+ elif session_type == SessionType.TEAM:
335
+ query = query.where(filter=FieldFilter("team_id", "==", component_id))
336
+ elif session_type == SessionType.WORKFLOW:
337
+ query = query.where(filter=FieldFilter("workflow_id", "==", component_id))
338
+ if start_timestamp is not None:
339
+ query = query.where(filter=FieldFilter("created_at", ">=", start_timestamp))
340
+ if end_timestamp is not None:
341
+ query = query.where(filter=FieldFilter("created_at", "<=", end_timestamp))
342
+ if session_name is not None:
343
+ query = query.where(filter=FieldFilter("session_data.session_name", "==", session_name))
344
+
345
+ # Apply sorting
346
+ query = apply_sorting(query, sort_by, sort_order)
347
+
348
+ # Get all documents for counting before pagination
349
+ all_docs = query.stream()
350
+ all_records = [doc.to_dict() for doc in all_docs]
351
+
352
+ if not all_records:
353
+ return [] if deserialize else ([], 0)
354
+
355
+ all_sessions_raw = [deserialize_session_json_fields(record) for record in all_records]
356
+
357
+ # Get total count before pagination
358
+ total_count = len(all_sessions_raw)
359
+
360
+ # Apply pagination to the results
361
+ if limit is not None and page is not None:
362
+ start_index = (page - 1) * limit
363
+ end_index = start_index + limit
364
+ sessions_raw = all_sessions_raw[start_index:end_index]
365
+ elif limit is not None:
366
+ sessions_raw = all_sessions_raw[:limit]
367
+ else:
368
+ sessions_raw = all_sessions_raw
369
+
370
+ if not deserialize:
371
+ return sessions_raw, total_count
372
+
373
+ sessions: List[Union[AgentSession, TeamSession, WorkflowSession]] = []
374
+ for session in sessions_raw:
375
+ if session["session_type"] == SessionType.AGENT.value:
376
+ agent_session = AgentSession.from_dict(session)
377
+ if agent_session is not None:
378
+ sessions.append(agent_session)
379
+ elif session["session_type"] == SessionType.TEAM.value:
380
+ team_session = TeamSession.from_dict(session)
381
+ if team_session is not None:
382
+ sessions.append(team_session)
383
+ elif session["session_type"] == SessionType.WORKFLOW.value:
384
+ workflow_session = WorkflowSession.from_dict(session)
385
+ if workflow_session is not None:
386
+ sessions.append(workflow_session)
387
+
388
+ if not sessions:
389
+ return [] if deserialize else ([], 0)
390
+
391
+ return sessions
392
+
393
+ except Exception as e:
394
+ log_error(f"Exception reading sessions: {e}")
395
+ return [] if deserialize else ([], 0)
396
+
397
+ def rename_session(
398
+ self, session_id: str, session_type: SessionType, session_name: str, deserialize: Optional[bool] = True
399
+ ) -> Optional[Union[Session, Dict[str, Any]]]:
400
+ """Rename a session in the database.
401
+
402
+ Args:
403
+ session_id (str): The ID of the session to rename.
404
+ session_type (SessionType): The type of session to rename.
405
+ session_name (str): The new name of the session.
406
+ deserialize (Optional[bool]): Whether to serialize the session. Defaults to True.
407
+
408
+ Returns:
409
+ Optional[Union[Session, Dict[str, Any]]]:
410
+ - When deserialize=True: Session object
411
+ - When deserialize=False: Session dictionary
412
+
413
+ Raises:
414
+ Exception: If there is an error renaming the session.
415
+ """
416
+ try:
417
+ collection_ref = self._get_collection(table_type="sessions")
418
+
419
+ docs = collection_ref.where(filter=FieldFilter("session_id", "==", session_id)).stream()
420
+ doc_ref = next((doc.reference for doc in docs), None)
421
+
422
+ if doc_ref is None:
423
+ return None
424
+
425
+ doc_ref.update({"session_data.session_name": session_name, "updated_at": int(time.time())})
426
+
427
+ updated_doc = doc_ref.get()
428
+ if not updated_doc.exists:
429
+ return None
430
+
431
+ result = updated_doc.to_dict()
432
+ if result is None:
433
+ return None
434
+ deserialized_session = deserialize_session_json_fields(result)
435
+
436
+ log_debug(f"Renamed session with id '{session_id}' to '{session_name}'")
437
+
438
+ if not deserialize:
439
+ return deserialized_session
440
+
441
+ if session_type == SessionType.AGENT:
442
+ return AgentSession.from_dict(deserialized_session)
443
+ elif session_type == SessionType.TEAM:
444
+ return TeamSession.from_dict(deserialized_session)
445
+ else:
446
+ return WorkflowSession.from_dict(deserialized_session)
447
+
448
+ except Exception as e:
449
+ log_error(f"Exception renaming session: {e}")
450
+ return None
451
+
452
+ def upsert_session(
453
+ self, session: Session, deserialize: Optional[bool] = True
454
+ ) -> Optional[Union[Session, Dict[str, Any]]]:
455
+ """Insert or update a session in the database.
456
+
457
+ Args:
458
+ session (Session): The session to upsert.
459
+
460
+ Returns:
461
+ Optional[Session]: The upserted session.
462
+
463
+ Raises:
464
+ Exception: If there is an error upserting the session.
465
+ """
466
+ try:
467
+ collection_ref = self._get_collection(table_type="sessions", create_collection_if_not_found=True)
468
+ serialized_session_dict = serialize_session_json_fields(session.to_dict())
469
+
470
+ if isinstance(session, AgentSession):
471
+ record = {
472
+ "session_id": serialized_session_dict.get("session_id"),
473
+ "session_type": SessionType.AGENT.value,
474
+ "agent_id": serialized_session_dict.get("agent_id"),
475
+ "user_id": serialized_session_dict.get("user_id"),
476
+ "runs": serialized_session_dict.get("runs"),
477
+ "agent_data": serialized_session_dict.get("agent_data"),
478
+ "session_data": serialized_session_dict.get("session_data"),
479
+ "summary": serialized_session_dict.get("summary"),
480
+ "metadata": serialized_session_dict.get("metadata"),
481
+ "created_at": serialized_session_dict.get("created_at"),
482
+ "updated_at": int(time.time()),
483
+ }
484
+
485
+ elif isinstance(session, TeamSession):
486
+ record = {
487
+ "session_id": serialized_session_dict.get("session_id"),
488
+ "session_type": SessionType.TEAM.value,
489
+ "team_id": serialized_session_dict.get("team_id"),
490
+ "user_id": serialized_session_dict.get("user_id"),
491
+ "runs": serialized_session_dict.get("runs"),
492
+ "team_data": serialized_session_dict.get("team_data"),
493
+ "session_data": serialized_session_dict.get("session_data"),
494
+ "summary": serialized_session_dict.get("summary"),
495
+ "metadata": serialized_session_dict.get("metadata"),
496
+ "created_at": serialized_session_dict.get("created_at"),
497
+ "updated_at": int(time.time()),
498
+ }
499
+
500
+ elif isinstance(session, WorkflowSession):
501
+ record = {
502
+ "session_id": serialized_session_dict.get("session_id"),
503
+ "session_type": SessionType.WORKFLOW.value,
504
+ "workflow_id": serialized_session_dict.get("workflow_id"),
505
+ "user_id": serialized_session_dict.get("user_id"),
506
+ "runs": serialized_session_dict.get("runs"),
507
+ "workflow_data": serialized_session_dict.get("workflow_data"),
508
+ "session_data": serialized_session_dict.get("session_data"),
509
+ "summary": serialized_session_dict.get("summary"),
510
+ "metadata": serialized_session_dict.get("metadata"),
511
+ "created_at": serialized_session_dict.get("created_at"),
512
+ "updated_at": int(time.time()),
513
+ }
514
+
515
+ # Find existing document or create new one
516
+ docs = collection_ref.where(filter=FieldFilter("session_id", "==", record["session_id"])).stream()
517
+ doc_ref = next((doc.reference for doc in docs), None)
518
+
519
+ if doc_ref is None:
520
+ # Create new document
521
+ doc_ref = collection_ref.document()
522
+
523
+ doc_ref.set(record, merge=True)
524
+
525
+ # Get the updated document
526
+ updated_doc = doc_ref.get()
527
+ if not updated_doc.exists:
528
+ return None
529
+
530
+ result = updated_doc.to_dict()
531
+ if result is None:
532
+ return None
533
+ deserialized_session = deserialize_session_json_fields(result)
534
+
535
+ if not deserialize:
536
+ return deserialized_session
537
+
538
+ if isinstance(session, AgentSession):
539
+ return AgentSession.from_dict(deserialized_session)
540
+ elif isinstance(session, TeamSession):
541
+ return TeamSession.from_dict(deserialized_session)
542
+ else:
543
+ return WorkflowSession.from_dict(deserialized_session)
544
+
545
+ except Exception as e:
546
+ log_error(f"Exception upserting session: {e}")
547
+ return None
548
+
549
+ # -- Memory methods --
550
+
551
+ def delete_user_memory(self, memory_id: str):
552
+ """Delete a user memory from the database.
553
+
554
+ Args:
555
+ memory_id (str): The ID of the memory to delete.
556
+
557
+ Returns:
558
+ bool: True if the memory was deleted, False otherwise.
559
+
560
+ Raises:
561
+ Exception: If there is an error deleting the memory.
562
+ """
563
+ try:
564
+ collection_ref = self._get_collection(table_type="memories")
565
+ docs = collection_ref.where(filter=FieldFilter("memory_id", "==", memory_id)).stream()
566
+
567
+ deleted_count = 0
568
+ for doc in docs:
569
+ doc.reference.delete()
570
+ deleted_count += 1
571
+
572
+ success = deleted_count > 0
573
+ if success:
574
+ log_debug(f"Successfully deleted user memory id: {memory_id}")
575
+ else:
576
+ log_debug(f"No user memory found with id: {memory_id}")
577
+
578
+ except Exception as e:
579
+ log_error(f"Error deleting user memory: {e}")
580
+
581
+ def delete_user_memories(self, memory_ids: List[str]) -> None:
582
+ """Delete user memories from the database.
583
+
584
+ Args:
585
+ memory_ids (List[str]): The IDs of the memories to delete.
586
+
587
+ Raises:
588
+ Exception: If there is an error deleting the memories.
589
+ """
590
+ try:
591
+ collection_ref = self._get_collection(table_type="memories")
592
+ batch = self.db_client.batch()
593
+ deleted_count = 0
594
+
595
+ for memory_id in memory_ids:
596
+ docs = collection_ref.where(filter=FieldFilter("memory_id", "==", memory_id)).stream()
597
+ for doc in docs:
598
+ batch.delete(doc.reference)
599
+ deleted_count += 1
600
+
601
+ batch.commit()
602
+
603
+ if deleted_count == 0:
604
+ log_info(f"No memories found with ids: {memory_ids}")
605
+ else:
606
+ log_info(f"Successfully deleted {deleted_count} memories")
607
+
608
+ except Exception as e:
609
+ log_error(f"Error deleting memories: {e}")
610
+
611
+ def get_all_memory_topics(self, create_collection_if_not_found: Optional[bool] = True) -> List[str]:
612
+ """Get all memory topics from the database.
613
+
614
+ Returns:
615
+ List[str]: The topics.
616
+
617
+ Raises:
618
+ Exception: If there is an error getting the topics.
619
+ """
620
+ try:
621
+ collection_ref = self._get_collection(table_type="memories")
622
+ if collection_ref is None:
623
+ return []
624
+
625
+ docs = collection_ref.stream()
626
+
627
+ all_topics = set()
628
+ for doc in docs:
629
+ data = doc.to_dict()
630
+ topics = data.get("topics", [])
631
+ if topics:
632
+ all_topics.update(topics)
633
+
634
+ return [topic for topic in all_topics if topic]
635
+
636
+ except Exception as e:
637
+ log_error(f"Exception reading from collection: {e}")
638
+ return []
639
+
640
+ def get_user_memory(self, memory_id: str, deserialize: Optional[bool] = True) -> Optional[UserMemory]:
641
+ """Get a memory from the database.
642
+
643
+ Args:
644
+ memory_id (str): The ID of the memory to get.
645
+ deserialize (Optional[bool]): Whether to serialize the memory. Defaults to True.
646
+
647
+ Returns:
648
+ Optional[UserMemory]:
649
+ - When deserialize=True: UserMemory object
650
+ - When deserialize=False: Memory dictionary
651
+
652
+ Raises:
653
+ Exception: If there is an error getting the memory.
654
+ """
655
+ try:
656
+ collection_ref = self._get_collection(table_type="memories")
657
+ docs = collection_ref.where(filter=FieldFilter("memory_id", "==", memory_id)).stream()
658
+
659
+ result = None
660
+ for doc in docs:
661
+ result = doc.to_dict()
662
+ break
663
+
664
+ if result is None or not deserialize:
665
+ return result
666
+
667
+ return UserMemory.from_dict(result)
668
+
669
+ except Exception as e:
670
+ log_error(f"Exception reading from collection: {e}")
671
+ return None
672
+
673
+ def get_user_memories(
674
+ self,
675
+ user_id: Optional[str] = None,
676
+ agent_id: Optional[str] = None,
677
+ team_id: Optional[str] = None,
678
+ topics: Optional[List[str]] = None,
679
+ search_content: Optional[str] = None,
680
+ limit: Optional[int] = None,
681
+ page: Optional[int] = None,
682
+ sort_by: Optional[str] = None,
683
+ sort_order: Optional[str] = None,
684
+ deserialize: Optional[bool] = True,
685
+ ) -> Union[List[UserMemory], Tuple[List[Dict[str, Any]], int]]:
686
+ """Get all memories from the database as UserMemory objects.
687
+
688
+ Args:
689
+ user_id (Optional[str]): The ID of the user to get the memories for.
690
+ agent_id (Optional[str]): The ID of the agent to get the memories for.
691
+ team_id (Optional[str]): The ID of the team to get the memories for.
692
+ topics (Optional[List[str]]): The topics to filter the memories by.
693
+ search_content (Optional[str]): The content to filter the memories by.
694
+ limit (Optional[int]): The limit of the memories to get.
695
+ page (Optional[int]): The page number to get.
696
+ sort_by (Optional[str]): The field to sort the memories by.
697
+ sort_order (Optional[str]): The order to sort the memories by.
698
+ deserialize (Optional[bool]): Whether to serialize the memories. Defaults to True.
699
+ create_table_if_not_found: Whether to create the index if it doesn't exist.
700
+
701
+ Returns:
702
+ Tuple[List[Dict[str, Any]], int]: A tuple containing the memories and the total count.
703
+
704
+ Raises:
705
+ Exception: If there is an error getting the memories.
706
+ """
707
+ try:
708
+ collection_ref = self._get_collection(table_type="memories")
709
+ if collection_ref is None:
710
+ return [] if deserialize else ([], 0)
711
+
712
+ query = collection_ref
713
+
714
+ if user_id is not None:
715
+ query = query.where(filter=FieldFilter("user_id", "==", user_id))
716
+ if agent_id is not None:
717
+ query = query.where(filter=FieldFilter("agent_id", "==", agent_id))
718
+ if team_id is not None:
719
+ query = query.where(filter=FieldFilter("team_id", "==", team_id))
720
+ if topics is not None and len(topics) > 0:
721
+ query = query.where(filter=FieldFilter("topics", "array_contains_any", topics))
722
+ if search_content is not None:
723
+ query = query.where(filter=FieldFilter("memory", "==", search_content))
724
+
725
+ # Apply sorting
726
+ query = apply_sorting(query, sort_by, sort_order)
727
+
728
+ # Get all documents
729
+ docs = query.stream()
730
+ all_records = [doc.to_dict() for doc in docs]
731
+
732
+ total_count = len(all_records)
733
+
734
+ # Apply pagination to the filtered results
735
+ if limit is not None and page is not None:
736
+ start_index = (page - 1) * limit
737
+ end_index = start_index + limit
738
+ records = all_records[start_index:end_index]
739
+ elif limit is not None:
740
+ records = all_records[:limit]
741
+ else:
742
+ records = all_records
743
+ if not deserialize:
744
+ return records, total_count
745
+
746
+ return [UserMemory.from_dict(record) for record in records]
747
+
748
+ except Exception as e:
749
+ log_error(f"Exception reading from collection: {e}")
750
+ return []
751
+
752
+ def get_user_memory_stats(
753
+ self,
754
+ limit: Optional[int] = None,
755
+ page: Optional[int] = None,
756
+ ) -> Tuple[List[Dict[str, Any]], int]:
757
+ """Get user memories stats.
758
+
759
+ Args:
760
+ limit (Optional[int]): The limit of the memories to get.
761
+ page (Optional[int]): The page number to get.
762
+
763
+ Returns:
764
+ Tuple[List[Dict[str, Any]], int]: A tuple containing the memories stats and the total count.
765
+
766
+ Raises:
767
+ Exception: If there is an error getting the memories stats.
768
+ """
769
+ try:
770
+ collection_ref = self._get_collection(table_type="memories")
771
+ docs = collection_ref.where(filter=FieldFilter("user_id", "!=", None)).stream()
772
+
773
+ user_stats = {}
774
+ for doc in docs:
775
+ data = doc.to_dict()
776
+ user_id = data.get("user_id")
777
+ if user_id:
778
+ if user_id not in user_stats:
779
+ user_stats[user_id] = {
780
+ "user_id": user_id,
781
+ "total_memories": 0,
782
+ "last_memory_updated_at": 0,
783
+ }
784
+ user_stats[user_id]["total_memories"] += 1
785
+ updated_at = data.get("updated_at", 0)
786
+ if updated_at > user_stats[user_id]["last_memory_updated_at"]:
787
+ user_stats[user_id]["last_memory_updated_at"] = updated_at
788
+
789
+ # Convert to list and sort
790
+ formatted_results = list(user_stats.values())
791
+ formatted_results.sort(key=lambda x: x["last_memory_updated_at"], reverse=True)
792
+
793
+ total_count = len(formatted_results)
794
+
795
+ # Apply pagination
796
+ if limit is not None:
797
+ start_idx = 0
798
+ if page is not None:
799
+ start_idx = (page - 1) * limit
800
+ formatted_results = formatted_results[start_idx : start_idx + limit]
801
+
802
+ return formatted_results, total_count
803
+
804
+ except Exception as e:
805
+ log_error(f"Exception getting user memory stats: {e}")
806
+ return [], 0
807
+
808
+ def upsert_user_memory(
809
+ self, memory: UserMemory, deserialize: Optional[bool] = True
810
+ ) -> Optional[Union[UserMemory, Dict[str, Any]]]:
811
+ """Upsert a user memory in the database.
812
+
813
+ Args:
814
+ memory (UserMemory): The memory to upsert.
815
+ deserialize (Optional[bool]): Whether to serialize the memory. Defaults to True.
816
+
817
+ Returns:
818
+ Optional[Union[UserMemory, Dict[str, Any]]]:
819
+ - When deserialize=True: UserMemory object
820
+ - When deserialize=False: Memory dictionary
821
+
822
+ Raises:
823
+ Exception: If there is an error upserting the memory.
824
+ """
825
+ try:
826
+ collection_ref = self._get_collection(table_type="memories", create_collection_if_not_found=True)
827
+ if collection_ref is None:
828
+ return None
829
+
830
+ if memory.memory_id is None:
831
+ memory.memory_id = str(uuid4())
832
+
833
+ update_doc = memory.to_dict()
834
+ update_doc["updated_at"] = int(time.time())
835
+
836
+ # Find existing document or create new one
837
+ docs = collection_ref.where("memory_id", "==", memory.memory_id).stream()
838
+ doc_ref = next((doc.reference for doc in docs), None)
839
+
840
+ if doc_ref is None:
841
+ doc_ref = collection_ref.document()
842
+
843
+ doc_ref.set(update_doc, merge=True)
844
+
845
+ if not deserialize:
846
+ return update_doc
847
+
848
+ return UserMemory.from_dict(update_doc)
849
+
850
+ except Exception as e:
851
+ log_error(f"Exception upserting user memory: {e}")
852
+ return None
853
+
854
+ def clear_memories(self) -> None:
855
+ """Delete all memories from the database.
856
+
857
+ Raises:
858
+ Exception: If an error occurs during deletion.
859
+ """
860
+ try:
861
+ collection_ref = self._get_collection(table_type="memories")
862
+
863
+ # Get all documents in the collection
864
+ docs = collection_ref.stream()
865
+
866
+ # Delete all documents in batches
867
+ batch = self.db_client.batch()
868
+ batch_count = 0
869
+
870
+ for doc in docs:
871
+ batch.delete(doc.reference)
872
+ batch_count += 1
873
+
874
+ # Firestore batch has a limit of 500 operations
875
+ if batch_count >= 500:
876
+ batch.commit()
877
+ batch = self.db_client.batch()
878
+ batch_count = 0
879
+
880
+ # Commit remaining operations
881
+ if batch_count > 0:
882
+ batch.commit()
883
+
884
+ except Exception as e:
885
+ from agno.utils.log import log_warning
886
+
887
+ log_warning(f"Exception deleting all memories: {e}")
888
+
889
+ # -- Metrics methods --
890
+
891
+ def _get_all_sessions_for_metrics_calculation(
892
+ self, start_timestamp: Optional[int] = None, end_timestamp: Optional[int] = None
893
+ ) -> List[Dict[str, Any]]:
894
+ """Get all sessions of all types for metrics calculation."""
895
+ try:
896
+ collection_ref = self._get_collection(table_type="sessions")
897
+
898
+ query = collection_ref
899
+ if start_timestamp is not None:
900
+ query = query.where(filter=FieldFilter("created_at", ">=", start_timestamp))
901
+ if end_timestamp is not None:
902
+ query = query.where(filter=FieldFilter("created_at", "<=", end_timestamp))
903
+
904
+ docs = query.stream()
905
+ results = []
906
+ for doc in docs:
907
+ data = doc.to_dict()
908
+ # Only include required fields for metrics
909
+ result = {
910
+ "user_id": data.get("user_id"),
911
+ "session_data": data.get("session_data"),
912
+ "runs": data.get("runs"),
913
+ "created_at": data.get("created_at"),
914
+ "session_type": data.get("session_type"),
915
+ }
916
+ results.append(result)
917
+
918
+ return results
919
+
920
+ except Exception as e:
921
+ log_error(f"Exception reading from sessions collection: {e}")
922
+ return []
923
+
924
+ def _get_metrics_calculation_starting_date(self, collection_ref) -> Optional[date]:
925
+ """Get the first date for which metrics calculation is needed."""
926
+ try:
927
+ query = collection_ref.order_by("date", direction="DESCENDING").limit(1)
928
+ docs = query.stream()
929
+
930
+ for doc in docs:
931
+ data = doc.to_dict()
932
+ result_date = datetime.strptime(data["date"], "%Y-%m-%d").date()
933
+ if data.get("completed"):
934
+ return result_date + timedelta(days=1)
935
+ else:
936
+ return result_date
937
+
938
+ # No metrics records. Return the date of the first recorded session.
939
+ first_session_result = self.get_sessions(sort_by="created_at", sort_order="asc", limit=1, deserialize=False)
940
+ first_session_date = None
941
+
942
+ if isinstance(first_session_result, list) and len(first_session_result) > 0:
943
+ first_session_date = first_session_result[0].created_at # type: ignore
944
+ elif isinstance(first_session_result, tuple) and len(first_session_result[0]) > 0:
945
+ first_session_date = first_session_result[0][0].get("created_at")
946
+
947
+ if first_session_date is None:
948
+ return None
949
+
950
+ return datetime.fromtimestamp(first_session_date, tz=timezone.utc).date()
951
+
952
+ except Exception as e:
953
+ log_error(f"Exception getting metrics calculation starting date: {e}")
954
+ return None
955
+
956
+ def calculate_metrics(self) -> Optional[list[dict]]:
957
+ """Calculate metrics for all dates without complete metrics."""
958
+ try:
959
+ collection_ref = self._get_collection(table_type="metrics", create_collection_if_not_found=True)
960
+
961
+ starting_date = self._get_metrics_calculation_starting_date(collection_ref)
962
+ if starting_date is None:
963
+ log_info("No session data found. Won't calculate metrics.")
964
+ return None
965
+
966
+ dates_to_process = get_dates_to_calculate_metrics_for(starting_date)
967
+ if not dates_to_process:
968
+ log_info("Metrics already calculated for all relevant dates.")
969
+ return None
970
+
971
+ start_timestamp = int(datetime.combine(dates_to_process[0], datetime.min.time()).timestamp())
972
+ end_timestamp = int(
973
+ datetime.combine(dates_to_process[-1] + timedelta(days=1), datetime.min.time()).timestamp()
974
+ )
975
+
976
+ sessions = self._get_all_sessions_for_metrics_calculation(
977
+ start_timestamp=start_timestamp, end_timestamp=end_timestamp
978
+ )
979
+ all_sessions_data = fetch_all_sessions_data(
980
+ sessions=sessions, dates_to_process=dates_to_process, start_timestamp=start_timestamp
981
+ )
982
+ if not all_sessions_data:
983
+ log_info("No new session data found. Won't calculate metrics.")
984
+ return None
985
+
986
+ results = []
987
+ metrics_records = []
988
+
989
+ for date_to_process in dates_to_process:
990
+ date_key = date_to_process.isoformat()
991
+ sessions_for_date = all_sessions_data.get(date_key, {})
992
+
993
+ # Skip dates with no sessions
994
+ if not any(len(sessions) > 0 for sessions in sessions_for_date.values()):
995
+ continue
996
+
997
+ metrics_record = calculate_date_metrics(date_to_process, sessions_for_date)
998
+ metrics_records.append(metrics_record)
999
+
1000
+ if metrics_records:
1001
+ results = bulk_upsert_metrics(collection_ref, metrics_records)
1002
+
1003
+ log_debug("Updated metrics calculations")
1004
+
1005
+ return results
1006
+
1007
+ except Exception as e:
1008
+ log_error(f"Exception calculating metrics: {e}")
1009
+ raise e
1010
+
1011
+ def get_metrics(
1012
+ self,
1013
+ starting_date: Optional[date] = None,
1014
+ ending_date: Optional[date] = None,
1015
+ ) -> Tuple[List[dict], Optional[int]]:
1016
+ """Get all metrics matching the given date range."""
1017
+ try:
1018
+ collection_ref = self._get_collection(table_type="metrics")
1019
+ if collection_ref is None:
1020
+ return [], None
1021
+
1022
+ query = collection_ref
1023
+ if starting_date:
1024
+ query = query.where(filter=FieldFilter("date", ">=", starting_date.isoformat()))
1025
+ if ending_date:
1026
+ query = query.where(filter=FieldFilter("date", "<=", ending_date.isoformat()))
1027
+
1028
+ docs = query.stream()
1029
+ records = []
1030
+ latest_updated_at = 0
1031
+
1032
+ for doc in docs:
1033
+ data = doc.to_dict()
1034
+ records.append(data)
1035
+ updated_at = data.get("updated_at", 0)
1036
+ if updated_at > latest_updated_at:
1037
+ latest_updated_at = updated_at
1038
+
1039
+ if not records:
1040
+ return [], None
1041
+
1042
+ return records, latest_updated_at
1043
+
1044
+ except Exception as e:
1045
+ log_error(f"Exception getting metrics: {e}")
1046
+ return [], None
1047
+
1048
+ # -- Knowledge methods --
1049
+
1050
+ def delete_knowledge_content(self, id: str):
1051
+ """Delete a knowledge row from the database.
1052
+
1053
+ Args:
1054
+ id (str): The ID of the knowledge row to delete.
1055
+
1056
+ Raises:
1057
+ Exception: If an error occurs during deletion.
1058
+ """
1059
+ try:
1060
+ collection_ref = self._get_collection(table_type="knowledge")
1061
+ docs = collection_ref.where(filter=FieldFilter("id", "==", id)).stream()
1062
+
1063
+ for doc in docs:
1064
+ doc.reference.delete()
1065
+
1066
+ except Exception as e:
1067
+ log_error(f"Error deleting knowledge source: {e}")
1068
+
1069
+ def get_knowledge_content(self, id: str) -> Optional[KnowledgeRow]:
1070
+ """Get a knowledge row from the database.
1071
+
1072
+ Args:
1073
+ id (str): The ID of the knowledge row to get.
1074
+
1075
+ Returns:
1076
+ Optional[KnowledgeRow]: The knowledge row, or None if it doesn't exist.
1077
+
1078
+ Raises:
1079
+ Exception: If an error occurs during retrieval.
1080
+ """
1081
+ try:
1082
+ collection_ref = self._get_collection(table_type="knowledge")
1083
+ docs = collection_ref.where(filter=FieldFilter("id", "==", id)).stream()
1084
+
1085
+ for doc in docs:
1086
+ data = doc.to_dict()
1087
+ return KnowledgeRow.model_validate(data)
1088
+
1089
+ return None
1090
+
1091
+ except Exception as e:
1092
+ log_error(f"Error getting knowledge source: {e}")
1093
+ return None
1094
+
1095
+ def get_knowledge_contents(
1096
+ self,
1097
+ limit: Optional[int] = None,
1098
+ page: Optional[int] = None,
1099
+ sort_by: Optional[str] = None,
1100
+ sort_order: Optional[str] = None,
1101
+ ) -> Tuple[List[KnowledgeRow], int]:
1102
+ """Get all knowledge contents from the database.
1103
+
1104
+ Args:
1105
+ limit (Optional[int]): The maximum number of knowledge contents to return.
1106
+ page (Optional[int]): The page number.
1107
+ sort_by (Optional[str]): The column to sort by.
1108
+ sort_order (Optional[str]): The order to sort by.
1109
+ create_table_if_not_found (Optional[bool]): Whether to create the table if it doesn't exist.
1110
+
1111
+ Returns:
1112
+ Tuple[List[KnowledgeRow], int]: The knowledge contents and total count.
1113
+
1114
+ Raises:
1115
+ Exception: If an error occurs during retrieval.
1116
+ """
1117
+ try:
1118
+ collection_ref = self._get_collection(table_type="knowledge")
1119
+ if collection_ref is None:
1120
+ return [], 0
1121
+
1122
+ query = collection_ref
1123
+
1124
+ # Apply sorting
1125
+ query = apply_sorting(query, sort_by, sort_order)
1126
+
1127
+ # Apply pagination
1128
+ query = apply_pagination(query, limit, page)
1129
+
1130
+ docs = query.stream()
1131
+ records = []
1132
+ for doc in docs:
1133
+ records.append(doc.to_dict())
1134
+
1135
+ knowledge_rows = [KnowledgeRow.model_validate(record) for record in records]
1136
+ total_count = len(knowledge_rows) # Simplified count
1137
+
1138
+ return knowledge_rows, total_count
1139
+
1140
+ except Exception as e:
1141
+ log_error(f"Error getting knowledge sources: {e}")
1142
+ return [], 0
1143
+
1144
+ def upsert_knowledge_content(self, knowledge_row: KnowledgeRow):
1145
+ """Upsert knowledge content in the database.
1146
+
1147
+ Args:
1148
+ knowledge_row (KnowledgeRow): The knowledge row to upsert.
1149
+
1150
+ Returns:
1151
+ Optional[KnowledgeRow]: The upserted knowledge row, or None if the operation fails.
1152
+ """
1153
+ try:
1154
+ collection_ref = self._get_collection(table_type="knowledge", create_collection_if_not_found=True)
1155
+ if collection_ref is None:
1156
+ return None
1157
+
1158
+ update_doc = knowledge_row.model_dump()
1159
+
1160
+ # Find existing document or create new one
1161
+ docs = collection_ref.where(filter=FieldFilter("id", "==", knowledge_row.id)).stream()
1162
+ doc_ref = next((doc.reference for doc in docs), None)
1163
+
1164
+ if doc_ref is None:
1165
+ doc_ref = collection_ref.document()
1166
+
1167
+ doc_ref.set(update_doc, merge=True)
1168
+
1169
+ return knowledge_row
1170
+
1171
+ except Exception as e:
1172
+ log_error(f"Error upserting knowledge document: {e}")
1173
+ return None
1174
+
1175
+ # -- Eval methods --
1176
+
1177
+ def create_eval_run(self, eval_run: EvalRunRecord) -> Optional[EvalRunRecord]:
1178
+ """Create an EvalRunRecord in the database."""
1179
+ try:
1180
+ collection_ref = self._get_collection(table_type="evals", create_collection_if_not_found=True)
1181
+
1182
+ current_time = int(time.time())
1183
+ eval_dict = eval_run.model_dump()
1184
+ eval_dict["created_at"] = current_time
1185
+ eval_dict["updated_at"] = current_time
1186
+
1187
+ doc_ref = collection_ref.document()
1188
+ doc_ref.set(eval_dict)
1189
+
1190
+ log_debug(f"Created eval run with id '{eval_run.run_id}'")
1191
+
1192
+ return eval_run
1193
+
1194
+ except Exception as e:
1195
+ log_error(f"Error creating eval run: {e}")
1196
+ return None
1197
+
1198
+ def delete_eval_run(self, eval_run_id: str) -> None:
1199
+ """Delete an eval run from the database."""
1200
+ try:
1201
+ collection_ref = self._get_collection(table_type="evals")
1202
+ docs = collection_ref.where(filter=FieldFilter("run_id", "==", eval_run_id)).stream()
1203
+
1204
+ deleted_count = 0
1205
+ for doc in docs:
1206
+ doc.reference.delete()
1207
+ deleted_count += 1
1208
+
1209
+ if deleted_count == 0:
1210
+ log_info(f"No eval run found with ID: {eval_run_id}")
1211
+ else:
1212
+ log_info(f"Deleted eval run with ID: {eval_run_id}")
1213
+
1214
+ except Exception as e:
1215
+ log_error(f"Error deleting eval run {eval_run_id}: {e}")
1216
+ raise
1217
+
1218
+ def delete_eval_runs(self, eval_run_ids: List[str]) -> None:
1219
+ """Delete multiple eval runs from the database.
1220
+
1221
+ Args:
1222
+ eval_run_ids (List[str]): The IDs of the eval runs to delete.
1223
+
1224
+ Raises:
1225
+ Exception: If there is an error deleting the eval runs.
1226
+ """
1227
+ try:
1228
+ collection_ref = self._get_collection(table_type="evals")
1229
+ batch = self.db_client.batch()
1230
+ deleted_count = 0
1231
+
1232
+ for eval_run_id in eval_run_ids:
1233
+ docs = collection_ref.where(filter=FieldFilter("run_id", "==", eval_run_id)).stream()
1234
+ for doc in docs:
1235
+ batch.delete(doc.reference)
1236
+ deleted_count += 1
1237
+
1238
+ batch.commit()
1239
+
1240
+ if deleted_count == 0:
1241
+ log_info(f"No eval runs found with IDs: {eval_run_ids}")
1242
+ else:
1243
+ log_info(f"Deleted {deleted_count} eval runs")
1244
+
1245
+ except Exception as e:
1246
+ log_error(f"Error deleting eval runs {eval_run_ids}: {e}")
1247
+ raise
1248
+
1249
+ def get_eval_run(
1250
+ self, eval_run_id: str, deserialize: Optional[bool] = True
1251
+ ) -> Optional[Union[EvalRunRecord, Dict[str, Any]]]:
1252
+ """Get an eval run from the database.
1253
+
1254
+ Args:
1255
+ eval_run_id (str): The ID of the eval run to get.
1256
+ deserialize (Optional[bool]): Whether to serialize the eval run. Defaults to True.
1257
+
1258
+ Returns:
1259
+ Optional[Union[EvalRunRecord, Dict[str, Any]]]:
1260
+ - When deserialize=True: EvalRunRecord object
1261
+ - When deserialize=False: EvalRun dictionary
1262
+
1263
+ Raises:
1264
+ Exception: If there is an error getting the eval run.
1265
+ """
1266
+ try:
1267
+ collection_ref = self._get_collection(table_type="evals")
1268
+ docs = collection_ref.where(filter=FieldFilter("run_id", "==", eval_run_id)).stream()
1269
+
1270
+ eval_run_raw = None
1271
+ for doc in docs:
1272
+ eval_run_raw = doc.to_dict()
1273
+ break
1274
+
1275
+ if not eval_run_raw:
1276
+ return None
1277
+
1278
+ if not deserialize:
1279
+ return eval_run_raw
1280
+
1281
+ return EvalRunRecord.model_validate(eval_run_raw)
1282
+
1283
+ except Exception as e:
1284
+ log_error(f"Exception getting eval run {eval_run_id}: {e}")
1285
+ return None
1286
+
1287
+ def get_eval_runs(
1288
+ self,
1289
+ limit: Optional[int] = None,
1290
+ page: Optional[int] = None,
1291
+ sort_by: Optional[str] = None,
1292
+ sort_order: Optional[str] = None,
1293
+ agent_id: Optional[str] = None,
1294
+ team_id: Optional[str] = None,
1295
+ workflow_id: Optional[str] = None,
1296
+ model_id: Optional[str] = None,
1297
+ filter_type: Optional[EvalFilterType] = None,
1298
+ eval_type: Optional[List[EvalType]] = None,
1299
+ deserialize: Optional[bool] = True,
1300
+ ) -> Union[List[EvalRunRecord], Tuple[List[Dict[str, Any]], int]]:
1301
+ """Get all eval runs from the database.
1302
+
1303
+ Args:
1304
+ limit (Optional[int]): The maximum number of eval runs to return.
1305
+ page (Optional[int]): The page number to return.
1306
+ sort_by (Optional[str]): The field to sort by.
1307
+ sort_order (Optional[str]): The order to sort by.
1308
+ agent_id (Optional[str]): The ID of the agent to filter by.
1309
+ team_id (Optional[str]): The ID of the team to filter by.
1310
+ workflow_id (Optional[str]): The ID of the workflow to filter by.
1311
+ model_id (Optional[str]): The ID of the model to filter by.
1312
+ eval_type (Optional[List[EvalType]]): The type of eval to filter by.
1313
+ filter_type (Optional[EvalFilterType]): The type of filter to apply.
1314
+ deserialize (Optional[bool]): Whether to serialize the eval runs. Defaults to True.
1315
+ create_table_if_not_found (Optional[bool]): Whether to create the table if it doesn't exist.
1316
+
1317
+ Returns:
1318
+ Union[List[EvalRunRecord], Tuple[List[Dict[str, Any]], int]]:
1319
+ - When deserialize=True: List of EvalRunRecord objects
1320
+ - When deserialize=False: List of eval run dictionaries and the total count
1321
+
1322
+ Raises:
1323
+ Exception: If there is an error getting the eval runs.
1324
+ """
1325
+ try:
1326
+ collection_ref = self._get_collection(table_type="evals")
1327
+ if collection_ref is None:
1328
+ return [] if deserialize else ([], 0)
1329
+
1330
+ query = collection_ref
1331
+
1332
+ if agent_id is not None:
1333
+ query = query.where(filter=FieldFilter("agent_id", "==", agent_id))
1334
+ if team_id is not None:
1335
+ query = query.where(filter=FieldFilter("team_id", "==", team_id))
1336
+ if workflow_id is not None:
1337
+ query = query.where(filter=FieldFilter("workflow_id", "==", workflow_id))
1338
+ if model_id is not None:
1339
+ query = query.where(filter=FieldFilter("model_id", "==", model_id))
1340
+ if eval_type is not None and len(eval_type) > 0:
1341
+ eval_values = [et.value for et in eval_type]
1342
+ query = query.where(filter=FieldFilter("eval_type", "in", eval_values))
1343
+ if filter_type is not None:
1344
+ if filter_type == EvalFilterType.AGENT:
1345
+ query = query.where(filter=FieldFilter("agent_id", "!=", None))
1346
+ elif filter_type == EvalFilterType.TEAM:
1347
+ query = query.where(filter=FieldFilter("team_id", "!=", None))
1348
+ elif filter_type == EvalFilterType.WORKFLOW:
1349
+ query = query.where(filter=FieldFilter("workflow_id", "!=", None))
1350
+
1351
+ # Apply default sorting by created_at desc if no sort parameters provided
1352
+ if sort_by is None:
1353
+ from google.cloud.firestore import Query
1354
+
1355
+ query = query.order_by("created_at", direction=Query.DESCENDING)
1356
+ else:
1357
+ query = apply_sorting(query, sort_by, sort_order)
1358
+
1359
+ # Get all documents for counting before pagination
1360
+ all_docs = query.stream()
1361
+ all_records = [doc.to_dict() for doc in all_docs]
1362
+
1363
+ if not all_records:
1364
+ return [] if deserialize else ([], 0)
1365
+
1366
+ # Get total count before pagination
1367
+ total_count = len(all_records)
1368
+
1369
+ # Apply pagination to the results
1370
+ if limit is not None and page is not None:
1371
+ start_index = (page - 1) * limit
1372
+ end_index = start_index + limit
1373
+ records = all_records[start_index:end_index]
1374
+ elif limit is not None:
1375
+ records = all_records[:limit]
1376
+ else:
1377
+ records = all_records
1378
+
1379
+ if not deserialize:
1380
+ return records, total_count
1381
+
1382
+ return [EvalRunRecord.model_validate(row) for row in records]
1383
+
1384
+ except Exception as e:
1385
+ log_error(f"Exception getting eval runs: {e}")
1386
+ return [] if deserialize else ([], 0)
1387
+
1388
+ def rename_eval_run(
1389
+ self, eval_run_id: str, name: str, deserialize: Optional[bool] = True
1390
+ ) -> Optional[Union[EvalRunRecord, Dict[str, Any]]]:
1391
+ """Update the name of an eval run in the database.
1392
+
1393
+ Args:
1394
+ eval_run_id (str): The ID of the eval run to update.
1395
+ name (str): The new name of the eval run.
1396
+ deserialize (Optional[bool]): Whether to serialize the eval run. Defaults to True.
1397
+
1398
+ Returns:
1399
+ Optional[Union[EvalRunRecord, Dict[str, Any]]]:
1400
+ - When deserialize=True: EvalRunRecord object
1401
+ - When deserialize=False: EvalRun dictionary
1402
+
1403
+ Raises:
1404
+ Exception: If there is an error updating the eval run.
1405
+ """
1406
+ try:
1407
+ collection_ref = self._get_collection(table_type="evals")
1408
+
1409
+ docs = collection_ref.where(filter=FieldFilter("run_id", "==", eval_run_id)).stream()
1410
+ doc_ref = next((doc.reference for doc in docs), None)
1411
+
1412
+ if doc_ref is None:
1413
+ return None
1414
+
1415
+ doc_ref.update({"name": name, "updated_at": int(time.time())})
1416
+
1417
+ updated_doc = doc_ref.get()
1418
+ if not updated_doc.exists:
1419
+ return None
1420
+
1421
+ result = updated_doc.to_dict()
1422
+
1423
+ log_debug(f"Renamed eval run with id '{eval_run_id}' to '{name}'")
1424
+
1425
+ if not result or not deserialize:
1426
+ return result
1427
+
1428
+ return EvalRunRecord.model_validate(result)
1429
+
1430
+ except Exception as e:
1431
+ log_error(f"Error updating eval run name {eval_run_id}: {e}")
1432
+ raise