agno 2.2.13__py3-none-any.whl → 2.4.3__py3-none-any.whl

agno/db/mongo/async_mongo.py

@@ -1,15 +1,19 @@
+import asyncio
 import time
 from datetime import date, datetime, timedelta, timezone
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 from uuid import uuid4
 
+if TYPE_CHECKING:
+    from agno.tracing.schemas import Span, Trace
+
 from agno.db.base import AsyncBaseDb, SessionType
 from agno.db.mongo.utils import (
     apply_pagination,
     apply_sorting,
     bulk_upsert_metrics,
     calculate_date_metrics,
-    create_collection_indexes,
+    create_collection_indexes_async,
     deserialize_cultural_knowledge_from_db,
     fetch_all_sessions_data,
     get_dates_to_calculate_metrics_for,
@@ -25,11 +29,26 @@ from agno.utils.log import log_debug, log_error, log_info
 from agno.utils.string import generate_id
 
 try:
-    import asyncio
-
     from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorCollection, AsyncIOMotorDatabase  # type: ignore
+
+    MOTOR_AVAILABLE = True
+except ImportError:
+    MOTOR_AVAILABLE = False
+    AsyncIOMotorClient = None  # type: ignore
+    AsyncIOMotorCollection = None  # type: ignore
+    AsyncIOMotorDatabase = None  # type: ignore
+
+try:
+    from pymongo import AsyncMongoClient  # type: ignore
+    from pymongo.collection import AsyncCollection  # type: ignore
+    from pymongo.database import AsyncDatabase  # type: ignore
+
+    PYMONGO_ASYNC_AVAILABLE = True
 except ImportError:
-    raise ImportError("`motor` not installed. Please install it using `pip install -U motor`")
+    PYMONGO_ASYNC_AVAILABLE = False
+    AsyncMongoClient = None  # type: ignore
+    AsyncDatabase = None  # type: ignore
+    AsyncCollection = None  # type: ignore
 
 try:
     from pymongo import ReturnDocument
@@ -37,11 +56,89 @@ try:
 except ImportError:
     raise ImportError("`pymongo` not installed. Please install it using `pip install -U pymongo`")
 
+# Ensure at least one async library is available
+if not MOTOR_AVAILABLE and not PYMONGO_ASYNC_AVAILABLE:
+    raise ImportError(
+        "Neither `motor` nor PyMongo async is installed. "
+        "Please install one of them using:\n"
+        "  - `pip install -U 'pymongo>=4.9'` (recommended)"
+        "  - `pip install -U motor` (legacy, deprecated)\n"
+    )
+
+# Create union types for client, database, and collection
+if TYPE_CHECKING:
+    if MOTOR_AVAILABLE and PYMONGO_ASYNC_AVAILABLE:
+        AsyncMongoClientType = Union[AsyncIOMotorClient, AsyncMongoClient]  # type: ignore
+        AsyncMongoDatabaseType = Union[AsyncIOMotorDatabase, AsyncDatabase]  # type: ignore
+        AsyncMongoCollectionType = Union[AsyncIOMotorCollection, AsyncCollection]  # type: ignore
+    elif MOTOR_AVAILABLE:
+        AsyncMongoClientType = AsyncIOMotorClient  # type: ignore
+        AsyncMongoDatabaseType = AsyncIOMotorDatabase  # type: ignore
+        AsyncMongoCollectionType = AsyncIOMotorCollection  # type: ignore
+    else:
+        AsyncMongoClientType = AsyncMongoClient  # type: ignore
+        AsyncMongoDatabaseType = AsyncDatabase  # type: ignore
+        AsyncMongoCollectionType = AsyncCollection  # type: ignore
+else:
+    # Runtime type - use Any to avoid import issues
+    AsyncMongoClientType = Any
+    AsyncMongoDatabaseType = Any
+    AsyncMongoCollectionType = Any
+
+
+# Client type constants (defined before class to allow use in _detect_client_type)
+_CLIENT_TYPE_MOTOR = "motor"
+_CLIENT_TYPE_PYMONGO_ASYNC = "pymongo_async"
+_CLIENT_TYPE_UNKNOWN = "unknown"
+
+
+def _detect_client_type(client: Any) -> str:
+    """Detect whether a client is Motor or PyMongo async."""
+    if client is None:
+        return _CLIENT_TYPE_UNKNOWN
+
+    # Check PyMongo async
+    if PYMONGO_ASYNC_AVAILABLE and AsyncMongoClient is not None:
+        try:
+            if isinstance(client, AsyncMongoClient):
+                return _CLIENT_TYPE_PYMONGO_ASYNC
+        except (TypeError, AttributeError):
+            pass  # Fall through to next check
+
+    if MOTOR_AVAILABLE and AsyncIOMotorClient is not None:
+        try:
+            if isinstance(client, AsyncIOMotorClient):
+                return _CLIENT_TYPE_MOTOR
+        except (TypeError, AttributeError):
+            pass  # Fall through to fallback
+
+    # Fallback to string matching only if isinstance fails
+    # (should rarely happen, but useful for edge cases)
+    client_type_name = type(client).__name__
+    if "Motor" in client_type_name or "AsyncIOMotor" in client_type_name:
+        return _CLIENT_TYPE_MOTOR
+    elif "AsyncMongo" in client_type_name:
+        return _CLIENT_TYPE_PYMONGO_ASYNC
+
+    # Last resort: check module name
+    module_name = type(client).__module__
+    if "motor" in module_name:
+        return _CLIENT_TYPE_MOTOR
+    elif "pymongo" in module_name:
+        return _CLIENT_TYPE_PYMONGO_ASYNC
+
+    return _CLIENT_TYPE_UNKNOWN
+
 
 class AsyncMongoDb(AsyncBaseDb):
+    # Client type constants (class-level access to module constants)
+    CLIENT_TYPE_MOTOR = _CLIENT_TYPE_MOTOR
+    CLIENT_TYPE_PYMONGO_ASYNC = _CLIENT_TYPE_PYMONGO_ASYNC
+    CLIENT_TYPE_UNKNOWN = _CLIENT_TYPE_UNKNOWN
+
     def __init__(
         self,
-        db_client: Optional[AsyncIOMotorClient] = None,
+        db_client: Optional[Union["AsyncIOMotorClient", "AsyncMongoClient"]] = None,
         db_name: Optional[str] = None,
         db_url: Optional[str] = None,
         session_collection: Optional[str] = None,
@@ -50,13 +147,22 @@ class AsyncMongoDb(AsyncBaseDb):
         eval_collection: Optional[str] = None,
         knowledge_collection: Optional[str] = None,
         culture_collection: Optional[str] = None,
+        traces_collection: Optional[str] = None,
+        spans_collection: Optional[str] = None,
+        learnings_collection: Optional[str] = None,
         id: Optional[str] = None,
     ):
         """
-        Async interface for interacting with a MongoDB database using Motor.
+        Async interface for interacting with a MongoDB database.
+
+        Supports both Motor (legacy) and PyMongo async (recommended) clients.
+        When both libraries are available, PyMongo async is preferred.
 
         Args:
-            db_client (Optional[AsyncIOMotorClient]): The MongoDB async client to use.
+            db_client (Optional[Union[AsyncIOMotorClient, AsyncMongoClient]]):
+                The MongoDB async client to use. Can be either Motor's AsyncIOMotorClient
+                or PyMongo's AsyncMongoClient. If not provided, a client will be created
+                from db_url using the preferred available library.
             db_name (Optional[str]): The name of the database to use.
             db_url (Optional[str]): The database URL to connect to.
             session_collection (Optional[str]): Name of the collection to store sessions.
@@ -65,10 +171,14 @@ class AsyncMongoDb(AsyncBaseDb):
             eval_collection (Optional[str]): Name of the collection to store evaluation runs.
             knowledge_collection (Optional[str]): Name of the collection to store knowledge documents.
             culture_collection (Optional[str]): Name of the collection to store cultural knowledge.
+            traces_collection (Optional[str]): Name of the collection to store traces.
+            spans_collection (Optional[str]): Name of the collection to store spans.
+            learnings_collection (Optional[str]): Name of the collection to store learnings.
             id (Optional[str]): ID of the database.
 
         Raises:
-            ValueError: If neither db_url nor db_client is provided.
+            ValueError: If neither db_url nor db_client is provided, or if db_client type is unsupported.
+            ImportError: If neither motor nor pymongo async is installed.
         """
         if id is None:
             base_seed = db_url or str(db_client)
@@ -84,10 +194,26 @@ class AsyncMongoDb(AsyncBaseDb):
             eval_table=eval_collection,
             knowledge_table=knowledge_collection,
             culture_table=culture_collection,
+            traces_table=traces_collection,
+            spans_table=spans_collection,
+            learnings_table=learnings_collection,
         )
 
