From ba672604a6fd3307ecb2b21c49e065a95ae4d7dd Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Thu, 16 Jan 2025 07:57:00 -0800
Subject: [PATCH] Implement input/output schemas for imperative api

---
 libs/langgraph/langgraph/func/__init__.py | 23 ++++++++++++++++++++++-
 libs/langgraph/tests/test_pregel.py       | 11 +++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py
index 4749c075f..0e3cd5793 100644
--- a/libs/langgraph/langgraph/func/__init__.py
+++ b/libs/langgraph/langgraph/func/__init__.py
@@ -112,6 +112,7 @@ def entrypoint(
     store: Optional[BaseStore] = None,
 ) -> Callable[[types.FunctionType], Pregel]:
     def _imp(func: types.FunctionType) -> Pregel:
+        # wrap generators in a function that writes to StreamWriter
         if inspect.isgeneratorfunction(func):
 
             def gen_wrapper(*args: Any, writer: StreamWriter, **kwargs: Any) -> Any:
@@ -134,6 +135,23 @@ def entrypoint(
             bound = get_runnable_for_func(func)
             stream_mode = "updates"
 
+        # get input and output types
+        sig = inspect.signature(func)
+        first_parameter_name = next(iter(sig.parameters.keys()), None)
+        if not first_parameter_name:
+            raise ValueError("Entrypoint function must have at least one parameter")
+        input_type = (
+            sig.parameters[first_parameter_name].annotation
+            if sig.parameters[first_parameter_name].annotation
+            is not inspect.Signature.empty
+            else Any
+        )
+        output_type = (
+            sig.return_annotation
+            if sig.return_annotation is not inspect.Signature.empty
+            else Any
+        )
+
         return Pregel(
             nodes={
                 func.__name__: PregelNode(
@@ -143,7 +161,10 @@ def entrypoint(
                     writers=[ChannelWrite([ChannelWriteEntry(END)], tags=[TAG_HIDDEN])],
                 )
             },
-            channels={START: EphemeralValue(Any), END: LastValue(Any, END)},
+            channels={
+                START: EphemeralValue(input_type),
+                END: LastValue(output_type, END),
+            },
             input_channels=START,
             output_channels=END,
             stream_channels=END,
diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py
index 64b5076de..905b58b2e 100644
--- a/libs/langgraph/tests/test_pregel.py
+++ b/libs/langgraph/tests/test_pregel.py
@@ -1450,6 +1450,17 @@ def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> Non
         answer = interrupt("question")
         return [m + answer for m in mapped]
 
+    assert graph.get_input_jsonschema() == {
+        "type": "array",
+        "items": {"type": "integer"},
+        "title": "LangGraphInput",
+    }
+    assert graph.get_output_jsonschema() == {
+        "type": "array",
+        "items": {"type": "string"},
+        "title": "LangGraphOutput",
+    }
+
     thread1 = {"configurable": {"thread_id": "1"}}
     assert [*graph.stream([0, 1], thread1)] == [
         {"mapper": "00"},