rasa-pro 3.12.0.dev1__py3-none-any.whl

rasa/llm_fine_tuning/annotation_module.py

@@ -0,0 +1,241 @@
+import asyncio
+from contextlib import contextmanager
+from typing import List, Generator, Optional, Tuple, Union
+
+import structlog
+
+from rasa.dialogue_understanding.commands import Command
+from rasa.e2e_test.e2e_test_case import TestSuite, TestCase, ActualStepOutput, TestStep
+from rasa.e2e_test.e2e_test_runner import E2ETestRunner, TEST_TURNS_TYPE
+from rasa.llm_fine_tuning.conversations import Conversation, ConversationStep
+from rasa.llm_fine_tuning.storage import StorageContext
+from rasa.shared.core.constants import USER
+from rasa.shared.core.trackers import DialogueStateTracker
+from rasa.shared.nlu.constants import LLM_PROMPT, LLM_COMMANDS
+from rasa.shared.utils.llm import tracker_as_readable_transcript
+
+ANNOTATION_MODULE_STORAGE_LOCATION = "1_command_annotations"
+
+preparing_fine_tuning_data = False
+
+structlogger = structlog.get_logger()
+
+
+@contextmanager
+def set_preparing_fine_tuning_data() -> Generator:
+    global preparing_fine_tuning_data
+    preparing_fine_tuning_data = True
+    try:
+        yield
+    finally:
+        preparing_fine_tuning_data = False
+
+
+def annotate_e2e_tests(
+    e2e_test_runner: E2ETestRunner,
+    test_suite: TestSuite,
+    storage_context: StorageContext,
+) -> List[Conversation]:
+    with set_preparing_fine_tuning_data():
+        converations = asyncio.run(
+            e2e_test_runner.run_tests_for_fine_tuning(
+                test_suite.test_cases,
+                test_suite.fixtures,
+                test_suite.metadata,
+            )
+        )
+
+    storage_context.write_conversations(
+        converations, ANNOTATION_MODULE_STORAGE_LOCATION
+    )
+
+    return converations
+
+
+def _get_previous_actual_step_output(
+    test_turns: TEST_TURNS_TYPE, i: int
+) -> Optional[ActualStepOutput]:
+    while i > 0:
+        i = i - 1
+        if isinstance(test_turns[i], ActualStepOutput):
+            return test_turns[i]  # type:ignore[return-value]
+    return None
+
+
+def generate_conversation(
+    test_turns: TEST_TURNS_TYPE,
+    test_case: TestCase,
+    tracker: DialogueStateTracker,
+    assertions_used: bool = False,
+) -> Optional[Conversation]:
+    """Generates a conversation object in case of e2e test passing.
+
+    Args:
+        test_turns: the turns that happened when running the test case or test step.
+        test_case: the `TestCase` instance.
+        tracker: the dialogue state tracker.
+        assertions_used: if True the e2e test format with assertions was used.
+
+    Returns:
+        Conversation.
+    """
+    steps = []
+
+    if assertions_used:
+        # we only have user steps, extract the bot response from the bot uttered
+        # events of the test turn
+        for i, original_step in enumerate(test_case.steps):
+            previous_turn = _get_previous_actual_step_output(test_turns, i)
+            steps.append(
+                _convert_to_conversation_step(
+                    original_step, test_turns[i], test_case.name, previous_turn
+                )
+            )
+            steps.extend(_create_bot_test_steps(test_turns[i]))
+    else:
+        for i, original_step in enumerate(test_case.steps):
+            if original_step.actor == USER:
+                previous_turn = _get_previous_actual_step_output(test_turns, i)
+                steps.append(
+                    _convert_to_conversation_step(
+                        original_step, test_turns[i], test_case.name, previous_turn
+                    )
+                )
+            else:
+                steps.append(original_step)
+
+    # Some messages in an e2e test case could be mapped to commands via
+    # 'NLUCommandAdapter', e.g. the message will not be annotated with a prompt and
+    # commands pair. Only convert steps that have a prompt and commands present into a
+    # ConversationStep.
+    # The conversation needs to have at least one 'ConversationStep' to be valid for
+    # fine-tuning.
+    if not any([isinstance(step, ConversationStep) for step in steps]):
+        structlogger.warning(
+            "annotation_module.skip_test_case.missing_llm_commands_and_prompts",
+            test_case=test_case.name,
+            file=test_case.file,
+        )
+        return None
+
+    transcript = tracker_as_readable_transcript(tracker, max_turns=None)
+
+    return Conversation(test_case.name, test_case, steps, transcript)
+
+
+def _create_bot_test_steps(current_turn: ActualStepOutput) -> List[TestStep]:
+    test_steps = []
+    for bot_event in current_turn.bot_uttered_events:
+        template = None
+        if "utter_action" in bot_event.metadata:
+            template = bot_event.metadata["utter_action"]
+
+        test_steps.append(TestStep(actor="bot", text=bot_event.text, template=template))
+
+    return test_steps
+
+
+def _convert_to_conversation_step(
+    current_step: TestStep,
+    current_turn: ActualStepOutput,
+    test_case_name: str,
+    previous_turn: Optional[ActualStepOutput],
+) -> Union[TestStep, ConversationStep]:
+    if not current_step.text == current_turn.text or not isinstance(
+        current_turn, ActualStepOutput
+    ):
+        # There should be a one to one mapping between test steps (steps read from file)
+        # and test turns (test result of e2e test). Verify that the current step is
+        # aligned with the current turn.
+        structlogger.debug(
+            "annotation_module.convert_to_conversation_step.skip_user_message",
+            test_case=test_case_name,
+            user_message=current_step.text,
+        )
+        return current_step
+
+    llm_prompt, llm_commands = _extract_llm_prompt_and_commands(current_turn)
+    if not llm_commands or not llm_prompt:
+        # If no commands or no prompt is present we cannot create a data point
+        # for fine-tuning, skipping this step.
+        structlogger.debug(
+            "annotation_module.convert_to_conversation_step.skip_user_message",
+            test_case=test_case_name,
+            user_message=current_step.text,
+            message="No commands/prompt associated with the message.",
+        )
+        return current_step
+
+    commands = [Command.command_from_json(data) for data in llm_commands]
+    rephrase = _should_be_rephrased(current_turn.text, previous_turn, test_case_name)
+
+    return ConversationStep(current_step, commands, llm_prompt, rephrase=rephrase)
+
+
+def _should_be_rephrased(
+    current_user_message: str,
+    previous_turn: Optional[ActualStepOutput],
+    test_case_name: str,
+) -> bool:
+    """Checks if the current user message should be rephrased or not.
+
+    A user message should not be rephrased in case the user message comes from a button
+    payload, i.e. the user clicked on a button.
+
+    Args:
+        current_user_message: The current user message.
+        previous_turn: The previous turn containing the bot uttered event that came
+            before.
+        test_case_name: The name of the test case.
+
+    Returns:
+        True, in case the user message should be rephrased, False otherwise.
+    """
+    # there is no previous turn, we are at the beginning of the conversation
+    if not previous_turn:
+        return True
+
+    buttons_present = (
+        previous_turn.bot_uttered_events
+        and "buttons" in previous_turn.bot_uttered_events[-1].data
+        and previous_turn.bot_uttered_events[-1].data["buttons"] is not None
+    )
+
+    if not buttons_present:
+        return True
+
+    # if the user utterance comes from a button payload we should not rephrase
+    # the user utterance in later steps
+    button_data = previous_turn.bot_uttered_events[-1].data["buttons"]
+    button_payloads = [data["payload"].lower() for data in button_data]
+    if current_user_message.lower() in button_payloads:
+        structlogger.debug(
+            "annotation_module.user_message_should_not_be_rephrased",
+            rephrase=False,
+            user_message=current_user_message,
+            test_case_name=test_case_name,
+        )
+        return False
+
+    return True
+
+
+def _extract_llm_prompt_and_commands(
+    turn: ActualStepOutput,
+) -> Tuple[Optional[str], Optional[str]]:
+    # There should be exactly one 'UserUttered' event
+    if not turn.user_uttered_events or len(turn.user_uttered_events) != 1:
+        return None, None
+
+    # Check if 'parse_data' contains the prompt and the commands
+    if (
+        not turn.user_uttered_events[0].parse_data
+        or LLM_PROMPT not in turn.user_uttered_events[0].parse_data
+        or LLM_COMMANDS not in turn.user_uttered_events[0].parse_data
+    ):
+        return None, None
+
+    return (
+        turn.user_uttered_events[0].parse_data[LLM_PROMPT],
+        turn.user_uttered_events[0].parse_data[LLM_COMMANDS],
+    )

