rasahq--rasa/tests/shared/nlu/training_data/test_entities_parser.py

from typing import Text, List, Dict, Any

import pytest

import rasa.shared.nlu.training_data.entities_parser as entities_parser
from rasa.shared.exceptions import InvalidEntityFormatException, SchemaValidationError
from rasa.shared.nlu.constants import TEXT


@pytest.mark.parametrize(
    "example, expected_entities, expected_text",
    [
        (
            "I need an [economy class](travel_flight_class:economy) ticket from "
            '[boston]{"entity": "city", "role": "from"} to [new york]{"entity": "city",'
            ' "role": "to"}, please.',
            [
                {
                    "start": 10,
                    "end": 23,
                    "value": "economy",
                    "entity": "travel_flight_class",
                },
                {
                    "start": 36,
                    "end": 42,
                    "value": "boston",
                    "entity": "city",
                    "role": "from",
                },
                {
                    "start": 46,
                    "end": 54,
                    "value": "new york",
                    "entity": "city",
                    "role": "to",
                },
            ],
            "I need an economy class ticket from boston to new york, please.",
        ),
        ("i'm looking for a place to eat", [], "i'm looking for a place to eat"),
        (
            "i'm looking for a place in the [north](loc-direction) of town",
            [{"start": 31, "end": 36, "value": "north", "entity": "loc-direction"}],
            "i'm looking for a place in the north of town",
        ),
        (
            "show me [chines](cuisine:chinese) restaurants",
            [{"start": 8, "end": 14, "value": "chinese", "entity": "cuisine"}],
            "show me chines restaurants",
        ),
        (
            'show me [italian]{"entity": "cuisine", "value": "22_ab-34*3.A:43er*+?df"} '
            "restaurants",
            [
                {
                    "start": 8,
                    "end": 15,
                    "value": "22_ab-34*3.A:43er*+?df",
                    "entity": "cuisine",
                }
            ],
            "show me italian restaurants",
        ),
        ("Do you know {ABC} club?", [], "Do you know {ABC} club?"),
        (
            "show me [chines](22_ab-34*3.A:43er*+?df) restaurants",
            [{"start": 8, "end": 14, "value": "43er*+?df", "entity": "22_ab-34*3.A"}],
            "show me chines restaurants",
        ),
        (
            'I want to fly from [Berlin]{"entity": "city", "role": "to"} to [LA]{'
            '"entity": "city", "role": "from", "value": "Los Angeles"}',
            [
                {
                    "start": 19,
                    "end": 25,
                    "value": "Berlin",
                    "entity": "city",
                    "role": "to",
                },
                {
                    "start": 29,
                    "end": 31,
                    "value": "Los Angeles",
                    "entity": "city",
                    "role": "from",
                },
            ],
            "I want to fly from Berlin to LA",
        ),
        (
            'I want to fly from [Berlin](city) to [LA]{"entity": "city", "role": '
            '"from", "value": "Los Angeles"}',
            [
                {"start": 19, "end": 25, "value": "Berlin", "entity": "city"},
                {
                    "start": 29,
                    "end": 31,
                    "value": "Los Angeles",
                    "entity": "city",
                    "role": "from",
                },
            ],
            "I want to fly from Berlin to LA",
        ),
        (
            'I want to fly from [Berlin](city) to [LA][{"entity": "city", "role": '
            '"from", "value": "Los Angeles"}, {"entity": "location", "value": '
            '"Los Angeles"}]',
            [
                {"start": 19, "end": 25, "value": "Berlin", "entity": "city"},
                {
                    "start": 29,
                    "end": 31,
                    "value": "Los Angeles",
                    "entity": "city",
                    "role": "from",
                },
                {
                    "start": 29,
                    "end": 31,
                    "value": "Los Angeles",
                    "entity": "location",
                },
            ],
            "I want to fly from Berlin to LA",
        ),
    ],
)
def test_markdown_entity_regex(
    example: Text, expected_entities: List[Dict[Text, Any]], expected_text: Text
):

    result = entities_parser.find_entities_in_training_example(example)
    assert result == expected_entities

    replaced_text = entities_parser.replace_entities(example)
    assert replaced_text == expected_text


def test_parse_training_example():
    message = entities_parser.parse_training_example("Hello!", intent="greet")
    assert message.get("intent") == "greet"
    assert message.get(TEXT) == "Hello!"


def test_parse_empty_example():
    message = entities_parser.parse_training_example("")
    assert message.get("intent") is None
    assert message.get(TEXT) == ""


def test_parse_training_example_with_entities():
    message = entities_parser.parse_training_example(
        "I am from [Berlin](city).", intent="inform"
    )
    assert message.get("intent") == "inform"
    assert message.get(TEXT) == "I am from Berlin."
    assert message.get("entities") == [
        {"start": 10, "end": 16, "value": "Berlin", "entity": "city"}
    ]


def test_markdown_entity_regex_error_handling_not_json():
    with pytest.raises(InvalidEntityFormatException):
        entities_parser.find_entities_in_training_example(
            # JSON syntax error: missing closing " for `role`
            'I want to fly from [Berlin]{"entity": "city", "role: "from"}'
        )


def test_markdown_entity_regex_error_handling_wrong_schema():
    with pytest.raises(SchemaValidationError):
        entities_parser.find_entities_in_training_example(
            # Schema error: "entiti" instead of "entity"
            'I want to fly from [Berlin]{"entiti": "city", "role": "from"}'
        )