Merge branch 'main' into erick/assign-input-params

tests/unit_tests/test_server_client.py

@@ -16,7 +16,7 @@ from langchain.prompts import PromptTemplate
 from langchain.prompts.base import StringPromptValue
 from langchain.schema.messages import HumanMessage, SystemMessage
 from langchain.schema.runnable import RunnableConfig, RunnablePassthrough
-from langchain.schema.runnable.base import RunnableLambda
+from langchain.schema.runnable.base import RunnableLambda, Runnable
 from langchain.schema.runnable.utils import ConfigurableField
 from pytest_mock import MockerFixture
 from typing_extensions import TypedDict
@@ -750,6 +750,33 @@ def test_rename_pydantic_model() -> None:
 
 
 @pytest.mark.asyncio
+async def test_runnable_assign(event_loop: AbstractEventLoop) -> None:
+    """Test serving multiple runnables."""
+
+    app = FastAPI()
+    chain = RunnablePassthrough().assign(a=RunnableLambda(lambda _: "a"))
+    # should only need "b" as input now
+    add_routes(app, chain, path="/assigned")
+
+    async with get_async_client(app, path="/assigned") as runnable:
+        assert await runnable.ainvoke({"b": "b"}) == {"a": "a", "b": "b"}
+        assert runnable.invoke({"b": "b"}) == {"a": "a", "b": "b"}
+
+
+@pytest.mark.asyncio
+async def test_runnable_assign_nopath(event_loop: AbstractEventLoop) -> None:
+    """Test serving multiple runnables."""
+
+    app = FastAPI()
+    chain = RunnablePassthrough().assign(a=RunnableLambda(lambda _: "a"))
+    # should only need "b" as input now
+    add_routes(app, chain)
+
+    async with get_async_client(app, path="/") as runnable:
+        assert await runnable.ainvoke({"b": "b"}) == {"a": "a", "b": "b"}
+        assert runnable.invoke({"b": "b"}) == {"a": "a", "b": "b"}
+
+
 async def test_input_schema_typed_dict() -> None:
     class InputType(TypedDict):
         foo: str