core[patch]: fix ChatGeneration.text with content blocks (#20294)

This commit is contained in:
Bagatur 2024-04-10 17:54:06 -05:00 committed by GitHub
parent 03b247cca1
commit cb25fa0d55
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 50 additions and 1 deletions

View File

@ -23,7 +23,24 @@ class ChatGeneration(Generation):
def set_text(cls, values: Dict[str, Any]) -> Dict[str, Any]: def set_text(cls, values: Dict[str, Any]) -> Dict[str, Any]:
"""Set the text attribute to be the contents of the message.""" """Set the text attribute to be the contents of the message."""
try: try:
values["text"] = values["message"].content text = ""
if isinstance(values["message"].content, str):
text = values["message"].content
# HACK: Assumes text in content blocks in OpenAI format.
# Uses first text block.
elif isinstance(values["message"].content, list):
for block in values["message"].content:
if isinstance(block, str):
text = block
break
elif isinstance(block, dict) and "text" in block:
text = block["text"]
break
else:
pass
else:
pass
values["text"] = text
except (KeyError, AttributeError) as e: except (KeyError, AttributeError) as e:
raise ValueError("Error while initializing ChatGeneration") from e raise ValueError("Error while initializing ChatGeneration") from e
return values return values

View File

@ -0,0 +1,32 @@
from typing import Union
import pytest
from langchain_core.messages import AIMessage
from langchain_core.outputs import ChatGeneration
@pytest.mark.parametrize(
"content",
[
"foo",
["foo"],
[{"text": "foo", "type": "text"}],
[
{"tool_use": {}, "type": "tool_use"},
{"text": "foo", "type": "text"},
"bar",
],
],
)
def test_msg_with_text(content: Union[str, list]) -> None:
expected = "foo"
actual = ChatGeneration(message=AIMessage(content=content)).text
assert actual == expected
@pytest.mark.parametrize("content", [[], [{"tool_use": {}, "type": "tool_use"}]])
def test_msg_no_text(content: Union[str, list]) -> None:
expected = ""
actual = ChatGeneration(message=AIMessage(content=content)).text
assert actual == expected