aixtools 0.2.8__py3-none-any.whl → 0.2.9__py3-none-any.whl

aixtools/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.2.8'
-__version_tuple__ = version_tuple = (0, 2, 8)
+__version__ = version = '0.2.9'
+__version_tuple__ = version_tuple = (0, 2, 9)
 
 __commit_id__ = commit_id = None

aixtools/a2a/google_sdk/pydantic_ai_adapter/agent_executor.py

@@ -1,3 +1,4 @@
+import asyncio
 from pathlib import Path
 
 from a2a.server.agent_execution import AgentExecutor, RequestContext
@@ -54,10 +55,23 @@ def _task_failed_event(text: str, context_id: str | None, task_id: str | None) -
     )
 
 
+def _task_cancelled_event(text: str, context_id: str | None, task_id: str | None) -> TaskStatusUpdateEvent:
+    """Creates a TaskStatusUpdateEvent indicating task cancellation."""
+    return TaskStatusUpdateEvent(
+        status=TaskStatus(
+            state=TaskState.canceled, message=new_agent_text_message(text=text, context_id=context_id, task_id=task_id)
+        ),
+        final=True,
+        context_id=context_id,
+        task_id=task_id,
+    )
+
+
 class PydanticAgentExecutor(AgentExecutor):
     def __init__(self, agent_parameters: AgentParameters):
         self._agent_parameters = agent_parameters
         self.history_storage = InMemoryHistoryStorage()
+        self._running_tasks: dict[str, asyncio.Task] = {}  # Track running agent tasks for cancellation
 
     def _convert_message_to_pydantic_parts(
         self,
@@ -102,11 +116,27 @@ class PydanticAgentExecutor(AgentExecutor):
         prompt = self._convert_message_to_pydantic_parts(session_tuple, message)
         history_message = self.history_storage.get(task.id)
 
-        try:
-            result = await agent.run(
+        # Create and track the agent run task for cancellation
+        agent_task = asyncio.create_task(
+            agent.run(
                 user_prompt=prompt,
                 message_history=history_message,
             )
+        )
+        self._running_tasks[task.id] = agent_task
+
+        try:
+            result = await agent_task
+        except asyncio.CancelledError:
+            # Task was cancelled, send cancellation event
+            await event_queue.enqueue_event(
+                _task_cancelled_event(
+                    text="Task was cancelled",
+                    context_id=context.context_id,
+                    task_id=task.id,
+                )
+            )
+            return
         except Exception as e:
             await event_queue.enqueue_event(
                 _task_failed_event(
@@ -116,6 +146,9 @@ class PydanticAgentExecutor(AgentExecutor):
                 )
             )
             return
+        finally:
+            # Clean up the task from tracking
+            self._running_tasks.pop(task.id, None)
 
         self.history_storage.store(task.id, result.all_messages())
 
@@ -185,9 +218,45 @@ class PydanticAgentExecutor(AgentExecutor):
             )
         )
 
-    async def cancel(self, ctx: RequestContext, event_queue: EventQueue) -> None:
-        """Cancel"""
-        raise Exception("cancel not supported")
+    async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
+        """Cancel the ongoing task identified by the task_id in the context.
+
+        Attempts to stop the running agent task and publishes a TaskStatusUpdateEvent
+        with state TaskState.canceled to the event_queue as per A2A SDK specification.
+
+        Args:
+            context: The request context containing the task ID to cancel.
+            event_queue: The queue to publish the cancellation status update to.
+        """
+        task = context.current_task
+        if not task:
+            logger.warning("No task to cancel in context")
+            return
+
+        task_id = task.id
+
+        # Check if we have a running task to cancel
+        if task_id in self._running_tasks:
+            agent_task = self._running_tasks[task_id]
+            if not agent_task.done():
+                logger.info("Cancelling running agent task: %s", task_id)
+                agent_task.cancel()
+                # The cancellation event will be sent by the execute method's except block
+                return
+
+        # If no running task found, check if task is already in terminal state
+        if is_in_terminal_state(task):
+            logger.info("Task %s is already in terminal state: %s", task_id, task.status.state)
+            return
+
+        # Send cancellation event for tasks that aren't currently running
+        await event_queue.enqueue_event(
+            _task_cancelled_event(
+                text="Task cancelled",
+                context_id=context.context_id,
+                task_id=task_id,
+            )
+        )
 
     def _build_agent(self, session_tuple: SessionIdTuple) -> Agent:
         params = self._agent_parameters

