EvoScientist-Multi/EvoScientist/utils.py

"""Utility functions for EvoScientist.

This module primarily contains helpers for displaying messages and prompts in
notebooks, and lightweight configuration loaders used by the agent runtime.
"""

import json
import logging
from pathlib import Path
from typing import Any

import yaml
from rich.console import Console
from rich.panel import Panel
from rich.text import Text

logger = logging.getLogger(__name__)
console = Console()


def format_message_content(message):
    """Convert message content to displayable string.

    Args:
        message: A LangChain message object with content attribute.

    Returns:
        Formatted string representation of the message content.
    """
    parts = []
    tool_calls_processed = False

    # Handle main content
    if isinstance(message.content, str):
        parts.append(message.content)
    elif isinstance(message.content, list):
        # Handle complex content like tool calls (Anthropic format)
        for item in message.content:
            if item.get("type") == "text":
                parts.append(item["text"])
            elif item.get("type") == "tool_use":
                parts.append(f"\n🔧 Tool Call: {item['name']}")
                parts.append(f"   Args: {json.dumps(item['input'], indent=2)}")
                parts.append(f"   ID: {item.get('id', 'N/A')}")
                tool_calls_processed = True
    else:
        parts.append(str(message.content))

    # Handle tool calls attached to the message (OpenAI format) - only if not already processed
    if (
        not tool_calls_processed
        and hasattr(message, "tool_calls")
        and message.tool_calls
    ):
        for tool_call in message.tool_calls:
            parts.append(f"\n🔧 Tool Call: {tool_call['name']}")
            parts.append(f"   Args: {json.dumps(tool_call['args'], indent=2)}")
            parts.append(f"   ID: {tool_call['id']}")

    return "\n".join(parts)


def format_messages(messages):
    """Format and display a list of messages with Rich formatting.

    Args:
        messages: List of LangChain message objects to display.
    """
    for m in messages:
        msg_type = m.__class__.__name__.replace("Message", "")
        content = format_message_content(m)

        if msg_type == "Human":
            console.print(Panel(content, title="🧑 Human", border_style="blue"))
        elif msg_type == "Ai":
            console.print(Panel(content, title="🤖 Assistant", border_style="green"))
        elif msg_type == "Tool":
            console.print(Panel(content, title="🔧 Tool Output", border_style="yellow"))
        else:
            console.print(Panel(content, title=f"📝 {msg_type}", border_style="white"))


def show_prompt(prompt_text: str, title: str = "Prompt", border_style: str = "blue"):
    """Display a prompt with rich formatting and XML tag highlighting.

    Args:
        prompt_text: The prompt string to display
        title: Title for the panel (default: "Prompt")
        border_style: Border color style (default: "blue")
    """
    # Create a formatted display of the prompt
    formatted_text = Text(prompt_text)
    formatted_text.highlight_regex(r"<[^>]+>", style="bold blue")  # Highlight XML tags
    formatted_text.highlight_regex(
        r"##[^#\n]+", style="bold magenta"
    )  # Highlight headers
    formatted_text.highlight_regex(
        r"###[^#\n]+", style="bold cyan"
    )  # Highlight sub-headers

    # Display in a panel for better presentation
    console.print(
        Panel(
            formatted_text,
            title=f"[bold green]{title}[/bold green]",
            border_style=border_style,
            padding=(1, 2),
        )
    )


