camel-ai 0.2.49__py3-none-any.whl → 0.2.50__py3-none-any.whl

camel/__init__.py

@@ -14,7 +14,7 @@
 
 from camel.logger import disable_logging, enable_logging, set_log_level
 
-__version__ = '0.2.49'
+__version__ = '0.2.50'
 
 __all__ = [
     '__version__',

camel/agents/chat_agent.py

@@ -28,6 +28,7 @@ from typing import (
     List,
     Optional,
     Set,
+    Tuple,
     Type,
     Union,
 )
@@ -78,7 +79,6 @@ from camel.utils import get_model_encoding
 if TYPE_CHECKING:
     from camel.terminators import ResponseTerminator
 
-
 logger = logging.getLogger(__name__)
 
 # AgentOps decorator setting
@@ -110,10 +110,17 @@ class ChatAgent(BaseAgent):
 
     Args:
         system_message (Union[BaseMessage, str], optional): The system message
-            for the chat agent.
-        model (BaseModelBackend, optional): The model backend to use for
-            generating responses. (default: :obj:`ModelPlatformType.DEFAULT`
-            with `ModelType.DEFAULT`)
+            for the chat agent. (default: :obj:`None`)
+        model (Union[BaseModelBackend, Tuple[str, str], str, ModelType,
+            Tuple[ModelPlatformType, ModelType], List[BaseModelBackend],
+            List[str], List[ModelType], List[Tuple[str, str]],
+            List[Tuple[ModelPlatformType, ModelType]]], optional):
+            The model backend(s) to use. Can be a single instance,
+            a specification (string, enum, tuple), or a list of instances
+            or specifications to be managed by `ModelManager`. If a list of
+            specifications (not `BaseModelBackend` instances) is provided,
+            they will be instantiated using `ModelFactory`. (default:
+            :obj:`ModelPlatformType.DEFAULT` with `ModelType.DEFAULT`)
         memory (AgentMemory, optional): The agent memory for managing chat
             messages. If `None`, a :obj:`ChatHistoryMemory` will be used.
             (default: :obj:`None`)
@@ -150,7 +157,18 @@ class ChatAgent(BaseAgent):
         self,
         system_message: Optional[Union[BaseMessage, str]] = None,
         model: Optional[
-            Union[BaseModelBackend, List[BaseModelBackend]]
+            Union[
+                BaseModelBackend,
+                Tuple[str, str],
+                str,
+                ModelType,
+                Tuple[ModelPlatformType, ModelType],
+                List[BaseModelBackend],
+                List[str],
+                List[ModelType],
+                List[Tuple[str, str]],
+                List[Tuple[ModelPlatformType, ModelType]],
+            ]
         ] = None,
         memory: Optional[AgentMemory] = None,
         message_window_size: Optional[int] = None,
@@ -165,19 +183,14 @@ class ChatAgent(BaseAgent):
         single_iteration: bool = False,
         agent_id: Optional[str] = None,
     ) -> None:
-        # Set up model backend
+        # Resolve model backends and set up model manager
+        resolved_models = self._resolve_models(model)
         self.model_backend = ModelManager(
-            (
-                model
-                if model is not None
-                else ModelFactory.create(
-                    model_platform=ModelPlatformType.DEFAULT,
-                    model_type=ModelType.DEFAULT,
-                )
-            ),
+            resolved_models,
             scheduling_strategy=scheduling_strategy,
         )
         self.model_type = self.model_backend.model_type
+
         # Assign unique ID
         self.agent_id = agent_id if agent_id else str(uuid.uuid4())
 
@@ -247,6 +260,137 @@ class ChatAgent(BaseAgent):
         for terminator in self.response_terminators:
             terminator.reset()
 
