rasa-pro 3.13.1a18__py3-none-any.whl → 3.13.1.dev1__py3-none-any.whl

rasa/__init__.py

@@ -5,5 +5,11 @@ from rasa import version
 # define the version before the other imports since these need it
 __version__ = version.__version__
 
+from litellm.integrations.langfuse.langfuse import LangFuseLogger
+
+from rasa.monkey_patches import litellm_langfuse_logger_init_fixed
+
+# Monkey-patch the init method as early as possible before the class is used
+LangFuseLogger.__init__ = litellm_langfuse_logger_init_fixed  # type: ignore
 
 logging.getLogger(__name__).addHandler(logging.NullHandler())

rasa/cli/scaffold.py

@@ -24,21 +24,10 @@ class ProjectTemplateName(Enum):
 
     DEFAULT = "default"
     TUTORIAL = "tutorial"
-    PLAIN = "plain"
-    FINANCE = "finance"
-    TELCO = "telco"
 
     def __str__(self) -> str:
         return self.value
 
-    @classmethod
-    def get_all_values(cls) -> List[str]:
-        return [name.value for name in cls]
-
-    @classmethod
-    def supported_values(cls) -> str:
-        return ", ".join(cls.get_all_values())
-
 
 template_domain_path = defaultdict(lambda: DEFAULT_DOMAIN_PATH)
 template_domain_path[ProjectTemplateName.DEFAULT] = "domain"
@@ -170,17 +159,9 @@ def create_initial_project(
     path: Text, template: ProjectTemplateName = ProjectTemplateName.DEFAULT
 ) -> None:
     """Creates directory structure and templates for initial project."""
-    import distutils.dir_util as dir_util
-
-    # clear the cache of the copy_tree function, this avoids issues if
-    # a project directory existed before and we removed folders in it
-    # with shutil.rmtree. see
-    # https://stackoverflow.com/questions/9160227/dir-util-copy-tree-fails-after-shutil-rmtree
-    if hasattr(dir_util, "_path_created"):
-        dir_util._path_created.clear()
-    else:
-        dir_util.SkipRepeatAbsolutePaths.clear()
-    dir_util.copy_tree(scaffold_path(template), path)
+    from distutils.dir_util import copy_tree
+
+    copy_tree(scaffold_path(template), path)
 
 
 def scaffold_path(template: ProjectTemplateName) -> Text:

rasa/core/actions/action.py

@@ -908,12 +908,10 @@ class RemoteAction(Action):
                 draft.setdefault("buttons", [])
                 draft["buttons"].extend(buttons)
 
+            # Avoid overwriting `draft` values with empty values
             response = {k: v for k, v in response.items() if v}
-
-            response.update(draft)
-            bot_messages.append(
-                create_bot_utterance(response, tracker.current_language)
-            )
+            draft.update(response)
+            bot_messages.append(create_bot_utterance(draft, tracker.current_language))
 
         return bot_messages
 

rasa/core/channels/studio_chat.py

@@ -42,14 +42,14 @@ if TYPE_CHECKING:
     from sanic import Sanic, Websocket  # type: ignore[attr-defined]
     from socketio import AsyncServer
 
-    from rasa.core.channels.channel import UserMessage
+    from rasa.core.channels.channel import InputChannel, UserMessage
     from rasa.shared.core.trackers import DialogueStateTracker
 
 
 structlogger = structlog.get_logger()
 
 
-def tracker_as_dump(tracker: "DialogueStateTracker") -> Dict[str, Any]:
+def tracker_as_dump(tracker: "DialogueStateTracker") -> str:
     """Create a dump of the tracker state."""
     from rasa.shared.core.trackers import get_trackers_for_conversation_sessions
 
@@ -60,9 +60,8 @@ def tracker_as_dump(tracker: "DialogueStateTracker") -> Dict[str, Any]:
     else:
         last_tracker = multiple_tracker_sessions[-1]
 
