deepset-ai--haystack/test/tools/skills/test_skill_toolset.py

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0


import threading
from concurrent.futures import ThreadPoolExecutor

import pytest

from haystack.components.agents.state import State
from haystack.components.agents.tool_calling import _run_tool
from haystack.dataclasses import ChatMessage, FileContent, ImageContent, ToolCall
from haystack.skill_stores.file_system.skill_store import FileSystemSkillStore
from haystack.tools import SkillToolset, Tool
from haystack.tools.errors import ToolInvocationError


def _get_tool(toolset, name):
    """Warm up the toolset and return its tool with the given name."""
    toolset.warm_up()
    return next(t for t in toolset if t.name == name)


def _write_skill(skills_dir, name, description=None, body="Instructions.", files=None):
    skill_dir = skills_dir / name
    skill_dir.mkdir(parents=True)
    frontmatter = f"---\nname: {name}\n"
    if description is not None:
        frontmatter += f"description: {description}\n"
    frontmatter += "---\n"
    (skill_dir / "SKILL.md").write_text(frontmatter + body, encoding="utf-8")
    for rel_path, content in (files or {}).items():
        target = skill_dir / rel_path
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_text(content, encoding="utf-8")
    return skill_dir


class TestSkillToolset:
    def test_tools_present_before_warm_up_without_io(self, tmp_path):
        _write_skill(tmp_path, "pdf-forms", description="Use to fill PDF forms.")
        toolset = SkillToolset(FileSystemSkillStore(tmp_path))

        # The (static) tool set is available immediately, with no store access required.
        assert toolset._is_warmed_up is False
        assert len(toolset) == 2
        assert {t.name for t in toolset} == {"load_skill", "read_skill_file"}
        assert "load_skill" in toolset
        assert toolset._is_warmed_up is False

    def test_scans_skills_on_warm_up(self, tmp_path):
        _write_skill(tmp_path, "pdf-forms", description="Use to fill PDF forms.")
        _write_skill(tmp_path, "excel", description="Use to edit spreadsheets.")

        toolset = SkillToolset(FileSystemSkillStore(tmp_path))

        # The catalog is only scanned on warm_up.
        assert toolset._is_warmed_up is False

        toolset.warm_up()

        assert toolset._is_warmed_up is True
        assert set(toolset.skills) == {"pdf-forms", "excel"}
        assert toolset.skills["pdf-forms"].description == "Use to fill PDF forms."

    def test_skills_property_warms_up_lazily(self, tmp_path):
        _write_skill(tmp_path, "pdf-forms", description="Use to fill PDF forms.")
        toolset = SkillToolset(FileSystemSkillStore(tmp_path))
        # Accessing `skills` without an explicit warm_up triggers it.
        assert set(toolset.skills) == {"pdf-forms"}
        assert toolset._is_warmed_up is True

    def test_warm_up_is_idempotent(self, tmp_path):
        _write_skill(tmp_path, "pdf-forms", description="Use to fill PDF forms.")
        toolset = SkillToolset(FileSystemSkillStore(tmp_path))
        toolset.warm_up()
        toolset.warm_up()
        assert set(toolset.skills) == {"pdf-forms"}

    def test_warm_up_warms_up_the_store(self, tmp_path):
        _write_skill(tmp_path, "pdf-forms", description="Use to fill PDF forms.")
        store = FileSystemSkillStore(tmp_path)
        toolset = SkillToolset(store)
        toolset.warm_up()
        assert store._is_warmed_up is True

    def test_concurrent_warm_up(self, tmp_path):
        # Concurrent first use (e.g. parallel requests hitting a shared Agent) must produce a complete,
        # consistent catalog in every thread.
        _write_skill(tmp_path, "pdf-forms", description="Use to fill PDF forms.")
        _write_skill(tmp_path, "excel", description="Use to edit spreadsheets.")
        toolset = SkillToolset(FileSystemSkillStore(tmp_path))

        num_threads = 8
        barrier = threading.Barrier(num_threads)

        def warm_up_and_list():
            barrier.wait()
            toolset.warm_up()
            return set(toolset.skills)

        with ThreadPoolExecutor(max_workers=num_threads) as executor:
            results = list(executor.map(lambda _: warm_up_and_list(), range(num_threads)))

        assert all(result == {"pdf-forms", "excel"} for result in results)
        assert "- pdf-forms: Use to fill PDF forms." in toolset._load_skill_tool.description

    def test_add_is_not_supported(self, tmp_path):
        _write_skill(tmp_path, "pdf-forms", description="Use to fill PDF forms.")
        toolset = SkillToolset(FileSystemSkillStore(tmp_path))
        extra = Tool(name="extra", description="d", parameters={"type": "object", "properties": {}}, function=len)
        with pytest.raises(NotImplementedError, match="does not support adding tools"):
            toolset.add(extra)

    def test_concat_is_not_supported(self, tmp_path):
        _write_skill(tmp_path, "pdf-forms", description="Use to fill PDF forms.")
        toolset = SkillToolset(FileSystemSkillStore(tmp_path))
        extra = Tool(name="extra", description="d", parameters={"type": "object", "properties": {}}, function=len)
        with pytest.raises(NotImplementedError, match="does not support concatenation"):
            _ = toolset + extra

    def test_accepts_skill_store_instance(self, tmp_path):
        _write_skill(tmp_path, "pdf-forms", description="Use to fill PDF forms.")
        store = FileSystemSkillStore(tmp_path)
        toolset = SkillToolset(store)
        toolset.warm_up()
        assert set(toolset.skills) == {"pdf-forms"}
        assert toolset._store is store

    def test_load_skill_description_lists_skills(self, tmp_path):
        _write_skill(tmp_path, "pdf-forms", description="Use to fill PDF forms.")
        load_skill = _get_tool(SkillToolset(FileSystemSkillStore(tmp_path)), "load_skill")
        assert "Available skills:" in load_skill.description
        assert "- pdf-forms: Use to fill PDF forms." in load_skill.description

    def test_load_skill_description_when_empty(self, tmp_path):
        load_skill = _get_tool(SkillToolset(FileSystemSkillStore(tmp_path)), "load_skill")
        assert "No skills are currently available." in load_skill.description

    def test_load_skill_returns_body_and_manifest(self, tmp_path):
        _write_skill(
            tmp_path,
            "pdf-forms",
            description="Use to fill PDF forms.",
            body="Step 1. Do the thing.",
            files={"reference/forms.md": "details"},
        )
        load_skill = _get_tool(SkillToolset(FileSystemSkillStore(tmp_path)), "load_skill")
        result = load_skill.invoke(name="pdf-forms")
        assert "Step 1. Do the thing." in result
        assert "reference/forms.md" in result

    def test_load_skill_unknown_raises(self, tmp_path):
        _write_skill(tmp_path, "pdf-forms", description="Use to fill PDF forms.")
        load_skill = _get_tool(SkillToolset(FileSystemSkillStore(tmp_path)), "load_skill")
        # The error propagates (wrapped by Tool.invoke) so the Agent can apply its tool-failure policy.
        with pytest.raises(ToolInvocationError, match="Unknown skill 'nope'"):
            load_skill.invoke(name="nope")

    def test_read_skill_file(self, tmp_path):
        _write_skill(tmp_path, "pdf-forms", description="d", files={"reference/forms.md": "form details"})
        read = _get_tool(SkillToolset(FileSystemSkillStore(tmp_path)), "read_skill_file")
        assert read.invoke(name="pdf-forms", path="reference/forms.md") == "form details"

    def test_read_skill_file_returns_image_as_content_part(self, tmp_path):
        skill_dir = _write_skill(tmp_path, "pdf-forms", description="d")
        (skill_dir / "logo.png").write_bytes(b"\x89PNG\r\n\x1a\n\x00\xff\xfe")
        read = _get_tool(SkillToolset(FileSystemSkillStore(tmp_path)), "read_skill_file")
        # raw_result keeps the ImageContent intact (no string conversion) and the tool wraps it in a list so it
        # rides back as a multimodal tool-result content part.
        assert read.outputs_to_string == {"raw_result": True}
        result = read.invoke(name="pdf-forms", path="logo.png")
        assert isinstance(result, list)
        assert isinstance(result[0], ImageContent)

    def test_read_skill_file_returns_pdf_as_content_part(self, tmp_path):
        skill_dir = _write_skill(tmp_path, "pdf-forms", description="d")
        (skill_dir / "guide.pdf").write_bytes(b"%PDF-1.4 fake pdf bytes")
        read = _get_tool(SkillToolset(FileSystemSkillStore(tmp_path)), "read_skill_file")
        result = read.invoke(name="pdf-forms", path="guide.pdf")
        assert isinstance(result, list)
        assert isinstance(result[0], FileContent)

    def test_read_skill_file_in_agent_loop_yields_expected_tool_message_parts(self, tmp_path):
        # Use the Agent tool-execution step (`_run_tool`, what the Agent loop calls) to check that multimodal content
        # from `read_skill_file` comes back as expected in the tool-result messages.
        skill_dir = _write_skill(tmp_path, "pdf-forms", description="d", files={"reference/forms.md": "form details"})
        (skill_dir / "logo.png").write_bytes(b"\x89PNG\r\n\x1a\n\x00\xff\xfe")  # PNG magic header + non-UTF-8 byte
        (skill_dir / "guide.pdf").write_bytes(b"%PDF-1.4 fake pdf bytes")  # minimal PDF header
        toolset = SkillToolset(FileSystemSkillStore(tmp_path))

        tool_calls = [
            ToolCall("read_skill_file", {"name": "pdf-forms", "path": "reference/forms.md"}, id="text"),
            ToolCall("read_skill_file", {"name": "pdf-forms", "path": "logo.png"}, id="image"),
            ToolCall("read_skill_file", {"name": "pdf-forms", "path": "guide.pdf"}, id="pdf"),
        ]
        tool_messages, _ = _run_tool(
            messages=[ChatMessage.from_assistant(tool_calls=tool_calls)], state=State(schema={}), tools=toolset
        )

        assert len(tool_messages) == 3
        results = {}
        for message in tool_messages:
            assert message.is_from("tool")
            tool_result = message.tool_call_results[0]
            assert tool_result.error is False
            results[tool_result.origin.id] = tool_result.result

        # Text comes back as a plain string; images/PDFs as a one-element list of the matching content part.
        assert results["text"] == "form details"
        assert isinstance(results["image"], list) and isinstance(results["image"][0], ImageContent)
        assert results["image"][0].mime_type == "image/png"
        assert isinstance(results["pdf"], list) and isinstance(results["pdf"][0], FileContent)
        assert results["pdf"][0].mime_type == "application/pdf"

    def test_read_skill_file_blocks_traversal(self, tmp_path):
        _write_skill(tmp_path, "pdf-forms", description="d")
        (tmp_path / "secret.txt").write_text("top secret")
        read = _get_tool(SkillToolset(FileSystemSkillStore(tmp_path)), "read_skill_file")
        with pytest.raises(ToolInvocationError, match="outside the skill directory") as exc:
            read.invoke(name="pdf-forms", path="../secret.txt")
        assert "top secret" not in str(exc.value)

    def test_read_skill_file_missing(self, tmp_path):
        _write_skill(tmp_path, "pdf-forms", description="d")
        read = _get_tool(SkillToolset(FileSystemSkillStore(tmp_path)), "read_skill_file")
        with pytest.raises(ToolInvocationError, match="not found"):
            read.invoke(name="pdf-forms", path="nope.md")

    def test_to_dict_and_from_dict(self, tmp_path):
        _write_skill(tmp_path, "pdf-forms", description="Use to fill PDF forms.")
        toolset = SkillToolset(FileSystemSkillStore(tmp_path))

        data = toolset.to_dict()
        assert data == {
            "type": "haystack.tools.skills.skill_toolset.SkillToolset",
            "data": {
                "store": {
                    "type": "haystack.skill_stores.file_system.skill_store.FileSystemSkillStore",
                    "init_parameters": {"skills_dir": str(tmp_path)},
                }
            },
        }

        restored = SkillToolset.from_dict(data)
        restored.warm_up()
        assert set(restored.skills) == {"pdf-forms"}