aixtools 0.3.2__py3-none-any.whl → 0.3.4__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.3.2'
-__version_tuple__ = version_tuple = (0, 3, 2)
+__version__ = version = '0.3.4'
+__version_tuple__ = version_tuple = (0, 3, 4)
 
 __commit_id__ = commit_id = None

aixtools/a2a/app.py

@@ -6,6 +6,9 @@ import json
 from functools import partial
 from typing import assert_never
 
+from a2a.server.apps import A2AStarletteApplication
+from a2a.server.request_handlers import RequestHandler
+from a2a.types import AgentCard
 from fasta2a.applications import FastA2A
 from fasta2a.broker import InMemoryBroker
 from fasta2a.schema import Part, TaskSendParams
@@ -26,6 +29,8 @@ from starlette.exceptions import HTTPException
 from starlette.requests import Request
 from starlette.responses import RedirectResponse
 
+from aixtools.a2a.auth_middleware import AuthMiddleware
+from aixtools.auth.auth import AccessTokenAuthProvider
 from aixtools.context import session_id_var, user_id_var
 
 
@@ -124,3 +129,18 @@ def fix_a2a_docs_pages(app: Starlette) -> None:
 
     app.router.add_route("/.well-known/agent.json", redirect_to_sub_agent, methods=["GET"])
     app.router.add_route("/", redirect_to_sub_agent, methods=["POST"])
+
+
+def build_a2a_starlette_app(
+    public_agent_card: AgentCard, request_handler: RequestHandler, auth_provider: AccessTokenAuthProvider | None = None
+) -> Starlette:
+    """Builds a A2AStarletteApplication with the given auth provider."""
+    server = A2AStarletteApplication(
+        agent_card=public_agent_card,
+        http_handler=request_handler,
+    )
+
+    app = server.build()
+    if auth_provider:
+        app.add_middleware(AuthMiddleware, provider=auth_provider)
+    return app

aixtools/a2a/auth_middleware.py

@@ -0,0 +1,46 @@
+"Auth module managing user authentication for A2A server"
+
+from logging import getLogger
+
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+
+from aixtools.auth.auth import AccessTokenAuthProvider, AuthTokenError
+from aixtools.utils import config
+
+logger = getLogger(__name__)
+
+
+class AuthMiddleware(BaseHTTPMiddleware):  # pylint: disable=too-few-public-methods
+    """
+    Middleware that enforces access token authentication for A2A route.
+    """
+
+    AGENT_CARD_PATH_SUFFIXES = [".well-known/agent.json", ".well-known/agent-card.json"]
+    APP_DEFAULT_SCOPE = config.APP_DEFAULT_SCOPE
+
+    def __init__(self, app, provider: AccessTokenAuthProvider):
+        super().__init__(app)
+        self.provider = provider
+
+    async def dispatch(self, request: Request, call_next):
+        """
+        Auth middleware function that checks whether the request has a valid access token.
+
+        :param request: The incoming http request
+        :param call_next: The next function in the chain
+        :raises HTTPException: If an access token isn't valid or missing
+        """
+
+        path = request.url.path
+
+        # allow agent cards to pass
+        if any(path.endswith(suffix) for suffix in self.AGENT_CARD_PATH_SUFFIXES):
+            return await call_next(request)
+
+        auth_header = request.headers.get("Authorization") or request.headers.get("authorization")
+        try:
+            await self.provider.verify_auth_header(auth_header)
+            return await call_next(request)
+        except AuthTokenError as e:
+            raise e.to_http_exception(self.APP_DEFAULT_SCOPE)

aixtools/a2a/google_sdk/remote_agent_connection.py

@@ -63,7 +63,7 @@ class RemoteAgentConnection:
         self,
         message: Message,
         *,
-        sleep_time: float = 0.2,
+        sleep_time: float = 2.0,
         max_iter=1000,
         on_task_submitted: Callable[[str], None] | None = None,
     ) -> Task | Message:

aixtools/a2a/google_sdk/utils.py

@@ -5,7 +5,7 @@ import asyncio
 import httpx
 from a2a.client import A2ACardResolver, ClientConfig, ClientFactory
 from a2a.server.agent_execution import RequestContext
-from a2a.types import AgentCard
+from a2a.types import AgentCard, PushNotificationConfig
 from a2a.utils import AGENT_CARD_WELL_KNOWN_PATH, PREV_AGENT_CARD_WELL_KNOWN_PATH
 
 from aixtools.a2a.google_sdk.remote_agent_connection import RemoteAgentConnection