def load_subagents(
    config_path: Path,
    *,
    tool_registry: dict[str, Any],
    prompt_refs: dict[str, str] | None = None,
    async_swap_pending: bool = False,
) -> list[dict[str, Any]]:
    """Load subagent definitions from a directory of YAML files and wire up tools.

    NOTE: This is a custom utility. deepagents does not natively load subagents
    from files - they're normally defined inline in the create_deep_agent() call.
    We externalize to YAML here to keep configuration separate from code.

    ``config_path`` must be a directory containing one ``<name>.yaml`` per
    sub-agent. All ``*.yaml`` files are merged into a single mapping. Files
    starting with ``.`` (dotfiles, editor swap files) or ``_`` (private /
    disabled) are ignored. ``.yml`` is intentionally not supported — keeps
    one canonical extension and avoids the dev-vs-wheel packaging mismatch.

    Each file's top level must be a mapping ``{<agent-name>: <spec>}``::

        planner-agent:
          description: "..."
          tools: [think_tool]
          system_prompt: |
            ...
        research-agent:
          description: "..."
          tools: [tavily_search, think_tool]
          system_prompt: |
            ...

    ``async_swap_pending`` tells the loader that an async-subagent swap will
    run downstream against a different tool registry (typically the deployed
    graph's own registry in ``subagents/_factory.py``). When True, a tool
    missing from ``tool_registry`` for an ``async: true`` spec is logged at
    DEBUG (it will be re-resolved by the async graph). When False, the
    caller IS the terminal registry — any missing tool is a genuine typo
    and logs at WARNING. Sync in-process callers
    (``EvoScientist._build_base_kwargs``, ``load_mcp_and_build_kwargs``)
    should pass True; the factory
    (``subagents/_factory.build_async_subagent_graph``) leaves it False.
    """
    prompt_refs = prompt_refs or {}

    if not config_path.is_dir():
        raise ValueError(
            f"{config_path}: sub-agent config must be a directory "
            f"containing one <name>.yaml per agent"
        )

    # Only ``.yaml`` is supported (canonical extension). Skip files starting
    # with ``_`` (private / disabled) or ``.`` (dotfiles, editor swap files
    # like ``.foo.yaml.swp``). ``.yml`` is intentionally not loaded — keeps
    # one canonical extension across the project, simplifies packaging
    # (no need for a parallel ``subagents/*.yml`` entry in ``pyproject.toml``
    # ``package-data``), and matches every existing yaml file in this repo.
    config: dict[str, Any] = {}
    for yml in sorted(config_path.glob("*.yaml")):
        if yml.name.startswith(".") or yml.name.startswith("_"):
            continue
        with yml.open(encoding="utf-8") as f:
            data = yaml.safe_load(f)
        if data is None:
            # Empty file — skip silently (matches dotfile/underscore skip behavior)
            continue
        if not isinstance(data, dict):
            raise ValueError(
                f"{yml}: top-level must be a mapping (one entry per sub-agent)"
            )
        # Detect duplicate keys across files
        for key in data:
            if key in config:
                raise ValueError(
                    f"Sub-agent {key!r} defined in multiple files; "
                    f"second occurrence in {yml.name}"
                )
        for key, spec in data.items():
            if not isinstance(spec, dict):
                raise ValueError(
                    f"{yml}: sub-agent {key!r} must map to a spec dict, "
                    f"got {type(spec).__name__}"
                )
        config.update(data)

    if not config:
        raise ValueError(
            f"{config_path}: no sub-agent definitions found "
            f"(expected one or more <name>.yaml files)"
        )

    subagents: list[dict[str, Any]] = []

    def _build_one(name: str, spec: dict[str, Any]) -> dict[str, Any]:
        subagent: dict[str, Any] = {
            "name": name,
            "description": spec.get("description", ""),
        }

        if "system_prompt_ref" in spec:
            ref = spec["system_prompt_ref"]
            if ref not in prompt_refs:
                raise ValueError(
                    f"Unknown system_prompt_ref '{ref}' for subagent '{name}'"
                )
            subagent["system_prompt"] = prompt_refs[ref]
        else:
            subagent["system_prompt"] = spec.get("system_prompt", "")

        if "model" in spec:
            subagent["model"] = spec["model"]

        if "skills" in spec:
            subagent["skills"] = spec["skills"]

        # Compute the async flag up-front so the tools-resolution block below
        # can pick the right log level for missing tools. Internal field:
        # carries the ``async:`` yaml flag through to
        # ``_maybe_swap_async_subagents`` so the swap doesn't need a second
        # yaml pass to discover async-flagged agents. Underscore prefix marks
        # it as internal — must be popped before passing to deepagents.
        async_val = spec.get("async", False)
        if not isinstance(async_val, bool):
            # Reject quoted-string yaml values like ``async: "false"`` —
            # ``bool("false")`` is ``True`` (non-empty string), which silently
            # flips the agent into async mode. Fail loud instead.
            raise ValueError(
                f"Subagent {name!r}: 'async' must be a boolean, "
                f"got {type(async_val).__name__}: {async_val!r}"
            )

        if "tools" in spec:
            resolved = []
            for t in spec["tools"]:
                if t in tool_registry:
                    resolved.append(tool_registry[t])
                elif async_val and async_swap_pending:
                    # Caller expects an async-subagent swap downstream against
                    # a different registry. The in-process registry is
                    # intentionally minimal for async agents — a miss here is
                    # not degraded state.
                    logger.debug(
                        "Subagent %r: tool %r not in the in-process registry; "
                        "resolved by the async graph",
                        name,
                        t,
                    )
                else:
                    # Terminal registry (factory) OR sync subagent — a missing
                    # tool is a genuine typo and won't be resolved anywhere.
                    logger.warning(
                        "Subagent %r: tool %r not in registry, skipping", name, t
                    )
            subagent["tools"] = resolved

        subagent["_async"] = async_val

        return subagent

    for name, spec in config.items():
        subagents.append(_build_one(name, spec))

    return subagents


def load_subagent(
    config_path: Path,
    name: str,
    *,
    tool_registry: dict[str, Any],
    prompt_refs: dict[str, str] | None = None,
    async_swap_pending: bool = False,
) -> dict[str, Any]:
    """Load a single sub-agent by name from YAML.

    See :func:`load_subagents` for ``async_swap_pending`` semantics.
    """
    for agent in load_subagents(
        config_path,
        tool_registry=tool_registry,
        prompt_refs=prompt_refs,
        async_swap_pending=async_swap_pending,
    ):
        if agent.get("name") == name:
            return agent
    raise KeyError(f"Sub-agent not found: {name}")