mirror of
https://github.com/hwchase17/langchain
synced 2024-11-11 19:11:02 +00:00
c2a3021bb0
Signed-off-by: ChengZi <chen.zhang@zilliz.com> Co-authored-by: Eugene Yurtsev <eyurtsev@gmail.com> Co-authored-by: Bagatur <22008038+baskaryan@users.noreply.github.com> Co-authored-by: Dan O'Donovan <dan.odonovan@gmail.com> Co-authored-by: Tom Daniel Grande <tomdgrande@gmail.com> Co-authored-by: Grande <Tom.Daniel.Grande@statsbygg.no> Co-authored-by: Bagatur <baskaryan@gmail.com> Co-authored-by: ccurme <chester.curme@gmail.com> Co-authored-by: Harrison Chase <hw.chase.17@gmail.com> Co-authored-by: Tomaz Bratanic <bratanic.tomaz@gmail.com> Co-authored-by: ZhangShenao <15201440436@163.com> Co-authored-by: Friso H. Kingma <fhkingma@gmail.com> Co-authored-by: ChengZi <chen.zhang@zilliz.com> Co-authored-by: Nuno Campos <nuno@langchain.dev> Co-authored-by: Morgante Pell <morgantep@google.com>
114 lines
2.9 KiB
Python
114 lines
2.9 KiB
Python
from typing import Any, List, Literal
|
|
|
|
from langchain_core.messages import AIMessage
|
|
from langchain_core.outputs import ChatGeneration
|
|
from pydantic import BaseModel
|
|
|
|
from langchain_anthropic.output_parsers import ToolsOutputParser
|
|
|
|
_CONTENT: List = [
|
|
{
|
|
"type": "text",
|
|
"text": "thought",
|
|
},
|
|
{"type": "tool_use", "input": {"bar": 0}, "id": "1", "name": "_Foo1"},
|
|
{
|
|
"type": "text",
|
|
"text": "thought",
|
|
},
|
|
{"type": "tool_use", "input": {"baz": "a"}, "id": "2", "name": "_Foo2"},
|
|
]
|
|
|
|
_RESULT: List = [ChatGeneration(message=AIMessage(_CONTENT))] # type: ignore[misc]
|
|
|
|
|
|
class _Foo1(BaseModel):
|
|
bar: int
|
|
|
|
|
|
class _Foo2(BaseModel):
|
|
baz: Literal["a", "b"]
|
|
|
|
|
|
def test_tools_output_parser() -> None:
|
|
output_parser = ToolsOutputParser()
|
|
expected = [
|
|
{
|
|
"name": "_Foo1",
|
|
"args": {"bar": 0},
|
|
"id": "1",
|
|
"index": 1,
|
|
"type": "tool_call",
|
|
},
|
|
{
|
|
"name": "_Foo2",
|
|
"args": {"baz": "a"},
|
|
"id": "2",
|
|
"index": 3,
|
|
"type": "tool_call",
|
|
},
|
|
]
|
|
actual = output_parser.parse_result(_RESULT)
|
|
assert expected == actual
|
|
|
|
|
|
def test_tools_output_parser_args_only() -> None:
|
|
output_parser = ToolsOutputParser(args_only=True)
|
|
expected = [
|
|
{"bar": 0},
|
|
{"baz": "a"},
|
|
]
|
|
actual = output_parser.parse_result(_RESULT)
|
|
assert expected == actual
|
|
|
|
expected = []
|
|
actual = output_parser.parse_result([ChatGeneration(message=AIMessage(""))]) # type: ignore[misc]
|
|
assert expected == actual
|
|
|
|
|
|
def test_tools_output_parser_first_tool_only() -> None:
|
|
output_parser = ToolsOutputParser(first_tool_only=True)
|
|
expected: Any = {
|
|
"name": "_Foo1",
|
|
"args": {"bar": 0},
|
|
"id": "1",
|
|
"index": 1,
|
|
"type": "tool_call",
|
|
}
|
|
actual = output_parser.parse_result(_RESULT)
|
|
assert expected == actual
|
|
|
|
expected = None
|
|
actual = output_parser.parse_result([ChatGeneration(message=AIMessage(""))]) # type: ignore[misc]
|
|
assert expected == actual
|
|
|
|
|
|
def test_tools_output_parser_pydantic() -> None:
|
|
output_parser = ToolsOutputParser(pydantic_schemas=[_Foo1, _Foo2])
|
|
expected = [_Foo1(bar=0), _Foo2(baz="a")]
|
|
actual = output_parser.parse_result(_RESULT)
|
|
assert expected == actual
|
|
|
|
|
|
def test_tools_output_parser_empty_content() -> None:
|
|
class ChartType(BaseModel):
|
|
chart_type: Literal["pie", "line", "bar"]
|
|
|
|
output_parser = ToolsOutputParser(
|
|
first_tool_only=True, pydantic_schemas=[ChartType]
|
|
)
|
|
message = AIMessage(
|
|
"",
|
|
tool_calls=[
|
|
{
|
|
"name": "ChartType",
|
|
"args": {"chart_type": "pie"},
|
|
"id": "foo",
|
|
"type": "tool_call",
|
|
}
|
|
],
|
|
)
|
|
actual = output_parser.invoke(message)
|
|
expected = ChartType(chart_type="pie")
|
|
assert expected == actual
|