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/firestore/firestore.py
CHANGED
|
@@ -1,8 +1,12 @@
|
|
|
1
|
+
import json
|
|
1
2
|
import time
|
|
2
3
|
from datetime import date, datetime, timedelta, timezone
|
|
3
|
-
from typing import Any, Dict, List, Optional, Tuple, Union
|
|
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.firestore.utils import (
|
|
8
12
|
apply_pagination,
|
|
@@ -45,6 +49,8 @@ class FirestoreDb(BaseDb):
|
|
|
45
49
|
eval_collection: Optional[str] = None,
|
|
46
50
|
knowledge_collection: Optional[str] = None,
|
|
47
51
|
culture_collection: Optional[str] = None,
|
|
52
|
+
traces_collection: Optional[str] = None,
|
|
53
|
+
spans_collection: Optional[str] = None,
|
|
48
54
|
id: Optional[str] = None,
|
|
49
55
|
):
|
|
50
56
|
"""
|
|
@@ -59,6 +65,8 @@ class FirestoreDb(BaseDb):
|
|
|
59
65
|
eval_collection (Optional[str]): Name of the collection to store evaluation runs.
|
|
60
66
|
knowledge_collection (Optional[str]): Name of the collection to store knowledge documents.
|
|
61
67
|
culture_collection (Optional[str]): Name of the collection to store cultural knowledge.
|
|
68
|
+
traces_collection (Optional[str]): Name of the collection to store traces.
|
|
69
|
+
spans_collection (Optional[str]): Name of the collection to store spans.
|
|
62
70
|
id (Optional[str]): ID of the database.
|
|
63
71
|
|
|
64
72
|
Raises:
|
|
@@ -76,6 +84,8 @@ class FirestoreDb(BaseDb):
|
|
|
76
84
|
eval_table=eval_collection,
|
|
77
85
|
knowledge_table=knowledge_collection,
|
|
78
86
|
culture_table=culture_collection,
|
|
87
|
+
traces_table=traces_collection,
|
|
88
|
+
spans_table=spans_collection,
|
|
79
89
|
)
|
|
80
90
|
|
|
81
91
|
_client: Optional[Client] = db_client
|
|
@@ -111,71 +121,88 @@ class FirestoreDb(BaseDb):
|
|
|
111
121
|
CollectionReference: The collection reference.
|
|
112
122
|
"""
|
|
113
123
|
if table_type == "sessions":
|
|
114
|
-
if
|
|
115
|
-
|
|
116
|
-
|
|
117
|
-
|
|
118
|
-
|
|
119
|
-
|
|
120
|
-
|
|
121
|
-
)
|
|
124
|
+
if self.session_table_name is None:
|
|
125
|
+
raise ValueError("Session collection was not provided on initialization")
|
|
126
|
+
self.session_collection = self._get_or_create_collection(
|
|
127
|
+
collection_name=self.session_table_name,
|
|
128
|
+
collection_type="sessions",
|
|
129
|
+
create_collection_if_not_found=create_collection_if_not_found,
|
|
130
|
+
)
|
|
122
131
|
return self.session_collection
|
|
123
132
|
|
|
124
133
|
if table_type == "memories":
|
|
125
|
-
if
|
|
126
|
-
|
|
127
|
-
|
|
128
|
-
|
|
129
|
-
|
|
130
|
-
|
|
131
|
-
|
|
132
|
-
)
|
|
134
|
+
if self.memory_table_name is None:
|
|
135
|
+
raise ValueError("Memory collection was not provided on initialization")
|
|
136
|
+
self.memory_collection = self._get_or_create_collection(
|
|
137
|
+
collection_name=self.memory_table_name,
|
|
138
|
+
collection_type="memories",
|
|
139
|
+
create_collection_if_not_found=create_collection_if_not_found,
|
|
140
|
+
)
|
|
133
141
|
return self.memory_collection
|
|
134
142
|
|
|
135
143
|
if table_type == "metrics":
|
|
136
|
-
if
|
|
137
|
-
|
|
138
|
-
|
|
139
|
-
|
|
140
|
-
|
|
141
|
-
|
|
142
|
-
|
|
143
|
-
)
|
|
144
|
+
if self.metrics_table_name is None:
|
|
145
|
+
raise ValueError("Metrics collection was not provided on initialization")
|
|
146
|
+
self.metrics_collection = self._get_or_create_collection(
|
|
147
|
+
collection_name=self.metrics_table_name,
|
|
148
|
+
collection_type="metrics",
|
|
149
|
+
create_collection_if_not_found=create_collection_if_not_found,
|
|
150
|
+
)
|
|
144
151
|
return self.metrics_collection
|
|
145
152
|
|
|
146
153
|
if table_type == "evals":
|
|
147
|
-
if
|
|
148
|
-
|
|
149
|
-
|
|
150
|
-
|
|
151
|
-
|
|
152
|
-
|
|
153
|
-
|
|
154
|
-
)
|
|
154
|
+
if self.eval_table_name is None:
|
|
155
|
+
raise ValueError("Eval collection was not provided on initialization")
|
|
156
|
+
self.eval_collection = self._get_or_create_collection(
|
|
157
|
+
collection_name=self.eval_table_name,
|
|
158
|
+
collection_type="evals",
|
|
159
|
+
create_collection_if_not_found=create_collection_if_not_found,
|
|
160
|
+
)
|
|
155
161
|
return self.eval_collection
|
|
156
162
|
|
|
157
163
|
if table_type == "knowledge":
|
|
158
|
-
if
|
|
159
|
-
|
|
160
|
-
|
|
161
|
-
|
|
162
|
-
|
|
163
|
-
|
|
164
|
-
|
|
165
|
-
)
|
|
164
|
+
if self.knowledge_table_name is None:
|
|
165
|
+
raise ValueError("Knowledge collection was not provided on initialization")
|
|
166
|
+
self.knowledge_collection = self._get_or_create_collection(
|
|
167
|
+
collection_name=self.knowledge_table_name,
|
|
168
|
+
collection_type="knowledge",
|
|
169
|
+
create_collection_if_not_found=create_collection_if_not_found,
|
|
170
|
+
)
|
|
166
171
|
return self.knowledge_collection
|
|
167
172
|
|
|
168
173
|
if table_type == "culture":
|
|
169
|
-
if
|
|
170
|
-
|
|
171
|
-
|
|
172
|
-
|
|
173
|
-
|
|
174
|
-
|
|
175
|
-
|
|
176
|
-
)
|
|
174
|
+
if self.culture_table_name is None:
|
|
175
|
+
raise ValueError("Culture collection was not provided on initialization")
|
|
176
|
+
self.culture_collection = self._get_or_create_collection(
|
|
177
|
+
collection_name=self.culture_table_name,
|
|
178
|
+
collection_type="culture",
|
|
179
|
+
create_collection_if_not_found=create_collection_if_not_found,
|
|
180
|
+
)
|
|
177
181
|
return self.culture_collection
|
|
178
182
|
|
|
183
|
+
if table_type == "traces":
|
|
184
|
+
if self.trace_table_name is None:
|
|
185
|
+
raise ValueError("Traces collection was not provided on initialization")
|
|
186
|
+
self.traces_collection = self._get_or_create_collection(
|
|
187
|
+
collection_name=self.trace_table_name,
|
|
188
|
+
collection_type="traces",
|
|
189
|
+
create_collection_if_not_found=create_collection_if_not_found,
|
|
190
|
+
)
|
|
191
|
+
return self.traces_collection
|
|
192
|
+
|
|
193
|
+
if table_type == "spans":
|
|
194
|
+
# Ensure traces collection exists first (spans reference traces)
|
|
195
|
+
if create_collection_if_not_found:
|
|
196
|
+
self._get_collection("traces", create_collection_if_not_found=True)
|
|
197
|
+
if self.span_table_name is None:
|
|
198
|
+
raise ValueError("Spans collection was not provided on initialization")
|
|
199
|
+
self.spans_collection = self._get_or_create_collection(
|
|
200
|
+
collection_name=self.span_table_name,
|
|
201
|
+
collection_type="spans",
|
|
202
|
+
create_collection_if_not_found=create_collection_if_not_found,
|
|
203
|
+
)
|
|
204
|
+
return self.spans_collection
|
|
205
|
+
|
|
179
206
|
raise ValueError(f"Unknown table type: {table_type}")
|
|
180
207
|
|
|
181
208
|
def _get_or_create_collection(
|
|
@@ -237,6 +264,14 @@ class FirestoreDb(BaseDb):
|
|
|
237
264
|
log_error(f"Error deleting session: {e}")
|
|
238
265
|
raise e
|
|
239
266
|
|
|
267
|
+
def get_latest_schema_version(self):
|
|
268
|
+
"""Get the latest version of the database schema."""
|
|
269
|
+
pass
|
|
270
|
+
|
|
271
|
+
def upsert_schema_version(self, version: str) -> None:
|
|
272
|
+
"""Upsert the schema version into the database."""
|
|
273
|
+
pass
|
|
274
|
+
|
|
240
275
|
def delete_sessions(self, session_ids: List[str]) -> None:
|
|
241
276
|
"""Delete multiple sessions from the database.
|
|
242
277
|
|
|
@@ -865,6 +900,7 @@ class FirestoreDb(BaseDb):
|
|
|
865
900
|
self,
|
|
866
901
|
limit: Optional[int] = None,
|
|
867
902
|
page: Optional[int] = None,
|
|
903
|
+
user_id: Optional[str] = None,
|
|
868
904
|
) -> Tuple[List[Dict[str, Any]], int]:
|
|
869
905
|
"""Get user memories stats.
|
|
870
906
|
|
|
@@ -881,7 +917,10 @@ class FirestoreDb(BaseDb):
|
|
|
881
917
|
try:
|
|
882
918
|
collection_ref = self._get_collection(table_type="memories")
|
|
883
919
|
|
|
884
|
-
|
|
920
|
+
if user_id:
|
|
921
|
+
query = collection_ref.where(filter=FieldFilter("user_id", "==", user_id))
|
|
922
|
+
else:
|
|
923
|
+
query = collection_ref.where(filter=FieldFilter("user_id", "!=", None))
|
|
885
924
|
|
|
886
925
|
docs = query.stream()
|
|
887
926
|
|
|
@@ -1793,3 +1832,588 @@ class FirestoreDb(BaseDb):
|
|
|
1793
1832
|
except Exception as e:
|
|
1794
1833
|
log_error(f"Error updating eval run name {eval_run_id}: {e}")
|
|
1795
1834
|
raise e
|
|
1835
|
+
|
|
1836
|
+
# --- Traces ---
|
|
1837
|
+
def upsert_trace(self, trace: "Trace") -> None:
|
|
1838
|
+
"""Create or update a single trace record in the database.
|
|
1839
|
+
|
|
1840
|
+
Args:
|
|
1841
|
+
trace: The Trace object to store (one per trace_id).
|
|
1842
|
+
"""
|
|
1843
|
+
try:
|
|
1844
|
+
collection_ref = self._get_collection(table_type="traces", create_collection_if_not_found=True)
|
|
1845
|
+
if collection_ref is None:
|
|
1846
|
+
return
|
|
1847
|
+
|
|
1848
|
+
# Check if trace already exists
|
|
1849
|
+
docs = collection_ref.where(filter=FieldFilter("trace_id", "==", trace.trace_id)).limit(1).stream()
|
|
1850
|
+
existing_doc = None
|
|
1851
|
+
existing_data = None
|
|
1852
|
+
for doc in docs:
|
|
1853
|
+
existing_doc = doc
|
|
1854
|
+
existing_data = doc.to_dict()
|
|
1855
|
+
break
|
|
1856
|
+
|
|
1857
|
+
if existing_data and existing_doc is not None:
|
|
1858
|
+
# Update existing trace
|
|
1859
|
+
def get_component_level(workflow_id, team_id, agent_id, name):
|
|
1860
|
+
is_root_name = ".run" in name or ".arun" in name
|
|
1861
|
+
if not is_root_name:
|
|
1862
|
+
return 0
|
|
1863
|
+
elif workflow_id:
|
|
1864
|
+
return 3
|
|
1865
|
+
elif team_id:
|
|
1866
|
+
return 2
|
|
1867
|
+
elif agent_id:
|
|
1868
|
+
return 1
|
|
1869
|
+
else:
|
|
1870
|
+
return 0
|
|
1871
|
+
|
|
1872
|
+
existing_level = get_component_level(
|
|
1873
|
+
existing_data.get("workflow_id"),
|
|
1874
|
+
existing_data.get("team_id"),
|
|
1875
|
+
existing_data.get("agent_id"),
|
|
1876
|
+
existing_data.get("name", ""),
|
|
1877
|
+
)
|
|
1878
|
+
new_level = get_component_level(trace.workflow_id, trace.team_id, trace.agent_id, trace.name)
|
|
1879
|
+
should_update_name = new_level > existing_level
|
|
1880
|
+
|
|
1881
|
+
# Parse existing start_time to calculate correct duration
|
|
1882
|
+
existing_start_time_str = existing_data.get("start_time")
|
|
1883
|
+
if isinstance(existing_start_time_str, str):
|
|
1884
|
+
existing_start_time = datetime.fromisoformat(existing_start_time_str.replace("Z", "+00:00"))
|
|
1885
|
+
else:
|
|
1886
|
+
existing_start_time = trace.start_time
|
|
1887
|
+
|
|
1888
|
+
recalculated_duration_ms = int((trace.end_time - existing_start_time).total_seconds() * 1000)
|
|
1889
|
+
|
|
1890
|
+
update_values: Dict[str, Any] = {
|
|
1891
|
+
"end_time": trace.end_time.isoformat(),
|
|
1892
|
+
"duration_ms": recalculated_duration_ms,
|
|
1893
|
+
"status": trace.status,
|
|
1894
|
+
}
|
|
1895
|
+
|
|
1896
|
+
if should_update_name:
|
|
1897
|
+
update_values["name"] = trace.name
|
|
1898
|
+
|
|
1899
|
+
# Update context fields only if new value is not None
|
|
1900
|
+
if trace.run_id is not None:
|
|
1901
|
+
update_values["run_id"] = trace.run_id
|
|
1902
|
+
if trace.session_id is not None:
|
|
1903
|
+
update_values["session_id"] = trace.session_id
|
|
1904
|
+
if trace.user_id is not None:
|
|
1905
|
+
update_values["user_id"] = trace.user_id
|
|
1906
|
+
if trace.agent_id is not None:
|
|
1907
|
+
update_values["agent_id"] = trace.agent_id
|
|
1908
|
+
if trace.team_id is not None:
|
|
1909
|
+
update_values["team_id"] = trace.team_id
|
|
1910
|
+
if trace.workflow_id is not None:
|
|
1911
|
+
update_values["workflow_id"] = trace.workflow_id
|
|
1912
|
+
|
|
1913
|
+
existing_doc.reference.update(update_values)
|
|
1914
|
+
else:
|
|
1915
|
+
# Create new trace with initialized counters
|
|
1916
|
+
trace_dict = trace.to_dict()
|
|
1917
|
+
trace_dict["total_spans"] = 0
|
|
1918
|
+
trace_dict["error_count"] = 0
|
|
1919
|
+
collection_ref.add(trace_dict)
|
|
1920
|
+
|
|
1921
|
+
except Exception as e:
|
|
1922
|
+
log_error(f"Error creating trace: {e}")
|
|
1923
|
+
|
|
1924
|
+
def get_trace(
|
|
1925
|
+
self,
|
|
1926
|
+
trace_id: Optional[str] = None,
|
|
1927
|
+
run_id: Optional[str] = None,
|
|
1928
|
+
):
|
|
1929
|
+
"""Get a single trace by trace_id or other filters.
|
|
1930
|
+
|
|
1931
|
+
Args:
|
|
1932
|
+
trace_id: The unique trace identifier.
|
|
1933
|
+
run_id: Filter by run ID (returns first match).
|
|
1934
|
+
|
|
1935
|
+
Returns:
|
|
1936
|
+
Optional[Trace]: The trace if found, None otherwise.
|
|
1937
|
+
|
|
1938
|
+
Note:
|
|
1939
|
+
If multiple filters are provided, trace_id takes precedence.
|
|
1940
|
+
For other filters, the most recent trace is returned.
|
|
1941
|
+
"""
|
|
1942
|
+
try:
|
|
1943
|
+
from agno.tracing.schemas import Trace
|
|
1944
|
+
|
|
1945
|
+
collection_ref = self._get_collection(table_type="traces")
|
|
1946
|
+
if collection_ref is None:
|
|
1947
|
+
return None
|
|
1948
|
+
|
|
1949
|
+
if trace_id:
|
|
1950
|
+
docs = collection_ref.where(filter=FieldFilter("trace_id", "==", trace_id)).limit(1).stream()
|
|
1951
|
+
elif run_id:
|
|
1952
|
+
from google.cloud.firestore import Query
|
|
1953
|
+
|
|
1954
|
+
docs = (
|
|
1955
|
+
collection_ref.where(filter=FieldFilter("run_id", "==", run_id))
|
|
1956
|
+
.order_by("start_time", direction=Query.DESCENDING)
|
|
1957
|
+
.limit(1)
|
|
1958
|
+
.stream()
|
|
1959
|
+
)
|
|
1960
|
+
else:
|
|
1961
|
+
log_debug("get_trace called without any filter parameters")
|
|
1962
|
+
return None
|
|
1963
|
+
|
|
1964
|
+
for doc in docs:
|
|
1965
|
+
trace_data = doc.to_dict()
|
|
1966
|
+
# Use stored values (default to 0 if not present)
|
|
1967
|
+
trace_data.setdefault("total_spans", 0)
|
|
1968
|
+
trace_data.setdefault("error_count", 0)
|
|
1969
|
+
return Trace.from_dict(trace_data)
|
|
1970
|
+
|
|
1971
|
+
return None
|
|
1972
|
+
|
|
1973
|
+
except Exception as e:
|
|
1974
|
+
log_error(f"Error getting trace: {e}")
|
|
1975
|
+
return None
|
|
1976
|
+
|
|
1977
|
+
def get_traces(
|
|
1978
|
+
self,
|
|
1979
|
+
run_id: Optional[str] = None,
|
|
1980
|
+
session_id: Optional[str] = None,
|
|
1981
|
+
user_id: Optional[str] = None,
|
|
1982
|
+
agent_id: Optional[str] = None,
|
|
1983
|
+
team_id: Optional[str] = None,
|
|
1984
|
+
workflow_id: Optional[str] = None,
|
|
1985
|
+
status: Optional[str] = None,
|
|
1986
|
+
start_time: Optional[datetime] = None,
|
|
1987
|
+
end_time: Optional[datetime] = None,
|
|
1988
|
+
limit: Optional[int] = 20,
|
|
1989
|
+
page: Optional[int] = 1,
|
|
1990
|
+
) -> tuple[List, int]:
|
|
1991
|
+
"""Get traces matching the provided filters.
|
|
1992
|
+
|
|
1993
|
+
Args:
|
|
1994
|
+
run_id: Filter by run ID.
|
|
1995
|
+
session_id: Filter by session ID.
|
|
1996
|
+
user_id: Filter by user ID.
|
|
1997
|
+
agent_id: Filter by agent ID.
|
|
1998
|
+
team_id: Filter by team ID.
|
|
1999
|
+
workflow_id: Filter by workflow ID.
|
|
2000
|
+
status: Filter by status (OK, ERROR, UNSET).
|
|
2001
|
+
start_time: Filter traces starting after this datetime.
|
|
2002
|
+
end_time: Filter traces ending before this datetime.
|
|
2003
|
+
limit: Maximum number of traces to return per page.
|
|
2004
|
+
page: Page number (1-indexed).
|
|
2005
|
+
|
|
2006
|
+
Returns:
|
|
2007
|
+
tuple[List[Trace], int]: Tuple of (list of matching traces, total count).
|
|
2008
|
+
"""
|
|
2009
|
+
try:
|
|
2010
|
+
from agno.tracing.schemas import Trace
|
|
2011
|
+
|
|
2012
|
+
collection_ref = self._get_collection(table_type="traces")
|
|
2013
|
+
if collection_ref is None:
|
|
2014
|
+
return [], 0
|
|
2015
|
+
|
|
2016
|
+
query = collection_ref
|
|
2017
|
+
|
|
2018
|
+
# Apply filters
|
|
2019
|
+
if run_id:
|
|
2020
|
+
query = query.where(filter=FieldFilter("run_id", "==", run_id))
|
|
2021
|
+
if session_id:
|
|
2022
|
+
query = query.where(filter=FieldFilter("session_id", "==", session_id))
|
|
2023
|
+
if user_id:
|
|
2024
|
+
query = query.where(filter=FieldFilter("user_id", "==", user_id))
|
|
2025
|
+
if agent_id:
|
|
2026
|
+
query = query.where(filter=FieldFilter("agent_id", "==", agent_id))
|
|
2027
|
+
if team_id:
|
|
2028
|
+
query = query.where(filter=FieldFilter("team_id", "==", team_id))
|
|
2029
|
+
if workflow_id:
|
|
2030
|
+
query = query.where(filter=FieldFilter("workflow_id", "==", workflow_id))
|
|
2031
|
+
if status:
|
|
2032
|
+
query = query.where(filter=FieldFilter("status", "==", status))
|
|
2033
|
+
if start_time:
|
|
2034
|
+
query = query.where(filter=FieldFilter("start_time", ">=", start_time.isoformat()))
|
|
2035
|
+
if end_time:
|
|
2036
|
+
query = query.where(filter=FieldFilter("end_time", "<=", end_time.isoformat()))
|
|
2037
|
+
|
|
2038
|
+
# Get all matching documents
|
|
2039
|
+
docs = query.stream()
|
|
2040
|
+
all_records = [doc.to_dict() for doc in docs]
|
|
2041
|
+
|
|
2042
|
+
# Sort by start_time descending
|
|
2043
|
+
all_records.sort(key=lambda x: x.get("start_time", ""), reverse=True)
|
|
2044
|
+
|
|
2045
|
+
# Get total count
|
|
2046
|
+
total_count = len(all_records)
|
|
2047
|
+
|
|
2048
|
+
# Apply pagination
|
|
2049
|
+
if limit and page:
|
|
2050
|
+
offset = (page - 1) * limit
|
|
2051
|
+
paginated_records = all_records[offset : offset + limit]
|
|
2052
|
+
elif limit:
|
|
2053
|
+
paginated_records = all_records[:limit]
|
|
2054
|
+
else:
|
|
2055
|
+
paginated_records = all_records
|
|
2056
|
+
|
|
2057
|
+
# Convert to Trace objects with stored span counts
|
|
2058
|
+
traces = []
|
|
2059
|
+
for trace_data in paginated_records:
|
|
2060
|
+
trace_data.setdefault("total_spans", 0)
|
|
2061
|
+
trace_data.setdefault("error_count", 0)
|
|
2062
|
+
traces.append(Trace.from_dict(trace_data))
|
|
2063
|
+
|
|
2064
|
+
return traces, total_count
|
|
2065
|
+
|
|
2066
|
+
except Exception as e:
|
|
2067
|
+
log_error(f"Error getting traces: {e}")
|
|
2068
|
+
return [], 0
|
|
2069
|
+
|
|
2070
|
+
def get_trace_stats(
|
|
2071
|
+
self,
|
|
2072
|
+
user_id: Optional[str] = None,
|
|
2073
|
+
agent_id: Optional[str] = None,
|
|
2074
|
+
team_id: Optional[str] = None,
|
|
2075
|
+
workflow_id: Optional[str] = None,
|
|
2076
|
+
start_time: Optional[datetime] = None,
|
|
2077
|
+
end_time: Optional[datetime] = None,
|
|
2078
|
+
limit: Optional[int] = 20,
|
|
2079
|
+
page: Optional[int] = 1,
|
|
2080
|
+
) -> tuple[List[Dict[str, Any]], int]:
|
|
2081
|
+
"""Get trace statistics grouped by session.
|
|
2082
|
+
|
|
2083
|
+
Args:
|
|
2084
|
+
user_id: Filter by user ID.
|
|
2085
|
+
agent_id: Filter by agent ID.
|
|
2086
|
+
team_id: Filter by team ID.
|
|
2087
|
+
workflow_id: Filter by workflow ID.
|
|
2088
|
+
start_time: Filter sessions with traces created after this datetime.
|
|
2089
|
+
end_time: Filter sessions with traces created before this datetime.
|
|
2090
|
+
limit: Maximum number of sessions to return per page.
|
|
2091
|
+
page: Page number (1-indexed).
|
|
2092
|
+
|
|
2093
|
+
Returns:
|
|
2094
|
+
tuple[List[Dict], int]: Tuple of (list of session stats dicts, total count).
|
|
2095
|
+
Each dict contains: session_id, user_id, agent_id, team_id, workflow_id, total_traces,
|
|
2096
|
+
first_trace_at, last_trace_at.
|
|
2097
|
+
"""
|
|
2098
|
+
try:
|
|
2099
|
+
collection_ref = self._get_collection(table_type="traces")
|
|
2100
|
+
if collection_ref is None:
|
|
2101
|
+
return [], 0
|
|
2102
|
+
|
|
2103
|
+
query = collection_ref
|
|
2104
|
+
|
|
2105
|
+
# Apply filters
|
|
2106
|
+
if user_id:
|
|
2107
|
+
query = query.where(filter=FieldFilter("user_id", "==", user_id))
|
|
2108
|
+
if agent_id:
|
|
2109
|
+
query = query.where(filter=FieldFilter("agent_id", "==", agent_id))
|
|
2110
|
+
if team_id:
|
|
2111
|
+
query = query.where(filter=FieldFilter("team_id", "==", team_id))
|
|
2112
|
+
if workflow_id:
|
|
2113
|
+
query = query.where(filter=FieldFilter("workflow_id", "==", workflow_id))
|
|
2114
|
+
if start_time:
|
|
2115
|
+
query = query.where(filter=FieldFilter("created_at", ">=", start_time.isoformat()))
|
|
2116
|
+
if end_time:
|
|
2117
|
+
query = query.where(filter=FieldFilter("created_at", "<=", end_time.isoformat()))
|
|
2118
|
+
|
|
2119
|
+
# Get all matching documents
|
|
2120
|
+
docs = query.stream()
|
|
2121
|
+
|
|
2122
|
+
# Aggregate by session_id
|
|
2123
|
+
session_stats: Dict[str, Dict[str, Any]] = {}
|
|
2124
|
+
for doc in docs:
|
|
2125
|
+
trace_data = doc.to_dict()
|
|
2126
|
+
session_id = trace_data.get("session_id")
|
|
2127
|
+
if not session_id:
|
|
2128
|
+
continue
|
|
2129
|
+
|
|
2130
|
+
if session_id not in session_stats:
|
|
2131
|
+
session_stats[session_id] = {
|
|
2132
|
+
"session_id": session_id,
|
|
2133
|
+
"user_id": trace_data.get("user_id"),
|
|
2134
|
+
"agent_id": trace_data.get("agent_id"),
|
|
2135
|
+
"team_id": trace_data.get("team_id"),
|
|
2136
|
+
"workflow_id": trace_data.get("workflow_id"),
|
|
2137
|
+
"total_traces": 0,
|
|
2138
|
+
"first_trace_at": trace_data.get("created_at"),
|
|
2139
|
+
"last_trace_at": trace_data.get("created_at"),
|
|
2140
|
+
}
|
|
2141
|
+
|
|
2142
|
+
session_stats[session_id]["total_traces"] += 1
|
|
2143
|
+
|
|
2144
|
+
created_at = trace_data.get("created_at")
|
|
2145
|
+
if (
|
|
2146
|
+
created_at
|
|
2147
|
+
and session_stats[session_id]["first_trace_at"]
|
|
2148
|
+
and session_stats[session_id]["last_trace_at"]
|
|
2149
|
+
):
|
|
2150
|
+
if created_at < session_stats[session_id]["first_trace_at"]:
|
|
2151
|
+
session_stats[session_id]["first_trace_at"] = created_at
|
|
2152
|
+
if created_at > session_stats[session_id]["last_trace_at"]:
|
|
2153
|
+
session_stats[session_id]["last_trace_at"] = created_at
|
|
2154
|
+
|
|
2155
|
+
# Convert to list and sort by last_trace_at descending
|
|
2156
|
+
stats_list = list(session_stats.values())
|
|
2157
|
+
stats_list.sort(key=lambda x: x.get("last_trace_at", ""), reverse=True)
|
|
2158
|
+
|
|
2159
|
+
# Convert datetime strings to datetime objects
|
|
2160
|
+
for stat in stats_list:
|
|
2161
|
+
first_trace_at = stat["first_trace_at"]
|
|
2162
|
+
last_trace_at = stat["last_trace_at"]
|
|
2163
|
+
if isinstance(first_trace_at, str):
|
|
2164
|
+
stat["first_trace_at"] = datetime.fromisoformat(first_trace_at.replace("Z", "+00:00"))
|
|
2165
|
+
if isinstance(last_trace_at, str):
|
|
2166
|
+
stat["last_trace_at"] = datetime.fromisoformat(last_trace_at.replace("Z", "+00:00"))
|
|
2167
|
+
|
|
2168
|
+
# Get total count
|
|
2169
|
+
total_count = len(stats_list)
|
|
2170
|
+
|
|
2171
|
+
# Apply pagination
|
|
2172
|
+
if limit and page:
|
|
2173
|
+
offset = (page - 1) * limit
|
|
2174
|
+
paginated_stats = stats_list[offset : offset + limit]
|
|
2175
|
+
elif limit:
|
|
2176
|
+
paginated_stats = stats_list[:limit]
|
|
2177
|
+
else:
|
|
2178
|
+
paginated_stats = stats_list
|
|
2179
|
+
|
|
2180
|
+
return paginated_stats, total_count
|
|
2181
|
+
|
|
2182
|
+
except Exception as e:
|
|
2183
|
+
log_error(f"Error getting trace stats: {e}")
|
|
2184
|
+
return [], 0
|
|
2185
|
+
|
|
2186
|
+
# --- Spans ---
|
|
2187
|
+
def create_span(self, span: "Span") -> None:
|
|
2188
|
+
"""Create a single span in the database.
|
|
2189
|
+
|
|
2190
|
+
Args:
|
|
2191
|
+
span: The Span object to store.
|
|
2192
|
+
"""
|
|
2193
|
+
try:
|
|
2194
|
+
collection_ref = self._get_collection(table_type="spans", create_collection_if_not_found=True)
|
|
2195
|
+
if collection_ref is None:
|
|
2196
|
+
return
|
|
2197
|
+
|
|
2198
|
+
span_dict = span.to_dict()
|
|
2199
|
+
# Serialize attributes as JSON string
|
|
2200
|
+
if "attributes" in span_dict and isinstance(span_dict["attributes"], dict):
|
|
2201
|
+
span_dict["attributes"] = json.dumps(span_dict["attributes"])
|
|
2202
|
+
|
|
2203
|
+
collection_ref.add(span_dict)
|
|
2204
|
+
|
|
2205
|
+
# Increment total_spans and error_count on trace
|
|
2206
|
+
traces_collection = self._get_collection(table_type="traces")
|
|
2207
|
+
if traces_collection:
|
|
2208
|
+
try:
|
|
2209
|
+
docs = (
|
|
2210
|
+
traces_collection.where(filter=FieldFilter("trace_id", "==", span.trace_id)).limit(1).stream()
|
|
2211
|
+
)
|
|
2212
|
+
for doc in docs:
|
|
2213
|
+
trace_data = doc.to_dict()
|
|
2214
|
+
current_total = trace_data.get("total_spans", 0)
|
|
2215
|
+
current_errors = trace_data.get("error_count", 0)
|
|
2216
|
+
|
|
2217
|
+
update_values = {"total_spans": current_total + 1}
|
|
2218
|
+
if span.status_code == "ERROR":
|
|
2219
|
+
update_values["error_count"] = current_errors + 1
|
|
2220
|
+
|
|
2221
|
+
doc.reference.update(update_values)
|
|
2222
|
+
break
|
|
2223
|
+
except Exception as update_error:
|
|
2224
|
+
log_debug(f"Could not update trace span counts: {update_error}")
|
|
2225
|
+
|
|
2226
|
+
except Exception as e:
|
|
2227
|
+
log_error(f"Error creating span: {e}")
|
|
2228
|
+
|
|
2229
|
+
def create_spans(self, spans: List) -> None:
|
|
2230
|
+
"""Create multiple spans in the database as a batch.
|
|
2231
|
+
|
|
2232
|
+
Args:
|
|
2233
|
+
spans: List of Span objects to store.
|
|
2234
|
+
"""
|
|
2235
|
+
if not spans:
|
|
2236
|
+
return
|
|
2237
|
+
|
|
2238
|
+
try:
|
|
2239
|
+
collection_ref = self._get_collection(table_type="spans", create_collection_if_not_found=True)
|
|
2240
|
+
if collection_ref is None:
|
|
2241
|
+
return
|
|
2242
|
+
|
|
2243
|
+
# Firestore batch has a limit of 500 operations
|
|
2244
|
+
batch = self.db_client.batch()
|
|
2245
|
+
batch_count = 0
|
|
2246
|
+
|
|
2247
|
+
for span in spans:
|
|
2248
|
+
span_dict = span.to_dict()
|
|
2249
|
+
# Serialize attributes as JSON string
|
|
2250
|
+
if "attributes" in span_dict and isinstance(span_dict["attributes"], dict):
|
|
2251
|
+
span_dict["attributes"] = json.dumps(span_dict["attributes"])
|
|
2252
|
+
|
|
2253
|
+
doc_ref = collection_ref.document()
|
|
2254
|
+
batch.set(doc_ref, span_dict)
|
|
2255
|
+
batch_count += 1
|
|
2256
|
+
|
|
2257
|
+
# Commit batch if reaching limit
|
|
2258
|
+
if batch_count >= 500:
|
|
2259
|
+
batch.commit()
|
|
2260
|
+
batch = self.db_client.batch()
|
|
2261
|
+
batch_count = 0
|
|
2262
|
+
|
|
2263
|
+
# Commit remaining operations
|
|
2264
|
+
if batch_count > 0:
|
|
2265
|
+
batch.commit()
|
|
2266
|
+
|
|
2267
|
+
# Update trace with total_spans and error_count
|
|
2268
|
+
trace_id = spans[0].trace_id
|
|
2269
|
+
spans_count = len(spans)
|
|
2270
|
+
error_count = sum(1 for s in spans if s.status_code == "ERROR")
|
|
2271
|
+
|
|
2272
|
+
traces_collection = self._get_collection(table_type="traces")
|
|
2273
|
+
if traces_collection:
|
|
2274
|
+
try:
|
|
2275
|
+
docs = traces_collection.where(filter=FieldFilter("trace_id", "==", trace_id)).limit(1).stream()
|
|
2276
|
+
for doc in docs:
|
|
2277
|
+
trace_data = doc.to_dict()
|
|
2278
|
+
current_total = trace_data.get("total_spans", 0)
|
|
2279
|
+
current_errors = trace_data.get("error_count", 0)
|
|
2280
|
+
|
|
2281
|
+
doc.reference.update(
|
|
2282
|
+
{
|
|
2283
|
+
"total_spans": current_total + spans_count,
|
|
2284
|
+
"error_count": current_errors + error_count,
|
|
2285
|
+
}
|
|
2286
|
+
)
|
|
2287
|
+
break
|
|
2288
|
+
except Exception as update_error:
|
|
2289
|
+
log_debug(f"Could not update trace span counts: {update_error}")
|
|
2290
|
+
|
|
2291
|
+
except Exception as e:
|
|
2292
|
+
log_error(f"Error creating spans batch: {e}")
|
|
2293
|
+
|
|
2294
|
+
def get_span(self, span_id: str):
|
|
2295
|
+
"""Get a single span by its span_id.
|
|
2296
|
+
|
|
2297
|
+
Args:
|
|
2298
|
+
span_id: The unique span identifier.
|
|
2299
|
+
|
|
2300
|
+
Returns:
|
|
2301
|
+
Optional[Span]: The span if found, None otherwise.
|
|
2302
|
+
"""
|
|
2303
|
+
try:
|
|
2304
|
+
from agno.tracing.schemas import Span
|
|
2305
|
+
|
|
2306
|
+
collection_ref = self._get_collection(table_type="spans")
|
|
2307
|
+
if collection_ref is None:
|
|
2308
|
+
return None
|
|
2309
|
+
|
|
2310
|
+
docs = collection_ref.where(filter=FieldFilter("span_id", "==", span_id)).limit(1).stream()
|
|
2311
|
+
|
|
2312
|
+
for doc in docs:
|
|
2313
|
+
span_data = doc.to_dict()
|
|
2314
|
+
# Deserialize attributes from JSON string
|
|
2315
|
+
if "attributes" in span_data and isinstance(span_data["attributes"], str):
|
|
2316
|
+
span_data["attributes"] = json.loads(span_data["attributes"])
|
|
2317
|
+
return Span.from_dict(span_data)
|
|
2318
|
+
|
|
2319
|
+
return None
|
|
2320
|
+
|
|
2321
|
+
except Exception as e:
|
|
2322
|
+
log_error(f"Error getting span: {e}")
|
|
2323
|
+
return None
|
|
2324
|
+
|
|
2325
|
+
def get_spans(
|
|
2326
|
+
self,
|
|
2327
|
+
trace_id: Optional[str] = None,
|
|
2328
|
+
parent_span_id: Optional[str] = None,
|
|
2329
|
+
limit: Optional[int] = 1000,
|
|
2330
|
+
) -> List:
|
|
2331
|
+
"""Get spans matching the provided filters.
|
|
2332
|
+
|
|
2333
|
+
Args:
|
|
2334
|
+
trace_id: Filter by trace ID.
|
|
2335
|
+
parent_span_id: Filter by parent span ID.
|
|
2336
|
+
limit: Maximum number of spans to return.
|
|
2337
|
+
|
|
2338
|
+
Returns:
|
|
2339
|
+
List[Span]: List of matching spans.
|
|
2340
|
+
"""
|
|
2341
|
+
try:
|
|
2342
|
+
from agno.tracing.schemas import Span
|
|
2343
|
+
|
|
2344
|
+
collection_ref = self._get_collection(table_type="spans")
|
|
2345
|
+
if collection_ref is None:
|
|
2346
|
+
return []
|
|
2347
|
+
|
|
2348
|
+
query = collection_ref
|
|
2349
|
+
|
|
2350
|
+
if trace_id:
|
|
2351
|
+
query = query.where(filter=FieldFilter("trace_id", "==", trace_id))
|
|
2352
|
+
if parent_span_id:
|
|
2353
|
+
query = query.where(filter=FieldFilter("parent_span_id", "==", parent_span_id))
|
|
2354
|
+
|
|
2355
|
+
if limit:
|
|
2356
|
+
query = query.limit(limit)
|
|
2357
|
+
|
|
2358
|
+
docs = query.stream()
|
|
2359
|
+
|
|
2360
|
+
spans = []
|
|
2361
|
+
for doc in docs:
|
|
2362
|
+
span_data = doc.to_dict()
|
|
2363
|
+
# Deserialize attributes from JSON string
|
|
2364
|
+
if "attributes" in span_data and isinstance(span_data["attributes"], str):
|
|
2365
|
+
span_data["attributes"] = json.loads(span_data["attributes"])
|
|
2366
|
+
spans.append(Span.from_dict(span_data))
|
|
2367
|
+
|
|
2368
|
+
return spans
|
|
2369
|
+
|
|
2370
|
+
except Exception as e:
|
|
2371
|
+
log_error(f"Error getting spans: {e}")
|
|
2372
|
+
return []
|
|
2373
|
+
|
|
2374
|
+
# -- Learning methods (stubs) --
|
|
2375
|
+
def get_learning(
|
|
2376
|
+
self,
|
|
2377
|
+
learning_type: str,
|
|
2378
|
+
user_id: Optional[str] = None,
|
|
2379
|
+
agent_id: Optional[str] = None,
|
|
2380
|
+
team_id: Optional[str] = None,
|
|
2381
|
+
session_id: Optional[str] = None,
|
|
2382
|
+
namespace: Optional[str] = None,
|
|
2383
|
+
entity_id: Optional[str] = None,
|
|
2384
|
+
entity_type: Optional[str] = None,
|
|
2385
|
+
) -> Optional[Dict[str, Any]]:
|
|
2386
|
+
raise NotImplementedError("Learning methods not yet implemented for FirestoreDb")
|
|
2387
|
+
|
|
2388
|
+
def upsert_learning(
|
|
2389
|
+
self,
|
|
2390
|
+
id: str,
|
|
2391
|
+
learning_type: str,
|
|
2392
|
+
content: Dict[str, Any],
|
|
2393
|
+
user_id: Optional[str] = None,
|
|
2394
|
+
agent_id: Optional[str] = None,
|
|
2395
|
+
team_id: Optional[str] = None,
|
|
2396
|
+
session_id: Optional[str] = None,
|
|
2397
|
+
namespace: Optional[str] = None,
|
|
2398
|
+
entity_id: Optional[str] = None,
|
|
2399
|
+
entity_type: Optional[str] = None,
|
|
2400
|
+
metadata: Optional[Dict[str, Any]] = None,
|
|
2401
|
+
) -> None:
|
|
2402
|
+
raise NotImplementedError("Learning methods not yet implemented for FirestoreDb")
|
|
2403
|
+
|
|
2404
|
+
def delete_learning(self, id: str) -> bool:
|
|
2405
|
+
raise NotImplementedError("Learning methods not yet implemented for FirestoreDb")
|
|
2406
|
+
|
|
2407
|
+
def get_learnings(
|
|
2408
|
+
self,
|
|
2409
|
+
learning_type: Optional[str] = None,
|
|
2410
|
+
user_id: Optional[str] = None,
|
|
2411
|
+
agent_id: Optional[str] = None,
|
|
2412
|
+
team_id: Optional[str] = None,
|
|
2413
|
+
session_id: Optional[str] = None,
|
|
2414
|
+
namespace: Optional[str] = None,
|
|
2415
|
+
entity_id: Optional[str] = None,
|
|
2416
|
+
entity_type: Optional[str] = None,
|
|
2417
|
+
limit: Optional[int] = None,
|
|
2418
|
+
) -> List[Dict[str, Any]]:
|
|
2419
|
+
raise NotImplementedError("Learning methods not yet implemented for FirestoreDb")
|