aip-agents-binary 0.6.0__py3-none-any.whl → 0.6.2__py3-none-any.whl

aip_agents/ptc/mcp/templates/mcp_client.py.template

@@ -0,0 +1,239 @@
+"""MCP Client for sandbox PTC execution.
+
+This module uses the official MCP Python SDK to call MCP tools from within
+the E2B sandbox. It provides sync wrappers around the async SDK.
+"""
+
+import asyncio
+import json
+import os
+from typing import Any
+
+import httpx
+from mcp import ClientSession
+from mcp.client.streamable_http import streamable_http_client
+from mcp.client.sse import sse_client
+
+
+def load_config() -> dict[str, Any]:
+    """Load PTC config from ptc_config.json."""
+    config_path = os.path.join(os.path.dirname(__file__), "..", "ptc_config.json")
+    with open(config_path) as f:
+        return json.load(f)
+
+
+_CONFIG: dict[str, Any] | None = None
+
+
+def get_config() -> dict[str, Any]:
+    """Get cached config."""
+    global _CONFIG
+    if _CONFIG is None:
+        _CONFIG = load_config()
+    return _CONFIG
+
+
+def _is_transient_error(exc: BaseException) -> bool:
+    """Return True for retryable network errors."""
+    if isinstance(exc, asyncio.CancelledError):
+        return False
+    # Network-related errors that are retryable
+    if isinstance(exc, (httpx.HTTPError, OSError, ConnectionError, TimeoutError)):
+        return True
+    # MCP SDK wraps network errors in ExceptionGroup from anyio TaskGroups
+    if isinstance(exc, BaseExceptionGroup):
+        return any(_is_transient_error(sub) for sub in exc.exceptions)
+    return False
+
+
+async def _execute_tool_call(
+    url: str,
+    transport: str,
+    headers: dict[str, str] | None,
+    timeout: float,
+    tool_name: str,
+    arguments: dict[str, Any],
+) -> Any:
+    """Execute a single MCP tool call (no retry)."""
+    if transport == "sse":
+        async with sse_client(
+            url=url,
+            timeout=timeout,
+            sse_read_timeout=timeout * 2,
+            headers=headers,
+        ) as (read_stream, write_stream):
+            async with ClientSession(read_stream, write_stream) as session:
+                await session.initialize()
+                result = await session.call_tool(tool_name, arguments)
+                return _normalize_result(result)
+
+    # Default: streamable-http
+    async with httpx.AsyncClient(
+        timeout=httpx.Timeout(timeout),
+        headers=headers,
+    ) as http_client:
+        async with streamable_http_client(
+            url=url,
+            http_client=http_client,
+        ) as (read_stream, write_stream, _):
+            async with ClientSession(read_stream, write_stream) as session:
+                await session.initialize()
+                result = await session.call_tool(tool_name, arguments)
+                return _normalize_result(result)
+
+
+async def _call_tool_async(
+    server_name: str,
+    tool_name: str,
+    arguments: dict[str, Any],
+) -> Any:
+    """Call an MCP tool asynchronously with retry.
+
+    Args:
+        server_name: Name of the MCP server.
+        tool_name: Name of the tool to call.
+        arguments: Arguments to pass to the tool.
+
+    Returns:
+        Normalized tool result.
+
+    Raises:
+        RuntimeError: If all retry attempts fail.
+    """
+    config = get_config()
+    server_config = config["servers"].get(server_name)
+    if not server_config:
+        raise RuntimeError(f"Server '{server_name}' not found in config")
+
+    allowed = server_config.get("allowed_tools")
+    if allowed and tool_name not in allowed:
+        raise PermissionError(f"Tool '{tool_name}' not allowed on server '{server_name}'")
+
+    url = server_config["url"]
+    headers = server_config.get("headers") or None
+    timeout = float(server_config.get("timeout", 60.0))
+    transport = server_config.get("transport", "streamable-http").lower().replace("_", "-")
+    max_retries = 3
+
+    last_error: Exception | None = None
+    backoff = 0.5
+
+    for attempt in range(max_retries):
+        try:
+            return await _execute_tool_call(url, transport, headers, timeout, tool_name, arguments)
+        except Exception as e:
+            last_error = e
+            if attempt < max_retries - 1 and _is_transient_error(e):
+                await asyncio.sleep(backoff)
+                backoff *= 2
+                continue
+            break
+
+    raise RuntimeError(f"MCP tool call failed: {last_error}") from last_error
+
+
+def _run_async(coro):
+    """Run an async coroutine, handling existing event loops (e.g., in Jupyter/E2B).
+
+    Args:
+        coro: The coroutine to run.
+
+    Returns:
+        The result of the coroutine.
+    """
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        # No running loop, use asyncio.run()
+        return asyncio.run(coro)
+
+    # There's a running loop - we need to run in a separate thread
+    import concurrent.futures
+
+    def run_in_new_loop():
+        """Run coroutine in a new event loop in a thread."""
+        new_loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(new_loop)
+        try:
+            result = new_loop.run_until_complete(coro)
+            new_loop.run_until_complete(new_loop.shutdown_asyncgens())
+            if hasattr(new_loop, "shutdown_default_executor"):
+                new_loop.run_until_complete(new_loop.shutdown_default_executor())
+            return result
+        finally:
+            new_loop.close()
+
+    with concurrent.futures.ThreadPoolExecutor() as executor:
+        future = executor.submit(run_in_new_loop)
+        return future.result()
+
+
+def call_tool(server_name: str, tool_name: str, arguments: dict[str, Any]) -> Any:
+    """Call an MCP tool (sync wrapper).
+
+    This is a synchronous wrapper around the async MCP SDK for easier use
+    in LLM-generated code. Handles both regular scripts and Jupyter/E2B environments.
+
+    Args:
+        server_name: Name of the MCP server.
+        tool_name: Name of the tool to call.
+        arguments: Arguments to pass to the tool.
+
+    Returns:
+        Normalized tool result.
+
+    Raises:
+        RuntimeError: If the tool call fails.
+    """
+    return _run_async(_call_tool_async(server_name, tool_name, arguments))
+
+
+def _normalize_result(result: Any) -> Any:
+    """Normalize MCP tool result.
+
+    Normalization rules:
+    - If result has content array with single text item: return text (JSON-parsed if possible)
+    - Otherwise return structured object with text and non_text lists
+
+    Args:
+        result: MCP CallToolResult object.
+
+    Returns:
+        Normalized result value.
+    """
+    # Handle MCP SDK result object
+    content = getattr(result, "content", [])
+    if not content:
+        return result
+
+    text_contents: list[str] = []
+    non_text_contents: list[Any] = []
+
+    for item in content:
+        item_type = getattr(item, "type", None)
+        if item_type == "text":
+            text_contents.append(getattr(item, "text", ""))
+        else:
+            # Convert non-text content to dict for serialization
+            non_text_contents.append({
+                "type": item_type,
+                **{k: v for k, v in vars(item).items() if k != "type"}
+            })
+
+    # Single text content, no non-text content
+    if len(text_contents) == 1 and not non_text_contents:
+        text = text_contents[0]
+        # Try to parse as JSON
+        stripped = text.strip()
+        if stripped.startswith(("{", "[")):
+            try:
+                return json.loads(text)
+            except json.JSONDecodeError:
+                pass
+        return text
+
+    # Return structured object for complex results
+    return {
+        "text": text_contents,
+        "non_text": non_text_contents,
+    }

