Incorrect `add_node` types for `Callable`

[chadhietala]

Checked other resources

• [X] I added a very descriptive title to this issue.
• [X] I searched the LangGraph/LangChain documentation with the integrated search.
• [X] I used the GitHub search to find a similar question and didn't find it.
• [X] I am sure that this is a bug in LangGraph/LangChain rather than my code.
• [X] I am sure this is better as an issue rather than a GitHub discussion, since this is a LangGraph bug and not a design question.

Example Code

from langgraph.graph import StateGraph
from langchain_core.messages import HumanMessage
from typing import TypedDict
from langchain_core.runnables.config import RunnableConfig

class MyState(TypedDict, total=False):
    input: HumanMessage

def some_node(state: MyState, config: RunnableConfig) -> MyState:
    return state

workflow = StateGraph(MyState)
workflow.add_node("some_node", some_node) # Will have a type error

Error Message and Stack Trace (if applicable)

error: Argument 2 to "add_node" of "StateGraph" has incompatible type "Callable[[MyState, RunnableConfig], MyState]"; expected "Runnable[Any, Any] | Callable[[Any], Any] | Callable[[Any], Awaitable[Any]] | Callable[[Iterator[Any]], Iterator[Any]] | Callable[[AsyncIterator[Any]], AsyncIterator[Any]] | Mapping[str, Any]"  [arg-type]

Description

add_node has a type of RunnableLike, however RunnableLike's type does not account for a Callable that takes both the state and the RunnableConfig. In the how-to guides for passing runtime values into tools, it mentions that the second argument is a RunnableConfig.

System Info

python=3.10.4
langgraph==0.1.8
langchain==0.2.8
langchain-community==0.2.7
langchain-core==0.2.19
langchain-experimental==0.0.61
langchain-openai==0.1.16
langchain-text-splitters==0.2.2

[hinthornw]

@vbarda think this is generally applicable to RunnableLambda too ya?

[Glinte]

I have previously opened a similar discussion https://github.com/langchain-ai/langgraph/discussions/1013 Changing the function to

def some_node(state: MyState, *, config: RunnableConfig | None = None) -> MyState:
    return state

will make the function has signature Callable[[MyState], MyState] which is compatible with RunnableLike``

[vbarda]

@vbarda think this is generally applicable to RunnableLambda too ya?

yea i believe so

[gbaian10]

I had the same problem, prompting type error. I once used RunnableLambda to wrap it and successfully eliminated the type warning.

But when I read the event "on_chain_stream" with name="xxx" it produces extra output. This bothered me a lot and I ended up removing RunnableLambda.

from typing import Annotated, TypedDict

from langchain_core.runnables import RunnableConfig, RunnableLambda
from langgraph.graph import StateGraph, add_messages

class MyState(TypedDict):
    messages: Annotated[list, add_messages]

def some_node(state: MyState, config: RunnableConfig) -> MyState:
    return {"messages": [("human", "Foo Bar!")]}

workflow = StateGraph(MyState)
workflow.add_node("some_node", some_node)
# ^ Will have a type error hint. (But it can run.)
workflow.set_entry_point("some_node")
app = workflow.compile()

workflow2 = StateGraph(MyState)
workflow2.add_node("some_node", RunnableLambda(some_node))
# No type error, but it generated some extra events in astream_event.
workflow2.set_entry_point("some_node")
app2 = workflow2.compile()

async def main() -> None:
    async for event in app.astream_events({"messages": []}, version="v2"):
        if event["name"] == "some_node" and event["event"] == "on_chain_stream":
            print(event["data"])  # This will print a message.

    print("-" * 20)

    async for event in app2.astream_events({"messages": []}, version="v2"):
        if event["name"] == "some_node" and event["event"] == "on_chain_stream":
            print(event["data"])  # This will print two messages.

if __name__ == "__main__":
    import asyncio

    asyncio.run(main())

Is it a wrong usage to wrap it with RunnableLambda here?

[vbarda]

@gbaian10 if the only reason to use RunnableLambda is to avoid type warnings, then definitely wrong