arcade-mcp/arcade/arcade/cli/utils.py

import importlib.util
import ipaddress
import os
import webbrowser
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Union, cast
from urllib.parse import urlparse

import idna
import typer
from arcadepy import NOT_GIVEN, APIConnectionError, APIStatusError, APITimeoutError, Arcade
from arcadepy.types import AuthorizationResponse
from openai import OpenAI
from openai.resources.chat.completions import ChatCompletionChunk, Stream
from openai.types.chat.chat_completion import Choice as ChatCompletionChoice
from openai.types.chat.chat_completion_chunk import Choice as ChatCompletionChunkChoice
from rich.console import Console
from rich.live import Live
from rich.markdown import Markdown
from rich.text import Text
from typer.core import TyperGroup
from typer.models import Context

from arcade.core.catalog import ToolCatalog
from arcade.core.config_model import Config
from arcade.core.errors import ToolkitLoadError
from arcade.core.schema import ToolDefinition
from arcade.core.toolkit import Toolkit

console = Console()


class OrderCommands(TyperGroup):
    def list_commands(self, ctx: Context) -> list[str]:  # type: ignore[override]
        """Return list of commands in the order appear."""
        return list(self.commands)  # get commands using self.commands


def create_cli_catalog(
    toolkit: str | None = None,
    show_toolkits: bool = False,
) -> ToolCatalog:
    """
    Load toolkits from the python environment.
    """
    if toolkit:
        toolkit = toolkit.lower()
        try:
            prefixed_toolkit = "arcade_" + toolkit
            toolkits = [Toolkit.from_package(prefixed_toolkit)]
        except ToolkitLoadError:
            try:  # try without prefix
                toolkits = [Toolkit.from_package(toolkit)]
            except ToolkitLoadError as e:
                console.print(f"❌ {e}", style="bold red")
                typer.Exit(code=1)
    else:
        toolkits = Toolkit.find_all_arcade_toolkits()

    if not toolkits:
        console.print("❌ No toolkits found or specified", style="bold red")
        typer.Exit(code=1)

    catalog = ToolCatalog()
    for loaded_toolkit in toolkits:
        if show_toolkits:
            console.print(f"Loading toolkit: {loaded_toolkit.name}", style="bold blue")
        catalog.add_toolkit(loaded_toolkit)
    return catalog


def compute_base_url(
    force_tls: bool,
    force_no_tls: bool,
    host: str,
    port: int | None,
) -> str:
    """
    Compute the base URL for the Arcade Engine from the provided overrides.

    force_no_tls takes precedence over force_tls. For example, if both are set to True,
    the resulting URL will use http.

    The port is included in the URL unless the host is a fully qualified domain name
    (excluding IP addresses) and no port is specified. Handles IPv4, IPv6, IDNs, and
    hostnames with underscores.

    This property exists to provide a consistent and correctly formatted URL for
    connecting to the Arcade Engine, taking into account various configuration
    options and edge cases. It ensures that:

    1. The correct protocol (http/https) is used based on the TLS setting.
    2. IPv4 and IPv6 addresses are properly formatted.
    3. Internationalized Domain Names (IDNs) are correctly encoded.
    4. Fully Qualified Domain Names (FQDNs) are identified and handled appropriately.
    5. Ports are included when necessary, respecting common conventions for FQDNs.
    6. Hostnames with underscores (common in development environments) are supported.
    7. Pre-existing port specifications in the host are respected.

    The resulting URL is always suffixed with api_version to specify the API version.

    Returns:
        str: The fully constructed URL for the Arcade Engine.
    """
    # Determine TLS setting based on input flags
    if force_no_tls:
        is_tls = False
    elif force_tls:
        is_tls = True
    else:
        is_tls = host != "localhost"

    # "localhost" defaults to dev port if not specified
    if host == "localhost" and port is None:
        port = 9099

    protocol = "https" if is_tls else "http"

    # Handle potential IDNs
    try:
        encoded_host = idna.encode(host).decode("ascii")
    except idna.IDNAError:
        encoded_host = host

    # Check if the host is a valid IP address (IPv4 or IPv6)
    try:
        ipaddress.ip_address(encoded_host)
        is_ip = True
    except ValueError:
        is_ip = False

    # Parse the host, handling potential IPv6 addresses
    host_for_parsing = f"[{encoded_host}]" if is_ip and ":" in encoded_host else encoded_host
    parsed_host = urlparse(f"//{host_for_parsing}")

    # Check if the host is a fully qualified domain name (excluding IP addresses)
    is_fqdn = "." in parsed_host.netloc and not is_ip and "_" not in parsed_host.netloc

    # Handle hosts that might already include a port
    if ":" in parsed_host.netloc and not is_ip:
        host, existing_port = parsed_host.netloc.rsplit(":", 1)
        if existing_port.isdigit():
            return f"{protocol}://{parsed_host.netloc}"

    if is_fqdn and port is None:
        return f"{protocol}://{encoded_host}"
    elif port is not None:
        return f"{protocol}://{encoded_host}:{port}"
    else:
        return f"{protocol}://{encoded_host}"