+        # Detect client type if provided
+        if db_client is not None:
+            self._client_type = _detect_client_type(db_client)
+            if self._client_type == self.CLIENT_TYPE_UNKNOWN:
+                raise ValueError(
+                    f"Unsupported MongoDB client type: {type(db_client).__name__}. "
+                    "Only Motor (AsyncIOMotorClient) or PyMongo async (AsyncMongoClient) are supported."
+                )
+        else:
+            # Auto-select preferred library when creating from URL
+            # Prefer PyMongo async if available, fallback to Motor
+            self._client_type = self.CLIENT_TYPE_PYMONGO_ASYNC if PYMONGO_ASYNC_AVAILABLE else self.CLIENT_TYPE_MOTOR
+
         # Store configuration for lazy initialization
-        self._provided_client: Optional[AsyncIOMotorClient] = db_client
+        self._provided_client: Optional[AsyncMongoClientType] = db_client
         self.db_url: Optional[str] = db_url
         self.db_name: str = db_name if db_name is not None else "agno"
 
@@ -95,8 +221,8 @@ class AsyncMongoDb(AsyncBaseDb):
             raise ValueError("One of db_url or db_client must be provided")
 
         # Client and database will be lazily initialized per event loop
-        self._client: Optional[AsyncIOMotorClient] = None
-        self._database: Optional[AsyncIOMotorDatabase] = None
+        self._client: Optional[AsyncMongoClientType] = None
+        self._database: Optional[AsyncMongoDatabaseType] = None
         self._event_loop: Optional[asyncio.AbstractEventLoop] = None
 
     async def table_exists(self, table_name: str) -> bool:
@@ -126,15 +252,27 @@ class AsyncMongoDb(AsyncBaseDb):
             if collection_name and not await self.table_exists(collection_name):
                 await self._get_collection(collection_type, create_collection_if_not_found=True)
 
-    def _ensure_client(self) -> AsyncIOMotorClient:
+    async def close(self) -> None:
+        """Close the MongoDB client connection.
+
+        Should be called during application shutdown to properly release
+        all database connections.
         """
-        Ensure the Motor client is valid for the current event loop.
+        if self._client is not None:
+            self._client.close()
+            self._client = None
+            self._database = None
 
