agno 2.2.13__py3-none-any.whl → 2.4.3__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/agent/__init__.py +6 -0
- agno/agent/agent.py +5252 -3145
- agno/agent/remote.py +525 -0
- agno/api/api.py +2 -0
- agno/client/__init__.py +3 -0
- agno/client/a2a/__init__.py +10 -0
- agno/client/a2a/client.py +554 -0
- agno/client/a2a/schemas.py +112 -0
- agno/client/a2a/utils.py +369 -0
- agno/client/os.py +2669 -0
- agno/compression/__init__.py +3 -0
- agno/compression/manager.py +247 -0
- agno/culture/manager.py +2 -2
- agno/db/base.py +927 -6
- agno/db/dynamo/dynamo.py +788 -2
- agno/db/dynamo/schemas.py +128 -0
- agno/db/dynamo/utils.py +26 -3
- agno/db/firestore/firestore.py +674 -50
- agno/db/firestore/schemas.py +41 -0
- agno/db/firestore/utils.py +25 -10
- agno/db/gcs_json/gcs_json_db.py +506 -3
- agno/db/gcs_json/utils.py +14 -2
- agno/db/in_memory/in_memory_db.py +203 -4
- agno/db/in_memory/utils.py +14 -2
- agno/db/json/json_db.py +498 -2
- agno/db/json/utils.py +14 -2
- agno/db/migrations/manager.py +199 -0
- agno/db/migrations/utils.py +19 -0
- agno/db/migrations/v1_to_v2.py +54 -16
- agno/db/migrations/versions/__init__.py +0 -0
- agno/db/migrations/versions/v2_3_0.py +977 -0
- agno/db/mongo/async_mongo.py +1013 -39
- agno/db/mongo/mongo.py +684 -4
- agno/db/mongo/schemas.py +48 -0
- agno/db/mongo/utils.py +17 -0
- agno/db/mysql/__init__.py +2 -1
- agno/db/mysql/async_mysql.py +2958 -0
- agno/db/mysql/mysql.py +722 -53
- agno/db/mysql/schemas.py +77 -11
- agno/db/mysql/utils.py +151 -8
- agno/db/postgres/async_postgres.py +1254 -137
- agno/db/postgres/postgres.py +2316 -93
- agno/db/postgres/schemas.py +153 -21
- agno/db/postgres/utils.py +22 -7
- agno/db/redis/redis.py +531 -3
- agno/db/redis/schemas.py +36 -0
- agno/db/redis/utils.py +31 -15
- agno/db/schemas/evals.py +1 -0
- agno/db/schemas/memory.py +20 -9
- agno/db/singlestore/schemas.py +70 -1
- agno/db/singlestore/singlestore.py +737 -74
- agno/db/singlestore/utils.py +13 -3
- agno/db/sqlite/async_sqlite.py +1069 -89
- agno/db/sqlite/schemas.py +133 -1
- agno/db/sqlite/sqlite.py +2203 -165
- agno/db/sqlite/utils.py +21 -11
- agno/db/surrealdb/models.py +25 -0
- agno/db/surrealdb/surrealdb.py +603 -1
- agno/db/utils.py +60 -0
- agno/eval/__init__.py +26 -3
- agno/eval/accuracy.py +25 -12
- agno/eval/agent_as_judge.py +871 -0
- agno/eval/base.py +29 -0
- agno/eval/performance.py +10 -4
- agno/eval/reliability.py +22 -13
- agno/eval/utils.py +2 -1
- agno/exceptions.py +42 -0
- agno/hooks/__init__.py +3 -0
- agno/hooks/decorator.py +164 -0
- agno/integrations/discord/client.py +13 -2
- agno/knowledge/__init__.py +4 -0
- agno/knowledge/chunking/code.py +90 -0
- agno/knowledge/chunking/document.py +65 -4
- agno/knowledge/chunking/fixed.py +4 -1
- agno/knowledge/chunking/markdown.py +102 -11
- agno/knowledge/chunking/recursive.py +2 -2
- agno/knowledge/chunking/semantic.py +130 -48
- agno/knowledge/chunking/strategy.py +18 -0
- agno/knowledge/embedder/azure_openai.py +0 -1
- agno/knowledge/embedder/google.py +1 -1
- agno/knowledge/embedder/mistral.py +1 -1
- agno/knowledge/embedder/nebius.py +1 -1
- agno/knowledge/embedder/openai.py +16 -12
- agno/knowledge/filesystem.py +412 -0
- agno/knowledge/knowledge.py +4261 -1199
- agno/knowledge/protocol.py +134 -0
- agno/knowledge/reader/arxiv_reader.py +3 -2
- agno/knowledge/reader/base.py +9 -7
- agno/knowledge/reader/csv_reader.py +91 -42
- agno/knowledge/reader/docx_reader.py +9 -10
- agno/knowledge/reader/excel_reader.py +225 -0
- agno/knowledge/reader/field_labeled_csv_reader.py +38 -48
- agno/knowledge/reader/firecrawl_reader.py +3 -2
- agno/knowledge/reader/json_reader.py +16 -22
- agno/knowledge/reader/markdown_reader.py +15 -14
- agno/knowledge/reader/pdf_reader.py +33 -28
- agno/knowledge/reader/pptx_reader.py +9 -10
- agno/knowledge/reader/reader_factory.py +135 -1
- agno/knowledge/reader/s3_reader.py +8 -16
- agno/knowledge/reader/tavily_reader.py +3 -3
- agno/knowledge/reader/text_reader.py +15 -14
- agno/knowledge/reader/utils/__init__.py +17 -0
- agno/knowledge/reader/utils/spreadsheet.py +114 -0
- agno/knowledge/reader/web_search_reader.py +8 -65
- agno/knowledge/reader/website_reader.py +16 -13
- agno/knowledge/reader/wikipedia_reader.py +36 -3
- agno/knowledge/reader/youtube_reader.py +3 -2
- agno/knowledge/remote_content/__init__.py +33 -0
- agno/knowledge/remote_content/config.py +266 -0
- agno/knowledge/remote_content/remote_content.py +105 -17
- agno/knowledge/utils.py +76 -22
- agno/learn/__init__.py +71 -0
- agno/learn/config.py +463 -0
- agno/learn/curate.py +185 -0
- agno/learn/machine.py +725 -0
- agno/learn/schemas.py +1114 -0
- agno/learn/stores/__init__.py +38 -0
- agno/learn/stores/decision_log.py +1156 -0
- agno/learn/stores/entity_memory.py +3275 -0
- agno/learn/stores/learned_knowledge.py +1583 -0
- agno/learn/stores/protocol.py +117 -0
- agno/learn/stores/session_context.py +1217 -0
- agno/learn/stores/user_memory.py +1495 -0
- agno/learn/stores/user_profile.py +1220 -0
- agno/learn/utils.py +209 -0
- agno/media.py +22 -6
- agno/memory/__init__.py +14 -1
- agno/memory/manager.py +223 -8
- 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/aimlapi.py +17 -0
- agno/models/anthropic/claude.py +434 -59
- agno/models/aws/bedrock.py +121 -20
- agno/models/aws/claude.py +131 -274
- agno/models/azure/ai_foundry.py +10 -6
- agno/models/azure/openai_chat.py +33 -10
- agno/models/base.py +1162 -561
- agno/models/cerebras/cerebras.py +120 -24
- agno/models/cerebras/cerebras_openai.py +21 -2
- agno/models/cohere/chat.py +65 -6
- agno/models/cometapi/cometapi.py +18 -1
- agno/models/dashscope/dashscope.py +2 -3
- agno/models/deepinfra/deepinfra.py +18 -1
- agno/models/deepseek/deepseek.py +69 -3
- agno/models/fireworks/fireworks.py +18 -1
- agno/models/google/gemini.py +959 -89
- agno/models/google/utils.py +22 -0
- agno/models/groq/groq.py +48 -18
- agno/models/huggingface/huggingface.py +17 -6
- agno/models/ibm/watsonx.py +16 -6
- agno/models/internlm/internlm.py +18 -1
- agno/models/langdb/langdb.py +13 -1
- agno/models/litellm/chat.py +88 -9
- agno/models/litellm/litellm_openai.py +18 -1
- agno/models/message.py +24 -5
- agno/models/meta/llama.py +40 -13
- agno/models/meta/llama_openai.py +22 -21
- agno/models/metrics.py +12 -0
- agno/models/mistral/mistral.py +8 -4
- agno/models/n1n/__init__.py +3 -0
- agno/models/n1n/n1n.py +57 -0
- agno/models/nebius/nebius.py +6 -7
- agno/models/nvidia/nvidia.py +20 -3
- agno/models/ollama/__init__.py +2 -0
- agno/models/ollama/chat.py +17 -6
- agno/models/ollama/responses.py +100 -0
- agno/models/openai/__init__.py +2 -0
- agno/models/openai/chat.py +117 -26
- agno/models/openai/open_responses.py +46 -0
- agno/models/openai/responses.py +110 -32
- agno/models/openrouter/__init__.py +2 -0
- agno/models/openrouter/openrouter.py +67 -2
- agno/models/openrouter/responses.py +146 -0
- agno/models/perplexity/perplexity.py +19 -1
- agno/models/portkey/portkey.py +7 -6
- agno/models/requesty/requesty.py +19 -2
- agno/models/response.py +20 -2
- agno/models/sambanova/sambanova.py +20 -3
- agno/models/siliconflow/siliconflow.py +19 -2
- agno/models/together/together.py +20 -3
- agno/models/vercel/v0.py +20 -3
- agno/models/vertexai/claude.py +124 -4
- agno/models/vllm/vllm.py +19 -14
- agno/models/xai/xai.py +19 -2
- agno/os/app.py +467 -137
- agno/os/auth.py +253 -5
- agno/os/config.py +22 -0
- agno/os/interfaces/a2a/a2a.py +7 -6
- agno/os/interfaces/a2a/router.py +635 -26
- agno/os/interfaces/a2a/utils.py +32 -33
- agno/os/interfaces/agui/agui.py +5 -3
- agno/os/interfaces/agui/router.py +26 -16
- agno/os/interfaces/agui/utils.py +97 -57
- agno/os/interfaces/base.py +7 -7
- agno/os/interfaces/slack/router.py +16 -7
- agno/os/interfaces/slack/slack.py +7 -7
- agno/os/interfaces/whatsapp/router.py +35 -7
- agno/os/interfaces/whatsapp/security.py +3 -1
- agno/os/interfaces/whatsapp/whatsapp.py +11 -8
- agno/os/managers.py +326 -0
- agno/os/mcp.py +652 -79
- agno/os/middleware/__init__.py +4 -0
- agno/os/middleware/jwt.py +718 -115
- agno/os/middleware/trailing_slash.py +27 -0
- agno/os/router.py +105 -1558
- agno/os/routers/agents/__init__.py +3 -0
- agno/os/routers/agents/router.py +655 -0
- agno/os/routers/agents/schema.py +288 -0
- agno/os/routers/components/__init__.py +3 -0
- agno/os/routers/components/components.py +475 -0
- agno/os/routers/database.py +155 -0
- agno/os/routers/evals/evals.py +111 -18
- agno/os/routers/evals/schemas.py +38 -5
- agno/os/routers/evals/utils.py +80 -11
- agno/os/routers/health.py +3 -3
- agno/os/routers/knowledge/knowledge.py +284 -35
- agno/os/routers/knowledge/schemas.py +14 -2
- agno/os/routers/memory/memory.py +274 -11
- agno/os/routers/memory/schemas.py +44 -3
- agno/os/routers/metrics/metrics.py +30 -15
- agno/os/routers/metrics/schemas.py +10 -6
- agno/os/routers/registry/__init__.py +3 -0
- agno/os/routers/registry/registry.py +337 -0
- agno/os/routers/session/session.py +143 -14
- agno/os/routers/teams/__init__.py +3 -0
- agno/os/routers/teams/router.py +550 -0
- agno/os/routers/teams/schema.py +280 -0
- agno/os/routers/traces/__init__.py +3 -0
- agno/os/routers/traces/schemas.py +414 -0
- agno/os/routers/traces/traces.py +549 -0
- agno/os/routers/workflows/__init__.py +3 -0
- agno/os/routers/workflows/router.py +757 -0
- agno/os/routers/workflows/schema.py +139 -0
- agno/os/schema.py +157 -584
- agno/os/scopes.py +469 -0
- agno/os/settings.py +3 -0
- agno/os/utils.py +574 -185
- agno/reasoning/anthropic.py +85 -1
- agno/reasoning/azure_ai_foundry.py +93 -1
- agno/reasoning/deepseek.py +102 -2
- agno/reasoning/default.py +6 -7
- agno/reasoning/gemini.py +87 -3
- agno/reasoning/groq.py +109 -2
- agno/reasoning/helpers.py +6 -7
- agno/reasoning/manager.py +1238 -0
- agno/reasoning/ollama.py +93 -1
- agno/reasoning/openai.py +115 -1
- agno/reasoning/vertexai.py +85 -1
- agno/registry/__init__.py +3 -0
- agno/registry/registry.py +68 -0
- agno/remote/__init__.py +3 -0
- agno/remote/base.py +581 -0
- agno/run/__init__.py +2 -4
- agno/run/agent.py +134 -19
- agno/run/base.py +49 -1
- agno/run/cancel.py +65 -52
- agno/run/cancellation_management/__init__.py +9 -0
- agno/run/cancellation_management/base.py +78 -0
- agno/run/cancellation_management/in_memory_cancellation_manager.py +100 -0
- agno/run/cancellation_management/redis_cancellation_manager.py +236 -0
- agno/run/requirement.py +181 -0
- agno/run/team.py +111 -19
- agno/run/workflow.py +2 -1
- agno/session/agent.py +57 -92
- agno/session/summary.py +1 -1
- agno/session/team.py +62 -115
- agno/session/workflow.py +353 -57
- agno/skills/__init__.py +17 -0
- agno/skills/agent_skills.py +377 -0
- agno/skills/errors.py +32 -0
- agno/skills/loaders/__init__.py +4 -0
- agno/skills/loaders/base.py +27 -0
- agno/skills/loaders/local.py +216 -0
- agno/skills/skill.py +65 -0
- agno/skills/utils.py +107 -0
- agno/skills/validator.py +277 -0
- agno/table.py +10 -0
- agno/team/__init__.py +5 -1
- agno/team/remote.py +447 -0
- agno/team/team.py +3769 -2202
- agno/tools/brandfetch.py +27 -18
- agno/tools/browserbase.py +225 -16
- agno/tools/crawl4ai.py +3 -0
- agno/tools/duckduckgo.py +25 -71
- agno/tools/exa.py +0 -21
- agno/tools/file.py +14 -13
- agno/tools/file_generation.py +12 -6
- agno/tools/firecrawl.py +15 -7
- agno/tools/function.py +94 -113
- agno/tools/google_bigquery.py +11 -2
- agno/tools/google_drive.py +4 -3
- agno/tools/knowledge.py +9 -4
- agno/tools/mcp/mcp.py +301 -18
- agno/tools/mcp/multi_mcp.py +269 -14
- agno/tools/mem0.py +11 -10
- agno/tools/memory.py +47 -46
- agno/tools/mlx_transcribe.py +10 -7
- agno/tools/models/nebius.py +5 -5
- agno/tools/models_labs.py +20 -10
- agno/tools/nano_banana.py +151 -0
- agno/tools/parallel.py +0 -7
- agno/tools/postgres.py +76 -36
- agno/tools/python.py +14 -6
- agno/tools/reasoning.py +30 -23
- agno/tools/redshift.py +406 -0
- agno/tools/shopify.py +1519 -0
- agno/tools/spotify.py +919 -0
- agno/tools/tavily.py +4 -1
- agno/tools/toolkit.py +253 -18
- agno/tools/websearch.py +93 -0
- agno/tools/website.py +1 -1
- agno/tools/wikipedia.py +1 -1
- agno/tools/workflow.py +56 -48
- agno/tools/yfinance.py +12 -11
- agno/tracing/__init__.py +12 -0
- agno/tracing/exporter.py +161 -0
- agno/tracing/schemas.py +276 -0
- agno/tracing/setup.py +112 -0
- agno/utils/agent.py +251 -10
- agno/utils/cryptography.py +22 -0
- agno/utils/dttm.py +33 -0
- agno/utils/events.py +264 -7
- agno/utils/hooks.py +111 -3
- agno/utils/http.py +161 -2
- agno/utils/mcp.py +49 -8
- agno/utils/media.py +22 -1
- agno/utils/models/ai_foundry.py +9 -2
- agno/utils/models/claude.py +20 -5
- agno/utils/models/cohere.py +9 -2
- agno/utils/models/llama.py +9 -2
- agno/utils/models/mistral.py +4 -2
- agno/utils/os.py +0 -0
- agno/utils/print_response/agent.py +99 -16
- agno/utils/print_response/team.py +223 -24
- agno/utils/print_response/workflow.py +0 -2
- agno/utils/prompts.py +8 -6
- agno/utils/remote.py +23 -0
- agno/utils/response.py +1 -13
- agno/utils/string.py +91 -2
- agno/utils/team.py +62 -12
- agno/utils/tokens.py +657 -0
- agno/vectordb/base.py +15 -2
- agno/vectordb/cassandra/cassandra.py +1 -1
- agno/vectordb/chroma/__init__.py +2 -1
- agno/vectordb/chroma/chromadb.py +468 -23
- agno/vectordb/clickhouse/clickhousedb.py +1 -1
- agno/vectordb/couchbase/couchbase.py +6 -2
- agno/vectordb/lancedb/lance_db.py +7 -38
- agno/vectordb/lightrag/lightrag.py +7 -6
- agno/vectordb/milvus/milvus.py +118 -84
- agno/vectordb/mongodb/__init__.py +2 -1
- agno/vectordb/mongodb/mongodb.py +14 -31
- agno/vectordb/pgvector/pgvector.py +120 -66
- agno/vectordb/pineconedb/pineconedb.py +2 -19
- agno/vectordb/qdrant/__init__.py +2 -1
- agno/vectordb/qdrant/qdrant.py +33 -56
- agno/vectordb/redis/__init__.py +2 -1
- agno/vectordb/redis/redisdb.py +19 -31
- agno/vectordb/singlestore/singlestore.py +17 -9
- agno/vectordb/surrealdb/surrealdb.py +2 -38
- agno/vectordb/weaviate/__init__.py +2 -1
- agno/vectordb/weaviate/weaviate.py +7 -3
- agno/workflow/__init__.py +5 -1
- agno/workflow/agent.py +2 -2
- agno/workflow/condition.py +12 -10
- agno/workflow/loop.py +28 -9
- agno/workflow/parallel.py +21 -13
- agno/workflow/remote.py +362 -0
- agno/workflow/router.py +12 -9
- agno/workflow/step.py +261 -36
- agno/workflow/steps.py +12 -8
- agno/workflow/types.py +40 -77
- agno/workflow/workflow.py +939 -213
- {agno-2.2.13.dist-info → agno-2.4.3.dist-info}/METADATA +134 -181
- agno-2.4.3.dist-info/RECORD +677 -0
- {agno-2.2.13.dist-info → agno-2.4.3.dist-info}/WHEEL +1 -1
- agno/tools/googlesearch.py +0 -98
- agno/tools/memori.py +0 -339
- agno-2.2.13.dist-info/RECORD +0 -575
- {agno-2.2.13.dist-info → agno-2.4.3.dist-info}/licenses/LICENSE +0 -0
- {agno-2.2.13.dist-info → agno-2.4.3.dist-info}/top_level.txt +0 -0
agno/db/mongo/mongo.py
CHANGED
|
@@ -1,8 +1,12 @@
|
|
|
1
1
|
import time
|
|
2
2
|
from datetime import date, datetime, timedelta, timezone
|
|
3
|
-
from
|
|
3
|
+
from importlib import metadata
|
|
4
|
+
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
|
|
4
5
|
from uuid import uuid4
|
|
5
6
|
|
|
7
|
+
if TYPE_CHECKING:
|
|
8
|
+
from agno.tracing.schemas import Span, Trace
|
|
9
|
+
|
|
6
10
|
from agno.db.base import BaseDb, SessionType
|
|
7
11
|
from agno.db.mongo.utils import (
|
|
8
12
|
apply_pagination,
|
|
@@ -28,10 +32,13 @@ try:
|
|
|
28
32
|
from pymongo import MongoClient, ReturnDocument
|
|
29
33
|
from pymongo.collection import Collection
|
|
30
34
|
from pymongo.database import Database
|
|
35
|
+
from pymongo.driver_info import DriverInfo
|
|
31
36
|
from pymongo.errors import OperationFailure
|
|
32
37
|
except ImportError:
|
|
33
38
|
raise ImportError("`pymongo` not installed. Please install it using `pip install pymongo`")
|
|
34
39
|
|
|
40
|
+
DRIVER_METADATA = DriverInfo(name="Agno", version=metadata.version("agno"))
|
|
41
|
+
|
|
35
42
|
|
|
36
43
|
class MongoDb(BaseDb):
|
|
37
44
|
def __init__(
|
|
@@ -45,6 +52,8 @@ class MongoDb(BaseDb):
|
|
|
45
52
|
eval_collection: Optional[str] = None,
|
|
46
53
|
knowledge_collection: Optional[str] = None,
|
|
47
54
|
culture_collection: Optional[str] = None,
|
|
55
|
+
traces_collection: Optional[str] = None,
|
|
56
|
+
spans_collection: Optional[str] = None,
|
|
48
57
|
id: Optional[str] = None,
|
|
49
58
|
):
|
|
50
59
|
"""
|
|
@@ -60,6 +69,8 @@ class MongoDb(BaseDb):
|
|
|
60
69
|
eval_collection (Optional[str]): Name of the collection to store evaluation runs.
|
|
61
70
|
knowledge_collection (Optional[str]): Name of the collection to store knowledge documents.
|
|
62
71
|
culture_collection (Optional[str]): Name of the collection to store cultural knowledge.
|
|
72
|
+
traces_collection (Optional[str]): Name of the collection to store traces.
|
|
73
|
+
spans_collection (Optional[str]): Name of the collection to store spans.
|
|
63
74
|
id (Optional[str]): ID of the database.
|
|
64
75
|
|
|
65
76
|
Raises:
|
|
@@ -79,20 +90,36 @@ class MongoDb(BaseDb):
|
|
|
79
90
|
eval_table=eval_collection,
|
|
80
91
|
knowledge_table=knowledge_collection,
|
|
81
92
|
culture_table=culture_collection,
|
|
93
|
+
traces_table=traces_collection,
|
|
94
|
+
spans_table=spans_collection,
|
|
82
95
|
)
|
|
83
96
|
|
|
84
97
|
_client: Optional[MongoClient] = db_client
|
|
85
98
|
if _client is None and db_url is not None:
|
|
86
|
-
_client = MongoClient(db_url)
|
|
99
|
+
_client = MongoClient(db_url, driver=DRIVER_METADATA)
|
|
87
100
|
if _client is None:
|
|
88
101
|
raise ValueError("One of db_url or db_client must be provided")
|
|
89
102
|
|
|
103
|
+
# append_metadata was added in PyMongo 4.14.0, but is a valid database name on earlier versions
|
|
104
|
+
if callable(_client.append_metadata):
|
|
105
|
+
_client.append_metadata(DRIVER_METADATA)
|
|
106
|
+
|
|
90
107
|
self.db_url: Optional[str] = db_url
|
|
91
108
|
self.db_client: MongoClient = _client
|
|
109
|
+
|
|
92
110
|
self.db_name: str = db_name if db_name is not None else "agno"
|
|
93
111
|
|
|
94
112
|
self._database: Optional[Database] = None
|
|
95
113
|
|
|
114
|
+
def close(self) -> None:
|
|
115
|
+
"""Close the MongoDB client connection.
|
|
116
|
+
|
|
117
|
+
Should be called during application shutdown to properly release
|
|
118
|
+
all database connections.
|
|
119
|
+
"""
|
|
120
|
+
if self.db_client is not None:
|
|
121
|
+
self.db_client.close()
|
|
122
|
+
|
|
96
123
|
@property
|
|
97
124
|
def database(self) -> Database:
|
|
98
125
|
if self._database is None:
|
|
@@ -203,6 +230,28 @@ class MongoDb(BaseDb):
|
|
|
203
230
|
)
|
|
204
231
|
return self.culture_collection
|
|
205
232
|
|
|
233
|
+
if table_type == "traces":
|
|
234
|
+
if not hasattr(self, "traces_collection"):
|
|
235
|
+
if self.trace_table_name is None:
|
|
236
|
+
raise ValueError("Traces collection was not provided on initialization")
|
|
237
|
+
self.traces_collection = self._get_or_create_collection(
|
|
238
|
+
collection_name=self.trace_table_name,
|
|
239
|
+
collection_type="traces",
|
|
240
|
+
create_collection_if_not_found=create_collection_if_not_found,
|
|
241
|
+
)
|
|
242
|
+
return self.traces_collection
|
|
243
|
+
|
|
244
|
+
if table_type == "spans":
|
|
245
|
+
if not hasattr(self, "spans_collection"):
|
|
246
|
+
if self.span_table_name is None:
|
|
247
|
+
raise ValueError("Spans collection was not provided on initialization")
|
|
248
|
+
self.spans_collection = self._get_or_create_collection(
|
|
249
|
+
collection_name=self.span_table_name,
|
|
250
|
+
collection_type="spans",
|
|
251
|
+
create_collection_if_not_found=create_collection_if_not_found,
|
|
252
|
+
)
|
|
253
|
+
return self.spans_collection
|
|
254
|
+
|
|
206
255
|
raise ValueError(f"Unknown table type: {table_type}")
|
|
207
256
|
|
|
208
257
|
def _get_or_create_collection(
|
|
@@ -236,6 +285,14 @@ class MongoDb(BaseDb):
|
|
|
236
285
|
log_error(f"Error getting collection {collection_name}: {e}")
|
|
237
286
|
raise
|
|
238
287
|
|
|
288
|
+
def get_latest_schema_version(self):
|
|
289
|
+
"""Get the latest version of the database schema."""
|
|
290
|
+
pass
|
|
291
|
+
|
|
292
|
+
def upsert_schema_version(self, version: str) -> None:
|
|
293
|
+
"""Upsert the schema version into the database."""
|
|
294
|
+
pass
|
|
295
|
+
|
|
239
296
|
# -- Session methods --
|
|
240
297
|
|
|
241
298
|
def delete_session(self, session_id: str) -> bool:
|
|
@@ -976,12 +1033,14 @@ class MongoDb(BaseDb):
|
|
|
976
1033
|
self,
|
|
977
1034
|
limit: Optional[int] = None,
|
|
978
1035
|
page: Optional[int] = None,
|
|
1036
|
+
user_id: Optional[str] = None,
|
|
979
1037
|
) -> Tuple[List[Dict[str, Any]], int]:
|
|
980
1038
|
"""Get user memories stats.
|
|
981
1039
|
|
|
982
1040
|
Args:
|
|
983
1041
|
limit (Optional[int]): The limit of the memories to get.
|
|
984
1042
|
page (Optional[int]): The page number to get.
|
|
1043
|
+
user_id (Optional[str]): User ID for filtering.
|
|
985
1044
|
|
|
986
1045
|
Returns:
|
|
987
1046
|
Tuple[List[Dict[str, Any]], int]: A tuple containing the memories stats and the total count.
|
|
@@ -994,9 +1053,11 @@ class MongoDb(BaseDb):
|
|
|
994
1053
|
if collection is None:
|
|
995
1054
|
return [], 0
|
|
996
1055
|
|
|
997
|
-
match_stage = {"user_id": {"$ne": None}}
|
|
1056
|
+
match_stage: Dict[str, Any] = {"user_id": {"$ne": None}}
|
|
1057
|
+
if user_id is not None:
|
|
1058
|
+
match_stage["user_id"] = user_id
|
|
998
1059
|
|
|
999
|
-
pipeline = [
|
|
1060
|
+
pipeline: List[Dict[str, Any]] = [
|
|
1000
1061
|
{"$match": match_stage},
|
|
1001
1062
|
{
|
|
1002
1063
|
"$group": {
|
|
@@ -1140,7 +1201,10 @@ class MongoDb(BaseDb):
|
|
|
1140
1201
|
"team_id": memory.team_id,
|
|
1141
1202
|
"memory_id": memory.memory_id,
|
|
1142
1203
|
"memory": memory.memory,
|
|
1204
|
+
"input": memory.input,
|
|
1205
|
+
"feedback": memory.feedback,
|
|
1143
1206
|
"topics": memory.topics,
|
|
1207
|
+
"created_at": memory.created_at,
|
|
1144
1208
|
"updated_at": updated_at,
|
|
1145
1209
|
}
|
|
1146
1210
|
|
|
@@ -1980,3 +2044,619 @@ class MongoDb(BaseDb):
|
|
|
1980
2044
|
for memory in memories:
|
|
1981
2045
|
self.upsert_user_memory(memory)
|
|
1982
2046
|
log_info(f"Migrated {len(memories)} memories to collection: {self.memory_table_name}")
|
|
2047
|
+
|
|
2048
|
+
# --- Traces ---
|
|
2049
|
+
def _get_component_level(
|
|
2050
|
+
self, workflow_id: Optional[str], team_id: Optional[str], agent_id: Optional[str], name: str
|
|
2051
|
+
) -> int:
|
|
2052
|
+
"""Get the component level for a trace based on its context.
|
|
2053
|
+
|
|
2054
|
+
Component levels (higher = more important):
|
|
2055
|
+
- 3: Workflow root (.run or .arun with workflow_id)
|
|
2056
|
+
- 2: Team root (.run or .arun with team_id)
|
|
2057
|
+
- 1: Agent root (.run or .arun with agent_id)
|
|
2058
|
+
- 0: Child span (not a root)
|
|
2059
|
+
|
|
2060
|
+
Args:
|
|
2061
|
+
workflow_id: The workflow ID of the trace.
|
|
2062
|
+
team_id: The team ID of the trace.
|
|
2063
|
+
agent_id: The agent ID of the trace.
|
|
2064
|
+
name: The name of the trace.
|
|
2065
|
+
|
|
2066
|
+
Returns:
|
|
2067
|
+
int: The component level (0-3).
|
|
2068
|
+
"""
|
|
2069
|
+
# Check if name indicates a root span
|
|
2070
|
+
is_root_name = ".run" in name or ".arun" in name
|
|
2071
|
+
|
|
2072
|
+
if not is_root_name:
|
|
2073
|
+
return 0 # Child span (not a root)
|
|
2074
|
+
elif workflow_id:
|
|
2075
|
+
return 3 # Workflow root
|
|
2076
|
+
elif team_id:
|
|
2077
|
+
return 2 # Team root
|
|
2078
|
+
elif agent_id:
|
|
2079
|
+
return 1 # Agent root
|
|
2080
|
+
else:
|
|
2081
|
+
return 0 # Unknown
|
|
2082
|
+
|
|
2083
|
+
def upsert_trace(self, trace: "Trace") -> None:
|
|
2084
|
+
"""Create or update a single trace record in the database.
|
|
2085
|
+
|
|
2086
|
+
Uses MongoDB's update_one with upsert=True and aggregation pipeline
|
|
2087
|
+
to handle concurrent inserts atomically and avoid race conditions.
|
|
2088
|
+
|
|
2089
|
+
Args:
|
|
2090
|
+
trace: The Trace object to store (one per trace_id).
|
|
2091
|
+
"""
|
|
2092
|
+
try:
|
|
2093
|
+
collection = self._get_collection(table_type="traces", create_collection_if_not_found=True)
|
|
2094
|
+
if collection is None:
|
|
2095
|
+
return
|
|
2096
|
+
|
|
2097
|
+
trace_dict = trace.to_dict()
|
|
2098
|
+
trace_dict.pop("total_spans", None)
|
|
2099
|
+
trace_dict.pop("error_count", None)
|
|
2100
|
+
|
|
2101
|
+
# Calculate the component level for the new trace
|
|
2102
|
+
new_level = self._get_component_level(trace.workflow_id, trace.team_id, trace.agent_id, trace.name)
|
|
2103
|
+
|
|
2104
|
+
# Use MongoDB aggregation pipeline update for atomic upsert
|
|
2105
|
+
# This allows conditional logic within a single atomic operation
|
|
2106
|
+
pipeline: List[Dict[str, Any]] = [
|
|
2107
|
+
{
|
|
2108
|
+
"$set": {
|
|
2109
|
+
# Always update these fields
|
|
2110
|
+
"status": trace.status,
|
|
2111
|
+
"created_at": {"$ifNull": ["$created_at", trace_dict.get("created_at")]},
|
|
2112
|
+
# Use $min for start_time (keep earliest)
|
|
2113
|
+
"start_time": {
|
|
2114
|
+
"$cond": {
|
|
2115
|
+
"if": {"$eq": [{"$type": "$start_time"}, "missing"]},
|
|
2116
|
+
"then": trace_dict.get("start_time"),
|
|
2117
|
+
"else": {"$min": ["$start_time", trace_dict.get("start_time")]},
|
|
2118
|
+
}
|
|
2119
|
+
},
|
|
2120
|
+
# Use $max for end_time (keep latest)
|
|
2121
|
+
"end_time": {
|
|
2122
|
+
"$cond": {
|
|
2123
|
+
"if": {"$eq": [{"$type": "$end_time"}, "missing"]},
|
|
2124
|
+
"then": trace_dict.get("end_time"),
|
|
2125
|
+
"else": {"$max": ["$end_time", trace_dict.get("end_time")]},
|
|
2126
|
+
}
|
|
2127
|
+
},
|
|
2128
|
+
# Preserve existing non-null context values using $ifNull
|
|
2129
|
+
"run_id": {"$ifNull": [trace.run_id, "$run_id"]},
|
|
2130
|
+
"session_id": {"$ifNull": [trace.session_id, "$session_id"]},
|
|
2131
|
+
"user_id": {"$ifNull": [trace.user_id, "$user_id"]},
|
|
2132
|
+
"agent_id": {"$ifNull": [trace.agent_id, "$agent_id"]},
|
|
2133
|
+
"team_id": {"$ifNull": [trace.team_id, "$team_id"]},
|
|
2134
|
+
"workflow_id": {"$ifNull": [trace.workflow_id, "$workflow_id"]},
|
|
2135
|
+
}
|
|
2136
|
+
},
|
|
2137
|
+
{
|
|
2138
|
+
"$set": {
|
|
2139
|
+
# Calculate duration_ms from the (potentially updated) start_time and end_time
|
|
2140
|
+
# MongoDB stores dates as strings in ISO format, so we need to parse them
|
|
2141
|
+
"duration_ms": {
|
|
2142
|
+
"$cond": {
|
|
2143
|
+
"if": {
|
|
2144
|
+
"$and": [
|
|
2145
|
+
{"$ne": [{"$type": "$start_time"}, "missing"]},
|
|
2146
|
+
{"$ne": [{"$type": "$end_time"}, "missing"]},
|
|
2147
|
+
]
|
|
2148
|
+
},
|
|
2149
|
+
"then": {
|
|
2150
|
+
"$subtract": [
|
|
2151
|
+
{"$toLong": {"$toDate": "$end_time"}},
|
|
2152
|
+
{"$toLong": {"$toDate": "$start_time"}},
|
|
2153
|
+
]
|
|
2154
|
+
},
|
|
2155
|
+
"else": trace_dict.get("duration_ms", 0),
|
|
2156
|
+
}
|
|
2157
|
+
},
|
|
2158
|
+
# Update name based on component level priority
|
|
2159
|
+
# Only update if new trace is from a higher-level component
|
|
2160
|
+
"name": {
|
|
2161
|
+
"$cond": {
|
|
2162
|
+
"if": {"$eq": [{"$type": "$name"}, "missing"]},
|
|
2163
|
+
"then": trace.name,
|
|
2164
|
+
"else": {
|
|
2165
|
+
"$cond": {
|
|
2166
|
+
"if": {
|
|
2167
|
+
"$gt": [
|
|
2168
|
+
new_level,
|
|
2169
|
+
{
|
|
2170
|
+
"$switch": {
|
|
2171
|
+
"branches": [
|
|
2172
|
+
# Check if existing name is a root span
|
|
2173
|
+
{
|
|
2174
|
+
"case": {
|
|
2175
|
+
"$not": {
|
|
2176
|
+
"$or": [
|
|
2177
|
+
{
|
|
2178
|
+
"$regexMatch": {
|
|
2179
|
+
"input": {"$ifNull": ["$name", ""]},
|
|
2180
|
+
"regex": "\\.run",
|
|
2181
|
+
}
|
|
2182
|
+
},
|
|
2183
|
+
{
|
|
2184
|
+
"$regexMatch": {
|
|
2185
|
+
"input": {"$ifNull": ["$name", ""]},
|
|
2186
|
+
"regex": "\\.arun",
|
|
2187
|
+
}
|
|
2188
|
+
},
|
|
2189
|
+
]
|
|
2190
|
+
}
|
|
2191
|
+
},
|
|
2192
|
+
"then": 0,
|
|
2193
|
+
},
|
|
2194
|
+
# Workflow root (level 3)
|
|
2195
|
+
{
|
|
2196
|
+
"case": {"$ne": ["$workflow_id", None]},
|
|
2197
|
+
"then": 3,
|
|
2198
|
+
},
|
|
2199
|
+
# Team root (level 2)
|
|
2200
|
+
{
|
|
2201
|
+
"case": {"$ne": ["$team_id", None]},
|
|
2202
|
+
"then": 2,
|
|
2203
|
+
},
|
|
2204
|
+
# Agent root (level 1)
|
|
2205
|
+
{
|
|
2206
|
+
"case": {"$ne": ["$agent_id", None]},
|
|
2207
|
+
"then": 1,
|
|
2208
|
+
},
|
|
2209
|
+
],
|
|
2210
|
+
"default": 0,
|
|
2211
|
+
}
|
|
2212
|
+
},
|
|
2213
|
+
]
|
|
2214
|
+
},
|
|
2215
|
+
"then": trace.name,
|
|
2216
|
+
"else": "$name",
|
|
2217
|
+
}
|
|
2218
|
+
},
|
|
2219
|
+
}
|
|
2220
|
+
},
|
|
2221
|
+
}
|
|
2222
|
+
},
|
|
2223
|
+
]
|
|
2224
|
+
|
|
2225
|
+
# Perform atomic upsert using aggregation pipeline
|
|
2226
|
+
collection.update_one(
|
|
2227
|
+
{"trace_id": trace.trace_id},
|
|
2228
|
+
pipeline,
|
|
2229
|
+
upsert=True,
|
|
2230
|
+
)
|
|
2231
|
+
|
|
2232
|
+
except Exception as e:
|
|
2233
|
+
log_error(f"Error creating trace: {e}")
|
|
2234
|
+
# Don't raise - tracing should not break the main application flow
|
|
2235
|
+
|
|
2236
|
+
def get_trace(
|
|
2237
|
+
self,
|
|
2238
|
+
trace_id: Optional[str] = None,
|
|
2239
|
+
run_id: Optional[str] = None,
|
|
2240
|
+
):
|
|
2241
|
+
"""Get a single trace by trace_id or other filters.
|
|
2242
|
+
|
|
2243
|
+
Args:
|
|
2244
|
+
trace_id: The unique trace identifier.
|
|
2245
|
+
run_id: Filter by run ID (returns first match).
|
|
2246
|
+
|
|
2247
|
+
Returns:
|
|
2248
|
+
Optional[Trace]: The trace if found, None otherwise.
|
|
2249
|
+
|
|
2250
|
+
Note:
|
|
2251
|
+
If multiple filters are provided, trace_id takes precedence.
|
|
2252
|
+
For other filters, the most recent trace is returned.
|
|
2253
|
+
"""
|
|
2254
|
+
try:
|
|
2255
|
+
from agno.tracing.schemas import Trace as TraceSchema
|
|
2256
|
+
|
|
2257
|
+
collection = self._get_collection(table_type="traces")
|
|
2258
|
+
if collection is None:
|
|
2259
|
+
return None
|
|
2260
|
+
|
|
2261
|
+
# Get spans collection for aggregation
|
|
2262
|
+
spans_collection = self._get_collection(table_type="spans")
|
|
2263
|
+
|
|
2264
|
+
query: Dict[str, Any] = {}
|
|
2265
|
+
if trace_id:
|
|
2266
|
+
query["trace_id"] = trace_id
|
|
2267
|
+
elif run_id:
|
|
2268
|
+
query["run_id"] = run_id
|
|
2269
|
+
else:
|
|
2270
|
+
log_debug("get_trace called without any filter parameters")
|
|
2271
|
+
return None
|
|
2272
|
+
|
|
2273
|
+
# Find trace with sorting by most recent
|
|
2274
|
+
result = collection.find_one(query, sort=[("start_time", -1)])
|
|
2275
|
+
|
|
2276
|
+
if result:
|
|
2277
|
+
# Calculate total_spans and error_count from spans collection
|
|
2278
|
+
total_spans = 0
|
|
2279
|
+
error_count = 0
|
|
2280
|
+
if spans_collection is not None:
|
|
2281
|
+
total_spans = spans_collection.count_documents({"trace_id": result["trace_id"]})
|
|
2282
|
+
error_count = spans_collection.count_documents(
|
|
2283
|
+
{"trace_id": result["trace_id"], "status_code": "ERROR"}
|
|
2284
|
+
)
|
|
2285
|
+
|
|
2286
|
+
result["total_spans"] = total_spans
|
|
2287
|
+
result["error_count"] = error_count
|
|
2288
|
+
# Remove MongoDB's _id field
|
|
2289
|
+
result.pop("_id", None)
|
|
2290
|
+
return TraceSchema.from_dict(result)
|
|
2291
|
+
return None
|
|
2292
|
+
|
|
2293
|
+
except Exception as e:
|
|
2294
|
+
log_error(f"Error getting trace: {e}")
|
|
2295
|
+
return None
|
|
2296
|
+
|
|
2297
|
+
def get_traces(
|
|
2298
|
+
self,
|
|
2299
|
+
run_id: Optional[str] = None,
|
|
2300
|
+
session_id: Optional[str] = None,
|
|
2301
|
+
user_id: Optional[str] = None,
|
|
2302
|
+
agent_id: Optional[str] = None,
|
|
2303
|
+
team_id: Optional[str] = None,
|
|
2304
|
+
workflow_id: Optional[str] = None,
|
|
2305
|
+
status: Optional[str] = None,
|
|
2306
|
+
start_time: Optional[datetime] = None,
|
|
2307
|
+
end_time: Optional[datetime] = None,
|
|
2308
|
+
limit: Optional[int] = 20,
|
|
2309
|
+
page: Optional[int] = 1,
|
|
2310
|
+
) -> tuple[List, int]:
|
|
2311
|
+
"""Get traces matching the provided filters with pagination.
|
|
2312
|
+
|
|
2313
|
+
Args:
|
|
2314
|
+
run_id: Filter by run ID.
|
|
2315
|
+
session_id: Filter by session ID.
|
|
2316
|
+
user_id: Filter by user ID.
|
|
2317
|
+
agent_id: Filter by agent ID.
|
|
2318
|
+
team_id: Filter by team ID.
|
|
2319
|
+
workflow_id: Filter by workflow ID.
|
|
2320
|
+
status: Filter by status (OK, ERROR, UNSET).
|
|
2321
|
+
start_time: Filter traces starting after this datetime.
|
|
2322
|
+
end_time: Filter traces ending before this datetime.
|
|
2323
|
+
limit: Maximum number of traces to return per page.
|
|
2324
|
+
page: Page number (1-indexed).
|
|
2325
|
+
|
|
2326
|
+
Returns:
|
|
2327
|
+
tuple[List[Trace], int]: Tuple of (list of matching traces, total count).
|
|
2328
|
+
"""
|
|
2329
|
+
try:
|
|
2330
|
+
from agno.tracing.schemas import Trace as TraceSchema
|
|
2331
|
+
|
|
2332
|
+
collection = self._get_collection(table_type="traces")
|
|
2333
|
+
if collection is None:
|
|
2334
|
+
log_debug("Traces collection not found")
|
|
2335
|
+
return [], 0
|
|
2336
|
+
|
|
2337
|
+
# Get spans collection for aggregation
|
|
2338
|
+
spans_collection = self._get_collection(table_type="spans")
|
|
2339
|
+
|
|
2340
|
+
# Build query
|
|
2341
|
+
query: Dict[str, Any] = {}
|
|
2342
|
+
if run_id:
|
|
2343
|
+
query["run_id"] = run_id
|
|
2344
|
+
if session_id:
|
|
2345
|
+
query["session_id"] = session_id
|
|
2346
|
+
if user_id:
|
|
2347
|
+
query["user_id"] = user_id
|
|
2348
|
+
if agent_id:
|
|
2349
|
+
query["agent_id"] = agent_id
|
|
2350
|
+
if team_id:
|
|
2351
|
+
query["team_id"] = team_id
|
|
2352
|
+
if workflow_id:
|
|
2353
|
+
query["workflow_id"] = workflow_id
|
|
2354
|
+
if status:
|
|
2355
|
+
query["status"] = status
|
|
2356
|
+
if start_time:
|
|
2357
|
+
query["start_time"] = {"$gte": start_time.isoformat()}
|
|
2358
|
+
if end_time:
|
|
2359
|
+
if "end_time" in query:
|
|
2360
|
+
query["end_time"]["$lte"] = end_time.isoformat()
|
|
2361
|
+
else:
|
|
2362
|
+
query["end_time"] = {"$lte": end_time.isoformat()}
|
|
2363
|
+
|
|
2364
|
+
# Get total count
|
|
2365
|
+
total_count = collection.count_documents(query)
|
|
2366
|
+
|
|
2367
|
+
# Apply pagination
|
|
2368
|
+
skip = ((page or 1) - 1) * (limit or 20)
|
|
2369
|
+
cursor = collection.find(query).sort("start_time", -1).skip(skip).limit(limit or 20)
|
|
2370
|
+
|
|
2371
|
+
results = list(cursor)
|
|
2372
|
+
|
|
2373
|
+
traces = []
|
|
2374
|
+
for row in results:
|
|
2375
|
+
# Calculate total_spans and error_count from spans collection
|
|
2376
|
+
total_spans = 0
|
|
2377
|
+
error_count = 0
|
|
2378
|
+
if spans_collection is not None:
|
|
2379
|
+
total_spans = spans_collection.count_documents({"trace_id": row["trace_id"]})
|
|
2380
|
+
error_count = spans_collection.count_documents(
|
|
2381
|
+
{"trace_id": row["trace_id"], "status_code": "ERROR"}
|
|
2382
|
+
)
|
|
2383
|
+
|
|
2384
|
+
row["total_spans"] = total_spans
|
|
2385
|
+
row["error_count"] = error_count
|
|
2386
|
+
# Remove MongoDB's _id field
|
|
2387
|
+
row.pop("_id", None)
|
|
2388
|
+
traces.append(TraceSchema.from_dict(row))
|
|
2389
|
+
|
|
2390
|
+
return traces, total_count
|
|
2391
|
+
|
|
2392
|
+
except Exception as e:
|
|
2393
|
+
log_error(f"Error getting traces: {e}")
|
|
2394
|
+
return [], 0
|
|
2395
|
+
|
|
2396
|
+
def get_trace_stats(
|
|
2397
|
+
self,
|
|
2398
|
+
user_id: Optional[str] = None,
|
|
2399
|
+
agent_id: Optional[str] = None,
|
|
2400
|
+
team_id: Optional[str] = None,
|
|
2401
|
+
workflow_id: Optional[str] = None,
|
|
2402
|
+
start_time: Optional[datetime] = None,
|
|
2403
|
+
end_time: Optional[datetime] = None,
|
|
2404
|
+
limit: Optional[int] = 20,
|
|
2405
|
+
page: Optional[int] = 1,
|
|
2406
|
+
) -> tuple[List[Dict[str, Any]], int]:
|
|
2407
|
+
"""Get trace statistics grouped by session.
|
|
2408
|
+
|
|
2409
|
+
Args:
|
|
2410
|
+
user_id: Filter by user ID.
|
|
2411
|
+
agent_id: Filter by agent ID.
|
|
2412
|
+
team_id: Filter by team ID.
|
|
2413
|
+
workflow_id: Filter by workflow ID.
|
|
2414
|
+
start_time: Filter sessions with traces created after this datetime.
|
|
2415
|
+
end_time: Filter sessions with traces created before this datetime.
|
|
2416
|
+
limit: Maximum number of sessions to return per page.
|
|
2417
|
+
page: Page number (1-indexed).
|
|
2418
|
+
|
|
2419
|
+
Returns:
|
|
2420
|
+
tuple[List[Dict], int]: Tuple of (list of session stats dicts, total count).
|
|
2421
|
+
Each dict contains: session_id, user_id, agent_id, team_id, total_traces,
|
|
2422
|
+
workflow_id, first_trace_at, last_trace_at.
|
|
2423
|
+
"""
|
|
2424
|
+
try:
|
|
2425
|
+
collection = self._get_collection(table_type="traces")
|
|
2426
|
+
if collection is None:
|
|
2427
|
+
log_debug("Traces collection not found")
|
|
2428
|
+
return [], 0
|
|
2429
|
+
|
|
2430
|
+
# Build match stage
|
|
2431
|
+
match_stage: Dict[str, Any] = {"session_id": {"$ne": None}}
|
|
2432
|
+
if user_id:
|
|
2433
|
+
match_stage["user_id"] = user_id
|
|
2434
|
+
if agent_id:
|
|
2435
|
+
match_stage["agent_id"] = agent_id
|
|
2436
|
+
if team_id:
|
|
2437
|
+
match_stage["team_id"] = team_id
|
|
2438
|
+
if workflow_id:
|
|
2439
|
+
match_stage["workflow_id"] = workflow_id
|
|
2440
|
+
if start_time:
|
|
2441
|
+
match_stage["created_at"] = {"$gte": start_time.isoformat()}
|
|
2442
|
+
if end_time:
|
|
2443
|
+
if "created_at" in match_stage:
|
|
2444
|
+
match_stage["created_at"]["$lte"] = end_time.isoformat()
|
|
2445
|
+
else:
|
|
2446
|
+
match_stage["created_at"] = {"$lte": end_time.isoformat()}
|
|
2447
|
+
|
|
2448
|
+
# Build aggregation pipeline
|
|
2449
|
+
pipeline: List[Dict[str, Any]] = [
|
|
2450
|
+
{"$match": match_stage},
|
|
2451
|
+
{
|
|
2452
|
+
"$group": {
|
|
2453
|
+
"_id": "$session_id",
|
|
2454
|
+
"user_id": {"$first": "$user_id"},
|
|
2455
|
+
"agent_id": {"$first": "$agent_id"},
|
|
2456
|
+
"team_id": {"$first": "$team_id"},
|
|
2457
|
+
"workflow_id": {"$first": "$workflow_id"},
|
|
2458
|
+
"total_traces": {"$sum": 1},
|
|
2459
|
+
"first_trace_at": {"$min": "$created_at"},
|
|
2460
|
+
"last_trace_at": {"$max": "$created_at"},
|
|
2461
|
+
}
|
|
2462
|
+
},
|
|
2463
|
+
{"$sort": {"last_trace_at": -1}},
|
|
2464
|
+
]
|
|
2465
|
+
|
|
2466
|
+
# Get total count
|
|
2467
|
+
count_pipeline = pipeline + [{"$count": "total"}]
|
|
2468
|
+
count_result = list(collection.aggregate(count_pipeline))
|
|
2469
|
+
total_count = count_result[0]["total"] if count_result else 0
|
|
2470
|
+
|
|
2471
|
+
# Apply pagination
|
|
2472
|
+
skip = ((page or 1) - 1) * (limit or 20)
|
|
2473
|
+
pipeline.append({"$skip": skip})
|
|
2474
|
+
pipeline.append({"$limit": limit or 20})
|
|
2475
|
+
|
|
2476
|
+
results = list(collection.aggregate(pipeline))
|
|
2477
|
+
|
|
2478
|
+
# Convert to list of dicts with datetime objects
|
|
2479
|
+
stats_list = []
|
|
2480
|
+
for row in results:
|
|
2481
|
+
# Convert ISO strings to datetime objects
|
|
2482
|
+
first_trace_at_str = row["first_trace_at"]
|
|
2483
|
+
last_trace_at_str = row["last_trace_at"]
|
|
2484
|
+
|
|
2485
|
+
# Parse ISO format strings to datetime objects
|
|
2486
|
+
first_trace_at = datetime.fromisoformat(first_trace_at_str.replace("Z", "+00:00"))
|
|
2487
|
+
last_trace_at = datetime.fromisoformat(last_trace_at_str.replace("Z", "+00:00"))
|
|
2488
|
+
|
|
2489
|
+
stats_list.append(
|
|
2490
|
+
{
|
|
2491
|
+
"session_id": row["_id"],
|
|
2492
|
+
"user_id": row["user_id"],
|
|
2493
|
+
"agent_id": row["agent_id"],
|
|
2494
|
+
"team_id": row["team_id"],
|
|
2495
|
+
"workflow_id": row["workflow_id"],
|
|
2496
|
+
"total_traces": row["total_traces"],
|
|
2497
|
+
"first_trace_at": first_trace_at,
|
|
2498
|
+
"last_trace_at": last_trace_at,
|
|
2499
|
+
}
|
|
2500
|
+
)
|
|
2501
|
+
|
|
2502
|
+
return stats_list, total_count
|
|
2503
|
+
|
|
2504
|
+
except Exception as e:
|
|
2505
|
+
log_error(f"Error getting trace stats: {e}")
|
|
2506
|
+
return [], 0
|
|
2507
|
+
|
|
2508
|
+
# --- Spans ---
|
|
2509
|
+
def create_span(self, span: "Span") -> None:
|
|
2510
|
+
"""Create a single span in the database.
|
|
2511
|
+
|
|
2512
|
+
Args:
|
|
2513
|
+
span: The Span object to store.
|
|
2514
|
+
"""
|
|
2515
|
+
try:
|
|
2516
|
+
collection = self._get_collection(table_type="spans", create_collection_if_not_found=True)
|
|
2517
|
+
if collection is None:
|
|
2518
|
+
return
|
|
2519
|
+
|
|
2520
|
+
collection.insert_one(span.to_dict())
|
|
2521
|
+
|
|
2522
|
+
except Exception as e:
|
|
2523
|
+
log_error(f"Error creating span: {e}")
|
|
2524
|
+
|
|
2525
|
+
def create_spans(self, spans: List) -> None:
|
|
2526
|
+
"""Create multiple spans in the database as a batch.
|
|
2527
|
+
|
|
2528
|
+
Args:
|
|
2529
|
+
spans: List of Span objects to store.
|
|
2530
|
+
"""
|
|
2531
|
+
if not spans:
|
|
2532
|
+
return
|
|
2533
|
+
|
|
2534
|
+
try:
|
|
2535
|
+
collection = self._get_collection(table_type="spans", create_collection_if_not_found=True)
|
|
2536
|
+
if collection is None:
|
|
2537
|
+
return
|
|
2538
|
+
|
|
2539
|
+
span_dicts = [span.to_dict() for span in spans]
|
|
2540
|
+
collection.insert_many(span_dicts)
|
|
2541
|
+
|
|
2542
|
+
except Exception as e:
|
|
2543
|
+
log_error(f"Error creating spans batch: {e}")
|
|
2544
|
+
|
|
2545
|
+
def get_span(self, span_id: str):
|
|
2546
|
+
"""Get a single span by its span_id.
|
|
2547
|
+
|
|
2548
|
+
Args:
|
|
2549
|
+
span_id: The unique span identifier.
|
|
2550
|
+
|
|
2551
|
+
Returns:
|
|
2552
|
+
Optional[Span]: The span if found, None otherwise.
|
|
2553
|
+
"""
|
|
2554
|
+
try:
|
|
2555
|
+
from agno.tracing.schemas import Span as SpanSchema
|
|
2556
|
+
|
|
2557
|
+
collection = self._get_collection(table_type="spans")
|
|
2558
|
+
if collection is None:
|
|
2559
|
+
return None
|
|
2560
|
+
|
|
2561
|
+
result = collection.find_one({"span_id": span_id})
|
|
2562
|
+
if result:
|
|
2563
|
+
# Remove MongoDB's _id field
|
|
2564
|
+
result.pop("_id", None)
|
|
2565
|
+
return SpanSchema.from_dict(result)
|
|
2566
|
+
return None
|
|
2567
|
+
|
|
2568
|
+
except Exception as e:
|
|
2569
|
+
log_error(f"Error getting span: {e}")
|
|
2570
|
+
return None
|
|
2571
|
+
|
|
2572
|
+
def get_spans(
|
|
2573
|
+
self,
|
|
2574
|
+
trace_id: Optional[str] = None,
|
|
2575
|
+
parent_span_id: Optional[str] = None,
|
|
2576
|
+
limit: Optional[int] = 1000,
|
|
2577
|
+
) -> List:
|
|
2578
|
+
"""Get spans matching the provided filters.
|
|
2579
|
+
|
|
2580
|
+
Args:
|
|
2581
|
+
trace_id: Filter by trace ID.
|
|
2582
|
+
parent_span_id: Filter by parent span ID.
|
|
2583
|
+
limit: Maximum number of spans to return.
|
|
2584
|
+
|
|
2585
|
+
Returns:
|
|
2586
|
+
List[Span]: List of matching spans.
|
|
2587
|
+
"""
|
|
2588
|
+
try:
|
|
2589
|
+
from agno.tracing.schemas import Span as SpanSchema
|
|
2590
|
+
|
|
2591
|
+
collection = self._get_collection(table_type="spans")
|
|
2592
|
+
if collection is None:
|
|
2593
|
+
return []
|
|
2594
|
+
|
|
2595
|
+
# Build query
|
|
2596
|
+
query: Dict[str, Any] = {}
|
|
2597
|
+
if trace_id:
|
|
2598
|
+
query["trace_id"] = trace_id
|
|
2599
|
+
if parent_span_id:
|
|
2600
|
+
query["parent_span_id"] = parent_span_id
|
|
2601
|
+
|
|
2602
|
+
cursor = collection.find(query).limit(limit or 1000)
|
|
2603
|
+
results = list(cursor)
|
|
2604
|
+
|
|
2605
|
+
spans = []
|
|
2606
|
+
for row in results:
|
|
2607
|
+
# Remove MongoDB's _id field
|
|
2608
|
+
row.pop("_id", None)
|
|
2609
|
+
spans.append(SpanSchema.from_dict(row))
|
|
2610
|
+
|
|
2611
|
+
return spans
|
|
2612
|
+
|
|
2613
|
+
except Exception as e:
|
|
2614
|
+
log_error(f"Error getting spans: {e}")
|
|
2615
|
+
return []
|
|
2616
|
+
|
|
2617
|
+
# -- Learning methods (stubs) --
|
|
2618
|
+
def get_learning(
|
|
2619
|
+
self,
|
|
2620
|
+
learning_type: str,
|
|
2621
|
+
user_id: Optional[str] = None,
|
|
2622
|
+
agent_id: Optional[str] = None,
|
|
2623
|
+
team_id: Optional[str] = None,
|
|
2624
|
+
session_id: Optional[str] = None,
|
|
2625
|
+
namespace: Optional[str] = None,
|
|
2626
|
+
entity_id: Optional[str] = None,
|
|
2627
|
+
entity_type: Optional[str] = None,
|
|
2628
|
+
) -> Optional[Dict[str, Any]]:
|
|
2629
|
+
raise NotImplementedError("Learning methods not yet implemented for MongoDb")
|
|
2630
|
+
|
|
2631
|
+
def upsert_learning(
|
|
2632
|
+
self,
|
|
2633
|
+
id: str,
|
|
2634
|
+
learning_type: str,
|
|
2635
|
+
content: Dict[str, Any],
|
|
2636
|
+
user_id: Optional[str] = None,
|
|
2637
|
+
agent_id: Optional[str] = None,
|
|
2638
|
+
team_id: Optional[str] = None,
|
|
2639
|
+
session_id: Optional[str] = None,
|
|
2640
|
+
namespace: Optional[str] = None,
|
|
2641
|
+
entity_id: Optional[str] = None,
|
|
2642
|
+
entity_type: Optional[str] = None,
|
|
2643
|
+
metadata: Optional[Dict[str, Any]] = None,
|
|
2644
|
+
) -> None:
|
|
2645
|
+
raise NotImplementedError("Learning methods not yet implemented for MongoDb")
|
|
2646
|
+
|
|
2647
|
+
def delete_learning(self, id: str) -> bool:
|
|
2648
|
+
raise NotImplementedError("Learning methods not yet implemented for MongoDb")
|
|
2649
|
+
|
|
2650
|
+
def get_learnings(
|
|
2651
|
+
self,
|
|
2652
|
+
learning_type: Optional[str] = None,
|
|
2653
|
+
user_id: Optional[str] = None,
|
|
2654
|
+
agent_id: Optional[str] = None,
|
|
2655
|
+
team_id: Optional[str] = None,
|
|
2656
|
+
session_id: Optional[str] = None,
|
|
2657
|
+
namespace: Optional[str] = None,
|
|
2658
|
+
entity_id: Optional[str] = None,
|
|
2659
|
+
entity_type: Optional[str] = None,
|
|
2660
|
+
limit: Optional[int] = None,
|
|
2661
|
+
) -> List[Dict[str, Any]]:
|
|
2662
|
+
raise NotImplementedError("Learning methods not yet implemented for MongoDb")
|