agno 1.8.2__py3-none-any.whl → 2.0.0__py3-none-any.whl

agno/run/base.py

@@ -1,18 +1,19 @@
 from dataclasses import asdict, dataclass
 from enum import Enum
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict
 
 from pydantic import BaseModel
 
-from agno.media import AudioArtifact, AudioResponse, ImageArtifact, VideoArtifact
+from agno.media import Audio, Image, Video
 from agno.models.message import Citations, Message, MessageReferences
+from agno.models.metrics import Metrics
 from agno.models.response import ToolExecution
 from agno.reasoning.step import ReasoningStep
 from agno.utils.log import log_error
 
 
 @dataclass
-class BaseRunResponseEvent:
+class BaseRunOutputEvent:
     def to_dict(self) -> Dict[str, Any]:
         _dict = {
             k: v
@@ -22,7 +23,7 @@ class BaseRunResponseEvent:
             not in [
                 "tools",
                 "tool",
-                "extra_data",
+                "metadata",
                 "image",
                 "images",
                 "videos",
@@ -30,13 +31,28 @@ class BaseRunResponseEvent:
                 "response_audio",
                 "citations",
                 "member_responses",
+                "reasoning_messages",
+                "reasoning_steps",
+                "references",
+                "additional_input",
+                "metrics",
             ]
         }
 
-        if hasattr(self, "extra_data") and self.extra_data is not None:
-            _dict["extra_data"] = (
-                self.extra_data.to_dict() if isinstance(self.extra_data, RunResponseExtraData) else self.extra_data
-            )
+        if hasattr(self, "metadata") and self.metadata is not None:
+            _dict["metadata"] = self.metadata
+
+        if hasattr(self, "additional_input") and self.additional_input is not None:
+            _dict["additional_input"] = [m.to_dict() for m in self.additional_input]
+
+        if hasattr(self, "reasoning_messages") and self.reasoning_messages is not None:
+            _dict["reasoning_messages"] = [m.to_dict() for m in self.reasoning_messages]
+
+        if hasattr(self, "reasoning_steps") and self.reasoning_steps is not None:
+            _dict["reasoning_steps"] = [rs.model_dump() for rs in self.reasoning_steps]
+
+        if hasattr(self, "references") and self.references is not None:
+            _dict["references"] = [r.model_dump() for r in self.references]
 
         if hasattr(self, "member_responses") and self.member_responses:
             _dict["member_responses"] = [response.to_dict() for response in self.member_responses]
@@ -44,21 +60,15 @@ class BaseRunResponseEvent:
         if hasattr(self, "images") and self.images is not None:
             _dict["images"] = []
             for img in self.images:
-                if isinstance(img, ImageArtifact):
+                if isinstance(img, Image):
                     _dict["images"].append(img.to_dict())
                 else:
                     _dict["images"].append(img)
 
-        if hasattr(self, "image") and self.image is not None:
-            if isinstance(self.image, ImageArtifact):
-                _dict["image"] = self.image.to_dict()
-            else:
-                _dict["image"] = self.image
-
         if hasattr(self, "videos") and self.videos is not None:
             _dict["videos"] = []
             for vid in self.videos:
-                if isinstance(vid, VideoArtifact):
+                if isinstance(vid, Video):
                     _dict["videos"].append(vid.to_dict())
                 else:
                     _dict["videos"].append(vid)
@@ -66,13 +76,13 @@ class BaseRunResponseEvent:
         if hasattr(self, "audio") and self.audio is not None:
             _dict["audio"] = []
             for aud in self.audio:
-                if isinstance(aud, AudioArtifact):
+                if isinstance(aud, Audio):
                     _dict["audio"].append(aud.to_dict())
                 else:
                     _dict["audio"].append(aud)
 
         if hasattr(self, "response_audio") and self.response_audio is not None:
-            if isinstance(self.response_audio, AudioResponse):
+            if isinstance(self.response_audio, Audio):
                 _dict["response_audio"] = self.response_audio.to_dict()
             else:
                 _dict["response_audio"] = self.response_audio
@@ -100,6 +110,9 @@ class BaseRunResponseEvent:
             else:
                 _dict["tool"] = self.tool
 
+        if hasattr(self, "metrics") and self.metrics is not None:
+            _dict["metrics"] = self.metrics.to_dict()
+
         return _dict
 
     def to_json(self) -> str:
@@ -121,31 +134,39 @@ class BaseRunResponseEvent:
 
         images = data.pop("images", None)
         if images:
-            data["images"] = [ImageArtifact.model_validate(image) for image in images]
-
-        image = data.pop("image", None)
-        if image:
-            data["image"] = ImageArtifact.model_validate(image)
+            data["images"] = [Image.model_validate(image) for image in images]
 
         videos = data.pop("videos", None)
         if videos:
-            data["videos"] = [VideoArtifact.model_validate(video) for video in videos]
+            data["videos"] = [Video.model_validate(video) for video in videos]
 
         audio = data.pop("audio", None)
         if audio:
-            data["audio"] = [AudioArtifact.model_validate(audio) for audio in audio]
+            data["audio"] = [Audio.model_validate(audio) for audio in audio]
 
         response_audio = data.pop("response_audio", None)
         if response_audio:
-            data["response_audio"] = AudioResponse.model_validate(response_audio)
+            data["response_audio"] = Audio.model_validate(response_audio)
+
+        additional_input = data.pop("additional_input", None)
+        if additional_input is not None:
+            data["additional_input"] = [Message.model_validate(message) for message in additional_input]
+
+        reasoning_steps = data.pop("reasoning_steps", None)
+        if reasoning_steps is not None:
+            data["reasoning_steps"] = [ReasoningStep.model_validate(step) for step in reasoning_steps]
 
