rasa-pro 3.12.4__py3-none-any.whl → 3.13.0.dev1__py3-none-any.whl

rasa/dialogue_understanding/generator/_jinja_filters.py

@@ -0,0 +1,9 @@
+import json
+
+
+def to_json_escaped_string(s: str) -> str:
+    """
+    Serializes the input string using JSON encoding to escape special characters
+    like newlines (\\n), tabs (\\t), and quotes.
+    """
+    return json.dumps(s, ensure_ascii=False)

rasa/dialogue_understanding/generator/constants.py

@@ -33,3 +33,7 @@ LLM_BASED_COMMAND_GENERATOR_CONFIG_FILE = "config.json"
 
 MODEL_NAME_GPT_4O_2024_11_20 = "gpt-4o-2024-11-20"
 MODEL_NAME_CLAUDE_3_5_SONNET_20240620 = "claude-3-5-sonnet-20240620"
+
+# JINJA template filters
+
+TO_JSON_ESCAPED_STRING_JINJA_FILTER = "to_json_escaped_string"

rasa/dialogue_understanding/generator/llm_based_command_generator.py

@@ -3,7 +3,7 @@ from functools import lru_cache
 from typing import Any, Dict, List, Optional, Set, Text, Tuple, Union
 
 import structlog
-from jinja2 import Template
+from jinja2 import Environment, Template, select_autoescape
 
 import rasa.dialogue_understanding.generator.utils
 import rasa.shared.utils.io
@@ -17,12 +17,14 @@ from rasa.dialogue_understanding.commands.handle_digressions_command import (
 )
 from rasa.dialogue_understanding.constants import KEY_MINIMIZE_NUM_CALLS
 from rasa.dialogue_understanding.generator import CommandGenerator
+from rasa.dialogue_understanding.generator._jinja_filters import to_json_escaped_string
 from rasa.dialogue_understanding.generator.constants import (
     DEFAULT_LLM_CONFIG,
     FLOW_RETRIEVAL_ACTIVE_KEY,
     FLOW_RETRIEVAL_FLOW_THRESHOLD,
     FLOW_RETRIEVAL_KEY,
     LLM_CONFIG_KEY,
+    TO_JSON_ESCAPED_STRING_JINJA_FILTER,
 )
 from rasa.dialogue_understanding.generator.flow_retrieval import FlowRetrieval
 from rasa.dialogue_understanding.stack.utils import top_flow_frame