+    def _resolve_models(
+        self,
+        model: Optional[
+            Union[
+                BaseModelBackend,
+                Tuple[str, str],
+                str,
+                ModelType,
+                Tuple[ModelPlatformType, ModelType],
+                List[BaseModelBackend],
+                List[str],
+                List[ModelType],
+                List[Tuple[str, str]],
+                List[Tuple[ModelPlatformType, ModelType]],
+            ]
+        ],
+    ) -> Union[BaseModelBackend, List[BaseModelBackend]]:
+        r"""Resolves model specifications into model backend instances.
+
+        This method handles various input formats for model specifications and
+        returns the appropriate model backend(s).
+
+        Args:
+            model: Model specification in various formats including single
+                model, list of models, or model type specifications.
+
+        Returns:
+            Union[BaseModelBackend, List[BaseModelBackend]]: Resolved model
+                backend(s).
+
+        Raises:
+            TypeError: If the model specification format is not supported.
+        """
+        if model is None:
+            # Default single model if none provided
+            return ModelFactory.create(
+                model_platform=ModelPlatformType.DEFAULT,
+                model_type=ModelType.DEFAULT,
+            )
+        elif isinstance(model, BaseModelBackend):
+            # Already a single pre-instantiated model
+            return model
+        elif isinstance(model, list):
+            return self._resolve_model_list(model)
+        elif isinstance(model, (ModelType, str)):
+            # Single string or ModelType -> use default platform
+            model_platform = ModelPlatformType.DEFAULT
+            model_type = model
+            logger.warning(
+                f"Model type '{model_type}' provided without a platform. "
+                f"Using platform '{model_platform}'. Note: platform "
+                "is not automatically inferred based on model type."
+            )
+            return ModelFactory.create(
+                model_platform=model_platform,
+                model_type=model_type,
+            )
+        elif isinstance(model, tuple) and len(model) == 2:
+            # Single tuple (platform, type)
+            model_platform, model_type = model  # type: ignore[assignment]
+            return ModelFactory.create(
+                model_platform=model_platform,
+                model_type=model_type,
+            )
+        else:
+            raise TypeError(
+                f"Unsupported type for model parameter: {type(model)}"
+            )
+
+    def _resolve_model_list(
+        self, model_list: list
+    ) -> Union[BaseModelBackend, List[BaseModelBackend]]:
+        r"""Resolves a list of model specifications into model backend
+        instances.
+
+        Args:
+            model_list (list): List of model specifications in various formats.
+
+        Returns:
+            Union[BaseModelBackend, List[BaseModelBackend]]: Resolved model
+                backend(s).
+
+        Raises:
+            TypeError: If the list elements format is not supported.
+        """
+        if not model_list:  # Handle empty list
+            logger.warning(
+                "Empty list provided for model, using default model."
+            )
+            return ModelFactory.create(
+                model_platform=ModelPlatformType.DEFAULT,
+                model_type=ModelType.DEFAULT,
+            )
+        elif isinstance(model_list[0], BaseModelBackend):
+            # List of pre-instantiated models
+            return model_list  # type: ignore[return-value]
+        elif isinstance(model_list[0], (str, ModelType)):
+            # List of strings or ModelTypes -> use default platform
+            model_platform = ModelPlatformType.DEFAULT
+            logger.warning(
+                f"List of model types {model_list} provided without "
+                f"platforms. Using platform '{model_platform}' for all. "
+                "Note: platform is not automatically inferred based on "
+                "model type."
+            )
+            resolved_models_list = []
+            for model_type_item in model_list:
+                resolved_models_list.append(
+                    ModelFactory.create(
+                        model_platform=model_platform,
+                        model_type=model_type_item,  # type: ignore[arg-type]
+                    )
+                )
+            return resolved_models_list
+        elif isinstance(model_list[0], tuple) and len(model_list[0]) == 2:
+            # List of tuples (platform, type)
+            resolved_models_list = []
+            for model_spec in model_list:
+                platform, type_ = model_spec[0], model_spec[1]  # type: ignore[index]
+                resolved_models_list.append(
+                    ModelFactory.create(
+                        model_platform=platform, model_type=type_
+                    )
+                )
+            return resolved_models_list
+        else:
+            raise TypeError(
+                "Unsupported type for list elements in model: "
+                f"{type(model_list[0])}"
+            )
+
     @property
     def system_message(self) -> Optional[BaseMessage]:
         r"""Returns the system message for the agent."""

camel/configs/__init__.py

@@ -26,6 +26,7 @@ from .mistral_config import MISTRAL_API_PARAMS, MistralConfig
 from .modelscope_config import MODELSCOPE_API_PARAMS, ModelScopeConfig
 from .moonshot_config import MOONSHOT_API_PARAMS, MoonshotConfig
 from .netmind_config import NETMIND_API_PARAMS, NetmindConfig
+from .novita_config import NOVITA_API_PARAMS, NovitaConfig
 from .nvidia_config import NVIDIA_API_PARAMS, NvidiaConfig
 from .ollama_config import OLLAMA_API_PARAMS, OllamaConfig
 from .openai_config import OPENAI_API_PARAMS, ChatGPTConfig
@@ -43,6 +44,7 @@ from .sglang_config import SGLANG_API_PARAMS, SGLangConfig
 from .siliconflow_config import SILICONFLOW_API_PARAMS, SiliconFlowConfig
 from .togetherai_config import TOGETHERAI_API_PARAMS, TogetherAIConfig
 from .vllm_config import VLLM_API_PARAMS, VLLMConfig
+from .watsonx_config import WATSONX_API_PARAMS, WatsonXConfig
 from .yi_config import YI_API_PARAMS, YiConfig
 from .zhipuai_config import ZHIPUAI_API_PARAMS, ZhipuAIConfig
 
