beaver_project/app-instance/backend/tests/unit/test_tool_assembler.py

from __future__ import annotations

import asyncio
from pathlib import Path
import subprocess
import sys
from types import SimpleNamespace

from beaver.engine.context import SkillContext
from beaver.foundation.embedding import EmbeddingRetriever
from beaver.skills.catalog.loader import SkillsLoader
from beaver.tools import BaseTool, ToolAssembler, ToolContext, ToolExecutor, ToolRegistry, ToolResult, ToolSpec


class DummyTool(BaseTool):
    def __init__(
        self,
        name: str,
        *,
        description: str | None = None,
        toolset: str = "test",
        always_available: bool = False,
    ) -> None:
        self._spec = ToolSpec(
            name=name,
            description=description or name,
            input_schema={"type": "object", "properties": {}},
            toolset=toolset,
            always_available=always_available,
        )

    @property
    def spec(self) -> ToolSpec:
        return self._spec

    async def invoke(self, arguments: dict, context: ToolContext) -> ToolResult:
        return ToolResult(success=True, content="ok", tool_name=self.spec.name)


class StaticRetriever:
    async def retrieve(self, **kwargs):
        candidates = kwargs["candidates"]
        top_k = kwargs["top_k"]
        preferred = ["search_files", "echo"]
        ordered = sorted(
            candidates,
            key=lambda item: preferred.index(item["name"]) if item["name"] in preferred else len(preferred),
        )
        return ordered[:top_k]


def test_tool_spec_exports_mcp_and_provider_schema() -> None:
    spec = ToolSpec(
        name="read_file",
        description="Read a file",
        input_schema={"type": "object", "properties": {"path": {"type": "string"}}},
        toolset="file",
    )

    assert spec.to_mcp_descriptor() == {
        "name": "read_file",
        "description": "Read a file",
        "inputSchema": {"type": "object", "properties": {"path": {"type": "string"}}},
    }
    assert spec.to_provider_schema()["function"]["parameters"] == spec.input_schema


def test_tool_assembler_merges_always_skill_hints_and_embedding(tmp_path: Path) -> None:
    skill_dir = tmp_path / "skills" / "docker-debug"
    skill_dir.mkdir(parents=True)
    (skill_dir / "SKILL.md").write_text(
        """---
name: docker-debug
description: Debug Docker issues.
tools:
  - terminal
---

# Docker Debug
""",
        encoding="utf-8",
    )

    registry = ToolRegistry()
    registry.register(DummyTool("memory", toolset="memory", always_available=True))
    registry.register(DummyTool("terminal", toolset="shell"))
    registry.register(DummyTool("search_files", toolset="file"))
    registry.register(DummyTool("echo", toolset="debug"))

    assembler = ToolAssembler(retriever=StaticRetriever())
    loader = SkillsLoader(tmp_path)
    selected = asyncio.run(
        assembler.assemble(
            task_description="排查 Docker 容器日志",
            registry=registry,
            skills_loader=loader,
            activated_skills=[SkillContext(name="docker-debug", content="")],
            top_k=1,
        )
    )

    assert [spec.name for spec in selected] == ["memory", "terminal", "search_files"]


def test_tool_assembler_uses_required_tools_section_when_frontmatter_omits_tools(tmp_path: Path) -> None:
    skill_dir = tmp_path / "skills" / "docker-debug"
    skill_dir.mkdir(parents=True)
    (skill_dir / "SKILL.md").write_text(
        """---
name: docker-debug
description: Debug Docker issues.
---

# Docker Debug

## Overview

Debug Docker issues.

## Required Tools

- `terminal`
- `search_files`

## Workflow

Inspect logs and search related files.
""",
        encoding="utf-8",
    )

    registry = ToolRegistry()
    registry.register(DummyTool("memory", toolset="memory", always_available=True))
    registry.register(DummyTool("terminal", toolset="shell"))
    registry.register(DummyTool("search_files", toolset="file"))
    registry.register(DummyTool("echo", toolset="debug"))

    assembler = ToolAssembler(retriever=StaticRetriever())
    loader = SkillsLoader(tmp_path)
    record = loader.get_skill_record("docker-debug")
    assert record is not None
    assert record.tool_hints == ["terminal", "search_files"]

    selected = asyncio.run(
        assembler.assemble(
            task_description="排查 Docker 容器日志",
            registry=registry,
            skills_loader=loader,
            activated_skills=[SkillContext(name="docker-debug", content="", tool_hints=record.tool_hints)],
            top_k=1,
        )
    )

    assert [spec.name for spec in selected] == ["memory", "terminal", "search_files", "echo"]


def test_embedding_fallback_can_return_all_or_top_k() -> None:
    candidates = [{"name": f"tool_{index}", "description": "", "input_schema": "{}"} for index in range(3)]
    retriever = EmbeddingRetriever(api_key_env="MISSING_EMBEDDING_KEY", api_base_env="MISSING_EMBEDDING_BASE")

    all_candidates = asyncio.run(
        retriever.retrieve(query="x", candidates=candidates, top_k=1, fallback_top_k=None)
    )
    top_candidate = asyncio.run(
        retriever.retrieve(query="x", candidates=candidates, top_k=1, fallback_top_k=1)
    )

    assert [item["name"] for item in all_candidates] == ["tool_0", "tool_1", "tool_2"]
    assert [item["name"] for item in top_candidate] == ["tool_0"]


def test_beaver_tools_import_does_not_load_provider_stack_with_socks_proxy() -> None:
    code = (
        "import beaver.tools\n"
        "from beaver.skills.catalog.loader import SkillsLoader\n"
        "print('ok')"
    )
    result = subprocess.run(
        [sys.executable, "-c", code],
        check=False,
        capture_output=True,
        text=True,
        env={
            "PYTHONPATH": str(Path(__file__).resolve().parents[2]),
            "HTTP_PROXY": "socks://127.0.0.1:7897/",
            "HTTPS_PROXY": "socks://127.0.0.1:7897/",
        },
    )

    assert result.returncode == 0, result.stderr
    assert result.stdout.strip() == "ok"


def test_tool_executor_parses_object_tool_call_string_arguments() -> None:
    tool_call = SimpleNamespace(name="echo", arguments='{"text": "hello"}')

    name, arguments = ToolExecutor._normalize_tool_call(tool_call)

    assert name == "echo"
    assert arguments == {"text": "hello"}