rasa/llm_fine_tuning/conversations.py

@@ -0,0 +1,144 @@
+from dataclasses import dataclass, field
+from typing import List, Union, Iterator, Any, Dict, Optional
+
+from rasa.dialogue_understanding.commands import (
+    Command,
+    StartFlowCommand,
+    SetSlotCommand,
+    CancelFlowCommand,
+    ChitChatAnswerCommand,
+    SkipQuestionCommand,
+    HumanHandoffCommand,
+    KnowledgeAnswerCommand,
+    ClarifyCommand,
+)
+from rasa.e2e_test.e2e_test_case import TestCase, TestStep
+from rasa.shared.core.constants import USER
+
+
+@dataclass
+class ConversationStep:
+    original_test_step: TestStep
+    llm_commands: List[Command]
+    llm_prompt: str
+    failed_rephrasings: List[str] = field(default_factory=list)
+    passed_rephrasings: List[str] = field(default_factory=list)
+    rephrase: bool = True
+
+    def as_dict(self) -> Dict[str, Any]:
+        data = {
+            "user": self.original_test_step.text,
+            "llm_commands": self._commands_to_str(),
+        }
+        if self.passed_rephrasings:
+            data["passing_rephrasings"] = self.passed_rephrasings
+        if self.failed_rephrasings:
+            data["failing_rephrasings"] = self.failed_rephrasings
+
+        return data
+
+    def _commands_to_str(self) -> List[str]:
+        output = []
+        for command in self.llm_commands:
+            if isinstance(command, StartFlowCommand):
+                output.append(f"StartFlow({command.flow})")
+            elif isinstance(command, SetSlotCommand):
+                output.append(f"SetSlot({command.name}, {command.value})")
+            elif isinstance(command, ClarifyCommand):
+                output.append(f"Clarify({command.options})")
+            elif isinstance(command, CancelFlowCommand):
+                output.append("CancelFlow()")
+            elif isinstance(command, ChitChatAnswerCommand):
+                output.append("ChitChat()")
+            elif isinstance(command, SkipQuestionCommand):
+                output.append("SkipQuestion()")
+            elif isinstance(command, KnowledgeAnswerCommand):
+                output.append("SearchAndReply()")
+            elif isinstance(command, HumanHandoffCommand):
+                output.append("HumanHandoff()")
+        return output
+
+    def commands_as_string(self) -> str:
+        return "\n".join(self._commands_to_str())
+
+
+@dataclass
+class Conversation:
+    name: str
+    original_e2e_test_case: TestCase
+    steps: List[Union[TestStep, ConversationStep]]
+    transcript: str
+
+    def iterate_over_annotated_user_steps(
+        self, rephrase: Optional[bool] = None
+    ) -> Iterator[ConversationStep]:
+        """Iterate over conversation steps.
+
+        Yield each step based on the rephrase parameter.
+
+        Args:
+            rephrase: Determines whether to yield steps based on their `rephrase`
+                attribute. Can be:
+                - None: Yield all ConversationStep instances regardless of their
+                    rephrase attribute.
+                - True: Yield only those ConversationStep instances where the
+                    rephrase attribute is True.
+                - False: Yield only those ConversationStep instances where the
+                    rephrase attribute is False.
+
+        Yields:
+            ConversationStep: The next conversation step that matches the specified
+            rephrase condition.
+        """
+        for step in self.steps:
+            if isinstance(step, ConversationStep):
+                if rephrase is None:
+                    yield step
+                elif rephrase is not None and step.rephrase == rephrase:
+                    yield step
+
+    def get_user_messages(self) -> List[str]:
+        return [
+            step.original_test_step.text
+            for step in self.iterate_over_annotated_user_steps()
+            if step.original_test_step.text
+        ]
+
+    def get_user_messages_to_rephrase(self) -> List[str]:
+        return [
+            step.original_test_step.text
+            for step in self.iterate_over_annotated_user_steps(rephrase=True)
+            if step.original_test_step.text
+        ]
+
+    def as_dict(self) -> Dict[str, Any]:
+        steps_data = []
+        for step in self.steps:
+            if isinstance(step, ConversationStep):
+                steps_data.append(step.as_dict())
+            elif isinstance(step, TestStep):
+                if step.text and step.actor == USER:
+                    steps_data.append({"user": step.text})
+                elif step.template:
+                    steps_data.append({"utter": step.template})
+                elif step.text:
+                    steps_data.append({"bot": step.text})
+
+        return {
+            "conversations": [
+                {
+                    "original_test_case": f"{self.original_e2e_test_case.file}::"
+                    f"{self.original_e2e_test_case.name}",
+                    "steps": steps_data,
+                }
+            ]
+        }
+
+    def get_number_of_rephrases(self, passing: bool) -> int:
+        return sum(
+            len(step.passed_rephrasings) if passing else len(step.failed_rephrasings)
+            for step in self.iterate_over_annotated_user_steps()
+        )
+
+    def get_full_name(self) -> str:
+        return f"{self.original_e2e_test_case.file}::{self.name}"

