Add validation for configurable keys passed to .with_config() (langchain-ai#11910)

- Fix some typing issues found while doing that

<!-- Thank you for contributing to LangChain!

Replace this entire comment with:
  - **Description:** a description of the change, 
  - **Issue:** the issue # it fixes (if applicable),
  - **Dependencies:** any dependencies required for this change,
- **Tag maintainer:** for a quicker response, tag the relevant
maintainer (see below),
- **Twitter handle:** we announce bigger features on Twitter. If your PR
gets announced, and you'd like a mention, we'll gladly shout you out!

Please make sure your PR is passing linting and testing before
submitting. Run `make format`, `make lint` and `make test` to check this
locally.

See contribution guidelines for more information on how to write/run
tests, lint, etc:

https://github.com/langchain-ai/langchain/blob/master/.github/CONTRIBUTING.md

If you're adding a new integration, please include:
1. a test for the integration, preferably unit tests that do not rely on
network access,
2. an example notebook showing its use. It lives in `docs/extras`
directory.

If no one reviews your PR within a few days, please @-mention one of
@baskaryan, @eyurtsev, @hwchase17.
 -->

Author: nfcampos

libs/langchain/langchain/runnables/openai_functions.py

@@ -4,7 +4,7 @@
 from typing_extensions import TypedDict
 
 from langchain.output_parsers.openai_functions import JsonOutputFunctionsParser
-from langchain.schema.output import ChatGeneration
+from langchain.schema.messages import BaseMessage
 from langchain.schema.runnable import RouterRunnable, Runnable, RunnableBinding
 
 
@@ -19,7 +19,7 @@ class OpenAIFunction(TypedDict):
     """The parameters to the function."""
 
 
-class OpenAIFunctionsRouter(RunnableBinding[ChatGeneration, Any]):
+class OpenAIFunctionsRouter(RunnableBinding[BaseMessage, Any]):
     """A runnable that routes to the selected function."""
 
     functions: Optional[List[OpenAIFunction]]

libs/langchain/langchain/schema/runnable/_locals.py

@@ -11,7 +11,7 @@
     Union,
 )
 
-from langchain.schema.runnable.base import Input, Output, RunnableSerializable
+from langchain.schema.runnable.base import Input, Other, Output, RunnableSerializable
 from langchain.schema.runnable.config import RunnableConfig
 from langchain.schema.runnable.passthrough import RunnablePassthrough
 
@@ -36,7 +36,7 @@ def __init__(self, key: Union[str, Mapping[str, str]], **kwargs: Any) -> None:
 
     def _concat_put(
         self,
-        input: Input,
+        input: Other,
         *,
         config: Optional[RunnableConfig] = None,
         replace: bool = False,
@@ -68,35 +68,35 @@ def _concat_put(
                 f"{(type(self.key))}."
             )
 
-    def invoke(self, input: Input, config: Optional[RunnableConfig] = None) -> Input:
+    def invoke(self, input: Other, config: Optional[RunnableConfig] = None) -> Other:
         self._concat_put(input, config=config, replace=True)
         return super().invoke(input, config=config)
 
     async def ainvoke(
         self,
-        input: Input,
+        input: Other,
         config: Optional[RunnableConfig] = None,
         **kwargs: Optional[Any],
-    ) -> Input:
+    ) -> Other:
         self._concat_put(input, config=config, replace=True)
         return await super().ainvoke(input, config=config)
 
     def transform(
         self,
-        input: Iterator[Input],
+        input: Iterator[Other],
         config: Optional[RunnableConfig] = None,
         **kwargs: Optional[Any],
-    ) -> Iterator[Input]:
+    ) -> Iterator[Other]:
         for chunk in super().transform(input, config=config):
             self._concat_put(chunk, config=config)
             yield chunk
 
     async def atransform(
         self,
-        input: AsyncIterator[Input],
+        input: AsyncIterator[Other],
         config: Optional[RunnableConfig] = None,
         **kwargs: Optional[Any],
-    ) -> AsyncIterator[Input]:
+    ) -> AsyncIterator[Other]:
         async for chunk in super().atransform(input, config=config):
             self._concat_put(chunk, config=config)
             yield chunk

libs/langchain/langchain/schema/runnable/base.py

@@ -2296,6 +2296,26 @@ class RunnableBinding(RunnableSerializable[Input, Output]):
     class Config:
         arbitrary_types_allowed = True
 
+    def __init__(
+        self,
+        *,
+        bound: Runnable[Input, Output],
+        kwargs: Mapping[str, Any],
+        config: Optional[Mapping[str, Any]] = None,
+        **other_kwargs: Any,
+    ) -> None:
+        config = config or {}
+        # config_specs contains the list of valid `configurable` keys
+        if configurable := config.get("configurable", None):
+            allowed_keys = set(s.id for s in bound.config_specs)
+            for key in configurable:
+                if key not in allowed_keys:
+                    raise ValueError(
+                        f"Configurable key '{key}' not found in runnable with"
+                        f" config keys: {allowed_keys}"
+                    )
+        super().__init__(bound=bound, kwargs=kwargs, config=config, **other_kwargs)
+
     @property
     def InputType(self) -> Type[Input]:
         return self.bound.InputType

libs/langchain/langchain/schema/runnable/passthrough.py

