rasa-pro 3.13.0.dev7__py3-none-any.whl → 3.13.0.dev9__py3-none-any.whl

rasa/shared/utils/llm.py

@@ -1,13 +1,17 @@
+from __future__ import annotations
+
 import importlib.resources
 import json
 import logging
 from copy import deepcopy
+from datetime import datetime
 from functools import wraps
 from typing import (
     TYPE_CHECKING,
     Any,
     Callable,
     Dict,
+    List,
     Literal,
     Optional,
     Text,
@@ -18,15 +22,23 @@ from typing import (
 )
 
 import structlog
+from pydantic import BaseModel, Field
 
 import rasa.shared.utils.io
-from rasa.core.utils import AvailableEndpoints
+from rasa.core.available_endpoints import AvailableEndpoints
 from rasa.shared.constants import (
+    CONFIG_NAME_KEY,
+    CONFIG_PIPELINE_KEY,
+    CONFIG_POLICIES_KEY,
     DEFAULT_PROMPT_PACKAGE_NAME,
+    LLM_CONFIG_KEY,
     MODEL_CONFIG_KEY,
     MODEL_GROUP_CONFIG_KEY,
     MODEL_GROUP_ID_CONFIG_KEY,
+    MODEL_GROUPS_CONFIG_KEY,
     MODELS_CONFIG_KEY,
+    PROMPT_CONFIG_KEY,
+    PROMPT_TEMPLATE_CONFIG_KEY,
     PROVIDER_CONFIG_KEY,
     RASA_PATTERN_INTERNAL_ERROR_USER_INPUT_EMPTY,
     RASA_PATTERN_INTERNAL_ERROR_USER_INPUT_TOO_LONG,
@@ -61,9 +73,11 @@ from rasa.shared.providers.mappings import (
     get_embedding_client_from_provider,
     get_llm_client_from_provider,
 )
+from rasa.shared.utils.common import all_subclasses
 from rasa.shared.utils.constants import LOG_COMPONENT_SOURCE_METHOD_INIT
 
 if TYPE_CHECKING:
+    from rasa.core.agent import Agent
     from rasa.shared.core.trackers import DialogueStateTracker
 
 
@@ -107,6 +121,18 @@ _CombineConfigs_F = TypeVar(
 )
 
 
+class SystemPrompts(BaseModel):
+    command_generator: str = Field(
+        ..., description="Prompt used by the LLM command generator."
+    )
+    enterprise_search: str = Field(
+        ..., description="Prompt for standard enterprise search requests."
+    )
+    contextual_response_rephraser: str = Field(
+        ..., description="Prompt used for re-phrasing assistant responses."
+    )
+
+
 def _compute_hash_for_cache_from_configs(
     config_x: Dict[str, Any], config_y: Dict[str, Any]
 ) -> int:
@@ -191,6 +217,7 @@ def tracker_as_readable_transcript(
     human_prefix: str = USER,
     ai_prefix: str = AI,
     max_turns: Optional[int] = 20,
+    turns_wrapper: Optional[Callable[[List[str]], List[str]]] = None,
 ) -> str:
     """Creates a readable dialogue from a tracker.
 
@@ -199,6 +226,7 @@ def tracker_as_readable_transcript(
         human_prefix: the prefix to use for human utterances
         ai_prefix: the prefix to use for ai utterances
         max_turns: the maximum number of turns to include in the transcript
+        turns_wrapper: optional function to wrap the turns in a custom way
 
     Example:
         >>> tracker = Tracker(
@@ -235,8 +263,11 @@ def tracker_as_readable_transcript(
         elif isinstance(event, BotUttered):
             transcript.append(f"{ai_prefix}: {sanitize_message_for_prompt(event.text)}")
 
-    if max_turns:
-        transcript = transcript[-max_turns:]
+    # turns_wrapper to count multiple utterances by bot/user as single turn
+    if turns_wrapper:
+        transcript = turns_wrapper(transcript)
+    # otherwise, just take the last `max_turns` lines of the transcript
+    transcript = transcript[-max_turns if max_turns is not None else None :]
 
     return "\n".join(transcript)
 
@@ -678,7 +709,6 @@ def get_prompt_template(
     Returns:
         The prompt template.
     """
-
     try:
         if jinja_file_path is not None:
             prompt_template = rasa.shared.utils.io.read_file(jinja_file_path)
@@ -814,7 +844,9 @@ def allowed_values_for_slot(slot: Slot) -> Union[str, None]:
 
 
 def resolve_model_client_config(
-    model_config: Optional[Dict[str, Any]], component_name: Optional[str] = None
+    model_config: Optional[Dict[str, Any]],
+    component_name: Optional[str] = None,
+    model_groups: Optional[List[Dict[str, Any]]] = None,
 ) -> Optional[Dict[str, Any]]:
     """Resolve the model group in the model config.
 
@@ -828,6 +860,7 @@ def resolve_model_client_config(
         model_config: The model config to be resolved.
         component_name: The name of the component.
         component_name: The method of the component.
+        model_groups: Model groups from endpoints.yml.
 
     Returns:
         The resolved llm config.
@@ -854,7 +887,12 @@ def resolve_model_client_config(
 
     model_group_id = model_config.get(MODEL_GROUP_CONFIG_KEY)
 
-    endpoints = AvailableEndpoints.get_instance()
+    # If `model_groups` is provided, use it to initialise `AvailableEndpoints`,
+    # since `get_instance()` reads from the local endpoints file instead.
+    if model_groups:
+        endpoints = AvailableEndpoints(model_groups=model_groups)
+    else:
+        endpoints = AvailableEndpoints.get_instance()
     if endpoints.model_groups is None:
         _raise_invalid_config_exception(
             reason=(
@@ -886,3 +924,150 @@ def resolve_model_client_config(
         )
 
     return model_group[0]
+
+
+def generate_sender_id(test_case_name: str) -> str:
+    # add timestamp suffix to ensure sender_id is unique
+    return f"{test_case_name}_{datetime.now()}"
+
+
+async def create_tracker_for_user_step(
+    step_sender_id: str,
+    agent: "Agent",
+    test_case_tracker: "DialogueStateTracker",
+    index_user_uttered_event: int,
+) -> None:
+    """Creates a tracker for the user step."""
+    tracker = test_case_tracker.copy()
+    # modify the sender id so that the original tracker is not overwritten
+    tracker.sender_id = step_sender_id
+
+    if tracker.events:
+        # get the timestamp of the event just before the user uttered event
+        timestamp = tracker.events[index_user_uttered_event - 1].timestamp
+        # revert the tracker to the event just before the user uttered event
+        tracker = tracker.travel_back_in_time(timestamp)
+
+    # store the tracker with the unique sender id
+    await agent.tracker_store.save(tracker)
+
+
+def check_prompt_config_keys_and_warn_if_deprecated(
+    config: dict, component_source: str
+) -> None:
+    """Checks and warns about deprecated config parameters."""
+    if PROMPT_CONFIG_KEY in config and PROMPT_TEMPLATE_CONFIG_KEY in config:
+        structlogger.warning(
+            f"{component_source}.init"
+            ".both_deprecated_and_non_deprecated_config_keys_used_at_the_same_time",
+            event_info=(
+                f"Both '{PROMPT_CONFIG_KEY}' and '{PROMPT_TEMPLATE_CONFIG_KEY}' "
+                f"are present in the config. '{PROMPT_CONFIG_KEY}' will be ignored "
+                f"in favor of {PROMPT_TEMPLATE_CONFIG_KEY}."
+            ),
+        )
+
+    # 'prompt' config key is deprecated in favor of 'prompt_template'
+    if PROMPT_CONFIG_KEY in config:
+        structlogger.warning(
+            f"{component_source}.init.deprecated_config_key",
+            event_info=(
+                f"The config parameter '{PROMPT_CONFIG_KEY}' is deprecated "
+                "and will be removed in Rasa 4.0.0. "
+                f"Please use the config parameter '{PROMPT_TEMPLATE_CONFIG_KEY}'"
+                f" instead. "
+            ),
+        )
+
+
+def _get_llm_command_generator_config(
+    config: Dict[Text, Any],
+) -> Optional[Dict[Text, Any]]:
+    """Get the llm command generator config from config.yml.
+
+    Args:
+        config: The config.yml file data.
+
+    Returns:
+        The llm command generator config.
+    """
+    from rasa.dialogue_understanding.generator import LLMBasedCommandGenerator
+
+    # Collect all LLM based Command Generator class names.
+    command_generator_subclasses = all_subclasses(LLMBasedCommandGenerator)
+    command_generator_class_names = [
+        command_generator.__name__ for command_generator in command_generator_subclasses
+    ]
+
+    # Read the LLM config of the Command Generator from the config.yml file.
+    pipelines = config.get(CONFIG_PIPELINE_KEY, [])
+    for pipeline in pipelines:
+        if pipeline.get(CONFIG_NAME_KEY) in command_generator_class_names:
+            return pipeline.get(LLM_CONFIG_KEY)
+
+    return None
+
+
+def _get_command_generator_prompt(
+    config: Dict[Text, Any], endpoints: Dict[Text, Any]
+) -> Text:
+    """Get the command generator prompt based on the config."""
+    from rasa.dialogue_understanding.generator.single_step.compact_llm_command_generator import (  # noqa: E501
+        DEFAULT_COMMAND_PROMPT_TEMPLATE_FILE_NAME,
+        FALLBACK_COMMAND_PROMPT_TEMPLATE_FILE_NAME,
+        MODEL_PROMPT_MAPPER,
+    )
+
+    model_config = _get_llm_command_generator_config(config)
+    llm_config = resolve_model_client_config(
+        model_config=model_config,
+        model_groups=endpoints.get(MODEL_GROUPS_CONFIG_KEY),
+    )
+    return get_default_prompt_template_based_on_model(
+        llm_config=llm_config,
+        model_prompt_mapping=MODEL_PROMPT_MAPPER,
+        default_prompt_path=DEFAULT_COMMAND_PROMPT_TEMPLATE_FILE_NAME,
+        fallback_prompt_path=FALLBACK_COMMAND_PROMPT_TEMPLATE_FILE_NAME,
+    )
+
+
+def _get_enterprise_search_prompt(config: Dict[Text, Any]) -> Text:
+    """Get the enterprise search prompt based on the config."""
+    from rasa.core.policies.enterprise_search_policy import EnterpriseSearchPolicy
+
+    def get_enterprise_search_config() -> Dict[Text, Any]:
+        policies = config.get(CONFIG_POLICIES_KEY, [])
+        for policy in policies:
+            if policy.get(CONFIG_NAME_KEY) == EnterpriseSearchPolicy.__name__:
+                return policy
+
+        return {}
+
+    enterprise_search_config = get_enterprise_search_config()
+    return EnterpriseSearchPolicy.get_system_default_prompt_based_on_config(
+        enterprise_search_config
+    )
+
+
+def get_system_default_prompts(
+    config: Dict[Text, Any], endpoints: Dict[Text, Any]
+) -> SystemPrompts:
+    """
+    Returns the system default prompts for the component.
+
+    Args:
+        config: The config.yml file data.
+        endpoints: The endpoints.yml file data.
+
+    Returns:
+        SystemPrompts: A Pydantic model containing all default prompts.
+    """
+    from rasa.core.nlg.contextual_response_rephraser import (
+        DEFAULT_RESPONSE_VARIATION_PROMPT_TEMPLATE,
+    )
+
+    return SystemPrompts(
+        command_generator=_get_command_generator_prompt(config, endpoints),
+        enterprise_search=_get_enterprise_search_prompt(config),
+        contextual_response_rephraser=DEFAULT_RESPONSE_VARIATION_PROMPT_TEMPLATE,
+    )

rasa/shared/utils/yaml.py

@@ -21,6 +21,7 @@ from ruamel.yaml import YAML, RoundTripRepresenter, YAMLError
 from ruamel.yaml.comments import CommentedMap, CommentedSeq
 from ruamel.yaml.constructor import BaseConstructor, DuplicateKeyError, ScalarNode
 from ruamel.yaml.loader import SafeLoader
+from ruamel.yaml.scalarstring import LiteralScalarString
 
 from rasa.shared.constants import (
     ASSERTIONS_SCHEMA_EXTENSIONS_FILE,
@@ -794,6 +795,25 @@ def write_yaml(
         should_preserve_key_order: Whether to force preserve key order in `data`.
         transform: A function to transform the data before writing it to the file.
     """
+
+    def multiline_str_representer(self: Any, value: str) -> Any:
+        """Dump multi-line strings as readable YAML block scalars where possible."""
+        if "\n" in value:
+            # First line after the newline decides: paragraph vs. snippet
+            first_line = value.split("\n", 1)[1]
+
+            # If the first line after the newline is not indented, treat the value
+            # as plain text. Indented text is likely pre-formatted YAML/JSON/etc.
+            if not first_line.startswith((" ", "\t")):
+                return self.represent_scalar(
+                    "tag:yaml.org,2002:str",
+                    LiteralScalarString(value),
+                    style="|",
+                )
+
+        # Fallback: keep default YAML scalar style (plain/quoted)
+        return self.represent_scalar("tag:yaml.org,2002:str", value)
+
     _enable_ordered_dict_yaml_dumping()
 
     if should_preserve_key_order:
@@ -808,6 +828,7 @@ def write_yaml(
         type(None),
         lambda self, _: self.represent_scalar("tag:yaml.org,2002:null", "null"),
     )
+    dumper.representer.add_representer(str, multiline_str_representer)
 
     if isinstance(target, StringIO):
         dumper.dump(data, target, transform=transform)
@@ -1025,6 +1046,17 @@ def validate_yaml_with_jsonschema(
     except (YAMLError, DuplicateKeyError) as e:
         raise YamlSyntaxException(underlying_yaml_exception=e)
 
+    validate_data_with_jsonschema(source_data, schema_content, humanize_error)
+
+
+def validate_data_with_jsonschema(
+    source_data: Any,
+    schema_content: Any,
+    humanize_error: Callable[
+        [jsonschema.ValidationError], str
+    ] = default_error_humanizer,
+) -> None:
+    """Validate Python object against the provided jsonschema content."""
     try:
         jsonschema.validate(source_data, schema_content)
     except jsonschema.ValidationError as error:

rasa/studio/data_handler.py

@@ -320,14 +320,14 @@ def create_new_flows_from_diff(
 
 
 def import_data_from_studio(
-    handler: StudioDataHandler, domain_path: Path, data_paths: List[Path]
+    handler: StudioDataHandler, domain_path: Path, data_path: Path
 ) -> Tuple[TrainingDataImporter, TrainingDataImporter]:
     """Construct TrainingDataImporter from Studio data and original data.
 
     Args:
         handler (StudioDataHandler): handler with data from studio
         domain_path (Path): Path to a domain file
-        data_paths (List[Path]): List of paths to training data files
+        data_path (List[Path]): List of paths to training data files
 
     Returns:
         Tuple[TrainingDataImporter, TrainingDataImporter]:
@@ -335,7 +335,7 @@ def import_data_from_studio(
     """
     tmp_dir = get_temp_dir_name()
     data_original = TrainingDataImporter.load_from_dict(
-        domain_path=domain_path, training_data_paths=data_paths
+        domain_path=str(domain_path), training_data_paths=[str(data_path)]
     )
 
     data_paths = []

rasa/studio/download/download.py

@@ -1,6 +1,6 @@
 import argparse
 from pathlib import Path
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, Optional, Tuple
 
 import questionary
 import structlog
@@ -46,7 +46,7 @@ def handle_download(args: argparse.Namespace) -> None:
     )
     handler.request_all_data()
 
-    domain_path, data_paths = _prepare_data_and_domain_paths(args)
+    domain_path, data_path = _prepare_data_and_domain_paths(args)
 
     # Handle config and endpoints.
     config_path, write_config = _handle_file_overwrite(
@@ -78,12 +78,12 @@ def handle_download(args: argparse.Namespace) -> None:
         structlogger.info("studio.download.config_endpoints", event_info=message)
 
     if not args.overwrite:
-        _handle_download_no_overwrite(handler, domain_path, data_paths)
+        _handle_download_no_overwrite(handler, domain_path, data_path)
     else:
-        _handle_download_with_overwrite(handler, domain_path, data_paths)
+        _handle_download_with_overwrite(handler, domain_path, data_path)
 
 
-def _prepare_data_and_domain_paths(args: argparse.Namespace) -> Tuple[Path, List[Path]]:
+def _prepare_data_and_domain_paths(args: argparse.Namespace) -> Tuple[Path, Path]:
     """Prepars the domain and data paths based on the provided arguments.
 
     Args:
@@ -115,28 +115,15 @@ def _prepare_data_and_domain_paths(args: argparse.Namespace) -> Tuple[Path, List
             domain_path = domain_path / STUDIO_DOMAIN_FILENAME
             domain_path.touch()
 
-    # Prepare data paths.
-    data_paths: List[Path] = []
-    for f in args.data:
-        data_path = rasa.cli.utils.get_validated_path(
-            f, "data", DEFAULT_DATA_PATH, none_is_valid=True
-        )
-
-        if data_path is None:
-            data_path = Path(f)
-            data_path.mkdir(parents=True, exist_ok=True)
-        else:
-            data_path = Path(data_path)
+    data_path = rasa.cli.utils.get_validated_path(
+        args.data[0], "data", DEFAULT_DATA_PATH, none_is_valid=True
+    )
 
-        if data_path.is_file() or data_path.is_dir():
-            data_paths.append(data_path)
-        else:
-            data_path.mkdir(parents=True, exist_ok=True)
-            data_paths.append(data_path)
+    data_path = Path(data_path or args.data[0])
+    if not (data_path.is_file() or data_path.is_dir()):
+        data_path.mkdir(parents=True, exist_ok=True)
 
-    # Remove duplicates while preserving order.
-    data_paths = list(dict.fromkeys(data_paths))
-    return domain_path, data_paths
+    return domain_path, data_path
 
 
 def _handle_file_overwrite(
@@ -177,7 +164,7 @@ def _handle_file_overwrite(
 
 
 def _handle_download_no_overwrite(
-    handler: StudioDataHandler, domain_path: Path, data_paths: List[Path]
+    handler: StudioDataHandler, domain_path: Path, data_path: Path
 ) -> None:
     """Handles downloading without overwriting existing files.
 
@@ -187,10 +174,10 @@ def _handle_download_no_overwrite(
         data_paths: The paths to the data files or directories.
     """
     data_from_studio, data_local = import_data_from_studio(
-        handler, domain_path, data_paths
+        handler, domain_path, data_path
     )
     _merge_domain_no_overwrite(domain_path, data_from_studio, data_local)
-    _merge_data_no_overwrite(data_paths, handler, data_from_studio, data_local)
+    _merge_data_no_overwrite(data_path, handler, data_from_studio, data_local)
 
 
 def _merge_domain_no_overwrite(
@@ -264,7 +251,7 @@ def _merge_file_domain(
 
 
 def _merge_data_no_overwrite(
-    data_paths: List[Path],
+    data_path: Path,
     handler: StudioDataHandler,
     data_from_studio: TrainingDataImporter,
     data_local: TrainingDataImporter,
@@ -272,38 +259,29 @@ def _merge_data_no_overwrite(
     """Merges NLU and flow data without overwriting existing data.
 
     Args:
-        data_paths: The paths to the data files or directories.
+        data_path: The paths to the data files or directories.
         handler: The StudioDataHandler instance.
         data_from_studio: The Studio data importer.
         data_local: The local data importer.
     """
-    if not data_paths:
+    if not data_path:
         structlogger.warning(
             "studio.download.merge_data_no_overwrite.no_path",
             event_info="No data paths provided. Skipping data merge.",
         )
         return
 
-    if len(data_paths) == 1:
-        data_path = data_paths[0]
-        if data_path.is_file():
-            _merge_file_data_no_overwrite(
-                data_path, handler, data_from_studio, data_local
-            )
-        elif data_path.is_dir():
-            _merge_dir_data_no_overwrite(
-                data_path, handler, data_from_studio, data_local
-            )
-        else:
-            structlogger.warning(
-                "studio.download.merge_data_no_overwrite.invalid_path",
-                event_info=(
-                    f"Provided path '{data_path}' is neither a file nor a directory."
-                ),
-            )
+    if data_path.is_file():
+        _merge_file_data_no_overwrite(data_path, handler, data_from_studio, data_local)
+    elif data_path.is_dir():
+        _merge_dir_data_no_overwrite(data_path, handler, data_from_studio, data_local)
     else:
-        # TODO: Handle multiple data paths.
-        raise NotImplementedError("Multiple data paths are not supported yet.")
+        structlogger.warning(
+            "studio.download.merge_data_no_overwrite.invalid_path",
+            event_info=(
+                f"Provided path '{data_path}' is neither a file nor a directory."
+            ),
+        )
 
 
 def _merge_file_data_no_overwrite(
@@ -353,25 +331,23 @@ def _merge_dir_data_no_overwrite(
 
 
 def _handle_download_with_overwrite(
-    handler: StudioDataHandler, domain_path: Path, data_paths: List[Path]
+    handler: StudioDataHandler, domain_path: Path, data_path: Path
 ) -> None:
     """Handles downloading and merging data when the user opts for overwrite.
 
     Args:
         handler: The StudioDataHandler instance.
         domain_path: The path to the domain file or directory.
-        data_paths: The paths to the data files or directories.
+        data_path: The paths to the data files or directories.
     """
     data_from_studio, data_local = import_data_from_studio(
-        handler, domain_path, data_paths
+        handler, domain_path, data_path
     )
     mapper = RasaPrimitiveStorageMapper(
-        domain_path=domain_path, training_data_paths=data_paths
+        domain_path=domain_path, training_data_paths=[data_path]
     )
     merge_domain_with_overwrite(data_from_studio, data_local, domain_path)
-    merge_flows_with_overwrite(
-        data_paths, handler, data_from_studio, data_local, mapper
-    )
+    merge_flows_with_overwrite(data_path, handler, data_from_studio, data_local, mapper)
 
 
 def _persist_nlu_diff(
@@ -432,8 +408,9 @@ def pretty_write_nlu_yaml(data: Dict, file: Path) -> None:
         file: The file to write to.
     """
     dumper = yaml.YAML()
-    for item in data["nlu"]:
-        if item.get("examples"):
-            item["examples"] = LiteralScalarString(item["examples"])
+    if nlu_data := data.get("nlu"):
+        for item in nlu_data:
+            if item.get("examples"):
+                item["examples"] = LiteralScalarString(item["examples"])
     with file.open("w", encoding="utf-8") as outfile:
         dumper.dump(data, outfile)

rasa/studio/download/flows.py

@@ -18,7 +18,7 @@ STUDIO_FLOWS_DIR_NAME = "studio_flows"
 
 
 def merge_flows_with_overwrite(
-    data_paths: List[Path],
+    data_path: Path,
     handler: Any,
     data_from_studio: TrainingDataImporter,
     data_local: TrainingDataImporter,
@@ -28,17 +28,12 @@ def merge_flows_with_overwrite(
     Merges flows data from a file or directory when overwrite is enabled.
 
     Args:
-        data_paths: List of paths to the training data.
+        data_path: List of paths to the training data.
         handler: The StudioDataHandler instance.
         data_from_studio: The TrainingDataImporter instance for Studio data.
         data_local: The TrainingDataImporter instance for local data.
         mapper: The RasaPrimitiveStorageMapper instance for mapping.
     """
-    if len(data_paths) != 1:
-        # TODO: Handle multiple data paths.
-        raise NotImplementedError("Multiple data paths are not supported yet.")
-
-    data_path = data_paths[0]
     if data_path.is_file():
         merge_training_data_file(handler, data_from_studio, data_local, data_path)
     elif data_path.is_dir():
@@ -132,7 +127,8 @@ def merge_nlu_in_directory(
         )
         nlu_data = nlu_data.merge(local_nlu.get_nlu_data())
 
-    pretty_write_nlu_yaml(read_yaml(nlu_data.nlu_as_yaml()), nlu_file_path)
+    if nlu_yaml := nlu_data.nlu_as_yaml():
+        pretty_write_nlu_yaml(read_yaml(nlu_yaml), nlu_file_path)
 
 
 def get_nlu_path(
@@ -211,14 +207,16 @@ def merge_flows_in_directory(
     local_flow_paths: Set[Path] = _get_local_flow_paths(local_flows, mapper)
 
     # Track updated flows and update local files with Studio flow data.
-    all_updated_flows: List[Flow] = []
+    all_updated_flows_ids: List[Text] = []
     for flow_file_path in local_flow_paths:
-        updated_file_flows = _update_flow_file(flow_file_path, studio_flow_map)
-        all_updated_flows.extend(updated_file_flows)
+        updated_flows_ids = _update_flow_file(flow_file_path, studio_flow_map)
+        all_updated_flows_ids.extend(updated_flows_ids)
 
     # Identify new Studio flows and save them as separate files in the directory.
     new_flows = [
-        flow for flow in studio_flow_map.values() if flow not in all_updated_flows
+        flow
+        for flow_id, flow in studio_flow_map.items()
+        if flow_id not in all_updated_flows_ids
     ]
     _dump_flows_as_separate_files(new_flows, data_path)
 
@@ -243,7 +241,7 @@ def _get_local_flow_paths(
 
 def _update_flow_file(
     flow_file_path: Path, studio_flows_map: Dict[Text, Any]
-) -> List[Flow]:
+) -> List[Text]:
     """
     Reads a flow file, updates outdated flows, and replaces them with studio versions.
 
@@ -252,31 +250,25 @@ def _update_flow_file(
         studio_flows_map: A dictionary mapping flow IDs to their updated versions.
 
     Returns:
-        A list of flows from the updated flow file.
+        A list of Flows IDs from the updated flow file.
     """
     file_flows = YAMLFlowsReader.read_from_file(flow_file_path, False)
-    updated_list: List[Any] = []
-    has_changes = False
-
-    for flow in file_flows.underlying_flows:
-        studio_flow = studio_flows_map.get(flow.id)
-        if studio_flow is not None and studio_flow != flow:
-            updated_list.append(studio_flow)
-            has_changes = True
-        else:
-            updated_list.append(flow)
-
-    if has_changes:
-        new_flows_list = FlowsList(underlying_flows=updated_list)
-        new_flows_list = strip_default_next_references(new_flows_list)
+
+    # Build a list of flows, replacing any outdated flow with its studio version
+    updated_flows = [
+        studio_flows_map.get(flow.id, flow) or flow
+        for flow in file_flows.underlying_flows
+    ]
+
+    # If the updated flows differ from the original file flows, write them back
+    if updated_flows != file_flows.underlying_flows:
         YamlFlowsWriter.dump(
-            flows=new_flows_list.underlying_flows,
+            flows=updated_flows,
             filename=flow_file_path,
             should_clean_json=True,
         )
-        return new_flows_list.underlying_flows
 
-    return file_flows.underlying_flows
+    return [flow.id for flow in updated_flows]
 
 
 def _dump_flows_as_separate_files(flows: List[Any], data_path: Path) -> None: