From 9d101511233e7942794f8df8c742dfaa1ceed4c0 Mon Sep 17 00:00:00 2001
From: ccurme <chester.curme@gmail.com>
Date: Tue, 1 Oct 2024 14:21:54 -0400
Subject: [PATCH] core[patch]: fix init of RunnableAssign (#26903)

Example in API ref currently raises ValidationError.

Resolves https://github.com/langchain-ai/langchain/issues/26862
---
 libs/core/langchain_core/runnables/passthrough.py    |  2 +-
 .../core/tests/unit_tests/runnables/test_runnable.py | 12 ++++++++++++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/libs/core/langchain_core/runnables/passthrough.py b/libs/core/langchain_core/runnables/passthrough.py
index 3fc82ec9febb9..7e7fa7b45bdf0 100644
--- a/libs/core/langchain_core/runnables/passthrough.py
+++ b/libs/core/langchain_core/runnables/passthrough.py
@@ -389,7 +389,7 @@ def add_ten(x: Dict[str, int]) -> Dict[str, int]:
             # returns {'input': 5, 'add_step': {'added': 15}}
     """
 
-    mapper: RunnableParallel[dict[str, Any]]
+    mapper: RunnableParallel
 
     def __init__(self, mapper: RunnableParallel[dict[str, Any]], **kwargs: Any) -> None:
         super().__init__(mapper=mapper, **kwargs)  # type: ignore[call-arg]
diff --git a/libs/core/tests/unit_tests/runnables/test_runnable.py b/libs/core/tests/unit_tests/runnables/test_runnable.py
index 211fab46833db..f9dff49c0b99d 100644
--- a/libs/core/tests/unit_tests/runnables/test_runnable.py
+++ b/libs/core/tests/unit_tests/runnables/test_runnable.py
@@ -63,6 +63,7 @@
     ConfigurableFieldSingleOption,
     RouterRunnable,
     Runnable,
+    RunnableAssign,
     RunnableBinding,
     RunnableBranch,
     RunnableConfig,
@@ -5413,3 +5414,14 @@ def test_schema_for_prompt_and_chat_model() -> None:
         "title": "PromptInput",
         "type": "object",
     }
+
+
+def test_runnable_assign() -> None:
+    def add_ten(x: dict[str, int]) -> dict[str, int]:
+        return {"added": x["input"] + 10}
+
+    mapper = RunnableParallel({"add_step": RunnableLambda(add_ten)})
+    runnable_assign = RunnableAssign(mapper)
+
+    result = runnable_assign.invoke({"input": 5})
+    assert result == {"input": 5, "add_step": {"added": 15}}