-        extra_data = data.pop("extra_data", None)
-        if extra_data:
-            data["extra_data"] = RunResponseExtraData.from_dict(extra_data)
+        reasoning_messages = data.pop("reasoning_messages", None)
+        if reasoning_messages is not None:
+            data["reasoning_messages"] = [Message.model_validate(message) for message in reasoning_messages]
+
+        references = data.pop("references", None)
+        if references is not None:
+            data["references"] = [MessageReferences.model_validate(reference) for reference in references]
 
-        # To make it backwards compatible
-        if "event" in data:
-            data.pop("event")
+        metrics = data.pop("metrics", None)
+        if metrics:
+            data["metrics"] = Metrics(**metrics)
 
         return cls(**data)
 
@@ -158,51 +179,6 @@ class BaseRunResponseEvent:
         return False
 
 
-@dataclass
-class RunResponseExtraData:
-    references: Optional[List[MessageReferences]] = None
-    add_messages: Optional[List[Message]] = None
-    reasoning_steps: Optional[List[ReasoningStep]] = None
-    reasoning_messages: Optional[List[Message]] = None
-
-    def to_dict(self) -> Dict[str, Any]:
-        _dict = {}
-        if self.add_messages is not None:
-            _dict["add_messages"] = [m.to_dict() for m in self.add_messages]
-        if self.reasoning_messages is not None:
-            _dict["reasoning_messages"] = [m.to_dict() for m in self.reasoning_messages]
-        if self.reasoning_steps is not None:
-            _dict["reasoning_steps"] = [rs.model_dump() for rs in self.reasoning_steps]
-        if self.references is not None:
-            _dict["references"] = [r.model_dump() for r in self.references]
-        return _dict
-
-    @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> "RunResponseExtraData":
-        add_messages = data.pop("add_messages", None)
-        if add_messages is not None:
-            add_messages = [Message.model_validate(message) for message in add_messages]
-
-        reasoning_steps = data.pop("reasoning_steps", None)
-        if reasoning_steps is not None:
-            reasoning_steps = [ReasoningStep.model_validate(step) for step in reasoning_steps]
-
-        reasoning_messages = data.pop("reasoning_messages", None)
-        if reasoning_messages is not None:
-            reasoning_messages = [Message.model_validate(message) for message in reasoning_messages]
-
-        references = data.pop("references", None)
-        if references is not None:
-            references = [MessageReferences.model_validate(reference) for reference in references]
-
-        return cls(
-            add_messages=add_messages,
-            reasoning_steps=reasoning_steps,
-            reasoning_messages=reasoning_messages,
-            references=references,
-        )
-
-
 class RunStatus(str, Enum):
     """State of the main run response"""
 

agno/run/cancel.py

@@ -0,0 +1,81 @@
+"""Run cancellation management."""
+
+import threading
+from typing import Dict
+
+from agno.exceptions import RunCancelledException
+from agno.utils.log import logger
+
+
+class RunCancellationManager:
+    """Manages cancellation state for agent runs."""
+
+    def __init__(self):
+        self._cancelled_runs: Dict[str, bool] = {}
+        self._lock = threading.Lock()
+
+    def register_run(self, run_id: str) -> None:
+        """Register a new run as not cancelled."""
+        with self._lock:
+            self._cancelled_runs[run_id] = False
+
+    def cancel_run(self, run_id: str) -> bool:
+        """Cancel a run by marking it as cancelled.
+
+        Returns:
+            bool: True if run was found and cancelled, False if run not found.
+        """
+        with self._lock:
+            if run_id in self._cancelled_runs:
+                self._cancelled_runs[run_id] = True
+                logger.info(f"Run {run_id} marked for cancellation")
+                return True
+            else:
+                logger.warning(f"Attempted to cancel unknown run {run_id}")
+                return False
+
+    def is_cancelled(self, run_id: str) -> bool:
+        """Check if a run is cancelled."""
+        with self._lock:
+            return self._cancelled_runs.get(run_id, False)
+
+    def cleanup_run(self, run_id: str) -> None:
+        """Remove a run from tracking (called when run completes)."""
+        with self._lock:
+            if run_id in self._cancelled_runs:
+                del self._cancelled_runs[run_id]
+
+    def raise_if_cancelled(self, run_id: str) -> None:
+        """Check if a run should be cancelled and raise exception if so."""
+        if self.is_cancelled(run_id):
+            logger.info(f"Cancelling run {run_id}")
+            raise RunCancelledException(f"Run {run_id} was cancelled")
+
+    def get_active_runs(self) -> Dict[str, bool]:
+        """Get all currently tracked runs and their cancellation status."""
+        with self._lock:
+            return self._cancelled_runs.copy()
+
+
+# Global cancellation manager instance
+_cancellation_manager = RunCancellationManager()
+
+
+def register_run(run_id: str) -> None:
+    """Register a new run for cancellation tracking."""
+    _cancellation_manager.register_run(run_id)
+
+
+def cancel_run(run_id: str) -> bool:
+    """Cancel a run."""
+    return _cancellation_manager.cancel_run(run_id)
+
+
+def cleanup_run(run_id: str) -> None:
+    """Clean up cancellation tracking for a completed run."""
+    _cancellation_manager.cleanup_run(run_id)
+
+
+def raise_if_cancelled(run_id: str) -> None:
+    """Check if a run should be cancelled and raise exception if so."""
+    _cancellation_manager.raise_if_cancelled(run_id)