rasahq--rasa/tests/engine/training/test_components.py

import dataclasses
from typing import Text
import uuid


from rasa.engine.caching import TrainingCache
from rasa.engine.graph import ExecutionContext, GraphNode, GraphSchema, SchemaNode
from rasa.engine.storage.resource import Resource
from rasa.engine.storage.storage import ModelStorage
from rasa.engine.training import fingerprinting
from rasa.engine.training.components import (
    PrecomputedValueProvider,
    FingerprintComponent,
    FingerprintStatus,
)
from tests.engine.graph_components_test_classes import CacheableText


def test_cached_component_returns_value_from_cache(default_model_storage: ModelStorage):

    cached_output = CacheableText("Cache me!!")

    node = GraphNode(
        node_name="cached",
        component_class=PrecomputedValueProvider,
        constructor_name="create",
        component_config={"output": cached_output},
        fn_name="get_value",
        inputs={},
        eager=False,
        model_storage=default_model_storage,
        resource=None,
        execution_context=ExecutionContext(GraphSchema({}), "1"),
    )

    node_name, returned_output = node()

    assert node_name == "cached"
    assert returned_output.text == "Cache me!!"


def test_cached_component_replace_schema_node():
    schema_node = SchemaNode(
        needs={"i1": "first_input", "i2": "second_input"},
        uses=FingerprintComponent,
        fn="add",
        constructor_name="load",
        config={"a": 1},
        eager=False,
        is_input=False,
        resource=Resource("hello"),
    )

    PrecomputedValueProvider.replace_schema_node(schema_node, 2)

    assert schema_node == SchemaNode(
        needs={"i1": "first_input", "i2": "second_input"},
        uses=PrecomputedValueProvider,
        fn="get_value",
        constructor_name="create",
        config={"output": 2},
        eager=False,
        is_input=False,
        resource=Resource("hello"),
    )


def test_fingerprint_component_replace_schema_node(temp_cache: TrainingCache):
    schema_node = SchemaNode(
        needs={"i1": "first_input", "i2": "second_input"},
        uses=PrecomputedValueProvider,
        fn="add",
        constructor_name="load",
        config={"a": 1},
        eager=False,
        is_input=False,
        resource=Resource("hello"),
    )

    FingerprintComponent.replace_schema_node(schema_node, temp_cache)

    assert schema_node == SchemaNode(
        needs={"i1": "first_input", "i2": "second_input"},
        uses=FingerprintComponent,
        fn="run",
        constructor_name="create",
        config={
            "config_of_replaced_component": {"a": 1},
            "cache": temp_cache,
            "graph_component_class": PrecomputedValueProvider,
        },
        eager=True,
        is_input=False,
        resource=Resource("hello"),
    )


@dataclasses.dataclass
class FingerprintableText:
    text: Text

    def fingerprint(self) -> Text:
        return self.text


def test_fingerprint_component_hit(
    default_model_storage: ModelStorage, temp_cache: TrainingCache
):

    cached_output = CacheableText("Cache me!!")
    output_fingerprint = uuid.uuid4().hex

    # We generate a fingerprint key that will match the one generated by the
    # `FingerprintComponent`.
    component_config = {"x": 1}
    fingerprint_key = fingerprinting.calculate_fingerprint_key(
        graph_component_class=PrecomputedValueProvider,
        config=component_config,
        inputs={
            "param_1": FingerprintableText("input_1"),
            "param_2": FingerprintableText("input_2"),
        },
    )
    # We cache the output using this fingerprint key.
    temp_cache.cache_output(
        fingerprint_key=fingerprint_key,
        output=cached_output,
        output_fingerprint=output_fingerprint,
        model_storage=default_model_storage,
    )

    # The node inputs and config match what we used to generate the fingerprint key.
    node = GraphNode(
        node_name="fingerprint_node",
        component_class=FingerprintComponent,
        constructor_name="create",
        component_config={
            "config_of_replaced_component": component_config,
            "cache": temp_cache,
            "graph_component_class": PrecomputedValueProvider,
        },
        fn_name="run",
        inputs={"param_1": "parent_node_1", "param_2": "parent_node_2"},
        eager=False,
        model_storage=default_model_storage,
        resource=None,
        execution_context=ExecutionContext(GraphSchema({}), "1"),
    )

    node_name, returned_output = node(
        ("parent_node_1", FingerprintableText("input_1")),
        ("parent_node_2", FingerprintStatus(is_hit=True, output_fingerprint="input_2")),
    )

    assert node_name == "fingerprint_node"
    assert returned_output.is_hit is True
    assert returned_output.output_fingerprint == output_fingerprint
    assert returned_output.output_fingerprint == returned_output.fingerprint()


def test_fingerprint_component_miss(
    default_model_storage: ModelStorage, temp_cache: TrainingCache
):

    component_config = {"x": 1}

    node = GraphNode(
        node_name="fingerprint_node",
        component_class=FingerprintComponent,
        constructor_name="create",
        component_config={
            "config_of_replaced_component": component_config,
            "cache": temp_cache,
            "graph_component_class": PrecomputedValueProvider,
        },
        fn_name="run",
        inputs={"param_1": "parent_node_1", "param_2": "parent_node_2"},
        eager=False,
        model_storage=default_model_storage,
        resource=None,
        execution_context=ExecutionContext(GraphSchema({}), "1"),
    )

    node_name, returned_output = node(
        ("parent_node_1", FingerprintableText("input_1")),
        ("parent_node_2", FingerprintStatus(is_hit=True, output_fingerprint="input_2")),
    )

    # As we didnt add anything to the cache, it cannot be a hit.
    assert node_name == "fingerprint_node"
    assert returned_output.is_hit is False
    assert returned_output.output_fingerprint is None
    assert returned_output.fingerprint() != returned_output.output_fingerprint
    assert returned_output.fingerprint() != returned_output.fingerprint()