def get_tools_from_engine(
    host: str,
    port: int | None = None,
    force_tls: bool = False,
    force_no_tls: bool = False,
    toolkit: str | None = None,
) -> list[ToolDefinition]:
    config = validate_and_get_config()
    base_url = compute_base_url(force_tls, force_no_tls, host, port)
    client = Arcade(api_key=config.api.key, base_url=base_url)

    tools = []
    page_iterator = client.tools.list(toolkit=toolkit or NOT_GIVEN)
    for tool in page_iterator:
        tools.append(ToolDefinition.model_validate(tool.model_dump()))

    return tools


def get_tool_messages(choice: dict) -> list[dict]:
    if hasattr(choice, "tool_messages") and choice.tool_messages:
        return choice.tool_messages  # type: ignore[no-any-return]
    return []


@dataclass
class StreamingResult:
    role: str
    full_message: str
    tool_messages: list
    tool_authorization: dict | None


def handle_streaming_content(stream: Stream[ChatCompletionChunk], model: str) -> StreamingResult:
    """
    Display the streamed markdown chunks as a single line.
    """
    from rich.live import Live

    full_message = ""
    tool_messages = []
    tool_authorization = None
    role = ""
    printed_role: bool = False

    with Live(console=console, refresh_per_second=10) as live:
        for chunk in stream:
            choice = chunk.choices[0]
            role = choice.delta.role or role

            # Display and get tool messages if they exist
            tool_messages += get_tool_messages(choice)  # type: ignore[arg-type]
            tool_authorization = get_tool_authorization(choice)

            chunk_message = choice.delta.content

            if role == "assistant" and tool_authorization:
                continue  # Skip the message if it's an auth request (handled later in handle_tool_authorization)

            if role == "assistant" and not printed_role:
                console.print(f"\n[blue][bold]Assistant[/bold] ({model}):[/blue] ")
                printed_role = True

            if chunk_message:
                full_message += chunk_message
                markdown_chunk = Markdown(full_message)
                live.update(markdown_chunk)

        # Markdownify URLs in the final message if applicable
        if role == "assistant":
            full_message = markdownify_urls(full_message)
            live.update(Markdown(full_message))

    return StreamingResult(role, full_message, tool_messages, tool_authorization)


def markdownify_urls(message: str) -> str:
    """
    Convert URLs in the message to markdown links.
    """
    import re

    # This regex will match URLs that are not already formatted as markdown links:
    # [Link text](https://example.com)
    url_pattern = r"(?<!\]\()https?://\S+"

    # Wrap all URLs in the message with markdown links
    return re.sub(url_pattern, r"[Link](\g<0>)", message)


def validate_and_get_config(
    validate_api: bool = True,
    validate_user: bool = True,
) -> Config:
    """
    Validates the configuration, user, and returns the Config object
    """
    from arcade.core.config import config

    if validate_api and (not config.api or not config.api.key):
        console.print(
            "❌ API configuration not found or key is missing. Please run `arcade login`.",
            style="bold red",
        )
        raise typer.Exit(code=1)

    if validate_user and (not config.user or not config.user.email):
        console.print(
            "❌ User email not found in configuration. Please run `arcade login`.", style="bold red"
        )
        raise typer.Exit(code=1)

    return config


