Stop propagating 'id' field from Message object to request payload (#10)

* Stop propagating 'id' field from Message object to request payload Signed-off-by: B-Step62 <[email protected]> * test fix Signed-off-by: B-Step62 <[email protected]> * Update libs/databricks/tests/unit_tests/test_chat_models.py Co-authored-by: Harutaka Kawamura <[email protected]> --------- Signed-off-by: B-Step62 <[email protected]> Co-authored-by: Harutaka Kawamura <[email protected]>
langchain-ai · Sep 14, 2024 · deee25b · deee25b
1 parent 2064238
commit deee25b
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/libs/databricks/langchain_databricks/chat_models.py b/libs/databricks/langchain_databricks/chat_models.py
@@ -422,9 +422,6 @@ def _convert_message_to_dict(message: BaseMessage) -> dict:
     if (name := message.name or message.additional_kwargs.get("name")) is not None:
         message_dict["name"] = name
 
-    if id := message.id:
-        message_dict["id"] = id
-
     if isinstance(message, ChatMessage):
         return {"role": message.role, **message_dict}
     elif isinstance(message, HumanMessage):

diff --git a/libs/databricks/tests/unit_tests/test_chat_models.py b/libs/databricks/tests/unit_tests/test_chat_models.py
@@ -231,6 +231,15 @@ def test_convert_message(role: str, expected_output: BaseMessage) -> None:
     assert dict_result == message
 
 
+def test_convert_message_not_propagate_id() -> None:
+    # The AIMessage returned by the model endpoint can contain "id" field,
+    # but it is not always supported for requests. Therefore, we should not
+    # propagate it to the request payload.
+    message = AIMessage(content="foo", id="some-id")
+    result = _convert_message_to_dict(message)
+    assert "id" not in result
+
+
 def test_convert_message_with_tool_calls() -> None:
     ID = "call_fb5f5e1a-bac0-4422-95e9-d06e6022ad12"
     tool_calls = [
@@ -267,6 +276,7 @@ def test_convert_message_with_tool_calls() -> None:
 
     # convert back
     dict_result = _convert_message_to_dict(result)
+    message_with_tools.pop("id")  # id is not propagated
     assert dict_result == message_with_tools
 
 
@@ -320,6 +330,7 @@ def test_convert_tool_message_chunk() -> None:
 
     # convert back
     dict_result = _convert_message_to_dict(result)
+    delta.pop("id")  # id is not propagated
     assert dict_result == delta