@@ -98,6 +100,8 @@ __all__ = [
     "MOONSHOT_API_PARAMS",
     'ModelScopeConfig',
     'MODELSCOPE_API_PARAMS',
+    'NovitaConfig',
+    'NOVITA_API_PARAMS',
     'SiliconFlowConfig',
     'SILICONFLOW_API_PARAMS',
     'AIMLConfig',
@@ -106,4 +110,6 @@ __all__ = [
     'OPENROUTER_API_PARAMS',
     'LMSTUDIO_API_PARAMS',
     'LMStudioConfig',
+    'WatsonXConfig',
+    'WATSONX_API_PARAMS',
 ]

camel/configs/modelscope_config.py

@@ -13,7 +13,7 @@
 # ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
 from __future__ import annotations
 
-from typing import Optional, Union
+from typing import Any, Dict, Optional, Union
 
 from camel.configs.base_config import BaseConfig
 
@@ -45,6 +45,8 @@ class ModelScopeConfig(BaseConfig):
             while higher values make it more diverse. (default: :obj:`0.3`)
         stream (bool, optional): If True, enables streaming output.
             (default: :obj:`None`)
+        extra_body (dict, optional): Extra body parameters to be passed to
+            the ModelScope API.
     """
 
     tool_choice: Optional[Union[dict[str, str], str]] = None
@@ -52,6 +54,7 @@ class ModelScopeConfig(BaseConfig):
     top_p: Optional[float] = None
     temperature: Optional[float] = None
     stream: Optional[bool] = None
+    extra_body: Optional[Dict[str, Any]] = None
 
 
 MODELSCOPE_API_PARAMS = {

camel/configs/novita_config.py

@@ -0,0 +1,102 @@
+# ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
+from __future__ import annotations
+
+from typing import Dict, Optional, Sequence, Type, Union
+
+from pydantic import BaseModel
+
+from camel.configs.base_config import BaseConfig
+
+
+class NovitaConfig(BaseConfig):
+    r"""Defines the parameters for generating chat completions using the
+    OpenAI API.
+
+    Args:
+        temperature (float, optional): Sampling temperature to use, between
+            :obj:`0` and :obj:`2`. Higher values make the output more random,
+            while lower values make it more focused and deterministic.
+            (default: :obj:`None`)
+        top_p (float, optional): An alternative to sampling with temperature,
+            called nucleus sampling, where the model considers the results of
+            the tokens with top_p probability mass. So :obj:`0.1` means only
+            the tokens comprising the top 10% probability mass are considered.
+            (default: :obj:`None`)
+        n (int, optional): How many chat completion choices to generate for
+            each input message. (default: :obj:`None`)
+        response_format (object, optional): An object specifying the format
+            that the model must output. Compatible with GPT-4 Turbo and all
+            GPT-3.5 Turbo models newer than gpt-3.5-turbo-1106. Setting to
+            {"type": "json_object"} enables JSON mode, which guarantees the
+            message the model generates is valid JSON. Important: when using
+            JSON mode, you must also instruct the model to produce JSON
+            yourself via a system or user message. Without this, the model
+            may generate an unending stream of whitespace until the generation
+            reaches the token limit, resulting in a long-running and seemingly
+            "stuck" request. Also note that the message content may be
+            partially cut off if finish_reason="length", which indicates the
+            generation exceeded max_tokens or the conversation exceeded the
+            max context length.
+        stream (bool, optional): If True, partial message deltas will be sent
+            as data-only server-sent events as they become available.
+            (default: :obj:`None`)
+        stop (str or list, optional): Up to :obj:`4` sequences where the API
+            will stop generating further tokens. (default: :obj:`None`)
+        max_tokens (int, optional): The maximum number of tokens to generate
+            in the chat completion. The total length of input tokens and
+            generated tokens is limited by the model's context length.
+            (default: :obj:`None`)
+        presence_penalty (float, optional): Number between :obj:`-2.0` and
+            :obj:`2.0`. Positive values penalize new tokens based on whether
+            they appear in the text so far, increasing the model's likelihood
+            to talk about new topics. See more information about frequency and
+            presence penalties. (default: :obj:`None`)
+        frequency_penalty (float, optional): Number between :obj:`-2.0` and
+            :obj:`2.0`. Positive values penalize new tokens based on their
+            existing frequency in the text so far, decreasing the model's
+            likelihood to repeat the same line verbatim. See more information
+            about frequency and presence penalties. (default: :obj:`None`)
+        logit_bias (dict, optional): Modify the likelihood of specified tokens
+            appearing in the completion. Accepts a json object that maps tokens
+            (specified by their token ID in the tokenizer) to an associated
+            bias value from :obj:`-100` to :obj:`100`. Mathematically, the bias
+            is added to the logits generated by the model prior to sampling.
+            The exact effect will vary per model, but values between:obj:` -1`
+            and :obj:`1` should decrease or increase likelihood of selection;
+            values like :obj:`-100` or :obj:`100` should result in a ban or
+            exclusive selection of the relevant token. (default: :obj:`None`)
+        user (str, optional): A unique identifier representing your end-user,
+            which can help OpenAI to monitor and detect abuse.
+            (default: :obj:`None`)
+        tools (list[FunctionTool], optional): A list of tools the model may
+            call. Currently, only functions are supported as a tool. Use this
+            to provide a list of functions the model may generate JSON inputs
+            for. A max of 128 functions are supported.
+    """
+
+    temperature: Optional[float] = None
+    top_p: Optional[float] = None
+    n: Optional[int] = None
+    stream: Optional[bool] = None
+    stop: Optional[Union[str, Sequence[str]]] = None
+    max_tokens: Optional[int] = None
+    presence_penalty: Optional[float] = None
+    response_format: Optional[Union[Type[BaseModel], Dict]] = None
+    frequency_penalty: Optional[float] = None
+    logit_bias: Optional[Dict] = None
+    user: Optional[str] = None
+
+
+NOVITA_API_PARAMS = {param for param in NovitaConfig.model_fields.keys()}

camel/configs/qwen_config.py

@@ -80,12 +80,5 @@ class QwenConfig(BaseConfig):
     stop: Optional[Union[str, List]] = None
     extra_body: Optional[Dict[str, Any]] = None
 
-    def __init__(self, include_usage: bool = True, **kwargs):
-        super().__init__(**kwargs)
-        # Only set stream_options when stream is True
-        # Otherwise, it will raise error when calling the API
-        if self.stream:
-            self.stream_options = {"include_usage": include_usage}
-
 
 QWEN_API_PARAMS = {param for param in QwenConfig.model_fields.keys()}

camel/configs/watsonx_config.py

@@ -0,0 +1,96 @@
+# ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
+from __future__ import annotations
+
+from typing import Dict, List, Literal, Optional
+
+from camel.configs.base_config import BaseConfig
+
+
+class WatsonXConfig(BaseConfig):
+    r"""Defines the parameters for generating chat completions using the
+    IBM WatsonX API.
+
+    See: https://ibm.github.io/watsonx-ai-python-sdk/fm_schema.html
+
+    Args:
+        frequency_penalty (float, optional): Number between -2.0 and 2.0.
+            Positive values penalize new tokens based on their existing
+            frequency in the text so far, decreasing the model's likelihood
+            to repeat the same line verbatim.
+            (default: :obj:`None`)
+        logprobs (bool, optional): Whether to return log probabilities of the
+            output tokens or not. If true, returns the log probabilities of
+            each output token returned in the content of message.
+            (default: :obj:`None`)
+        top_logprobs (int, optional): An integer between 0 and 20 specifying
+            the number of most likely tokens to return at each token position,
+            each with an associated log probability.
+            (default: :obj:`None`)
+        presence_penalty (float, optional): Number between -2.0 and 2.0.
+            Positive values penalize new tokens based on whether they appear
+            in the text so far, increasing the model's likelihood to talk
+            about new topics.
+            (default: :obj:`None`)
+        temperature (float, optional): What sampling temperature to use,
+            between 0 and 2. Higher values like 0.8 will make the output
+            more random, while lower values like 0.2 will make it more focused
+            and deterministic. We generally recommend altering this or top_p
+            but not both.
+            (default: :obj:`None`)
+        max_tokens (int, optional): The maximum number of tokens to generate
+            in the chat completion. The total length of input tokens and
+            generated tokens is limited by the model's context length.
+            (default: :obj:`None`)
+        time_limit (int, optional): The maximum amount of time in seconds
+            that the API will spend generating a response.
+            (default: :obj:`None`)
+        top_p (float, optional): Controls the randomness of the generated
+            results. Lower values lead to less randomness, while higher
+            values increase randomness. (default: :obj:`None`)
+        n (int, optional): How many chat completion choices to generate for
+            each input message. Note that you will be charged based on the
+            total number of tokens generated.
+            (default: :obj:`None`)
+        logit_biaslogit_bias (Optional[dict], optional): Modify probability
+            of specific tokens appearing in the completion.
+            (default: :obj:`None`)
+        seed (int, optional): If specified, the system will make a best effort
+            to sample deterministically, such that repeated requests with the
+            same seed and parameters should return the same result.
+            (default: :obj:`None`)
+        stop (List[str], optional): Up to 4 sequences where the API will stop
+            generating further tokens.
+            (default: :obj:`None`)
+        tool_choice_options (Literal["none", "auto"], optional): The options
+            for the tool choice.
+            (default: :obj:`"auto"`)
+    """
+
+    frequency_penalty: Optional[float] = None
+    logprobs: Optional[bool] = None
+    top_logprobs: Optional[int] = None
+    presence_penalty: Optional[float] = None
+    temperature: Optional[float] = None
+    max_tokens: Optional[int] = None
+    time_limit: Optional[int] = None
+    top_p: Optional[float] = None
+    n: Optional[int] = None
+    logit_bias: Optional[Dict] = None
+    seed: Optional[int] = None
+    stop: Optional[List[str]] = None
+    tool_choice_options: Literal["none", "auto"] = "auto"
+
+
+WATSONX_API_PARAMS = {param for param in WatsonXConfig.model_fields.keys()}

camel/models/__init__.py

@@ -32,6 +32,7 @@ from .modelscope_model import ModelScopeModel
 from .moonshot_model import MoonshotModel
 from .nemotron_model import NemotronModel
 from .netmind_model import NetmindModel
+from .novita_model import NovitaModel
 from .nvidia_model import NvidiaModel
 from .ollama_model import OllamaModel
 from .openai_audio_models import OpenAIAudioModels
@@ -48,6 +49,7 @@ from .stub_model import StubModel
 from .togetherai_model import TogetherAIModel
 from .vllm_model import VLLMModel
 from .volcano_model import VolcanoModel
+from .watsonx_model import WatsonXModel
 from .yi_model import YiModel
 from .zhipuai_model import ZhipuAIModel
 
@@ -68,6 +70,7 @@ __all__ = [
     'OpenAIAudioModels',
     'NetmindModel',
     'NemotronModel',
+    'NovitaModel',
     'NvidiaModel',
     'OllamaModel',
     'VLLMModel',
@@ -92,4 +95,5 @@ __all__ = [
     'SiliconFlowModel',
     'VolcanoModel',
     'LMStudioModel',
+    'WatsonXModel',
 ]

camel/models/model_factory.py

@@ -32,6 +32,7 @@ from camel.models.mistral_model import MistralModel
 from camel.models.modelscope_model import ModelScopeModel
 from camel.models.moonshot_model import MoonshotModel
 from camel.models.netmind_model import NetmindModel
+from camel.models.novita_model import NovitaModel
 from camel.models.nvidia_model import NvidiaModel
 from camel.models.ollama_model import OllamaModel
 from camel.models.openai_compatible_model import OpenAICompatibleModel
@@ -47,6 +48,7 @@ from camel.models.stub_model import StubModel
 from camel.models.togetherai_model import TogetherAIModel
 from camel.models.vllm_model import VLLMModel
 from camel.models.volcano_model import VolcanoModel
+from camel.models.watsonx_model import WatsonXModel
 from camel.models.yi_model import YiModel
 from camel.models.zhipuai_model import ZhipuAIModel
 from camel.types import ModelPlatformType, ModelType, UnifiedModelType
@@ -62,8 +64,8 @@ class ModelFactory:
 
     @staticmethod
     def create(
-        model_platform: ModelPlatformType,
-        model_type: Union[ModelType, str],
+        model_platform: Union[ModelPlatformType, str],
+        model_type: Union[ModelType, str, UnifiedModelType],
         model_config_dict: Optional[Dict] = None,
         token_counter: Optional[BaseTokenCounter] = None,
         api_key: Optional[str] = None,
@@ -73,10 +75,12 @@ class ModelFactory:
         r"""Creates an instance of `BaseModelBackend` of the specified type.
 
         Args:
-            model_platform (ModelPlatformType): Platform from which the model
-                originates.
-            model_type (Union[ModelType, str]): Model for which a
-                backend is created. Can be a `str` for open source platforms.
+            model_platform (Union[ModelPlatformType, str]): Platform from
+                which the model originates. Can be a string or
+                ModelPlatformType enum.
+            model_type (Union[ModelType, str, UnifiedModelType]): Model for
+                which a backend is created. Can be a string, ModelType enum, or
+                UnifiedModelType.
             model_config_dict (Optional[Dict]): A dictionary that will be fed
                 into the backend constructor. (default: :obj:`None`)
             token_counter (Optional[BaseTokenCounter], optional): Token
@@ -97,6 +101,21 @@ class ModelFactory:
         Raises:
             ValueError: If there is no backend for the model.
         """
+        # Convert string to ModelPlatformType enum if needed
+        if isinstance(model_platform, str):
+            try:
+                model_platform = ModelPlatformType(model_platform)
+            except ValueError:
+                raise ValueError(f"Unknown model platform: {model_platform}")
+
+        # Convert string to ModelType enum or UnifiedModelType if needed
+        if isinstance(model_type, str):
+            try:
+                model_type = ModelType(model_type)
+            except ValueError:
+                # If not in ModelType, create a UnifiedModelType
+                model_type = UnifiedModelType(model_type)
+
         model_class: Optional[Type[BaseModelBackend]] = None
         model_type = UnifiedModelType(model_type)
 
@@ -163,8 +182,12 @@ class ModelFactory:
             model_class = MoonshotModel
         elif model_platform.is_modelscope:
             model_class = ModelScopeModel
+        elif model_platform.is_novita:
+            model_class = NovitaModel
         elif model_type == ModelType.STUB:
             model_class = StubModel
+        elif model_type.is_watsonx:
+            model_class = WatsonXModel
 
         if model_class is None:
             raise ValueError(