From b199ddc54c7c04fcbc5e06cf7e6db5cf0cd17896 Mon Sep 17 00:00:00 2001
From: Bagatur <baskaryan@gmail.com>
Date: Thu, 8 Feb 2024 17:24:40 -0800
Subject: [PATCH] core[patch]: chain decorator typing

---
 libs/core/langchain_core/runnables/base.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/libs/core/langchain_core/runnables/base.py b/libs/core/langchain_core/runnables/base.py
index 42f5f33c31c..82488cf9e50 100644
--- a/libs/core/langchain_core/runnables/base.py
+++ b/libs/core/langchain_core/runnables/base.py
@@ -23,6 +23,7 @@ from typing import (
     List,
     Mapping,
     Optional,
+    Protocol,
     Sequence,
     Set,
     Tuple,
@@ -4373,6 +4374,11 @@ def coerce_to_runnable(thing: RunnableLike) -> Runnable[Input, Output]:
         )
 
 
+class _CallableWithKwargs(Protocol, Generic[Input, Output]):
+    def __call__(self, input: Input, **kwargs: Any) -> Output:
+        ...
+
+
 @overload
 def chain(
     func: Callable[[Input], Coroutine[Any, Any, Output]],
@@ -4407,6 +4413,10 @@ def chain(
         Callable[[Input], Iterator[Output]],
         Callable[[Input], Coroutine[Any, Any, Output]],
         Callable[[Input], AsyncIterator[Output]],
+        _CallableWithKwargs[Input, Output],
+        _CallableWithKwargs[Input, Iterator[Output]],
+        _CallableWithKwargs[Input, Coroutine[Any, Any, Output]],
+        _CallableWithKwargs[Input, AsyncIterator[Output]],
     ],
 ) -> Runnable[Input, Output]:
     """Decorate a function to make it a Runnable.