mirror of https://github.com/hwchase17/langchain
core[patch]: include tool_calls in ai msg chunk serialization (#20291)
parent
0fa551c278
commit
03b247cca1
@ -0,0 +1,67 @@
|
|||||||
|
from langchain_core.load import dumpd, load
|
||||||
|
from langchain_core.messages import (
|
||||||
|
AIMessage,
|
||||||
|
AIMessageChunk,
|
||||||
|
InvalidToolCall,
|
||||||
|
ToolCall,
|
||||||
|
ToolCallChunk,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_serdes_message() -> None:
|
||||||
|
msg = AIMessage(
|
||||||
|
content=[{"text": "blah", "type": "text"}],
|
||||||
|
tool_calls=[ToolCall(name="foo", args={"bar": 1}, id="baz")],
|
||||||
|
invalid_tool_calls=[
|
||||||
|
InvalidToolCall(name="foobad", args="blah", id="booz", error="bad")
|
||||||
|
],
|
||||||
|
)
|
||||||
|
expected = {
|
||||||
|
"lc": 1,
|
||||||
|
"type": "constructor",
|
||||||
|
"id": ["langchain", "schema", "messages", "AIMessage"],
|
||||||
|
"kwargs": {
|
||||||
|
"content": [{"text": "blah", "type": "text"}],
|
||||||
|
"tool_calls": [{"name": "foo", "args": {"bar": 1}, "id": "baz"}],
|
||||||
|
"invalid_tool_calls": [
|
||||||
|
{"name": "foobad", "args": "blah", "id": "booz", "error": "bad"}
|
||||||
|
],
|
||||||
|
},
|
||||||
|
}
|
||||||
|
actual = dumpd(msg)
|
||||||
|
assert actual == expected
|
||||||
|
assert load(actual) == msg
|
||||||
|
|
||||||
|
|
||||||
|
def test_serdes_message_chunk() -> None:
|
||||||
|
chunk = AIMessageChunk(
|
||||||
|
content=[{"text": "blah", "type": "text"}],
|
||||||
|
tool_call_chunks=[
|
||||||
|
ToolCallChunk(name="foo", args='{"bar": 1}', id="baz", index=0),
|
||||||
|
ToolCallChunk(name="foobad", args="blah", id="booz", index=1),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
expected = {
|
||||||
|
"lc": 1,
|
||||||
|
"type": "constructor",
|
||||||
|
"id": ["langchain", "schema", "messages", "AIMessageChunk"],
|
||||||
|
"kwargs": {
|
||||||
|
"content": [{"text": "blah", "type": "text"}],
|
||||||
|
"tool_calls": [{"name": "foo", "args": {"bar": 1}, "id": "baz"}],
|
||||||
|
"invalid_tool_calls": [
|
||||||
|
{
|
||||||
|
"name": "foobad",
|
||||||
|
"args": "blah",
|
||||||
|
"id": "booz",
|
||||||
|
"error": "Malformed args.",
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"tool_call_chunks": [
|
||||||
|
{"name": "foo", "args": '{"bar": 1}', "id": "baz", "index": 0},
|
||||||
|
{"name": "foobad", "args": "blah", "id": "booz", "index": 1},
|
||||||
|
],
|
||||||
|
},
|
||||||
|
}
|
||||||
|
actual = dumpd(chunk)
|
||||||
|
assert actual == expected
|
||||||
|
assert load(actual) == chunk
|
File diff suppressed because one or more lines are too long
Loading…
Reference in New Issue