rasa-pro 3.14.0.dev9__py3-none-any.whl → 3.14.0.dev11__py3-none-any.whl

rasa/nlu/classifiers/mitie_intent_classifier.py

@@ -4,6 +4,9 @@ import logging
 import typing
 from typing import Any, Dict, List, Optional, Text, Type
 
+# importing mitie at module level to ensure error is raised if mitie is not installed
+import mitie
+
 from rasa.engine.graph import ExecutionContext, GraphComponent
 from rasa.engine.recipes.default_recipe import DefaultV1Recipe
 from rasa.engine.storage.resource import Resource

rasa/nlu/classifiers/sklearn_intent_classifier.py

@@ -6,6 +6,7 @@ import warnings
 from typing import Any, Dict, List, Optional, Text, Tuple, Type
 
 import numpy as np
+import skops.io as sio
 
 import rasa.shared.utils.io
 from rasa.engine.graph import ExecutionContext, GraphComponent
@@ -266,8 +267,6 @@ class SklearnIntentClassifier(GraphComponent, IntentClassifier):
 
     def persist(self) -> None:
         """Persist this model into the passed directory."""
-        import skops.io as sio
-
         with self._model_storage.write_to(self._resource) as model_dir:
             file_name = self.__class__.__name__
             classifier_file_name = model_dir / f"{file_name}_classifier.skops"
@@ -291,7 +290,6 @@ class SklearnIntentClassifier(GraphComponent, IntentClassifier):
         **kwargs: Any,
     ) -> SklearnIntentClassifier:
         """Loads trained component (see parent class for full docstring)."""
-        import skops.io as sio
         from sklearn.preprocessing import LabelEncoder
 
         try:

rasa/nlu/extractors/crf_entity_extractor.py

@@ -10,6 +10,9 @@ from typing import Any, Callable, Dict, List, Optional, Text, Tuple, Type
 
 import numpy as np
 
+# importing sklearn_crfsuite at module level to raise error if not installed
+import sklearn_crfsuite
+
 import rasa.nlu.utils.bilou_utils as bilou_utils
 import rasa.shared.utils.io
 import rasa.utils.train_utils
@@ -20,7 +23,6 @@ from rasa.engine.storage.storage import ModelStorage
 from rasa.nlu.constants import TOKENS_NAMES
 from rasa.nlu.extractors.extractor import EntityExtractorMixin
 from rasa.nlu.test import determine_token_labels
-from rasa.nlu.tokenizers.spacy_tokenizer import POS_TAG_KEY
 from rasa.nlu.tokenizers.tokenizer import Token, Tokenizer
 from rasa.shared.constants import DOCS_URL_COMPONENTS
 from rasa.shared.nlu.constants import (
@@ -37,6 +39,8 @@ from rasa.shared.nlu.training_data.message import Message
 from rasa.shared.nlu.training_data.training_data import TrainingData
 from rasa.utils.tensorflow.constants import BILOU_FLAG, FEATURIZERS
 
+POS_TAG_KEY = "pos"
+
 logger = logging.getLogger(__name__)
 
 if typing.TYPE_CHECKING:
@@ -445,13 +449,11 @@ class CRFEntityExtractor(GraphComponent, EntityExtractorMixin):
     def _load_taggers(
         cls, model_dir: Path, config: Dict[Text, Any]
     ) -> Dict[str, "CRF"]:
-        """
-        Load taggers from model directory that persists trained binary
-        `model.crfsuite` files.
-        """
+        """Load taggers from model directory.
 
+        It persists trained binary `model.crfsuite` files.
+        """
         import pycrfsuite
-        import sklearn_crfsuite
 
         # Get tagger directories
         taggers_base = model_dir / TAGGERS_DIR
@@ -613,8 +615,7 @@ class CRFEntityExtractor(GraphComponent, EntityExtractorMixin):
 
     @staticmethod
     def _pattern_of_token(message: Message, idx: int) -> Dict[Text, bool]:
-        """Get the patterns of the token at the given index extracted by the
-        'RegexFeaturizer'.
+        """Get patterns of token at the given index extracted by 'RegexFeaturizer'.
 
         The 'RegexFeaturizer' adds all patterns listed in the training data to the
         token. The pattern name is mapped to either 'True' (pattern applies to token) or
@@ -728,8 +729,6 @@ class CRFEntityExtractor(GraphComponent, EntityExtractorMixin):
         crf_order: List[str],
     ) -> OrderedDict[str, CRF]:
         """Train the crf tagger based on the training data."""
