agno 2.3.10__py3-none-any.whl → 2.3.11__py3-none-any.whl

agno/db/base.py

@@ -308,8 +308,8 @@ class BaseDb(ABC):
 
     # --- Traces ---
     @abstractmethod
-    def create_trace(self, trace: "Trace") -> None:
-        """Create a single trace record in the database.
+    def upsert_trace(self, trace: "Trace") -> None:
+        """Create or update a single trace record in the database.
 
         Args:
             trace: The Trace object to store (one per trace_id).
@@ -759,11 +759,11 @@ class AsyncBaseDb(ABC):
 
     # --- Traces ---
     @abstractmethod
-    async def create_trace(self, trace) -> None:
-        """Create a single trace record in the database.
+    async def upsert_trace(self, trace) -> None:
+        """Create or update a single trace record in the database.
 
         Args:
-            trace: The Trace object to store (one per trace_id).
+            trace: The Trace object to update (one per trace_id).
         """
         raise NotImplementedError
 

agno/db/dynamo/dynamo.py

@@ -2076,8 +2076,8 @@ class DynamoDb(BaseDb):
             raise e
 
     # --- Traces ---
-    def create_trace(self, trace: "Trace") -> None:
-        """Create a single trace record in the database.
+    def upsert_trace(self, trace: "Trace") -> None:
+        """Create or update a single trace record in the database.
 
         Args:
             trace: The Trace object to store (one per trace_id).

agno/db/firestore/firestore.py

@@ -1841,8 +1841,8 @@ class FirestoreDb(BaseDb):
             raise e
 
     # --- Traces ---
-    def create_trace(self, trace: "Trace") -> None:
-        """Create a single trace record in the database.
+    def upsert_trace(self, trace: "Trace") -> None:
+        """Create or update a single trace record in the database.
 
         Args:
             trace: The Trace object to store (one per trace_id).

agno/db/gcs_json/gcs_json_db.py

@@ -1355,8 +1355,8 @@ class GcsJsonDb(BaseDb):
             raise e
 
     # --- Traces ---
-    def create_trace(self, trace: "Trace") -> None:
-        """Create a single trace record in the database.
+    def upsert_trace(self, trace: "Trace") -> None:
+        """Create or update a single trace record in the database.
 
         Args:
             trace: The Trace object to store (one per trace_id).

agno/db/in_memory/in_memory_db.py

@@ -1173,8 +1173,8 @@ class InMemoryDb(BaseDb):
             raise e
 
     # --- Traces ---
-    def create_trace(self, trace: "Trace") -> None:
-        """Create a single trace record in the database.
+    def upsert_trace(self, trace: "Trace") -> None:
+        """Create or update a single trace record in the database.
 
         Args:
             trace: The Trace object to store (one per trace_id).

agno/db/json/json_db.py

@@ -1349,8 +1349,8 @@ class JsonDb(BaseDb):
             raise e
 
     # --- Traces ---
-    def create_trace(self, trace: "Trace") -> None:
-        """Create a single trace record in the database.
+    def upsert_trace(self, trace: "Trace") -> None:
+        """Create or update a single trace record in the database.
 
         Args:
             trace: The Trace object to store (one per trace_id).

agno/db/mongo/async_mongo.py

@@ -2191,8 +2191,45 @@ class AsyncMongoDb(AsyncBaseDb):
             raise e
 
     # --- Traces ---
-    async def create_trace(self, trace: "Trace") -> None:
-        """Create a single trace record in the database.
+    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).
@@ -2202,75 +2239,140 @@ class AsyncMongoDb(AsyncBaseDb):
             if collection is None:
                 return
 
-            # Check if trace already exists
-            existing = await collection.find_one({"trace_id": trace.trace_id})
-
-            if existing:
-                # workflow (level 3) > team (level 2) > agent (level 1) > child/unknown (level 0)
-                def get_component_level(
-                    workflow_id: Optional[str], team_id: Optional[str], agent_id: Optional[str], name: str
-                ) -> int:
-                    # 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
+            trace_dict = trace.to_dict()
+            trace_dict.pop("total_spans", None)
+            trace_dict.pop("error_count", None)
 
-                existing_level = get_component_level(
-                    existing.get("workflow_id"),
-                    existing.get("team_id"),
-                    existing.get("agent_id"),
-                    existing.get("name", ""),
-                )
-                new_level = get_component_level(trace.workflow_id, trace.team_id, trace.agent_id, trace.name)
-
-                # Only update name if new trace is from a higher or equal level
-                should_update_name = new_level > existing_level
-
-                # Parse existing start_time to calculate correct duration
-                existing_start_time_str = existing.get("start_time")
-                if isinstance(existing_start_time_str, str):
-                    existing_start_time = datetime.fromisoformat(existing_start_time_str.replace("Z", "+00:00"))
-                else:
-                    existing_start_time = trace.start_time
+            # 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)
 
-                recalculated_duration_ms = int((trace.end_time - existing_start_time).total_seconds() * 1000)
-
-                update_values: Dict[str, Any] = {
-                    "end_time": trace.end_time.isoformat(),
-                    "duration_ms": recalculated_duration_ms,
-                    "status": trace.status,
-                    "name": trace.name if should_update_name else existing.get("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",
+                                    }
+                                },
+                            }
+                        },
+                    }
+                },
+            ]
 
-                # Update context fields ONLY if new value is not None (preserve non-null values)
-                if trace.run_id is not None:
-                    update_values["run_id"] = trace.run_id
-                if trace.session_id is not None:
-                    update_values["session_id"] = trace.session_id
-                if trace.user_id is not None:
-                    update_values["user_id"] = trace.user_id
-                if trace.agent_id is not None:
-                    update_values["agent_id"] = trace.agent_id
-                if trace.team_id is not None:
-                    update_values["team_id"] = trace.team_id
-                if trace.workflow_id is not None:
-                    update_values["workflow_id"] = trace.workflow_id
-
-                await collection.update_one({"trace_id": trace.trace_id}, {"$set": update_values})
-            else:
-                trace_dict = trace.to_dict()
-                trace_dict.pop("total_spans", None)
-                trace_dict.pop("error_count", None)
-                await collection.insert_one(trace_dict)
+            # 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}")
@@ -2655,4 +2757,4 @@ class AsyncMongoDb(AsyncBaseDb):
 
         except Exception as e:
             log_error(f"Error getting spans: {e}")
-            return []
+            return []

agno/db/mongo/mongo.py

@@ -2028,8 +2028,45 @@ class MongoDb(BaseDb):
             log_info(f"Migrated {len(memories)} memories to collection: {self.memory_table_name}")
 
     # --- Traces ---
-    def create_trace(self, trace: "Trace") -> None:
-        """Create a single trace record in the database.
+    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
+
+    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).
@@ -2039,83 +2076,140 @@ class MongoDb(BaseDb):
             if collection is None:
                 return
 
-            # Check if trace already exists
-            existing = collection.find_one({"trace_id": trace.trace_id})
-
-            if existing:
-                # workflow (level 3) > team (level 2) > agent (level 1) > child/unknown (level 0)
-                def get_component_level(
-                    workflow_id: Optional[str], team_id: Optional[str], agent_id: Optional[str], name: str
-                ) -> int:
-                    # 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
-
-                existing_level = get_component_level(
-                    existing.get("workflow_id"),
-                    existing.get("team_id"),
-                    existing.get("agent_id"),
-                    existing.get("name", ""),
-                )
-                new_level = get_component_level(trace.workflow_id, trace.team_id, trace.agent_id, trace.name)
-
-                # Only update name if new trace is from a higher or equal level
-                should_update_name = new_level > existing_level
-
-                # Parse existing start_time to calculate correct duration
-                existing_start_time_str = existing.get("start_time")
-                if isinstance(existing_start_time_str, str):
-                    existing_start_time = datetime.fromisoformat(existing_start_time_str.replace("Z", "+00:00"))
-                else:
-                    existing_start_time = trace.start_time
+            trace_dict = trace.to_dict()
+            trace_dict.pop("total_spans", None)
+            trace_dict.pop("error_count", None)
 
-                recalculated_duration_ms = int((trace.end_time - existing_start_time).total_seconds() * 1000)
+            # 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)
 
-                update_values: Dict[str, Any] = {
-                    "end_time": trace.end_time.isoformat(),
-                    "duration_ms": recalculated_duration_ms,
-                    "status": trace.status,
-                    "name": trace.name if should_update_name else existing.get("name"),
-                }
-
-                # Update context fields ONLY if new value is not None (preserve non-null values)
-                if trace.run_id is not None:
-                    update_values["run_id"] = trace.run_id
-                if trace.session_id is not None:
-                    update_values["session_id"] = trace.session_id
-                if trace.user_id is not None:
-                    update_values["user_id"] = trace.user_id
-                if trace.agent_id is not None:
-                    update_values["agent_id"] = trace.agent_id
-                if trace.team_id is not None:
-                    update_values["team_id"] = trace.team_id
-                if trace.workflow_id is not None:
-                    update_values["workflow_id"] = trace.workflow_id
-
-                log_debug(
-                    f"  Updating trace with context: run_id={update_values.get('run_id', 'unchanged')}, "
-                    f"session_id={update_values.get('session_id', 'unchanged')}, "
-                    f"user_id={update_values.get('user_id', 'unchanged')}, "
-                    f"agent_id={update_values.get('agent_id', 'unchanged')}, "
-                    f"team_id={update_values.get('team_id', 'unchanged')}, "
-                )
+            # 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",
+                                    }
+                                },
+                            }
+                        },
+                    }
+                },
+            ]
 
-                collection.update_one({"trace_id": trace.trace_id}, {"$set": update_values})
-            else:
-                trace_dict = trace.to_dict()
-                trace_dict.pop("total_spans", None)
-                trace_dict.pop("error_count", None)
-                collection.insert_one(trace_dict)
+            # Perform atomic upsert using aggregation pipeline
+            collection.update_one(
+                {"trace_id": trace.trace_id},
+                pipeline,
+                upsert=True,
+            )
 
         except Exception as e:
             log_error(f"Error creating trace: {e}")
@@ -2500,4 +2594,4 @@ class MongoDb(BaseDb):
 
         except Exception as e:
             log_error(f"Error getting spans: {e}")
-            return []
+            return []