fix docstring and clean up (#53)

* fix docstring and clean up

* fix linting

* fix linter

Author: Mateusz-Switala

libs/ibm/langchain_ibm/chat_models.py

@@ -401,19 +401,17 @@ class ChatWatsonx(BaseChatModel):
     Example:
         .. code-block:: python
 
-            from ibm_watsonx_ai.metanames import GenTextParamsMetaNames
-            parameters = {
-                GenTextParamsMetaNames.DECODING_METHOD: "sample",
-                GenTextParamsMetaNames.MAX_NEW_TOKENS: 100,
-                GenTextParamsMetaNames.MIN_NEW_TOKENS: 1,
-                GenTextParamsMetaNames.TEMPERATURE: 0.5,
-                GenTextParamsMetaNames.TOP_K: 50,
-                GenTextParamsMetaNames.TOP_P: 1,
-            }
+            from ibm_watsonx_ai.foundation_models.schema import TextChatParameters
+
+            parameters = TextChatParameters(
+                max_tokens=100,
+                temperature=0.5,
+                top_p=1,
+                )
 
             from langchain_ibm import ChatWatsonx
             watsonx_llm = ChatWatsonx(
-                model_id="meta-llama/llama-3-70b-instruct",
+                model_id="meta-llama/llama-3-3-70b-instruct",
                 url="https://us-south.ml.cloud.ibm.com",
                 apikey="*****",
                 project_id="*****",
@@ -527,6 +525,18 @@ class ChatWatsonx(BaseChatModel):
     """Time limit in milliseconds - if not completed within this time, 
     generation will stop."""
 
+    logit_bias: Optional[dict] = None
+    """Increasing or decreasing probability of tokens being selected 
+    during generation."""
+
+    seed: Optional[int] = None
+    """Random number generator seed to use in sampling mode 
+    for experimental repeatability."""
+
+    stop: Optional[list[str]] = None
+    """Stop sequences are one or more strings which will cause the text generation 
+    to stop if/when they are produced as part of the output."""
+
     verify: Union[str, bool, None] = None
     """You can pass one of following as verify:
         * the path to a CA_BUNDLE file
@@ -602,16 +612,8 @@ def validate_environment(self) -> Self:
             {
                 k: v
                 for k, v in {
-                    "frequency_penalty": self.frequency_penalty,
-                    "logprobs": self.logprobs,
-                    "top_logprobs": self.top_logprobs,
-                    "max_tokens": self.max_tokens,
-                    "n": self.n,
-                    "presence_penalty": self.presence_penalty,
-                    "response_format": self.response_format,
-                    "temperature": self.temperature,
-                    "top_p": self.top_p,
-                    "time_limit": self.time_limit,
+                    param: getattr(self, param)
+                    for param in ChatWatsonx._get_supported_chat_params()
                 }.items()
                 if v is not None
             }
@@ -768,18 +770,7 @@ def _stream(
     @staticmethod
     def _merge_params(params: dict, kwargs: dict) -> dict:
         param_updates = {}
-        for k in [
-            "frequency_penalty",
-            "logprobs",
-            "top_logprobs",
-            "max_tokens",
-            "n",
-            "presence_penalty",
-            "response_format",
-            "temperature",
-            "top_p",
-            "time_limit",
-        ]:
+        for k in ChatWatsonx._get_supported_chat_params():
             if kwargs.get(k) is not None:
                 param_updates[k] = kwargs.pop(k)
 
@@ -837,6 +828,25 @@ def _create_chat_result(
 
         return ChatResult(generations=generations, llm_output=llm_output)
 
+    @staticmethod
+    def _get_supported_chat_params() -> list[str]:
+        # watsonx.ai Chat API doc: https://cloud.ibm.com/apidocs/watsonx-ai#text-chat
+        return [
+            "frequency_penalty",
+            "logprobs",
+            "top_logprobs",
+            "max_tokens",
+            "n",
+            "presence_penalty",
+            "response_format",
+            "temperature",
+            "top_p",
+            "time_limit",
+            "logit_bias",
+            "seed",
+            "stop",
+        ]
+
     def bind_functions(
         self,
         functions: Sequence[Union[Dict[str, Any], Type[BaseModel], Callable, BaseTool]],

libs/ibm/tests/unit_tests/test_chat_models.py

@@ -1,6 +1,7 @@
 """Test ChatWatsonx API wrapper."""
 
 import os
+from typing import Any
 
 from langchain_ibm import ChatWatsonx
 
@@ -82,3 +83,55 @@ def test_initialize_chat_watsonx_cpd_bad_path_without_instance_id() -> None:
     except ValueError as e:
         assert "instance_id" in e.__str__()
         assert "WATSONX_INSTANCE_ID" in e.__str__()
+
+
+def test_initialize_chat_watsonx_with_all_supported_params(mocker: Any) -> None:
+    # All params values are taken from
+    # ibm_watsonx_ai.foundation_models.schema.TextChatParameters.get_sample_params()
+
+    from ibm_watsonx_ai.foundation_models.schema import (  # type: ignore[import-untyped]
+        TextChatParameters,
+    )
+
+    TOP_P = 0.8
+
+    def mock_modelinference_chat(*args: Any, **kwargs: Any) -> dict:
+        """Mock ModelInference.chat method"""
+
+        assert kwargs.get("params", None) == (
+            TextChatParameters.get_sample_params()
+            | dict(
+                logit_bias={"1003": -100, "1004": -100}, seed=41, stop=["this", "the"]
+            )
+            | dict(top_p=TOP_P)
+        )
+        # logit_bias, seed and stop available in sdk since 1.2.7
+        return {"id": "123", "choices": [{"message": dict(content="Hi", role="ai")}]}
+
+    with mocker.patch(
+        "ibm_watsonx_ai.foundation_models.ModelInference.__init__", return_value=None
+    ), mocker.patch(
+        "ibm_watsonx_ai.foundation_models.ModelInference.chat",
+        side_effect=mock_modelinference_chat,
+    ):
+        chat = ChatWatsonx(
+            model_id="google/flan-ul2",
+            url="https://us-south.ml.cloud.ibm.com",  # type: ignore[arg-type]
+            apikey="test_apikey",  # type: ignore[arg-type]
+            frequency_penalty=0.5,
+            logprobs=True,
+            top_logprobs=3,
+            presence_penalty=0.3,
+            response_format={"type": "json_object"},
+            temperature=0.7,
+            max_tokens=100,
+            time_limit=600000,
+            top_p=0.9,
+            n=1,
+            logit_bias={"1003": -100, "1004": -100},
+            seed=41,
+            stop=["this", "the"],
+        )
+
+        # change only top_n
+        chat.invoke("Hello", top_p=TOP_P)