Fix for chat stream tool choice (#51)

* partially revert changes from PR 50

* improve condition for setting name in tool_call_chunk

* fix mypy

Author: Mateusz-Switala

libs/ibm/langchain_ibm/chat_models.py

@@ -274,6 +274,7 @@ def _convert_delta_to_message_chunk(
     _dict: Mapping[str, Any],
     default_class: Type[BaseMessageChunk],
     call_id: str,
+    is_first_tool_chunk: bool,
 ) -> BaseMessageChunk:
     id_ = call_id
     role = cast(str, _dict.get("role"))
@@ -290,8 +291,12 @@ def _convert_delta_to_message_chunk(
         try:
             tool_call_chunks = [
                 tool_call_chunk(
-                    name=rtc["function"].get("name"),
+                    name=rtc["function"].get("name")
+                    if is_first_tool_chunk or (rtc.get("id") is not None)
+                    else None,
                     args=rtc["function"].get("arguments"),
+                    # `id` is provided only for the first delta with unique tool_calls
+                    # (multiple tool calls scenario)
                     id=rtc.get("id"),
                     index=rtc["index"],
                 )
@@ -328,6 +333,7 @@ def _convert_chunk_to_generation_chunk(
     default_chunk_class: Type,
     base_generation_info: Optional[Dict],
     is_first_chunk: bool,
+    is_first_tool_chunk: bool,
 ) -> Optional[ChatGenerationChunk]:
     token_usage = chunk.get("usage")
     choices = chunk.get("choices", [])
@@ -348,7 +354,7 @@ def _convert_chunk_to_generation_chunk(
         return None
 
     message_chunk = _convert_delta_to_message_chunk(
-        choice["delta"], default_chunk_class, chunk["id"]
+        choice["delta"], default_chunk_class, chunk["id"], is_first_tool_chunk
     )
     generation_info = {**base_generation_info} if base_generation_info else {}
 
@@ -722,6 +728,7 @@ def _stream(
         base_generation_info: dict = {}
 
         is_first_chunk = True
+        is_first_tool_chunk = True
 
         for chunk in self.watsonx_model.chat_stream(
             messages=message_dicts, **(kwargs | {"params": updated_params})
@@ -733,6 +740,7 @@ def _stream(
                 default_chunk_class,
                 base_generation_info if is_first_chunk else {},
                 is_first_chunk,
+                is_first_tool_chunk,
             )
             if generation_chunk is None:
                 continue
@@ -742,6 +750,16 @@ def _stream(
                 run_manager.on_llm_new_token(
                     generation_chunk.text, chunk=generation_chunk, logprobs=logprobs
                 )
+            if hasattr(generation_chunk.message, "tool_calls") and isinstance(
+                generation_chunk.message.tool_calls, list
+            ):
+                first_tool_call = (
+                    generation_chunk.message.tool_calls[0]
+                    if generation_chunk.message.tool_calls
+                    else None
+                )
+                if isinstance(first_tool_call, dict) and first_tool_call.get("name"):
+                    is_first_tool_chunk = False
 
             is_first_chunk = False
 

libs/ibm/tests/integration_tests/test_chat_models.py

@@ -667,7 +667,7 @@ def get_weather(city: Literal["nyc"]) -> str:
         if ai_message is None:
             ai_message = chunk
         else:
-            ai_message += chunk
+            ai_message += chunk  # type: ignore[assignment]
         print(chunk.id, type(chunk.id))
         assert isinstance(chunk, AIMessageChunk)
         assert chunk.content == ""

libs/ibm/tests/integration_tests/test_chat_models_standard.py

@@ -60,6 +60,7 @@ def chat_model_params(self) -> dict:
             "url": URL,
             "apikey": WX_APIKEY,
             "project_id": WX_PROJECT_ID,
+            "temperature": 0,
         }
 
     @pytest.mark.xfail(reason="Supported for vision model.")