-        Motor's AsyncIOMotorClient is tied to the event loop it was created in.
-        If we detect a new event loop, we need to refresh the client.
+    def _ensure_client(self) -> AsyncMongoClientType:
+        """
+        Ensure the MongoDB async client is valid for the current event loop.
+
+        Both Motor's AsyncIOMotorClient and PyMongo's AsyncMongoClient are tied to
+        the event loop they were created in. If we detect a new event loop, we need
+        to refresh the client.
 
         Returns:
-            AsyncIOMotorClient: A valid client for the current event loop.
+            Union[AsyncIOMotorClient, AsyncMongoClient]: A valid client for the current event loop.
         """
         try:
             current_loop = asyncio.get_running_loop()
@@ -144,8 +282,13 @@ class AsyncMongoDb(AsyncBaseDb):
                 if self._provided_client is not None:
                     self._client = self._provided_client
                 elif self.db_url is not None:
-                    self._client = AsyncIOMotorClient(self.db_url)
-                    log_debug("Created AsyncIOMotorClient outside event loop")
+                    # Create client based on detected type
+                    if self._client_type == self.CLIENT_TYPE_PYMONGO_ASYNC and PYMONGO_ASYNC_AVAILABLE:
+                        self._client = AsyncMongoClient(self.db_url)  # type: ignore
+                    elif self._client_type == self.CLIENT_TYPE_MOTOR and MOTOR_AVAILABLE:
+                        self._client = AsyncIOMotorClient(self.db_url)  # type: ignore
+                    else:
+                        raise RuntimeError(f"Client type '{self._client_type}' not available")
             return self._client  # type: ignore
 
         # Check if we're in a different event loop
@@ -153,43 +296,47 @@ class AsyncMongoDb(AsyncBaseDb):
             # New event loop detected, create new client
             if self._provided_client is not None:
                 # User provided a client, use it but warn them
+                client_type_name = (
+                    "AsyncMongoClient" if self._client_type == self.CLIENT_TYPE_PYMONGO_ASYNC else "AsyncIOMotorClient"
+                )
                 log_debug(
-                    "New event loop detected. Using provided AsyncIOMotorClient, "
+                    f"New event loop detected. Using provided {client_type_name}, "
                     "which may cause issues if it was created in a different event loop."
                 )
                 self._client = self._provided_client
             elif self.db_url is not None:
-                # Create a new client for this event loop
-                old_loop_id = id(self._event_loop) if self._event_loop else "None"
-                new_loop_id = id(current_loop)
-                log_debug(f"Event loop changed from {old_loop_id} to {new_loop_id}, creating new AsyncIOMotorClient")
-                self._client = AsyncIOMotorClient(self.db_url)
+                if self._client_type == self.CLIENT_TYPE_PYMONGO_ASYNC and PYMONGO_ASYNC_AVAILABLE:
+                    self._client = AsyncMongoClient(self.db_url)  # type: ignore
+                elif self._client_type == self.CLIENT_TYPE_MOTOR and MOTOR_AVAILABLE:
+                    self._client = AsyncIOMotorClient(self.db_url)  # type: ignore
+                else:
+                    raise RuntimeError(f"Client type '{self._client_type}' not available")
 
             self._event_loop = current_loop
             self._database = None  # Reset database reference
-            # Clear collection caches when switching event loops
+            # Clear collection caches and initialization flags when switching event loops
             for attr in list(vars(self).keys()):
-                if attr.endswith("_collection"):
+                if attr.endswith("_collection") or attr.endswith("_initialized"):
                     delattr(self, attr)
 
         return self._client  # type: ignore
 
     @property
-    def db_client(self) -> AsyncIOMotorClient:
+    def db_client(self) -> AsyncMongoClientType:
         """Get the MongoDB client, ensuring it's valid for the current event loop."""
         return self._ensure_client()
 
     @property
-    def database(self) -> AsyncIOMotorDatabase:
+    def database(self) -> AsyncMongoDatabaseType:
         """Get the MongoDB database, ensuring it's valid for the current event loop."""
         try:
             current_loop = asyncio.get_running_loop()
             if self._database is None or self._event_loop != current_loop:
-                self._database = self.db_client[self.db_name]
+                self._database = self.db_client[self.db_name]  # type: ignore
         except RuntimeError:
             # No running loop - fallback to existing database or create new one
             if self._database is None:
-                self._database = self.db_client[self.db_name]
+                self._database = self.db_client[self.db_name]  # type: ignore
         return self._database
 
     # -- DB methods --
@@ -204,7 +351,7 @@ class AsyncMongoDb(AsyncBaseDb):
 
     async def _get_collection(
         self, table_type: str, create_collection_if_not_found: Optional[bool] = True
-    ) -> Optional[AsyncIOMotorCollection]:
+    ) -> Optional[AsyncMongoCollectionType]:
         """Get or create a collection based on table type.
 
         Args:
@@ -212,7 +359,7 @@ class AsyncMongoDb(AsyncBaseDb):
             create_collection_if_not_found (Optional[bool]): Whether to create the collection if it doesn't exist.
 
         Returns:
-            AsyncIOMotorCollection: The collection object.
+            Union[AsyncIOMotorCollection, AsyncCollection]: The collection object.
         """
         # Ensure client is valid for current event loop before accessing collections
         _ = self.db_client  # This triggers _ensure_client()
@@ -286,11 +433,44 @@ class AsyncMongoDb(AsyncBaseDb):
                 )
             return self.culture_collection
 