def log_engine_health(client: Arcade) -> None:
    try:
        result = client.health.check(timeout=2)
        if result.healthy:
            return

        console.print(
            "⚠️ Warning: Arcade Engine is unhealthy",
            style="bold yellow",
        )

    except APIConnectionError:
        console.print(
            "⚠️ Warning: Arcade Engine was unreachable. (Is it running?)",
            style="bold yellow",
        )

    except APIStatusError as e:
        console.print(
            "[bold][yellow]⚠️ Warning: "
            + str(e)
            + " ("
            + "[/yellow]"
            + "[red]"
            + str(e.status_code)
            + "[/red]"
            + "[yellow])[/yellow][/bold]"
        )


@dataclass
class ChatInteractionResult:
    history: list[dict]
    tool_messages: list[dict]
    tool_authorization: dict | None


def handle_chat_interaction(
    client: OpenAI, model: str, history: list[dict], user_email: str | None, stream: bool = False
) -> ChatInteractionResult:
    """
    Handle a single chat-request/chat-response interaction for both streamed and non-streamed responses.
    Handling the chat response includes:
    - Streaming the response if the stream flag is set
    - Displaying the response in the console
    - Getting the tool messages and tool authorization from the response
    - Updating the history with the response, tool calls, and tool responses
    """
    if stream:
        # TODO Fix this in the client so users don't deal with these
        # typing issues
        response = client.chat.completions.create(  # type: ignore[call-overload]
            model=model,
            messages=history,
            tool_choice="generate",
            user=user_email,
            stream=True,
        )
        streaming_result = handle_streaming_content(response, model)
        role, message_content = streaming_result.role, streaming_result.full_message
        tool_messages, tool_authorization = (
            streaming_result.tool_messages,
            streaming_result.tool_authorization,
        )
    else:
        response = client.chat.completions.create(  # type: ignore[call-overload]
            model=model,
            messages=history,
            tool_choice="generate",
            user=user_email,
            stream=False,
        )
        message_content = response.choices[0].message.content or ""

        # Get extra fields from the response
        tool_messages = get_tool_messages(response.choices[0])
        tool_authorization = get_tool_authorization(response.choices[0])

        role = response.choices[0].message.role

        if role == "assistant" and tool_authorization:
            pass  # Skip the message if it's an auth request (handled later in handle_tool_authorization)
        elif role == "assistant":
            message_content = markdownify_urls(message_content)
            console.print(
                f"\n[blue][bold]Assistant[/bold] ({model}):[/blue] ", Markdown(message_content)
            )
        else:
            console.print(f"\n[bold]{role}:[/bold] {message_content}")

    history += tool_messages
    history.append({"role": role, "content": message_content})

    return ChatInteractionResult(history, tool_messages, tool_authorization)


def handle_tool_authorization(
    arcade_client: Arcade,
    tool_authorization: AuthorizationResponse,
    history: list[dict[str, Any]],
    openai_client: OpenAI,
    model: str,
    user_email: str | None,
    stream: bool,
) -> ChatInteractionResult:
    with Live(console=console, refresh_per_second=4) as live:
        if tool_authorization.authorization_url:
            authorization_url = str(tool_authorization.authorization_url)
            webbrowser.open(authorization_url)
            message = (
                "You'll need to authorize this action in your browser.\n\n"
                f"If a browser doesn't open automatically, click [this link]({authorization_url}) "
                f"or copy this URL and paste it into your browser:\n\n{authorization_url}"
            )
            live.update(Markdown(message, style="dim"))

        wait_for_authorization_completion(arcade_client, tool_authorization)

        message = "Thanks for authorizing the action! Sending your request..."
        live.update(Text(message, style="dim"))

    history.pop()
    return handle_chat_interaction(openai_client, model, history, user_email, stream)


def wait_for_authorization_completion(
    client: Arcade, tool_authorization: AuthorizationResponse | None
) -> None:
    """
    Wait for the authorization for a tool call to complete i.e., wait for the user to click on
    the approval link and authorize Arcade.
    """
    if tool_authorization is None:
        return

    auth_response = AuthorizationResponse.model_validate(tool_authorization)

    while auth_response.status != "completed":
        try:
            auth_response = client.auth.status(
                authorization_id=cast(str, auth_response.authorization_id),
                scopes=" ".join(auth_response.scopes) if auth_response.scopes else NOT_GIVEN,
                wait=59,
            )
        except APITimeoutError:
            continue