aixtools/a2a/google_sdk/remote_agent_connection.py

@@ -1,10 +1,12 @@
 import asyncio
+from typing import Callable
 
 from a2a.client import Client
 from a2a.types import (
     AgentCard,
     Message,
     Task,
+    TaskIdParams,
     TaskQueryParams,
     TaskState,
 )
@@ -63,6 +65,7 @@ class RemoteAgentConnection:
         *,
         sleep_time: float = 0.2,
         max_iter=1000,
+        on_task_submitted: Callable[[str], None] | None = None,
     ) -> Task | Message:
         """
         Sends a message to the remote agent and polls for the task status at regular intervals.
@@ -75,6 +78,9 @@ class RemoteAgentConnection:
         if isinstance(last_task, Message):
             return last_task
 
+        if on_task_submitted:
+            on_task_submitted(last_task.id)
+
         if is_in_terminal_or_interrupted_state(last_task):
             return last_task
         task_id = last_task.id
@@ -86,3 +92,13 @@ class RemoteAgentConnection:
 
         timeout_seconds = max_iter * sleep_time
         raise Exception(f"Task did not complete in {timeout_seconds} seconds")  # pylint: disable=broad-exception-raised
+
+    async def cancel_task(self, task_id: str) -> Task:
+        """
+        Cancels a task by its ID.
+        """
+        try:
+            return await self._client.cancel_task(TaskIdParams(id=task_id))
+        except Exception as e:
+            logger.error("Exception found in cancel_task: %s", str(e))
+            raise e

aixtools/a2a/google_sdk/utils.py

@@ -40,7 +40,7 @@ class _AgentCardResolver:
 
     def __init__(self, client: httpx.AsyncClient):
         self._httpx_client = client
-        self._a2a_client_factory = ClientFactory(ClientConfig(httpx_client=self._httpx_client))
+        self._a2a_client_factory = ClientFactory(ClientConfig(httpx_client=self._httpx_client, polling=True))
         self.clients: dict[str, RemoteAgentConnection] = {}
 
     def register_agent_card(self, card: AgentCard):

{aixtools-0.2.8.dist-info → aixtools-0.2.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aixtools
-Version: 0.2.8
+Version: 0.2.9
 Summary: Tools for AI exploration and debugging
 Requires-Python: >=3.11.2
 Description-Content-Type: text/markdown

{aixtools-0.2.8.dist-info → aixtools-0.2.9.dist-info}/RECORD

@@ -1,5 +1,5 @@
 aixtools/__init__.py,sha256=9NGHm7LjsQmsvjTZvw6QFJexSvAU4bCoN_KBk9SCa00,260
-aixtools/_version.py,sha256=NCr4lkV1m1Jmisv_F-sazmV2gGdDIaVwqLyPwmjayqM,704
+aixtools/_version.py,sha256=051on7ZmwGNyKvbO1AXKoElw7RjLuRmeJqVOApytNd4,704
 aixtools/app.py,sha256=JzQ0nrv_bjDQokllIlGHOV0HEb-V8N6k_nGQH-TEsVU,5227
 aixtools/chainlit.md,sha256=yC37Ly57vjKyiIvK4oUvf4DYxZCwH7iocTlx7bLeGLU,761
 aixtools/context.py,sha256=I_MD40ZnvRm5WPKAKqBUAdXIf8YaurkYUUHSVVy-QvU,598
@@ -20,9 +20,9 @@ aixtools/.chainlit/translations/zh-CN.json,sha256=EWxhT2_6CW9z0F6SI2llr3RsaL2omH
 aixtools/a2a/app.py,sha256=p18G7fAInl9dcNYq6RStBjv1C3aD6oilQq3WXtBuk30,5069
 aixtools/a2a/utils.py,sha256=EHr3IyyBJn23ni-JcfAf6i3VpQmPs0g1TSnAZazvY_8,4039
 aixtools/a2a/google_sdk/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-aixtools/a2a/google_sdk/remote_agent_connection.py,sha256=oDCRSN3gfONY1Ibp8BrtysIVqfQQ-lWe5N7lr1ymHxY,2819
-aixtools/a2a/google_sdk/utils.py,sha256=7ru263peolaMnMLiWvCDWYimZs9aE9DFel2SyokCGjY,3501
-aixtools/a2a/google_sdk/pydantic_ai_adapter/agent_executor.py,sha256=VRwPNi_T4ShCp3vfcvIbB57gPC-vumuLboPhlBiU3lc,7068
+aixtools/a2a/google_sdk/remote_agent_connection.py,sha256=gFsIja_nIQsAzxElv3A7_hG9GpDI5pkV-B_KDlcYlnc,3329
+aixtools/a2a/google_sdk/utils.py,sha256=4VIPV2GtG5IRY-KSZN6iRMS3ntxG2uKgd_d5tQvnZHw,3515
+aixtools/a2a/google_sdk/pydantic_ai_adapter/agent_executor.py,sha256=8VuU2WXeSHUK3_rRm_mjX6elqdC9NA2uz1aELzeC8BU,9784
 aixtools/a2a/google_sdk/pydantic_ai_adapter/storage.py,sha256=nGoVL7MPoZJW7iVR71laqpUYP308yFKZIifJtvUgpiU,878
 aixtools/agents/__init__.py,sha256=MAW196S2_G7uGqv-VNjvlOETRfuV44WlU1leO7SiR0A,282
 aixtools/agents/agent.py,sha256=tceQByn-RGBIhW8BOjKoP0yhNzZLwAa6CxwhPhRe3PU,7270
@@ -88,8 +88,8 @@ aixtools/utils/chainlit/cl_agent_show.py,sha256=vaRuowp4BRvhxEr5hw0zHEJ7iaSF_5bo
 aixtools/utils/chainlit/cl_utils.py,sha256=fxaxdkcZg6uHdM8uztxdPowg3a2f7VR7B26VPY4t-3c,5738
 aixtools/vault/__init__.py,sha256=fsr_NuX3GZ9WZ7dGfe0gp_5-z3URxAfwVRXw7Xyc0dU,141
 aixtools/vault/vault.py,sha256=9dZLWdZQk9qN_Q9Djkofw9LUKnJqnrX5H0fGusVLBhA,6037
-aixtools-0.2.8.dist-info/METADATA,sha256=Ev-Wdm53oRwRAXJrf7Kn47DoJ2_WeiWae2CSasDVs_o,27229
-aixtools-0.2.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-aixtools-0.2.8.dist-info/entry_points.txt,sha256=q8412TG4T0S8K0SKeWp2vkVPIDYQs0jNoHqcQ7qxOiA,155
-aixtools-0.2.8.dist-info/top_level.txt,sha256=wBn-rw9bCtxrR4AYEYgjilNCUVmKY0LWby9Zan2PRJM,9
-aixtools-0.2.8.dist-info/RECORD,,
+aixtools-0.2.9.dist-info/METADATA,sha256=jWcF3SX9eos7TLkvD4UaurdbD4x9JEcQApycD2u_kB8,27229
+aixtools-0.2.9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+aixtools-0.2.9.dist-info/entry_points.txt,sha256=q8412TG4T0S8K0SKeWp2vkVPIDYQs0jNoHqcQ7qxOiA,155
+aixtools-0.2.9.dist-info/top_level.txt,sha256=wBn-rw9bCtxrR4AYEYgjilNCUVmKY0LWby9Zan2PRJM,9
+aixtools-0.2.9.dist-info/RECORD,,