@@ -42,9 +42,15 @@ async def get_agent_card(client: httpx.AsyncClient, address: str) -> AgentCard:
 class _AgentCardResolver:
     """Helper class to resolve and manage agent cards and their connections."""
 
-    def __init__(self, client: httpx.AsyncClient):
+    def __init__(
+        self, client: httpx.AsyncClient, push_notification_configs: list[PushNotificationConfig] | None = None
+    ):
         self._httpx_client = client
-        self._a2a_client_factory = ClientFactory(ClientConfig(httpx_client=self._httpx_client, polling=True))
+        self._a2a_client_factory = ClientFactory(
+            ClientConfig(
+                httpx_client=self._httpx_client, polling=True, push_notification_configs=push_notification_configs or []
+            )
+        )
         self.clients: dict[str, RemoteAgentConnection] = {}
 
     def register_agent_card(self, card: AgentCard):
@@ -72,13 +78,20 @@ async def get_a2a_clients(
     agent_hosts: list[str],
     session_id_tuple: SessionIdTuple,
     auth_token: str = None,
+    push_notification_config_url: str | None = None,
     *,
     timeout: float = DEFAULT_A2A_TIMEOUT,
 ) -> dict[str, RemoteAgentConnection]:
     """Get A2A clients for all agents defined in the configuration."""
+    push_notification_configs = (
+        [PushNotificationConfig(url=push_notification_config_url)] if push_notification_config_url else None
+    )
+
     headers = create_session_headers(session_id_tuple, auth_token)
     httpx_client = httpx.AsyncClient(headers=headers, timeout=timeout, follow_redirects=True)
-    clients = await _AgentCardResolver(httpx_client).get_a2a_clients(agent_hosts)
+    clients = await _AgentCardResolver(
+        httpx_client, push_notification_configs=push_notification_configs
+    ).get_a2a_clients(agent_hosts)
     for client in clients.values():
         logger.info("Using A2A server at: %s", client.get_agent_card().url)
     return clients

aixtools/auth/auth.py

@@ -120,6 +120,7 @@ class AccessTokenVerifier:
         except InvalidSignatureError as e:
             raise AuthTokenError(AuthTokenErrorCode.INVALID_SIGNATURE) from e
         except jwt.exceptions.PyJWTError as e:
+            logger.exception("Unable to check JWT token: %s", token)
             raise AuthTokenError(AuthTokenErrorCode.JWT_ERROR) from e
 
     def authorize_claims(self, claims: dict, expected_scope: str):
@@ -148,10 +149,11 @@ class AccessTokenVerifier:
             logger.info("Authorized JWT token, against %s", groups)
             return
 
-        logger.error("Could not find any group in JWT token, matching: %s", self.authorized_groups)
+        email = claims.get("email")
+        logger.warning("User %s group %s does not match configured groups %s", email, groups, self.authorized_groups)
         raise AuthTokenError(
             AuthTokenErrorCode.MISSING_GROUPS_ERROR,
-            f"Could not find any group in JWT token, matching: {self.authorized_groups}",
+            f"User {email} group {groups} does not match configured groups {self.authorized_groups}",
         )
 
 

{aixtools-0.3.2.dist-info → aixtools-0.3.4.dist-info}/METADATA

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

{aixtools-0.3.2.dist-info → aixtools-0.3.4.dist-info}/RECORD

@@ -1,5 +1,5 @@
 aixtools/__init__.py,sha256=9NGHm7LjsQmsvjTZvw6QFJexSvAU4bCoN_KBk9SCa00,260
-aixtools/_version.py,sha256=e8NqPtZ8fggRgk3GPrqZ_U_BDV8aSULw1u_Gn9NNbnk,704
+aixtools/_version.py,sha256=3nDaC5e0d_scBB1bUEKPlItbvbY0PmXNNyyOTNFNWNI,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
@@ -17,11 +17,12 @@ aixtools/.chainlit/translations/nl.json,sha256=R3e-WxkQXAiuQgnnXjFWhwzpn1EA9xJ8g
 aixtools/.chainlit/translations/ta.json,sha256=pxa2uLEEDjiGiT6MFcCJ_kNh5KoFViHFptcJjc79Llc,17224
 aixtools/.chainlit/translations/te.json,sha256=0qGj-ODEHVOcxfVVX5IszS1QBCKSXuU1okANP_EbvBQ,16885
 aixtools/.chainlit/translations/zh-CN.json,sha256=EWxhT2_6CW9z0F6SI2llr3RsaL2omH1QZWHVG2n5POA,8664
-aixtools/a2a/app.py,sha256=p18G7fAInl9dcNYq6RStBjv1C3aD6oilQq3WXtBuk30,5069
+aixtools/a2a/app.py,sha256=ugx9FR8QadpEIgy79V6vx1zoamm6ldds6PYKePPQ3sA,5809
+aixtools/a2a/auth_middleware.py,sha256=Q9JmM0JZyJinQBm_BbSizhKpUgIvyWXEegCUmm204AI,1635
 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=gFsIja_nIQsAzxElv3A7_hG9GpDI5pkV-B_KDlcYlnc,3329
-aixtools/a2a/google_sdk/utils.py,sha256=k5VtTr4XLxsSKQDqly9QtWcZ5xfm03VbPpQ8OwoXtiQ,4369
+aixtools/a2a/google_sdk/remote_agent_connection.py,sha256=wJXIHD76yL3nlI4b9IA7FQlpKRxrFDsSh78TIm1Fy4E,3329
+aixtools/a2a/google_sdk/utils.py,sha256=eEjYBeSYS7U8RsRxC-rTiw0VXo25CIR9VyRoqdvzGY4,4859
 aixtools/a2a/google_sdk/pydantic_ai_adapter/agent_executor.py,sha256=xamLIKl-FeB8PElVKgiIKR-MhFKmPjywHmC_ow-fGAA,9904
 aixtools/a2a/google_sdk/pydantic_ai_adapter/storage.py,sha256=nGoVL7MPoZJW7iVR71laqpUYP308yFKZIifJtvUgpiU,878
 aixtools/agents/__init__.py,sha256=MAW196S2_G7uGqv-VNjvlOETRfuV44WlU1leO7SiR0A,282
@@ -33,7 +34,7 @@ aixtools/agents/nodes_to_str.py,sha256=UkOu5Nry827J4H_ohQU3tPBfJxtr3p6FfCfWoUy5u
 aixtools/agents/print_nodes.py,sha256=wVTngNfqM0As845WTRz6G3Rei_Gr3HuBlvu-G_eXuig,1665
 aixtools/agents/prompt.py,sha256=oZl6_3SelyoSysLpF6AAmLHLHhwyPYCtX8hJ2pRUnhw,7396
 aixtools/auth/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-aixtools/auth/auth.py,sha256=rCvkcZCDNo-f3lsuMkLlcfpTO5GlRsHY9qhCnB7ayhU,8123
+aixtools/auth/auth.py,sha256=0GAelLK6MaLq-iHm4jE76UPvLplEfodlxb0VYRttLNk,8258
 aixtools/compliance/__init__.py,sha256=vnw0zEdySIJWvDAJ8DCRRaWmY_agEOz1qlpAdhmtiuo,191
 aixtools/compliance/private_data.py,sha256=OOM9mIp3_w0fNgj3VAEWBl7-jrPc19_Ls1pC5dfF5UY,5323
 aixtools/db/__init__.py,sha256=b8vRhme3egV-aUZbAntnOaDkSXB8UT0Xy5oqQhU_z0Q,399
@@ -95,8 +96,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.3.2.dist-info/METADATA,sha256=FnBXFfGny-BUh1L6fkR_JUZy7OjZPbjRhhoDIgfg5jM,28014
-aixtools-0.3.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-aixtools-0.3.2.dist-info/entry_points.txt,sha256=q8412TG4T0S8K0SKeWp2vkVPIDYQs0jNoHqcQ7qxOiA,155
-aixtools-0.3.2.dist-info/top_level.txt,sha256=wBn-rw9bCtxrR4AYEYgjilNCUVmKY0LWby9Zan2PRJM,9
-aixtools-0.3.2.dist-info/RECORD,,
+aixtools-0.3.4.dist-info/METADATA,sha256=tIIys8pgVbNONjTKYW0fVbbEtkhwQP8F18hYuPQff1c,28014
+aixtools-0.3.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+aixtools-0.3.4.dist-info/entry_points.txt,sha256=q8412TG4T0S8K0SKeWp2vkVPIDYQs0jNoHqcQ7qxOiA,155
+aixtools-0.3.4.dist-info/top_level.txt,sha256=wBn-rw9bCtxrR4AYEYgjilNCUVmKY0LWby9Zan2PRJM,9
+aixtools-0.3.4.dist-info/RECORD,,