Langgraph Help: Writing awaitable actions in nodes

[areddy294]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

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


class ReturnNodeValue:
    def __init__(self, assistantId: str):
        self.assistantId = assistantId

    async def __call__(self, state: State, **kwargs) -> Any:
        response = await self.custom_api_call(state["messages"][-1], self.assistantId)
        return {"messages": [response]}

    async def custom_api_call(self, message: str, assistantId: str) -> str:
        return await callApi(message, assistantId)


class GraphBuilder:
    def __init__(self):
        self.graph_builder = StateGraph(State)
        self.buildGraph()

    def buildGraph(self, assistantId1, assistantId2):
        self.graph_builder.add_edge(START, "node1")
        self.graph_builder.add_edge("node1", "node2")
        self.graph_builder.add_edge("node2", END)

        self.graph_builder.add_node("node1", ReturnNodeValue(assistantId1))
        self.graph_builder.add_node("node2", ReturnNodeValue(assistantId2))

    async def executeGraph(self):
        try:
            initial_state = {"messages": ["Hellow"]}
            app = self.graph_builder.compile()
            result = await app.ainvoke(initial_state)
            # Error on above line: "Dispatcher.dispatch_forever.<locals>.<lambda>() got an unexpected keyword argument 'context'"
            return result
        except Exception as e:
            print(f"An error occurred: {e}")
            raise


graph = GraphBuilder()
graph.buildGraph("abc123", "abc234")
result = graph.executeGraph()

Description

Hello, I'm trying to follow the langgraph documentation and create a basic graph as below. The main caveat of my graph is that I don't have an LLM, but just make API calls that aggregate as a response. The problem I'm facing is with writing awaitable Actions to make the API calls in the graph nodes.
NOTE: This same code works if I change it to its non-awaited counterparts.
Reproducible code is attached.

Error: Dispatcher.dispatch_forever..() got an unexpected keyword argument 'context'

The error happens when I hit the app.ainvoke(), and as mentioned in the note, the code works perfectly fine if I don't await anything.

Any help would greatly be appreciated, Thank you :)

System Info

System Information

OS: Windows
OS Version: 10.0.19045
Python Version: 3.11.5 | packaged by Anaconda, Inc. | (main, Sep 11 2023, 13:26:23) [MSC v.1916 64 bit (AMD64)]

Package Information

langchain_core: 0.3.18
langchain: 0.3.0
langsmith: 0.1.128
langchain_openai: 0.2.8
langchain_text_splitters: 0.3.0
langgraph: 0.2.48

Optional packages not installed

langserve

Other Dependencies

aiohttp: 3.8.5
async-timeout: 4.0.2
httpx: 0.27.2
jsonpatch: 1.33
langgraph-checkpoint: 2.0.4
langgraph-sdk: 0.1.36
numpy: 1.24.3
openai: 1.54.4
orjson: 3.10.7
packaging: 24.1
pydantic: 2.9.2
PyYAML: 6.0
requests: 2.31.0
SQLAlchemy: 1.4.39
tenacity: 8.2.2
tiktoken: 0.8.0
typing-extensions: 4.12.2