aip_agents/ptc/naming.py

@@ -0,0 +1,184 @@
+"""Naming utilities for PTC module.
+
+Shared naming helpers for sanitizing module, function, and parameter names
+used in sandbox code generation and prompt building.
+
+Authors:
+    Putu Ravindra Wiguna (putu.r.wiguna@gdplabs.id)
+"""
+
+import keyword
+import re
+from typing import Any
+
+# Reserved module names that cannot be used for server packages.
+RESERVED_MODULE_NAMES = frozenset({"ptc_helper", "mcp_client"})
+
+# Default placeholder for example values.
+DEFAULT_EXAMPLE_PLACEHOLDER = '"example"'
+
+
+def sanitize_module_name(name: str) -> str:
+    """Sanitize server name for use as Python module name.
+
+    Args:
+        name: Original server name.
+
+    Returns:
+        Valid Python module name (lowercase, underscored).
+    """
+    # Replace non-alphanumeric chars with underscore
+    sanitized = re.sub(r"[^a-zA-Z0-9]", "_", name)
+    # Remove leading digits
+    sanitized = re.sub(r"^\d+", "", sanitized)
+    # Remove consecutive underscores
+    sanitized = re.sub(r"_+", "_", sanitized)
+    # Strip leading/trailing underscores
+    sanitized = sanitized.strip("_")
+    # Ensure not empty
+    if not sanitized:
+        sanitized = "server"
+    sanitized = sanitized.lower()
+    # Avoid Python keywords
+    if keyword.iskeyword(sanitized):
+        sanitized = f"{sanitized}_"
+    return sanitized
+
+
+def sanitize_module_name_with_reserved(name: str) -> str:
+    """Sanitize server name, avoiding reserved module names.
+
+    If the sanitized name collides with a reserved name (e.g., ptc_helper),
+    appends '_mcp' suffix to avoid collision.
+
+    Args:
+        name: Original server name.
+
+    Returns:
+        Valid Python module name that does not collide with reserved names.
+    """
+    sanitized = sanitize_module_name(name)
+    if sanitized in RESERVED_MODULE_NAMES:
+        return f"{sanitized}_mcp"
+    return sanitized
+
+
+def sanitize_function_name(name: str) -> str:
+    """Sanitize tool name for use as Python function name.
+
+    Args:
+        name: Original tool name.
+
+    Returns:
+        Valid Python function name (lowercase, underscored).
+    """
+    return sanitize_module_name(name)
+
+
+def sanitize_param_name(name: str) -> str:
+    """Sanitize parameter name for use as Python parameter.
+
+    Args:
+        name: Original parameter name (e.g., userId, user-id).
+
+    Returns:
+        Valid Python parameter name (lowercase, underscored).
+    """
+    return sanitize_module_name(name)
+
+
+def json_type_to_python(json_type: str | list) -> str:
+    """Convert JSON schema type to Python type hint.
+
+    Args:
+        json_type: JSON schema type.
+
+    Returns:
+        Python type hint string.
+    """
+    if isinstance(json_type, list):
+        return "Any"
+
+    type_map = {
+        "string": "str",
+        "integer": "int",
+        "number": "float",
+        "boolean": "bool",
+        "array": "list",
+        "object": "dict",
+        "null": "None",
+    }
+    return type_map.get(json_type, "Any")
+
+
+def schema_to_params(schema: dict[str, Any]) -> str:
+    """Convert JSON schema to Python function parameters.
+
+    Args:
+        schema: JSON schema for tool input.
+
+    Returns:
+        Function parameter string.
+    """
+    properties = schema.get("properties", {})
+    required = set(schema.get("required", []))
+
+    if not properties:
+        return "**kwargs: Any"
+
+    params: list[str] = []
+    optional_params: list[str] = []
+
+    for prop_name, prop_schema in properties.items():
+        safe_name = sanitize_param_name(prop_name)
+        type_hint = json_type_to_python(prop_schema.get("type", "any"))
+
+        if prop_name in required:
+            params.append(f"{safe_name}: {type_hint}")
+        else:
+            optional_params.append(f"{safe_name}: {type_hint} | None = None")
+
+    # Required params first, then optional
+    all_params = params + optional_params
+    return ", ".join(all_params) if all_params else "**kwargs: Any"
+
+
+def example_value_from_schema(
+    prop_schema: dict[str, Any],
+    default_placeholder: str = DEFAULT_EXAMPLE_PLACEHOLDER,
+) -> str:
+    """Return an example value for a schema property.
+
+    Prefers schema-provided examples, defaults, or enums, and falls back
+    to type-based placeholders.
+
+    Args:
+        prop_schema: Property schema from JSON schema.
+        default_placeholder: Placeholder value to use for unknown types.
+
+    Returns:
+        Example value as a Python literal string.
+    """
+    if prop_schema.get("examples"):
+        return repr(prop_schema["examples"][0])
+    if "default" in prop_schema:
+        return repr(prop_schema["default"])
+    if prop_schema.get("enum"):
+        return repr(prop_schema["enum"][0])
+
+    prop_type = prop_schema.get("type", "string")
+    type_placeholders = {
+        "integer": "1",
+        "number": "1.0",
+        "boolean": "True",
+        "array": "[]",
+        "object": "{}",
+    }
+
+    if prop_type in type_placeholders:
+        return type_placeholders[prop_type]
+
+    if prop_type == "string":
+        return default_placeholder
+
+    return default_placeholder

aip_agents/ptc/payload.py

@@ -0,0 +1,26 @@
+"""Shared payload types for PTC.
+
+This module defines the payload structure used across PTC implementations
+(MCP, custom tools, etc.).
+
+Authors:
+    Putu Ravindra Wiguna (putu.r.wiguna@gdplabs.id)
+"""
+
+from dataclasses import dataclass, field
+
+
+@dataclass
+class SandboxPayload:
+    """Payload to upload to the sandbox for PTC execution.
+
+    Attributes:
+        files: Mapping of file path -> content to upload (static bundle).
+        per_run_files: Mapping of file path -> content to upload on every run.
+            These files are always re-uploaded, even when the static bundle is cached.
+        env: Environment variables to set in the sandbox.
+    """
+
+    files: dict[str, str] = field(default_factory=dict)
+    per_run_files: dict[str, str] = field(default_factory=dict)
+    env: dict[str, str] = field(default_factory=dict)