aip-agents-binary 0.0.0__py3-none-any.whl

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

@@ -0,0 +1,240 @@
+"""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 | None,
+    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=None if timeout is None else 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
+    raw_timeout = server_config.get("timeout")
+    timeout = None if raw_timeout is None else float(raw_timeout)
+    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,430 @@
+"""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 dataclasses import dataclass
+from typing import Any
+
+# Reserved module names that cannot be used for server packages.
+RESERVED_MODULE_NAMES = frozenset({"ptc_helper", "mcp_client", "custom"})
+
+# Default placeholder for example values.
+DEFAULT_EXAMPLE_PLACEHOLDER = '"example"'
+
+# Shared kwargs signature fragment.
+KWARGS_ANY_PARAMS = "**kwargs: Any"
+RESERVED_CUSTOM_PARAM_NAMES = frozenset({"kwargs"})
+
+
+@dataclass(frozen=True)
+class CustomToolFieldSpec:
+    """Projected field metadata for a custom tool schema property."""
+
+    schema_key: str
+    param_name: str
+    type_hint: str
+    required: bool
+
+
+@dataclass(frozen=True)
+class CustomToolRuntimeParamSpec:
+    """Runtime signature metadata for inspect.Signature reconstruction."""
+
+    name: str
+    kind: str
+    annotation_expr: str
+    has_default: bool = False
+    default_expr: str | None = None
+
+
+@dataclass(frozen=True)
+class CustomToolSignatureProjection:
+    """Canonical custom-tool signature projection derived from JSON schema."""
+
+    field_specs: list[CustomToolFieldSpec]
+    canonical_params: str
+    runtime_params: list[CustomToolRuntimeParamSpec]
+    schema_to_param: dict[str, str]
+    param_to_schema: dict[str, str]
+
+
+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 is_valid_identifier(name: str) -> bool:
+    """Check if a string is a valid Python identifier and not a keyword.
+
+    Args:
+        name: Candidate identifier.
+
+    Returns:
+        True if name is a valid identifier and not a keyword.
+    """
+    return name.isidentifier() and not keyword.iskeyword(name)
+
+
+def _json_scalar_type_to_python(json_type: Any) -> str:
+    """Convert a scalar JSON schema type token into a Python type hint token."""
+    type_map = {
+        "string": "str",
+        "integer": "int",
+        "number": "float",
+        "boolean": "bool",
+        "array": "list",
+        "object": "dict",
+        "null": "None",
+    }
+    return type_map.get(json_type, "Any") if isinstance(json_type, str) else "Any"
+
+
+def _json_union_type_to_python(json_types: list[Any]) -> str:
+    """Convert JSON schema union type list into a Python type hint expression."""
+    if not json_types:
+        return "Any"
+
+    converted_types: list[str] = []
+    seen: set[str] = set()
+    has_none = False
+
+    for json_type in json_types:
+        python_type = _json_scalar_type_to_python(json_type)
+        if python_type == "None":
+            has_none = True
+            continue
+        if python_type in seen:
+            continue
+        seen.add(python_type)
+        converted_types.append(python_type)
+
+    if has_none:
+        converted_types.append("None")
+
+    if not converted_types:
+        return "Any"
+
+    return " | ".join(converted_types)
+
+
+def json_type_to_python(json_type: str | list) -> str:
+    """Convert JSON schema type to Python type hint.
+
+    Handles union types (list form) by converting each type and joining with |.
+    For example, ["string", "null"] -> "str | None".
+
+    Edge cases handled:
+    - Empty list returns "Any"
+    - Duplicate types are deduplicated while preserving order
+    - "null" is normalized to end of union (e.g., ["null", "string"] -> "str | None")
+    - Unknown types map to "Any" (deduplicated if multiple)
+
+    Args:
+        json_type: JSON schema type (string or list of strings for union types).
+
+    Returns:
+        Python type hint string.
+    """
+    if isinstance(json_type, list):
+        return _json_union_type_to_python(json_type)
+
+    return _json_scalar_type_to_python(json_type)
+
+
+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
+
+    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}: {_optional_type_hint(type_hint)} = None")
+
+    # Required params first, then optional
+    all_params = params + optional_params
+    return ", ".join(all_params) if all_params else KWARGS_ANY_PARAMS
+
+
+def _optional_type_hint(type_hint: str) -> str:
+    """Return nullable hint while preserving explicit None type.
+
+    Avoids duplicating None if the type hint already includes it
+    (e.g., "str | None" stays as "str | None", not "str | None | None").
+
+    Also deduplicates if None appears multiple times in the union
+    (e.g., "str | None | None" becomes "str | None").
+
+    Args:
+        type_hint: The base type hint string.
+
+    Returns:
+        Type hint string that includes None.
+    """
+    normalized = type_hint.strip()
+    if not normalized:
+        return "Any | None"
+
+    if normalized == "None":
+        return "None"
+
+    parts = [p.strip() for p in normalized.split("|") if p.strip()]
+    if not parts:
+        return "Any | None"
+
+    # If None is already in the union, deduplicate and return
+    if "None" in parts:
+        seen = set()
+        deduplicated: list[str] = []
+        for part in parts:
+            if part not in seen:
+                seen.add(part)
+                deduplicated.append(part)
+        return " | ".join(deduplicated)
+
+    return " | ".join([*parts, "None"])
+
+
+def _empty_custom_tool_signature_projection() -> CustomToolSignatureProjection:
+    """Return default projection for schema without properties."""
+    return CustomToolSignatureProjection(
+        field_specs=[],
+        canonical_params=KWARGS_ANY_PARAMS,
+        runtime_params=[
+            CustomToolRuntimeParamSpec(
+                name="kwargs",
+                kind="VAR_KEYWORD",
+                annotation_expr="Any",
+            )
+        ],
+        schema_to_param={},
+        param_to_schema={},
+    )
+
+
+def _validate_custom_param_name(
+    schema_key: str,
+    param_name: str,
+    param_to_schema: dict[str, str],
+) -> None:
+    """Validate reserved-name and sanitization-collision constraints."""
+    if param_name in RESERVED_CUSTOM_PARAM_NAMES:
+        raise ValueError(f"Custom tool schema parameter '{schema_key}' maps to reserved name '{param_name}'")
+
+    prior_schema_key = param_to_schema.get(param_name)
+    if prior_schema_key and prior_schema_key != schema_key:
+        raise ValueError(
+            "Custom tool schema parameter sanitization collision: "
+            f"'{prior_schema_key}' and '{schema_key}' both map to '{param_name}'"
+        )
+
+
+def _build_canonical_custom_params(field_specs: list[CustomToolFieldSpec]) -> str:
+    """Build keyword-only canonical parameter signature string."""
+    canonical_parts = ["*"]
+    for field in field_specs:
+        if field.required:
+            canonical_parts.append(f"{field.param_name}: {field.type_hint}")
+            continue
+        canonical_parts.append(f"{field.param_name}: {_optional_type_hint(field.type_hint)} = None")
+    canonical_parts.append(KWARGS_ANY_PARAMS)
+    return ", ".join(canonical_parts)
+
+
+def project_custom_tool_signature(schema: dict[str, Any]) -> CustomToolSignatureProjection:
+    """Project schema into a canonical custom-tool signature contract.
+
+    This helper is custom-tool specific and intentionally independent from
+    `schema_to_params(...)` so MCP behavior remains unchanged.
+
+    Args:
+        schema: JSON schema for custom tool input.
+
+    Returns:
+        Canonical projection for rendering and runtime introspection.
+
+    Raises:
+        ValueError: If parameter sanitization causes collisions.
+    """
+    properties = schema.get("properties", {})
+    required = set(schema.get("required", []))
+
+    if not properties:
+        return _empty_custom_tool_signature_projection()
+
+    field_specs: list[CustomToolFieldSpec] = []
+    runtime_params: list[CustomToolRuntimeParamSpec] = []
+    schema_to_param: dict[str, str] = {}
+    param_to_schema: dict[str, str] = {}
+
+    for schema_key, prop_schema in properties.items():
+        param_name = sanitize_param_name(schema_key)
+        _validate_custom_param_name(schema_key, param_name, param_to_schema)
+
+        type_hint = json_type_to_python(prop_schema.get("type", "any"))
+        is_required = schema_key in required
+
+        field_specs.append(
+            CustomToolFieldSpec(
+                schema_key=schema_key,
+                param_name=param_name,
+                type_hint=type_hint,
+                required=is_required,
+            )
+        )
+        runtime_params.append(
+            CustomToolRuntimeParamSpec(
+                name=param_name,
+                kind="KEYWORD_ONLY",
+                annotation_expr=type_hint if is_required else _optional_type_hint(type_hint),
+                has_default=not is_required,
+                default_expr="None" if not is_required else None,
+            )
+        )
+        schema_to_param[schema_key] = param_name
+        param_to_schema[param_name] = schema_key
+
+    runtime_params.append(
+        CustomToolRuntimeParamSpec(
+            name="kwargs",
+            kind="VAR_KEYWORD",
+            annotation_expr="Any",
+        )
+    )
+
+    return CustomToolSignatureProjection(
+        field_specs=field_specs,
+        canonical_params=_build_canonical_custom_params(field_specs),
+        runtime_params=runtime_params,
+        schema_to_param=schema_to_param,
+        param_to_schema=param_to_schema,
+    )
+
+
+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)