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

aip_agents/ptc/custom_tools_payload.py

@@ -0,0 +1,499 @@
+"""Custom LangChain tools payload generation for PTC sandbox.
+
+This module generates the sandbox payload (wrapper modules, registry, sources)
+that allows LLM-generated code to call custom LangChain tools inside the sandbox.
+
+The payload includes:
+- tools/custom/__init__.py: Package init re-exporting all tool wrappers
+- tools/custom/<tool_name>.py: Per-tool wrapper module with sync call function
+- tools/custom_sources/__init__.py: Package init for file tool sources
+- tools/custom_sources/<tool_name>/__init__.py: File tool source (copied from user)
+- tools/custom_registry.py: Factory functions to build tool instances
+- tools/custom_invoke.py: Safe invoke helper for calling tool.ainvoke
+- tools/custom_defaults.json: Per-run tool config (uploaded each run)
+
+Authors:
+    Putu Ravindra Wiguna (putu.r.wiguna@gdplabs.id)
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from aip_agents.ptc.custom_tools import (
+    PTCCustomToolConfig,
+    PTCToolDef,
+)
+from aip_agents.ptc.naming import (
+    CustomToolRuntimeParamSpec,
+    project_custom_tool_signature,
+    sanitize_function_name,
+)
+from aip_agents.ptc.payload import SandboxPayload
+from aip_agents.ptc.template_utils import render_template
+from aip_agents.utils.logger import get_logger
+
+_TEMPLATE_PACKAGE = "aip_agents.ptc.custom_tools_templates"
+
+logger = get_logger(__name__)
+
+_SAFE_ANNOTATIONS = frozenset({"str", "int", "float", "bool", "list", "dict", "Any", "None"})
+
+
+@dataclass
+class CustomToolPayloadResult:
+    """Result of custom tool payload generation.
+
+    Attributes:
+        payload: The sandbox payload with files to upload.
+        tool_names: List of sanitized tool names available in sandbox.
+    """
+
+    payload: SandboxPayload
+    tool_names: list[str] = field(default_factory=list)
+
+
+def _validate_tool_configs(
+    tool_configs: dict[str, dict[str, object]] | None,
+    configured_tools: list[PTCToolDef],
+) -> dict[str, dict[str, object]]:
+    """Validate tool_configs and return sanitized values.
+
+    Args:
+        tool_configs: Optional per-tool config values to validate.
+        configured_tools: List of configured tool definitions.
+
+    Returns:
+        Validated tool_configs with only known tools and JSON-serializable dict values.
+    """
+    if not tool_configs:
+        return {}
+
+    if not isinstance(tool_configs, dict):
+        logger.warning("tool_configs must be a dict mapping tool names to dict values")
+        return {}
+
+    configured_names = {t.get("name", "") for t in configured_tools}
+    validated: dict[str, dict[str, object]] = {}
+
+    for tool_name, tool_config in tool_configs.items():
+        if not isinstance(tool_name, str):
+            logger.warning(f"tool_configs contains non-string tool name: {tool_name!r}")
+            continue
+        if tool_name not in configured_names:
+            logger.warning(f"tool_configs contains config for unknown tool: {tool_name}")
+            continue
+        if not isinstance(tool_config, dict):
+            logger.warning(f"tool_configs for tool '{tool_name}' must be a dict")
+            continue
+        if any(not isinstance(key, str) for key in tool_config.keys()):
+            logger.warning(f"tool_configs for tool '{tool_name}' must use string keys")
+            continue
+        try:
+            json.dumps(tool_config)
+        except (TypeError, ValueError) as exc:
+            logger.warning(f"tool_configs for tool '{tool_name}' must be JSON-serializable: {exc}")
+            continue
+        validated[tool_name] = tool_config
+
+    return validated
+
+
+def build_custom_tools_payload(
+    config: PTCCustomToolConfig,
+    tool_configs: dict[str, dict[str, object]] | None = None,
+) -> CustomToolPayloadResult:
+    """Build sandbox payload for custom LangChain tools.
+
+    Args:
+        config: Custom tool configuration with tool definitions.
+        tool_configs: Optional per-tool config values to write to defaults.json.
+
+    Returns:
+        CustomToolPayloadResult with payload and available tool names.
+    """
+    result = CustomToolPayloadResult(payload=SandboxPayload())
+
+    if not config.enabled or not config.tools:
+        return result
+
+    # Collect sanitized tool names
+    sanitized_tools: list[tuple[str, PTCToolDef]] = []
+    for tool in config.tools:
+        name = tool.get("name", "")
+        sanitized_name = sanitize_function_name(name)
+        sanitized_tools.append((sanitized_name, tool))
+
+    # Sort for deterministic output
+    sanitized_tools.sort(key=lambda x: x[0])
+    result.tool_names = [name for name, _ in sanitized_tools]
+
+    validated_tool_configs = _validate_tool_configs(tool_configs, config.tools)
+
+    # Generate tools/custom/__init__.py
+    result.payload.files["tools/custom/__init__.py"] = _generate_custom_init(sanitized_tools)
+
+    # Generate tools/custom/<tool_name>.py for each tool
+    for sanitized_name, tool in sanitized_tools:
+        wrapper_content = _generate_tool_wrapper(sanitized_name, tool)
+        result.payload.files[f"tools/custom/{sanitized_name}.py"] = wrapper_content
+
+    # Generate tools/custom_sources/__init__.py
+    file_tools = [(name, tool) for name, tool in sanitized_tools if tool.get("kind") == "file"]
+    result.payload.files["tools/custom_sources/__init__.py"] = _generate_custom_sources_init(file_tools)
+
+    # Generate tools/custom_sources/<tool_name>/__init__.py for file tools
+    for sanitized_name, tool in file_tools:
+        source_content = _read_file_tool_source(tool)
+        result.payload.files[f"tools/custom_sources/{sanitized_name}/__init__.py"] = source_content
+
+    # Generate tools/custom_registry.py
+    result.payload.files["tools/custom_registry.py"] = _generate_registry(sanitized_tools)
+
+    # Generate tools/custom_invoke.py
+    result.payload.files["tools/custom_invoke.py"] = _generate_invoke_helper()
+
+    # Generate tools/custom_defaults.json (per-run file, always re-uploaded).
+    # Key the JSON by the sanitized tool name (matching the name baked into
+    # the generated wrapper and ``build_runtime_config`` lookup) so tools
+    # configured with hyphens or other non-Python-identifier characters
+    # still resolve to their defaults at runtime.
+    defaults = {sanitize_function_name(name): value for name, value in validated_tool_configs.items()}
+    result.payload.per_run_files["tools/custom_defaults.json"] = json.dumps(defaults, indent=2)
+
+    # Bundle package sources from package_path
+    for sanitized_name, tool in sanitized_tools:
+        if tool.get("kind") == "package":
+            package_path = tool.get("package_path")
+            if package_path:
+                _bundle_package_source(result.payload, package_path, config.bundle_roots)
+
+    return result
+
+
+def _generate_custom_init(sanitized_tools: list[tuple[str, PTCToolDef]]) -> str:
+    """Generate tools/custom/__init__.py content.
+
+    Args:
+        sanitized_tools: List of (sanitized_name, tool_def) tuples.
+
+    Returns:
+        Python source code for the __init__.py module.
+    """
+    # Sort for deterministic output
+    sorted_names = sorted(name for name, _ in sanitized_tools)
+
+    all_list = ", ".join(f'"{name}"' for name in sorted_names)
+
+    return render_template(
+        _TEMPLATE_PACKAGE,
+        "custom_init.py.template",
+        {
+            "all_list": all_list,
+        },
+    )
+
+
+def _generate_tool_wrapper(sanitized_name: str, tool: PTCToolDef) -> str:
+    """Generate tools/custom/<tool_name>.py wrapper module.
+
+    Args:
+        sanitized_name: Sanitized tool name for function/module.
+        tool: Tool definition.
+
+    Returns:
+        Python source code for the wrapper module.
+    """
+    original_name = tool.get("name", sanitized_name)
+    schema = tool.get("input_schema", {"type": "object", "properties": {}})
+    projection = project_custom_tool_signature(schema)
+    required_params = tuple(field.param_name for field in projection.field_specs if field.required)
+    optional_params = tuple(field.param_name for field in projection.field_specs if not field.required)
+
+    return render_template(
+        _TEMPLATE_PACKAGE,
+        "custom_wrapper.py.template",
+        {
+            "original_name": original_name,
+            "sanitized_name": sanitized_name,
+            "canonical_params": projection.canonical_params,
+            "schema_to_param": repr(projection.schema_to_param),
+            "param_to_schema": repr(projection.param_to_schema),
+            "required_params": repr(required_params),
+            "optional_params": repr(optional_params),
+            "runtime_signature_params": _render_runtime_signature_params(projection.runtime_params),
+        },
+    )
+
+
+def _is_safe_annotation(expr: str) -> bool:
+    """Check if annotation expression is safe.
+
+    Safe annotations are:
+    - Single safe type: str, int, float, bool, list, dict, Any, None
+    - Union of safe types using | (e.g., str | None, str | int, int | str | None)
+
+    This validates the ENTIRE expression, not just a prefix, to prevent
+    bypasses like "str | __import__('os').system('x')".
+    """
+    if expr in _SAFE_ANNOTATIONS:
+        return True
+
+    parts = [part.strip() for part in expr.split("|")]
+    if len(parts) <= 1 or any(not part for part in parts):
+        return False
+
+    return all(part in _SAFE_ANNOTATIONS for part in parts)
+
+
+def _sanitize_annotation(param: CustomToolRuntimeParamSpec) -> str:
+    """Return safe annotation expression for a runtime parameter."""
+    annotation = param.annotation_expr
+    if _is_safe_annotation(annotation):
+        return annotation
+
+    logger.warning(f"Potentially unsafe annotation '{annotation}' in parameter '{param.name}', using 'Any'")
+    return "Any"
+
+
+def _is_safe_numeric_literal(default_expr: str) -> bool:
+    """Return True if expression is a safe signed integer or float literal."""
+    stripped = default_expr.replace("-", "", 1).replace("+", "", 1)
+    return stripped.count(".") <= 1 and stripped.replace(".", "").isdigit()
+
+
+def _sanitize_default_fragment(param: CustomToolRuntimeParamSpec) -> str:
+    """Return rendered default fragment for inspect.Parameter declaration."""
+    if not param.has_default:
+        return ""
+
+    default_expr = param.default_expr if param.default_expr is not None else "None"
+    if default_expr not in ("None", "True", "False") and not _is_safe_numeric_literal(default_expr):
+        logger.warning(f"Potentially unsafe default '{default_expr}' in parameter '{param.name}', using 'None'")
+        default_expr = "None"
+
+    return f", default={default_expr}"
+
+
+def _render_runtime_signature_params(runtime_params: list[CustomToolRuntimeParamSpec]) -> str:
+    """Render inspect.Parameter entries for generated wrapper signatures.
+
+    Sanitizes annotation and default expressions to prevent code injection.
+    Unsafe values are replaced with safe defaults and logged.
+    """
+    rendered_params: list[str] = []
+
+    for param in runtime_params:
+        annotation = _sanitize_annotation(param)
+        default_fragment = _sanitize_default_fragment(param)
+
+        rendered_params.append(
+            "inspect.Parameter("
+            f'"{param.name}", '
+            f"inspect.Parameter.{param.kind}, "
+            f"annotation={annotation}{default_fragment}"
+            ")"
+        )
+
+    return ",\n        ".join(rendered_params)
+
+
+def _generate_custom_sources_init(file_tools: list[tuple[str, PTCToolDef]]) -> str:
+    """Generate tools/custom_sources/__init__.py content.
+
+    Args:
+        file_tools: List of (sanitized_name, tool_def) tuples for file tools.
+
+    Returns:
+        Python source code for the __init__.py module.
+    """
+    sorted_names = sorted(name for name, _ in file_tools)
+    all_list = ", ".join(f'"{name}"' for name in sorted_names)
+
+    return render_template(
+        _TEMPLATE_PACKAGE,
+        "custom_sources_init.py.template",
+        {"all_list": all_list},
+    )
+
+
+def _read_file_tool_source(tool: PTCToolDef) -> str:
+    """Read source code from a file tool.
+
+    Args:
+        tool: File tool definition with file_path.
+
+    Returns:
+        Source code content.
+
+    Raises:
+        FileNotFoundError: If file does not exist.
+    """
+    file_path = tool.get("file_path", "")
+    path = Path(file_path)
+    return path.read_text(encoding="utf-8")
+
+
+def _generate_registry(sanitized_tools: list[tuple[str, PTCToolDef]]) -> str:
+    """Generate tools/custom_registry.py content.
+
+    Args:
+        sanitized_tools: List of (sanitized_name, tool_def) tuples.
+
+    Returns:
+        Python source code for the registry module.
+    """
+    build_functions: list[str] = []
+
+    for sanitized_name, tool in sorted(sanitized_tools, key=lambda x: x[0]):
+        kind = tool.get("kind")
+        class_name = tool.get("class_name", "")
+
+        if kind == "package":
+            import_path = tool.get("import_path", "")
+            import_statement = f"from {import_path} import {class_name}"
+        elif kind == "file":
+            import_statement = f"from tools.custom_sources.{sanitized_name} import {class_name}"
+        else:
+            import_statement = ""
+
+        # Generate build function
+        build_func = _generate_build_function(sanitized_name, class_name, import_statement)
+        build_functions.append(build_func)
+
+    build_functions_str = "\n\n".join(build_functions)
+
+    return render_template(
+        _TEMPLATE_PACKAGE,
+        "custom_registry.py.template",
+        {
+            "build_functions": build_functions_str,
+        },
+    )
+
+
+def _generate_build_function(sanitized_name: str, class_name: str, import_statement: str = "") -> str:
+    """Generate a build function for a single tool.
+
+    Args:
+        sanitized_name: Sanitized tool name.
+        class_name: Tool class name.
+        import_statement: Import statement for the tool class, rendered inside the build function.
+
+    Returns:
+        Python source code for the build function.
+    """
+    return render_template(
+        _TEMPLATE_PACKAGE,
+        "custom_build_function.py.template",
+        {
+            "sanitized_name": sanitized_name,
+            "class_name": class_name,
+            "import_statement": import_statement,
+        },
+    )
+
+
+def _generate_invoke_helper() -> str:
+    """Generate tools/custom_invoke.py content.
+
+    The structured-output constants are injected from the host-side executor so the
+    sandbox helper and the host reader/cleanup cannot drift apart.
+    """
+    from aip_agents.ptc.executor import STRUCTURED_OUTPUT_DIR, STRUCTURED_OUTPUT_ENV
+
+    return render_template(
+        _TEMPLATE_PACKAGE,
+        "custom_invoke.py.template",
+        {
+            "structured_output_env": STRUCTURED_OUTPUT_ENV,
+            "structured_output_dir": STRUCTURED_OUTPUT_DIR,
+        },
+    )
+
+
+def _bundle_package_source(
+    payload: SandboxPayload,
+    package_path: str,
+    bundle_roots: list[str],
+) -> None:
+    """Bundle package source files into payload.
+
+    Copies Python files from package_path into the payload, preserving
+    the directory structure relative to the source root.
+
+    For standard layouts, files are relative to the bundle root.
+    For src/ layouts, files are relative to the src/ directory so that
+    imports work correctly (e.g., `from my_pkg import ...`).
+
+    Args:
+        payload: Sandbox payload to add files to.
+        package_path: Path to the package directory.
+        bundle_roots: List of allowed bundle roots.
+    """
+    pkg_path = Path(package_path).resolve()
+
+    # Find which bundle root contains this path
+    bundle_root = None
+    for root in bundle_roots:
+        root_path = Path(root).resolve()
+        try:
+            pkg_path.relative_to(root_path)
+            bundle_root = root_path
+            break
+        except ValueError:
+            continue
+
+    if bundle_root is None:
+        # Path validation should have caught this earlier
+        return
+
+    # Check for src/ subdirectory (common package layout)
+    # When src/ exists, use it as the base for relative paths so imports work
+    src_path = pkg_path / "src"
+    if src_path.is_dir():
+        # For src/ layout: copy from src/ and calculate paths relative to src/
+        # This ensures `from my_pkg import ...` works when packages/ is on sys.path
+        _copy_python_files(payload, src_path, src_path)
+    else:
+        # For flat layout: copy from pkg_path relative to bundle_root
+        _copy_python_files(payload, pkg_path, bundle_root)
+
+
+def _copy_python_files(
+    payload: SandboxPayload,
+    source_dir: Path,
+    bundle_root: Path,
+) -> None:
+    """Copy Python files from source directory to payload.
+
+    Args:
+        payload: Sandbox payload to add files to.
+        source_dir: Source directory to copy from.
+        bundle_root: Bundle root for relative path calculation.
+    """
+    for py_file in source_dir.rglob("*.py"):
+        # Skip __pycache__ and similar
+        if "__pycache__" in py_file.parts:
+            continue
+        if ".egg-info" in str(py_file):
+            continue
+
+        # Calculate relative path from bundle root
+        try:
+            rel_path = py_file.relative_to(bundle_root)
+        except ValueError:
+            # File is not under bundle root
+            continue
+
+        # Read and add to payload
+        try:
+            content = py_file.read_text(encoding="utf-8")
+            # Use packages/ prefix for bundled sources
+            payload_path = f"packages/{rel_path.as_posix()}"
+            payload.files[payload_path] = content
+        except Exception as exc:
+            logger.warning(f"Skipping unreadable file '{py_file}': {exc}")
+            continue

aip_agents/ptc/custom_tools_templates/__init__.py

@@ -0,0 +1 @@
+"""Templates for custom LangChain tool payload generation."""

aip_agents/ptc/custom_tools_templates/custom_build_function.py.template

@@ -0,0 +1,24 @@
+def build_$sanitized_name() -> BaseTool:
+    """Build an instance of $class_name.
+
+    Returns:
+        Configured tool instance.
+    """
+    $import_statement
+    # Check if the class is already an instance (module-level singleton)
+    tool_or_class = $class_name
+    if isinstance(tool_or_class, BaseTool):
+        tool = tool_or_class
+    else:
+        try:
+            tool = tool_or_class()
+        except TypeError as e:
+            raise TypeError(
+                f"Failed to instantiate $class_name(). "
+                "Constructor args are not supported in MVP. "
+                "Use tool_config_schema for runtime configuration. "
+                f"Original error: {e}"
+            ) from e
+
+    _configure_tool(tool, "$sanitized_name")
+    return tool

aip_agents/ptc/custom_tools_templates/custom_init.py.template

@@ -0,0 +1,26 @@
+"""Generated custom tools package for PTC sandbox execution.
+
+This package provides access to custom LangChain tools configured for this agent.
+Import tools directly:
+
+    from tools.custom import tool_name
+
+Or from specific modules:
+
+    from tools.custom.tool_name import tool_name
+"""
+
+from importlib import import_module
+from typing import Any
+
+__all__ = [$all_list]
+
+
+def __getattr__(name: str) -> Any:
+    """Lazy-load custom tool wrapper functions on first access."""
+    if name in __all__:
+        module = import_module(f"tools.custom.{name}")
+        tool = getattr(module, name)
+        globals()[name] = tool
+        return tool
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

aip_agents/ptc/custom_tools_templates/custom_invoke.py.template

@@ -0,0 +1,149 @@
+"""Safe invoke helper for custom tools.
+
+This module provides a sync wrapper around tool.ainvoke for use in PTC code.
+"""
+
+import asyncio
+import json
+import os
+import sys
+import threading
+from typing import Any
+
+from langchain_core.tools import BaseTool
+
+try:
+    from langgraph.types import Command
+except Exception:  # pragma: no cover - sandbox dependency guard
+    Command = None
+
+_STRUCTURED_OUTPUT_ENV = "$structured_output_env"
+_STRUCTURED_OUTPUT_DIR = "$structured_output_dir"
+_STRUCTURED_OUTPUT_LOCK = threading.Lock()
+_SUPPORTED_UPDATE_FIELDS = ("metadata", "artifacts", "pii_mapping", "usage_metadata")
+_INTENTIONALLY_LOCAL_FIELDS = ("result",)
+
+
+def run_tool(tool: BaseTool, kwargs: dict[str, Any], config: dict[str, Any] | None = None) -> Any:
+    """Run a tool's ainvoke method synchronously.
+
+    Args:
+        tool: Tool instance to invoke.
+        kwargs: Arguments to pass to the tool.
+        config: Optional runtime config (e.g. {"metadata": ...}) forwarded
+            to tool.ainvoke.
+
+    Returns:
+        Tool execution result.
+    """
+    result = _run_async_safely(tool.ainvoke(kwargs, config=config))
+    _write_command_update(result)
+    return result
+
+
+def _write_command_update(result: Any) -> None:
+    """Write supported Command.update fields to the host-provided JSONL side channel."""
+    if Command is None or not isinstance(result, Command):
+        return
+
+    update = getattr(result, "update", None)
+    if not isinstance(update, dict):
+        return
+
+    unsupported_fields = sorted(set(update) - set(_SUPPORTED_UPDATE_FIELDS) - set(_INTENTIONALLY_LOCAL_FIELDS))
+    if unsupported_fields:
+        fields = ", ".join(unsupported_fields)
+        print(
+            f"Warning: PTC custom tool Command.update fields are not supported and were not bridged: {fields}",
+            file=sys.stderr,
+        )
+
+    serializable_update: dict[str, Any] = {}
+    for field in _SUPPORTED_UPDATE_FIELDS:
+        if field not in update:
+            continue
+        try:
+            json.dumps(update[field])
+        except (TypeError, ValueError):
+            print(
+                f"Warning: PTC custom tool Command.update field is not JSON-serializable and was not bridged: {field}",
+                file=sys.stderr,
+            )
+            continue
+        serializable_update[field] = update[field]
+
+    if not serializable_update:
+        return
+
+    path = _get_structured_output_path()
+    if not path:
+        return
+
+    try:
+        # E2B's command API may read/cleanup from a different sandbox user than
+        # execute_code. Keep the side-channel non-world-writable while allowing
+        # that command plane to traverse/read the executor-generated file.
+        os.makedirs(_STRUCTURED_OUTPUT_DIR, mode=0o755, exist_ok=True)
+        os.chmod(_STRUCTURED_OUTPUT_DIR, 0o755)
+        line = json.dumps(serializable_update, separators=(",", ":")) + "\n"
+        with _STRUCTURED_OUTPUT_LOCK:
+            with open(path, "a", encoding="utf-8") as file:
+                file.write(line)
+            os.chmod(path, 0o644)
+    except OSError as exc:
+        print(f"Warning: failed to write PTC structured output side-channel: {exc}", file=sys.stderr)
+
+
+def _get_structured_output_path() -> str | None:
+    """Return the current host-provided side-channel path when it is reserved for PTC."""
+    path = os.environ.get(_STRUCTURED_OUTPUT_ENV)
+    if not path:
+        return None
+    normalized_path = os.path.normpath(path)
+    if os.path.dirname(normalized_path) != _STRUCTURED_OUTPUT_DIR:
+        return None
+    name = os.path.basename(normalized_path)
+    if not name.endswith(".jsonl"):
+        return None
+    stem = name[:-len(".jsonl")]
+    if len(stem) != 32 or any(char not in "0123456789abcdef" for char in stem):
+        return None
+    return normalized_path
+
+
+def _run_async_safely(coro: Any) -> Any:
+    """Run an async coroutine safely in sync context.
+
+    Handles the case where we may or may not be in an async context.
+
+    Args:
+        coro: Coroutine to run.
+
+    Returns:
+        Result of the coroutine.
+    """
+    try:
+        asyncio.get_running_loop()
+    except RuntimeError:
+        # No running loop - create new one
+        return asyncio.run(coro)
+
+    # Already in async context - run in a separate thread with a new loop
+    import concurrent.futures
+
+    def run_in_new_loop() -> Any:
+        """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()