+        if table_type == "traces":
+            if reset_cache or not hasattr(self, "traces_collection"):
+                if self.trace_table_name is None:
+                    raise ValueError("Traces collection was not provided on initialization")
+                self.traces_collection = await self._get_or_create_collection(
+                    collection_name=self.trace_table_name,
+                    collection_type="traces",
+                    create_collection_if_not_found=create_collection_if_not_found,
+                )
+            return self.traces_collection
+
+        if table_type == "spans":
+            if reset_cache or not hasattr(self, "spans_collection"):
+                if self.span_table_name is None:
+                    raise ValueError("Spans collection was not provided on initialization")
+                self.spans_collection = await self._get_or_create_collection(
+                    collection_name=self.span_table_name,
+                    collection_type="spans",
+                    create_collection_if_not_found=create_collection_if_not_found,
+                )
+            return self.spans_collection
+
+        if table_type == "learnings":
+            if reset_cache or not hasattr(self, "learnings_collection"):
+                if self.learnings_table_name is None:
+                    raise ValueError("Learnings collection was not provided on initialization")
+                self.learnings_collection = await self._get_or_create_collection(
+                    collection_name=self.learnings_table_name,
+                    collection_type="learnings",
+                    create_collection_if_not_found=create_collection_if_not_found,
+                )
+            return self.learnings_collection
+
         raise ValueError(f"Unknown table type: {table_type}")
 
     async def _get_or_create_collection(
         self, collection_name: str, collection_type: str, create_collection_if_not_found: Optional[bool] = True
-    ) -> Optional[AsyncIOMotorCollection]:
+    ) -> Optional[AsyncMongoCollectionType]:
         """Get or create a collection with proper indexes.
 
         Args:
@@ -299,7 +479,7 @@ class AsyncMongoDb(AsyncBaseDb):
             create_collection_if_not_found (Optional[bool]): Whether to create the collection if it doesn't exist.
 
         Returns:
-            Optional[AsyncIOMotorCollection]: The collection object.
+            Union[AsyncIOMotorCollection, AsyncCollection]: The collection object.
         """
         try:
             collection = self.database[collection_name]
@@ -307,9 +487,8 @@ class AsyncMongoDb(AsyncBaseDb):
             if not hasattr(self, f"_{collection_name}_initialized"):
                 if not create_collection_if_not_found:
                     return None
-                # Note: Motor doesn't have sync create_index, so we use it as-is
-                # The indexes are created in the background
-                create_collection_indexes(collection, collection_type)  # type: ignore
+                # Create indexes asynchronously for async MongoDB collections
+                await create_collection_indexes_async(collection, collection_type)
                 setattr(self, f"_{collection_name}_initialized", True)
                 log_debug(f"Initialized collection '{collection_name}'")
             else:
@@ -321,6 +500,14 @@ class AsyncMongoDb(AsyncBaseDb):
             log_error(f"Error getting collection {collection_name}: {e}")
             raise
 
+    def get_latest_schema_version(self):
+        """Get the latest version of the database schema."""
+        pass
+
+    def upsert_schema_version(self, version: str) -> None:
+        """Upsert the schema version into the database."""
+        pass
+
     # -- Session methods --
 
     async def delete_session(self, session_id: str) -> bool:
@@ -1241,6 +1428,9 @@ class AsyncMongoDb(AsyncBaseDb):
                     "memory_id": memory.memory_id,
                     "memory": memory.memory,
                     "topics": memory.topics,
+                    "input": memory.input,
+                    "feedback": memory.feedback,
+                    "created_at": memory.created_at,
                     "updated_at": updated_at,
                 }
 
@@ -1533,7 +1723,7 @@ class AsyncMongoDb(AsyncBaseDb):
             log_error(f"Exception reading from sessions collection: {e}")
             return []
 
-    async def _get_metrics_calculation_starting_date(self, collection: AsyncIOMotorCollection) -> Optional[date]:
+    async def _get_metrics_calculation_starting_date(self, collection: AsyncMongoCollectionType) -> Optional[date]:
         """Get the first date for which metrics calculation is needed."""
         try:
             result = await collection.find_one({}, sort=[("date", -1)], limit=1)
@@ -2024,3 +2214,787 @@ class AsyncMongoDb(AsyncBaseDb):
         except Exception as e:
             log_error(f"Error updating eval run name {eval_run_id}: {e}")
             raise e