@@ -226,12 +228,25 @@ class LLMBasedCommandGenerator(
 
     @lru_cache
     def compile_template(self, template: str) -> Template:
-        """Compile the prompt template.
+        """
+        Compile the prompt template and register custom filters.
 
         Compiling the template is an expensive operation,
         so we cache the result.
         """
-        return Template(template)
+        # Create an environment
+        # Autoescaping disabled explicitly for LLM prompt templates rendered from
+        # strings (safe, not HTML)
+        env = Environment(
+            autoescape=select_autoescape(
+                disabled_extensions=["jinja2"], default_for_string=False, default=True
+            )
+        )
+        # Register filters
+        env.filters[TO_JSON_ESCAPED_STRING_JINJA_FILTER] = to_json_escaped_string
+
+        # Return the template which can leverage registered filters
+        return env.from_string(template)
 
     @classmethod
     def load_prompt_template_from_model_storage(

rasa/dialogue_understanding/generator/nlu_command_adapter.py

@@ -140,7 +140,7 @@ class NLUCommandAdapter(GraphComponent, CommandGenerator):
 
         if commands:
             commands = clean_up_commands(
-                commands, tracker, flows, self._execution_context
+                commands, tracker, flows, self._execution_context, domain
             )
             log_llm(
                 logger=structlogger,

rasa/dialogue_understanding/generator/prompt_templates/command_prompt_v2_claude_3_5_sonnet_20240620_template.jinja2

@@ -8,7 +8,7 @@ Your task is to analyze the current conversation context and generate a list of
 * `set slot slot_name slot_value`: Slot setting. For example, `set slot transfer_money_recipient Freddy`. Can be used to correct and change previously set values.
 * `cancel flow`: Cancelling the current flow.
 * `disambiguate flows flow_name1 flow_name2 ... flow_name_n`: Disambiguate which flow should be started when user input is ambiguous by listing the potential flows as options. For example, `disambiguate flows list_contacts add_contact remove_contact ...` if the user just wrote "contacts".
-* `provide info`: Responding to the user's questions by supplying relevant information, such as answering FAQs or explaining services.
+* `search and reply`: Responding to the user's questions by supplying relevant information, such as answering FAQs or explaining services.
 * `offtopic reply`: Responding to casual or social user messages that are unrelated to any flows, engaging in friendly conversation and addressing off-topic remarks.
 * `hand over`: Handing over to a human, in case the user seems frustrated or explicitly asks to speak to one.
 
@@ -33,7 +33,7 @@ Your task is to analyze the current conversation context and generate a list of
 ## Available Flows and Slots
 Use the following structured data:
 ```json
-{"flows":[{% for flow in available_flows %}{"name":"{{ flow.name }}","description":"{{ flow.description }}"{% if flow.slots %},"slots":[{% for slot in flow.slots %}{"name":"{{ slot.name }}"{% if slot.description %},"description":"{{ slot.description }}"{% endif %}{% if slot.allowed_values %},"allowed_values":{{ slot.allowed_values }}{% endif %}}{% if not loop.last %},{% endif %}{% endfor %}]{% endif %}}{% if not loop.last %},{% endif %}{% endfor %}]}
+{"flows":[{% for flow in available_flows %}{"name":"{{ flow.name }}","description":{{ flow.description | to_json_escaped_string }}{% if flow.slots %},"slots":[{% for slot in flow.slots %}{"name":"{{ slot.name }}"{% if slot.description %},"description":{{ slot.description | to_json_escaped_string }}{% endif %}{% if slot.allowed_values %},"allowed_values":{{ slot.allowed_values }}{% endif %}}{% if not loop.last %},{% endif %}{% endfor %}]{% endif %}}{% if not loop.last %},{% endif %}{% endfor %}]}
 ```
 
 --
@@ -41,7 +41,7 @@ Use the following structured data:
 ## Current State
 {% if current_flow != None %}Use the following structured data:
 ```json
-{"active_flow":"{{ current_flow }}","current_step":{"requested_slot":"{{ current_slot }}","requested_slot_description":"{{ current_slot_description }}"},"slots":[{% for slot in flow_slots %}{"name":"{{ slot.name }}","value":"{{ slot.value }}","type":"{{ slot.type }}"{% if slot.description %},"description":"{{ slot.description }}"{% endif %}{% if slot.allowed_values %},"allowed_values":"{{ slot.allowed_values }}"{% endif %}}{% if not loop.last %},{% endif %}{% endfor %}]}
+{"active_flow":"{{ current_flow }}","current_step":{"requested_slot":"{{ current_slot }}","requested_slot_description":{{ current_slot_description | to_json_escaped_string }}},"slots":[{% for slot in flow_slots %}{"name":"{{ slot.name }}","value":"{{ slot.value }}","type":"{{ slot.type }}"{% if slot.description %},"description":{{ slot.description | to_json_escaped_string }}{% endif %}{% if slot.allowed_values %},"allowed_values":"{{ slot.allowed_values }}"{% endif %}}{% if not loop.last %},{% endif %}{% endfor %}]}
 ```{% else %}
 You are currently not inside any flow.{% endif %}
 

rasa/dialogue_understanding/generator/prompt_templates/command_prompt_v2_gpt_4o_2024_11_20_template.jinja2

@@ -6,7 +6,7 @@ Your task is to analyze the current conversation context and generate a list of
 ## Available Flows and Slots
 Use the following structured data:
 ```json
-{"flows":[{% for flow in available_flows %}{"name":"{{ flow.name }}","description":"{{ flow.description }}"{% if flow.slots %},"slots":[{% for slot in flow.slots %}{"name":"{{ slot.name }}"{% if slot.description %},"description":"{{ slot.description }}"{% endif %}{% if slot.allowed_values %},"allowed_values":{{ slot.allowed_values }}{% endif %}}{% if not loop.last %},{% endif %}{% endfor %}]{% endif %}}{% if not loop.last %},{% endif %}{% endfor %}]}
+{"flows":[{% for flow in available_flows %}{"name":"{{ flow.name }}","description":{{ flow.description | to_json_escaped_string }}{% if flow.slots %},"slots":[{% for slot in flow.slots %}{"name":"{{ slot.name }}"{% if slot.description %},"description":{{ slot.description | to_json_escaped_string }}{% endif %}{% if slot.allowed_values %},"allowed_values":{{ slot.allowed_values }}{% endif %}}{% if not loop.last %},{% endif %}{% endfor %}]{% endif %}}{% if not loop.last %},{% endif %}{% endfor %}]}
 ```
 
 ---
@@ -16,7 +16,7 @@ Use the following structured data:
 * `set slot slot_name slot_value`: Slot setting. For example, `set slot transfer_money_recipient Freddy`. Can be used to correct and change previously set values.
 * `cancel flow`: Cancelling the current flow.
 * `disambiguate flows flow_name1 flow_name2 ... flow_name_n`: Disambiguate which flow should be started when user input is ambiguous by listing the potential flows as options. For example, `disambiguate flows list_contacts add_contact remove_contact ...` if the user just wrote "contacts".
-* `provide info`: Responding to the user's questions by supplying relevant information, such as answering FAQs or explaining services.
+* `search and reply`: Responding to the user's message by accessing and supplying relevant information from the knowledge base to address their inquiry effectively.
 * `offtopic reply`: Responding to casual or social user messages that are unrelated to any flows, engaging in friendly conversation and addressing off-topic remarks.
 * `hand over`: Handing over to a human, in case the user seems frustrated or explicitly asks to speak to one.
 
@@ -27,8 +27,11 @@ Use the following structured data:
 * For categorical slots try to match the user message with allowed slot values. Use "other" if you cannot match it.
 * Set the boolean slots based on the user response. Map positive responses to `True`, and negative to `False`.
 * Extract text slot values exactly as provided by the user. Avoid assumptions, format changes, or partial extractions.
-* Only use information provided by the user.
 * Use clarification in ambiguous cases.
+* Use `disambiguate flows` only when multiple flows could fit the same message (e.g., "card" could mean `block_card` or `replace_card`).
+* A user asking a question does not automatically imply that they want `search and reply`. The objective is to help them complete a business process if its possible to do so via a flow.
+* **Flow Priority**: If a user message can be addressed by starting a flow (even if it looks like a general question), ALWAYS start the flow first. Example: If the user says "How do I activate my card?", use `start flow activate_card` instead of `search and reply`. Only use `search and reply` if no flow matches the request.
+* Only use information provided by the user.
 * Multiple flows can be started. If a user wants to digress into a second flow, you do not need to cancel the current flow.
 * Do not cancel the flow unless the user explicitly requests it.
 * Strictly adhere to the provided action format.
@@ -40,7 +43,7 @@ Use the following structured data:
 ## Current State
 {% if current_flow != None %}Use the following structured data:
 ```json
-{"active_flow":"{{ current_flow }}","current_step":{"requested_slot":"{{ current_slot }}","requested_slot_description":"{{ current_slot_description }}"},"slots":[{% for slot in flow_slots %}{"name":"{{ slot.name }}","value":"{{ slot.value }}","type":"{{ slot.type }}"{% if slot.description %},"description":"{{ slot.description }}"{% endif %}{% if slot.allowed_values %},"allowed_values":"{{ slot.allowed_values }}"{% endif %}}{% if not loop.last %},{% endif %}{% endfor %}]}
+{"active_flow":"{{ current_flow }}","current_step":{"requested_slot":"{{ current_slot }}","requested_slot_description":{{ current_slot_description | to_json_escaped_string }}},"slots":[{% for slot in flow_slots %}{"name":"{{ slot.name }}","value":"{{ slot.value }}","type":"{{ slot.type }}"{% if slot.description %},"description":{{ slot.description | to_json_escaped_string }}{% endif %}{% if slot.allowed_values %},"allowed_values":"{{ slot.allowed_values }}"{% endif %}}{% if not loop.last %},{% endif %}{% endfor %}]}
 ```{% else %}
 You are currently not inside any flow.{% endif %}
 

rasa/dialogue_understanding/processor/command_processor.py

@@ -54,9 +54,11 @@ from rasa.shared.core.constants import (
     FLOW_HASHES_SLOT,
     SlotMappingType,
 )
+from rasa.shared.core.domain import Domain
 from rasa.shared.core.events import Event, SlotSet
 from rasa.shared.core.flows import FlowsList
 from rasa.shared.core.flows.steps.collect import CollectInformationFlowStep
+from rasa.shared.core.policies.utils import contains_intentless_policy_responses
 from rasa.shared.core.slot_mappings import SlotMapping
 from rasa.shared.core.slots import Slot
 from rasa.shared.core.trackers import DialogueStateTracker
@@ -197,6 +199,7 @@ def execute_commands(
     all_flows: FlowsList,
     execution_context: ExecutionContext,
     story_graph: Optional[StoryGraph] = None,
+    domain: Optional[Domain] = None,
 ) -> List[Event]:
     """Executes a list of commands.
 
@@ -206,6 +209,7 @@ def execute_commands(
         all_flows: All flows.
         execution_context: Information about the single graph run.
         story_graph: StoryGraph object with stories available for training.
+        domain: The domain of the bot.
 
     Returns:
         A list of the events that were created.
@@ -214,7 +218,7 @@ def execute_commands(
     original_tracker = tracker.copy()
 
     commands = clean_up_commands(
-        commands, tracker, all_flows, execution_context, story_graph
+        commands, tracker, all_flows, execution_context, story_graph, domain
     )
 
     updated_flows = find_updated_flows(tracker, all_flows)
@@ -381,6 +385,7 @@ def clean_up_commands(
     all_flows: FlowsList,
     execution_context: ExecutionContext,
     story_graph: Optional[StoryGraph] = None,
+    domain: Optional[Domain] = None,
 ) -> List[Command]:
     """Clean up a list of commands.
 
@@ -396,10 +401,13 @@ def clean_up_commands(
         all_flows: All flows.
         execution_context: Information about a single graph run.
         story_graph: StoryGraph object with stories available for training.
+        domain: The domain of the bot.
 
     Returns:
     The cleaned up commands.
     """
+    domain = domain if domain else Domain.empty()
+
     slots_so_far, active_flow = filled_slots_for_active_flow(tracker, all_flows)
 
     clean_commands: List[Command] = []
@@ -465,7 +473,12 @@ def clean_up_commands(
         # handle chitchat command differently from other free-form answer commands
         elif isinstance(command, ChitChatAnswerCommand):
             clean_commands = clean_up_chitchat_command(
-                clean_commands, command, all_flows, execution_context, story_graph
+                clean_commands,
+                command,
+                all_flows,
+                execution_context,
+                domain,
+                story_graph,
             )
 
         elif isinstance(command, FreeFormAnswerCommand):
@@ -708,6 +721,7 @@ def clean_up_chitchat_command(
     command: ChitChatAnswerCommand,
     flows: FlowsList,
     execution_context: ExecutionContext,
+    domain: Domain,
     story_graph: Optional[StoryGraph] = None,
 ) -> List[Command]:
     """Clean up a chitchat answer command.
@@ -721,6 +735,8 @@ def clean_up_chitchat_command(
         flows: All flows.
         execution_context: Information about a single graph run.
         story_graph: StoryGraph object with stories available for training.
+        domain: The domain of the bot.
+
     Returns:
         The cleaned up commands.
     """
@@ -746,10 +762,9 @@ def clean_up_chitchat_command(
     )
     defines_intentless_policy = execution_context.has_node(IntentlessPolicy)
 
-    has_e2e_stories = True if (story_graph and story_graph.has_e2e_stories()) else False
-
     if (has_action_trigger_chitchat and not defines_intentless_policy) or (
-        defines_intentless_policy and not has_e2e_stories
+        defines_intentless_policy
+        and not contains_intentless_policy_responses(flows, domain, story_graph)
     ):
         resulting_commands.insert(
             0, CannotHandleCommand(RASA_PATTERN_CANNOT_HANDLE_CHITCHAT)

rasa/dialogue_understanding/processor/command_processor_component.py

@@ -6,6 +6,7 @@ import rasa.dialogue_understanding.processor.command_processor
 from rasa.engine.graph import ExecutionContext, GraphComponent
 from rasa.engine.storage.resource import Resource
 from rasa.engine.storage.storage import ModelStorage
+from rasa.shared.core.domain import Domain
 from rasa.shared.core.events import Event
 from rasa.shared.core.flows import FlowsList
 from rasa.shared.core.trackers import DialogueStateTracker
@@ -15,7 +16,8 @@ from rasa.shared.core.training_data.structures import StoryGraph
 class CommandProcessorComponent(GraphComponent):
     """Processes commands by issuing events to modify a tracker.
 
-    Minimal component that applies commands to a tracker."""
+    Minimal component that applies commands to a tracker.
+    """
 
     def __init__(self, execution_context: ExecutionContext):
         self._execution_context = execution_context
@@ -36,8 +38,9 @@ class CommandProcessorComponent(GraphComponent):
         tracker: DialogueStateTracker,
         flows: FlowsList,
         story_graph: StoryGraph,
+        domain: Domain,
     ) -> List[Event]:
         """Execute commands to update tracker state."""
         return rasa.dialogue_understanding.processor.command_processor.execute_commands(
-            tracker, flows, self._execution_context, story_graph
+            tracker, flows, self._execution_context, story_graph, domain
         )

rasa/dialogue_understanding_test/command_metric_calculation.py

@@ -1,54 +1,21 @@
+import typing
 from collections import defaultdict
 from typing import Dict, List
 
-from pydantic import BaseModel
-
 from rasa.dialogue_understanding.commands import Command
 from rasa.dialogue_understanding_test.command_comparison import (
     is_command_present_in_list,
 )
-from rasa.dialogue_understanding_test.du_test_result import (
-    DialogueUnderstandingTestResult,
-)
-
-
-class CommandMetrics(BaseModel):
-    tp: int
-    fp: int
-    fn: int
-    total_count: int
-
-    @staticmethod
-    def _safe_divide(numerator: float, denominator: float) -> float:
-        """Safely perform division, returning 0.0 if the denominator is zero."""
-        return numerator / denominator if denominator > 0 else 0.0
+from rasa.dialogue_understanding_test.command_metrics import CommandMetrics
 
-    def get_precision(self) -> float:
-        return self._safe_divide(self.tp, self.tp + self.fp)
-
-    def get_recall(self) -> float:
-        return self._safe_divide(self.tp, self.tp + self.fn)
-
-    def get_f1_score(self) -> float:
-        precision = self.get_precision()
-        recall = self.get_recall()
-
-        return self._safe_divide(2 * precision * recall, precision + recall)
-
-    def as_dict(self) -> Dict[str, float]:
-        return {
-            "tp": self.tp,
-            "fp": self.fp,
-            "fn": self.fn,
-            "precision": self.get_precision(),
-            "recall": self.get_recall(),
-            "f1_score": self.get_f1_score(),
-            "total_count": self.total_count,
-        }
+if typing.TYPE_CHECKING:
+    from rasa.dialogue_understanding_test.du_test_result import (
+        DialogueUnderstandingTestResult,
+    )
 
 
 def calculate_command_metrics(
-    test_results: List[DialogueUnderstandingTestResult],
+    test_results: List["DialogueUnderstandingTestResult"],
 ) -> Dict[str, CommandMetrics]:
     """Calculate the command metrics for the test result."""
     metrics: Dict[str, CommandMetrics] = defaultdict(

rasa/dialogue_understanding_test/command_metrics.py

@@ -0,0 +1,38 @@
+from typing import Dict
+
+from pydantic import BaseModel
+
+
+class CommandMetrics(BaseModel):
+    tp: int
+    fp: int
+    fn: int
+    total_count: int
+
+    @staticmethod
+    def _safe_divide(numerator: float, denominator: float) -> float:
+        """Safely perform division, returning 0.0 if the denominator is zero."""
+        return numerator / denominator if denominator > 0 else 0.0
+
+    def get_precision(self) -> float:
+        return self._safe_divide(self.tp, self.tp + self.fp)
+
+    def get_recall(self) -> float:
+        return self._safe_divide(self.tp, self.tp + self.fn)
+
+    def get_f1_score(self) -> float:
+        precision = self.get_precision()
+        recall = self.get_recall()
+
+        return self._safe_divide(2 * precision * recall, precision + recall)
+
+    def as_dict(self) -> Dict[str, float]:
+        return {
+            "tp": self.tp,
+            "fp": self.fp,
+            "fn": self.fn,
+            "precision": self.get_precision(),
+            "recall": self.get_recall(),
+            "f1_score": self.get_f1_score(),
+            "total_count": self.total_count,
+        }

rasa/dialogue_understanding_test/du_test_case.py

@@ -1,7 +1,11 @@
+from collections import defaultdict
 from typing import Any, Dict, Iterator, List, Optional, Tuple
 
 from pydantic import BaseModel, Field
 
+from rasa.core import IntentlessPolicy
+from rasa.core.nlg.contextual_response_rephraser import ContextualResponseRephraser
+from rasa.core.policies.enterprise_search_policy import EnterpriseSearchPolicy
 from rasa.dialogue_understanding.commands.prompt_command import PromptCommand
 from rasa.dialogue_understanding.generator.command_parser import parse_commands
 from rasa.dialogue_understanding_test.command_comparison import are_command_lists_equal
@@ -69,6 +73,8 @@ class DialogueUnderstandingOutput(BaseModel):
     commands: Dict[str, List[PromptCommand]]
     # List of prompts
     prompts: Optional[List[Dict[str, Any]]] = None
+    # Latency of the full message roundtrip
+    latency: Optional[float] = None
 
     class Config:
         """Skip validation for PromptCommand protocol as pydantic does not know how to
@@ -88,27 +94,41 @@ class DialogueUnderstandingOutput(BaseModel):
     def get_component_names_that_predicted_commands_or_have_llm_response(
         self,
     ) -> List[str]:
-        """Get all component names that have predicted commands or recieved
+        """Get all relevant component names.
+
+        Components are relevant if they have predicted commands or received a
         non-empty response from LLM.
         """
+        # Exclude components that are not related to Dialogue Understanding
+        component_names_to_exclude = [
+            EnterpriseSearchPolicy.__name__,
+            IntentlessPolicy.__name__,
+            ContextualResponseRephraser.__name__,
+        ]
+
         component_names_that_predicted_commands = (
             [
                 component_name
                 for component_name, predicted_commands in self.commands.items()
                 if predicted_commands
+                and component_name not in component_names_to_exclude
             ]
             if self.commands
             else []
         )
+
         components_with_prompts = (
             [
                 str(prompt.get(KEY_COMPONENT_NAME, None))
                 for prompt in self.prompts
                 if prompt.get(KEY_LLM_RESPONSE_METADATA, None)
+                and prompt.get(KEY_COMPONENT_NAME, None)
+                not in component_names_to_exclude
             ]
             if self.prompts
             else []
         )
+
         return list(
             set(component_names_that_predicted_commands + components_with_prompts)
         )
@@ -290,41 +310,54 @@ class DialogueUnderstandingTestStep(BaseModel):
 
         return ""
 
-    def get_latencies(self) -> List[float]:
+    def get_latencies(self) -> Dict[str, List[float]]:
         if self.dialogue_understanding_output is None:
-            return []
+            return {}
 
-        prompts = self.dialogue_understanding_output.get_component_name_to_prompt_info()
+        component_name_to_prompt_info = (
+            self.dialogue_understanding_output.get_component_name_to_prompt_info()
+        )
 
-        return [
-            prompt_data.get(KEY_LATENCY, 0.0)
-            for prompt in prompts.values()
-            for prompt_data in prompt
-        ]
+        latencies = defaultdict(list)
+        for component_name, prompt_info_list in component_name_to_prompt_info.items():
+            for prompt_info in prompt_info_list:
+                latencies[component_name].append(prompt_info.get(KEY_LATENCY, 0.0))
 
-    def get_completion_tokens(self) -> List[int]:
+        return latencies
+
+    def get_completion_tokens(self) -> Dict[str, List[float]]:
         if self.dialogue_understanding_output is None:
-            return []
+            return {}
 
-        prompts = self.dialogue_understanding_output.get_component_name_to_prompt_info()
+        component_name_to_prompt_info = (
+            self.dialogue_understanding_output.get_component_name_to_prompt_info()
+        )
 
-        return [
-            prompt_data.get(KEY_COMPLETION_TOKENS, 0)
-            for prompt in prompts.values()
-            for prompt_data in prompt
-        ]
+        completion_tokens = defaultdict(list)
+        for component_name, prompt_info_list in component_name_to_prompt_info.items():
+            for prompt_info in prompt_info_list:
+                completion_tokens[component_name].append(
+                    prompt_info.get(KEY_COMPLETION_TOKENS, 0.0)
+                )
+
+        return completion_tokens
 
-    def get_prompt_tokens(self) -> List[int]:
+    def get_prompt_tokens(self) -> Dict[str, List[float]]:
         if self.dialogue_understanding_output is None:
-            return []
+            return {}
 
-        prompts = self.dialogue_understanding_output.get_component_name_to_prompt_info()
+        component_name_to_prompt_info = (
+            self.dialogue_understanding_output.get_component_name_to_prompt_info()
+        )
 
-        return [
-            prompt_data.get(KEY_PROMPT_TOKENS, 0)
-            for prompt in prompts.values()
-            for prompt_data in prompt
-        ]
+        prompt_tokens = defaultdict(list)
+        for component_name, prompt_info_list in component_name_to_prompt_info.items():
+            for prompt_info in prompt_info_list:
+                prompt_tokens[component_name].append(
+                    prompt_info.get(KEY_PROMPT_TOKENS, 0.0)
+                )
+
+        return prompt_tokens
 
 
 class DialogueUnderstandingTestCase(BaseModel):