tools.py

import json
import uuid
from typing import Any

from config import INVOKE_RE, PARAM_RE, FUNC_CALLS_BLOCK_RE, ANTI_P5JS_PROMPT


def build_tools_system_prompt(tools: list[dict] | None) -> str:
    if not tools:
        return ""
    tool_descriptions = []
    for t in tools:
        if "function" in t:
            fn = t["function"]
            name = fn.get("name", "")
            desc = fn.get("description", "")
            schema = fn.get("parameters", {})
        else:
            name = t.get("name", "")
            desc = t.get("description", "")
            schema = t.get("input_schema", {})
        tool_descriptions.append(
            f"<tool>\n<name>{name}</name>\n<description>{desc}</description>\n"
            f"<parameters>{json.dumps(schema, ensure_ascii=False)}</parameters>\n</tool>"
        )
    tools_xml = "\n".join(tool_descriptions)
    return (
        "In this environment you have access to a set of tools you can use to answer the user's question. "
        "When you need to call a tool, you MUST emit it in EXACTLY this XML format — and nothing else until the tool result arrives:\n"
        "<function_calls>\n"
        "<invoke name=\"TOOL_NAME\">\n"
        "<parameter name=\"PARAM_NAME\">PARAM_VALUE</parameter>\n"
        "...\n"
        "</invoke>\n"
        "</function_calls>\n\n"
        "Rules:\n"
        "- Emit the XML exactly as shown, with the literal tags <function_calls>, <invoke>, <parameter>.\n"
        "- One <invoke> per tool call. You can emit multiple <invoke> blocks inside one <function_calls>.\n"
        "- Do NOT wrap the XML in markdown code fences.\n"
        "- After emitting </function_calls>, stop. Do not add any trailing text — wait for the tool result.\n"
        "- Parameter values must be raw text (for objects/arrays use compact JSON).\n\n"
        f"Available tools:\n<tools>\n{tools_xml}\n</tools>\n"
    )


def extract_tool_results_from_content(content: Any) -> tuple[str, list[dict]]:
    """Returns (plain_text, tool_results) from an Anthropic-style content array."""
    text_parts: list[str] = []
    tool_results: list[dict] = []
    if isinstance(content, list):
        for block in content:
            if isinstance(block, dict):
                btype = block.get("type")
                if btype == "tool_result":
                    tool_results.append({
                        "id": block.get("tool_use_id", ""),
                        "content": normalize_content(block.get("content", "")),
                    })
                elif btype == "text":
                    text_parts.append(block.get("text", ""))
                elif btype == "tool_use":
                    pass
    elif isinstance(content, str):
        text_parts.append(content)
    return ("\n".join(p for p in text_parts if p), tool_results)


def extract_tool_uses_from_content(content: Any) -> tuple[str, list[dict]]:
    """Returns (plain_text, tool_uses) from an Anthropic-style assistant content array."""
    text_parts: list[str] = []
    tool_uses: list[dict] = []
    if isinstance(content, list):
        for block in content:
            if isinstance(block, dict):
                btype = block.get("type")
                if btype == "text":
                    text_parts.append(block.get("text", ""))
                elif btype == "tool_use":
                    tool_uses.append({
                        "id": block.get("id", ""),
                        "name": block.get("name", ""),
                        "input": block.get("input", {}),
                    })
    elif isinstance(content, str):
        text_parts.append(content)
    return ("\n".join(p for p in text_parts if p), tool_uses)


def render_assistant_tool_uses_as_xml(text: str, tool_uses: list[dict]) -> str:
    if not tool_uses:
        return text
    parts = []
    if text:
        parts.append(text)
    invokes = []
    for tu in tool_uses:
        params = []
        for k, v in (tu.get("input") or {}).items():
            v_str = v if isinstance(v, str) else json.dumps(v, ensure_ascii=False)
            params.append(f'<parameter name="{k}">{v_str}</parameter>')
        invokes.append(f'<invoke name="{tu["name"]}">\n' + "\n".join(params) + "\n</invoke>")
    parts.append("<function_calls>\n" + "\n".join(invokes) + "\n</function_calls>")
    return "\n".join(parts)


def render_tool_results_as_xml(tool_results: list[dict]) -> str:
    if not tool_results:
        return ""
    items = []
    for tr in tool_results:
        items.append(
            f'<result tool_use_id="{tr["id"]}">\n{tr["content"]}\n</result>'
        )
    return "<function_results>\n" + "\n".join(items) + "\n</function_results>"


def normalize_content(content: Any) -> str:
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for block in content:
            if isinstance(block, dict):
                btype = block.get("type")
                if btype == "text":
                    parts.append(block.get("text", ""))
                elif btype == "image_url":
                    url = block.get("image_url", {})
                    parts.append(f"[image: {url.get('url', '') if isinstance(url, dict) else url}]")
            elif isinstance(block, str):
                parts.append(block)
        return "\n".join(p for p in parts if p)
    if content is None:
        return ""
    return str(content)


def parse_function_calls_text(text: str) -> list[dict]:
    """Extract tool_use records from assistant text containing <function_calls> blocks."""
    tool_uses: list[dict] = []
    for block_match in FUNC_CALLS_BLOCK_RE.finditer(text):
        inner = block_match.group(1)
        for inv in INVOKE_RE.finditer(inner):
            name = inv.group(1).strip()
            body = inv.group(2)
            input_obj: dict[str, Any] = {}
            for p in PARAM_RE.finditer(body):
                pname = p.group(1).strip()
                pval = p.group(2).strip()
                try:
                    parsed = json.loads(pval)
                    input_obj[pname] = parsed
                except Exception:
                    input_obj[pname] = pval
            tool_uses.append({
                "id": f"toolu_{uuid.uuid4().hex[:24]}",
                "name": name,
                "input": input_obj,
            })
    return tool_uses


def split_text_and_tools(text: str) -> tuple[str, list[dict]]:
    """Return (clean_text_without_xml, tool_use_list)."""
    tool_uses = parse_function_calls_text(text)
    cleaned = FUNC_CALLS_BLOCK_RE.sub("", text).strip()
    return cleaned, tool_uses