@@ -22,7 +22,7 @@
 
 from langchain.pydantic_v1 import BaseModel, create_model
 from langchain.schema.runnable.base import (
-    Input,
+    Other,
     Runnable,
     RunnableParallel,
     RunnableSerializable,
@@ -33,17 +33,17 @@
 from langchain.utils.iter import safetee
 
 
-def identity(x: Input) -> Input:
+def identity(x: Other) -> Other:
     """An identity function"""
     return x
 
 
-async def aidentity(x: Input) -> Input:
+async def aidentity(x: Other) -> Other:
     """An async identity function"""
     return x
 
 
-class RunnablePassthrough(RunnableSerializable[Input, Input]):
+class RunnablePassthrough(RunnableSerializable[Other, Other]):
     """A runnable to passthrough inputs unchanged or with additional keys.
 
     This runnable behaves almost like the identity function, except that it
@@ -100,20 +100,20 @@ def fake_llm(prompt: str) -> str: # Fake LLM for the example
             # {'llm1': 'completion', 'llm2': 'completion', 'total_chars': 20}
     """
 
-    input_type: Optional[Type[Input]] = None
+    input_type: Optional[Type[Other]] = None
 
-    func: Optional[Callable[[Input], None]] = None
+    func: Optional[Callable[[Other], None]] = None
 
-    afunc: Optional[Callable[[Input], Awaitable[None]]] = None
+    afunc: Optional[Callable[[Other], Awaitable[None]]] = None
 
     def __init__(
         self,
         func: Optional[
-            Union[Callable[[Input], None], Callable[[Input], Awaitable[None]]]
+            Union[Callable[[Other], None], Callable[[Other], Awaitable[None]]]
         ] = None,
-        afunc: Optional[Callable[[Input], Awaitable[None]]] = None,
+        afunc: Optional[Callable[[Other], Awaitable[None]]] = None,
         *,
-        input_type: Optional[Type[Input]] = None,
+        input_type: Optional[Type[Other]] = None,
         **kwargs: Any,
     ) -> None:
         if inspect.iscoroutinefunction(func):
@@ -161,17 +161,17 @@ def assign(
         """
         return RunnableAssign(RunnableParallel(kwargs))
 
-    def invoke(self, input: Input, config: Optional[RunnableConfig] = None) -> Input:
+    def invoke(self, input: Other, config: Optional[RunnableConfig] = None) -> Other:
         if self.func is not None:
             self.func(input)
         return self._call_with_config(identity, input, config)
 
     async def ainvoke(
         self,
-        input: Input,
+        input: Other,
         config: Optional[RunnableConfig] = None,
         **kwargs: Optional[Any],
-    ) -> Input:
+    ) -> Other:
         if self.afunc is not None:
             await self.afunc(input, **kwargs)
         elif self.func is not None:
@@ -180,10 +180,10 @@ async def ainvoke(
 
     def transform(
         self,
-        input: Iterator[Input],
+        input: Iterator[Other],
         config: Optional[RunnableConfig] = None,
         **kwargs: Any,
-    ) -> Iterator[Input]:
+    ) -> Iterator[Other]:
         if self.func is None:
             for chunk in self._transform_stream_with_config(input, identity, config):
                 yield chunk
@@ -202,10 +202,10 @@ def transform(
 
     async def atransform(
         self,
-        input: AsyncIterator[Input],
+        input: AsyncIterator[Other],
         config: Optional[RunnableConfig] = None,
         **kwargs: Any,
-    ) -> AsyncIterator[Input]:
+    ) -> AsyncIterator[Other]:
         if self.afunc is None and self.func is None:
             async for chunk in self._atransform_stream_with_config(
                 input, identity, config
@@ -231,19 +231,19 @@ async def atransform(
 
     def stream(
         self,
-        input: Input,
+        input: Other,
         config: Optional[RunnableConfig] = None,
         **kwargs: Any,
-    ) -> Iterator[Input]:
+    ) -> Iterator[Other]:
         return self.transform(iter([input]), config, **kwargs)
 
     async def astream(
         self,
-        input: Input,
+        input: Other,
         config: Optional[RunnableConfig] = None,
         **kwargs: Any,
-    ) -> AsyncIterator[Input]:
-        async def input_aiter() -> AsyncIterator[Input]:
+    ) -> AsyncIterator[Other]:
+        async def input_aiter() -> AsyncIterator[Other]:
             yield input
 
         async for chunk in self.atransform(input_aiter(), config, **kwargs):

libs/langchain/langchain/schema/runnable/utils.py

@@ -24,9 +24,9 @@
     Union,
 )
 
-Input = TypeVar("Input")
+Input = TypeVar("Input", contravariant=True)
 # Output type should implement __concat__, as eg str, list, dict do
-Output = TypeVar("Output")
+Output = TypeVar("Output", covariant=True)
 
 
 async def gated_coro(semaphore: asyncio.Semaphore, coro: Coroutine) -> Any:

libs/langchain/tests/unit_tests/schema/runnable/test_runnable.py

@@ -1002,6 +1002,9 @@ def test_configurable_fields_example() -> None:
         },
     }
 
+    with pytest.raises(ValueError):
+        chain_configurable.with_config(configurable={"llm123": "chat"})
+
     assert (
         chain_configurable.with_config(configurable={"llm": "chat"}).invoke(
             {"name": "John"}