fix: Handle response metadata in merge_messages_runs (#28453)

This commit is contained in:
William FH 2024-12-02 13:56:23 -08:00 committed by GitHub
parent 60021e54b5
commit ecee41ab72
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 20 additions and 0 deletions

View File

@ -556,6 +556,8 @@ def merge_message_runs(
else: else:
last_chunk = _msg_to_chunk(last) last_chunk = _msg_to_chunk(last)
curr_chunk = _msg_to_chunk(curr) curr_chunk = _msg_to_chunk(curr)
if curr_chunk.response_metadata:
curr_chunk.response_metadata.clear()
if ( if (
isinstance(last_chunk.content, str) isinstance(last_chunk.content, str)
and isinstance(curr_chunk.content, str) and isinstance(curr_chunk.content, str)

View File

@ -59,6 +59,24 @@ def test_merge_message_runs_str_without_separator(
assert messages == messages_model_copy assert messages == messages_model_copy
def test_merge_message_runs_response_metadata() -> None:
messages = [
AIMessage("foo", id="1", response_metadata={"input_tokens": 1}),
AIMessage("bar", id="2", response_metadata={"input_tokens": 2}),
]
expected = [
AIMessage(
"foo\nbar",
id="1",
response_metadata={"input_tokens": 1},
)
]
actual = merge_message_runs(messages)
assert actual == expected
# Check it's not mutated
assert messages[1].response_metadata == {"input_tokens": 2}
def test_merge_message_runs_content() -> None: def test_merge_message_runs_content() -> None:
messages = [ messages = [
AIMessage("foo", id="1"), AIMessage("foo", id="1"),