aixtools 0.2.8__tar.gz → 0.2.9__tar.gz

{aixtools-0.2.8 → aixtools-0.2.9}/PKG-INFO

@@ -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 → aixtools-0.2.9}/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 = 'gac214deb6'
+__commit_id__ = commit_id = 'g9ef985eb7'

{aixtools-0.2.8 → aixtools-0.2.9}/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-0.2.8 → aixtools-0.2.9}/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-0.2.8 → aixtools-0.2.9}/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):