agno 0.1.2__py3-none-any.whl → 2.3.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.
- agno/__init__.py +8 -0
- agno/agent/__init__.py +44 -5
- agno/agent/agent.py +10531 -2975
- agno/api/agent.py +14 -53
- agno/api/api.py +7 -46
- agno/api/evals.py +22 -0
- agno/api/os.py +17 -0
- agno/api/routes.py +6 -25
- agno/api/schemas/__init__.py +9 -0
- agno/api/schemas/agent.py +6 -9
- agno/api/schemas/evals.py +16 -0
- agno/api/schemas/os.py +14 -0
- agno/api/schemas/team.py +10 -10
- agno/api/schemas/utils.py +21 -0
- agno/api/schemas/workflows.py +16 -0
- agno/api/settings.py +53 -0
- agno/api/team.py +22 -26
- agno/api/workflow.py +28 -0
- agno/cloud/aws/base.py +214 -0
- agno/cloud/aws/s3/__init__.py +2 -0
- agno/cloud/aws/s3/api_client.py +43 -0
- agno/cloud/aws/s3/bucket.py +195 -0
- agno/cloud/aws/s3/object.py +57 -0
- agno/compression/__init__.py +3 -0
- agno/compression/manager.py +247 -0
- agno/culture/__init__.py +3 -0
- agno/culture/manager.py +956 -0
- agno/db/__init__.py +24 -0
- agno/db/async_postgres/__init__.py +3 -0
- agno/db/base.py +946 -0
- agno/db/dynamo/__init__.py +3 -0
- agno/db/dynamo/dynamo.py +2781 -0
- agno/db/dynamo/schemas.py +442 -0
- agno/db/dynamo/utils.py +743 -0
- agno/db/firestore/__init__.py +3 -0
- agno/db/firestore/firestore.py +2379 -0
- agno/db/firestore/schemas.py +181 -0
- agno/db/firestore/utils.py +376 -0
- agno/db/gcs_json/__init__.py +3 -0
- agno/db/gcs_json/gcs_json_db.py +1791 -0
- agno/db/gcs_json/utils.py +228 -0
- agno/db/in_memory/__init__.py +3 -0
- agno/db/in_memory/in_memory_db.py +1312 -0
- agno/db/in_memory/utils.py +230 -0
- agno/db/json/__init__.py +3 -0
- agno/db/json/json_db.py +1777 -0
- agno/db/json/utils.py +230 -0
- agno/db/migrations/manager.py +199 -0
- agno/db/migrations/v1_to_v2.py +635 -0
- agno/db/migrations/versions/v2_3_0.py +938 -0
- agno/db/mongo/__init__.py +17 -0
- agno/db/mongo/async_mongo.py +2760 -0
- agno/db/mongo/mongo.py +2597 -0
- agno/db/mongo/schemas.py +119 -0
- agno/db/mongo/utils.py +276 -0
- agno/db/mysql/__init__.py +4 -0
- agno/db/mysql/async_mysql.py +2912 -0
- agno/db/mysql/mysql.py +2923 -0
- agno/db/mysql/schemas.py +186 -0
- agno/db/mysql/utils.py +488 -0
- agno/db/postgres/__init__.py +4 -0
- agno/db/postgres/async_postgres.py +2579 -0
- agno/db/postgres/postgres.py +2870 -0
- agno/db/postgres/schemas.py +187 -0
- agno/db/postgres/utils.py +442 -0
- agno/db/redis/__init__.py +3 -0
- agno/db/redis/redis.py +2141 -0
- agno/db/redis/schemas.py +159 -0
- agno/db/redis/utils.py +346 -0
- agno/db/schemas/__init__.py +4 -0
- agno/db/schemas/culture.py +120 -0
- agno/db/schemas/evals.py +34 -0
- agno/db/schemas/knowledge.py +40 -0
- agno/db/schemas/memory.py +61 -0
- agno/db/singlestore/__init__.py +3 -0
- agno/db/singlestore/schemas.py +179 -0
- agno/db/singlestore/singlestore.py +2877 -0
- agno/db/singlestore/utils.py +384 -0
- agno/db/sqlite/__init__.py +4 -0
- agno/db/sqlite/async_sqlite.py +2911 -0
- agno/db/sqlite/schemas.py +181 -0
- agno/db/sqlite/sqlite.py +2908 -0
- agno/db/sqlite/utils.py +429 -0
- agno/db/surrealdb/__init__.py +3 -0
- agno/db/surrealdb/metrics.py +292 -0
- agno/db/surrealdb/models.py +334 -0
- agno/db/surrealdb/queries.py +71 -0
- agno/db/surrealdb/surrealdb.py +1908 -0
- agno/db/surrealdb/utils.py +147 -0
- agno/db/utils.py +118 -0
- agno/eval/__init__.py +24 -0
- agno/eval/accuracy.py +666 -276
- agno/eval/agent_as_judge.py +861 -0
- agno/eval/base.py +29 -0
- agno/eval/performance.py +779 -0
- agno/eval/reliability.py +241 -62
- agno/eval/utils.py +120 -0
- agno/exceptions.py +143 -1
- agno/filters.py +354 -0
- agno/guardrails/__init__.py +6 -0
- agno/guardrails/base.py +19 -0
- agno/guardrails/openai.py +144 -0
- agno/guardrails/pii.py +94 -0
- agno/guardrails/prompt_injection.py +52 -0
- agno/hooks/__init__.py +3 -0
- agno/hooks/decorator.py +164 -0
- agno/integrations/discord/__init__.py +3 -0
- agno/integrations/discord/client.py +203 -0
- agno/knowledge/__init__.py +5 -1
- agno/{document → knowledge}/chunking/agentic.py +22 -14
- agno/{document → knowledge}/chunking/document.py +2 -2
- agno/{document → knowledge}/chunking/fixed.py +7 -6
- agno/knowledge/chunking/markdown.py +151 -0
- agno/{document → knowledge}/chunking/recursive.py +15 -3
- agno/knowledge/chunking/row.py +39 -0
- agno/knowledge/chunking/semantic.py +91 -0
- agno/knowledge/chunking/strategy.py +165 -0
- agno/knowledge/content.py +74 -0
- agno/knowledge/document/__init__.py +5 -0
- agno/{document → knowledge/document}/base.py +12 -2
- agno/knowledge/embedder/__init__.py +5 -0
- agno/knowledge/embedder/aws_bedrock.py +343 -0
- agno/knowledge/embedder/azure_openai.py +210 -0
- agno/{embedder → knowledge/embedder}/base.py +8 -0
- agno/knowledge/embedder/cohere.py +323 -0
- agno/knowledge/embedder/fastembed.py +62 -0
- agno/{embedder → knowledge/embedder}/fireworks.py +1 -1
- agno/knowledge/embedder/google.py +258 -0
- agno/knowledge/embedder/huggingface.py +94 -0
- agno/knowledge/embedder/jina.py +182 -0
- agno/knowledge/embedder/langdb.py +22 -0
- agno/knowledge/embedder/mistral.py +206 -0
- agno/knowledge/embedder/nebius.py +13 -0
- agno/knowledge/embedder/ollama.py +154 -0
- agno/knowledge/embedder/openai.py +195 -0
- agno/knowledge/embedder/sentence_transformer.py +63 -0
- agno/{embedder → knowledge/embedder}/together.py +1 -1
- agno/knowledge/embedder/vllm.py +262 -0
- agno/knowledge/embedder/voyageai.py +165 -0
- agno/knowledge/knowledge.py +3006 -0
- agno/knowledge/reader/__init__.py +7 -0
- agno/knowledge/reader/arxiv_reader.py +81 -0
- agno/knowledge/reader/base.py +95 -0
- agno/knowledge/reader/csv_reader.py +164 -0
- agno/knowledge/reader/docx_reader.py +82 -0
- agno/knowledge/reader/field_labeled_csv_reader.py +290 -0
- agno/knowledge/reader/firecrawl_reader.py +201 -0
- agno/knowledge/reader/json_reader.py +88 -0
- agno/knowledge/reader/markdown_reader.py +137 -0
- agno/knowledge/reader/pdf_reader.py +431 -0
- agno/knowledge/reader/pptx_reader.py +101 -0
- agno/knowledge/reader/reader_factory.py +313 -0
- agno/knowledge/reader/s3_reader.py +89 -0
- agno/knowledge/reader/tavily_reader.py +193 -0
- agno/knowledge/reader/text_reader.py +127 -0
- agno/knowledge/reader/web_search_reader.py +325 -0
- agno/knowledge/reader/website_reader.py +455 -0
- agno/knowledge/reader/wikipedia_reader.py +91 -0
- agno/knowledge/reader/youtube_reader.py +78 -0
- agno/knowledge/remote_content/remote_content.py +88 -0
- agno/knowledge/reranker/__init__.py +3 -0
- agno/{reranker → knowledge/reranker}/base.py +1 -1
- agno/{reranker → knowledge/reranker}/cohere.py +2 -2
- agno/knowledge/reranker/infinity.py +195 -0
- agno/knowledge/reranker/sentence_transformer.py +54 -0
- agno/knowledge/types.py +39 -0
- agno/knowledge/utils.py +234 -0
- agno/media.py +439 -95
- agno/memory/__init__.py +16 -3
- agno/memory/manager.py +1474 -123
- agno/memory/strategies/__init__.py +15 -0
- agno/memory/strategies/base.py +66 -0
- agno/memory/strategies/summarize.py +196 -0
- agno/memory/strategies/types.py +37 -0
- agno/models/aimlapi/__init__.py +5 -0
- agno/models/aimlapi/aimlapi.py +62 -0
- agno/models/anthropic/__init__.py +4 -0
- agno/models/anthropic/claude.py +960 -496
- agno/models/aws/__init__.py +15 -0
- agno/models/aws/bedrock.py +686 -451
- agno/models/aws/claude.py +190 -183
- agno/models/azure/__init__.py +18 -1
- agno/models/azure/ai_foundry.py +489 -0
- agno/models/azure/openai_chat.py +89 -40
- agno/models/base.py +2477 -550
- agno/models/cerebras/__init__.py +12 -0
- agno/models/cerebras/cerebras.py +565 -0
- agno/models/cerebras/cerebras_openai.py +131 -0
- agno/models/cohere/__init__.py +4 -0
- agno/models/cohere/chat.py +306 -492
- agno/models/cometapi/__init__.py +5 -0
- agno/models/cometapi/cometapi.py +74 -0
- agno/models/dashscope/__init__.py +5 -0
- agno/models/dashscope/dashscope.py +90 -0
- agno/models/deepinfra/__init__.py +5 -0
- agno/models/deepinfra/deepinfra.py +45 -0
- agno/models/deepseek/__init__.py +4 -0
- agno/models/deepseek/deepseek.py +110 -9
- agno/models/fireworks/__init__.py +4 -0
- agno/models/fireworks/fireworks.py +19 -22
- agno/models/google/__init__.py +3 -7
- agno/models/google/gemini.py +1717 -662
- agno/models/google/utils.py +22 -0
- agno/models/groq/__init__.py +4 -0
- agno/models/groq/groq.py +391 -666
- agno/models/huggingface/__init__.py +4 -0
- agno/models/huggingface/huggingface.py +266 -538
- agno/models/ibm/__init__.py +5 -0
- agno/models/ibm/watsonx.py +432 -0
- agno/models/internlm/__init__.py +3 -0
- agno/models/internlm/internlm.py +20 -3
- agno/models/langdb/__init__.py +1 -0
- agno/models/langdb/langdb.py +60 -0
- agno/models/litellm/__init__.py +14 -0
- agno/models/litellm/chat.py +503 -0
- agno/models/litellm/litellm_openai.py +42 -0
- agno/models/llama_cpp/__init__.py +5 -0
- agno/models/llama_cpp/llama_cpp.py +22 -0
- agno/models/lmstudio/__init__.py +5 -0
- agno/models/lmstudio/lmstudio.py +25 -0
- agno/models/message.py +361 -39
- agno/models/meta/__init__.py +12 -0
- agno/models/meta/llama.py +502 -0
- agno/models/meta/llama_openai.py +79 -0
- agno/models/metrics.py +120 -0
- agno/models/mistral/__init__.py +4 -0
- agno/models/mistral/mistral.py +293 -393
- agno/models/nebius/__init__.py +3 -0
- agno/models/nebius/nebius.py +53 -0
- agno/models/nexus/__init__.py +3 -0
- agno/models/nexus/nexus.py +22 -0
- agno/models/nvidia/__init__.py +4 -0
- agno/models/nvidia/nvidia.py +22 -3
- agno/models/ollama/__init__.py +4 -2
- agno/models/ollama/chat.py +257 -492
- agno/models/openai/__init__.py +7 -0
- agno/models/openai/chat.py +725 -770
- agno/models/openai/like.py +16 -2
- agno/models/openai/responses.py +1121 -0
- agno/models/openrouter/__init__.py +4 -0
- agno/models/openrouter/openrouter.py +62 -5
- agno/models/perplexity/__init__.py +5 -0
- agno/models/perplexity/perplexity.py +203 -0
- agno/models/portkey/__init__.py +3 -0
- agno/models/portkey/portkey.py +82 -0
- agno/models/requesty/__init__.py +5 -0
- agno/models/requesty/requesty.py +69 -0
- agno/models/response.py +177 -7
- agno/models/sambanova/__init__.py +4 -0
- agno/models/sambanova/sambanova.py +23 -4
- agno/models/siliconflow/__init__.py +5 -0
- agno/models/siliconflow/siliconflow.py +42 -0
- agno/models/together/__init__.py +4 -0
- agno/models/together/together.py +21 -164
- agno/models/utils.py +266 -0
- agno/models/vercel/__init__.py +3 -0
- agno/models/vercel/v0.py +43 -0
- agno/models/vertexai/__init__.py +0 -1
- agno/models/vertexai/claude.py +190 -0
- agno/models/vllm/__init__.py +3 -0
- agno/models/vllm/vllm.py +83 -0
- agno/models/xai/__init__.py +2 -0
- agno/models/xai/xai.py +111 -7
- agno/os/__init__.py +3 -0
- agno/os/app.py +1027 -0
- agno/os/auth.py +244 -0
- agno/os/config.py +126 -0
- agno/os/interfaces/__init__.py +1 -0
- agno/os/interfaces/a2a/__init__.py +3 -0
- agno/os/interfaces/a2a/a2a.py +42 -0
- agno/os/interfaces/a2a/router.py +249 -0
- agno/os/interfaces/a2a/utils.py +924 -0
- agno/os/interfaces/agui/__init__.py +3 -0
- agno/os/interfaces/agui/agui.py +47 -0
- agno/os/interfaces/agui/router.py +147 -0
- agno/os/interfaces/agui/utils.py +574 -0
- agno/os/interfaces/base.py +25 -0
- agno/os/interfaces/slack/__init__.py +3 -0
- agno/os/interfaces/slack/router.py +148 -0
- agno/os/interfaces/slack/security.py +30 -0
- agno/os/interfaces/slack/slack.py +47 -0
- agno/os/interfaces/whatsapp/__init__.py +3 -0
- agno/os/interfaces/whatsapp/router.py +210 -0
- agno/os/interfaces/whatsapp/security.py +55 -0
- agno/os/interfaces/whatsapp/whatsapp.py +36 -0
- agno/os/mcp.py +293 -0
- agno/os/middleware/__init__.py +9 -0
- agno/os/middleware/jwt.py +797 -0
- agno/os/router.py +258 -0
- agno/os/routers/__init__.py +3 -0
- agno/os/routers/agents/__init__.py +3 -0
- agno/os/routers/agents/router.py +599 -0
- agno/os/routers/agents/schema.py +261 -0
- agno/os/routers/evals/__init__.py +3 -0
- agno/os/routers/evals/evals.py +450 -0
- agno/os/routers/evals/schemas.py +174 -0
- agno/os/routers/evals/utils.py +231 -0
- agno/os/routers/health.py +31 -0
- agno/os/routers/home.py +52 -0
- agno/os/routers/knowledge/__init__.py +3 -0
- agno/os/routers/knowledge/knowledge.py +1008 -0
- agno/os/routers/knowledge/schemas.py +178 -0
- agno/os/routers/memory/__init__.py +3 -0
- agno/os/routers/memory/memory.py +661 -0
- agno/os/routers/memory/schemas.py +88 -0
- agno/os/routers/metrics/__init__.py +3 -0
- agno/os/routers/metrics/metrics.py +190 -0
- agno/os/routers/metrics/schemas.py +47 -0
- agno/os/routers/session/__init__.py +3 -0
- agno/os/routers/session/session.py +997 -0
- agno/os/routers/teams/__init__.py +3 -0
- agno/os/routers/teams/router.py +512 -0
- agno/os/routers/teams/schema.py +257 -0
- agno/os/routers/traces/__init__.py +3 -0
- agno/os/routers/traces/schemas.py +414 -0
- agno/os/routers/traces/traces.py +499 -0
- agno/os/routers/workflows/__init__.py +3 -0
- agno/os/routers/workflows/router.py +624 -0
- agno/os/routers/workflows/schema.py +75 -0
- agno/os/schema.py +534 -0
- agno/os/scopes.py +469 -0
- agno/{playground → os}/settings.py +7 -15
- agno/os/utils.py +973 -0
- agno/reasoning/anthropic.py +80 -0
- agno/reasoning/azure_ai_foundry.py +67 -0
- agno/reasoning/deepseek.py +63 -0
- agno/reasoning/default.py +97 -0
- agno/reasoning/gemini.py +73 -0
- agno/reasoning/groq.py +71 -0
- agno/reasoning/helpers.py +24 -1
- agno/reasoning/ollama.py +67 -0
- agno/reasoning/openai.py +86 -0
- agno/reasoning/step.py +2 -1
- agno/reasoning/vertexai.py +76 -0
- agno/run/__init__.py +6 -0
- agno/run/agent.py +822 -0
- agno/run/base.py +247 -0
- agno/run/cancel.py +81 -0
- agno/run/requirement.py +181 -0
- agno/run/team.py +767 -0
- agno/run/workflow.py +708 -0
- agno/session/__init__.py +10 -0
- agno/session/agent.py +260 -0
- agno/session/summary.py +265 -0
- agno/session/team.py +342 -0
- agno/session/workflow.py +501 -0
- agno/table.py +10 -0
- agno/team/__init__.py +37 -0
- agno/team/team.py +9536 -0
- agno/tools/__init__.py +7 -0
- agno/tools/agentql.py +120 -0
- agno/tools/airflow.py +22 -12
- agno/tools/api.py +122 -0
- agno/tools/apify.py +276 -83
- agno/tools/{arxiv_toolkit.py → arxiv.py} +20 -12
- agno/tools/aws_lambda.py +28 -7
- agno/tools/aws_ses.py +66 -0
- agno/tools/baidusearch.py +11 -4
- agno/tools/bitbucket.py +292 -0
- agno/tools/brandfetch.py +213 -0
- agno/tools/bravesearch.py +106 -0
- agno/tools/brightdata.py +367 -0
- agno/tools/browserbase.py +209 -0
- agno/tools/calcom.py +32 -23
- agno/tools/calculator.py +24 -37
- agno/tools/cartesia.py +187 -0
- agno/tools/{clickup_tool.py → clickup.py} +17 -28
- agno/tools/confluence.py +91 -26
- agno/tools/crawl4ai.py +139 -43
- agno/tools/csv_toolkit.py +28 -22
- agno/tools/dalle.py +36 -22
- agno/tools/daytona.py +475 -0
- agno/tools/decorator.py +169 -14
- agno/tools/desi_vocal.py +23 -11
- agno/tools/discord.py +32 -29
- agno/tools/docker.py +716 -0
- agno/tools/duckdb.py +76 -81
- agno/tools/duckduckgo.py +43 -40
- agno/tools/e2b.py +703 -0
- agno/tools/eleven_labs.py +65 -54
- agno/tools/email.py +13 -5
- agno/tools/evm.py +129 -0
- agno/tools/exa.py +324 -42
- agno/tools/fal.py +39 -35
- agno/tools/file.py +196 -30
- agno/tools/file_generation.py +356 -0
- agno/tools/financial_datasets.py +288 -0
- agno/tools/firecrawl.py +108 -33
- agno/tools/function.py +960 -122
- agno/tools/giphy.py +34 -12
- agno/tools/github.py +1294 -97
- agno/tools/gmail.py +922 -0
- agno/tools/google_bigquery.py +117 -0
- agno/tools/google_drive.py +271 -0
- agno/tools/google_maps.py +253 -0
- agno/tools/googlecalendar.py +607 -107
- agno/tools/googlesheets.py +377 -0
- agno/tools/hackernews.py +20 -12
- agno/tools/jina.py +24 -14
- agno/tools/jira.py +48 -19
- agno/tools/knowledge.py +218 -0
- agno/tools/linear.py +82 -43
- agno/tools/linkup.py +58 -0
- agno/tools/local_file_system.py +15 -7
- agno/tools/lumalab.py +41 -26
- agno/tools/mcp/__init__.py +10 -0
- agno/tools/mcp/mcp.py +331 -0
- agno/tools/mcp/multi_mcp.py +347 -0
- agno/tools/mcp/params.py +24 -0
- agno/tools/mcp_toolbox.py +284 -0
- agno/tools/mem0.py +193 -0
- agno/tools/memory.py +419 -0
- agno/tools/mlx_transcribe.py +11 -9
- agno/tools/models/azure_openai.py +190 -0
- agno/tools/models/gemini.py +203 -0
- agno/tools/models/groq.py +158 -0
- agno/tools/models/morph.py +186 -0
- agno/tools/models/nebius.py +124 -0
- agno/tools/models_labs.py +163 -82
- agno/tools/moviepy_video.py +18 -13
- agno/tools/nano_banana.py +151 -0
- agno/tools/neo4j.py +134 -0
- agno/tools/newspaper.py +15 -4
- agno/tools/newspaper4k.py +19 -6
- agno/tools/notion.py +204 -0
- agno/tools/openai.py +181 -17
- agno/tools/openbb.py +27 -20
- agno/tools/opencv.py +321 -0
- agno/tools/openweather.py +233 -0
- agno/tools/oxylabs.py +385 -0
- agno/tools/pandas.py +25 -15
- agno/tools/parallel.py +314 -0
- agno/tools/postgres.py +238 -185
- agno/tools/pubmed.py +125 -13
- agno/tools/python.py +48 -35
- agno/tools/reasoning.py +283 -0
- agno/tools/reddit.py +207 -29
- agno/tools/redshift.py +406 -0
- agno/tools/replicate.py +69 -26
- agno/tools/resend.py +11 -6
- agno/tools/scrapegraph.py +179 -19
- agno/tools/searxng.py +23 -31
- agno/tools/serpapi.py +15 -10
- agno/tools/serper.py +255 -0
- agno/tools/shell.py +23 -12
- agno/tools/shopify.py +1519 -0
- agno/tools/slack.py +56 -14
- agno/tools/sleep.py +8 -6
- agno/tools/spider.py +35 -11
- agno/tools/spotify.py +919 -0
- agno/tools/sql.py +34 -19
- agno/tools/tavily.py +158 -8
- agno/tools/telegram.py +18 -8
- agno/tools/todoist.py +218 -0
- agno/tools/toolkit.py +134 -9
- agno/tools/trafilatura.py +388 -0
- agno/tools/trello.py +25 -28
- agno/tools/twilio.py +18 -9
- agno/tools/user_control_flow.py +78 -0
- agno/tools/valyu.py +228 -0
- agno/tools/visualization.py +467 -0
- agno/tools/webbrowser.py +28 -0
- agno/tools/webex.py +76 -0
- agno/tools/website.py +23 -19
- agno/tools/webtools.py +45 -0
- agno/tools/whatsapp.py +286 -0
- agno/tools/wikipedia.py +28 -19
- agno/tools/workflow.py +285 -0
- agno/tools/{twitter.py → x.py} +142 -46
- agno/tools/yfinance.py +41 -39
- agno/tools/youtube.py +34 -17
- agno/tools/zendesk.py +15 -5
- agno/tools/zep.py +454 -0
- agno/tools/zoom.py +86 -37
- agno/tracing/__init__.py +12 -0
- agno/tracing/exporter.py +157 -0
- agno/tracing/schemas.py +276 -0
- agno/tracing/setup.py +111 -0
- agno/utils/agent.py +938 -0
- agno/utils/audio.py +37 -1
- agno/utils/certs.py +27 -0
- agno/utils/code_execution.py +11 -0
- agno/utils/common.py +103 -20
- agno/utils/cryptography.py +22 -0
- agno/utils/dttm.py +33 -0
- agno/utils/events.py +700 -0
- agno/utils/functions.py +107 -37
- agno/utils/gemini.py +426 -0
- agno/utils/hooks.py +171 -0
- agno/utils/http.py +185 -0
- agno/utils/json_schema.py +159 -37
- agno/utils/knowledge.py +36 -0
- agno/utils/location.py +19 -0
- agno/utils/log.py +221 -8
- agno/utils/mcp.py +214 -0
- agno/utils/media.py +335 -14
- agno/utils/merge_dict.py +22 -1
- agno/utils/message.py +77 -2
- agno/utils/models/ai_foundry.py +50 -0
- agno/utils/models/claude.py +373 -0
- agno/utils/models/cohere.py +94 -0
- agno/utils/models/llama.py +85 -0
- agno/utils/models/mistral.py +100 -0
- agno/utils/models/openai_responses.py +140 -0
- agno/utils/models/schema_utils.py +153 -0
- agno/utils/models/watsonx.py +41 -0
- agno/utils/openai.py +257 -0
- agno/utils/pickle.py +1 -1
- agno/utils/pprint.py +124 -8
- agno/utils/print_response/agent.py +930 -0
- agno/utils/print_response/team.py +1914 -0
- agno/utils/print_response/workflow.py +1668 -0
- agno/utils/prompts.py +111 -0
- agno/utils/reasoning.py +108 -0
- agno/utils/response.py +163 -0
- agno/utils/serialize.py +32 -0
- agno/utils/shell.py +4 -4
- agno/utils/streamlit.py +487 -0
- agno/utils/string.py +204 -51
- agno/utils/team.py +139 -0
- agno/utils/timer.py +9 -2
- agno/utils/tokens.py +657 -0
- agno/utils/tools.py +19 -1
- agno/utils/whatsapp.py +305 -0
- agno/utils/yaml_io.py +3 -3
- agno/vectordb/__init__.py +2 -0
- agno/vectordb/base.py +87 -9
- agno/vectordb/cassandra/__init__.py +5 -1
- agno/vectordb/cassandra/cassandra.py +383 -27
- agno/vectordb/chroma/__init__.py +4 -0
- agno/vectordb/chroma/chromadb.py +748 -83
- agno/vectordb/clickhouse/__init__.py +7 -1
- agno/vectordb/clickhouse/clickhousedb.py +554 -53
- agno/vectordb/couchbase/__init__.py +3 -0
- agno/vectordb/couchbase/couchbase.py +1446 -0
- agno/vectordb/lancedb/__init__.py +5 -0
- agno/vectordb/lancedb/lance_db.py +730 -98
- agno/vectordb/langchaindb/__init__.py +5 -0
- agno/vectordb/langchaindb/langchaindb.py +163 -0
- agno/vectordb/lightrag/__init__.py +5 -0
- agno/vectordb/lightrag/lightrag.py +388 -0
- agno/vectordb/llamaindex/__init__.py +3 -0
- agno/vectordb/llamaindex/llamaindexdb.py +166 -0
- agno/vectordb/milvus/__init__.py +3 -0
- agno/vectordb/milvus/milvus.py +966 -78
- agno/vectordb/mongodb/__init__.py +9 -1
- agno/vectordb/mongodb/mongodb.py +1175 -172
- agno/vectordb/pgvector/__init__.py +8 -0
- agno/vectordb/pgvector/pgvector.py +599 -115
- agno/vectordb/pineconedb/__init__.py +5 -1
- agno/vectordb/pineconedb/pineconedb.py +406 -43
- agno/vectordb/qdrant/__init__.py +4 -0
- agno/vectordb/qdrant/qdrant.py +914 -61
- agno/vectordb/redis/__init__.py +9 -0
- agno/vectordb/redis/redisdb.py +682 -0
- agno/vectordb/singlestore/__init__.py +8 -1
- agno/vectordb/singlestore/singlestore.py +771 -0
- agno/vectordb/surrealdb/__init__.py +3 -0
- agno/vectordb/surrealdb/surrealdb.py +663 -0
- agno/vectordb/upstashdb/__init__.py +5 -0
- agno/vectordb/upstashdb/upstashdb.py +718 -0
- agno/vectordb/weaviate/__init__.py +8 -0
- agno/vectordb/weaviate/index.py +15 -0
- agno/vectordb/weaviate/weaviate.py +1009 -0
- agno/workflow/__init__.py +23 -1
- agno/workflow/agent.py +299 -0
- agno/workflow/condition.py +759 -0
- agno/workflow/loop.py +756 -0
- agno/workflow/parallel.py +853 -0
- agno/workflow/router.py +723 -0
- agno/workflow/step.py +1564 -0
- agno/workflow/steps.py +613 -0
- agno/workflow/types.py +556 -0
- agno/workflow/workflow.py +4327 -514
- agno-2.3.13.dist-info/METADATA +639 -0
- agno-2.3.13.dist-info/RECORD +613 -0
- {agno-0.1.2.dist-info → agno-2.3.13.dist-info}/WHEEL +1 -1
- agno-2.3.13.dist-info/licenses/LICENSE +201 -0
- agno/api/playground.py +0 -91
- agno/api/schemas/playground.py +0 -22
- agno/api/schemas/user.py +0 -22
- agno/api/schemas/workspace.py +0 -46
- agno/api/user.py +0 -160
- agno/api/workspace.py +0 -151
- agno/cli/auth_server.py +0 -118
- agno/cli/config.py +0 -275
- agno/cli/console.py +0 -88
- agno/cli/credentials.py +0 -23
- agno/cli/entrypoint.py +0 -571
- agno/cli/operator.py +0 -355
- agno/cli/settings.py +0 -85
- agno/cli/ws/ws_cli.py +0 -817
- agno/constants.py +0 -13
- agno/document/__init__.py +0 -1
- agno/document/chunking/semantic.py +0 -47
- agno/document/chunking/strategy.py +0 -31
- agno/document/reader/__init__.py +0 -1
- agno/document/reader/arxiv_reader.py +0 -41
- agno/document/reader/base.py +0 -22
- agno/document/reader/csv_reader.py +0 -84
- agno/document/reader/docx_reader.py +0 -46
- agno/document/reader/firecrawl_reader.py +0 -99
- agno/document/reader/json_reader.py +0 -43
- agno/document/reader/pdf_reader.py +0 -219
- agno/document/reader/s3/pdf_reader.py +0 -46
- agno/document/reader/s3/text_reader.py +0 -51
- agno/document/reader/text_reader.py +0 -41
- agno/document/reader/website_reader.py +0 -175
- agno/document/reader/youtube_reader.py +0 -50
- agno/embedder/__init__.py +0 -1
- agno/embedder/azure_openai.py +0 -86
- agno/embedder/cohere.py +0 -72
- agno/embedder/fastembed.py +0 -37
- agno/embedder/google.py +0 -73
- agno/embedder/huggingface.py +0 -54
- agno/embedder/mistral.py +0 -80
- agno/embedder/ollama.py +0 -57
- agno/embedder/openai.py +0 -74
- agno/embedder/sentence_transformer.py +0 -38
- agno/embedder/voyageai.py +0 -64
- agno/eval/perf.py +0 -201
- agno/file/__init__.py +0 -1
- agno/file/file.py +0 -16
- agno/file/local/csv.py +0 -32
- agno/file/local/txt.py +0 -19
- agno/infra/app.py +0 -240
- agno/infra/base.py +0 -144
- agno/infra/context.py +0 -20
- agno/infra/db_app.py +0 -52
- agno/infra/resource.py +0 -205
- agno/infra/resources.py +0 -55
- agno/knowledge/agent.py +0 -230
- agno/knowledge/arxiv.py +0 -22
- agno/knowledge/combined.py +0 -22
- agno/knowledge/csv.py +0 -28
- agno/knowledge/csv_url.py +0 -19
- agno/knowledge/document.py +0 -20
- agno/knowledge/docx.py +0 -30
- agno/knowledge/json.py +0 -28
- agno/knowledge/langchain.py +0 -71
- agno/knowledge/llamaindex.py +0 -66
- agno/knowledge/pdf.py +0 -28
- agno/knowledge/pdf_url.py +0 -26
- agno/knowledge/s3/base.py +0 -60
- agno/knowledge/s3/pdf.py +0 -21
- agno/knowledge/s3/text.py +0 -23
- agno/knowledge/text.py +0 -30
- agno/knowledge/website.py +0 -88
- agno/knowledge/wikipedia.py +0 -31
- agno/knowledge/youtube.py +0 -22
- agno/memory/agent.py +0 -392
- agno/memory/classifier.py +0 -104
- agno/memory/db/__init__.py +0 -1
- agno/memory/db/base.py +0 -42
- agno/memory/db/mongodb.py +0 -189
- agno/memory/db/postgres.py +0 -203
- agno/memory/db/sqlite.py +0 -193
- agno/memory/memory.py +0 -15
- agno/memory/row.py +0 -36
- agno/memory/summarizer.py +0 -192
- agno/memory/summary.py +0 -19
- agno/memory/workflow.py +0 -38
- agno/models/google/gemini_openai.py +0 -26
- agno/models/ollama/hermes.py +0 -221
- agno/models/ollama/tools.py +0 -362
- agno/models/vertexai/gemini.py +0 -595
- agno/playground/__init__.py +0 -3
- agno/playground/async_router.py +0 -421
- agno/playground/deploy.py +0 -249
- agno/playground/operator.py +0 -92
- agno/playground/playground.py +0 -91
- agno/playground/schemas.py +0 -76
- agno/playground/serve.py +0 -55
- agno/playground/sync_router.py +0 -405
- agno/reasoning/agent.py +0 -68
- agno/run/response.py +0 -112
- agno/storage/agent/__init__.py +0 -0
- agno/storage/agent/base.py +0 -38
- agno/storage/agent/dynamodb.py +0 -350
- agno/storage/agent/json.py +0 -92
- agno/storage/agent/mongodb.py +0 -228
- agno/storage/agent/postgres.py +0 -367
- agno/storage/agent/session.py +0 -79
- agno/storage/agent/singlestore.py +0 -303
- agno/storage/agent/sqlite.py +0 -357
- agno/storage/agent/yaml.py +0 -93
- agno/storage/workflow/__init__.py +0 -0
- agno/storage/workflow/base.py +0 -40
- agno/storage/workflow/mongodb.py +0 -233
- agno/storage/workflow/postgres.py +0 -366
- agno/storage/workflow/session.py +0 -60
- agno/storage/workflow/sqlite.py +0 -359
- agno/tools/googlesearch.py +0 -88
- agno/utils/defaults.py +0 -57
- agno/utils/filesystem.py +0 -39
- agno/utils/git.py +0 -52
- agno/utils/json_io.py +0 -30
- agno/utils/load_env.py +0 -19
- agno/utils/py_io.py +0 -19
- agno/utils/pyproject.py +0 -18
- agno/utils/resource_filter.py +0 -31
- agno/vectordb/singlestore/s2vectordb.py +0 -390
- agno/vectordb/singlestore/s2vectordb2.py +0 -355
- agno/workspace/__init__.py +0 -0
- agno/workspace/config.py +0 -325
- agno/workspace/enums.py +0 -6
- agno/workspace/helpers.py +0 -48
- agno/workspace/operator.py +0 -758
- agno/workspace/settings.py +0 -63
- agno-0.1.2.dist-info/LICENSE +0 -375
- agno-0.1.2.dist-info/METADATA +0 -502
- agno-0.1.2.dist-info/RECORD +0 -352
- agno-0.1.2.dist-info/entry_points.txt +0 -3
- /agno/{cli → db/migrations}/__init__.py +0 -0
- /agno/{cli/ws → db/migrations/versions}/__init__.py +0 -0
- /agno/{document/chunking/__init__.py → db/schemas/metrics.py} +0 -0
- /agno/{document/reader/s3 → integrations}/__init__.py +0 -0
- /agno/{file/local → knowledge/chunking}/__init__.py +0 -0
- /agno/{infra → knowledge/remote_content}/__init__.py +0 -0
- /agno/{knowledge/s3 → tools/models}/__init__.py +0 -0
- /agno/{reranker → utils/models}/__init__.py +0 -0
- /agno/{storage → utils/print_response}/__init__.py +0 -0
- {agno-0.1.2.dist-info → agno-2.3.13.dist-info}/top_level.txt +0 -0
agno/db/sqlite/sqlite.py
ADDED
|
@@ -0,0 +1,2908 @@
|
|
|
1
|
+
import time
|
|
2
|
+
from datetime import date, datetime, timedelta, timezone
|
|
3
|
+
from pathlib import Path
|
|
4
|
+
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple, Union, cast
|
|
5
|
+
from uuid import uuid4
|
|
6
|
+
|
|
7
|
+
if TYPE_CHECKING:
|
|
8
|
+
from agno.tracing.schemas import Span, Trace
|
|
9
|
+
|
|
10
|
+
from agno.db.base import BaseDb, SessionType
|
|
11
|
+
from agno.db.migrations.manager import MigrationManager
|
|
12
|
+
from agno.db.schemas.culture import CulturalKnowledge
|
|
13
|
+
from agno.db.schemas.evals import EvalFilterType, EvalRunRecord, EvalType
|
|
14
|
+
from agno.db.schemas.knowledge import KnowledgeRow
|
|
15
|
+
from agno.db.schemas.memory import UserMemory
|
|
16
|
+
from agno.db.sqlite.schemas import get_table_schema_definition
|
|
17
|
+
from agno.db.sqlite.utils import (
|
|
18
|
+
apply_sorting,
|
|
19
|
+
bulk_upsert_metrics,
|
|
20
|
+
calculate_date_metrics,
|
|
21
|
+
deserialize_cultural_knowledge_from_db,
|
|
22
|
+
fetch_all_sessions_data,
|
|
23
|
+
get_dates_to_calculate_metrics_for,
|
|
24
|
+
is_table_available,
|
|
25
|
+
is_valid_table,
|
|
26
|
+
serialize_cultural_knowledge_for_db,
|
|
27
|
+
)
|
|
28
|
+
from agno.db.utils import deserialize_session_json_fields, serialize_session_json_fields
|
|
29
|
+
from agno.session import AgentSession, Session, TeamSession, WorkflowSession
|
|
30
|
+
from agno.utils.log import log_debug, log_error, log_info, log_warning
|
|
31
|
+
from agno.utils.string import generate_id
|
|
32
|
+
|
|
33
|
+
try:
|
|
34
|
+
from sqlalchemy import Column, MetaData, String, Table, func, select, text
|
|
35
|
+
from sqlalchemy.dialects import sqlite
|
|
36
|
+
from sqlalchemy.engine import Engine, create_engine
|
|
37
|
+
from sqlalchemy.orm import scoped_session, sessionmaker
|
|
38
|
+
from sqlalchemy.schema import ForeignKey, Index, UniqueConstraint
|
|
39
|
+
except ImportError:
|
|
40
|
+
raise ImportError("`sqlalchemy` not installed. Please install it using `pip install sqlalchemy`")
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
class SqliteDb(BaseDb):
|
|
44
|
+
def __init__(
|
|
45
|
+
self,
|
|
46
|
+
db_file: Optional[str] = None,
|
|
47
|
+
db_engine: Optional[Engine] = None,
|
|
48
|
+
db_url: Optional[str] = None,
|
|
49
|
+
session_table: Optional[str] = None,
|
|
50
|
+
culture_table: Optional[str] = None,
|
|
51
|
+
memory_table: Optional[str] = None,
|
|
52
|
+
metrics_table: Optional[str] = None,
|
|
53
|
+
eval_table: Optional[str] = None,
|
|
54
|
+
knowledge_table: Optional[str] = None,
|
|
55
|
+
traces_table: Optional[str] = None,
|
|
56
|
+
spans_table: Optional[str] = None,
|
|
57
|
+
versions_table: Optional[str] = None,
|
|
58
|
+
id: Optional[str] = None,
|
|
59
|
+
):
|
|
60
|
+
"""
|
|
61
|
+
Interface for interacting with a SQLite database.
|
|
62
|
+
|
|
63
|
+
The following order is used to determine the database connection:
|
|
64
|
+
1. Use the db_engine
|
|
65
|
+
2. Use the db_url
|
|
66
|
+
3. Use the db_file
|
|
67
|
+
4. Create a new database in the current directory
|
|
68
|
+
|
|
69
|
+
Args:
|
|
70
|
+
db_file (Optional[str]): The database file to connect to.
|
|
71
|
+
db_engine (Optional[Engine]): The SQLAlchemy database engine to use.
|
|
72
|
+
db_url (Optional[str]): The database URL to connect to.
|
|
73
|
+
session_table (Optional[str]): Name of the table to store Agent, Team and Workflow sessions.
|
|
74
|
+
culture_table (Optional[str]): Name of the table to store cultural notions.
|
|
75
|
+
memory_table (Optional[str]): Name of the table to store user memories.
|
|
76
|
+
metrics_table (Optional[str]): Name of the table to store metrics.
|
|
77
|
+
eval_table (Optional[str]): Name of the table to store evaluation runs data.
|
|
78
|
+
knowledge_table (Optional[str]): Name of the table to store knowledge documents data.
|
|
79
|
+
traces_table (Optional[str]): Name of the table to store run traces.
|
|
80
|
+
spans_table (Optional[str]): Name of the table to store span events.
|
|
81
|
+
versions_table (Optional[str]): Name of the table to store schema versions.
|
|
82
|
+
id (Optional[str]): ID of the database.
|
|
83
|
+
|
|
84
|
+
Raises:
|
|
85
|
+
ValueError: If none of the tables are provided.
|
|
86
|
+
"""
|
|
87
|
+
if id is None:
|
|
88
|
+
seed = db_url or db_file or str(db_engine.url) if db_engine else "sqlite:///agno.db"
|
|
89
|
+
id = generate_id(seed)
|
|
90
|
+
|
|
91
|
+
super().__init__(
|
|
92
|
+
id=id,
|
|
93
|
+
session_table=session_table,
|
|
94
|
+
culture_table=culture_table,
|
|
95
|
+
memory_table=memory_table,
|
|
96
|
+
metrics_table=metrics_table,
|
|
97
|
+
eval_table=eval_table,
|
|
98
|
+
knowledge_table=knowledge_table,
|
|
99
|
+
traces_table=traces_table,
|
|
100
|
+
spans_table=spans_table,
|
|
101
|
+
versions_table=versions_table,
|
|
102
|
+
)
|
|
103
|
+
|
|
104
|
+
_engine: Optional[Engine] = db_engine
|
|
105
|
+
if _engine is None:
|
|
106
|
+
if db_url is not None:
|
|
107
|
+
_engine = create_engine(db_url)
|
|
108
|
+
elif db_file is not None:
|
|
109
|
+
db_path = Path(db_file).resolve()
|
|
110
|
+
db_path.parent.mkdir(parents=True, exist_ok=True)
|
|
111
|
+
db_file = str(db_path)
|
|
112
|
+
_engine = create_engine(f"sqlite:///{db_path}")
|
|
113
|
+
else:
|
|
114
|
+
# If none of db_engine, db_url, or db_file are provided, create a db in the current directory
|
|
115
|
+
default_db_path = Path("./agno.db").resolve()
|
|
116
|
+
_engine = create_engine(f"sqlite:///{default_db_path}")
|
|
117
|
+
db_file = str(default_db_path)
|
|
118
|
+
log_debug(f"Created SQLite database: {default_db_path}")
|
|
119
|
+
|
|
120
|
+
self.db_engine: Engine = _engine
|
|
121
|
+
self.db_url: Optional[str] = db_url
|
|
122
|
+
self.db_file: Optional[str] = db_file
|
|
123
|
+
self.metadata: MetaData = MetaData()
|
|
124
|
+
|
|
125
|
+
# Initialize database session
|
|
126
|
+
self.Session: scoped_session = scoped_session(sessionmaker(bind=self.db_engine))
|
|
127
|
+
|
|
128
|
+
# -- DB methods --
|
|
129
|
+
def table_exists(self, table_name: str) -> bool:
|
|
130
|
+
"""Check if a table with the given name exists in the SQLite database.
|
|
131
|
+
|
|
132
|
+
Args:
|
|
133
|
+
table_name: Name of the table to check
|
|
134
|
+
|
|
135
|
+
Returns:
|
|
136
|
+
bool: True if the table exists in the database, False otherwise
|
|
137
|
+
"""
|
|
138
|
+
with self.Session() as sess:
|
|
139
|
+
return is_table_available(session=sess, table_name=table_name)
|
|
140
|
+
|
|
141
|
+
def _create_all_tables(self):
|
|
142
|
+
"""Create all tables for the database."""
|
|
143
|
+
tables_to_create = [
|
|
144
|
+
(self.session_table_name, "sessions"),
|
|
145
|
+
(self.memory_table_name, "memories"),
|
|
146
|
+
(self.metrics_table_name, "metrics"),
|
|
147
|
+
(self.eval_table_name, "evals"),
|
|
148
|
+
(self.knowledge_table_name, "knowledge"),
|
|
149
|
+
(self.versions_table_name, "versions"),
|
|
150
|
+
]
|
|
151
|
+
|
|
152
|
+
for table_name, table_type in tables_to_create:
|
|
153
|
+
self._get_or_create_table(table_name=table_name, table_type=table_type, create_table_if_not_found=True)
|
|
154
|
+
|
|
155
|
+
def _create_table(self, table_name: str, table_type: str) -> Table:
|
|
156
|
+
"""
|
|
157
|
+
Create a table with the appropriate schema based on the table type.
|
|
158
|
+
|
|
159
|
+
Args:
|
|
160
|
+
table_name (str): Name of the table to create
|
|
161
|
+
table_type (str): Type of table (used to get schema definition)
|
|
162
|
+
|
|
163
|
+
Returns:
|
|
164
|
+
Table: SQLAlchemy Table object
|
|
165
|
+
"""
|
|
166
|
+
try:
|
|
167
|
+
table_schema = get_table_schema_definition(table_type).copy()
|
|
168
|
+
|
|
169
|
+
columns: List[Column] = []
|
|
170
|
+
indexes: List[str] = []
|
|
171
|
+
unique_constraints: List[str] = []
|
|
172
|
+
schema_unique_constraints = table_schema.pop("_unique_constraints", [])
|
|
173
|
+
|
|
174
|
+
# Get the columns, indexes, and unique constraints from the table schema
|
|
175
|
+
for col_name, col_config in table_schema.items():
|
|
176
|
+
column_args = [col_name, col_config["type"]()]
|
|
177
|
+
column_kwargs = {}
|
|
178
|
+
|
|
179
|
+
if col_config.get("primary_key", False):
|
|
180
|
+
column_kwargs["primary_key"] = True
|
|
181
|
+
if "nullable" in col_config:
|
|
182
|
+
column_kwargs["nullable"] = col_config["nullable"]
|
|
183
|
+
if col_config.get("index", False):
|
|
184
|
+
indexes.append(col_name)
|
|
185
|
+
if col_config.get("unique", False):
|
|
186
|
+
column_kwargs["unique"] = True
|
|
187
|
+
unique_constraints.append(col_name)
|
|
188
|
+
|
|
189
|
+
# Handle foreign key constraint
|
|
190
|
+
if "foreign_key" in col_config:
|
|
191
|
+
fk_ref = col_config["foreign_key"]
|
|
192
|
+
# For spans table, dynamically replace the traces table reference
|
|
193
|
+
# with the actual trace table name configured for this db instance
|
|
194
|
+
if table_type == "spans" and "trace_id" in fk_ref:
|
|
195
|
+
fk_ref = f"{self.trace_table_name}.trace_id"
|
|
196
|
+
column_args.append(ForeignKey(fk_ref))
|
|
197
|
+
|
|
198
|
+
columns.append(Column(*column_args, **column_kwargs)) # type: ignore
|
|
199
|
+
|
|
200
|
+
# Create the table object
|
|
201
|
+
table = Table(table_name, self.metadata, *columns)
|
|
202
|
+
|
|
203
|
+
# Add multi-column unique constraints with table-specific names
|
|
204
|
+
for constraint in schema_unique_constraints:
|
|
205
|
+
constraint_name = f"{table_name}_{constraint['name']}"
|
|
206
|
+
constraint_columns = constraint["columns"]
|
|
207
|
+
table.append_constraint(UniqueConstraint(*constraint_columns, name=constraint_name))
|
|
208
|
+
|
|
209
|
+
# Add indexes to the table definition
|
|
210
|
+
for idx_col in indexes:
|
|
211
|
+
idx_name = f"idx_{table_name}_{idx_col}"
|
|
212
|
+
table.append_constraint(Index(idx_name, idx_col))
|
|
213
|
+
|
|
214
|
+
# Create table
|
|
215
|
+
table_created = False
|
|
216
|
+
if not self.table_exists(table_name):
|
|
217
|
+
table.create(self.db_engine, checkfirst=True)
|
|
218
|
+
log_debug(f"Successfully created table '{table_name}'")
|
|
219
|
+
table_created = True
|
|
220
|
+
else:
|
|
221
|
+
log_debug(f"Table '{table_name}' already exists, skipping creation")
|
|
222
|
+
|
|
223
|
+
# Create indexes
|
|
224
|
+
for idx in table.indexes:
|
|
225
|
+
try:
|
|
226
|
+
# Check if index already exists
|
|
227
|
+
with self.Session() as sess:
|
|
228
|
+
exists_query = text("SELECT 1 FROM sqlite_master WHERE type = 'index' AND name = :index_name")
|
|
229
|
+
exists = sess.execute(exists_query, {"index_name": idx.name}).scalar() is not None
|
|
230
|
+
if exists:
|
|
231
|
+
log_debug(f"Index {idx.name} already exists in table {table_name}, skipping creation")
|
|
232
|
+
continue
|
|
233
|
+
|
|
234
|
+
idx.create(self.db_engine)
|
|
235
|
+
|
|
236
|
+
log_debug(f"Created index: {idx.name} for table {table_name}")
|
|
237
|
+
except Exception as e:
|
|
238
|
+
log_warning(f"Error creating index {idx.name}: {e}")
|
|
239
|
+
|
|
240
|
+
# Store the schema version for the created table
|
|
241
|
+
if table_name != self.versions_table_name and table_created:
|
|
242
|
+
latest_schema_version = MigrationManager(self).latest_schema_version
|
|
243
|
+
self.upsert_schema_version(table_name=table_name, version=latest_schema_version.public)
|
|
244
|
+
|
|
245
|
+
return table
|
|
246
|
+
|
|
247
|
+
except Exception as e:
|
|
248
|
+
from traceback import print_exc
|
|
249
|
+
|
|
250
|
+
print_exc()
|
|
251
|
+
log_error(f"Could not create table '{table_name}': {e}")
|
|
252
|
+
raise e
|
|
253
|
+
|
|
254
|
+
def _get_table(self, table_type: str, create_table_if_not_found: Optional[bool] = False) -> Optional[Table]:
|
|
255
|
+
if table_type == "sessions":
|
|
256
|
+
self.session_table = self._get_or_create_table(
|
|
257
|
+
table_name=self.session_table_name,
|
|
258
|
+
table_type=table_type,
|
|
259
|
+
create_table_if_not_found=create_table_if_not_found,
|
|
260
|
+
)
|
|
261
|
+
return self.session_table
|
|
262
|
+
|
|
263
|
+
elif table_type == "memories":
|
|
264
|
+
self.memory_table = self._get_or_create_table(
|
|
265
|
+
table_name=self.memory_table_name,
|
|
266
|
+
table_type="memories",
|
|
267
|
+
create_table_if_not_found=create_table_if_not_found,
|
|
268
|
+
)
|
|
269
|
+
return self.memory_table
|
|
270
|
+
|
|
271
|
+
elif table_type == "metrics":
|
|
272
|
+
self.metrics_table = self._get_or_create_table(
|
|
273
|
+
table_name=self.metrics_table_name,
|
|
274
|
+
table_type="metrics",
|
|
275
|
+
create_table_if_not_found=create_table_if_not_found,
|
|
276
|
+
)
|
|
277
|
+
return self.metrics_table
|
|
278
|
+
|
|
279
|
+
elif table_type == "evals":
|
|
280
|
+
self.eval_table = self._get_or_create_table(
|
|
281
|
+
table_name=self.eval_table_name,
|
|
282
|
+
table_type="evals",
|
|
283
|
+
create_table_if_not_found=create_table_if_not_found,
|
|
284
|
+
)
|
|
285
|
+
|
|
286
|
+
return self.eval_table
|
|
287
|
+
|
|
288
|
+
elif table_type == "knowledge":
|
|
289
|
+
self.knowledge_table = self._get_or_create_table(
|
|
290
|
+
table_name=self.knowledge_table_name,
|
|
291
|
+
table_type="knowledge",
|
|
292
|
+
create_table_if_not_found=create_table_if_not_found,
|
|
293
|
+
)
|
|
294
|
+
return self.knowledge_table
|
|
295
|
+
|
|
296
|
+
elif table_type == "traces":
|
|
297
|
+
self.traces_table = self._get_or_create_table(
|
|
298
|
+
table_name=self.trace_table_name,
|
|
299
|
+
table_type="traces",
|
|
300
|
+
create_table_if_not_found=create_table_if_not_found,
|
|
301
|
+
)
|
|
302
|
+
return self.traces_table
|
|
303
|
+
|
|
304
|
+
elif table_type == "spans":
|
|
305
|
+
# Ensure traces table exists first (spans has FK to traces)
|
|
306
|
+
if create_table_if_not_found:
|
|
307
|
+
self._get_table(table_type="traces", create_table_if_not_found=True)
|
|
308
|
+
|
|
309
|
+
self.spans_table = self._get_or_create_table(
|
|
310
|
+
table_name=self.span_table_name,
|
|
311
|
+
table_type="spans",
|
|
312
|
+
create_table_if_not_found=create_table_if_not_found,
|
|
313
|
+
)
|
|
314
|
+
return self.spans_table
|
|
315
|
+
|
|
316
|
+
elif table_type == "culture":
|
|
317
|
+
self.culture_table = self._get_or_create_table(
|
|
318
|
+
table_name=self.culture_table_name,
|
|
319
|
+
table_type="culture",
|
|
320
|
+
create_table_if_not_found=create_table_if_not_found,
|
|
321
|
+
)
|
|
322
|
+
return self.culture_table
|
|
323
|
+
|
|
324
|
+
elif table_type == "versions":
|
|
325
|
+
self.versions_table = self._get_or_create_table(
|
|
326
|
+
table_name=self.versions_table_name,
|
|
327
|
+
table_type="versions",
|
|
328
|
+
create_table_if_not_found=create_table_if_not_found,
|
|
329
|
+
)
|
|
330
|
+
return self.versions_table
|
|
331
|
+
|
|
332
|
+
else:
|
|
333
|
+
raise ValueError(f"Unknown table type: '{table_type}'")
|
|
334
|
+
|
|
335
|
+
def _get_or_create_table(
|
|
336
|
+
self,
|
|
337
|
+
table_name: str,
|
|
338
|
+
table_type: str,
|
|
339
|
+
create_table_if_not_found: Optional[bool] = False,
|
|
340
|
+
) -> Optional[Table]:
|
|
341
|
+
"""
|
|
342
|
+
Check if the table exists and is valid, else create it.
|
|
343
|
+
|
|
344
|
+
Args:
|
|
345
|
+
table_name (str): Name of the table to get or create
|
|
346
|
+
table_type (str): Type of table (used to get schema definition)
|
|
347
|
+
|
|
348
|
+
Returns:
|
|
349
|
+
Table: SQLAlchemy Table object
|
|
350
|
+
"""
|
|
351
|
+
with self.Session() as sess, sess.begin():
|
|
352
|
+
table_is_available = is_table_available(session=sess, table_name=table_name)
|
|
353
|
+
|
|
354
|
+
if not table_is_available:
|
|
355
|
+
if not create_table_if_not_found:
|
|
356
|
+
return None
|
|
357
|
+
return self._create_table(table_name=table_name, table_type=table_type)
|
|
358
|
+
|
|
359
|
+
# SQLite version of table validation (no schema)
|
|
360
|
+
if not is_valid_table(db_engine=self.db_engine, table_name=table_name, table_type=table_type):
|
|
361
|
+
raise ValueError(f"Table {table_name} has an invalid schema")
|
|
362
|
+
|
|
363
|
+
try:
|
|
364
|
+
table = Table(table_name, self.metadata, autoload_with=self.db_engine)
|
|
365
|
+
return table
|
|
366
|
+
|
|
367
|
+
except Exception as e:
|
|
368
|
+
log_error(f"Error loading existing table {table_name}: {e}")
|
|
369
|
+
raise e
|
|
370
|
+
|
|
371
|
+
def get_latest_schema_version(self, table_name: str):
|
|
372
|
+
"""Get the latest version of the database schema."""
|
|
373
|
+
table = self._get_table(table_type="versions", create_table_if_not_found=True)
|
|
374
|
+
if table is None:
|
|
375
|
+
return "2.0.0"
|
|
376
|
+
with self.Session() as sess:
|
|
377
|
+
stmt = select(table)
|
|
378
|
+
# Latest version for the given table
|
|
379
|
+
stmt = stmt.where(table.c.table_name == table_name)
|
|
380
|
+
stmt = stmt.order_by(table.c.version.desc()).limit(1)
|
|
381
|
+
result = sess.execute(stmt).fetchone()
|
|
382
|
+
if result is None:
|
|
383
|
+
return "2.0.0"
|
|
384
|
+
version_dict = dict(result._mapping)
|
|
385
|
+
return version_dict.get("version") or "2.0.0"
|
|
386
|
+
|
|
387
|
+
def upsert_schema_version(self, table_name: str, version: str) -> None:
|
|
388
|
+
"""Upsert the schema version into the database."""
|
|
389
|
+
table = self._get_table(table_type="versions", create_table_if_not_found=True)
|
|
390
|
+
if table is None:
|
|
391
|
+
return
|
|
392
|
+
current_datetime = datetime.now().isoformat()
|
|
393
|
+
with self.Session() as sess, sess.begin():
|
|
394
|
+
stmt = sqlite.insert(table).values(
|
|
395
|
+
table_name=table_name,
|
|
396
|
+
version=version,
|
|
397
|
+
created_at=current_datetime, # Store as ISO format string
|
|
398
|
+
updated_at=current_datetime,
|
|
399
|
+
)
|
|
400
|
+
# Update version if table_name already exists
|
|
401
|
+
stmt = stmt.on_conflict_do_update(
|
|
402
|
+
index_elements=["table_name"],
|
|
403
|
+
set_=dict(version=version, updated_at=current_datetime),
|
|
404
|
+
)
|
|
405
|
+
sess.execute(stmt)
|
|
406
|
+
|
|
407
|
+
# -- Session methods --
|
|
408
|
+
|
|
409
|
+
def delete_session(self, session_id: str) -> bool:
|
|
410
|
+
"""
|
|
411
|
+
Delete a session from the database.
|
|
412
|
+
|
|
413
|
+
Args:
|
|
414
|
+
session_id (str): ID of the session to delete
|
|
415
|
+
|
|
416
|
+
Raises:
|
|
417
|
+
Exception: If an error occurs during deletion.
|
|
418
|
+
"""
|
|
419
|
+
try:
|
|
420
|
+
table = self._get_table(table_type="sessions")
|
|
421
|
+
if table is None:
|
|
422
|
+
return False
|
|
423
|
+
|
|
424
|
+
with self.Session() as sess, sess.begin():
|
|
425
|
+
delete_stmt = table.delete().where(table.c.session_id == session_id)
|
|
426
|
+
result = sess.execute(delete_stmt)
|
|
427
|
+
if result.rowcount == 0:
|
|
428
|
+
log_debug(f"No session found to deletewith session_id: {session_id}")
|
|
429
|
+
return False
|
|
430
|
+
else:
|
|
431
|
+
log_debug(f"Successfully deleted session with session_id: {session_id}")
|
|
432
|
+
return True
|
|
433
|
+
|
|
434
|
+
except Exception as e:
|
|
435
|
+
log_error(f"Error deleting session: {e}")
|
|
436
|
+
raise e
|
|
437
|
+
|
|
438
|
+
def delete_sessions(self, session_ids: List[str]) -> None:
|
|
439
|
+
"""Delete all given sessions from the database.
|
|
440
|
+
Can handle multiple session types in the same run.
|
|
441
|
+
|
|
442
|
+
Args:
|
|
443
|
+
session_ids (List[str]): The IDs of the sessions to delete.
|
|
444
|
+
|
|
445
|
+
Raises:
|
|
446
|
+
Exception: If an error occurs during deletion.
|
|
447
|
+
"""
|
|
448
|
+
try:
|
|
449
|
+
table = self._get_table(table_type="sessions")
|
|
450
|
+
if table is None:
|
|
451
|
+
return
|
|
452
|
+
|
|
453
|
+
with self.Session() as sess, sess.begin():
|
|
454
|
+
delete_stmt = table.delete().where(table.c.session_id.in_(session_ids))
|
|
455
|
+
result = sess.execute(delete_stmt)
|
|
456
|
+
|
|
457
|
+
log_debug(f"Successfully deleted {result.rowcount} sessions")
|
|
458
|
+
|
|
459
|
+
except Exception as e:
|
|
460
|
+
log_error(f"Error deleting sessions: {e}")
|
|
461
|
+
raise e
|
|
462
|
+
|
|
463
|
+
def get_session(
|
|
464
|
+
self,
|
|
465
|
+
session_id: str,
|
|
466
|
+
session_type: SessionType,
|
|
467
|
+
user_id: Optional[str] = None,
|
|
468
|
+
deserialize: Optional[bool] = True,
|
|
469
|
+
) -> Optional[Union[Session, Dict[str, Any]]]:
|
|
470
|
+
"""
|
|
471
|
+
Read a session from the database.
|
|
472
|
+
|
|
473
|
+
Args:
|
|
474
|
+
session_id (str): ID of the session to read.
|
|
475
|
+
session_type (SessionType): Type of session to get.
|
|
476
|
+
user_id (Optional[str]): User ID to filter by. Defaults to None.
|
|
477
|
+
deserialize (Optional[bool]): Whether to serialize the session. Defaults to True.
|
|
478
|
+
|
|
479
|
+
Returns:
|
|
480
|
+
Optional[Union[Session, Dict[str, Any]]]:
|
|
481
|
+
- When deserialize=True: Session object
|
|
482
|
+
- When deserialize=False: Session dictionary
|
|
483
|
+
|
|
484
|
+
Raises:
|
|
485
|
+
Exception: If an error occurs during retrieval.
|
|
486
|
+
"""
|
|
487
|
+
try:
|
|
488
|
+
table = self._get_table(table_type="sessions")
|
|
489
|
+
if table is None:
|
|
490
|
+
return None
|
|
491
|
+
|
|
492
|
+
with self.Session() as sess, sess.begin():
|
|
493
|
+
stmt = select(table).where(table.c.session_id == session_id)
|
|
494
|
+
|
|
495
|
+
# Filtering
|
|
496
|
+
if user_id is not None:
|
|
497
|
+
stmt = stmt.where(table.c.user_id == user_id)
|
|
498
|
+
|
|
499
|
+
result = sess.execute(stmt).fetchone()
|
|
500
|
+
if result is None:
|
|
501
|
+
return None
|
|
502
|
+
|
|
503
|
+
session_raw = deserialize_session_json_fields(dict(result._mapping))
|
|
504
|
+
if not session_raw or not deserialize:
|
|
505
|
+
return session_raw
|
|
506
|
+
|
|
507
|
+
if session_type == SessionType.AGENT:
|
|
508
|
+
return AgentSession.from_dict(session_raw)
|
|
509
|
+
elif session_type == SessionType.TEAM:
|
|
510
|
+
return TeamSession.from_dict(session_raw)
|
|
511
|
+
elif session_type == SessionType.WORKFLOW:
|
|
512
|
+
return WorkflowSession.from_dict(session_raw)
|
|
513
|
+
else:
|
|
514
|
+
raise ValueError(f"Invalid session type: {session_type}")
|
|
515
|
+
|
|
516
|
+
except Exception as e:
|
|
517
|
+
log_debug(f"Exception reading from sessions table: {e}")
|
|
518
|
+
raise e
|
|
519
|
+
|
|
520
|
+
def get_sessions(
|
|
521
|
+
self,
|
|
522
|
+
session_type: Optional[SessionType] = None,
|
|
523
|
+
user_id: Optional[str] = None,
|
|
524
|
+
component_id: Optional[str] = None,
|
|
525
|
+
session_name: Optional[str] = None,
|
|
526
|
+
start_timestamp: Optional[int] = None,
|
|
527
|
+
end_timestamp: Optional[int] = None,
|
|
528
|
+
limit: Optional[int] = None,
|
|
529
|
+
page: Optional[int] = None,
|
|
530
|
+
sort_by: Optional[str] = None,
|
|
531
|
+
sort_order: Optional[str] = None,
|
|
532
|
+
deserialize: Optional[bool] = True,
|
|
533
|
+
) -> Union[List[Session], Tuple[List[Dict[str, Any]], int]]:
|
|
534
|
+
"""
|
|
535
|
+
Get all sessions in the given table. Can filter by user_id and entity_id.
|
|
536
|
+
Args:
|
|
537
|
+
session_type (Optional[SessionType]): The type of session to get.
|
|
538
|
+
user_id (Optional[str]): The ID of the user to filter by.
|
|
539
|
+
component_id (Optional[str]): The ID of the agent / workflow to filter by.
|
|
540
|
+
session_name (Optional[str]): The name of the session to filter by.
|
|
541
|
+
start_timestamp (Optional[int]): The start timestamp to filter by.
|
|
542
|
+
end_timestamp (Optional[int]): The end timestamp to filter by.
|
|
543
|
+
limit (Optional[int]): The maximum number of sessions to return. Defaults to None.
|
|
544
|
+
page (Optional[int]): The page number to return. Defaults to None.
|
|
545
|
+
sort_by (Optional[str]): The field to sort by. Defaults to None.
|
|
546
|
+
sort_order (Optional[str]): The sort order. Defaults to None.
|
|
547
|
+
deserialize (Optional[bool]): Whether to serialize the sessions. Defaults to True.
|
|
548
|
+
create_table_if_not_found (Optional[bool]): Whether to create the table if it doesn't exist.
|
|
549
|
+
|
|
550
|
+
Returns:
|
|
551
|
+
List[Session]:
|
|
552
|
+
- When deserialize=True: List of Session objects matching the criteria.
|
|
553
|
+
- When deserialize=False: List of Session dictionaries matching the criteria.
|
|
554
|
+
|
|
555
|
+
Raises:
|
|
556
|
+
Exception: If an error occurs during retrieval.
|
|
557
|
+
"""
|
|
558
|
+
try:
|
|
559
|
+
table = self._get_table(table_type="sessions")
|
|
560
|
+
if table is None:
|
|
561
|
+
return [] if deserialize else ([], 0)
|
|
562
|
+
|
|
563
|
+
with self.Session() as sess, sess.begin():
|
|
564
|
+
stmt = select(table)
|
|
565
|
+
|
|
566
|
+
# Filtering
|
|
567
|
+
if user_id is not None:
|
|
568
|
+
stmt = stmt.where(table.c.user_id == user_id)
|
|
569
|
+
if component_id is not None:
|
|
570
|
+
if session_type == SessionType.AGENT:
|
|
571
|
+
stmt = stmt.where(table.c.agent_id == component_id)
|
|
572
|
+
elif session_type == SessionType.TEAM:
|
|
573
|
+
stmt = stmt.where(table.c.team_id == component_id)
|
|
574
|
+
elif session_type == SessionType.WORKFLOW:
|
|
575
|
+
stmt = stmt.where(table.c.workflow_id == component_id)
|
|
576
|
+
if start_timestamp is not None:
|
|
577
|
+
stmt = stmt.where(table.c.created_at >= start_timestamp)
|
|
578
|
+
if end_timestamp is not None:
|
|
579
|
+
stmt = stmt.where(table.c.created_at <= end_timestamp)
|
|
580
|
+
if session_name is not None:
|
|
581
|
+
stmt = stmt.where(table.c.session_data.like(f"%{session_name}%"))
|
|
582
|
+
if session_type is not None:
|
|
583
|
+
stmt = stmt.where(table.c.session_type == session_type.value)
|
|
584
|
+
|
|
585
|
+
# Getting total count
|
|
586
|
+
count_stmt = select(func.count()).select_from(stmt.alias())
|
|
587
|
+
total_count = sess.execute(count_stmt).scalar()
|
|
588
|
+
|
|
589
|
+
# Sorting
|
|
590
|
+
stmt = apply_sorting(stmt, table, sort_by, sort_order)
|
|
591
|
+
|
|
592
|
+
# Paginating
|
|
593
|
+
if limit is not None:
|
|
594
|
+
stmt = stmt.limit(limit)
|
|
595
|
+
if page is not None:
|
|
596
|
+
stmt = stmt.offset((page - 1) * limit)
|
|
597
|
+
|
|
598
|
+
records = sess.execute(stmt).fetchall()
|
|
599
|
+
if records is None:
|
|
600
|
+
return [] if deserialize else ([], 0)
|
|
601
|
+
|
|
602
|
+
sessions_raw = [deserialize_session_json_fields(dict(record._mapping)) for record in records]
|
|
603
|
+
if not deserialize:
|
|
604
|
+
return sessions_raw, total_count
|
|
605
|
+
if not sessions_raw:
|
|
606
|
+
return []
|
|
607
|
+
|
|
608
|
+
if session_type == SessionType.AGENT:
|
|
609
|
+
return [AgentSession.from_dict(record) for record in sessions_raw] # type: ignore
|
|
610
|
+
elif session_type == SessionType.TEAM:
|
|
611
|
+
return [TeamSession.from_dict(record) for record in sessions_raw] # type: ignore
|
|
612
|
+
elif session_type == SessionType.WORKFLOW:
|
|
613
|
+
return [WorkflowSession.from_dict(record) for record in sessions_raw] # type: ignore
|
|
614
|
+
else:
|
|
615
|
+
raise ValueError(f"Invalid session type: {session_type}")
|
|
616
|
+
|
|
617
|
+
except Exception as e:
|
|
618
|
+
log_debug(f"Exception reading from sessions table: {e}")
|
|
619
|
+
raise e
|
|
620
|
+
|
|
621
|
+
def rename_session(
|
|
622
|
+
self,
|
|
623
|
+
session_id: str,
|
|
624
|
+
session_type: SessionType,
|
|
625
|
+
session_name: str,
|
|
626
|
+
deserialize: Optional[bool] = True,
|
|
627
|
+
) -> Optional[Union[Session, Dict[str, Any]]]:
|
|
628
|
+
"""
|
|
629
|
+
Rename a session in the database.
|
|
630
|
+
|
|
631
|
+
Args:
|
|
632
|
+
session_id (str): The ID of the session to rename.
|
|
633
|
+
session_type (SessionType): The type of session to rename.
|
|
634
|
+
session_name (str): The new name for the session.
|
|
635
|
+
deserialize (Optional[bool]): Whether to serialize the session. Defaults to True.
|
|
636
|
+
|
|
637
|
+
Returns:
|
|
638
|
+
Optional[Union[Session, Dict[str, Any]]]:
|
|
639
|
+
- When deserialize=True: Session object
|
|
640
|
+
- When deserialize=False: Session dictionary
|
|
641
|
+
|
|
642
|
+
Raises:
|
|
643
|
+
Exception: If an error occurs during renaming.
|
|
644
|
+
"""
|
|
645
|
+
try:
|
|
646
|
+
# Get the current session as a deserialized object
|
|
647
|
+
# Get the session record
|
|
648
|
+
session = self.get_session(session_id, session_type, deserialize=True)
|
|
649
|
+
if session is None:
|
|
650
|
+
return None
|
|
651
|
+
|
|
652
|
+
session = cast(Session, session)
|
|
653
|
+
# Update the session name
|
|
654
|
+
if session.session_data is None:
|
|
655
|
+
session.session_data = {}
|
|
656
|
+
session.session_data["session_name"] = session_name
|
|
657
|
+
|
|
658
|
+
# Upsert the updated session back to the database
|
|
659
|
+
return self.upsert_session(session, deserialize=deserialize)
|
|
660
|
+
|
|
661
|
+
except Exception as e:
|
|
662
|
+
log_error(f"Exception renaming session: {e}")
|
|
663
|
+
raise e
|
|
664
|
+
|
|
665
|
+
def upsert_session(
|
|
666
|
+
self, session: Session, deserialize: Optional[bool] = True
|
|
667
|
+
) -> Optional[Union[Session, Dict[str, Any]]]:
|
|
668
|
+
"""
|
|
669
|
+
Insert or update a session in the database.
|
|
670
|
+
|
|
671
|
+
Args:
|
|
672
|
+
session (Session): The session data to upsert.
|
|
673
|
+
deserialize (Optional[bool]): Whether to serialize the session. Defaults to True.
|
|
674
|
+
|
|
675
|
+
Returns:
|
|
676
|
+
Optional[Session]:
|
|
677
|
+
- When deserialize=True: Session object
|
|
678
|
+
- When deserialize=False: Session dictionary
|
|
679
|
+
|
|
680
|
+
Raises:
|
|
681
|
+
Exception: If an error occurs during upserting.
|
|
682
|
+
"""
|
|
683
|
+
try:
|
|
684
|
+
table = self._get_table(table_type="sessions", create_table_if_not_found=True)
|
|
685
|
+
if table is None:
|
|
686
|
+
return None
|
|
687
|
+
|
|
688
|
+
serialized_session = serialize_session_json_fields(session.to_dict())
|
|
689
|
+
|
|
690
|
+
if isinstance(session, AgentSession):
|
|
691
|
+
with self.Session() as sess, sess.begin():
|
|
692
|
+
stmt = sqlite.insert(table).values(
|
|
693
|
+
session_id=serialized_session.get("session_id"),
|
|
694
|
+
session_type=SessionType.AGENT.value,
|
|
695
|
+
agent_id=serialized_session.get("agent_id"),
|
|
696
|
+
user_id=serialized_session.get("user_id"),
|
|
697
|
+
agent_data=serialized_session.get("agent_data"),
|
|
698
|
+
session_data=serialized_session.get("session_data"),
|
|
699
|
+
metadata=serialized_session.get("metadata"),
|
|
700
|
+
runs=serialized_session.get("runs"),
|
|
701
|
+
summary=serialized_session.get("summary"),
|
|
702
|
+
created_at=serialized_session.get("created_at"),
|
|
703
|
+
updated_at=serialized_session.get("created_at"),
|
|
704
|
+
)
|
|
705
|
+
stmt = stmt.on_conflict_do_update(
|
|
706
|
+
index_elements=["session_id"],
|
|
707
|
+
set_=dict(
|
|
708
|
+
agent_id=serialized_session.get("agent_id"),
|
|
709
|
+
user_id=serialized_session.get("user_id"),
|
|
710
|
+
runs=serialized_session.get("runs"),
|
|
711
|
+
summary=serialized_session.get("summary"),
|
|
712
|
+
agent_data=serialized_session.get("agent_data"),
|
|
713
|
+
session_data=serialized_session.get("session_data"),
|
|
714
|
+
metadata=serialized_session.get("metadata"),
|
|
715
|
+
updated_at=int(time.time()),
|
|
716
|
+
),
|
|
717
|
+
)
|
|
718
|
+
stmt = stmt.returning(*table.columns) # type: ignore
|
|
719
|
+
result = sess.execute(stmt)
|
|
720
|
+
row = result.fetchone()
|
|
721
|
+
|
|
722
|
+
session_raw = deserialize_session_json_fields(dict(row._mapping)) if row else None
|
|
723
|
+
if session_raw is None or not deserialize:
|
|
724
|
+
return session_raw
|
|
725
|
+
return AgentSession.from_dict(session_raw)
|
|
726
|
+
|
|
727
|
+
elif isinstance(session, TeamSession):
|
|
728
|
+
with self.Session() as sess, sess.begin():
|
|
729
|
+
stmt = sqlite.insert(table).values(
|
|
730
|
+
session_id=serialized_session.get("session_id"),
|
|
731
|
+
session_type=SessionType.TEAM.value,
|
|
732
|
+
team_id=serialized_session.get("team_id"),
|
|
733
|
+
user_id=serialized_session.get("user_id"),
|
|
734
|
+
runs=serialized_session.get("runs"),
|
|
735
|
+
summary=serialized_session.get("summary"),
|
|
736
|
+
created_at=serialized_session.get("created_at"),
|
|
737
|
+
updated_at=serialized_session.get("created_at"),
|
|
738
|
+
team_data=serialized_session.get("team_data"),
|
|
739
|
+
session_data=serialized_session.get("session_data"),
|
|
740
|
+
metadata=serialized_session.get("metadata"),
|
|
741
|
+
)
|
|
742
|
+
|
|
743
|
+
stmt = stmt.on_conflict_do_update(
|
|
744
|
+
index_elements=["session_id"],
|
|
745
|
+
set_=dict(
|
|
746
|
+
team_id=serialized_session.get("team_id"),
|
|
747
|
+
user_id=serialized_session.get("user_id"),
|
|
748
|
+
summary=serialized_session.get("summary"),
|
|
749
|
+
runs=serialized_session.get("runs"),
|
|
750
|
+
team_data=serialized_session.get("team_data"),
|
|
751
|
+
session_data=serialized_session.get("session_data"),
|
|
752
|
+
metadata=serialized_session.get("metadata"),
|
|
753
|
+
updated_at=int(time.time()),
|
|
754
|
+
),
|
|
755
|
+
)
|
|
756
|
+
stmt = stmt.returning(*table.columns) # type: ignore
|
|
757
|
+
result = sess.execute(stmt)
|
|
758
|
+
row = result.fetchone()
|
|
759
|
+
|
|
760
|
+
session_raw = deserialize_session_json_fields(dict(row._mapping)) if row else None
|
|
761
|
+
if session_raw is None or not deserialize:
|
|
762
|
+
return session_raw
|
|
763
|
+
return TeamSession.from_dict(session_raw)
|
|
764
|
+
|
|
765
|
+
else:
|
|
766
|
+
with self.Session() as sess, sess.begin():
|
|
767
|
+
stmt = sqlite.insert(table).values(
|
|
768
|
+
session_id=serialized_session.get("session_id"),
|
|
769
|
+
session_type=SessionType.WORKFLOW.value,
|
|
770
|
+
workflow_id=serialized_session.get("workflow_id"),
|
|
771
|
+
user_id=serialized_session.get("user_id"),
|
|
772
|
+
runs=serialized_session.get("runs"),
|
|
773
|
+
summary=serialized_session.get("summary"),
|
|
774
|
+
created_at=serialized_session.get("created_at") or int(time.time()),
|
|
775
|
+
updated_at=serialized_session.get("updated_at") or int(time.time()),
|
|
776
|
+
workflow_data=serialized_session.get("workflow_data"),
|
|
777
|
+
session_data=serialized_session.get("session_data"),
|
|
778
|
+
metadata=serialized_session.get("metadata"),
|
|
779
|
+
)
|
|
780
|
+
stmt = stmt.on_conflict_do_update(
|
|
781
|
+
index_elements=["session_id"],
|
|
782
|
+
set_=dict(
|
|
783
|
+
workflow_id=serialized_session.get("workflow_id"),
|
|
784
|
+
user_id=serialized_session.get("user_id"),
|
|
785
|
+
summary=serialized_session.get("summary"),
|
|
786
|
+
runs=serialized_session.get("runs"),
|
|
787
|
+
workflow_data=serialized_session.get("workflow_data"),
|
|
788
|
+
session_data=serialized_session.get("session_data"),
|
|
789
|
+
metadata=serialized_session.get("metadata"),
|
|
790
|
+
updated_at=int(time.time()),
|
|
791
|
+
),
|
|
792
|
+
)
|
|
793
|
+
stmt = stmt.returning(*table.columns) # type: ignore
|
|
794
|
+
result = sess.execute(stmt)
|
|
795
|
+
row = result.fetchone()
|
|
796
|
+
|
|
797
|
+
session_raw = deserialize_session_json_fields(dict(row._mapping)) if row else None
|
|
798
|
+
if session_raw is None or not deserialize:
|
|
799
|
+
return session_raw
|
|
800
|
+
return WorkflowSession.from_dict(session_raw)
|
|
801
|
+
|
|
802
|
+
except Exception as e:
|
|
803
|
+
log_warning(f"Exception upserting into table: {e}")
|
|
804
|
+
raise e
|
|
805
|
+
|
|
806
|
+
def upsert_sessions(
|
|
807
|
+
self,
|
|
808
|
+
sessions: List[Session],
|
|
809
|
+
deserialize: Optional[bool] = True,
|
|
810
|
+
preserve_updated_at: bool = False,
|
|
811
|
+
) -> List[Union[Session, Dict[str, Any]]]:
|
|
812
|
+
"""
|
|
813
|
+
Bulk upsert multiple sessions for improved performance on large datasets.
|
|
814
|
+
|
|
815
|
+
Args:
|
|
816
|
+
sessions (List[Session]): List of sessions to upsert.
|
|
817
|
+
deserialize (Optional[bool]): Whether to deserialize the sessions. Defaults to True.
|
|
818
|
+
preserve_updated_at (bool): If True, preserve the updated_at from the session object.
|
|
819
|
+
|
|
820
|
+
Returns:
|
|
821
|
+
List[Union[Session, Dict[str, Any]]]: List of upserted sessions.
|
|
822
|
+
|
|
823
|
+
Raises:
|
|
824
|
+
Exception: If an error occurs during bulk upsert.
|
|
825
|
+
"""
|
|
826
|
+
if not sessions:
|
|
827
|
+
return []
|
|
828
|
+
|
|
829
|
+
try:
|
|
830
|
+
table = self._get_table(table_type="sessions", create_table_if_not_found=True)
|
|
831
|
+
if table is None:
|
|
832
|
+
log_info("Sessions table not available, falling back to individual upserts")
|
|
833
|
+
return [
|
|
834
|
+
result
|
|
835
|
+
for session in sessions
|
|
836
|
+
if session is not None
|
|
837
|
+
for result in [self.upsert_session(session, deserialize=deserialize)]
|
|
838
|
+
if result is not None
|
|
839
|
+
]
|
|
840
|
+
|
|
841
|
+
# Group sessions by type for batch processing
|
|
842
|
+
agent_sessions = []
|
|
843
|
+
team_sessions = []
|
|
844
|
+
workflow_sessions = []
|
|
845
|
+
|
|
846
|
+
for session in sessions:
|
|
847
|
+
if isinstance(session, AgentSession):
|
|
848
|
+
agent_sessions.append(session)
|
|
849
|
+
elif isinstance(session, TeamSession):
|
|
850
|
+
team_sessions.append(session)
|
|
851
|
+
elif isinstance(session, WorkflowSession):
|
|
852
|
+
workflow_sessions.append(session)
|
|
853
|
+
|
|
854
|
+
results: List[Union[Session, Dict[str, Any]]] = []
|
|
855
|
+
|
|
856
|
+
with self.Session() as sess, sess.begin():
|
|
857
|
+
# Bulk upsert agent sessions
|
|
858
|
+
if agent_sessions:
|
|
859
|
+
agent_data = []
|
|
860
|
+
for session in agent_sessions:
|
|
861
|
+
serialized_session = serialize_session_json_fields(session.to_dict())
|
|
862
|
+
# Use preserved updated_at if flag is set and value exists, otherwise use current time
|
|
863
|
+
updated_at = serialized_session.get("updated_at") if preserve_updated_at else int(time.time())
|
|
864
|
+
agent_data.append(
|
|
865
|
+
{
|
|
866
|
+
"session_id": serialized_session.get("session_id"),
|
|
867
|
+
"session_type": SessionType.AGENT.value,
|
|
868
|
+
"agent_id": serialized_session.get("agent_id"),
|
|
869
|
+
"user_id": serialized_session.get("user_id"),
|
|
870
|
+
"agent_data": serialized_session.get("agent_data"),
|
|
871
|
+
"session_data": serialized_session.get("session_data"),
|
|
872
|
+
"metadata": serialized_session.get("metadata"),
|
|
873
|
+
"runs": serialized_session.get("runs"),
|
|
874
|
+
"summary": serialized_session.get("summary"),
|
|
875
|
+
"created_at": serialized_session.get("created_at"),
|
|
876
|
+
"updated_at": updated_at,
|
|
877
|
+
}
|
|
878
|
+
)
|
|
879
|
+
|
|
880
|
+
if agent_data:
|
|
881
|
+
stmt = sqlite.insert(table)
|
|
882
|
+
stmt = stmt.on_conflict_do_update(
|
|
883
|
+
index_elements=["session_id"],
|
|
884
|
+
set_=dict(
|
|
885
|
+
agent_id=stmt.excluded.agent_id,
|
|
886
|
+
user_id=stmt.excluded.user_id,
|
|
887
|
+
agent_data=stmt.excluded.agent_data,
|
|
888
|
+
session_data=stmt.excluded.session_data,
|
|
889
|
+
metadata=stmt.excluded.metadata,
|
|
890
|
+
runs=stmt.excluded.runs,
|
|
891
|
+
summary=stmt.excluded.summary,
|
|
892
|
+
updated_at=stmt.excluded.updated_at,
|
|
893
|
+
),
|
|
894
|
+
)
|
|
895
|
+
sess.execute(stmt, agent_data)
|
|
896
|
+
|
|
897
|
+
# Fetch the results for agent sessions
|
|
898
|
+
agent_ids = [session.session_id for session in agent_sessions]
|
|
899
|
+
select_stmt = select(table).where(table.c.session_id.in_(agent_ids))
|
|
900
|
+
result = sess.execute(select_stmt).fetchall()
|
|
901
|
+
|
|
902
|
+
for row in result:
|
|
903
|
+
session_dict = deserialize_session_json_fields(dict(row._mapping))
|
|
904
|
+
if deserialize:
|
|
905
|
+
deserialized_agent_session = AgentSession.from_dict(session_dict)
|
|
906
|
+
if deserialized_agent_session is None:
|
|
907
|
+
continue
|
|
908
|
+
results.append(deserialized_agent_session)
|
|
909
|
+
else:
|
|
910
|
+
results.append(session_dict)
|
|
911
|
+
|
|
912
|
+
# Bulk upsert team sessions
|
|
913
|
+
if team_sessions:
|
|
914
|
+
team_data = []
|
|
915
|
+
for session in team_sessions:
|
|
916
|
+
serialized_session = serialize_session_json_fields(session.to_dict())
|
|
917
|
+
# Use preserved updated_at if flag is set and value exists, otherwise use current time
|
|
918
|
+
updated_at = serialized_session.get("updated_at") if preserve_updated_at else int(time.time())
|
|
919
|
+
team_data.append(
|
|
920
|
+
{
|
|
921
|
+
"session_id": serialized_session.get("session_id"),
|
|
922
|
+
"session_type": SessionType.TEAM.value,
|
|
923
|
+
"team_id": serialized_session.get("team_id"),
|
|
924
|
+
"user_id": serialized_session.get("user_id"),
|
|
925
|
+
"runs": serialized_session.get("runs"),
|
|
926
|
+
"summary": serialized_session.get("summary"),
|
|
927
|
+
"created_at": serialized_session.get("created_at"),
|
|
928
|
+
"updated_at": updated_at,
|
|
929
|
+
"team_data": serialized_session.get("team_data"),
|
|
930
|
+
"session_data": serialized_session.get("session_data"),
|
|
931
|
+
"metadata": serialized_session.get("metadata"),
|
|
932
|
+
}
|
|
933
|
+
)
|
|
934
|
+
|
|
935
|
+
if team_data:
|
|
936
|
+
stmt = sqlite.insert(table)
|
|
937
|
+
stmt = stmt.on_conflict_do_update(
|
|
938
|
+
index_elements=["session_id"],
|
|
939
|
+
set_=dict(
|
|
940
|
+
team_id=stmt.excluded.team_id,
|
|
941
|
+
user_id=stmt.excluded.user_id,
|
|
942
|
+
team_data=stmt.excluded.team_data,
|
|
943
|
+
session_data=stmt.excluded.session_data,
|
|
944
|
+
metadata=stmt.excluded.metadata,
|
|
945
|
+
runs=stmt.excluded.runs,
|
|
946
|
+
summary=stmt.excluded.summary,
|
|
947
|
+
updated_at=stmt.excluded.updated_at,
|
|
948
|
+
),
|
|
949
|
+
)
|
|
950
|
+
sess.execute(stmt, team_data)
|
|
951
|
+
|
|
952
|
+
# Fetch the results for team sessions
|
|
953
|
+
team_ids = [session.session_id for session in team_sessions]
|
|
954
|
+
select_stmt = select(table).where(table.c.session_id.in_(team_ids))
|
|
955
|
+
result = sess.execute(select_stmt).fetchall()
|
|
956
|
+
|
|
957
|
+
for row in result:
|
|
958
|
+
session_dict = deserialize_session_json_fields(dict(row._mapping))
|
|
959
|
+
if deserialize:
|
|
960
|
+
deserialized_team_session = TeamSession.from_dict(session_dict)
|
|
961
|
+
if deserialized_team_session is None:
|
|
962
|
+
continue
|
|
963
|
+
results.append(deserialized_team_session)
|
|
964
|
+
else:
|
|
965
|
+
results.append(session_dict)
|
|
966
|
+
|
|
967
|
+
# Bulk upsert workflow sessions
|
|
968
|
+
if workflow_sessions:
|
|
969
|
+
workflow_data = []
|
|
970
|
+
for session in workflow_sessions:
|
|
971
|
+
serialized_session = serialize_session_json_fields(session.to_dict())
|
|
972
|
+
# Use preserved updated_at if flag is set and value exists, otherwise use current time
|
|
973
|
+
updated_at = serialized_session.get("updated_at") if preserve_updated_at else int(time.time())
|
|
974
|
+
workflow_data.append(
|
|
975
|
+
{
|
|
976
|
+
"session_id": serialized_session.get("session_id"),
|
|
977
|
+
"session_type": SessionType.WORKFLOW.value,
|
|
978
|
+
"workflow_id": serialized_session.get("workflow_id"),
|
|
979
|
+
"user_id": serialized_session.get("user_id"),
|
|
980
|
+
"runs": serialized_session.get("runs"),
|
|
981
|
+
"summary": serialized_session.get("summary"),
|
|
982
|
+
"created_at": serialized_session.get("created_at"),
|
|
983
|
+
"updated_at": updated_at,
|
|
984
|
+
"workflow_data": serialized_session.get("workflow_data"),
|
|
985
|
+
"session_data": serialized_session.get("session_data"),
|
|
986
|
+
"metadata": serialized_session.get("metadata"),
|
|
987
|
+
}
|
|
988
|
+
)
|
|
989
|
+
|
|
990
|
+
if workflow_data:
|
|
991
|
+
stmt = sqlite.insert(table)
|
|
992
|
+
stmt = stmt.on_conflict_do_update(
|
|
993
|
+
index_elements=["session_id"],
|
|
994
|
+
set_=dict(
|
|
995
|
+
workflow_id=stmt.excluded.workflow_id,
|
|
996
|
+
user_id=stmt.excluded.user_id,
|
|
997
|
+
workflow_data=stmt.excluded.workflow_data,
|
|
998
|
+
session_data=stmt.excluded.session_data,
|
|
999
|
+
metadata=stmt.excluded.metadata,
|
|
1000
|
+
runs=stmt.excluded.runs,
|
|
1001
|
+
summary=stmt.excluded.summary,
|
|
1002
|
+
updated_at=stmt.excluded.updated_at,
|
|
1003
|
+
),
|
|
1004
|
+
)
|
|
1005
|
+
sess.execute(stmt, workflow_data)
|
|
1006
|
+
|
|
1007
|
+
# Fetch the results for workflow sessions
|
|
1008
|
+
workflow_ids = [session.session_id for session in workflow_sessions]
|
|
1009
|
+
select_stmt = select(table).where(table.c.session_id.in_(workflow_ids))
|
|
1010
|
+
result = sess.execute(select_stmt).fetchall()
|
|
1011
|
+
|
|
1012
|
+
for row in result:
|
|
1013
|
+
session_dict = deserialize_session_json_fields(dict(row._mapping))
|
|
1014
|
+
if deserialize:
|
|
1015
|
+
deserialized_workflow_session = WorkflowSession.from_dict(session_dict)
|
|
1016
|
+
if deserialized_workflow_session is None:
|
|
1017
|
+
continue
|
|
1018
|
+
results.append(deserialized_workflow_session)
|
|
1019
|
+
else:
|
|
1020
|
+
results.append(session_dict)
|
|
1021
|
+
|
|
1022
|
+
return results
|
|
1023
|
+
|
|
1024
|
+
except Exception as e:
|
|
1025
|
+
log_error(f"Exception during bulk session upsert, falling back to individual upserts: {e}")
|
|
1026
|
+
# Fallback to individual upserts
|
|
1027
|
+
return [
|
|
1028
|
+
result
|
|
1029
|
+
for session in sessions
|
|
1030
|
+
if session is not None
|
|
1031
|
+
for result in [self.upsert_session(session, deserialize=deserialize)]
|
|
1032
|
+
if result is not None
|
|
1033
|
+
]
|
|
1034
|
+
|
|
1035
|
+
# -- Memory methods --
|
|
1036
|
+
|
|
1037
|
+
def delete_user_memory(self, memory_id: str, user_id: Optional[str] = None):
|
|
1038
|
+
"""Delete a user memory from the database.
|
|
1039
|
+
|
|
1040
|
+
Args:
|
|
1041
|
+
memory_id (str): The ID of the memory to delete.
|
|
1042
|
+
user_id (Optional[str]): The user ID to filter by. Defaults to None.
|
|
1043
|
+
|
|
1044
|
+
Returns:
|
|
1045
|
+
bool: True if deletion was successful, False otherwise.
|
|
1046
|
+
|
|
1047
|
+
Raises:
|
|
1048
|
+
Exception: If an error occurs during deletion.
|
|
1049
|
+
"""
|
|
1050
|
+
try:
|
|
1051
|
+
table = self._get_table(table_type="memories")
|
|
1052
|
+
if table is None:
|
|
1053
|
+
return
|
|
1054
|
+
|
|
1055
|
+
with self.Session() as sess, sess.begin():
|
|
1056
|
+
delete_stmt = table.delete().where(table.c.memory_id == memory_id)
|
|
1057
|
+
if user_id is not None:
|
|
1058
|
+
delete_stmt = delete_stmt.where(table.c.user_id == user_id)
|
|
1059
|
+
result = sess.execute(delete_stmt)
|
|
1060
|
+
|
|
1061
|
+
success = result.rowcount > 0
|
|
1062
|
+
if success:
|
|
1063
|
+
log_debug(f"Successfully deleted user memory id: {memory_id}")
|
|
1064
|
+
else:
|
|
1065
|
+
log_debug(f"No user memory found with id: {memory_id}")
|
|
1066
|
+
|
|
1067
|
+
except Exception as e:
|
|
1068
|
+
log_error(f"Error deleting user memory: {e}")
|
|
1069
|
+
raise e
|
|
1070
|
+
|
|
1071
|
+
def delete_user_memories(self, memory_ids: List[str], user_id: Optional[str] = None) -> None:
|
|
1072
|
+
"""Delete user memories from the database.
|
|
1073
|
+
|
|
1074
|
+
Args:
|
|
1075
|
+
memory_ids (List[str]): The IDs of the memories to delete.
|
|
1076
|
+
user_id (Optional[str]): The user ID to filter by. Defaults to None.
|
|
1077
|
+
|
|
1078
|
+
Raises:
|
|
1079
|
+
Exception: If an error occurs during deletion.
|
|
1080
|
+
"""
|
|
1081
|
+
try:
|
|
1082
|
+
table = self._get_table(table_type="memories")
|
|
1083
|
+
if table is None:
|
|
1084
|
+
return
|
|
1085
|
+
|
|
1086
|
+
with self.Session() as sess, sess.begin():
|
|
1087
|
+
delete_stmt = table.delete().where(table.c.memory_id.in_(memory_ids))
|
|
1088
|
+
if user_id is not None:
|
|
1089
|
+
delete_stmt = delete_stmt.where(table.c.user_id == user_id)
|
|
1090
|
+
result = sess.execute(delete_stmt)
|
|
1091
|
+
if result.rowcount == 0:
|
|
1092
|
+
log_debug(f"No user memories found with ids: {memory_ids}")
|
|
1093
|
+
|
|
1094
|
+
except Exception as e:
|
|
1095
|
+
log_error(f"Error deleting user memories: {e}")
|
|
1096
|
+
raise e
|
|
1097
|
+
|
|
1098
|
+
def get_all_memory_topics(self) -> List[str]:
|
|
1099
|
+
"""Get all memory topics from the database.
|
|
1100
|
+
|
|
1101
|
+
Returns:
|
|
1102
|
+
List[str]: List of memory topics.
|
|
1103
|
+
"""
|
|
1104
|
+
try:
|
|
1105
|
+
table = self._get_table(table_type="memories")
|
|
1106
|
+
if table is None:
|
|
1107
|
+
return []
|
|
1108
|
+
|
|
1109
|
+
with self.Session() as sess, sess.begin():
|
|
1110
|
+
# Select topics from all results
|
|
1111
|
+
stmt = select(table.c.topics)
|
|
1112
|
+
result = sess.execute(stmt).fetchall()
|
|
1113
|
+
return list(set([record[0] for record in result]))
|
|
1114
|
+
|
|
1115
|
+
except Exception as e:
|
|
1116
|
+
log_debug(f"Exception reading from memory table: {e}")
|
|
1117
|
+
raise e
|
|
1118
|
+
|
|
1119
|
+
def get_user_memory(
|
|
1120
|
+
self,
|
|
1121
|
+
memory_id: str,
|
|
1122
|
+
deserialize: Optional[bool] = True,
|
|
1123
|
+
user_id: Optional[str] = None,
|
|
1124
|
+
) -> Optional[Union[UserMemory, Dict[str, Any]]]:
|
|
1125
|
+
"""Get a memory from the database.
|
|
1126
|
+
|
|
1127
|
+
Args:
|
|
1128
|
+
memory_id (str): The ID of the memory to get.
|
|
1129
|
+
deserialize (Optional[bool]): Whether to serialize the memory. Defaults to True.
|
|
1130
|
+
user_id (Optional[str]): The user ID to filter by. Defaults to None.
|
|
1131
|
+
|
|
1132
|
+
Returns:
|
|
1133
|
+
Optional[Union[UserMemory, Dict[str, Any]]]:
|
|
1134
|
+
- When deserialize=True: UserMemory object
|
|
1135
|
+
- When deserialize=False: Memory dictionary
|
|
1136
|
+
|
|
1137
|
+
Raises:
|
|
1138
|
+
Exception: If an error occurs during retrieval.
|
|
1139
|
+
"""
|
|
1140
|
+
try:
|
|
1141
|
+
table = self._get_table(table_type="memories")
|
|
1142
|
+
if table is None:
|
|
1143
|
+
return None
|
|
1144
|
+
|
|
1145
|
+
with self.Session() as sess, sess.begin():
|
|
1146
|
+
stmt = select(table).where(table.c.memory_id == memory_id)
|
|
1147
|
+
if user_id is not None:
|
|
1148
|
+
stmt = stmt.where(table.c.user_id == user_id)
|
|
1149
|
+
result = sess.execute(stmt).fetchone()
|
|
1150
|
+
if result is None:
|
|
1151
|
+
return None
|
|
1152
|
+
|
|
1153
|
+
memory_raw = dict(result._mapping)
|
|
1154
|
+
if not memory_raw or not deserialize:
|
|
1155
|
+
return memory_raw
|
|
1156
|
+
|
|
1157
|
+
return UserMemory.from_dict(memory_raw)
|
|
1158
|
+
|
|
1159
|
+
except Exception as e:
|
|
1160
|
+
log_debug(f"Exception reading from memorytable: {e}")
|
|
1161
|
+
raise e
|
|
1162
|
+
|
|
1163
|
+
def get_user_memories(
|
|
1164
|
+
self,
|
|
1165
|
+
user_id: Optional[str] = None,
|
|
1166
|
+
agent_id: Optional[str] = None,
|
|
1167
|
+
team_id: Optional[str] = None,
|
|
1168
|
+
topics: Optional[List[str]] = None,
|
|
1169
|
+
search_content: Optional[str] = None,
|
|
1170
|
+
limit: Optional[int] = None,
|
|
1171
|
+
page: Optional[int] = None,
|
|
1172
|
+
sort_by: Optional[str] = None,
|
|
1173
|
+
sort_order: Optional[str] = None,
|
|
1174
|
+
deserialize: Optional[bool] = True,
|
|
1175
|
+
) -> Union[List[UserMemory], Tuple[List[Dict[str, Any]], int]]:
|
|
1176
|
+
"""Get all memories from the database as UserMemory objects.
|
|
1177
|
+
|
|
1178
|
+
Args:
|
|
1179
|
+
user_id (Optional[str]): The ID of the user to filter by.
|
|
1180
|
+
agent_id (Optional[str]): The ID of the agent to filter by.
|
|
1181
|
+
team_id (Optional[str]): The ID of the team to filter by.
|
|
1182
|
+
topics (Optional[List[str]]): The topics to filter by.
|
|
1183
|
+
search_content (Optional[str]): The content to search for.
|
|
1184
|
+
limit (Optional[int]): The maximum number of memories to return.
|
|
1185
|
+
page (Optional[int]): The page number.
|
|
1186
|
+
sort_by (Optional[str]): The column to sort by.
|
|
1187
|
+
sort_order (Optional[str]): The order to sort by.
|
|
1188
|
+
deserialize (Optional[bool]): Whether to serialize the memories. Defaults to True.
|
|
1189
|
+
|
|
1190
|
+
|
|
1191
|
+
Returns:
|
|
1192
|
+
Union[List[UserMemory], Tuple[List[Dict[str, Any]], int]]:
|
|
1193
|
+
- When deserialize=True: List of UserMemory objects
|
|
1194
|
+
- When deserialize=False: List of UserMemory dictionaries and total count
|
|
1195
|
+
|
|
1196
|
+
Raises:
|
|
1197
|
+
Exception: If an error occurs during retrieval.
|
|
1198
|
+
"""
|
|
1199
|
+
try:
|
|
1200
|
+
table = self._get_table(table_type="memories")
|
|
1201
|
+
if table is None:
|
|
1202
|
+
return [] if deserialize else ([], 0)
|
|
1203
|
+
|
|
1204
|
+
with self.Session() as sess, sess.begin():
|
|
1205
|
+
stmt = select(table)
|
|
1206
|
+
|
|
1207
|
+
# Filtering
|
|
1208
|
+
if user_id is not None:
|
|
1209
|
+
stmt = stmt.where(table.c.user_id == user_id)
|
|
1210
|
+
if agent_id is not None:
|
|
1211
|
+
stmt = stmt.where(table.c.agent_id == agent_id)
|
|
1212
|
+
if team_id is not None:
|
|
1213
|
+
stmt = stmt.where(table.c.team_id == team_id)
|
|
1214
|
+
if topics is not None:
|
|
1215
|
+
for topic in topics:
|
|
1216
|
+
stmt = stmt.where(func.cast(table.c.topics, String).like(f'%"{topic}"%'))
|
|
1217
|
+
if search_content is not None:
|
|
1218
|
+
stmt = stmt.where(table.c.memory.ilike(f"%{search_content}%"))
|
|
1219
|
+
|
|
1220
|
+
# Get total count after applying filtering
|
|
1221
|
+
count_stmt = select(func.count()).select_from(stmt.alias())
|
|
1222
|
+
total_count = sess.execute(count_stmt).scalar()
|
|
1223
|
+
|
|
1224
|
+
# Sorting
|
|
1225
|
+
stmt = apply_sorting(stmt, table, sort_by, sort_order)
|
|
1226
|
+
# Paginating
|
|
1227
|
+
if limit is not None:
|
|
1228
|
+
stmt = stmt.limit(limit)
|
|
1229
|
+
if page is not None:
|
|
1230
|
+
stmt = stmt.offset((page - 1) * limit)
|
|
1231
|
+
|
|
1232
|
+
result = sess.execute(stmt).fetchall()
|
|
1233
|
+
if not result:
|
|
1234
|
+
return [] if deserialize else ([], 0)
|
|
1235
|
+
|
|
1236
|
+
memories_raw = [record._mapping for record in result]
|
|
1237
|
+
|
|
1238
|
+
if not deserialize:
|
|
1239
|
+
return memories_raw, total_count
|
|
1240
|
+
|
|
1241
|
+
return [UserMemory.from_dict(record) for record in memories_raw]
|
|
1242
|
+
|
|
1243
|
+
except Exception as e:
|
|
1244
|
+
log_error(f"Error reading from memory table: {e}")
|
|
1245
|
+
raise e
|
|
1246
|
+
|
|
1247
|
+
def get_user_memory_stats(
|
|
1248
|
+
self,
|
|
1249
|
+
limit: Optional[int] = None,
|
|
1250
|
+
page: Optional[int] = None,
|
|
1251
|
+
user_id: Optional[str] = None,
|
|
1252
|
+
) -> Tuple[List[Dict[str, Any]], int]:
|
|
1253
|
+
"""Get user memories stats.
|
|
1254
|
+
|
|
1255
|
+
Args:
|
|
1256
|
+
limit (Optional[int]): The maximum number of user stats to return.
|
|
1257
|
+
page (Optional[int]): The page number.
|
|
1258
|
+
user_id (Optional[str]): User ID for filtering.
|
|
1259
|
+
|
|
1260
|
+
Returns:
|
|
1261
|
+
Tuple[List[Dict[str, Any]], int]: A list of dictionaries containing user stats and total count.
|
|
1262
|
+
|
|
1263
|
+
Example:
|
|
1264
|
+
(
|
|
1265
|
+
[
|
|
1266
|
+
{
|
|
1267
|
+
"user_id": "123",
|
|
1268
|
+
"total_memories": 10,
|
|
1269
|
+
"last_memory_updated_at": 1714560000,
|
|
1270
|
+
},
|
|
1271
|
+
],
|
|
1272
|
+
total_count: 1,
|
|
1273
|
+
)
|
|
1274
|
+
"""
|
|
1275
|
+
try:
|
|
1276
|
+
table = self._get_table(table_type="memories")
|
|
1277
|
+
if table is None:
|
|
1278
|
+
return [], 0
|
|
1279
|
+
|
|
1280
|
+
with self.Session() as sess, sess.begin():
|
|
1281
|
+
stmt = select(
|
|
1282
|
+
table.c.user_id,
|
|
1283
|
+
func.count(table.c.memory_id).label("total_memories"),
|
|
1284
|
+
func.max(table.c.updated_at).label("last_memory_updated_at"),
|
|
1285
|
+
)
|
|
1286
|
+
if user_id is not None:
|
|
1287
|
+
stmt = stmt.where(table.c.user_id == user_id)
|
|
1288
|
+
else:
|
|
1289
|
+
stmt = stmt.where(table.c.user_id.is_not(None))
|
|
1290
|
+
stmt = stmt.group_by(table.c.user_id)
|
|
1291
|
+
stmt = stmt.order_by(func.max(table.c.updated_at).desc())
|
|
1292
|
+
|
|
1293
|
+
count_stmt = select(func.count()).select_from(stmt.alias())
|
|
1294
|
+
total_count = sess.execute(count_stmt).scalar() or 0
|
|
1295
|
+
|
|
1296
|
+
# Pagination
|
|
1297
|
+
if limit is not None:
|
|
1298
|
+
stmt = stmt.limit(limit)
|
|
1299
|
+
if page is not None:
|
|
1300
|
+
stmt = stmt.offset((page - 1) * limit)
|
|
1301
|
+
|
|
1302
|
+
result = sess.execute(stmt).fetchall()
|
|
1303
|
+
if not result:
|
|
1304
|
+
return [], 0
|
|
1305
|
+
|
|
1306
|
+
return [
|
|
1307
|
+
{
|
|
1308
|
+
"user_id": record.user_id, # type: ignore
|
|
1309
|
+
"total_memories": record.total_memories,
|
|
1310
|
+
"last_memory_updated_at": record.last_memory_updated_at,
|
|
1311
|
+
}
|
|
1312
|
+
for record in result
|
|
1313
|
+
], total_count
|
|
1314
|
+
|
|
1315
|
+
except Exception as e:
|
|
1316
|
+
log_error(f"Error getting user memory stats: {e}")
|
|
1317
|
+
raise e
|
|
1318
|
+
|
|
1319
|
+
def upsert_user_memory(
|
|
1320
|
+
self, memory: UserMemory, deserialize: Optional[bool] = True
|
|
1321
|
+
) -> Optional[Union[UserMemory, Dict[str, Any]]]:
|
|
1322
|
+
"""Upsert a user memory in the database.
|
|
1323
|
+
|
|
1324
|
+
Args:
|
|
1325
|
+
memory (UserMemory): The user memory to upsert.
|
|
1326
|
+
deserialize (Optional[bool]): Whether to serialize the memory. Defaults to True.
|
|
1327
|
+
|
|
1328
|
+
Returns:
|
|
1329
|
+
Optional[Union[UserMemory, Dict[str, Any]]]:
|
|
1330
|
+
- When deserialize=True: UserMemory object
|
|
1331
|
+
- When deserialize=False: UserMemory dictionary
|
|
1332
|
+
|
|
1333
|
+
Raises:
|
|
1334
|
+
Exception: If an error occurs during upsert.
|
|
1335
|
+
"""
|
|
1336
|
+
try:
|
|
1337
|
+
table = self._get_table(table_type="memories", create_table_if_not_found=True)
|
|
1338
|
+
if table is None:
|
|
1339
|
+
return None
|
|
1340
|
+
|
|
1341
|
+
if memory.memory_id is None:
|
|
1342
|
+
memory.memory_id = str(uuid4())
|
|
1343
|
+
|
|
1344
|
+
current_time = int(time.time())
|
|
1345
|
+
|
|
1346
|
+
with self.Session() as sess, sess.begin():
|
|
1347
|
+
stmt = sqlite.insert(table).values(
|
|
1348
|
+
user_id=memory.user_id,
|
|
1349
|
+
agent_id=memory.agent_id,
|
|
1350
|
+
team_id=memory.team_id,
|
|
1351
|
+
memory_id=memory.memory_id,
|
|
1352
|
+
memory=memory.memory,
|
|
1353
|
+
topics=memory.topics,
|
|
1354
|
+
input=memory.input,
|
|
1355
|
+
feedback=memory.feedback,
|
|
1356
|
+
created_at=memory.created_at,
|
|
1357
|
+
updated_at=memory.created_at,
|
|
1358
|
+
)
|
|
1359
|
+
stmt = stmt.on_conflict_do_update( # type: ignore
|
|
1360
|
+
index_elements=["memory_id"],
|
|
1361
|
+
set_=dict(
|
|
1362
|
+
memory=memory.memory,
|
|
1363
|
+
topics=memory.topics,
|
|
1364
|
+
input=memory.input,
|
|
1365
|
+
agent_id=memory.agent_id,
|
|
1366
|
+
team_id=memory.team_id,
|
|
1367
|
+
feedback=memory.feedback,
|
|
1368
|
+
updated_at=current_time,
|
|
1369
|
+
# Preserve created_at on update - don't overwrite existing value
|
|
1370
|
+
created_at=table.c.created_at,
|
|
1371
|
+
),
|
|
1372
|
+
).returning(table)
|
|
1373
|
+
|
|
1374
|
+
result = sess.execute(stmt)
|
|
1375
|
+
row = result.fetchone()
|
|
1376
|
+
|
|
1377
|
+
if row is None:
|
|
1378
|
+
return None
|
|
1379
|
+
|
|
1380
|
+
memory_raw = row._mapping
|
|
1381
|
+
if not memory_raw or not deserialize:
|
|
1382
|
+
return memory_raw
|
|
1383
|
+
|
|
1384
|
+
return UserMemory.from_dict(memory_raw)
|
|
1385
|
+
|
|
1386
|
+
except Exception as e:
|
|
1387
|
+
log_error(f"Error upserting user memory: {e}")
|
|
1388
|
+
raise e
|
|
1389
|
+
|
|
1390
|
+
def upsert_memories(
|
|
1391
|
+
self,
|
|
1392
|
+
memories: List[UserMemory],
|
|
1393
|
+
deserialize: Optional[bool] = True,
|
|
1394
|
+
preserve_updated_at: bool = False,
|
|
1395
|
+
) -> List[Union[UserMemory, Dict[str, Any]]]:
|
|
1396
|
+
"""
|
|
1397
|
+
Bulk upsert multiple user memories for improved performance on large datasets.
|
|
1398
|
+
|
|
1399
|
+
Args:
|
|
1400
|
+
memories (List[UserMemory]): List of memories to upsert.
|
|
1401
|
+
deserialize (Optional[bool]): Whether to deserialize the memories. Defaults to True.
|
|
1402
|
+
|
|
1403
|
+
Returns:
|
|
1404
|
+
List[Union[UserMemory, Dict[str, Any]]]: List of upserted memories.
|
|
1405
|
+
|
|
1406
|
+
Raises:
|
|
1407
|
+
Exception: If an error occurs during bulk upsert.
|
|
1408
|
+
"""
|
|
1409
|
+
if not memories:
|
|
1410
|
+
return []
|
|
1411
|
+
|
|
1412
|
+
try:
|
|
1413
|
+
table = self._get_table(table_type="memories", create_table_if_not_found=True)
|
|
1414
|
+
if table is None:
|
|
1415
|
+
log_info("Memories table not available, falling back to individual upserts")
|
|
1416
|
+
return [
|
|
1417
|
+
result
|
|
1418
|
+
for memory in memories
|
|
1419
|
+
if memory is not None
|
|
1420
|
+
for result in [self.upsert_user_memory(memory, deserialize=deserialize)]
|
|
1421
|
+
if result is not None
|
|
1422
|
+
]
|
|
1423
|
+
# Prepare bulk data
|
|
1424
|
+
bulk_data = []
|
|
1425
|
+
current_time = int(time.time())
|
|
1426
|
+
|
|
1427
|
+
for memory in memories:
|
|
1428
|
+
if memory.memory_id is None:
|
|
1429
|
+
memory.memory_id = str(uuid4())
|
|
1430
|
+
|
|
1431
|
+
# Use preserved updated_at if flag is set and value exists, otherwise use current time
|
|
1432
|
+
updated_at = memory.updated_at if preserve_updated_at else current_time
|
|
1433
|
+
|
|
1434
|
+
bulk_data.append(
|
|
1435
|
+
{
|
|
1436
|
+
"user_id": memory.user_id,
|
|
1437
|
+
"agent_id": memory.agent_id,
|
|
1438
|
+
"team_id": memory.team_id,
|
|
1439
|
+
"memory_id": memory.memory_id,
|
|
1440
|
+
"memory": memory.memory,
|
|
1441
|
+
"topics": memory.topics,
|
|
1442
|
+
"input": memory.input,
|
|
1443
|
+
"feedback": memory.feedback,
|
|
1444
|
+
"created_at": memory.created_at,
|
|
1445
|
+
"updated_at": updated_at,
|
|
1446
|
+
}
|
|
1447
|
+
)
|
|
1448
|
+
|
|
1449
|
+
results: List[Union[UserMemory, Dict[str, Any]]] = []
|
|
1450
|
+
|
|
1451
|
+
with self.Session() as sess, sess.begin():
|
|
1452
|
+
# Bulk upsert memories using SQLite ON CONFLICT DO UPDATE
|
|
1453
|
+
stmt = sqlite.insert(table)
|
|
1454
|
+
stmt = stmt.on_conflict_do_update(
|
|
1455
|
+
index_elements=["memory_id"],
|
|
1456
|
+
set_=dict(
|
|
1457
|
+
memory=stmt.excluded.memory,
|
|
1458
|
+
topics=stmt.excluded.topics,
|
|
1459
|
+
input=stmt.excluded.input,
|
|
1460
|
+
agent_id=stmt.excluded.agent_id,
|
|
1461
|
+
team_id=stmt.excluded.team_id,
|
|
1462
|
+
feedback=stmt.excluded.feedback,
|
|
1463
|
+
updated_at=stmt.excluded.updated_at,
|
|
1464
|
+
# Preserve created_at on update
|
|
1465
|
+
created_at=table.c.created_at,
|
|
1466
|
+
),
|
|
1467
|
+
)
|
|
1468
|
+
sess.execute(stmt, bulk_data)
|
|
1469
|
+
|
|
1470
|
+
# Fetch results
|
|
1471
|
+
memory_ids = [memory.memory_id for memory in memories if memory.memory_id]
|
|
1472
|
+
select_stmt = select(table).where(table.c.memory_id.in_(memory_ids))
|
|
1473
|
+
result = sess.execute(select_stmt).fetchall()
|
|
1474
|
+
|
|
1475
|
+
for row in result:
|
|
1476
|
+
memory_dict = dict(row._mapping)
|
|
1477
|
+
if deserialize:
|
|
1478
|
+
results.append(UserMemory.from_dict(memory_dict))
|
|
1479
|
+
else:
|
|
1480
|
+
results.append(memory_dict)
|
|
1481
|
+
|
|
1482
|
+
return results
|
|
1483
|
+
|
|
1484
|
+
except Exception as e:
|
|
1485
|
+
log_error(f"Exception during bulk memory upsert, falling back to individual upserts: {e}")
|
|
1486
|
+
|
|
1487
|
+
# Fallback to individual upserts
|
|
1488
|
+
return [
|
|
1489
|
+
result
|
|
1490
|
+
for memory in memories
|
|
1491
|
+
if memory is not None
|
|
1492
|
+
for result in [self.upsert_user_memory(memory, deserialize=deserialize)]
|
|
1493
|
+
if result is not None
|
|
1494
|
+
]
|
|
1495
|
+
|
|
1496
|
+
def clear_memories(self) -> None:
|
|
1497
|
+
"""Delete all memories from the database.
|
|
1498
|
+
|
|
1499
|
+
Raises:
|
|
1500
|
+
Exception: If an error occurs during deletion.
|
|
1501
|
+
"""
|
|
1502
|
+
try:
|
|
1503
|
+
table = self._get_table(table_type="memories")
|
|
1504
|
+
if table is None:
|
|
1505
|
+
return
|
|
1506
|
+
|
|
1507
|
+
with self.Session() as sess, sess.begin():
|
|
1508
|
+
sess.execute(table.delete())
|
|
1509
|
+
|
|
1510
|
+
except Exception as e:
|
|
1511
|
+
from agno.utils.log import log_warning
|
|
1512
|
+
|
|
1513
|
+
log_warning(f"Exception deleting all memories: {e}")
|
|
1514
|
+
raise e
|
|
1515
|
+
|
|
1516
|
+
# -- Metrics methods --
|
|
1517
|
+
|
|
1518
|
+
def _get_all_sessions_for_metrics_calculation(
|
|
1519
|
+
self, start_timestamp: Optional[int] = None, end_timestamp: Optional[int] = None
|
|
1520
|
+
) -> List[Dict[str, Any]]:
|
|
1521
|
+
"""
|
|
1522
|
+
Get all sessions of all types (agent, team, workflow) as raw dictionaries.
|
|
1523
|
+
|
|
1524
|
+
Args:
|
|
1525
|
+
start_timestamp (Optional[int]): The start timestamp to filter by. Defaults to None.
|
|
1526
|
+
end_timestamp (Optional[int]): The end timestamp to filter by. Defaults to None.
|
|
1527
|
+
|
|
1528
|
+
Returns:
|
|
1529
|
+
List[Dict[str, Any]]: List of session dictionaries with session_type field.
|
|
1530
|
+
|
|
1531
|
+
Raises:
|
|
1532
|
+
Exception: If an error occurs during retrieval.
|
|
1533
|
+
"""
|
|
1534
|
+
try:
|
|
1535
|
+
table = self._get_table(table_type="sessions")
|
|
1536
|
+
if table is None:
|
|
1537
|
+
return []
|
|
1538
|
+
|
|
1539
|
+
stmt = select(
|
|
1540
|
+
table.c.user_id,
|
|
1541
|
+
table.c.session_data,
|
|
1542
|
+
table.c.runs,
|
|
1543
|
+
table.c.created_at,
|
|
1544
|
+
table.c.session_type,
|
|
1545
|
+
)
|
|
1546
|
+
|
|
1547
|
+
if start_timestamp is not None:
|
|
1548
|
+
stmt = stmt.where(table.c.created_at >= start_timestamp)
|
|
1549
|
+
if end_timestamp is not None:
|
|
1550
|
+
stmt = stmt.where(table.c.created_at <= end_timestamp)
|
|
1551
|
+
|
|
1552
|
+
with self.Session() as sess:
|
|
1553
|
+
result = sess.execute(stmt).fetchall()
|
|
1554
|
+
return [record._mapping for record in result]
|
|
1555
|
+
|
|
1556
|
+
except Exception as e:
|
|
1557
|
+
log_error(f"Error reading from sessions table: {e}")
|
|
1558
|
+
raise e
|
|
1559
|
+
|
|
1560
|
+
def _get_metrics_calculation_starting_date(self, table: Table) -> Optional[date]:
|
|
1561
|
+
"""Get the first date for which metrics calculation is needed:
|
|
1562
|
+
|
|
1563
|
+
1. If there are metrics records, return the date of the first day without a complete metrics record.
|
|
1564
|
+
2. If there are no metrics records, return the date of the first recorded session.
|
|
1565
|
+
3. If there are no metrics records and no sessions records, return None.
|
|
1566
|
+
|
|
1567
|
+
Args:
|
|
1568
|
+
table (Table): The table to get the starting date for.
|
|
1569
|
+
|
|
1570
|
+
Returns:
|
|
1571
|
+
Optional[date]: The starting date for which metrics calculation is needed.
|
|
1572
|
+
"""
|
|
1573
|
+
with self.Session() as sess:
|
|
1574
|
+
stmt = select(table).order_by(table.c.date.desc()).limit(1)
|
|
1575
|
+
result = sess.execute(stmt).fetchone()
|
|
1576
|
+
|
|
1577
|
+
# 1. Return the date of the first day without a complete metrics record.
|
|
1578
|
+
if result is not None:
|
|
1579
|
+
if result.completed:
|
|
1580
|
+
return result._mapping["date"] + timedelta(days=1)
|
|
1581
|
+
else:
|
|
1582
|
+
return result._mapping["date"]
|
|
1583
|
+
|
|
1584
|
+
# 2. No metrics records. Return the date of the first recorded session.
|
|
1585
|
+
first_session, _ = self.get_sessions(sort_by="created_at", sort_order="asc", limit=1, deserialize=False)
|
|
1586
|
+
first_session_date = first_session[0]["created_at"] if first_session else None # type: ignore
|
|
1587
|
+
|
|
1588
|
+
# 3. No metrics records and no sessions records. Return None.
|
|
1589
|
+
if not first_session_date:
|
|
1590
|
+
return None
|
|
1591
|
+
|
|
1592
|
+
return datetime.fromtimestamp(first_session_date, tz=timezone.utc).date()
|
|
1593
|
+
|
|
1594
|
+
def calculate_metrics(self) -> Optional[list[dict]]:
|
|
1595
|
+
"""Calculate metrics for all dates without complete metrics.
|
|
1596
|
+
|
|
1597
|
+
Returns:
|
|
1598
|
+
Optional[list[dict]]: The calculated metrics.
|
|
1599
|
+
|
|
1600
|
+
Raises:
|
|
1601
|
+
Exception: If an error occurs during metrics calculation.
|
|
1602
|
+
"""
|
|
1603
|
+
try:
|
|
1604
|
+
table = self._get_table(table_type="metrics", create_table_if_not_found=True)
|
|
1605
|
+
if table is None:
|
|
1606
|
+
return None
|
|
1607
|
+
|
|
1608
|
+
starting_date = self._get_metrics_calculation_starting_date(table)
|
|
1609
|
+
if starting_date is None:
|
|
1610
|
+
log_info("No session data found. Won't calculate metrics.")
|
|
1611
|
+
return None
|
|
1612
|
+
|
|
1613
|
+
dates_to_process = get_dates_to_calculate_metrics_for(starting_date)
|
|
1614
|
+
if not dates_to_process:
|
|
1615
|
+
log_info("Metrics already calculated for all relevant dates.")
|
|
1616
|
+
return None
|
|
1617
|
+
|
|
1618
|
+
start_timestamp = int(
|
|
1619
|
+
datetime.combine(dates_to_process[0], datetime.min.time()).replace(tzinfo=timezone.utc).timestamp()
|
|
1620
|
+
)
|
|
1621
|
+
end_timestamp = int(
|
|
1622
|
+
datetime.combine(dates_to_process[-1] + timedelta(days=1), datetime.min.time())
|
|
1623
|
+
.replace(tzinfo=timezone.utc)
|
|
1624
|
+
.timestamp()
|
|
1625
|
+
)
|
|
1626
|
+
|
|
1627
|
+
sessions = self._get_all_sessions_for_metrics_calculation(
|
|
1628
|
+
start_timestamp=start_timestamp, end_timestamp=end_timestamp
|
|
1629
|
+
)
|
|
1630
|
+
all_sessions_data = fetch_all_sessions_data(
|
|
1631
|
+
sessions=sessions,
|
|
1632
|
+
dates_to_process=dates_to_process,
|
|
1633
|
+
start_timestamp=start_timestamp,
|
|
1634
|
+
)
|
|
1635
|
+
if not all_sessions_data:
|
|
1636
|
+
log_info("No new session data found. Won't calculate metrics.")
|
|
1637
|
+
return None
|
|
1638
|
+
|
|
1639
|
+
results = []
|
|
1640
|
+
metrics_records = []
|
|
1641
|
+
|
|
1642
|
+
for date_to_process in dates_to_process:
|
|
1643
|
+
date_key = date_to_process.isoformat()
|
|
1644
|
+
sessions_for_date = all_sessions_data.get(date_key, {})
|
|
1645
|
+
|
|
1646
|
+
# Skip dates with no sessions
|
|
1647
|
+
if not any(len(sessions) > 0 for sessions in sessions_for_date.values()):
|
|
1648
|
+
continue
|
|
1649
|
+
|
|
1650
|
+
metrics_record = calculate_date_metrics(date_to_process, sessions_for_date)
|
|
1651
|
+
metrics_records.append(metrics_record)
|
|
1652
|
+
|
|
1653
|
+
if metrics_records:
|
|
1654
|
+
with self.Session() as sess, sess.begin():
|
|
1655
|
+
results = bulk_upsert_metrics(session=sess, table=table, metrics_records=metrics_records)
|
|
1656
|
+
|
|
1657
|
+
log_debug("Updated metrics calculations")
|
|
1658
|
+
|
|
1659
|
+
return results
|
|
1660
|
+
|
|
1661
|
+
except Exception as e:
|
|
1662
|
+
log_error(f"Error refreshing metrics: {e}")
|
|
1663
|
+
raise e
|
|
1664
|
+
|
|
1665
|
+
def get_metrics(
|
|
1666
|
+
self,
|
|
1667
|
+
starting_date: Optional[date] = None,
|
|
1668
|
+
ending_date: Optional[date] = None,
|
|
1669
|
+
) -> Tuple[List[dict], Optional[int]]:
|
|
1670
|
+
"""Get all metrics matching the given date range.
|
|
1671
|
+
|
|
1672
|
+
Args:
|
|
1673
|
+
starting_date (Optional[date]): The starting date to filter metrics by.
|
|
1674
|
+
ending_date (Optional[date]): The ending date to filter metrics by.
|
|
1675
|
+
|
|
1676
|
+
Returns:
|
|
1677
|
+
Tuple[List[dict], Optional[int]]: A tuple containing the metrics and the timestamp of the latest update.
|
|
1678
|
+
|
|
1679
|
+
Raises:
|
|
1680
|
+
Exception: If an error occurs during retrieval.
|
|
1681
|
+
"""
|
|
1682
|
+
try:
|
|
1683
|
+
table = self._get_table(table_type="metrics", create_table_if_not_found=True)
|
|
1684
|
+
if table is None:
|
|
1685
|
+
return [], None
|
|
1686
|
+
|
|
1687
|
+
with self.Session() as sess, sess.begin():
|
|
1688
|
+
stmt = select(table)
|
|
1689
|
+
if starting_date:
|
|
1690
|
+
stmt = stmt.where(table.c.date >= starting_date)
|
|
1691
|
+
if ending_date:
|
|
1692
|
+
stmt = stmt.where(table.c.date <= ending_date)
|
|
1693
|
+
result = sess.execute(stmt).fetchall()
|
|
1694
|
+
if not result:
|
|
1695
|
+
return [], None
|
|
1696
|
+
|
|
1697
|
+
# Get the latest updated_at
|
|
1698
|
+
latest_stmt = select(func.max(table.c.updated_at))
|
|
1699
|
+
latest_updated_at = sess.execute(latest_stmt).scalar()
|
|
1700
|
+
|
|
1701
|
+
return [row._mapping for row in result], latest_updated_at
|
|
1702
|
+
|
|
1703
|
+
except Exception as e:
|
|
1704
|
+
log_error(f"Error getting metrics: {e}")
|
|
1705
|
+
raise e
|
|
1706
|
+
|
|
1707
|
+
# -- Knowledge methods --
|
|
1708
|
+
|
|
1709
|
+
def delete_knowledge_content(self, id: str):
|
|
1710
|
+
"""Delete a knowledge row from the database.
|
|
1711
|
+
|
|
1712
|
+
Args:
|
|
1713
|
+
id (str): The ID of the knowledge row to delete.
|
|
1714
|
+
|
|
1715
|
+
Raises:
|
|
1716
|
+
Exception: If an error occurs during deletion.
|
|
1717
|
+
"""
|
|
1718
|
+
table = self._get_table(table_type="knowledge")
|
|
1719
|
+
if table is None:
|
|
1720
|
+
return
|
|
1721
|
+
|
|
1722
|
+
try:
|
|
1723
|
+
with self.Session() as sess, sess.begin():
|
|
1724
|
+
stmt = table.delete().where(table.c.id == id)
|
|
1725
|
+
sess.execute(stmt)
|
|
1726
|
+
|
|
1727
|
+
except Exception as e:
|
|
1728
|
+
log_error(f"Error deleting knowledge content: {e}")
|
|
1729
|
+
raise e
|
|
1730
|
+
|
|
1731
|
+
def get_knowledge_content(self, id: str) -> Optional[KnowledgeRow]:
|
|
1732
|
+
"""Get a knowledge row from the database.
|
|
1733
|
+
|
|
1734
|
+
Args:
|
|
1735
|
+
id (str): The ID of the knowledge row to get.
|
|
1736
|
+
|
|
1737
|
+
Returns:
|
|
1738
|
+
Optional[KnowledgeRow]: The knowledge row, or None if it doesn't exist.
|
|
1739
|
+
|
|
1740
|
+
Raises:
|
|
1741
|
+
Exception: If an error occurs during retrieval.
|
|
1742
|
+
"""
|
|
1743
|
+
table = self._get_table(table_type="knowledge")
|
|
1744
|
+
if table is None:
|
|
1745
|
+
return None
|
|
1746
|
+
|
|
1747
|
+
try:
|
|
1748
|
+
with self.Session() as sess, sess.begin():
|
|
1749
|
+
stmt = select(table).where(table.c.id == id)
|
|
1750
|
+
result = sess.execute(stmt).fetchone()
|
|
1751
|
+
if result is None:
|
|
1752
|
+
return None
|
|
1753
|
+
|
|
1754
|
+
return KnowledgeRow.model_validate(result._mapping)
|
|
1755
|
+
|
|
1756
|
+
except Exception as e:
|
|
1757
|
+
log_error(f"Error getting knowledge content: {e}")
|
|
1758
|
+
raise e
|
|
1759
|
+
|
|
1760
|
+
def get_knowledge_contents(
|
|
1761
|
+
self,
|
|
1762
|
+
limit: Optional[int] = None,
|
|
1763
|
+
page: Optional[int] = None,
|
|
1764
|
+
sort_by: Optional[str] = None,
|
|
1765
|
+
sort_order: Optional[str] = None,
|
|
1766
|
+
) -> Tuple[List[KnowledgeRow], int]:
|
|
1767
|
+
"""Get all knowledge contents from the database.
|
|
1768
|
+
|
|
1769
|
+
Args:
|
|
1770
|
+
limit (Optional[int]): The maximum number of knowledge contents to return.
|
|
1771
|
+
page (Optional[int]): The page number.
|
|
1772
|
+
sort_by (Optional[str]): The column to sort by.
|
|
1773
|
+
sort_order (Optional[str]): The order to sort by.
|
|
1774
|
+
|
|
1775
|
+
Returns:
|
|
1776
|
+
Tuple[List[KnowledgeRow], int]: The knowledge contents and total count.
|
|
1777
|
+
|
|
1778
|
+
Raises:
|
|
1779
|
+
Exception: If an error occurs during retrieval.
|
|
1780
|
+
"""
|
|
1781
|
+
table = self._get_table(table_type="knowledge")
|
|
1782
|
+
if table is None:
|
|
1783
|
+
return [], 0
|
|
1784
|
+
|
|
1785
|
+
try:
|
|
1786
|
+
with self.Session() as sess, sess.begin():
|
|
1787
|
+
stmt = select(table)
|
|
1788
|
+
|
|
1789
|
+
# Apply sorting
|
|
1790
|
+
if sort_by is not None:
|
|
1791
|
+
stmt = stmt.order_by(getattr(table.c, sort_by) * (1 if sort_order == "asc" else -1))
|
|
1792
|
+
|
|
1793
|
+
# Get total count before applying limit and pagination
|
|
1794
|
+
count_stmt = select(func.count()).select_from(stmt.alias())
|
|
1795
|
+
total_count = sess.execute(count_stmt).scalar()
|
|
1796
|
+
|
|
1797
|
+
# Apply pagination after count
|
|
1798
|
+
if limit is not None:
|
|
1799
|
+
stmt = stmt.limit(limit)
|
|
1800
|
+
if page is not None:
|
|
1801
|
+
stmt = stmt.offset((page - 1) * limit)
|
|
1802
|
+
|
|
1803
|
+
result = sess.execute(stmt).fetchall()
|
|
1804
|
+
return [KnowledgeRow.model_validate(record._mapping) for record in result], total_count
|
|
1805
|
+
|
|
1806
|
+
except Exception as e:
|
|
1807
|
+
log_error(f"Error getting knowledge contents: {e}")
|
|
1808
|
+
raise e
|
|
1809
|
+
|
|
1810
|
+
def upsert_knowledge_content(self, knowledge_row: KnowledgeRow):
|
|
1811
|
+
"""Upsert knowledge content in the database.
|
|
1812
|
+
|
|
1813
|
+
Args:
|
|
1814
|
+
knowledge_row (KnowledgeRow): The knowledge row to upsert.
|
|
1815
|
+
|
|
1816
|
+
Returns:
|
|
1817
|
+
Optional[KnowledgeRow]: The upserted knowledge row, or None if the operation fails.
|
|
1818
|
+
"""
|
|
1819
|
+
try:
|
|
1820
|
+
table = self._get_table(table_type="knowledge", create_table_if_not_found=True)
|
|
1821
|
+
if table is None:
|
|
1822
|
+
return None
|
|
1823
|
+
|
|
1824
|
+
with self.Session() as sess, sess.begin():
|
|
1825
|
+
update_fields = {
|
|
1826
|
+
k: v
|
|
1827
|
+
for k, v in {
|
|
1828
|
+
"name": knowledge_row.name,
|
|
1829
|
+
"description": knowledge_row.description,
|
|
1830
|
+
"metadata": knowledge_row.metadata,
|
|
1831
|
+
"type": knowledge_row.type,
|
|
1832
|
+
"size": knowledge_row.size,
|
|
1833
|
+
"linked_to": knowledge_row.linked_to,
|
|
1834
|
+
"access_count": knowledge_row.access_count,
|
|
1835
|
+
"status": knowledge_row.status,
|
|
1836
|
+
"status_message": knowledge_row.status_message,
|
|
1837
|
+
"created_at": knowledge_row.created_at,
|
|
1838
|
+
"updated_at": knowledge_row.updated_at,
|
|
1839
|
+
"external_id": knowledge_row.external_id,
|
|
1840
|
+
}.items()
|
|
1841
|
+
# Filtering out None fields if updating
|
|
1842
|
+
if v is not None
|
|
1843
|
+
}
|
|
1844
|
+
|
|
1845
|
+
stmt = (
|
|
1846
|
+
sqlite.insert(table)
|
|
1847
|
+
.values(knowledge_row.model_dump())
|
|
1848
|
+
.on_conflict_do_update(index_elements=["id"], set_=update_fields)
|
|
1849
|
+
)
|
|
1850
|
+
sess.execute(stmt)
|
|
1851
|
+
|
|
1852
|
+
return knowledge_row
|
|
1853
|
+
|
|
1854
|
+
except Exception as e:
|
|
1855
|
+
log_error(f"Error upserting knowledge content: {e}")
|
|
1856
|
+
raise e
|
|
1857
|
+
|
|
1858
|
+
# -- Eval methods --
|
|
1859
|
+
|
|
1860
|
+
def create_eval_run(self, eval_run: EvalRunRecord) -> Optional[EvalRunRecord]:
|
|
1861
|
+
"""Create an EvalRunRecord in the database.
|
|
1862
|
+
|
|
1863
|
+
Args:
|
|
1864
|
+
eval_run (EvalRunRecord): The eval run to create.
|
|
1865
|
+
|
|
1866
|
+
Returns:
|
|
1867
|
+
Optional[EvalRunRecord]: The created eval run, or None if the operation fails.
|
|
1868
|
+
|
|
1869
|
+
Raises:
|
|
1870
|
+
Exception: If an error occurs during creation.
|
|
1871
|
+
"""
|
|
1872
|
+
try:
|
|
1873
|
+
table = self._get_table(table_type="evals", create_table_if_not_found=True)
|
|
1874
|
+
if table is None:
|
|
1875
|
+
return None
|
|
1876
|
+
|
|
1877
|
+
with self.Session() as sess, sess.begin():
|
|
1878
|
+
current_time = int(time.time())
|
|
1879
|
+
stmt = sqlite.insert(table).values(
|
|
1880
|
+
{
|
|
1881
|
+
"created_at": current_time,
|
|
1882
|
+
"updated_at": current_time,
|
|
1883
|
+
**eval_run.model_dump(),
|
|
1884
|
+
}
|
|
1885
|
+
)
|
|
1886
|
+
sess.execute(stmt)
|
|
1887
|
+
sess.commit()
|
|
1888
|
+
|
|
1889
|
+
log_debug(f"Created eval run with id '{eval_run.run_id}'")
|
|
1890
|
+
|
|
1891
|
+
return eval_run
|
|
1892
|
+
|
|
1893
|
+
except Exception as e:
|
|
1894
|
+
log_error(f"Error creating eval run: {e}")
|
|
1895
|
+
raise e
|
|
1896
|
+
|
|
1897
|
+
def delete_eval_run(self, eval_run_id: str) -> None:
|
|
1898
|
+
"""Delete an eval run from the database.
|
|
1899
|
+
|
|
1900
|
+
Args:
|
|
1901
|
+
eval_run_id (str): The ID of the eval run to delete.
|
|
1902
|
+
"""
|
|
1903
|
+
try:
|
|
1904
|
+
table = self._get_table(table_type="evals")
|
|
1905
|
+
if table is None:
|
|
1906
|
+
return
|
|
1907
|
+
|
|
1908
|
+
with self.Session() as sess, sess.begin():
|
|
1909
|
+
stmt = table.delete().where(table.c.run_id == eval_run_id)
|
|
1910
|
+
result = sess.execute(stmt)
|
|
1911
|
+
if result.rowcount == 0:
|
|
1912
|
+
log_warning(f"No eval run found with ID: {eval_run_id}")
|
|
1913
|
+
else:
|
|
1914
|
+
log_debug(f"Deleted eval run with ID: {eval_run_id}")
|
|
1915
|
+
|
|
1916
|
+
except Exception as e:
|
|
1917
|
+
log_error(f"Error deleting eval run {eval_run_id}: {e}")
|
|
1918
|
+
raise e
|
|
1919
|
+
|
|
1920
|
+
def delete_eval_runs(self, eval_run_ids: List[str]) -> None:
|
|
1921
|
+
"""Delete multiple eval runs from the database.
|
|
1922
|
+
|
|
1923
|
+
Args:
|
|
1924
|
+
eval_run_ids (List[str]): List of eval run IDs to delete.
|
|
1925
|
+
"""
|
|
1926
|
+
try:
|
|
1927
|
+
table = self._get_table(table_type="evals")
|
|
1928
|
+
if table is None:
|
|
1929
|
+
return
|
|
1930
|
+
|
|
1931
|
+
with self.Session() as sess, sess.begin():
|
|
1932
|
+
stmt = table.delete().where(table.c.run_id.in_(eval_run_ids))
|
|
1933
|
+
result = sess.execute(stmt)
|
|
1934
|
+
if result.rowcount == 0:
|
|
1935
|
+
log_debug(f"No eval runs found with IDs: {eval_run_ids}")
|
|
1936
|
+
else:
|
|
1937
|
+
log_debug(f"Deleted {result.rowcount} eval runs")
|
|
1938
|
+
|
|
1939
|
+
except Exception as e:
|
|
1940
|
+
log_error(f"Error deleting eval runs {eval_run_ids}: {e}")
|
|
1941
|
+
raise e
|
|
1942
|
+
|
|
1943
|
+
def get_eval_run(
|
|
1944
|
+
self, eval_run_id: str, deserialize: Optional[bool] = True
|
|
1945
|
+
) -> Optional[Union[EvalRunRecord, Dict[str, Any]]]:
|
|
1946
|
+
"""Get an eval run from the database.
|
|
1947
|
+
|
|
1948
|
+
Args:
|
|
1949
|
+
eval_run_id (str): The ID of the eval run to get.
|
|
1950
|
+
deserialize (Optional[bool]): Whether to serialize the eval run. Defaults to True.
|
|
1951
|
+
|
|
1952
|
+
Returns:
|
|
1953
|
+
Optional[Union[EvalRunRecord, Dict[str, Any]]]:
|
|
1954
|
+
- When deserialize=True: EvalRunRecord object
|
|
1955
|
+
- When deserialize=False: EvalRun dictionary
|
|
1956
|
+
|
|
1957
|
+
Raises:
|
|
1958
|
+
Exception: If an error occurs during retrieval.
|
|
1959
|
+
"""
|
|
1960
|
+
try:
|
|
1961
|
+
table = self._get_table(table_type="evals")
|
|
1962
|
+
if table is None:
|
|
1963
|
+
return None
|
|
1964
|
+
|
|
1965
|
+
with self.Session() as sess, sess.begin():
|
|
1966
|
+
stmt = select(table).where(table.c.run_id == eval_run_id)
|
|
1967
|
+
result = sess.execute(stmt).fetchone()
|
|
1968
|
+
if result is None:
|
|
1969
|
+
return None
|
|
1970
|
+
|
|
1971
|
+
eval_run_raw = result._mapping
|
|
1972
|
+
if not eval_run_raw or not deserialize:
|
|
1973
|
+
return eval_run_raw
|
|
1974
|
+
|
|
1975
|
+
return EvalRunRecord.model_validate(eval_run_raw)
|
|
1976
|
+
|
|
1977
|
+
except Exception as e:
|
|
1978
|
+
log_error(f"Exception getting eval run {eval_run_id}: {e}")
|
|
1979
|
+
raise e
|
|
1980
|
+
|
|
1981
|
+
def get_eval_runs(
|
|
1982
|
+
self,
|
|
1983
|
+
limit: Optional[int] = None,
|
|
1984
|
+
page: Optional[int] = None,
|
|
1985
|
+
sort_by: Optional[str] = None,
|
|
1986
|
+
sort_order: Optional[str] = None,
|
|
1987
|
+
agent_id: Optional[str] = None,
|
|
1988
|
+
team_id: Optional[str] = None,
|
|
1989
|
+
workflow_id: Optional[str] = None,
|
|
1990
|
+
model_id: Optional[str] = None,
|
|
1991
|
+
filter_type: Optional[EvalFilterType] = None,
|
|
1992
|
+
eval_type: Optional[List[EvalType]] = None,
|
|
1993
|
+
deserialize: Optional[bool] = True,
|
|
1994
|
+
) -> Union[List[EvalRunRecord], Tuple[List[Dict[str, Any]], int]]:
|
|
1995
|
+
"""Get all eval runs from the database.
|
|
1996
|
+
|
|
1997
|
+
Args:
|
|
1998
|
+
limit (Optional[int]): The maximum number of eval runs to return.
|
|
1999
|
+
page (Optional[int]): The page number.
|
|
2000
|
+
sort_by (Optional[str]): The column to sort by.
|
|
2001
|
+
sort_order (Optional[str]): The order to sort by.
|
|
2002
|
+
agent_id (Optional[str]): The ID of the agent to filter by.
|
|
2003
|
+
team_id (Optional[str]): The ID of the team to filter by.
|
|
2004
|
+
workflow_id (Optional[str]): The ID of the workflow to filter by.
|
|
2005
|
+
model_id (Optional[str]): The ID of the model to filter by.
|
|
2006
|
+
eval_type (Optional[List[EvalType]]): The type(s) of eval to filter by.
|
|
2007
|
+
filter_type (Optional[EvalFilterType]): Filter by component type (agent, team, workflow).
|
|
2008
|
+
deserialize (Optional[bool]): Whether to serialize the eval runs. Defaults to True.
|
|
2009
|
+
create_table_if_not_found (Optional[bool]): Whether to create the table if it doesn't exist.
|
|
2010
|
+
|
|
2011
|
+
Returns:
|
|
2012
|
+
Union[List[EvalRunRecord], Tuple[List[Dict[str, Any]], int]]:
|
|
2013
|
+
- When deserialize=True: List of EvalRunRecord objects
|
|
2014
|
+
- When deserialize=False: List of EvalRun dictionaries and total count
|
|
2015
|
+
|
|
2016
|
+
Raises:
|
|
2017
|
+
Exception: If an error occurs during retrieval.
|
|
2018
|
+
"""
|
|
2019
|
+
try:
|
|
2020
|
+
table = self._get_table(table_type="evals")
|
|
2021
|
+
if table is None:
|
|
2022
|
+
return [] if deserialize else ([], 0)
|
|
2023
|
+
|
|
2024
|
+
with self.Session() as sess, sess.begin():
|
|
2025
|
+
stmt = select(table)
|
|
2026
|
+
|
|
2027
|
+
# Filtering
|
|
2028
|
+
if agent_id is not None:
|
|
2029
|
+
stmt = stmt.where(table.c.agent_id == agent_id)
|
|
2030
|
+
if team_id is not None:
|
|
2031
|
+
stmt = stmt.where(table.c.team_id == team_id)
|
|
2032
|
+
if workflow_id is not None:
|
|
2033
|
+
stmt = stmt.where(table.c.workflow_id == workflow_id)
|
|
2034
|
+
if model_id is not None:
|
|
2035
|
+
stmt = stmt.where(table.c.model_id == model_id)
|
|
2036
|
+
if eval_type is not None and len(eval_type) > 0:
|
|
2037
|
+
stmt = stmt.where(table.c.eval_type.in_(eval_type))
|
|
2038
|
+
if filter_type is not None:
|
|
2039
|
+
if filter_type == EvalFilterType.AGENT:
|
|
2040
|
+
stmt = stmt.where(table.c.agent_id.is_not(None))
|
|
2041
|
+
elif filter_type == EvalFilterType.TEAM:
|
|
2042
|
+
stmt = stmt.where(table.c.team_id.is_not(None))
|
|
2043
|
+
elif filter_type == EvalFilterType.WORKFLOW:
|
|
2044
|
+
stmt = stmt.where(table.c.workflow_id.is_not(None))
|
|
2045
|
+
|
|
2046
|
+
# Get total count after applying filtering
|
|
2047
|
+
count_stmt = select(func.count()).select_from(stmt.alias())
|
|
2048
|
+
total_count = sess.execute(count_stmt).scalar()
|
|
2049
|
+
|
|
2050
|
+
# Sorting - apply default sort by created_at desc if no sort parameters provided
|
|
2051
|
+
if sort_by is None:
|
|
2052
|
+
stmt = stmt.order_by(table.c.created_at.desc())
|
|
2053
|
+
else:
|
|
2054
|
+
stmt = apply_sorting(stmt, table, sort_by, sort_order)
|
|
2055
|
+
# Paginating
|
|
2056
|
+
if limit is not None:
|
|
2057
|
+
stmt = stmt.limit(limit)
|
|
2058
|
+
if page is not None:
|
|
2059
|
+
stmt = stmt.offset((page - 1) * limit)
|
|
2060
|
+
|
|
2061
|
+
result = sess.execute(stmt).fetchall()
|
|
2062
|
+
if not result:
|
|
2063
|
+
return [] if deserialize else ([], 0)
|
|
2064
|
+
|
|
2065
|
+
eval_runs_raw = [row._mapping for row in result]
|
|
2066
|
+
if not deserialize:
|
|
2067
|
+
return eval_runs_raw, total_count
|
|
2068
|
+
|
|
2069
|
+
return [EvalRunRecord.model_validate(row) for row in eval_runs_raw]
|
|
2070
|
+
|
|
2071
|
+
except Exception as e:
|
|
2072
|
+
log_error(f"Exception getting eval runs: {e}")
|
|
2073
|
+
raise e
|
|
2074
|
+
|
|
2075
|
+
def rename_eval_run(
|
|
2076
|
+
self, eval_run_id: str, name: str, deserialize: Optional[bool] = True
|
|
2077
|
+
) -> Optional[Union[EvalRunRecord, Dict[str, Any]]]:
|
|
2078
|
+
"""Upsert the name of an eval run in the database, returning raw dictionary.
|
|
2079
|
+
|
|
2080
|
+
Args:
|
|
2081
|
+
eval_run_id (str): The ID of the eval run to update.
|
|
2082
|
+
name (str): The new name of the eval run.
|
|
2083
|
+
deserialize (Optional[bool]): Whether to serialize the eval run. Defaults to True.
|
|
2084
|
+
|
|
2085
|
+
Returns:
|
|
2086
|
+
Optional[Union[EvalRunRecord, Dict[str, Any]]]:
|
|
2087
|
+
- When deserialize=True: EvalRunRecord object
|
|
2088
|
+
- When deserialize=False: EvalRun dictionary
|
|
2089
|
+
|
|
2090
|
+
Raises:
|
|
2091
|
+
Exception: If an error occurs during update.
|
|
2092
|
+
"""
|
|
2093
|
+
try:
|
|
2094
|
+
table = self._get_table(table_type="evals")
|
|
2095
|
+
if table is None:
|
|
2096
|
+
return None
|
|
2097
|
+
|
|
2098
|
+
with self.Session() as sess, sess.begin():
|
|
2099
|
+
stmt = (
|
|
2100
|
+
table.update().where(table.c.run_id == eval_run_id).values(name=name, updated_at=int(time.time()))
|
|
2101
|
+
)
|
|
2102
|
+
sess.execute(stmt)
|
|
2103
|
+
|
|
2104
|
+
eval_run_raw = self.get_eval_run(eval_run_id=eval_run_id, deserialize=deserialize)
|
|
2105
|
+
|
|
2106
|
+
log_debug(f"Renamed eval run with id '{eval_run_id}' to '{name}'")
|
|
2107
|
+
|
|
2108
|
+
if not eval_run_raw or not deserialize:
|
|
2109
|
+
return eval_run_raw
|
|
2110
|
+
|
|
2111
|
+
return EvalRunRecord.model_validate(eval_run_raw)
|
|
2112
|
+
|
|
2113
|
+
except Exception as e:
|
|
2114
|
+
log_error(f"Error renaming eval run {eval_run_id}: {e}")
|
|
2115
|
+
raise e
|
|
2116
|
+
|
|
2117
|
+
# -- Trace methods --
|
|
2118
|
+
|
|
2119
|
+
def _get_traces_base_query(self, table: Table, spans_table: Optional[Table] = None):
|
|
2120
|
+
"""Build base query for traces with aggregated span counts.
|
|
2121
|
+
|
|
2122
|
+
Args:
|
|
2123
|
+
table: The traces table.
|
|
2124
|
+
spans_table: The spans table (optional).
|
|
2125
|
+
|
|
2126
|
+
Returns:
|
|
2127
|
+
SQLAlchemy select statement with total_spans and error_count calculated dynamically.
|
|
2128
|
+
"""
|
|
2129
|
+
from sqlalchemy import case, func, literal
|
|
2130
|
+
|
|
2131
|
+
if spans_table is not None:
|
|
2132
|
+
# JOIN with spans table to calculate total_spans and error_count
|
|
2133
|
+
return (
|
|
2134
|
+
select(
|
|
2135
|
+
table,
|
|
2136
|
+
func.coalesce(func.count(spans_table.c.span_id), 0).label("total_spans"),
|
|
2137
|
+
func.coalesce(func.sum(case((spans_table.c.status_code == "ERROR", 1), else_=0)), 0).label(
|
|
2138
|
+
"error_count"
|
|
2139
|
+
),
|
|
2140
|
+
)
|
|
2141
|
+
.select_from(table.outerjoin(spans_table, table.c.trace_id == spans_table.c.trace_id))
|
|
2142
|
+
.group_by(table.c.trace_id)
|
|
2143
|
+
)
|
|
2144
|
+
else:
|
|
2145
|
+
# Fallback if spans table doesn't exist
|
|
2146
|
+
return select(table, literal(0).label("total_spans"), literal(0).label("error_count"))
|
|
2147
|
+
|
|
2148
|
+
def _get_trace_component_level_expr(self, workflow_id_col, team_id_col, agent_id_col, name_col):
|
|
2149
|
+
"""Build a SQL CASE expression that returns the component level for a trace.
|
|
2150
|
+
|
|
2151
|
+
Component levels (higher = more important):
|
|
2152
|
+
- 3: Workflow root (.run or .arun with workflow_id)
|
|
2153
|
+
- 2: Team root (.run or .arun with team_id)
|
|
2154
|
+
- 1: Agent root (.run or .arun with agent_id)
|
|
2155
|
+
- 0: Child span (not a root)
|
|
2156
|
+
|
|
2157
|
+
Args:
|
|
2158
|
+
workflow_id_col: SQL column/expression for workflow_id
|
|
2159
|
+
team_id_col: SQL column/expression for team_id
|
|
2160
|
+
agent_id_col: SQL column/expression for agent_id
|
|
2161
|
+
name_col: SQL column/expression for name
|
|
2162
|
+
|
|
2163
|
+
Returns:
|
|
2164
|
+
SQLAlchemy CASE expression returning the component level as an integer.
|
|
2165
|
+
"""
|
|
2166
|
+
from sqlalchemy import and_, case, or_
|
|
2167
|
+
|
|
2168
|
+
is_root_name = or_(name_col.contains(".run"), name_col.contains(".arun"))
|
|
2169
|
+
|
|
2170
|
+
return case(
|
|
2171
|
+
# Workflow root (level 3)
|
|
2172
|
+
(and_(workflow_id_col.isnot(None), is_root_name), 3),
|
|
2173
|
+
# Team root (level 2)
|
|
2174
|
+
(and_(team_id_col.isnot(None), is_root_name), 2),
|
|
2175
|
+
# Agent root (level 1)
|
|
2176
|
+
(and_(agent_id_col.isnot(None), is_root_name), 1),
|
|
2177
|
+
# Child span or unknown (level 0)
|
|
2178
|
+
else_=0,
|
|
2179
|
+
)
|
|
2180
|
+
|
|
2181
|
+
def upsert_trace(self, trace: "Trace") -> None:
|
|
2182
|
+
"""Create or update a single trace record in the database.
|
|
2183
|
+
|
|
2184
|
+
Uses INSERT ... ON CONFLICT DO UPDATE (upsert) to handle concurrent inserts
|
|
2185
|
+
atomically and avoid race conditions.
|
|
2186
|
+
|
|
2187
|
+
Args:
|
|
2188
|
+
trace: The Trace object to store (one per trace_id).
|
|
2189
|
+
"""
|
|
2190
|
+
from sqlalchemy import case
|
|
2191
|
+
|
|
2192
|
+
try:
|
|
2193
|
+
table = self._get_table(table_type="traces", create_table_if_not_found=True)
|
|
2194
|
+
if table is None:
|
|
2195
|
+
return
|
|
2196
|
+
|
|
2197
|
+
trace_dict = trace.to_dict()
|
|
2198
|
+
trace_dict.pop("total_spans", None)
|
|
2199
|
+
trace_dict.pop("error_count", None)
|
|
2200
|
+
|
|
2201
|
+
with self.Session() as sess, sess.begin():
|
|
2202
|
+
# Use upsert to handle concurrent inserts atomically
|
|
2203
|
+
# On conflict, update fields while preserving existing non-null context values
|
|
2204
|
+
# and keeping the earliest start_time
|
|
2205
|
+
insert_stmt = sqlite.insert(table).values(trace_dict)
|
|
2206
|
+
|
|
2207
|
+
# Build component level expressions for comparing trace priority
|
|
2208
|
+
new_level = self._get_trace_component_level_expr(
|
|
2209
|
+
insert_stmt.excluded.workflow_id,
|
|
2210
|
+
insert_stmt.excluded.team_id,
|
|
2211
|
+
insert_stmt.excluded.agent_id,
|
|
2212
|
+
insert_stmt.excluded.name,
|
|
2213
|
+
)
|
|
2214
|
+
existing_level = self._get_trace_component_level_expr(
|
|
2215
|
+
table.c.workflow_id,
|
|
2216
|
+
table.c.team_id,
|
|
2217
|
+
table.c.agent_id,
|
|
2218
|
+
table.c.name,
|
|
2219
|
+
)
|
|
2220
|
+
|
|
2221
|
+
# Build the ON CONFLICT DO UPDATE clause
|
|
2222
|
+
# Use MIN for start_time, MAX for end_time to capture full trace duration
|
|
2223
|
+
# SQLite stores timestamps as ISO strings, so string comparison works for ISO format
|
|
2224
|
+
# Duration is calculated as: (MAX(end_time) - MIN(start_time)) in milliseconds
|
|
2225
|
+
# SQLite doesn't have epoch extraction, so we calculate duration using julianday
|
|
2226
|
+
upsert_stmt = insert_stmt.on_conflict_do_update(
|
|
2227
|
+
index_elements=["trace_id"],
|
|
2228
|
+
set_={
|
|
2229
|
+
"end_time": func.max(table.c.end_time, insert_stmt.excluded.end_time),
|
|
2230
|
+
"start_time": func.min(table.c.start_time, insert_stmt.excluded.start_time),
|
|
2231
|
+
# Calculate duration in milliseconds using julianday (SQLite-specific)
|
|
2232
|
+
# julianday returns days, so multiply by 86400000 to get milliseconds
|
|
2233
|
+
"duration_ms": (
|
|
2234
|
+
func.julianday(func.max(table.c.end_time, insert_stmt.excluded.end_time))
|
|
2235
|
+
- func.julianday(func.min(table.c.start_time, insert_stmt.excluded.start_time))
|
|
2236
|
+
)
|
|
2237
|
+
* 86400000,
|
|
2238
|
+
"status": insert_stmt.excluded.status,
|
|
2239
|
+
# Update name only if new trace is from a higher-level component
|
|
2240
|
+
# Priority: workflow (3) > team (2) > agent (1) > child spans (0)
|
|
2241
|
+
"name": case(
|
|
2242
|
+
(new_level > existing_level, insert_stmt.excluded.name),
|
|
2243
|
+
else_=table.c.name,
|
|
2244
|
+
),
|
|
2245
|
+
# Preserve existing non-null context values using COALESCE
|
|
2246
|
+
"run_id": func.coalesce(insert_stmt.excluded.run_id, table.c.run_id),
|
|
2247
|
+
"session_id": func.coalesce(insert_stmt.excluded.session_id, table.c.session_id),
|
|
2248
|
+
"user_id": func.coalesce(insert_stmt.excluded.user_id, table.c.user_id),
|
|
2249
|
+
"agent_id": func.coalesce(insert_stmt.excluded.agent_id, table.c.agent_id),
|
|
2250
|
+
"team_id": func.coalesce(insert_stmt.excluded.team_id, table.c.team_id),
|
|
2251
|
+
"workflow_id": func.coalesce(insert_stmt.excluded.workflow_id, table.c.workflow_id),
|
|
2252
|
+
},
|
|
2253
|
+
)
|
|
2254
|
+
sess.execute(upsert_stmt)
|
|
2255
|
+
|
|
2256
|
+
except Exception as e:
|
|
2257
|
+
log_error(f"Error creating trace: {e}")
|
|
2258
|
+
# Don't raise - tracing should not break the main application flow
|
|
2259
|
+
|
|
2260
|
+
def get_trace(
|
|
2261
|
+
self,
|
|
2262
|
+
trace_id: Optional[str] = None,
|
|
2263
|
+
run_id: Optional[str] = None,
|
|
2264
|
+
):
|
|
2265
|
+
"""Get a single trace by trace_id or other filters.
|
|
2266
|
+
|
|
2267
|
+
Args:
|
|
2268
|
+
trace_id: The unique trace identifier.
|
|
2269
|
+
run_id: Filter by run ID (returns first match).
|
|
2270
|
+
|
|
2271
|
+
Returns:
|
|
2272
|
+
Optional[Trace]: The trace if found, None otherwise.
|
|
2273
|
+
|
|
2274
|
+
Note:
|
|
2275
|
+
If multiple filters are provided, trace_id takes precedence.
|
|
2276
|
+
For other filters, the most recent trace is returned.
|
|
2277
|
+
"""
|
|
2278
|
+
try:
|
|
2279
|
+
from agno.tracing.schemas import Trace
|
|
2280
|
+
|
|
2281
|
+
table = self._get_table(table_type="traces")
|
|
2282
|
+
if table is None:
|
|
2283
|
+
return None
|
|
2284
|
+
|
|
2285
|
+
# Get spans table for JOIN
|
|
2286
|
+
spans_table = self._get_table(table_type="spans")
|
|
2287
|
+
|
|
2288
|
+
with self.Session() as sess:
|
|
2289
|
+
# Build query with aggregated span counts
|
|
2290
|
+
stmt = self._get_traces_base_query(table, spans_table)
|
|
2291
|
+
|
|
2292
|
+
if trace_id:
|
|
2293
|
+
stmt = stmt.where(table.c.trace_id == trace_id)
|
|
2294
|
+
elif run_id:
|
|
2295
|
+
stmt = stmt.where(table.c.run_id == run_id)
|
|
2296
|
+
else:
|
|
2297
|
+
log_debug("get_trace called without any filter parameters")
|
|
2298
|
+
return None
|
|
2299
|
+
|
|
2300
|
+
# Order by most recent and get first result
|
|
2301
|
+
stmt = stmt.order_by(table.c.start_time.desc()).limit(1)
|
|
2302
|
+
result = sess.execute(stmt).fetchone()
|
|
2303
|
+
|
|
2304
|
+
if result:
|
|
2305
|
+
return Trace.from_dict(dict(result._mapping))
|
|
2306
|
+
return None
|
|
2307
|
+
|
|
2308
|
+
except Exception as e:
|
|
2309
|
+
log_error(f"Error getting trace: {e}")
|
|
2310
|
+
return None
|
|
2311
|
+
|
|
2312
|
+
def get_traces(
|
|
2313
|
+
self,
|
|
2314
|
+
run_id: Optional[str] = None,
|
|
2315
|
+
session_id: Optional[str] = None,
|
|
2316
|
+
user_id: Optional[str] = None,
|
|
2317
|
+
agent_id: Optional[str] = None,
|
|
2318
|
+
team_id: Optional[str] = None,
|
|
2319
|
+
workflow_id: Optional[str] = None,
|
|
2320
|
+
status: Optional[str] = None,
|
|
2321
|
+
start_time: Optional[datetime] = None,
|
|
2322
|
+
end_time: Optional[datetime] = None,
|
|
2323
|
+
limit: Optional[int] = 20,
|
|
2324
|
+
page: Optional[int] = 1,
|
|
2325
|
+
) -> tuple[List, int]:
|
|
2326
|
+
"""Get traces matching the provided filters with pagination.
|
|
2327
|
+
|
|
2328
|
+
Args:
|
|
2329
|
+
run_id: Filter by run ID.
|
|
2330
|
+
session_id: Filter by session ID.
|
|
2331
|
+
user_id: Filter by user ID.
|
|
2332
|
+
agent_id: Filter by agent ID.
|
|
2333
|
+
team_id: Filter by team ID.
|
|
2334
|
+
workflow_id: Filter by workflow ID.
|
|
2335
|
+
status: Filter by status (OK, ERROR, UNSET).
|
|
2336
|
+
start_time: Filter traces starting after this datetime.
|
|
2337
|
+
end_time: Filter traces ending before this datetime.
|
|
2338
|
+
limit: Maximum number of traces to return per page.
|
|
2339
|
+
page: Page number (1-indexed).
|
|
2340
|
+
|
|
2341
|
+
Returns:
|
|
2342
|
+
tuple[List[Trace], int]: Tuple of (list of matching traces, total count).
|
|
2343
|
+
"""
|
|
2344
|
+
try:
|
|
2345
|
+
from sqlalchemy import func
|
|
2346
|
+
|
|
2347
|
+
from agno.tracing.schemas import Trace
|
|
2348
|
+
|
|
2349
|
+
log_debug(
|
|
2350
|
+
f"get_traces called with filters: run_id={run_id}, session_id={session_id}, user_id={user_id}, agent_id={agent_id}, page={page}, limit={limit}"
|
|
2351
|
+
)
|
|
2352
|
+
|
|
2353
|
+
table = self._get_table(table_type="traces")
|
|
2354
|
+
if table is None:
|
|
2355
|
+
log_debug(" Traces table not found")
|
|
2356
|
+
return [], 0
|
|
2357
|
+
|
|
2358
|
+
# Get spans table for JOIN
|
|
2359
|
+
spans_table = self._get_table(table_type="spans")
|
|
2360
|
+
|
|
2361
|
+
with self.Session() as sess:
|
|
2362
|
+
# Build base query with aggregated span counts
|
|
2363
|
+
base_stmt = self._get_traces_base_query(table, spans_table)
|
|
2364
|
+
|
|
2365
|
+
# Apply filters
|
|
2366
|
+
if run_id:
|
|
2367
|
+
base_stmt = base_stmt.where(table.c.run_id == run_id)
|
|
2368
|
+
if session_id:
|
|
2369
|
+
base_stmt = base_stmt.where(table.c.session_id == session_id)
|
|
2370
|
+
if user_id:
|
|
2371
|
+
base_stmt = base_stmt.where(table.c.user_id == user_id)
|
|
2372
|
+
if agent_id:
|
|
2373
|
+
base_stmt = base_stmt.where(table.c.agent_id == agent_id)
|
|
2374
|
+
if team_id:
|
|
2375
|
+
base_stmt = base_stmt.where(table.c.team_id == team_id)
|
|
2376
|
+
if workflow_id:
|
|
2377
|
+
base_stmt = base_stmt.where(table.c.workflow_id == workflow_id)
|
|
2378
|
+
if status:
|
|
2379
|
+
base_stmt = base_stmt.where(table.c.status == status)
|
|
2380
|
+
if start_time:
|
|
2381
|
+
# Convert datetime to ISO string for comparison
|
|
2382
|
+
base_stmt = base_stmt.where(table.c.start_time >= start_time.isoformat())
|
|
2383
|
+
if end_time:
|
|
2384
|
+
# Convert datetime to ISO string for comparison
|
|
2385
|
+
base_stmt = base_stmt.where(table.c.end_time <= end_time.isoformat())
|
|
2386
|
+
|
|
2387
|
+
# Get total count
|
|
2388
|
+
count_stmt = select(func.count()).select_from(base_stmt.alias())
|
|
2389
|
+
total_count = sess.execute(count_stmt).scalar() or 0
|
|
2390
|
+
|
|
2391
|
+
# Apply pagination
|
|
2392
|
+
offset = (page - 1) * limit if page and limit else 0
|
|
2393
|
+
paginated_stmt = base_stmt.order_by(table.c.start_time.desc()).limit(limit).offset(offset)
|
|
2394
|
+
|
|
2395
|
+
results = sess.execute(paginated_stmt).fetchall()
|
|
2396
|
+
|
|
2397
|
+
traces = [Trace.from_dict(dict(row._mapping)) for row in results]
|
|
2398
|
+
return traces, total_count
|
|
2399
|
+
|
|
2400
|
+
except Exception as e:
|
|
2401
|
+
log_error(f"Error getting traces: {e}")
|
|
2402
|
+
return [], 0
|
|
2403
|
+
|
|
2404
|
+
def get_trace_stats(
|
|
2405
|
+
self,
|
|
2406
|
+
user_id: Optional[str] = None,
|
|
2407
|
+
agent_id: Optional[str] = None,
|
|
2408
|
+
team_id: Optional[str] = None,
|
|
2409
|
+
workflow_id: Optional[str] = None,
|
|
2410
|
+
start_time: Optional[datetime] = None,
|
|
2411
|
+
end_time: Optional[datetime] = None,
|
|
2412
|
+
limit: Optional[int] = 20,
|
|
2413
|
+
page: Optional[int] = 1,
|
|
2414
|
+
) -> tuple[List[Dict[str, Any]], int]:
|
|
2415
|
+
"""Get trace statistics grouped by session.
|
|
2416
|
+
|
|
2417
|
+
Args:
|
|
2418
|
+
user_id: Filter by user ID.
|
|
2419
|
+
agent_id: Filter by agent ID.
|
|
2420
|
+
team_id: Filter by team ID.
|
|
2421
|
+
workflow_id: Filter by workflow ID.
|
|
2422
|
+
start_time: Filter sessions with traces created after this datetime.
|
|
2423
|
+
end_time: Filter sessions with traces created before this datetime.
|
|
2424
|
+
limit: Maximum number of sessions to return per page.
|
|
2425
|
+
page: Page number (1-indexed).
|
|
2426
|
+
|
|
2427
|
+
Returns:
|
|
2428
|
+
tuple[List[Dict], int]: Tuple of (list of session stats dicts, total count).
|
|
2429
|
+
"""
|
|
2430
|
+
try:
|
|
2431
|
+
from sqlalchemy import func
|
|
2432
|
+
|
|
2433
|
+
table = self._get_table(table_type="traces")
|
|
2434
|
+
if table is None:
|
|
2435
|
+
log_debug("Traces table not found")
|
|
2436
|
+
return [], 0
|
|
2437
|
+
|
|
2438
|
+
with self.Session() as sess:
|
|
2439
|
+
# Build base query grouped by session_id
|
|
2440
|
+
base_stmt = (
|
|
2441
|
+
select(
|
|
2442
|
+
table.c.session_id,
|
|
2443
|
+
table.c.user_id,
|
|
2444
|
+
table.c.agent_id,
|
|
2445
|
+
table.c.team_id,
|
|
2446
|
+
table.c.workflow_id,
|
|
2447
|
+
func.count(table.c.trace_id).label("total_traces"),
|
|
2448
|
+
func.min(table.c.created_at).label("first_trace_at"),
|
|
2449
|
+
func.max(table.c.created_at).label("last_trace_at"),
|
|
2450
|
+
)
|
|
2451
|
+
.where(table.c.session_id.isnot(None)) # Only sessions with session_id
|
|
2452
|
+
.group_by(
|
|
2453
|
+
table.c.session_id, table.c.user_id, table.c.agent_id, table.c.team_id, table.c.workflow_id
|
|
2454
|
+
)
|
|
2455
|
+
)
|
|
2456
|
+
|
|
2457
|
+
# Apply filters
|
|
2458
|
+
if user_id:
|
|
2459
|
+
base_stmt = base_stmt.where(table.c.user_id == user_id)
|
|
2460
|
+
if workflow_id:
|
|
2461
|
+
base_stmt = base_stmt.where(table.c.workflow_id == workflow_id)
|
|
2462
|
+
if team_id:
|
|
2463
|
+
base_stmt = base_stmt.where(table.c.team_id == team_id)
|
|
2464
|
+
if agent_id:
|
|
2465
|
+
base_stmt = base_stmt.where(table.c.agent_id == agent_id)
|
|
2466
|
+
if start_time:
|
|
2467
|
+
# Convert datetime to ISO string for comparison
|
|
2468
|
+
base_stmt = base_stmt.where(table.c.created_at >= start_time.isoformat())
|
|
2469
|
+
if end_time:
|
|
2470
|
+
# Convert datetime to ISO string for comparison
|
|
2471
|
+
base_stmt = base_stmt.where(table.c.created_at <= end_time.isoformat())
|
|
2472
|
+
|
|
2473
|
+
# Get total count of sessions
|
|
2474
|
+
count_stmt = select(func.count()).select_from(base_stmt.alias())
|
|
2475
|
+
total_count = sess.execute(count_stmt).scalar() or 0
|
|
2476
|
+
|
|
2477
|
+
# Apply pagination and ordering
|
|
2478
|
+
offset = (page - 1) * limit if page and limit else 0
|
|
2479
|
+
paginated_stmt = base_stmt.order_by(func.max(table.c.created_at).desc()).limit(limit).offset(offset)
|
|
2480
|
+
|
|
2481
|
+
results = sess.execute(paginated_stmt).fetchall()
|
|
2482
|
+
|
|
2483
|
+
# Convert to list of dicts with datetime objects
|
|
2484
|
+
from datetime import datetime
|
|
2485
|
+
|
|
2486
|
+
stats_list = []
|
|
2487
|
+
for row in results:
|
|
2488
|
+
# Convert ISO strings to datetime objects
|
|
2489
|
+
first_trace_at_str = row.first_trace_at
|
|
2490
|
+
last_trace_at_str = row.last_trace_at
|
|
2491
|
+
|
|
2492
|
+
# Parse ISO format strings to datetime objects
|
|
2493
|
+
first_trace_at = datetime.fromisoformat(first_trace_at_str.replace("Z", "+00:00"))
|
|
2494
|
+
last_trace_at = datetime.fromisoformat(last_trace_at_str.replace("Z", "+00:00"))
|
|
2495
|
+
|
|
2496
|
+
stats_list.append(
|
|
2497
|
+
{
|
|
2498
|
+
"session_id": row.session_id,
|
|
2499
|
+
"user_id": row.user_id,
|
|
2500
|
+
"agent_id": row.agent_id,
|
|
2501
|
+
"team_id": row.team_id,
|
|
2502
|
+
"workflow_id": row.workflow_id,
|
|
2503
|
+
"total_traces": row.total_traces,
|
|
2504
|
+
"first_trace_at": first_trace_at,
|
|
2505
|
+
"last_trace_at": last_trace_at,
|
|
2506
|
+
}
|
|
2507
|
+
)
|
|
2508
|
+
|
|
2509
|
+
return stats_list, total_count
|
|
2510
|
+
|
|
2511
|
+
except Exception as e:
|
|
2512
|
+
log_error(f"Error getting trace stats: {e}")
|
|
2513
|
+
return [], 0
|
|
2514
|
+
|
|
2515
|
+
# -- Span methods --
|
|
2516
|
+
|
|
2517
|
+
def create_span(self, span: "Span") -> None:
|
|
2518
|
+
"""Create a single span in the database.
|
|
2519
|
+
|
|
2520
|
+
Args:
|
|
2521
|
+
span: The Span object to store.
|
|
2522
|
+
"""
|
|
2523
|
+
try:
|
|
2524
|
+
table = self._get_table(table_type="spans", create_table_if_not_found=True)
|
|
2525
|
+
if table is None:
|
|
2526
|
+
return
|
|
2527
|
+
|
|
2528
|
+
with self.Session() as sess, sess.begin():
|
|
2529
|
+
stmt = sqlite.insert(table).values(span.to_dict())
|
|
2530
|
+
sess.execute(stmt)
|
|
2531
|
+
|
|
2532
|
+
except Exception as e:
|
|
2533
|
+
log_error(f"Error creating span: {e}")
|
|
2534
|
+
|
|
2535
|
+
def create_spans(self, spans: List) -> None:
|
|
2536
|
+
"""Create multiple spans in the database as a batch.
|
|
2537
|
+
|
|
2538
|
+
Args:
|
|
2539
|
+
spans: List of Span objects to store.
|
|
2540
|
+
"""
|
|
2541
|
+
if not spans:
|
|
2542
|
+
return
|
|
2543
|
+
|
|
2544
|
+
try:
|
|
2545
|
+
table = self._get_table(table_type="spans", create_table_if_not_found=True)
|
|
2546
|
+
if table is None:
|
|
2547
|
+
return
|
|
2548
|
+
|
|
2549
|
+
with self.Session() as sess, sess.begin():
|
|
2550
|
+
for span in spans:
|
|
2551
|
+
stmt = sqlite.insert(table).values(span.to_dict())
|
|
2552
|
+
sess.execute(stmt)
|
|
2553
|
+
|
|
2554
|
+
except Exception as e:
|
|
2555
|
+
log_error(f"Error creating spans batch: {e}")
|
|
2556
|
+
|
|
2557
|
+
def get_span(self, span_id: str):
|
|
2558
|
+
"""Get a single span by its span_id.
|
|
2559
|
+
|
|
2560
|
+
Args:
|
|
2561
|
+
span_id: The unique span identifier.
|
|
2562
|
+
|
|
2563
|
+
Returns:
|
|
2564
|
+
Optional[Span]: The span if found, None otherwise.
|
|
2565
|
+
"""
|
|
2566
|
+
try:
|
|
2567
|
+
from agno.tracing.schemas import Span
|
|
2568
|
+
|
|
2569
|
+
table = self._get_table(table_type="spans")
|
|
2570
|
+
if table is None:
|
|
2571
|
+
return None
|
|
2572
|
+
|
|
2573
|
+
with self.Session() as sess:
|
|
2574
|
+
stmt = table.select().where(table.c.span_id == span_id)
|
|
2575
|
+
result = sess.execute(stmt).fetchone()
|
|
2576
|
+
if result:
|
|
2577
|
+
return Span.from_dict(dict(result._mapping))
|
|
2578
|
+
return None
|
|
2579
|
+
|
|
2580
|
+
except Exception as e:
|
|
2581
|
+
log_error(f"Error getting span: {e}")
|
|
2582
|
+
return None
|
|
2583
|
+
|
|
2584
|
+
def get_spans(
|
|
2585
|
+
self,
|
|
2586
|
+
trace_id: Optional[str] = None,
|
|
2587
|
+
parent_span_id: Optional[str] = None,
|
|
2588
|
+
) -> List:
|
|
2589
|
+
"""Get spans matching the provided filters.
|
|
2590
|
+
|
|
2591
|
+
Args:
|
|
2592
|
+
trace_id: Filter by trace ID.
|
|
2593
|
+
parent_span_id: Filter by parent span ID.
|
|
2594
|
+
|
|
2595
|
+
Returns:
|
|
2596
|
+
List[Span]: List of matching spans.
|
|
2597
|
+
"""
|
|
2598
|
+
try:
|
|
2599
|
+
from agno.tracing.schemas import Span
|
|
2600
|
+
|
|
2601
|
+
table = self._get_table(table_type="spans")
|
|
2602
|
+
if table is None:
|
|
2603
|
+
return []
|
|
2604
|
+
|
|
2605
|
+
with self.Session() as sess:
|
|
2606
|
+
stmt = table.select()
|
|
2607
|
+
|
|
2608
|
+
# Apply filters
|
|
2609
|
+
if trace_id:
|
|
2610
|
+
stmt = stmt.where(table.c.trace_id == trace_id)
|
|
2611
|
+
if parent_span_id:
|
|
2612
|
+
stmt = stmt.where(table.c.parent_span_id == parent_span_id)
|
|
2613
|
+
|
|
2614
|
+
results = sess.execute(stmt).fetchall()
|
|
2615
|
+
return [Span.from_dict(dict(row._mapping)) for row in results]
|
|
2616
|
+
|
|
2617
|
+
except Exception as e:
|
|
2618
|
+
log_error(f"Error getting spans: {e}")
|
|
2619
|
+
return []
|
|
2620
|
+
|
|
2621
|
+
# -- Migrations --
|
|
2622
|
+
|
|
2623
|
+
def migrate_table_from_v1_to_v2(self, v1_db_schema: str, v1_table_name: str, v1_table_type: str):
|
|
2624
|
+
"""Migrate all content in the given table to the right v2 table"""
|
|
2625
|
+
|
|
2626
|
+
from agno.db.migrations.v1_to_v2 import (
|
|
2627
|
+
get_all_table_content,
|
|
2628
|
+
parse_agent_sessions,
|
|
2629
|
+
parse_memories,
|
|
2630
|
+
parse_team_sessions,
|
|
2631
|
+
parse_workflow_sessions,
|
|
2632
|
+
)
|
|
2633
|
+
|
|
2634
|
+
# Get all content from the old table
|
|
2635
|
+
old_content: list[dict[str, Any]] = get_all_table_content(
|
|
2636
|
+
db=self,
|
|
2637
|
+
db_schema=v1_db_schema,
|
|
2638
|
+
table_name=v1_table_name,
|
|
2639
|
+
)
|
|
2640
|
+
if not old_content:
|
|
2641
|
+
log_info(f"No content to migrate from table {v1_table_name}")
|
|
2642
|
+
return
|
|
2643
|
+
|
|
2644
|
+
# Parse the content into the new format
|
|
2645
|
+
memories: List[UserMemory] = []
|
|
2646
|
+
sessions: Sequence[Union[AgentSession, TeamSession, WorkflowSession]] = []
|
|
2647
|
+
if v1_table_type == "agent_sessions":
|
|
2648
|
+
sessions = parse_agent_sessions(old_content)
|
|
2649
|
+
elif v1_table_type == "team_sessions":
|
|
2650
|
+
sessions = parse_team_sessions(old_content)
|
|
2651
|
+
elif v1_table_type == "workflow_sessions":
|
|
2652
|
+
sessions = parse_workflow_sessions(old_content)
|
|
2653
|
+
elif v1_table_type == "memories":
|
|
2654
|
+
memories = parse_memories(old_content)
|
|
2655
|
+
else:
|
|
2656
|
+
raise ValueError(f"Invalid table type: {v1_table_type}")
|
|
2657
|
+
|
|
2658
|
+
# Insert the new content into the new table
|
|
2659
|
+
if v1_table_type == "agent_sessions":
|
|
2660
|
+
for session in sessions:
|
|
2661
|
+
self.upsert_session(session)
|
|
2662
|
+
log_info(f"Migrated {len(sessions)} Agent sessions to table: {self.session_table_name}")
|
|
2663
|
+
|
|
2664
|
+
elif v1_table_type == "team_sessions":
|
|
2665
|
+
for session in sessions:
|
|
2666
|
+
self.upsert_session(session)
|
|
2667
|
+
log_info(f"Migrated {len(sessions)} Team sessions to table: {self.session_table_name}")
|
|
2668
|
+
|
|
2669
|
+
elif v1_table_type == "workflow_sessions":
|
|
2670
|
+
for session in sessions:
|
|
2671
|
+
self.upsert_session(session)
|
|
2672
|
+
log_info(f"Migrated {len(sessions)} Workflow sessions to table: {self.session_table_name}")
|
|
2673
|
+
|
|
2674
|
+
elif v1_table_type == "memories":
|
|
2675
|
+
for memory in memories:
|
|
2676
|
+
self.upsert_user_memory(memory)
|
|
2677
|
+
log_info(f"Migrated {len(memories)} memories to table: {self.memory_table}")
|
|
2678
|
+
|
|
2679
|
+
# -- Culture methods --
|
|
2680
|
+
|
|
2681
|
+
def clear_cultural_knowledge(self) -> None:
|
|
2682
|
+
"""Delete all cultural artifacts from the database.
|
|
2683
|
+
|
|
2684
|
+
Raises:
|
|
2685
|
+
Exception: If an error occurs during deletion.
|
|
2686
|
+
"""
|
|
2687
|
+
try:
|
|
2688
|
+
table = self._get_table(table_type="culture")
|
|
2689
|
+
if table is None:
|
|
2690
|
+
return
|
|
2691
|
+
|
|
2692
|
+
with self.Session() as sess, sess.begin():
|
|
2693
|
+
sess.execute(table.delete())
|
|
2694
|
+
|
|
2695
|
+
except Exception as e:
|
|
2696
|
+
from agno.utils.log import log_warning
|
|
2697
|
+
|
|
2698
|
+
log_warning(f"Exception deleting all cultural artifacts: {e}")
|
|
2699
|
+
raise e
|
|
2700
|
+
|
|
2701
|
+
def delete_cultural_knowledge(self, id: str) -> None:
|
|
2702
|
+
"""Delete a cultural artifact from the database.
|
|
2703
|
+
|
|
2704
|
+
Args:
|
|
2705
|
+
id (str): The ID of the cultural artifact to delete.
|
|
2706
|
+
|
|
2707
|
+
Raises:
|
|
2708
|
+
Exception: If an error occurs during deletion.
|
|
2709
|
+
"""
|
|
2710
|
+
try:
|
|
2711
|
+
table = self._get_table(table_type="culture")
|
|
2712
|
+
if table is None:
|
|
2713
|
+
return
|
|
2714
|
+
|
|
2715
|
+
with self.Session() as sess, sess.begin():
|
|
2716
|
+
delete_stmt = table.delete().where(table.c.id == id)
|
|
2717
|
+
result = sess.execute(delete_stmt)
|
|
2718
|
+
|
|
2719
|
+
success = result.rowcount > 0
|
|
2720
|
+
if success:
|
|
2721
|
+
log_debug(f"Successfully deleted cultural artifact id: {id}")
|
|
2722
|
+
else:
|
|
2723
|
+
log_debug(f"No cultural artifact found with id: {id}")
|
|
2724
|
+
|
|
2725
|
+
except Exception as e:
|
|
2726
|
+
log_error(f"Error deleting cultural artifact: {e}")
|
|
2727
|
+
raise e
|
|
2728
|
+
|
|
2729
|
+
def get_cultural_knowledge(
|
|
2730
|
+
self, id: str, deserialize: Optional[bool] = True
|
|
2731
|
+
) -> Optional[Union[CulturalKnowledge, Dict[str, Any]]]:
|
|
2732
|
+
"""Get a cultural artifact from the database.
|
|
2733
|
+
|
|
2734
|
+
Args:
|
|
2735
|
+
id (str): The ID of the cultural artifact to get.
|
|
2736
|
+
deserialize (Optional[bool]): Whether to serialize the cultural artifact. Defaults to True.
|
|
2737
|
+
|
|
2738
|
+
Returns:
|
|
2739
|
+
Optional[CulturalKnowledge]: The cultural artifact, or None if it doesn't exist.
|
|
2740
|
+
|
|
2741
|
+
Raises:
|
|
2742
|
+
Exception: If an error occurs during retrieval.
|
|
2743
|
+
"""
|
|
2744
|
+
try:
|
|
2745
|
+
table = self._get_table(table_type="culture")
|
|
2746
|
+
if table is None:
|
|
2747
|
+
return None
|
|
2748
|
+
|
|
2749
|
+
with self.Session() as sess, sess.begin():
|
|
2750
|
+
stmt = select(table).where(table.c.id == id)
|
|
2751
|
+
result = sess.execute(stmt).fetchone()
|
|
2752
|
+
if result is None:
|
|
2753
|
+
return None
|
|
2754
|
+
|
|
2755
|
+
db_row = dict(result._mapping)
|
|
2756
|
+
if not db_row or not deserialize:
|
|
2757
|
+
return db_row
|
|
2758
|
+
|
|
2759
|
+
return deserialize_cultural_knowledge_from_db(db_row)
|
|
2760
|
+
|
|
2761
|
+
except Exception as e:
|
|
2762
|
+
log_error(f"Exception reading from cultural artifacts table: {e}")
|
|
2763
|
+
raise e
|
|
2764
|
+
|
|
2765
|
+
def get_all_cultural_knowledge(
|
|
2766
|
+
self,
|
|
2767
|
+
name: Optional[str] = None,
|
|
2768
|
+
agent_id: Optional[str] = None,
|
|
2769
|
+
team_id: Optional[str] = None,
|
|
2770
|
+
limit: Optional[int] = None,
|
|
2771
|
+
page: Optional[int] = None,
|
|
2772
|
+
sort_by: Optional[str] = None,
|
|
2773
|
+
sort_order: Optional[str] = None,
|
|
2774
|
+
deserialize: Optional[bool] = True,
|
|
2775
|
+
) -> Union[List[CulturalKnowledge], Tuple[List[Dict[str, Any]], int]]:
|
|
2776
|
+
"""Get all cultural artifacts from the database as CulturalNotion objects.
|
|
2777
|
+
|
|
2778
|
+
Args:
|
|
2779
|
+
name (Optional[str]): The name of the cultural artifact to filter by.
|
|
2780
|
+
agent_id (Optional[str]): The ID of the agent to filter by.
|
|
2781
|
+
team_id (Optional[str]): The ID of the team to filter by.
|
|
2782
|
+
limit (Optional[int]): The maximum number of cultural artifacts to return.
|
|
2783
|
+
page (Optional[int]): The page number.
|
|
2784
|
+
sort_by (Optional[str]): The column to sort by.
|
|
2785
|
+
sort_order (Optional[str]): The order to sort by.
|
|
2786
|
+
deserialize (Optional[bool]): Whether to serialize the cultural artifacts. Defaults to True.
|
|
2787
|
+
|
|
2788
|
+
Returns:
|
|
2789
|
+
Union[List[CulturalKnowledge], Tuple[List[Dict[str, Any]], int]]:
|
|
2790
|
+
- When deserialize=True: List of CulturalNotion objects
|
|
2791
|
+
- When deserialize=False: List of CulturalNotion dictionaries and total count
|
|
2792
|
+
|
|
2793
|
+
Raises:
|
|
2794
|
+
Exception: If an error occurs during retrieval.
|
|
2795
|
+
"""
|
|
2796
|
+
try:
|
|
2797
|
+
table = self._get_table(table_type="culture")
|
|
2798
|
+
if table is None:
|
|
2799
|
+
return [] if deserialize else ([], 0)
|
|
2800
|
+
|
|
2801
|
+
with self.Session() as sess, sess.begin():
|
|
2802
|
+
stmt = select(table)
|
|
2803
|
+
|
|
2804
|
+
# Filtering
|
|
2805
|
+
if name is not None:
|
|
2806
|
+
stmt = stmt.where(table.c.name == name)
|
|
2807
|
+
if agent_id is not None:
|
|
2808
|
+
stmt = stmt.where(table.c.agent_id == agent_id)
|
|
2809
|
+
if team_id is not None:
|
|
2810
|
+
stmt = stmt.where(table.c.team_id == team_id)
|
|
2811
|
+
|
|
2812
|
+
# Get total count after applying filtering
|
|
2813
|
+
count_stmt = select(func.count()).select_from(stmt.alias())
|
|
2814
|
+
total_count = sess.execute(count_stmt).scalar()
|
|
2815
|
+
|
|
2816
|
+
# Sorting
|
|
2817
|
+
stmt = apply_sorting(stmt, table, sort_by, sort_order)
|
|
2818
|
+
# Paginating
|
|
2819
|
+
if limit is not None:
|
|
2820
|
+
stmt = stmt.limit(limit)
|
|
2821
|
+
if page is not None:
|
|
2822
|
+
stmt = stmt.offset((page - 1) * limit)
|
|
2823
|
+
|
|
2824
|
+
result = sess.execute(stmt).fetchall()
|
|
2825
|
+
if not result:
|
|
2826
|
+
return [] if deserialize else ([], 0)
|
|
2827
|
+
|
|
2828
|
+
db_rows = [dict(record._mapping) for record in result]
|
|
2829
|
+
|
|
2830
|
+
if not deserialize:
|
|
2831
|
+
return db_rows, total_count
|
|
2832
|
+
|
|
2833
|
+
return [deserialize_cultural_knowledge_from_db(row) for row in db_rows]
|
|
2834
|
+
|
|
2835
|
+
except Exception as e:
|
|
2836
|
+
log_error(f"Error reading from cultural artifacts table: {e}")
|
|
2837
|
+
raise e
|
|
2838
|
+
|
|
2839
|
+
def upsert_cultural_knowledge(
|
|
2840
|
+
self, cultural_knowledge: CulturalKnowledge, deserialize: Optional[bool] = True
|
|
2841
|
+
) -> Optional[Union[CulturalKnowledge, Dict[str, Any]]]:
|
|
2842
|
+
"""Upsert a cultural artifact into the database.
|
|
2843
|
+
|
|
2844
|
+
Args:
|
|
2845
|
+
cultural_knowledge (CulturalKnowledge): The cultural artifact to upsert.
|
|
2846
|
+
deserialize (Optional[bool]): Whether to serialize the cultural artifact. Defaults to True.
|
|
2847
|
+
|
|
2848
|
+
Returns:
|
|
2849
|
+
Optional[Union[CulturalNotion, Dict[str, Any]]]:
|
|
2850
|
+
- When deserialize=True: CulturalNotion object
|
|
2851
|
+
- When deserialize=False: CulturalNotion dictionary
|
|
2852
|
+
|
|
2853
|
+
Raises:
|
|
2854
|
+
Exception: If an error occurs during upsert.
|
|
2855
|
+
"""
|
|
2856
|
+
try:
|
|
2857
|
+
table = self._get_table(table_type="culture", create_table_if_not_found=True)
|
|
2858
|
+
if table is None:
|
|
2859
|
+
return None
|
|
2860
|
+
|
|
2861
|
+
if cultural_knowledge.id is None:
|
|
2862
|
+
cultural_knowledge.id = str(uuid4())
|
|
2863
|
+
|
|
2864
|
+
# Serialize content, categories, and notes into a JSON string for DB storage (SQLite requires strings)
|
|
2865
|
+
content_json_str = serialize_cultural_knowledge_for_db(cultural_knowledge)
|
|
2866
|
+
|
|
2867
|
+
with self.Session() as sess, sess.begin():
|
|
2868
|
+
stmt = sqlite.insert(table).values(
|
|
2869
|
+
id=cultural_knowledge.id,
|
|
2870
|
+
name=cultural_knowledge.name,
|
|
2871
|
+
summary=cultural_knowledge.summary,
|
|
2872
|
+
content=content_json_str,
|
|
2873
|
+
metadata=cultural_knowledge.metadata,
|
|
2874
|
+
input=cultural_knowledge.input,
|
|
2875
|
+
created_at=cultural_knowledge.created_at,
|
|
2876
|
+
updated_at=int(time.time()),
|
|
2877
|
+
agent_id=cultural_knowledge.agent_id,
|
|
2878
|
+
team_id=cultural_knowledge.team_id,
|
|
2879
|
+
)
|
|
2880
|
+
stmt = stmt.on_conflict_do_update( # type: ignore
|
|
2881
|
+
index_elements=["id"],
|
|
2882
|
+
set_=dict(
|
|
2883
|
+
name=cultural_knowledge.name,
|
|
2884
|
+
summary=cultural_knowledge.summary,
|
|
2885
|
+
content=content_json_str,
|
|
2886
|
+
metadata=cultural_knowledge.metadata,
|
|
2887
|
+
input=cultural_knowledge.input,
|
|
2888
|
+
updated_at=int(time.time()),
|
|
2889
|
+
agent_id=cultural_knowledge.agent_id,
|
|
2890
|
+
team_id=cultural_knowledge.team_id,
|
|
2891
|
+
),
|
|
2892
|
+
).returning(table)
|
|
2893
|
+
|
|
2894
|
+
result = sess.execute(stmt)
|
|
2895
|
+
row = result.fetchone()
|
|
2896
|
+
|
|
2897
|
+
if row is None:
|
|
2898
|
+
return None
|
|
2899
|
+
|
|
2900
|
+
db_row: Dict[str, Any] = dict(row._mapping)
|
|
2901
|
+
if not db_row or not deserialize:
|
|
2902
|
+
return db_row
|
|
2903
|
+
|
|
2904
|
+
return deserialize_cultural_knowledge_from_db(db_row)
|
|
2905
|
+
|
|
2906
|
+
except Exception as e:
|
|
2907
|
+
log_error(f"Error upserting cultural knowledge: {e}")
|
|
2908
|
+
raise e
|