-        import sklearn_crfsuite
-
         entity_taggers = OrderedDict()
 
         for tag_name in crf_order:

rasa/nlu/extractors/mitie_entity_extractor.py

@@ -4,6 +4,9 @@ import logging
 import typing
 from typing import Any, Dict, List, Optional, Text, Type
 
+# importing mitie at module level to ensure error is raised if mitie is not installed
+import mitie
+
 import rasa.shared.utils.io
 from rasa.engine.graph import ExecutionContext, GraphComponent
 from rasa.engine.recipes.default_recipe import DefaultV1Recipe

rasa/nlu/extractors/spacy_entity_extractor.py

@@ -1,6 +1,9 @@
 import typing
 from typing import Any, Dict, List, Text, Type
 
+# importing spacy at module level to ensure error is raised if spacy is not installed
+import spacy  # noqa: F401
+
 from rasa.engine.graph import ExecutionContext, GraphComponent
 from rasa.engine.recipes.default_recipe import DefaultV1Recipe
 from rasa.engine.storage.resource import Resource

rasa/nlu/featurizers/dense_featurizer/mitie_featurizer.py

@@ -4,6 +4,8 @@ import logging
 import typing
 from typing import Any, Dict, List, Text, Tuple, Type
 
+# importing mitie at module level to ensure error is raised if mitie is not installed
+import mitie
 import numpy as np
 
 from rasa.engine.graph import ExecutionContext, GraphComponent

rasa/nlu/featurizers/dense_featurizer/spacy_featurizer.py

@@ -4,6 +4,9 @@ from typing import Any, Dict, List, Text, Type
 
 import numpy as np
 
+# importing spacy at module level to ensure error is raised if spacy is not installed
+import spacy  # noqa: F401
+
 from rasa.engine.graph import ExecutionContext, GraphComponent
 from rasa.engine.recipes.default_recipe import DefaultV1Recipe
 from rasa.engine.storage.resource import Resource

rasa/nlu/featurizers/sparse_featurizer/count_vectors_featurizer.py

@@ -21,12 +21,14 @@ from rasa.nlu.constants import (
 )
 from rasa.nlu.featurizers.sparse_featurizer.sparse_featurizer import SparseFeaturizer
 from rasa.nlu.tokenizers.tokenizer import Tokenizer
-from rasa.nlu.utils.spacy_utils import SpacyModel
 from rasa.shared.constants import DOCS_URL_COMPONENTS
 from rasa.shared.exceptions import FileIOException, RasaException
 from rasa.shared.nlu.constants import ACTION_NAME, INTENT, INTENT_RESPONSE_KEY, TEXT
 from rasa.shared.nlu.training_data.message import Message
 from rasa.shared.nlu.training_data.training_data import TrainingData
+from rasa.shared.utils.common import conditional_import
+
+SpacyModel, _ = conditional_import("rasa.nlu.utils.spacy_utils", "SpacyModel")
 
 BUFFER_SLOTS_PREFIX = "buf_"
 
@@ -603,7 +605,7 @@ class CountVectorsFeaturizer(SparseFeaturizer, GraphComponent):
             return [], []
 
     def train(
-        self, training_data: TrainingData, model: Optional[SpacyModel] = None
+        self, training_data: TrainingData, model: Optional[Any] = None
     ) -> Resource:
         """Trains the featurizer.
 

rasa/nlu/featurizers/sparse_featurizer/lexical_syntactic_featurizer.py

@@ -18,6 +18,10 @@ from typing import (
 import numpy as np
 import scipy.sparse
 
+# importing spacy at module level to ensure error is raised if spacy is not installed
+# spacy is needed when importing from rasa.nlu.tokenizers.spacy_tokenizer
+import spacy  # noqa: F401
+
 import rasa.shared.utils.io
 import rasa.utils.io
 from rasa.engine.graph import ExecutionContext, GraphComponent

rasa/nlu/tokenizers/jieba_tokenizer.py

@@ -6,6 +6,9 @@ import os
 import shutil
 from typing import Any, Dict, List, Optional, Text
 
+# importing jieba at module level to ensure error is raised if jieba is not installed
+import jieba
+
 from rasa.engine.graph import ExecutionContext
 from rasa.engine.recipes.default_recipe import DefaultV1Recipe
 from rasa.engine.storage.resource import Resource
@@ -81,8 +84,6 @@ class JiebaTokenizer(Tokenizer):
         be found in the documentation of jieba.
         https://github.com/fxsjy/jieba#load-dictionary
         """
-        import jieba
-
         jieba_userdicts = glob.glob(f"{path}/*")
         for jieba_userdict in jieba_userdicts:
             logger.info(f"Loading Jieba User Dictionary at {jieba_userdict}")
@@ -95,8 +96,6 @@ class JiebaTokenizer(Tokenizer):
 
     def tokenize(self, message: Message, attribute: Text) -> List[Token]:
         """Tokenizes the text of the provided attribute of the incoming message."""
-        import jieba
-
         text = message.get(attribute)
 
         tokenized = jieba.tokenize(text)

rasa/nlu/tokenizers/mitie_tokenizer.py

@@ -2,6 +2,9 @@ from __future__ import annotations
 
 from typing import Any, Dict, List, Text
 
+# importing mitie at module level to ensure error is raised if mitie is not installed
+import mitie
+
 from rasa.engine.graph import ExecutionContext
 from rasa.engine.recipes.default_recipe import DefaultV1Recipe
 from rasa.engine.storage.resource import Resource
@@ -49,8 +52,6 @@ class MitieTokenizer(Tokenizer):
 
     def tokenize(self, message: Message, attribute: Text) -> List[Token]:
         """Tokenizes the text of the provided attribute of the incoming message."""
-        import mitie
-
         text = message.get(attribute)
 
         encoded_sentence = text.encode(DEFAULT_ENCODING)

rasa/nlu/tokenizers/spacy_tokenizer.py

@@ -1,6 +1,9 @@
 import typing
 from typing import Any, Dict, List, Optional, Text, Type
 
+# importing spacy at module level to ensure error is raised if spacy is not installed
+import spacy
+
 from rasa.engine.recipes.default_recipe import DefaultV1Recipe
 from rasa.nlu.constants import SPACY_DOCS
 from rasa.nlu.tokenizers.tokenizer import Token, Tokenizer
@@ -64,8 +67,6 @@ class SpacyTokenizer(Tokenizer):
 
     @staticmethod
     def _tag_of_token(token: Any) -> Text:
-        import spacy
-
         if spacy.about.__version__ > "2" and token._.has("tag"):
             return token._.get("tag")
         else:

rasa/nlu/utils/mitie_utils.py

@@ -4,6 +4,9 @@ import typing
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Text
 
+# importing mitie at module level to ensure error is raised if mitie is not installed
+import mitie
+
 from rasa.engine.graph import ExecutionContext, GraphComponent
 from rasa.engine.recipes.default_recipe import DefaultV1Recipe
 from rasa.engine.storage.resource import Resource

rasa/nlu/utils/spacy_utils.py

@@ -5,6 +5,9 @@ import logging
 import typing
 from typing import Any, Dict, List, Optional, Text, Tuple
 
+# importing spacy at module level to ensure error is raised if spacy is not installed
+import spacy
+
 from rasa.engine.graph import ExecutionContext, GraphComponent
 from rasa.engine.recipes.default_recipe import DefaultV1Recipe
 from rasa.engine.storage.resource import Resource
@@ -76,8 +79,6 @@ class SpacyNLP(GraphComponent):
     @staticmethod
     def load_model(spacy_model_name: Text) -> SpacyModel:
         """Try loading the model, catching the OSError if missing."""
-        import spacy
-
         if not spacy_model_name:
             raise InvalidModelError(
                 f"Missing model configuration for `SpacyNLP` in `config.yml`.\n"

rasa/shared/agents/auth/agent_auth_manager.py

@@ -51,7 +51,9 @@ class AgentAuthManager:
         )
 
     @classmethod
-    def load_auth(cls, config: Dict[str, Any]) -> "AgentAuthManager":
+    def load_auth(
+        cls, config: Optional[Dict[str, Any]]
+    ) -> Optional["AgentAuthManager"]:
         """Connect to authentication using specified strategy type and persist
         the auth instance to the manager in a ready-to-use state.
 
@@ -61,7 +63,8 @@ class AgentAuthManager:
         Raises:
             AgentAuthInitializationException: If the authentication connection fails.
         """
-        config = config or {}
+        if not config:
+            return None
         try:
             auth_type = AgentAuthManager.detect_auth_type(config)
 
@@ -83,4 +86,4 @@ class AgentAuthManager:
                 event_info=event_info,
                 config=config,
             )
-            raise AgentAuthInitializationException(e)
+            raise AgentAuthInitializationException(e) from e

rasa/shared/agents/auth/auth_strategy/oauth2_auth_strategy.py

@@ -10,6 +10,7 @@ from rasa.shared.agents.auth.auth_strategy import AgentAuthStrategy
 from rasa.shared.agents.auth.constants import (
     CONFIG_CLIENT_ID_KEY,
     CONFIG_CLIENT_SECRET_KEY,
+    CONFIG_OAUTH_KEY,
     CONFIG_SCOPE_KEY,
     CONFIG_TIMEOUT_KEY,
     CONFIG_TOKEN_URL_KEY,
@@ -56,11 +57,18 @@ class OAuth2AuthStrategy(AgentAuthStrategy):
     @classmethod
     def from_config(cls, config: Dict[str, Any]) -> "OAuth2AuthStrategy":
         """Create OAuth2AuthStrategy from configuration dictionary."""
-        token_url = config.get(CONFIG_TOKEN_URL_KEY)
-        client_id = config.get(CONFIG_CLIENT_ID_KEY)
-        client_secret = config.get(CONFIG_CLIENT_SECRET_KEY)
-        scope = config.get(CONFIG_SCOPE_KEY)
-        timeout = config.get(CONFIG_TIMEOUT_KEY) or cls.DEFAULT_ACCESS_TOKEN_TIMEOUT
+        # Extract OAuth2 config from nested "oauth" key if present
+        oauth_config = config.get(CONFIG_OAUTH_KEY, config)
+        if not oauth_config:
+            raise ValueError("OAuth2 configuration is required")
+
+        token_url = oauth_config.get(CONFIG_TOKEN_URL_KEY)
+        client_id = oauth_config.get(CONFIG_CLIENT_ID_KEY)
+        client_secret = oauth_config.get(CONFIG_CLIENT_SECRET_KEY)
+        scope = oauth_config.get(CONFIG_SCOPE_KEY)
+        timeout = (
+            oauth_config.get(CONFIG_TIMEOUT_KEY) or cls.DEFAULT_ACCESS_TOKEN_TIMEOUT
+        )
 
         if not token_url:
             raise ValueError("Token URL is required for OAuth2 authentication")

rasa/shared/core/constants.py

@@ -129,7 +129,7 @@ GLOBAL_SILENCE_TIMEOUT_KEY = "global_silence_timeout"
 SILENCE_TIMEOUT_SLOT = "silence_timeout"
 SLOT_CONSECUTIVE_SILENCE_TIMEOUTS = "consecutive_silence_timeouts"
 GLOBAL_SILENCE_TIMEOUT_DEFAULT_VALUE = 7.0
-SILENCE_SLOTS = [SILENCE_TIMEOUT_SLOT, SLOT_CONSECUTIVE_SILENCE_TIMEOUTS]
+SILENCE_SLOTS = {SILENCE_TIMEOUT_SLOT, SLOT_CONSECUTIVE_SILENCE_TIMEOUTS}
 # slots for knowledge base
 SLOT_LISTED_ITEMS = "knowledge_base_listed_objects"
 SLOT_LAST_OBJECT = "knowledge_base_last_object"
@@ -148,6 +148,10 @@ DEFAULT_SLOT_NAMES = {
     FLOW_HASHES_SLOT,
 }
 
+SLOTS_EXCLUDED_FOR_AGENT = (
+    SILENCE_SLOTS | DEFAULT_SLOT_NAMES | KNOWLEDGE_BASE_SLOT_NAMES
+)
+
 SLOT_MAPPINGS = "mappings"
 MAPPING_CONDITIONS = "conditions"
 KEY_MAPPING_TYPE = "type"

rasa/shared/utils/mcp/server_connection.py

@@ -67,7 +67,7 @@ class MCPServerConnection:
     def from_config(cls, server_config: Dict[str, Any]) -> "MCPServerConnection":
         """Initialize the MCP server connection from a configuration dictionary."""
         auth_config = server_config.get("additional_params")
-        _auth_manager = AgentAuthManager.load_auth(auth_config) if auth_config else None
+        _auth_manager = AgentAuthManager.load_auth(auth_config)
         return cls(
             server_config["name"],
             server_config["url"],

rasa/telemetry.py

@@ -50,6 +50,7 @@ from rasa.shared.constants import (
     UTTER_ASK_PREFIX,
 )
 from rasa.shared.core.flows import Flow
+from rasa.shared.core.flows.flows_list import FlowsList
 from rasa.shared.core.flows.steps import (
     CallFlowStep,
     CollectInformationFlowStep,
@@ -1083,6 +1084,7 @@ def track_model_training(
         "num_regexes": len(nlu_data.regex_features),
         "is_finetuning": is_finetuning,
         "recipe": config.get(CONFIG_RECIPE_KEY),
+        "agents": _collect_agent_configuration(flows),
     }
 
     flow_statistics = _collect_flow_statistics(flows.underlying_flows)
@@ -1125,6 +1127,67 @@ def track_model_training(
     )
 
 
+def _collect_agent_configuration(flows: FlowsList) -> Dict[str, Any]:
+    agent_data: Dict[str, Any] = {}
+
+    from rasa.core.available_agents import AvailableAgents
+    from rasa.core.available_endpoints import AvailableEndpoints
+
+    agents = AvailableAgents.get_instance().agents
+    mcp_servers = AvailableEndpoints.get_instance().mcp_servers
+
+    if not agents and not mcp_servers:
+        return agent_data
+
+    agent_data["usage"] = []
+
+    for flow in flows.underlying_flows:
+        for step in flow.steps:
+            if isinstance(step, CallFlowStep):
+                if flows.flow_by_id(step.call) is not None:
+                    continue
+
+                if step.is_calling_mcp_tool():
+                    agent_data["usage"].append(
+                        {
+                            "flow": flow.id,
+                            "mcp_tool": step.call,
+                            "mcp_server": step.mcp_server,
+                            "mapping": step.mapping,
+                        }
+                    )
+
+                if step.is_calling_agent():
+                    if step.exit_if:
+                        agent_data["usage"].append(
+                            {
+                                "flow": flow.id,
+                                "agent": step.call,
+                                "exit_if": step.exit_if,
+                            }
+                        )
+                    else:
+                        agent_data["usage"].append(
+                            {
+                                "flow": flow.id,
+                                "agent": step.call,
+                            }
+                        )
+
+    # dump mcp_servers and agents as json and exclude any values that are None
+    if mcp_servers:
+        agent_data["mcp_servers"] = [
+            mcp_server.model_dump(exclude_none=True) for mcp_server in mcp_servers
+        ]
+    if agents:
+        agent_data["agents"] = [
+            {agent_name: agent_info.model_dump(exclude_none=True)}
+            for agent_name, agent_info in agents.items()
+        ]
+
+    return agent_data
+
+
 def _collect_flow_statistics(flows: List[Flow]) -> Dict[str, Any]:
     """Collects some statistics about the flows, such as number of specific steps."""
     data = {

rasa/utils/log_utils.py

@@ -19,19 +19,110 @@ from rasa.shared.constants import (
 )
 
 FORCE_JSON_LOGGING = os.environ.get("FORCE_JSON_LOGGING")
+ANSI_CYAN_BOLD = "\033[1;36m"
+ANSI_RESET = "\033[0m"
 
 
 class HumanConsoleRenderer(ConsoleRenderer):
     """Console renderer that outputs human-readable logs."""
 
     def __call__(self, logger: WrappedLogger, name: str, event_dict: EventDict) -> str:
+        should_highlight = event_dict.get("highlight", False)
+        terminal_width = self._get_terminal_width()
+
+        # Use event_info as title for the log entry
         if "event_info" in event_dict:
             event_key = event_dict["event"]
             event_dict["event"] = event_dict["event_info"]
             event_dict["event_key"] = event_key
-            del event_dict["event_info"]
-
-        return super().__call__(logger, name, event_dict)
+            event_dict.pop("event_info", None)
+
+        # In case the log entry should be highlighted
+        # make sure to surround the log entry with ===
+        event_dict = self._highlight_log_entry(
+            event_dict, terminal_width, should_highlight
+        )
+
+        # Format JSON data for better readability
+        event_dict = self._format_json_data(event_dict)
+
+        # Render the log entry first
+        result = super().__call__(logger, name, event_dict)
+
+        # ensure that newlines are properly rendered
+        result = "\n".join(result.split("\\n"))
+
+        # Add closing === if we highlighted this entry
+        if should_highlight:
+            result += f"\n{'=' * terminal_width}"
+
+        return result
+
+    def _highlight_log_entry(
+        self, event_dict: EventDict, terminal_width: int, should_highlight: bool
+    ) -> EventDict:
+        if should_highlight:
+            # Only highlight if log level is DEBUG
+            # structlog passes log level as 'level' or 'levelname'
+            level = event_dict.get("level", event_dict.get("levelname", "")).upper()
+            if level == "DEBUG":
+                # Add opening === before the event info
+                if "event" in event_dict and isinstance(event_dict["event"], str):
+                    event_info = event_dict["event"]
+                    event_dict["event"] = (
+                        f"\n{'=' * terminal_width}\n"
+                        f"{ANSI_CYAN_BOLD}{event_info}{ANSI_RESET}\n"
+                    )
+
+            event_dict.pop("highlight", None)
+
+        return event_dict
+
+    def _format_json_data(self, event_dict: EventDict) -> EventDict:
+        """Format JSON data in the event dict for better readability."""
+        # Get the list of fields to format from the event dict
+        fields_to_format = event_dict.get("json_formatting", [])
+
+        if not fields_to_format:
+            return event_dict
+
+        import json
+
+        # Format only the specified fields
+        for key in fields_to_format:
+            if key in event_dict:
+                value = event_dict[key]
+
+                try:
+                    # Try to parse as JSON if it's a string
+                    if isinstance(value, str):
+                        parsed = json.loads(value)
+                        # If it's a dict or list, format it nicely
+                        if isinstance(parsed, (dict, list)):
+                            formatted = json.dumps(parsed, indent=2, ensure_ascii=False)
+                            event_dict[key] = formatted
+                    elif isinstance(value, (dict, list)):
+                        # Format JSON with indentation for better readability
+                        formatted = json.dumps(value, indent=2, ensure_ascii=False)
+                        event_dict[key] = formatted
+                except (TypeError, ValueError, json.JSONDecodeError):
+                    # If it's not JSON serializable or if it's not valid JSON,
+                    # leave it as is
+                    pass
+
+        # Remove the json_formatting key from the output
+        event_dict.pop("json_formatting", None)
+
+        return event_dict
+
+    def _get_terminal_width(self) -> int:
+        """Get the width of the terminal."""
+        import shutil
+
+        try:
+            return shutil.get_terminal_size((80, 20)).columns
+        except Exception:
+            return 80
 
 
 def configure_structlog(
@@ -52,7 +143,7 @@ def configure_structlog(
         level=log_level,
     )
 
-    shared_processors = [
+    shared_processors: List[structlog.typing.Processor] = [
         # Processors that have nothing to do with output,
         # e.g., add timestamps or log level names.
         # If log level is too low, abort pipeline and throw away log entry.

rasa/utils/tensorflow/__init__.py

@@ -1,25 +1,3 @@
-"""TensorFlow utilities for Rasa components.
-
-This module provides TensorFlow-related utilities and components.
-TensorFlow is an optional dependency and should be installed separately
-if you want to use components that require it. These are:
-- DIETClassifier
-- TEDPolicy
-- UnexpecTEDIntentPolicy
-- ResponseSelector
-- ConveRTFeaturizer
-- LanguageModelFeaturizer
-
-To install Rasa with TensorFlow support:
-`pip install "rasa[tensorflow]"`
-
-To install it with poetry:
-`poetry install --extras tensorflow`
-
-For macOS with Apple Silicon (M1/M2) (platform-specific TensorFlow installation):
-`pip install "rasa[tensorflow,metal]"`
-"""
-
 import importlib.util
 import logging
 

rasa/utils/tensorflow/callback.py

@@ -67,15 +67,24 @@ class RasaModelCheckpoint(tf.keras.callbacks.Callback):
         """
         if self._does_model_improve(logs):
             logger.debug(f"Creating model checkpoint at epoch={epoch + 1} ...")
-            self.model.save_weights(
-                self.checkpoint_file, overwrite=True, save_format="tf"
-            )
+            # Ensure model is built before saving weights
+            if not self.model.built:
+                # Build the model with dummy data to ensure it's built
+                import tensorflow as tf
+
+                dummy_input = tf.zeros((1, 1))
+                _ = self.model(dummy_input)
+                # Ensure the directory exists before saving
+                import os
+
+                os.makedirs(os.path.dirname(self.checkpoint_file), exist_ok=True)
+                self.model.save_weights(self.checkpoint_file, overwrite=True)
 
     def _does_model_improve(self, curr_results: Dict[Text, Any]) -> bool:
         """Checks whether the current results are better than the best so far.
 
-        Results are considered better if each metric is equal or better than the best so
-        far, and at least one is better.
+        Results are considered better if each metric is
+        equal or better than the best so far, and at least one is better.
 
         Args:
             curr_results: The training metrics for this epoch.
@@ -85,7 +94,8 @@ class RasaModelCheckpoint(tf.keras.callbacks.Callback):
             for k in curr_results.keys()
             if k.startswith("val") and (k.endswith("_acc") or k.endswith("_f1"))
         ]
-        # the "val" prefix is prepended to metrics in fit if _should_eval returns true
+        # the "val" prefix is prepended to metrics in fit
+        # if _should_eval returns true
         # for this particular epoch
         if len(curr_metric_names) == 0:
             # the metrics are not validation metrics
@@ -94,7 +104,7 @@ class RasaModelCheckpoint(tf.keras.callbacks.Callback):
         if not self.best_metrics_so_far:
             for metric_name in curr_metric_names:
                 self.best_metrics_so_far[metric_name] = float(curr_results[metric_name])
-            return True
+                return True
 
         at_least_one_improved = False
         improved_metrics = {}