aixtools 0.2.20__tar.gz → 0.2.22__tar.gz

{aixtools-0.2.20 → aixtools-0.2.22}/PKG-INFO

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

{aixtools-0.2.20 → aixtools-0.2.22}/aixtools/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.2.20'
-__version_tuple__ = version_tuple = (0, 2, 20)
+__version__ = version = '0.2.22'
+__version_tuple__ = version_tuple = (0, 2, 22)
 
-__commit_id__ = commit_id = 'ge7f4f14bb'
+__commit_id__ = commit_id = 'g15d95015f'

{aixtools-0.2.20 → aixtools-0.2.22}/aixtools/a2a/google_sdk/pydantic_ai_adapter/agent_executor.py

@@ -20,7 +20,7 @@ from pydantic_ai import Agent, BinaryContent
 
 from aixtools.a2a.google_sdk.pydantic_ai_adapter.storage import InMemoryHistoryStorage
 from aixtools.a2a.google_sdk.remote_agent_connection import is_in_terminal_state
-from aixtools.a2a.google_sdk.utils import get_session_id_tuple
+from aixtools.a2a.google_sdk.utils import get_auth_token, get_session_id_tuple
 from aixtools.agents import get_agent
 from aixtools.agents.prompt import build_user_input
 from aixtools.context import SessionIdTuple
@@ -100,7 +100,8 @@ class PydanticAgentExecutor(AgentExecutor):
             event_queue (EventQueue): The event queue to enqueue events.
         """
         session_tuple = get_session_id_tuple(context)
-        agent = self._build_agent(session_tuple)
+        auth_token = get_auth_token(context)
+        agent = self._build_agent(session_tuple, auth_token)
         if context.message is None:
             raise ValueError("No message provided")
 
@@ -258,9 +259,9 @@ class PydanticAgentExecutor(AgentExecutor):
             )
         )
 
-    def _build_agent(self, session_tuple: SessionIdTuple) -> Agent:
+    def _build_agent(self, session_tuple: SessionIdTuple, auth_token: str | None) -> Agent:
         params = self._agent_parameters
-        mcp_servers = get_configured_mcp_servers(session_tuple, params.mcp_servers)
+        mcp_servers = get_configured_mcp_servers(session_tuple, params.mcp_servers, auth_token=auth_token)
         return get_agent(
             system_prompt=params.system_prompt,
             toolsets=mcp_servers,

{aixtools-0.2.20 → aixtools-0.2.22}/aixtools/a2a/google_sdk/utils.py

@@ -96,3 +96,15 @@ def get_session_id_tuple(context: RequestContext) -> SessionIdTuple:
     """Get the user_id, session_id tuple from the request context."""
     headers = context.call_context.state.get("headers", {})
     return headers.get("user-id", DEFAULT_USER_ID), headers.get("session-id", DEFAULT_SESSION_ID)
+
+
+def get_auth_token(context: RequestContext) -> str | None:
+    """Get the Authorization bearer token from the request context."""
+    headers = context.call_context.state.get("headers", {})
+    auth_header = headers.get("Authorization") or headers.get("authorization")
+    auth_token = None
+    if auth_header and auth_header.lower().startswith("bearer "):
+        auth_token = auth_header.split(" ", 1)[1].strip()
+    else:
+        logger.warning("Cannot find bearer token in Authorization header: %s", auth_header)
+    return auth_token

{aixtools-0.2.20 → aixtools-0.2.22}/aixtools/mcp/client.py

@@ -23,6 +23,7 @@ from pydantic_ai.toolsets.abstract import ToolsetTool
 from aixtools.context import SessionIdTuple
 from aixtools.logging.logging_config import get_logger
 from aixtools.server.utils import create_session_headers
+from aixtools.utils.config import MCP_TOOLS_MAX_RETRIES
 
 MCP_TOOL_CACHE_TTL = 300  # 5 minutes
 DEFAULT_MCP_CONNECTION_TIMEOUT = 30
@@ -31,7 +32,6 @@ CACHE_KEY = "TOOL_LIST"
 
 logger = get_logger(__name__)
 
-
 # Default log_handler for MCP clients
 LOGGING_LEVEL_MAP = logging.getLevelNamesMapping()
 
@@ -87,6 +87,7 @@ def get_mcp_servers(
     mcp_configs: list[MCPConfig],
     session_id_tuple: SessionIdTuple,
     auth_token: str = None,
+    max_retries: int = MCP_TOOLS_MAX_RETRIES,
     *,
     timeout: float = DEFAULT_MCP_CONNECTION_TIMEOUT,
 ):
@@ -115,7 +116,7 @@ def get_mcp_servers(
     servers = []
     for config in mcp_configs:
         server = CachedMCPServerStreamableHTTP(
-            url=config.url, headers=headers, timeout=timeout, read_timeout=config.read_timeout
+            url=config.url, headers=headers, timeout=timeout, read_timeout=config.read_timeout, max_retries=max_retries
         )
         logger.info("Using MCP server at %s", config.url)
         servers.append(server)
@@ -123,10 +124,16 @@ def get_mcp_servers(
 
 
 def get_configured_mcp_servers(
-    session_id_tuple: SessionIdTuple, mcp_urls: list[str], timeout: int = DEFAULT_MCP_CONNECTION_TIMEOUT
+    session_id_tuple: SessionIdTuple,
+    mcp_urls: list[str],
+    *,
+    timeout: int = DEFAULT_MCP_CONNECTION_TIMEOUT,
+    auth_token: str | None = None,
 ):
     """Create MCP server instances from a list of URLs."""
-    return get_mcp_servers([MCPConfig(url=url) for url in mcp_urls], session_id_tuple, timeout=timeout)
+    return get_mcp_servers(
+        [MCPConfig(url=url) for url in mcp_urls], session_id_tuple, timeout=timeout, auth_token=auth_token
+    )
 
 
 class CachedMCPServerStreamableHTTP(MCPServerStreamableHTTP):

{aixtools-0.2.20 → aixtools-0.2.22}/aixtools/utils/config.py

@@ -144,3 +144,5 @@ SKIP_MCP_AUTHORIZATION = str2bool(get_variable_env("SKIP_MCP_AUTHORIZATION", Tru
 APP_DEFAULT_SCOPE = get_variable_env("APP_DEFAULT_SCOPE", allow_empty=True)
 
 AUTH_TEST_TOKEN = get_variable_env("AUTH_TEST_TOKEN", allow_empty=True)
+
+MCP_TOOLS_MAX_RETRIES = int(get_variable_env("MCP_TOOLS_MAX_RETRIES", default=10))