-    # TODO: this is a bug: the bridge converts this back to json, but it
-    # should be json in the first place
-    return last_tracker.current_state(EventVerbosity.AFTER_RESTART)
+    state = last_tracker.current_state(EventVerbosity.AFTER_RESTART)
+    return json.dumps(state)
 
 
 def does_need_action_prediction(tracker: "DialogueStateTracker") -> bool:
@@ -149,7 +148,6 @@ class StudioChatInput(SocketIOInput, VoiceInputChannel):
         from rasa.core.agent import Agent
 
         self.agent: Optional[Agent] = None
-        self.latest_tracker_session_id = None
 
         # Initialize the SocketIO input channel
         SocketIOInput.__init__(
@@ -181,9 +179,7 @@ class StudioChatInput(SocketIOInput, VoiceInputChannel):
         self._register_tracker_update_hook()
 
     @classmethod
-    def from_credentials(
-        cls, credentials: Optional[Dict[Text, Any]]
-    ) -> "StudioChatInput":
+    def from_credentials(cls, credentials: Optional[Dict[Text, Any]]) -> "InputChannel":
         """Creates a StudioChatInput channel from credentials."""
         credentials = credentials or {}
 
@@ -215,11 +211,6 @@ class StudioChatInput(SocketIOInput, VoiceInputChannel):
         if not self.sio:
             structlogger.error("studio_chat.on_tracker_updated.sio_not_initialized")
             return
-
-        # we need the latest session id to use it for the llm helper to get the
-        # most recent conversation the user had with the bot.
-        self.latest_tracker_session_id = sender_id
-
         await self.sio.emit("tracker", tracker_dump, room=sender_id)
 
     async def on_message_proxy(
@@ -231,14 +222,7 @@ class StudioChatInput(SocketIOInput, VoiceInputChannel):
 
         Triggers a tracker update notification after processing the message.
         """
-        try:
-            await on_new_message(message)
-        except Exception as e:
-            structlogger.exception(
-                "studio_chat.on_new_message.error",
-                error=str(e),
-                sender_id=message.sender_id,
-            )
+        await on_new_message(message)
 
         if not self.agent:
             structlogger.error("studio_chat.on_message_proxy.agent_not_initialized")
@@ -291,7 +275,7 @@ class StudioChatInput(SocketIOInput, VoiceInputChannel):
     async def collect_call_parameters(
         self, channel_websocket: "Websocket"
     ) -> Optional[CallParameters]:
-        """Voice method to collect call parameters."""
+        """Voice method to collect call parameters"""
         session_id = channel_websocket.session_id
         return CallParameters(session_id, "local", "local", stream_id=session_id)
 
@@ -321,7 +305,7 @@ class StudioChatInput(SocketIOInput, VoiceInputChannel):
     def create_output_channel(
         self, voice_websocket: "Websocket", tts_engine: TTSEngine
     ) -> VoiceOutputChannel:
-        """Create a voice output channel."""
+        """Create a voice output channel"""
         return StudioVoiceOutputChannel(
             voice_websocket,
             tts_engine,
@@ -448,16 +432,8 @@ class StudioChatInput(SocketIOInput, VoiceInputChannel):
                     ws.put_message(data)
                 return
 
-            try:
-                # Handle text messages
-                await self.handle_user_message(sid, data, on_new_message)
-            except Exception as e:
-                structlogger.exception(
-                    "studio_chat.sio.handle_message.error",
-                    error=str(e),
-                    sid=sid,
-                )
-                await self.sio.emit("error", str(e), room=sid)
+            # Handle text messages
+            await self.handle_user_message(sid, data, on_new_message)
 
         @self.sio.on("update_tracker", namespace=self.namespace)
         async def on_update_tracker(sid: Text, data: Dict) -> None:

rasa/core/policies/enterprise_search_policy.py

@@ -784,19 +784,17 @@ class EnterpriseSearchPolicy(LLMHealthCheckMixin, EmbeddingsHealthCheckMixin, Po
         if not os.path.exists(docs_folder) or not os.path.isdir(docs_folder):
             error_message = (
                 f"Document source directory does not exist or is not a "
-                f"directory: '{os.path.abspath(docs_folder)}'. "
+                f"directory: '{docs_folder}'. "
                 "Please specify a valid path to the documents source directory in the "
                 "vector_store configuration."
             )
             structlogger.error(
                 "enterprise_search_policy.train.faiss.invalid_source_directory",
                 message=error_message,
-                docs_folder={os.path.abspath(docs_folder)},
-                configuration_value=docs_folder,
             )
             print_error_and_exit(error_message)
 
-        docs = glob.glob(os.path.join(docs_folder, "**/*.txt"), recursive=True)
+        docs = glob.glob(os.path.join(docs_folder, "*.txt"), recursive=True)
         if not docs or len(docs) < 1:
             error_message = (
                 f"Document source directory is empty: '{docs_folder}'. "
@@ -1135,7 +1133,8 @@ class EnterpriseSearchPolicy(LLMHealthCheckMixin, EmbeddingsHealthCheckMixin, Po
         embeddings_config: Dict[Text, Any],
         log_source_method: str,
     ) -> None:
-        """Perform the health checks using resolved LLM and embeddings configurations.
+        """
+        Perform the health checks using resolved LLM and embeddings configurations.
         Resolved means the configuration is either:
         - A reference to a model group that has already been expanded into
           its corresponding configuration using the information from
@@ -1164,7 +1163,8 @@ class EnterpriseSearchPolicy(LLMHealthCheckMixin, EmbeddingsHealthCheckMixin, Po
 
     @classmethod
     def get_system_default_prompt_based_on_config(cls, config: Dict[str, Any]) -> str:
-        """Resolves the default prompt template for Enterprise Search Policy based on
+        """
+        Resolves the default prompt template for Enterprise Search Policy based on
         the component's configuration.
 
         - The old prompt is selected when both citation and relevancy check are either
@@ -1195,7 +1195,8 @@ class EnterpriseSearchPolicy(LLMHealthCheckMixin, EmbeddingsHealthCheckMixin, Po
         relevancy_check_enabled: bool,
         citation_enabled: bool,
     ) -> str:
-        """Returns the appropriate default prompt template based on the feature flags.
+        """
+        Returns the appropriate default prompt template based on the feature flags.
 
         The selection follows this priority:
         1. If relevancy check is enabled, return the prompt that includes both relevancy

rasa/core/policies/flows/flow_executor.py

@@ -740,14 +740,7 @@ def _run_action_step(
             # do not log about non-existing validation actions of collect steps
             utter_action_name = render_template_variables("{{context.utter}}", context)
             if utter_action_name not in available_actions:
-                structlogger.warning(
-                    "flow.step.run.action.unknown",
-                    action=action_name,
-                    event_info=(
-                        f"The action '{action_name}' is not defined in the domain but "
-                        f"getting triggered by the flow '{step.flow_id}'."
-                    ),
-                )
+                structlogger.warning("flow.step.run.action.unknown", action=action_name)
         return ContinueFlowWithNextStep(events=initial_events)
 
 

rasa/dialogue_understanding/generator/flow_retrieval.py

@@ -219,6 +219,11 @@ class FlowRetrieval(EmbeddingsHealthCheckMixin):
         if self.vector_store is not None:
             with self._model_storage.write_to(self._resource) as model_path:
                 self.vector_store.save_local(model_path)
+        else:
+            structlogger.warning(
+                "flow_retrieval.persist_vector_store.not_initialized",
+                event_info="Vector store is None, not persisted.",
+            )
 
     def _persist_config(self) -> None:
         with self._model_storage.write_to(self._resource) as path:
@@ -244,16 +249,6 @@ class FlowRetrieval(EmbeddingsHealthCheckMixin):
         )
 
         flows_to_embedd = flows.exclude_link_only_flows()
-
-        if not flows_to_embedd:
-            structlogger.debug(
-                "flow_retrieval.populate_vector_store.no_flows_to_embed",
-                event_info=(
-                    "No flows to embed in the vector store, skipping population."
-                ),
-            )
-            return
-
         embeddings = self._create_embedder(self.config)
         documents = self._generate_flow_documents(flows_to_embedd, domain)
         try:
@@ -425,6 +420,10 @@ class FlowRetrieval(EmbeddingsHealthCheckMixin):
             The top k documents with similarity scores.
         """
         if self.vector_store is None:
+            structlogger.error(
+                "flow_retrieval.query_vector_store.vector_store_not_configured",
+                event_info="Vector store is not configured",
+            )
             return []
         try:
             documents_with_scores = (

rasa/dialogue_understanding/generator/llm_based_command_generator.py

@@ -1,6 +1,8 @@
 from abc import ABC, abstractmethod
+from asyncio import Lock
 from functools import lru_cache
 from typing import Any, Dict, List, Optional, Set, Text, Tuple, Union
+from uuid import UUID, uuid4
 
 import structlog
 from jinja2 import Environment, Template, select_autoescape
@@ -89,6 +91,9 @@ class LLMBasedCommandGenerator(
         else:
             self.flow_retrieval = None
 
+        self.sender_id_to_session_id_mapping: Dict[str, UUID] = {}
+        self._lock = Lock()
+
     ### Abstract methods
     @staticmethod
     @abstractmethod
@@ -225,8 +230,7 @@ class LLMBasedCommandGenerator(
 
     @lru_cache
     def compile_template(self, template: str) -> Template:
-        """
-        Compile the prompt template and register custom filters.
+        """Compile the prompt template and register custom filters.
 
         Compiling the template is an expensive operation,
         so we cache the result.
@@ -327,7 +331,9 @@ class LLMBasedCommandGenerator(
 
     @measure_llm_latency
     async def invoke_llm(
-        self, prompt: Union[List[dict], List[str], str]
+        self,
+        prompt: Union[List[dict], List[str], str],
+        metadata: Optional[Dict[str, Any]] = None,
     ) -> Optional[LLMResponse]:
         """Use LLM to generate a response.
 
@@ -340,6 +346,7 @@ class LLMBasedCommandGenerator(
                 - a list of messages. Each message is a string and will be formatted
                     as a user message.
                 - a single message as a string which will be formatted as user message.
+            metadata: Optional metadata to be passed to the LLM call.
 
         Returns:
             An LLMResponse object.
@@ -351,7 +358,7 @@ class LLMBasedCommandGenerator(
             self.config.get(LLM_CONFIG_KEY), self.get_default_llm_config()
         )
         try:
-            return await llm.acompletion(prompt)
+            return await llm.acompletion(prompt, metadata)
         except Exception as e:
             # unfortunately, langchain does not wrap LLM exceptions which means
             # we have to catch all exceptions here
@@ -654,3 +661,7 @@ class LLMBasedCommandGenerator(
     def get_default_llm_config() -> Dict[str, Any]:
         """Get the default LLM config for the command generator."""
         return DEFAULT_LLM_CONFIG
+
+    async def _get_or_create_session_id(self, sender_id: str) -> UUID:
+        async with self._lock:
+            return self.sender_id_to_session_id_mapping.setdefault(sender_id, uuid4())

rasa/dialogue_understanding/generator/llm_command_generator.py

@@ -55,7 +55,9 @@ class LLMCommandGenerator(SingleStepLLMCommandGenerator):
         )
 
     async def invoke_llm(
-        self, prompt: Union[List[dict], List[str], str]
+        self,
+        prompt: Union[List[dict], List[str], str],
+        metadata: Optional[Dict[str, Any]] = None,
     ) -> Optional[LLMResponse]:
         try:
             return await super().invoke_llm(prompt)

rasa/dialogue_understanding/generator/multi_step/multi_step_llm_command_generator.py

@@ -42,6 +42,9 @@ from rasa.engine.storage.resource import Resource
 from rasa.engine.storage.storage import ModelStorage
 from rasa.shared.constants import (
     EMBEDDINGS_CONFIG_KEY,
+    LANGFUSE_CUSTOM_METADATA_DICT,
+    LANGFUSE_METADATA_SESSION_ID,
+    LANGFUSE_TAGS,
     RASA_PATTERN_CANNOT_HANDLE_NOT_SUPPORTED,
     ROUTE_TO_CALM_SLOT,
 )
@@ -493,7 +496,21 @@ class MultiStepLLMCommandGenerator(LLMBasedCommandGenerator):
             prompt=prompt,
         )
 
-        response = await self.invoke_llm(prompt)
+        if tracker:
+            session_id = str(await self._get_or_create_session_id(tracker.sender_id))
+        else:
+            session_id = "unknown"
+
+        metadata = {
+            LANGFUSE_METADATA_SESSION_ID: session_id,
+            LANGFUSE_CUSTOM_METADATA_DICT: {
+                "component": self.__class__.__name__,
+                "function": "_predict_commands_for_active_flow",
+            },
+            LANGFUSE_TAGS: [self.__class__.__name__],
+        }
+
+        response = await self.invoke_llm(prompt, metadata)
         llm_response = LLMResponse.ensure_llm_response(response)
         actions = None
         if llm_response and llm_response.choices:
@@ -548,7 +565,20 @@ class MultiStepLLMCommandGenerator(LLMBasedCommandGenerator):
             prompt=prompt,
         )
 
-        response = await self.invoke_llm(prompt)
+        if tracker:
+            session_id = str(await self._get_or_create_session_id(tracker.sender_id))
+        else:
+            session_id = "unknown"
+
+        metadata = {
+            LANGFUSE_METADATA_SESSION_ID: session_id,
+            LANGFUSE_CUSTOM_METADATA_DICT: {
+                "component": self.__class__.__name__,
+                "function": "_predict_commands_for_handling_flows",
+            },
+            LANGFUSE_TAGS: [self.__class__.__name__],
+        }
+        response = await self.invoke_llm(prompt, metadata)
         llm_response = LLMResponse.ensure_llm_response(response)
         actions = None
         if llm_response and llm_response.choices:
@@ -638,7 +668,22 @@ class MultiStepLLMCommandGenerator(LLMBasedCommandGenerator):
             prompt=prompt,
         )
 
-        response = await self.invoke_llm(prompt)
+        if tracker:
+            session_id = str(await self._get_or_create_session_id(tracker.sender_id))
+
+        else:
+            session_id = "unknown"
+
+        metadata = {
+            LANGFUSE_METADATA_SESSION_ID: session_id,
+            LANGFUSE_CUSTOM_METADATA_DICT: {
+                "component": self.__class__.__name__,
+                "function": "_predict_commands_for_newly_started_flow",
+            },
+            LANGFUSE_TAGS: [self.__class__.__name__],
+        }
+
+        response = await self.invoke_llm(prompt, metadata)
         llm_response = LLMResponse.ensure_llm_response(response)
         actions = None
         if llm_response and llm_response.choices:

rasa/dialogue_understanding/generator/single_step/single_step_based_llm_command_generator.py

@@ -41,6 +41,9 @@ from rasa.engine.storage.resource import Resource
 from rasa.engine.storage.storage import ModelStorage
 from rasa.shared.constants import (
     EMBEDDINGS_CONFIG_KEY,
+    LANGFUSE_CUSTOM_METADATA_DICT,
+    LANGFUSE_METADATA_SESSION_ID,
+    LANGFUSE_TAGS,
     PROMPT_TEMPLATE_CONFIG_KEY,
     ROUTE_TO_CALM_SLOT,
 )
@@ -269,7 +272,18 @@ class SingleStepBasedLLMCommandGenerator(LLMBasedCommandGenerator, ABC):
             prompt=flow_prompt,
         )
 
-        response = await self.invoke_llm(flow_prompt)
+        if tracker:
+            session_id = str(await self._get_or_create_session_id(tracker.sender_id))
+        else:
+            session_id = "unknown"
+
+        metadata = {
+            LANGFUSE_METADATA_SESSION_ID: session_id,
+            LANGFUSE_CUSTOM_METADATA_DICT: {"component": self.__class__.__name__},
+            LANGFUSE_TAGS: [self.__class__.__name__],
+        }
+
+        response = await self.invoke_llm(flow_prompt, metadata)
         llm_response = LLMResponse.ensure_llm_response(response)
         # The check for 'None' maintains compatibility with older versions
         # of LLMCommandGenerator. In previous implementations, 'invoke_llm'

rasa/hooks.py

@@ -1,8 +1,20 @@
 import argparse
 import logging
+import os
 from typing import TYPE_CHECKING, List, Optional, Text, Union
 
+import litellm
 import pluggy
+import structlog
+
+from rasa.shared.providers.constants import (
+    LANGFUSE_CALLBACK_NAME,
+    LANGFUSE_HOST_ENV_VAR,
+    LANGFUSE_PROJECT_ID_ENV_VAR,
+    LANGFUSE_PUBLIC_KEY_ENV_VAR,
+    LANGFUSE_SECRET_KEY_ENV_VAR,
+    RASA_LANGFUSE_INTEGRATION_ENABLED_ENV_VAR,
+)
 
 # IMPORTANT: do not import anything from rasa here - use scoped imports
 #  this avoids circular imports, as the hooks are used in different places
@@ -17,6 +29,7 @@ if TYPE_CHECKING:
 
 hookimpl = pluggy.HookimplMarker("rasa")
 logger = logging.getLogger(__name__)
+structlogger = structlog.get_logger()
 
 
 @hookimpl  # type: ignore[misc]
@@ -56,6 +69,8 @@ def configure_commandline(cmdline_arguments: argparse.Namespace) -> Optional[Tex
         config.configure_tracing(tracer_provider)
         config.configure_metrics(endpoints_file)
 
+    _init_langfuse_integration()
+
     return endpoints_file
 
 
@@ -87,3 +102,43 @@ def create_tracker_store(
             endpoint_config=endpoint_config, domain=domain, event_broker=event_broker
         )
     return endpoint_config
+
+
+def _is_langfuse_integration_enabled() -> bool:
+    return (
+        os.environ.get(RASA_LANGFUSE_INTEGRATION_ENABLED_ENV_VAR, "false").lower()
+        == "true"
+    )
+
+
+def _init_langfuse_integration() -> None:
+    if not _is_langfuse_integration_enabled():
+        structlogger.info(
+            "hooks._init_langfuse_integration.disabled",
+            event_info="Langfuse integration is disabled.",
+        )
+        return
+
+    if (
+        not os.environ.get(LANGFUSE_HOST_ENV_VAR)
+        or not os.environ.get(LANGFUSE_PROJECT_ID_ENV_VAR)
+        or not os.environ.get(LANGFUSE_PUBLIC_KEY_ENV_VAR)
+        or not os.environ.get(LANGFUSE_SECRET_KEY_ENV_VAR)
+    ):
+        structlogger.warning(
+            "hooks._init_langfuse_integration.missing_langfuse_keys",
+            event_info=(
+                "Langfuse integration is enabled, but some environment variables "
+                "are missing. Please set LANGFUSE_HOST, LANGFUSE_PROJECT_ID, "
+                "LANGFUSE_PUBLIC_KEY and LANGFUSE_SECRET_KEY environment "
+                "variables to use Langfuse integration."
+            ),
+        )
+        return
+
+    litellm.success_callback = [LANGFUSE_CALLBACK_NAME]
+    litellm.failure_callback = [LANGFUSE_CALLBACK_NAME]
+    structlogger.info(
+        "hooks.langfuse_callbacks_initialized",
+        event_info="Langfuse integration initialized.",
+    )

rasa/llm_fine_tuning/utils.py

@@ -1,6 +1,6 @@
 from contextlib import contextmanager
 from datetime import datetime
-from typing import Callable, Generator, List, Union
+from typing import Any, Callable, Dict, Generator, List, Optional, Union
 
 import structlog
 
@@ -24,7 +24,9 @@ def make_mock_invoke_llm(commands: str) -> Callable:
     """
 
     async def _mock_invoke_llm(
-        self: LLMBasedCommandGenerator, prompt: Union[List[dict], List[str], str]
+        self: LLMBasedCommandGenerator,
+        prompt: Union[List[dict], List[str], str],
+        metadata: Optional[Dict[str, Any]] = None,
     ) -> LLMResponse:
         structlogger.debug(
             f"LLM call intercepted, response mocked. "

rasa/model_manager/model_api.py

@@ -39,9 +39,11 @@ from rasa.model_manager.trainer_service import (
     update_training_status,
 )
 from rasa.model_manager.utils import (
+    InvalidPathException,
     get_logs_content,
     logs_base_path,
     models_base_path,
+    subpath,
 )
 from rasa.model_manager.warm_rasa_process import (
     initialize_warm_rasa_process,
@@ -51,7 +53,6 @@ from rasa.server import ErrorResponse
 from rasa.shared.exceptions import InvalidConfigException
 from rasa.shared.utils.yaml import dump_obj_as_yaml_to_string
 from rasa.studio.upload import build_calm_import_parts
-from rasa.utils.io import InvalidPathException, subpath
 
 dotenv.load_dotenv()
 
@@ -104,7 +105,7 @@ async def update_status_of_all_bots() -> None:
         await update_bot_status(bot)
 
 
-def base_server_url(request: Optional[Request]) -> str:
+def base_server_url(request: Request) -> str:
     """Return the base URL of the server."""
     if SERVER_BASE_URL:
         return SERVER_BASE_URL.rstrip("/")

rasa/model_manager/runner_service.py

@@ -15,11 +15,11 @@ from rasa.model_manager import config
 from rasa.model_manager.utils import (
     logs_path,
     models_base_path,
+    subpath,
     write_encoded_data_to_file,
 )
 from rasa.model_manager.warm_rasa_process import start_rasa_process
 from rasa.studio.prompts import handle_prompts
-from rasa.utils.io import subpath
 
 structlogger = structlog.get_logger()
 

rasa/model_manager/trainer_service.py

@@ -14,6 +14,7 @@ from rasa.model_manager.utils import (
     ensure_base_directory_exists,
     logs_path,
     models_base_path,
+    subpath,
     write_encoded_data_to_file,
 )
 from rasa.model_manager.warm_rasa_process import (
@@ -21,7 +22,6 @@ from rasa.model_manager.warm_rasa_process import (
 )
 from rasa.model_training import generate_random_model_name
 from rasa.studio.prompts import handle_prompts
-from rasa.utils.io import subpath
 
 structlogger = structlog.get_logger()
 
@@ -53,15 +53,6 @@ class TrainingSession(BaseModel):
         """Check if the training is running."""
         return self.status == TrainingSessionStatus.RUNNING
 
-    def has_just_finished(self) -> bool:
-        if not self.is_status_indicating_alive():
-            # skip if the training is not running
-            return False
-        if self.process.poll() is None:
-            # process is still running
-            return False
-        return True
-
     def model_path(self) -> str:
         """Return the path to the model."""
         return subpath(models_base_path(), f"{self.model_name}.tar.gz")
@@ -98,8 +89,14 @@ def terminate_training(training: TrainingSession) -> None:
 
 
 def update_training_status(training: TrainingSession) -> None:
-    if training.has_just_finished():
-        complete_training(training)
+    if not training.is_status_indicating_alive():
+        # skip if the training is not running
+        return
+    if training.process.poll() is None:
+        # process is still running
+        return
+
+    complete_training(training)
 
 
 def complete_training(training: TrainingSession) -> None: