stack/tests/prisma/test_tools.py

"""Tests for prisma.tools."""

from __future__ import annotations

import pytest

from prisma.project import STARTERS
from prisma.tools import (
    NOTES_FIELD,
    build_extract_tool,
    parse_extraction_fields,
)

_PALLIATIVE_TEMPLATE = STARTERS["palliative-rfi"][1]
_SKIN_SUBS_TEMPLATE = STARTERS["skin-subs"][1]


class TestImports:
    def test_importable(self):
        from prisma.tools import LLMTool

        assert LLMTool is not None


class TestParseExtractionFields:
    def test_reads_field_lines(self):
        props = parse_extraction_fields(
            "## Study metadata\n"
            "- **authors** — lead + et al.\n"
            "- **year** — publication year (YYYY)\n"
        )
        assert set(props) == {"authors", "year"}
        assert props["authors"]["type"] == "string"
        assert props["authors"]["description"] == "lead + et al."

    def test_ignores_prose_and_headings(self):
        props = parse_extraction_fields(
            "# Data extraction template\n\n"
            "The LLM will fill this form for every included study.\n\n"
            "## Population\n"
            "- **n_total** — total enrolled\n"
        )
        assert list(props) == ["n_total"]

    def test_bare_field_without_description(self):
        props = parse_extraction_fields("- **comparator_detail**\n")
        assert props["comparator_detail"] == {"type": "string"}

    def test_wrapped_description_is_joined(self):
        props = parse_extraction_fields(
            "- **design** — RCT / quasi-experimental / cohort /\n"
            "  systematic review / economic\n"
        )
        assert props["design"]["enum"] == [
            "RCT",
            "quasi-experimental",
            "cohort",
            "systematic review",
            "economic",
            "",
        ]

    def test_blank_line_ends_a_field(self):
        props = parse_extraction_fields(
            "- **n_total** — total enrolled\n"
            "\n"
            "  stray indented prose after a blank line\n"
        )
        assert props["n_total"]["description"] == "total enrolled"

    def test_heading_ends_a_field(self):
        props = parse_extraction_fields(
            "- **n_total** — total enrolled\n## Intervention\n"
        )
        assert props["n_total"]["description"] == "total enrolled"

    def test_later_definition_wins(self):
        props = parse_extraction_fields(
            "- **year** — first\n- **year** — second\n",
        )
        assert props["year"]["description"] == "second"


class TestEnumDetection:
    def test_short_closed_list_becomes_enum(self):
        props = parse_extraction_fields("- **overall** — low / some concerns / high\n")
        assert props["overall"]["enum"] == ["low", "some concerns", "high", ""]

    def test_empty_string_always_allowed(self):
        props = parse_extraction_fields("- **coi_declared** — yes / no / partial\n")
        assert "" in props["coi_declared"]["enum"]

    def test_long_members_stay_free_text(self):
        # "mean diff with 95% CI" is prose, not a code.
        props = parse_extraction_fields(
            "- **effect_size** — OR / RR / HR / mean diff with 95% CI\n"
        )
        assert "enum" not in props["effect_size"]

    def test_parenthetical_slash_is_not_a_separator(self):
        props = parse_extraction_fields(
            "- **medicare_relevance** — direct (Medicare data/population) / "
            "transferable / weak\n"
        )
        assert "enum" not in props["medicare_relevance"]

    def test_two_part_description_is_prose(self):
        props = parse_extraction_fields(
            "- **instrument_performance** — accuracy / discrimination as reported\n"
        )
        assert "enum" not in props["instrument_performance"]

    def test_no_slash_is_free_text(self):
        props = parse_extraction_fields("- **n_total** — total enrolled\n")
        assert "enum" not in props["n_total"]


class TestBuildExtractTool:
    def test_palliative_schema_follows_its_own_template(self):
        """The bug: palliative-rfi used to inherit skin-subs' fields."""
        tool = build_extract_tool(_PALLIATIVE_TEMPLATE, project="palliative-rfi")
        props = tool.schema["properties"]
        assert {"rfi_parts", "eligibility_basis", "fwa_findings"} <= set(props)
        assert not {"wound_type", "product_category", "cost_per_healed"} & set(props)

    def test_skin_subs_schema_still_has_its_fields(self):
        tool = build_extract_tool(_SKIN_SUBS_TEMPLATE, project="skin-subs")
        props = tool.schema["properties"]
        assert {"wound_type", "product_name", "n_total"} <= set(props)
        assert props["wound_type"]["enum"] == [
            "DFU",
            "VLU",
            "pressure",
            "burn",
            "mixed",
            "",
        ]

    def test_notes_field_present_and_required(self):
        tool = build_extract_tool(_PALLIATIVE_TEMPLATE, project="palliative-rfi")
        assert NOTES_FIELD in tool.schema["properties"]
        assert tool.schema["required"] == [NOTES_FIELD]

    def test_template_notes_field_is_not_clobbered(self):
        tool = build_extract_tool(
            "- **extraction_notes** — reviewer caveats only\n", project="x"
        )
        assert (
            tool.schema["properties"][NOTES_FIELD]["description"]
            == "reviewer caveats only"
        )

    def test_tool_name_is_stable(self):
        tool = build_extract_tool(_PALLIATIVE_TEMPLATE, project="palliative-rfi")
        assert tool.name == "extract_study_data"
        assert "palliative-rfi" in tool.description

    def test_empty_template_raises(self):
        with pytest.raises(ValueError, match="defines no"):
            build_extract_tool("# Nothing here\n\nJust prose.\n", project="broken")