aixtools 0.2.8__tar.gz → 0.2.10__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aixtools
-Version: 0.2.8
+Version: 0.2.10
 Summary: Tools for AI exploration and debugging
 Requires-Python: >=3.11.2
 Description-Content-Type: text/markdown
@@ -9,7 +9,7 @@ Requires-Dist: cachebox>=5.0.1
 Requires-Dist: chainlit>=2.5.5
 Requires-Dist: colorlog>=6.9.0
 Requires-Dist: fasta2a>=0.5.0
-Requires-Dist: fastmcp>=2.10.2
+Requires-Dist: fastmcp>=2.11.0
 Requires-Dist: hvac>=2.3.0
 Requires-Dist: ipykernel>=6.29.5
 Requires-Dist: langchain-chroma>=0.2.3

{aixtools-0.2.8 → aixtools-0.2.10}/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.10'
+__version_tuple__ = version_tuple = (0, 2, 10)
 
-__commit_id__ = commit_id = 'gac214deb6'
+__commit_id__ = commit_id = 'ge6b757983'

{aixtools-0.2.8 → aixtools-0.2.10}/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.10}/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.10}/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 → aixtools-0.2.10}/aixtools/auth/auth.py

@@ -7,7 +7,11 @@ import logging
 
 import jwt
 from fastapi import HTTPException
+from fastmcp.server.auth.auth import AuthProvider
 from jwt import ExpiredSignatureError, InvalidAudienceError, InvalidIssuerError, InvalidSignatureError, PyJWKClient
+from mcp.server.auth.provider import (
+    AccessToken,
+)
 
 from aixtools.utils import config
 
@@ -147,3 +151,28 @@ class AccessTokenVerifier:
             AuthTokenErrorCode.MISSING_GROUPS_ERROR,
             f"Could not find any group in JWT token, matching: {self.authorized_groups}",
         )
+
+
+class AccessTokenAuthProvider(AuthProvider):
+    """Authentication provider for MCP servers for validating, authorizing and extracting access tokens."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.token_verifier = AccessTokenVerifier()
+
+    async def verify_token(self, token: str) -> AccessToken:
+        """Verify the access token and return an AccessToken object."""
+
+        logger.info("Received verify token request")
+        claims = self.token_verifier.verify(token)
+        scopes = claims.get("scp", "")
+        self.token_verifier.authorize_claims(claims, scopes)
+
+        scopes_arr = []
+        if scopes:
+            scopes_arr = scopes.split(" ")
+
+        logger.info("Authorized the token")
+        return AccessToken(
+            token=token, client_id=self.token_verifier.api_id, scopes=scopes_arr, expires_at=claims.get("exp", None)
+        )

{aixtools-0.2.8 → aixtools-0.2.10}/aixtools/utils/config.py

@@ -138,3 +138,6 @@ APP_TENANT_ID = get_variable_env("APP_TENANT_ID")
 
 # used for token authorization check
 APP_AUTHORIZED_GROUPS = get_variable_env("APP_AUTHORIZED_GROUPS", allow_empty=True)
+
+# used to skip authorization in local tests if required.
+SKIP_MCP_AUTHORIZATION = str2bool(get_variable_env("SKIP_MCP_AUTHORIZATION", True, False))

{aixtools-0.2.8 → aixtools-0.2.10}/pyproject.toml

@@ -17,7 +17,7 @@ dependencies = [
     "chainlit>=2.5.5",
     "colorlog>=6.9.0",
     "fasta2a>=0.5.0",
-    "fastmcp>=2.10.2",
+    "fastmcp>=2.11.0",
     "hvac>=2.3.0",
     "ipykernel>=6.29.5",
     "langchain-chroma>=0.2.3",