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

aip_agents/ptc/custom_tools.py

@@ -0,0 +1,511 @@
+"""Custom LangChain tools configuration for PTC sandbox.
+
+This module provides configuration classes for defining custom LangChain tools
+that can be used inside the PTC sandbox alongside MCP tools.
+
+Authors:
+    Putu Ravindra Wiguna (putu.r.wiguna@gdplabs.id)
+"""
+
+from __future__ import annotations
+
+import ast
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Literal, TypedDict
+
+from aip_agents.ptc.exceptions import PTCError
+from aip_agents.ptc.naming import (
+    is_valid_identifier,
+    project_custom_tool_signature,
+    sanitize_function_name,
+)
+from aip_agents.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class PTCPackageToolDef(TypedDict, total=False):
+    """Definition for a package-based tool.
+
+    Required fields:
+        name: Tool name (used for imports and config lookup).
+        kind: Must be "package".
+        import_path: Dotted import path (e.g., "aip_agents.tools.time_tool").
+        class_name: Class name to import (e.g., "TimeTool").
+
+    Optional fields:
+        package_path: Path to source directory to bundle.
+            If omitted, the tool must already be available in sandbox (e.g., from site-packages).
+        description: Tool description (derived from tool object at construction time).
+        input_schema: JSON schema for tool input (derived from tool object at construction time).
+    """
+
+    name: str
+    kind: Literal["package"]
+    import_path: str
+    class_name: str
+    package_path: str  # Optional
+    description: str  # Optional - derived from tool object
+    input_schema: dict  # Optional - derived from tool object
+
+
+class PTCFileToolDef(TypedDict, total=False):
+    """Definition for a single-file tool.
+
+    Required fields:
+        name: Tool name (used for imports and config lookup).
+        kind: Must be "file".
+        file_path: Path to the Python file containing the tool.
+        class_name: Class name to import from the file.
+
+    Optional fields:
+        description: Tool description (derived from tool object at construction time).
+        input_schema: JSON schema for tool input (derived from tool object at construction time).
+    """
+
+    name: str
+    kind: Literal["file"]
+    file_path: str
+    class_name: str
+    description: str  # Optional - derived from tool object
+    input_schema: dict  # Optional - derived from tool object
+
+
+# Union type for tool definitions
+PTCToolDef = PTCPackageToolDef | PTCFileToolDef
+
+
+class PTCCustomToolValidationError(PTCError):
+    """Error raised when custom tool configuration is invalid."""
+
+    pass
+
+
+@dataclass
+class PTCCustomToolConfig:
+    """Configuration for custom LangChain tools in PTC sandbox.
+
+    This config controls which user-defined LangChain tools are available
+    in the sandbox for use by PTC code.
+
+    Attributes:
+        enabled: Whether custom tools are enabled.
+        bundle_roots: List of allowed root directories for bundling tool sources.
+            Paths outside these roots will cause validation errors.
+        requirements: List of pip requirements to install in sandbox.
+        tools: List of tool definitions (package or file tools).
+
+    Example:
+        >>> config = PTCCustomToolConfig(
+        ...     enabled=True,
+        ...     bundle_roots=["/app/tools"],
+        ...     requirements=["pydantic>=2.0"],
+        ...     tools=[
+        ...         {
+        ...             "name": "time_tool",
+        ...             "kind": "package",
+        ...             "import_path": "aip_agents.tools.time_tool",
+        ...             "class_name": "TimeTool",
+        ...         },
+        ...     ],
+        ... )
+    """
+
+    enabled: bool = False
+    bundle_roots: list[str] = field(default_factory=list)
+    requirements: list[str] = field(default_factory=list)
+    tools: list[PTCToolDef] = field(default_factory=list)
+
+
+def _validate_class_name(class_name: str, tool_name: str) -> None:
+    """Validate a class name used in generated code.
+
+    Args:
+        class_name: Class name to validate.
+        tool_name: Tool name for error messages.
+
+    Raises:
+        PTCCustomToolValidationError: If class name is invalid.
+    """
+    if not is_valid_identifier(class_name):
+        raise PTCCustomToolValidationError(
+            f"Tool '{tool_name}' has invalid class_name '{class_name}'. Must be a valid Python identifier."
+        )
+
+
+def _validate_package_tool(tool: PTCToolDef, name: str) -> None:
+    """Validate a package tool definition.
+
+    Args:
+        tool: Package tool definition to validate.
+        name: Tool name for error messages.
+
+    Raises:
+        PTCCustomToolValidationError: If package tool definition is invalid.
+    """
+    import_path = tool.get("import_path", "")
+    if not import_path:
+        raise PTCCustomToolValidationError(f"Package tool '{name}' missing required 'import_path' field")
+    if import_path.startswith("."):
+        raise PTCCustomToolValidationError(f"Package tool '{name}' import_path must be absolute: '{import_path}'")
+    if any(not part.isidentifier() for part in import_path.split(".")):
+        raise PTCCustomToolValidationError(f"Package tool '{name}' has invalid import_path: '{import_path}'")
+    class_name = tool.get("class_name", "")
+    if not class_name:
+        raise PTCCustomToolValidationError(f"Package tool '{name}' missing required 'class_name' field")
+    _validate_class_name(class_name, name)
+
+
+def _validate_file_tool(tool: PTCToolDef, name: str) -> None:
+    """Validate a file tool definition.
+
+    Args:
+        tool: File tool definition to validate.
+        name: Tool name for error messages.
+
+    Raises:
+        PTCCustomToolValidationError: If file tool definition is invalid.
+    """
+    if not tool.get("file_path"):
+        raise PTCCustomToolValidationError(f"File tool '{name}' missing required 'file_path' field")
+    class_name = tool.get("class_name", "")
+    if not class_name:
+        raise PTCCustomToolValidationError(f"File tool '{name}' missing required 'class_name' field")
+    _validate_class_name(class_name, name)
+
+
+def validate_tool_def(tool: PTCToolDef) -> None:
+    """Validate a tool definition has required fields.
+
+    Args:
+        tool: Tool definition to validate.
+
+    Raises:
+        PTCCustomToolValidationError: If tool definition is invalid.
+    """
+    name = tool.get("name")
+    kind = tool.get("kind")
+
+    if not name:
+        raise PTCCustomToolValidationError("Tool definition missing required 'name' field")
+
+    if kind not in ("package", "file"):
+        raise PTCCustomToolValidationError(f"Tool '{name}' has invalid kind '{kind}'. Must be 'package' or 'file'.")
+
+    if kind == "package":
+        _validate_package_tool(tool, name)
+    elif kind == "file":
+        _validate_file_tool(tool, name)
+
+
+def _validate_input_schema(schema: dict, tool_name: str) -> None:
+    """Validate input_schema by projecting it into a canonical signature.
+
+    This provides fail-fast validation at config time for:
+    - Parameter sanitization collisions (e.g., 'user-id' and 'user_id' map to same name)
+    - Reserved name collisions (e.g., schema key sanitizing to 'kwargs')
+
+    Args:
+        schema: The input_schema JSON schema to validate.
+        tool_name: Tool name for error messages.
+
+    Raises:
+        PTCCustomToolValidationError: If input_schema has invalid parameter names.
+    """
+    try:
+        project_custom_tool_signature(schema)
+    except ValueError as e:
+        raise PTCCustomToolValidationError(f"Tool '{tool_name}' has invalid input_schema: {e}") from e
+
+
+def validate_path_within_bundle_roots(
+    path: str | Path,
+    bundle_roots: list[str],
+    tool_name: str,
+) -> None:
+    """Validate that a path is within one of the bundle roots.
+
+    Args:
+        path: Path to validate.
+        bundle_roots: List of allowed root directories.
+        tool_name: Tool name for error messages.
+
+    Raises:
+        PTCCustomToolValidationError: If path is outside all bundle roots.
+    """
+    if not bundle_roots:
+        raise PTCCustomToolValidationError(
+            f"Tool '{tool_name}' requires bundle_roots to be set when using package_path or file_path"
+        )
+
+    # Resolve real path to prevent symlink escapes
+    real_path = Path(path).resolve()
+
+    for root in bundle_roots:
+        real_root = Path(root).resolve()
+        try:
+            real_path.relative_to(real_root)
+            return  # Path is within this root
+        except ValueError:
+            continue
+
+    raise PTCCustomToolValidationError(
+        f"Tool '{tool_name}' path '{path}' is outside all bundle_roots. Allowed roots: {bundle_roots}"
+    )
+
+
+def detect_relative_imports(file_path: str | Path) -> list[str]:
+    """Detect relative imports in a Python file using AST.
+
+    Args:
+        file_path: Path to the Python file to scan.
+
+    Returns:
+        List of relative import statements found.
+
+    Raises:
+        PTCCustomToolValidationError: If file cannot be read or parsed.
+    """
+    path = Path(file_path)
+    if not path.exists():
+        raise PTCCustomToolValidationError(f"File not found: {file_path}")
+
+    try:
+        source = path.read_text(encoding="utf-8")
+    except Exception as exc:
+        raise PTCCustomToolValidationError(f"Cannot read file '{file_path}': {exc}") from exc
+
+    try:
+        tree = ast.parse(source, filename=str(path))
+    except SyntaxError as exc:
+        raise PTCCustomToolValidationError(f"Syntax error in file '{file_path}': {exc}") from exc
+
+    relative_imports: list[str] = []
+    for node in ast.walk(tree):
+        if isinstance(node, ast.ImportFrom):
+            # level > 0 indicates relative import (e.g., from . import X, from ..foo import Y)
+            if node.level > 0:
+                module = node.module or ""
+                dots = "." * node.level
+                import_str = f"from {dots}{module} import ..."
+                relative_imports.append(import_str)
+
+    return relative_imports
+
+
+def check_name_collisions(tools: list[PTCToolDef]) -> None:
+    """Check for tool name collisions after sanitization.
+
+    Args:
+        tools: List of tool definitions.
+
+    Raises:
+        PTCCustomToolValidationError: If two tools sanitize to the same name.
+    """
+    sanitized_to_original: dict[str, list[str]] = {}
+
+    for tool in tools:
+        name = tool.get("name", "")
+        sanitized = sanitize_function_name(name)
+
+        if sanitized not in sanitized_to_original:
+            sanitized_to_original[sanitized] = []
+        sanitized_to_original[sanitized].append(name)
+
+    # Check for collisions
+    for sanitized, originals in sanitized_to_original.items():
+        if len(originals) > 1:
+            raise PTCCustomToolValidationError(
+                f"Tool name collision: tools {originals} all sanitize to '{sanitized}'. Please use unique tool names."
+            )
+
+
+def _validate_tool_paths(tool: dict, bundle_roots: list[str]) -> None:
+    """Validate paths for a single tool and check for relative imports.
+
+    Args:
+        tool: Tool definition dict with kind, name, and path fields.
+        bundle_roots: List of allowed bundle root directories.
+
+    Raises:
+        PTCCustomToolValidationError: If paths are invalid or file uses relative imports.
+    """
+    kind = tool.get("kind")
+    name = tool.get("name", "")
+
+    if kind == "package":
+        package_path = tool.get("package_path")
+        if package_path:
+            path = Path(package_path)
+            if not path.exists() or not path.is_dir():
+                raise PTCCustomToolValidationError(
+                    f"Package tool '{name}' package_path does not exist or is not a directory: '{package_path}'"
+                )
+            validate_path_within_bundle_roots(package_path, bundle_roots, name)
+
+    elif kind == "file":
+        file_path = tool.get("file_path")
+        if file_path:
+            validate_path_within_bundle_roots(file_path, bundle_roots, name)
+
+            relative_imports = detect_relative_imports(file_path)
+            if relative_imports:
+                raise PTCCustomToolValidationError(
+                    f"File tool '{name}' uses relative imports which are not supported: "
+                    f"{relative_imports}. Use a package tool with package_path instead."
+                )
+
+
+def validate_custom_tool_config(config: PTCCustomToolConfig) -> None:
+    """Validate the complete custom tool configuration.
+
+    This runs all validation checks:
+    - Tool definitions have required fields
+    - Input schemas have valid parameter names (no collisions/reserved names)
+    - Paths are within bundle roots
+    - File tools don't use relative imports
+    - No name collisions
+
+    Args:
+        config: Custom tool configuration to validate.
+
+    Raises:
+        PTCCustomToolValidationError: If configuration is invalid.
+    """
+    if not config.enabled or not config.tools:
+        return
+
+    # Validate each tool definition
+    for tool in config.tools:
+        validate_tool_def(tool)
+
+    # Validate input_schema for each tool (fail-fast for param name issues)
+    for tool in config.tools:
+        if "input_schema" not in tool:
+            continue
+
+        name = tool.get("name", "")
+        raw_schema = tool.get("input_schema")
+        if not isinstance(raw_schema, dict):
+            raise PTCCustomToolValidationError(
+                f"Tool '{name}' has invalid input_schema: expected dict, got {type(raw_schema).__name__}"
+            )
+        _validate_input_schema(raw_schema, name)
+
+    # Check for name collisions
+    check_name_collisions(config.tools)
+
+    # Validate paths and check for relative imports
+    for tool in config.tools:
+        _validate_tool_paths(tool, config.bundle_roots)
+
+
+def extract_tool_metadata(tool: object) -> dict:
+    """Extract metadata from a LangChain BaseTool instance.
+
+    Extracts name, description, and input_schema from the tool object.
+    This is called at agent construction time to populate tool definitions.
+
+    Args:
+        tool: A LangChain BaseTool instance (or compatible object).
+
+    Returns:
+        Dict with keys:
+        - name: str (tool name)
+        - description: str (tool description, empty if not available)
+        - input_schema: dict (JSON schema, empty object schema if not available)
+
+    Note:
+        If schema inference fails, returns an empty object schema
+        ({"type": "object", "properties": {}}) so prompt/index output
+        falls back to tool(**kwargs).
+    """
+    result: dict = {
+        "name": "",
+        "description": "",
+        "input_schema": {"type": "object", "properties": {}},
+    }
+
+    # Extract name
+    if hasattr(tool, "name"):
+        result["name"] = str(tool.name)
+
+    # Extract description
+    if hasattr(tool, "description"):
+        desc = tool.description
+        if desc:
+            # Clean up multiline description - normalize whitespace
+            result["description"] = " ".join(str(desc).split())
+
+    # Extract input_schema from args_schema (Pydantic model)
+    if hasattr(tool, "args_schema") and tool.args_schema is not None:
+        try:
+            args_schema = tool.args_schema
+            # Pydantic v2 uses model_json_schema()
+            if hasattr(args_schema, "model_json_schema"):
+                schema = args_schema.model_json_schema()
+                # Remove Pydantic-specific keys that aren't needed
+                schema.pop("title", None)
+                schema.pop("$defs", None)
+                schema.pop("definitions", None)
+                result["input_schema"] = schema
+            # Pydantic v1 fallback uses schema()
+            elif hasattr(args_schema, "schema"):
+                schema = args_schema.schema()
+                schema.pop("title", None)
+                schema.pop("definitions", None)
+                result["input_schema"] = schema
+        except Exception:
+            # If schema extraction fails, keep the empty object schema
+            pass
+
+    return result
+
+
+def enrich_tool_def_with_metadata(tool_def: PTCToolDef, tool: object) -> PTCToolDef:
+    """Enrich a tool definition with metadata extracted from the tool object.
+
+    Always derives description and input_schema from the tool object, overwriting
+    any user-supplied values. This ensures the tool object is the source of truth.
+    Called at agent construction time.
+
+    Args:
+        tool_def: The tool definition to enrich.
+        tool: The LangChain BaseTool instance.
+
+    Returns:
+        The tool definition with description and input_schema populated from the tool object.
+
+    Raises:
+        PTCCustomToolValidationError: If tool name in metadata does not match tool_def name.
+    """
+    metadata = extract_tool_metadata(tool)
+
+    # Create a mutable copy
+    enriched: dict = dict(tool_def)
+
+    # Validate name matches if metadata provides one
+    if metadata["name"] and metadata["name"] != enriched.get("name"):
+        raise PTCCustomToolValidationError(
+            f"Tool name mismatch: tool_def name '{enriched.get('name')}' does not match "
+            f"tool object name '{metadata['name']}'"
+        )
+
+    # Log if we're overwriting existing values
+    if "description" in enriched and enriched["description"]:
+        logger.debug(
+            f"Overwriting user-supplied description for tool '{enriched.get('name', 'unknown')}' "
+            f"with derived metadata from tool object"
+        )
+    if "input_schema" in enriched and enriched["input_schema"]:
+        logger.debug(
+            f"Overwriting user-supplied input_schema for tool '{enriched.get('name', 'unknown')}' "
+            f"with derived metadata from tool object"
+        )
+
+    # Always overwrite with derived metadata (tool object is source of truth)
+    enriched["description"] = metadata["description"]
+    enriched["input_schema"] = metadata["input_schema"]
+
+    return enriched  # type: ignore[return-value]