Custom Tools with Remote Agent Server

"""Example: Using custom tools with remote agent server.

This example demonstrates how to use custom tools with a remote agent server
by building a custom base image that includes the tool implementation.

Prerequisites:
    1. Build the custom base image first:
       cd examples/02_remote_agent_server/05_custom_tool
       ./build_custom_image.sh

    2. Set LLM_API_KEY environment variable

The workflow is:
1. Define a custom tool (LogDataTool for logging structured data to JSON)
2. Create a simple Dockerfile that copies the tool into the base image
3. Build the custom base image
4. Use DockerDevWorkspace with base_image pointing to the custom image
5. DockerDevWorkspace builds the agent server on top of the custom base image
6. The server dynamically registers tools when the client creates a conversation
7. The agent can use the custom tool during execution
8. Verify the logged data by reading the JSON file from the workspace

This pattern is useful for:
- Collecting structured data during agent runs (logs, metrics, events)
- Implementing custom integrations with external systems
- Adding domain-specific operations to the agent
"""

import os
import platform
import subprocess
import sys
import time
from pathlib import Path

from pydantic import SecretStr

from openhands.sdk import (
    LLM,
    Conversation,
    RemoteConversation,
    Tool,
    get_logger,
)
from openhands.workspace import DockerDevWorkspace


logger = get_logger(__name__)

# 1) Ensure we have LLM API key
api_key = os.getenv("LLM_API_KEY")
assert api_key is not None, "LLM_API_KEY environment variable is not set."

llm = LLM(
    usage_id="agent",
    model=os.getenv("LLM_MODEL", "anthropic/claude-sonnet-4-5-20250929"),
    base_url=os.getenv("LLM_BASE_URL"),
    api_key=SecretStr(api_key),
)


def detect_platform():
    """Detects the correct Docker platform string."""
    machine = platform.machine().lower()
    if "arm" in machine or "aarch64" in machine:
        return "linux/arm64"
    return "linux/amd64"


# Get the directory containing this script
example_dir = Path(__file__).parent.absolute()

# Custom base image tag (contains custom tools, agent server built on top)
CUSTOM_BASE_IMAGE_TAG = "custom-base-image:latest"

# 2) Check if custom base image exists, build if not
logger.info(f"🔍 Checking for custom base image: {CUSTOM_BASE_IMAGE_TAG}")
result = subprocess.run(
    ["docker", "images", "-q", CUSTOM_BASE_IMAGE_TAG],
    capture_output=True,
    text=True,
    check=False,
)

if not result.stdout.strip():
    logger.info("⚠️  Custom base image not found. Building...")
    logger.info("📦 Building custom base image with custom tools...")
    build_script = example_dir / "build_custom_image.sh"
    try:
        subprocess.run(
            [str(build_script), CUSTOM_BASE_IMAGE_TAG],
            cwd=str(example_dir),
            check=True,
        )
        logger.info("✅ Custom base image built successfully!")
    except subprocess.CalledProcessError as e:
        logger.error(f"❌ Failed to build custom base image: {e}")
        logger.error("Please run ./build_custom_image.sh manually and fix any errors.")
        sys.exit(1)
else:
    logger.info(f"✅ Custom base image found: {CUSTOM_BASE_IMAGE_TAG}")

# 3) Create a DockerDevWorkspace with the custom base image
#    DockerDevWorkspace will build the agent server on top of this base image
logger.info("🚀 Building and starting agent server with custom tools...")
logger.info("📦 This may take a few minutes on first run...")

with DockerDevWorkspace(
    base_image=CUSTOM_BASE_IMAGE_TAG,
    host_port=8011,
    platform=detect_platform(),
    target="source",  # NOTE: "binary" target does not work with custom tools
) as workspace:
    logger.info("✅ Custom agent server started!")

    # 4) Import custom tools to register them in the client's registry
    #    This allows the client to send the module qualname to the server
    #    The server will then import the same module and execute the tool
    import custom_tools.log_data  # noqa: F401

    # 5) Create agent with custom tools
    #    Note: We specify the tool here, but it's actually executed on the server
    #    Get default tools and add our custom tool
    from openhands.sdk import Agent
    from openhands.tools.preset.default import get_default_condenser, get_default_tools

    tools = get_default_tools(enable_browser=False)
    # Add our custom tool!
    tools.append(Tool(name="LogDataTool"))

    agent = Agent(
        llm=llm,
        tools=tools,
        system_prompt_kwargs={"cli_mode": True},
        condenser=get_default_condenser(
            llm=llm.model_copy(update={"usage_id": "condenser"})
        ),
    )

    # 6) Set up callback collection
    received_events: list = []
    last_event_time = {"ts": time.time()}

    def event_callback(event) -> None:
        event_type = type(event).__name__
        logger.info(f"🔔 Callback received event: {event_type}\n{event}")
        received_events.append(event)
        last_event_time["ts"] = time.time()

    # 7) Test the workspace with a simple command
    result = workspace.execute_command(
        "echo 'Custom agent server ready!' && python --version"
    )
    logger.info(
        f"Command '{result.command}' completed with exit code {result.exit_code}"
    )
    logger.info(f"Output: {result.stdout}")

    # 8) Create conversation with the custom agent
    conversation = Conversation(
        agent=agent,
        workspace=workspace,
        callbacks=[event_callback],
    )
    assert isinstance(conversation, RemoteConversation)

    try:
        logger.info(f"\n📋 Conversation ID: {conversation.state.id}")

        logger.info("📝 Sending task to analyze files and log findings...")
        conversation.send_message(
            "Please analyze the Python files in the current directory. "
            "Use the LogDataTool to log your findings as you work. "
            "For example:\n"
            "- Log when you start analyzing a file (level: info)\n"
            "- Log any interesting patterns you find (level: info)\n"
            "- Log any potential issues (level: warning)\n"
            "- Include relevant data like file names, line numbers, etc.\n\n"
            "Make at least 3 log entries using the LogDataTool."
        )
        logger.info("🚀 Running conversation...")
        conversation.run()
        logger.info("✅ Task completed!")
        logger.info(f"Agent status: {conversation.state.execution_status}")

        # Wait for events to settle (no events for 2 seconds)
        logger.info("⏳ Waiting for events to stop...")
        while time.time() - last_event_time["ts"] < 2.0:
            time.sleep(0.1)
        logger.info("✅ Events have stopped")

        # 9) Read the logged data from the JSON file using file_download API
        logger.info("\n📊 Logged Data Summary:")
        logger.info("=" * 80)

        # Download the log file from the workspace using the file download API
        import json
        import tempfile

        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".json", delete=False
        ) as tmp_file:
            local_path = tmp_file.name

        download_result = workspace.file_download(
            source_path="/tmp/agent_data.json",
            destination_path=local_path,
        )

        if download_result.success:
            try:
                with open(local_path) as f:
                    log_entries = json.load(f)
                logger.info(f"Found {len(log_entries)} log entries:\n")
                for i, entry in enumerate(log_entries, 1):
                    logger.info(f"Entry {i}:")
                    logger.info(f"  Timestamp: {entry.get('timestamp', 'N/A')}")
                    logger.info(f"  Level: {entry.get('level', 'N/A')}")
                    logger.info(f"  Message: {entry.get('message', 'N/A')}")
                    if entry.get("data"):
                        logger.info(f"  Data: {json.dumps(entry['data'], indent=4)}")
                    logger.info("")
            except json.JSONDecodeError:
                logger.info("Log file exists but couldn't parse JSON")
                with open(local_path) as f:
                    logger.info(f"Raw content: {f.read()}")
            finally:
                # Clean up the temporary file
                Path(local_path).unlink(missing_ok=True)
        else:
            logger.info("No log file found (agent may not have used the tool)")
            if download_result.error:
                logger.debug(f"Download error: {download_result.error}")

        logger.info("=" * 80)

        cost = conversation.conversation_stats.get_combined_metrics().accumulated_cost
        print(f"\nEXAMPLE_COST: {cost}")

    finally:
        logger.info("\n🧹 Cleaning up conversation...")
        conversation.close()

logger.info("\n✅ Example completed successfully!")
logger.info("\nThis example demonstrated how to:")
logger.info("1. Create a custom tool that logs structured data to JSON")
logger.info("2. Build a simple base image with the custom tool")
logger.info("3. Use DockerDevWorkspace with base_image to build agent server on top")
logger.info("4. Enable dynamic tool registration on the server")
logger.info("5. Use the custom tool during agent execution")
logger.info("6. Read the logged data back from the workspace")

"""Log Data Tool - Example custom tool for logging structured data to JSON.

This tool demonstrates how to create a custom tool that logs structured data
to a local JSON file during agent execution. The data can be retrieved and
verified after the agent completes.
"""

import json
from collections.abc import Sequence
from datetime import UTC, datetime
from enum import Enum
from pathlib import Path
from typing import Any