rasa/llm_fine_tuning/llm_data_preparation_module.py

@@ -0,0 +1,178 @@
+from dataclasses import dataclass
+from typing import List, Dict, Any, Optional
+
+import structlog
+from tqdm import tqdm
+
+from rasa.llm_fine_tuning.conversations import Conversation, ConversationStep
+from rasa.llm_fine_tuning.storage import StorageContext
+
+LLM_DATA_PREPARATION_MODULE_STORAGE_LOCATION = "3_llm_finetune_data/llm_ft_data.jsonl"
+
+structlogger = structlog.get_logger()
+
+
+@dataclass
+class LLMDataExample:
+    prompt: str
+    output: str
+    original_test_name: str
+    original_user_utterance: str
+    rephrased_user_utterance: str
+
+    def as_dict(self) -> Dict[str, Any]:
+        return {
+            "prompt": self.prompt,
+            "output": self.output,
+            "original_test_name": self.original_test_name,
+            "original_user_utterance": self.original_user_utterance,
+            "rephrased_user_utterance": self.rephrased_user_utterance,
+        }
+
+
+def _create_data_point(
+    prompt: str,
+    step: ConversationStep,
+    conversation: Conversation,
+    rephrased_user_message: Optional[str] = None,
+) -> LLMDataExample:
+    return LLMDataExample(
+        prompt,
+        step.commands_as_string(),
+        conversation.get_full_name(),
+        step.original_test_step.text,
+        rephrased_user_message,
+    )
+
+
+def _update_prompt(
+    prompt: str,
+    original_user_steps: List[ConversationStep],
+    rephrased_user_steps: List[str],
+) -> Optional[str]:
+    if len(original_user_steps) != len(rephrased_user_steps):
+        structlogger.debug(
+            "llm_fine_tuning.llm_data_preparation_module.failed_to_update_prompt",
+            original_user_steps=[
+                step.original_test_step.text for step in original_user_steps
+            ],
+            rephrased_user_steps=rephrased_user_steps,
+        )
+        return None
+
+    updated_prompt = prompt
+    for user_step, rephrased_message in zip(original_user_steps, rephrased_user_steps):
+        # replace all occurrences of the original user message with the rephrased user
+        # message in the conversation history mentioned in the prompt
+        updated_prompt = updated_prompt.replace(
+            f"USER: {user_step.original_test_step.text}", f"USER: {rephrased_message}"
+        )
+
+    # replace the latest user message mentioned in the prompt
+    updated_prompt = updated_prompt.replace(
+        f"'''{original_user_steps[-1].original_test_step.text}'''",
+        f"'''{rephrased_user_steps[-1]}'''",
+    )
+
+    return updated_prompt
+
+
+def _convert_conversation_into_llm_data(
+    conversation: Conversation,
+) -> List[LLMDataExample]:
+    data = []
+
+    # construct new conversations from the rephrasings
+    new_conversations = _construct_new_conversations(conversation)
+
+    original_user_steps = [
+        step for step in conversation.iterate_over_annotated_user_steps()
+    ]
+
+    for i, step in enumerate(original_user_steps):
+        # create data point for the original e2e test case
+        data.append(_create_data_point(step.llm_prompt, step, conversation))
+
+        # create data points using the rephrasings, e.g. 'new_conversations'
+        for rephrased_user_steps in new_conversations:
+            # +1 to include the current user turn
+            prompt = _update_prompt(
+                step.llm_prompt,
+                original_user_steps[: i + 1],
+                rephrased_user_steps[: i + 1],
+            )
+            if prompt:
+                data.append(
+                    _create_data_point(
+                        prompt, step, conversation, rephrased_user_steps[i]
+                    )
+                )
+
+    return data
+
+
+def _construct_new_conversations(conversation: Conversation) -> List[List[str]]:
+    """Construct new conversations from the rephrasings.
+
+    In general, we will combine the passing rephrasings at the same index position to
+    construct a new conversation. If for one particular user turn no other passing
+    rephrasing exists, we reset the index and take the first passing rephrasing again.
+
+    Args:
+        conversation: The conversation.
+
+    Returns:
+        A list of new conversations (only rephrased user turns).
+    """
+    max_passed_rephrasings = max(
+        [
+            len(step.passed_rephrasings)
+            for step in conversation.iterate_over_annotated_user_steps()
+        ]
+    )
+
+    new_conversations = []
+    for i in range(0, max_passed_rephrasings):
+        current_conversation = []
+        for step in conversation.iterate_over_annotated_user_steps():
+            # take the orginial user message in case no passing rephrasings exist
+            if not step.passed_rephrasings and step.original_test_step.text:
+                structlogger.debug(
+                    "llm_fine_tuning.llm_data_preparation_module."
+                    "construct_new_conversations.no_passed_rephrasings",
+                    conversation=conversation.get_full_name(),
+                    step=step.original_test_step.text,
+                    message="Take original user message instead of rephrasing.",
+                )
+                current_conversation.append(step.original_test_step.text)
+                continue
+
+            # some user steps might have less rephrasings than others
+            # loop over the rephrasings
+            index = i % len(step.passed_rephrasings)
+            current_conversation.append(step.passed_rephrasings[index])
+        if current_conversation:
+            new_conversations.append(current_conversation)
+
+    structlogger.debug(
+        "llm_fine_tuning.llm_data_preparation_module.construct_new_conversations",
+        conversation=conversation.get_full_name(),
+        new_conversations=new_conversations,
+    )
+
+    return new_conversations
+
+
+def convert_to_fine_tuning_data(
+    conversations: List[Conversation], storage_context: StorageContext
+) -> List[LLMDataExample]:
+    llm_data = []
+
+    for i in tqdm(range(len(conversations))):
+        llm_data.extend(_convert_conversation_into_llm_data(conversations[i]))
+
+    storage_context.write_llm_data(
+        llm_data, LLM_DATA_PREPARATION_MODULE_STORAGE_LOCATION
+    )
+
+    return llm_data