aixtools 0.3.3__tar.gz → 0.3.5__tar.gz

{aixtools-0.3.3 → aixtools-0.3.5}/PKG-INFO

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

{aixtools-0.3.3 → aixtools-0.3.5}/aixtools/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.3.3'
-__version_tuple__ = version_tuple = (0, 3, 3)
+__version__ = version = '0.3.5'
+__version_tuple__ = version_tuple = (0, 3, 5)
 
-__commit_id__ = commit_id = 'g7da95a013'
+__commit_id__ = commit_id = 'g9eb31448f'

{aixtools-0.3.3 → aixtools-0.3.5}/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-0.3.5/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-0.3.3 → aixtools-0.3.5}/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.3 → aixtools-0.3.5}/aixtools/testing/aix_test_model.py

@@ -19,6 +19,20 @@ from ..utils.utils import async_iter
 FINAL_RESULT_TOOL_NAME = "final_result"
 
 
+def _convert_to_parts(items: list[str | TextPart | ToolCallPart]) -> list[TextPart | ToolCallPart]:
+    """Convert a list of strings, TextParts, or ToolCallParts to a list of Parts."""
+    parts: list[TextPart | ToolCallPart] = []
+    for item in items:
+        match item:
+            case str():
+                parts.append(TextPart(item))
+            case TextPart() | ToolCallPart():
+                parts.append(item)
+            case _:
+                raise ValueError(f"Invalid item type in list: {type(item)}, item: {item}")
+    return parts
+
+
 def final_result_tool(result: BaseModel | dict) -> ToolCallPart:
     """Create a ToolCallPart for the final result."""
     if isinstance(result, BaseModel):
@@ -33,7 +47,11 @@ class AixTestModel(Model):
     Test model, returns a specified list of answers, including messages or tool calls
     This is used for testing the agent and model interaction with the rest of the system.
 
-    responses: Is a list of strings or ToolCallPart objects that the model will return in order.
+    responses: Is a list of response items that the model will return in order. Each item can be:
+        - str: A simple text response
+        - TextPart: A text part
+        - ToolCallPart: A tool call
+        - list[str | TextPart | ToolCallPart]: Multiple parts in a single response
 
     Note: The agent will continue to invoke 'request()' until it returns a txt (i.e. not a ToolCallPart).
 
@@ -69,11 +87,25 @@ class AixTestModel(Model):
                     ]
                 )
     ```
+
+    Example with multiple parts in a single response:
+    ```
+        model = AixTestModel(
+                    responses=[
+                        [
+                            "First, let me think about this...",
+                            ToolCallPart(tool_name='analyze', args={'data': 'xyz'}),
+                            TextPart("Analysis complete."),
+                        ],
+                        "Final answer after analysis.",
+                    ]
+                )
+    ```
     """
 
     def __init__(  # pylint: disable=super-init-not-called
         self,
-        responses: list[str | TextPart | ToolCallPart] | AsyncGeneratorType,
+        responses: list[str | TextPart | ToolCallPart | list[str | TextPart | ToolCallPart]] | AsyncGeneratorType,
         sleep_time: float | None = None,
     ):
         self.response_iter = responses(self) if callable(responses) else async_iter(responses)
@@ -141,15 +173,14 @@ class AixTestModel(Model):
         if callable(res):
             res = res(self, messages)
         match res:
-            case str():
-                return ModelResponse(parts=[TextPart(res)], model_name=self.model_name)
-            case ToolCallPart():
-                return ModelResponse(parts=[res], model_name=self.model_name)
-            case TextPart():
-                return ModelResponse(parts=[res], model_name=self.model_name)
             case Exception():
                 raise res
+            case list():
+                parts = _convert_to_parts(res)
+            case str() | TextPart() | ToolCallPart():
+                parts = _convert_to_parts([res])
             case _:
                 raise ValueError(f"Invalid response type: {type(res)}, response: {res}")
         if self.sleep_time:
             await asyncio.sleep(self.sleep_time)
+        return ModelResponse(parts=parts, model_name=self.model_name)

{aixtools-0.3.3 → aixtools-0.3.5}/aixtools.egg-info/SOURCES.txt

@@ -20,6 +20,7 @@ aixtools/.chainlit/translations/ta.json
 aixtools/.chainlit/translations/te.json
 aixtools/.chainlit/translations/zh-CN.json
 aixtools/a2a/app.py
+aixtools/a2a/auth_middleware.py
 aixtools/a2a/utils.py
 aixtools/a2a/google_sdk/__init__.py
 aixtools/a2a/google_sdk/remote_agent_connection.py