def get_tool_authorization(
    choice: Union[ChatCompletionChoice, ChatCompletionChunkChoice],
) -> dict | None:
    """
    Get the tool authorization from a chat response's choice.
    """
    if hasattr(choice, "tool_authorizations") and choice.tool_authorizations:
        return choice.tool_authorizations[0]  # type: ignore[no-any-return]
    return None


def is_authorization_pending(tool_authorization: dict | None) -> bool:
    """
    Check if the authorization for a tool call is pending.
    Expects a chat response's choice.tool_authorizations as input.
    """
    is_auth_pending = (
        tool_authorization is not None and tool_authorization.get("status", "") == "pending"
    )
    return is_auth_pending


def get_eval_files(directory: str) -> list[Path]:
    """
    Get a list of evaluation files starting with 'eval_' and ending with '.py' in the given directory.

    Args:
        directory: The directory to search for evaluation files.

    Returns:
        A list of Paths to the evaluation files. Returns an empty list if no files are found.
    """
    directory_path = Path(directory).resolve()

    if directory_path.is_dir():
        eval_files = [
            f
            for f in directory_path.iterdir()
            if f.is_file() and f.name.startswith("eval_") and f.name.endswith(".py")
        ]
    elif directory_path.is_file():
        eval_files = (
            [directory_path]
            if directory_path.name.startswith("eval_") and directory_path.name.endswith(".py")
            else []
        )
    else:
        console.print(f"Path not found: {directory_path}", style="bold red")
        return []

    if not eval_files:
        console.print(
            "No evaluation files found. Filenames must start with 'eval_' and end with '.py'.",
            style="bold yellow",
        )
        return []

    return eval_files


def load_eval_suites(eval_files: list[Path]) -> list[Callable]:
    """
    Load evaluation suites from the given eval_files by importing the modules
    and extracting functions decorated with `@tool_eval`.

    Args:
        eval_files: A list of Paths to evaluation files.

    Returns:
        A list of callable evaluation suite functions.
    """
    eval_suites = []
    for eval_file_path in eval_files:
        module_name = eval_file_path.stem  # filename without extension

        # Now we need to load the module from eval_file_path
        file_path_str = str(eval_file_path)
        module_name_str = module_name

        # Load using importlib
        spec = importlib.util.spec_from_file_location(module_name_str, file_path_str)
        if spec is None:
            console.print(f"Failed to load {eval_file_path}", style="bold red")
            continue

        module = importlib.util.module_from_spec(spec)
        if spec.loader is not None:
            spec.loader.exec_module(module)
        else:
            console.print(f"Failed to load module: {module_name}", style="bold red")
            continue

        eval_suite_funcs = [
            obj
            for name, obj in module.__dict__.items()
            if callable(obj) and hasattr(obj, "__tool_eval__")
        ]

        if not eval_suite_funcs:
            console.print(f"No @tool_eval functions found in {eval_file_path}", style="bold yellow")
            continue

        eval_suites.extend(eval_suite_funcs)

    return eval_suites


def create_new_env_file() -> None:
    """
    Create a new env file if one doesn't already exist.
    """
    env_file = os.path.expanduser("~/.arcade/arcade.env")
    if not os.path.exists(env_file):
        template_path = os.path.join(
            os.path.dirname(__file__), "..", "templates", "arcade.template.env"
        )
        os.makedirs(os.path.dirname(env_file), exist_ok=True)

        with open(template_path) as template_file, open(env_file, "w") as new_env_file:
            template_contents = template_file.read()
            new_env_file.write(template_contents)

        console.print(f"Created new environment file at {env_file}", style="bold green")


def is_config_file_deprecated() -> bool:
    """
    Check if the user is using the deprecated config file.

    Returns:
        bool: True if the user is using the deprecated config file, False otherwise.
    """
    deprecated_config_file_path = os.path.expanduser("~/.arcade/arcade.toml")
    if os.path.exists(deprecated_config_file_path):
        console.print(
            f"Deprecation Notice: You are using a deprecated config file at {deprecated_config_file_path}. Please migrate to the new format by running,\n\n\t$ arcade logout && arcade login\n",
            style="bold yellow",
        )
        return True
    return False


def delete_deprecated_config_file() -> None:
    """
    Delete the deprecated config file if it exists.
    """
    deprecated_config_file_path = os.path.expanduser("~/.arcade/arcade.toml")

    if os.path.exists(deprecated_config_file_path):
        os.remove(deprecated_config_file_path)