from pydantic import Field

from openhands.sdk import (
    Action,
    ImageContent,
    Observation,
    TextContent,
    ToolDefinition,
)
from openhands.sdk.tool import ToolExecutor, register_tool


# --- Enums and Models ---


class LogLevel(str, Enum):
    """Log level for entries."""

    DEBUG = "debug"
    INFO = "info"
    WARNING = "warning"
    ERROR = "error"


class LogDataAction(Action):
    """Action to log structured data to a JSON file."""

    message: str = Field(description="The log message")
    level: LogLevel = Field(
        default=LogLevel.INFO,
        description="Log level (debug, info, warning, error)",
    )
    data: dict[str, Any] = Field(
        default_factory=dict,
        description="Additional structured data to include in the log entry",
    )


class LogDataObservation(Observation):
    """Observation returned after logging data."""

    success: bool = Field(description="Whether the data was successfully logged")
    log_file: str = Field(description="Path to the log file")
    entry_count: int = Field(description="Total number of entries in the log file")

    @property
    def to_llm_content(self) -> Sequence[TextContent | ImageContent]:
        """Convert observation to LLM content."""
        if self.success:
            return [
                TextContent(
                    text=(
                        f"✅ Data logged successfully to {self.log_file}\n"
                        f"Total entries: {self.entry_count}"
                    )
                )
            ]
        return [TextContent(text="❌ Failed to log data")]


# --- Executor ---

# Default log file path
DEFAULT_LOG_FILE = "/tmp/agent_data.json"


class LogDataExecutor(ToolExecutor[LogDataAction, LogDataObservation]):
    """Executor that logs structured data to a JSON file."""

    def __init__(self, log_file: str = DEFAULT_LOG_FILE):
        """Initialize the log data executor.

        Args:
            log_file: Path to the JSON log file
        """
        self.log_file = Path(log_file)

    def __call__(
        self,
        action: LogDataAction,
        conversation=None,  # noqa: ARG002
    ) -> LogDataObservation:
        """Execute the log data action.

        Args:
            action: The log data action
            conversation: Optional conversation context (not used)

        Returns:
            LogDataObservation with the result
        """
        # Load existing entries or start fresh
        entries: list[dict[str, Any]] = []
        if self.log_file.exists():
            try:
                with open(self.log_file) as f:
                    entries = json.load(f)
            except (json.JSONDecodeError, OSError):
                entries = []

        # Create new entry with timestamp
        entry = {
            "timestamp": datetime.now(UTC).isoformat(),
            "level": action.level.value,
            "message": action.message,
            "data": action.data,
        }
        entries.append(entry)

        # Write back to file
        self.log_file.parent.mkdir(parents=True, exist_ok=True)
        with open(self.log_file, "w") as f:
            json.dump(entries, f, indent=2)

        return LogDataObservation(
            success=True,
            log_file=str(self.log_file),
            entry_count=len(entries),
        )


# --- Tool Definition ---

_LOG_DATA_DESCRIPTION = """Log structured data to a JSON file.

Use this tool to record information, findings, or events during your work.
Each log entry includes a timestamp and can contain arbitrary structured data.

Parameters:
* message: A descriptive message for the log entry
* level: Log level - one of 'debug', 'info', 'warning', 'error' (default: info)
* data: Optional dictionary of additional structured data to include

Example usage:
- Log a finding: message="Found potential issue", level="warning", data={"file": "app.py", "line": 42}
- Log progress: message="Completed analysis", level="info", data={"files_checked": 10}
"""  # noqa: E501


class LogDataTool(ToolDefinition[LogDataAction, LogDataObservation]):
    """Tool for logging structured data to a JSON file."""

    @classmethod
    def create(cls, conv_state, **params) -> Sequence[ToolDefinition]:  # noqa: ARG003
        """Create LogDataTool instance.

        Args:
            conv_state: Conversation state (not used in this example)
            **params: Additional parameters:
                - log_file: Path to the JSON log file (default: /tmp/agent_data.json)

        Returns:
            A sequence containing a single LogDataTool instance
        """
        log_file = params.get("log_file", DEFAULT_LOG_FILE)
        executor = LogDataExecutor(log_file=log_file)

        return [
            cls(
                description=_LOG_DATA_DESCRIPTION,
                action_type=LogDataAction,
                observation_type=LogDataObservation,
                executor=executor,
            )
        ]


# Auto-register the tool when this module is imported
# This is what enables dynamic tool registration in the remote agent server
register_tool("LogDataTool", LogDataTool)