From fb2a134d7cfa6ed96d883667d2105ebf9fd07a58 Mon Sep 17 00:00:00 2001 From: William Zhu Date: Mon, 18 Nov 2024 17:40:25 -0500 Subject: [PATCH] run make format/lint --- .../tests/unit_tests/runnables/test_runnable.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/libs/core/tests/unit_tests/runnables/test_runnable.py b/libs/core/tests/unit_tests/runnables/test_runnable.py index 2f49c71a89c..dbd8f748241 100644 --- a/libs/core/tests/unit_tests/runnables/test_runnable.py +++ b/libs/core/tests/unit_tests/runnables/test_runnable.py @@ -5466,13 +5466,15 @@ def test_runnable_assign() -> None: result = runnable_assign.invoke({"input": 5}) assert result == {"input": 5, "add_step": {"added": 15}} + def test_runnable_typed_dict_schema() -> None: """Testing that the schema is generated properly(not empty) when using TypedDict subclasses to annotate the arguments of a RunnableParallel children. """ from typing_extensions import TypedDict - from langchain_core.runnables import RunnableParallel, RunnableLambda + + from langchain_core.runnables import RunnableLambda, RunnableParallel class Foo(TypedDict): foo: str @@ -5480,16 +5482,14 @@ def test_runnable_typed_dict_schema() -> None: class InputData(Foo): bar: str - def forward_foo(input_data: InputData): + def forward_foo(input_data: InputData) -> str: return input_data["foo"] - def transform_input(input_data: InputData): + def transform_input(input_data: InputData) -> dict[str, str]: foo = input_data["foo"] bar = input_data["bar"] - return { - "transformed": foo + bar - } + return {"transformed": foo + bar} foo_runnable = RunnableLambda(forward_foo) other_runnable = RunnableLambda(transform_input) @@ -5498,4 +5498,7 @@ def test_runnable_typed_dict_schema() -> None: foo=foo_runnable, other=other_runnable, ) - assert(repr(parallel.input_schema.validate({ "foo": "Y", "bar": "Z" })) == "RunnableParallelInput(root={'foo': 'Y', 'bar': 'Z'})") + assert ( + repr(parallel.input_schema.validate({"foo": "Y", "bar": "Z"})) + == "RunnableParallelInput(root={'foo': 'Y', 'bar': 'Z'})" + )