+
+    # --- Traces ---
+    def _get_component_level(
+        self, workflow_id: Optional[str], team_id: Optional[str], agent_id: Optional[str], name: str
+    ) -> int:
+        """Get the component level for a trace based on its context.
+
+        Component levels (higher = more important):
+            - 3: Workflow root (.run or .arun with workflow_id)
+            - 2: Team root (.run or .arun with team_id)
+            - 1: Agent root (.run or .arun with agent_id)
+            - 0: Child span (not a root)
+
+        Args:
+            workflow_id: The workflow ID of the trace.
+            team_id: The team ID of the trace.
+            agent_id: The agent ID of the trace.
+            name: The name of the trace.
+
+        Returns:
+            int: The component level (0-3).
+        """
+        # Check if name indicates a root span
+        is_root_name = ".run" in name or ".arun" in name
+
+        if not is_root_name:
+            return 0  # Child span (not a root)
+        elif workflow_id:
+            return 3  # Workflow root
+        elif team_id:
+            return 2  # Team root
+        elif agent_id:
+            return 1  # Agent root
+        else:
+            return 0  # Unknown
+
+    async def upsert_trace(self, trace: "Trace") -> None:
+        """Create or update a single trace record in the database.
+
+        Uses MongoDB's update_one with upsert=True and aggregation pipeline
+        to handle concurrent inserts atomically and avoid race conditions.
+
+        Args:
+            trace: The Trace object to store (one per trace_id).
+        """
+        try:
+            collection = await self._get_collection(table_type="traces", create_collection_if_not_found=True)
+            if collection is None:
+                return
+
+            trace_dict = trace.to_dict()
+            trace_dict.pop("total_spans", None)
+            trace_dict.pop("error_count", None)
+
+            # Calculate the component level for the new trace
+            new_level = self._get_component_level(trace.workflow_id, trace.team_id, trace.agent_id, trace.name)
+
+            # Use MongoDB aggregation pipeline update for atomic upsert
+            # This allows conditional logic within a single atomic operation
+            pipeline: List[Dict[str, Any]] = [
+                {
+                    "$set": {
+                        # Always update these fields
+                        "status": trace.status,
+                        "created_at": {"$ifNull": ["$created_at", trace_dict.get("created_at")]},
+                        # Use $min for start_time (keep earliest)
+                        "start_time": {
+                            "$cond": {
+                                "if": {"$eq": [{"$type": "$start_time"}, "missing"]},
+                                "then": trace_dict.get("start_time"),
+                                "else": {"$min": ["$start_time", trace_dict.get("start_time")]},
+                            }
+                        },
+                        # Use $max for end_time (keep latest)
+                        "end_time": {
+                            "$cond": {
+                                "if": {"$eq": [{"$type": "$end_time"}, "missing"]},
+                                "then": trace_dict.get("end_time"),
+                                "else": {"$max": ["$end_time", trace_dict.get("end_time")]},
+                            }
+                        },
+                        # Preserve existing non-null context values using $ifNull
+                        "run_id": {"$ifNull": [trace.run_id, "$run_id"]},
+                        "session_id": {"$ifNull": [trace.session_id, "$session_id"]},
+                        "user_id": {"$ifNull": [trace.user_id, "$user_id"]},
+                        "agent_id": {"$ifNull": [trace.agent_id, "$agent_id"]},
+                        "team_id": {"$ifNull": [trace.team_id, "$team_id"]},
+                        "workflow_id": {"$ifNull": [trace.workflow_id, "$workflow_id"]},
+                    }
+                },
+                {
+                    "$set": {
+                        # Calculate duration_ms from the (potentially updated) start_time and end_time
+                        # MongoDB stores dates as strings in ISO format, so we need to parse them
+                        "duration_ms": {
+                            "$cond": {
+                                "if": {
+                                    "$and": [
+                                        {"$ne": [{"$type": "$start_time"}, "missing"]},
+                                        {"$ne": [{"$type": "$end_time"}, "missing"]},
+                                    ]
+                                },
+                                "then": {
+                                    "$subtract": [
+                                        {"$toLong": {"$toDate": "$end_time"}},
+                                        {"$toLong": {"$toDate": "$start_time"}},
+                                    ]
+                                },
+                                "else": trace_dict.get("duration_ms", 0),
+                            }
+                        },
+                        # Update name based on component level priority
+                        # Only update if new trace is from a higher-level component
+                        "name": {
+                            "$cond": {
+                                "if": {"$eq": [{"$type": "$name"}, "missing"]},
+                                "then": trace.name,
+                                "else": {
+                                    "$cond": {
+                                        "if": {
+                                            "$gt": [
+                                                new_level,
+                                                {
+                                                    "$switch": {
+                                                        "branches": [
+                                                            # Check if existing name is a root span
+                                                            {
+                                                                "case": {
+                                                                    "$not": {
+                                                                        "$or": [
+                                                                            {
+                                                                                "$regexMatch": {
+                                                                                    "input": {"$ifNull": ["$name", ""]},
+                                                                                    "regex": "\\.run",
+                                                                                }
+                                                                            },
+                                                                            {
+                                                                                "$regexMatch": {
+                                                                                    "input": {"$ifNull": ["$name", ""]},
+                                                                                    "regex": "\\.arun",
+                                                                                }
+                                                                            },
+                                                                        ]
+                                                                    }
+                                                                },
+                                                                "then": 0,
+                                                            },
+                                                            # Workflow root (level 3)
+                                                            {
+                                                                "case": {"$ne": ["$workflow_id", None]},
+                                                                "then": 3,
+                                                            },
+                                                            # Team root (level 2)
+                                                            {
+                                                                "case": {"$ne": ["$team_id", None]},
+                                                                "then": 2,
+                                                            },
+                                                            # Agent root (level 1)
+                                                            {
+                                                                "case": {"$ne": ["$agent_id", None]},
+                                                                "then": 1,
+                                                            },
+                                                        ],
+                                                        "default": 0,
+                                                    }
+                                                },
+                                            ]
+                                        },
+                                        "then": trace.name,
+                                        "else": "$name",
+                                    }
+                                },
+                            }
+                        },
+                    }
+                },
+            ]
+
+            # Perform atomic upsert using aggregation pipeline
+            await collection.update_one(
+                {"trace_id": trace.trace_id},
+                pipeline,
+                upsert=True,
+            )
+
+        except Exception as e:
+            log_error(f"Error creating trace: {e}")
+            # Don't raise - tracing should not break the main application flow
+
+    async def get_trace(
+        self,
+        trace_id: Optional[str] = None,
+        run_id: Optional[str] = None,
+    ):
+        """Get a single trace by trace_id or other filters.
+
+        Args:
+            trace_id: The unique trace identifier.
+            run_id: Filter by run ID (returns first match).
+
+        Returns:
+            Optional[Trace]: The trace if found, None otherwise.
+
+        Note:
+            If multiple filters are provided, trace_id takes precedence.
+            For other filters, the most recent trace is returned.
+        """
+        try:
+            from agno.tracing.schemas import Trace as TraceSchema
+
+            collection = await self._get_collection(table_type="traces")
+            if collection is None:
+                return None
+
+            # Get spans collection for aggregation
+            spans_collection = await self._get_collection(table_type="spans")
+
+            query: Dict[str, Any] = {}
+            if trace_id:
+                query["trace_id"] = trace_id
+            elif run_id:
+                query["run_id"] = run_id
+            else:
+                log_debug("get_trace called without any filter parameters")
+                return None
+
+            # Find trace with sorting by most recent
+            result = await collection.find_one(query, sort=[("start_time", -1)])
+
+            if result:
+                # Calculate total_spans and error_count from spans collection
+                total_spans = 0
+                error_count = 0
+                if spans_collection is not None:
+                    total_spans = await spans_collection.count_documents({"trace_id": result["trace_id"]})
+                    error_count = await spans_collection.count_documents(
+                        {"trace_id": result["trace_id"], "status_code": "ERROR"}
+                    )
+
+                result["total_spans"] = total_spans
+                result["error_count"] = error_count
+                # Remove MongoDB's _id field
+                result.pop("_id", None)
+                return TraceSchema.from_dict(result)
+            return None
+
+        except Exception as e:
+            log_error(f"Error getting trace: {e}")
+            return None
+
+    async def get_traces(
+        self,
+        run_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        user_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+        team_id: Optional[str] = None,
+        workflow_id: Optional[str] = None,
+        status: Optional[str] = None,
+        start_time: Optional[datetime] = None,
+        end_time: Optional[datetime] = None,
+        limit: Optional[int] = 20,
+        page: Optional[int] = 1,
+    ) -> tuple[List, int]:
+        """Get traces matching the provided filters with pagination.
+
+        Args:
+            run_id: Filter by run ID.
+            session_id: Filter by session ID.
+            user_id: Filter by user ID.
+            agent_id: Filter by agent ID.
+            team_id: Filter by team ID.
+            workflow_id: Filter by workflow ID.
+            status: Filter by status (OK, ERROR, UNSET).
+            start_time: Filter traces starting after this datetime.
+            end_time: Filter traces ending before this datetime.
+            limit: Maximum number of traces to return per page.
+            page: Page number (1-indexed).
+
+        Returns:
+            tuple[List[Trace], int]: Tuple of (list of matching traces, total count).
+        """
+        try:
+            from agno.tracing.schemas import Trace as TraceSchema
+
+            collection = await self._get_collection(table_type="traces")
+            if collection is None:
+                log_debug("Traces collection not found")
+                return [], 0
+
+            # Get spans collection for aggregation
+            spans_collection = await self._get_collection(table_type="spans")
+
+            # Build query
+            query: Dict[str, Any] = {}
+            if run_id:
+                query["run_id"] = run_id
+            if session_id:
+                query["session_id"] = session_id
+            if user_id:
+                query["user_id"] = user_id
+            if agent_id:
+                query["agent_id"] = agent_id
+            if team_id:
+                query["team_id"] = team_id
+            if workflow_id:
+                query["workflow_id"] = workflow_id
+            if status:
+                query["status"] = status
+            if start_time:
+                query["start_time"] = {"$gte": start_time.isoformat()}
+            if end_time:
+                if "end_time" in query:
+                    query["end_time"]["$lte"] = end_time.isoformat()
+                else:
+                    query["end_time"] = {"$lte": end_time.isoformat()}
+
+            # Get total count
+            total_count = await collection.count_documents(query)
+
+            # Apply pagination
+            skip = ((page or 1) - 1) * (limit or 20)
+            cursor = collection.find(query).sort("start_time", -1).skip(skip).limit(limit or 20)
+
+            results = await cursor.to_list(length=None)
+
+            traces = []
+            for row in results:
+                # Calculate total_spans and error_count from spans collection
+                total_spans = 0
+                error_count = 0
+                if spans_collection is not None:
+                    total_spans = await spans_collection.count_documents({"trace_id": row["trace_id"]})
+                    error_count = await spans_collection.count_documents(
+                        {"trace_id": row["trace_id"], "status_code": "ERROR"}
+                    )
+
+                row["total_spans"] = total_spans
+                row["error_count"] = error_count
+                # Remove MongoDB's _id field
+                row.pop("_id", None)
+                traces.append(TraceSchema.from_dict(row))
+
+            return traces, total_count
+
+        except Exception as e:
+            log_error(f"Error getting traces: {e}")
+            return [], 0
+
+    async def get_trace_stats(
+        self,
+        user_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+        team_id: Optional[str] = None,
+        workflow_id: Optional[str] = None,
+        start_time: Optional[datetime] = None,
+        end_time: Optional[datetime] = None,
+        limit: Optional[int] = 20,
+        page: Optional[int] = 1,
+    ) -> tuple[List[Dict[str, Any]], int]:
+        """Get trace statistics grouped by session.
+
+        Args:
+            user_id: Filter by user ID.
+            agent_id: Filter by agent ID.
+            team_id: Filter by team ID.
+            workflow_id: Filter by workflow ID.
+            start_time: Filter sessions with traces created after this datetime.
+            end_time: Filter sessions with traces created before this datetime.
+            limit: Maximum number of sessions to return per page.
+            page: Page number (1-indexed).
+
+        Returns:
+            tuple[List[Dict], int]: Tuple of (list of session stats dicts, total count).
+                Each dict contains: session_id, user_id, agent_id, team_id, total_traces,
+                workflow_id, first_trace_at, last_trace_at.
+        """
+        try:
+            collection = await self._get_collection(table_type="traces")
+            if collection is None:
+                log_debug("Traces collection not found")
+                return [], 0
+
+            # Build match stage
+            match_stage: Dict[str, Any] = {"session_id": {"$ne": None}}
+            if user_id:
+                match_stage["user_id"] = user_id
+            if agent_id:
+                match_stage["agent_id"] = agent_id
+            if team_id:
+                match_stage["team_id"] = team_id
+            if workflow_id:
+                match_stage["workflow_id"] = workflow_id
+            if start_time:
+                match_stage["created_at"] = {"$gte": start_time.isoformat()}
+            if end_time:
+                if "created_at" in match_stage:
+                    match_stage["created_at"]["$lte"] = end_time.isoformat()
+                else:
+                    match_stage["created_at"] = {"$lte": end_time.isoformat()}
+
+            # Build aggregation pipeline
+            pipeline: List[Dict[str, Any]] = [
+                {"$match": match_stage},
+                {
+                    "$group": {
+                        "_id": "$session_id",
+                        "user_id": {"$first": "$user_id"},
+                        "agent_id": {"$first": "$agent_id"},
+                        "team_id": {"$first": "$team_id"},
+                        "workflow_id": {"$first": "$workflow_id"},
+                        "total_traces": {"$sum": 1},
+                        "first_trace_at": {"$min": "$created_at"},
+                        "last_trace_at": {"$max": "$created_at"},
+                    }
+                },
+                {"$sort": {"last_trace_at": -1}},
+            ]
+
+            # Get total count
+            count_pipeline = pipeline + [{"$count": "total"}]
+            count_result = await collection.aggregate(count_pipeline).to_list(length=1)
+            total_count = count_result[0]["total"] if count_result else 0
+
+            # Apply pagination
+            skip = ((page or 1) - 1) * (limit or 20)
+            pipeline.append({"$skip": skip})
+            pipeline.append({"$limit": limit or 20})
+
+            results = await collection.aggregate(pipeline).to_list(length=None)
+
+            # Convert to list of dicts with datetime objects
+            stats_list = []
+            for row in results:
+                # Convert ISO strings to datetime objects
+                first_trace_at_str = row["first_trace_at"]
+                last_trace_at_str = row["last_trace_at"]
+
+                # Parse ISO format strings to datetime objects
+                first_trace_at = datetime.fromisoformat(first_trace_at_str.replace("Z", "+00:00"))
+                last_trace_at = datetime.fromisoformat(last_trace_at_str.replace("Z", "+00:00"))
+
+                stats_list.append(
+                    {
+                        "session_id": row["_id"],
+                        "user_id": row["user_id"],
+                        "agent_id": row["agent_id"],
+                        "team_id": row["team_id"],
+                        "workflow_id": row["workflow_id"],
+                        "total_traces": row["total_traces"],
+                        "first_trace_at": first_trace_at,
+                        "last_trace_at": last_trace_at,
+                    }
+                )
+
+            return stats_list, total_count
+
+        except Exception as e:
+            log_error(f"Error getting trace stats: {e}")
+            return [], 0
+
+    # --- Spans ---
+    async def create_span(self, span: "Span") -> None:
+        """Create a single span in the database.
+
+        Args:
+            span: The Span object to store.
+        """
+        try:
+            collection = await self._get_collection(table_type="spans", create_collection_if_not_found=True)
+            if collection is None:
+                return
+
+            await collection.insert_one(span.to_dict())
+
+        except Exception as e:
+            log_error(f"Error creating span: {e}")
+
+    async def create_spans(self, spans: List) -> None:
+        """Create multiple spans in the database as a batch.
+
+        Args:
+            spans: List of Span objects to store.
+        """
+        if not spans:
+            return
+
+        try:
+            collection = await self._get_collection(table_type="spans", create_collection_if_not_found=True)
+            if collection is None:
+                return
+
+            span_dicts = [span.to_dict() for span in spans]
+            await collection.insert_many(span_dicts)
+
+        except Exception as e:
+            log_error(f"Error creating spans batch: {e}")
+
+    async def get_span(self, span_id: str):
+        """Get a single span by its span_id.
+
+        Args:
+            span_id: The unique span identifier.
+
+        Returns:
+            Optional[Span]: The span if found, None otherwise.
+        """
+        try:
+            from agno.tracing.schemas import Span as SpanSchema
+
+            collection = await self._get_collection(table_type="spans")
+            if collection is None:
+                return None
+
+            result = await collection.find_one({"span_id": span_id})
+            if result:
+                # Remove MongoDB's _id field
+                result.pop("_id", None)
+                return SpanSchema.from_dict(result)
+            return None
+
+        except Exception as e:
+            log_error(f"Error getting span: {e}")
+            return None
+
+    async def get_spans(
+        self,
+        trace_id: Optional[str] = None,
+        parent_span_id: Optional[str] = None,
+        limit: Optional[int] = 1000,
+    ) -> List:
+        """Get spans matching the provided filters.
+
+        Args:
+            trace_id: Filter by trace ID.
+            parent_span_id: Filter by parent span ID.
+            limit: Maximum number of spans to return.
+
+        Returns:
+            List[Span]: List of matching spans.
+        """
+        try:
+            from agno.tracing.schemas import Span as SpanSchema
+
+            collection = await self._get_collection(table_type="spans")
+            if collection is None:
+                return []
+
+            # Build query
+            query: Dict[str, Any] = {}
+            if trace_id:
+                query["trace_id"] = trace_id
+            if parent_span_id:
+                query["parent_span_id"] = parent_span_id
+
+            cursor = collection.find(query).limit(limit or 1000)
+            results = await cursor.to_list(length=None)
+
+            spans = []
+            for row in results:
+                # Remove MongoDB's _id field
+                row.pop("_id", None)
+                spans.append(SpanSchema.from_dict(row))
+
+            return spans
+
+        except Exception as e:
+            log_error(f"Error getting spans: {e}")
+            return []
+
+    # -- Learning methods --
+    async def get_learning(
+        self,
+        learning_type: str,
+        user_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+        team_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        namespace: Optional[str] = None,
+        entity_id: Optional[str] = None,
+        entity_type: Optional[str] = None,
+    ) -> Optional[Dict[str, Any]]:
+        """Retrieve a learning record.
+
+        Args:
+            learning_type: Type of learning ('user_profile', 'session_context', etc.)
+            user_id: Filter by user ID.
+            agent_id: Filter by agent ID.
+            team_id: Filter by team ID.
+            session_id: Filter by session ID.
+            namespace: Filter by namespace ('user', 'global', or custom).
+            entity_id: Filter by entity ID (for entity-specific learnings).
+            entity_type: Filter by entity type ('person', 'company', etc.).
+
+        Returns:
+            Dict with 'content' key containing the learning data, or None.
+        """
+        try:
+            collection = await self._get_collection(table_type="learnings", create_collection_if_not_found=False)
+            if collection is None:
+                return None
+
+            # Build query
+            query: Dict[str, Any] = {"learning_type": learning_type}
+            if user_id is not None:
+                query["user_id"] = user_id
+            if agent_id is not None:
+                query["agent_id"] = agent_id
+            if team_id is not None:
+                query["team_id"] = team_id
+            if session_id is not None:
+                query["session_id"] = session_id
+            if namespace is not None:
+                query["namespace"] = namespace
+            if entity_id is not None:
+                query["entity_id"] = entity_id
+            if entity_type is not None:
+                query["entity_type"] = entity_type
+
+            result = await collection.find_one(query)
+            if result is None:
+                return None
+
+            # Remove MongoDB's _id field
+            result.pop("_id", None)
+            return {"content": result.get("content")}
+
+        except Exception as e:
+            log_debug(f"Error retrieving learning: {e}")
+            return None
+
+    async def upsert_learning(
+        self,
+        id: str,
+        learning_type: str,
+        content: Dict[str, Any],
+        user_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+        team_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        namespace: Optional[str] = None,
+        entity_id: Optional[str] = None,
+        entity_type: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """Insert or update a learning record.
+
+        Args:
+            id: Unique identifier for the learning.
+            learning_type: Type of learning ('user_profile', 'session_context', etc.)
+            content: The learning content as a dict.
+            user_id: Associated user ID.
+            agent_id: Associated agent ID.
+            team_id: Associated team ID.
+            session_id: Associated session ID.
+            namespace: Namespace for scoping ('user', 'global', or custom).
+            entity_id: Associated entity ID (for entity-specific learnings).
+            entity_type: Entity type ('person', 'company', etc.).
+            metadata: Optional metadata.
+        """
+        try:
+            collection = await self._get_collection(table_type="learnings", create_collection_if_not_found=True)
+            if collection is None:
+                return
+
+            current_time = int(time.time())
+
+            document = {
+                "learning_id": id,
+                "learning_type": learning_type,
+                "namespace": namespace,
+                "user_id": user_id,
+                "agent_id": agent_id,
+                "team_id": team_id,
+                "session_id": session_id,
+                "entity_id": entity_id,
+                "entity_type": entity_type,
+                "content": content,
+                "metadata": metadata,
+                "updated_at": current_time,
+            }
+
+            # Use upsert to insert or update
+            await collection.update_one(
+                {"learning_id": id},
+                {"$set": document, "$setOnInsert": {"created_at": current_time}},
+                upsert=True,
+            )
+
+            log_debug(f"Upserted learning: {id}")
+
+        except Exception as e:
+            log_debug(f"Error upserting learning: {e}")
+
+    async def delete_learning(self, id: str) -> bool:
+        """Delete a learning record.
+
+        Args:
+            id: The learning ID to delete.
+
+        Returns:
+            True if deleted, False otherwise.
+        """
+        try:
+            collection = await self._get_collection(table_type="learnings", create_collection_if_not_found=False)
+            if collection is None:
+                return False
+
+            result = await collection.delete_one({"learning_id": id})
+            return result.deleted_count > 0
+
+        except Exception as e:
+            log_debug(f"Error deleting learning: {e}")
+            return False
+
+    async def get_learnings(
+        self,
+        learning_type: Optional[str] = None,
+        user_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+        team_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        namespace: Optional[str] = None,
+        entity_id: Optional[str] = None,
+        entity_type: Optional[str] = None,
+        limit: Optional[int] = None,
+    ) -> List[Dict[str, Any]]:
+        """Get multiple learning records.
+
+        Args:
+            learning_type: Filter by learning type.
+            user_id: Filter by user ID.
+            agent_id: Filter by agent ID.
+            team_id: Filter by team ID.
+            session_id: Filter by session ID.
+            namespace: Filter by namespace ('user', 'global', or custom).
+            entity_id: Filter by entity ID (for entity-specific learnings).
+            entity_type: Filter by entity type ('person', 'company', etc.).
+            limit: Maximum number of records to return.
+
+        Returns:
+            List of learning records.
+        """
+        try:
+            collection = await self._get_collection(table_type="learnings", create_collection_if_not_found=False)
+            if collection is None:
+                return []
+
+            # Build query
+            query: Dict[str, Any] = {}
+            if learning_type is not None:
+                query["learning_type"] = learning_type
+            if user_id is not None:
+                query["user_id"] = user_id
+            if agent_id is not None:
+                query["agent_id"] = agent_id
+            if team_id is not None:
+                query["team_id"] = team_id
+            if session_id is not None:
+                query["session_id"] = session_id
+            if namespace is not None:
+                query["namespace"] = namespace
+            if entity_id is not None:
+                query["entity_id"] = entity_id
+            if entity_type is not None:
+                query["entity_type"] = entity_type
+
+            cursor = collection.find(query)
+            if limit is not None:
+                cursor = cursor.limit(limit)
+
+            results = await cursor.to_list(length=None)
+
+            learnings = []
+            for row in results:
+                # Remove MongoDB's _id field
+                row.pop("_id", None)
+                learnings.append(row)
+
+            return learnings
+
+        except Exception as e:
+            log_debug(f"Error getting learnings: {e}")
+            return []