Using langgraph Interrupts in AWS Lambda with WebSockets

[wuodar]

Hello,

I'm developing an assistant using Langchain + LangGraph, deployed on AWS Lambda, with communication handled via Amazon API Gateway WebSocket API.

I want to incorporate human-in-the-loop functionality, but before that, I need to implement a checkpointer for my chosen database, which seems like a significant amount of work. Therefore, I'd like to explore whether langgraph interrupts could work in my environment.

The main challenge I see is that AWS Lambda is serverless, meaning that if I interrupt the graph to send a message to the user via WebSocket, the Lambda instance handling the execution will terminate before the user responds. When the user replies, their response will be received by a new Lambda instance.

Given this, is it possible to implement langgraph interrupt in such an environment? If so, could you provide an example or guidance on how to make it work?

The documentation provides the following example:

from langgraph.types import interrupt

def human_node(state: State):
    value = interrupt(
        # Any JSON serializable value to surface to the human.
        # For example, a question or a piece of text or a set of keys in the state
       {
          "text_to_revise": state["some_text"]
       }
    )
    # Update the state with the human's input or route the graph based on the input.
    return {
        "some_text": value
    }

graph = graph_builder.compile(
    checkpointer=checkpointer # Required for `interrupt` to work
)

# Run the graph until the interrupt
thread_config = {"configurable": {"thread_id": "some_id"}}
graph.invoke(some_input, config=thread_config)

# Resume the graph with the human's input
graph.invoke(Command(resume=value_from_human), config=thread_config)

This works if everything happens within a single instance, but in AWS Lambda, after sending a message to the user, the instance will be terminated. Then, when the user responds, their input will be processed by a new Lambda instance.

Can langgraph handle this scenario? If so, how should I structure the workflow to properly resume execution?

Thanks in advance for your insights!

[JoergSchultz-TWT]

Hi!
I implemented this by writing the checkpoint into Dynamodb with a unique identifier. When the answer from the human comes back, the frontend also sends the identifier. The lambda function can then read the checkpoint from the Dynamodb and I can add the human input and kick start the langgraph from there.

[wuodar]

Question - can't we just load checkpoint using thread_id as the identifier?

[wuodar]

Could you provide some pseudo-code how your solution works, based on the example I mentioned?

[wolfman-numba1]

See if this might help: https://github.com/aaronsu11/langgraph-checkpoint-dynamodb
Think it was written for the exact use case you have in mind of writing LangGraph in lambda, interrupting and picking it up again later.

[JoergSchultz-TWT]

And here we go. Admittedly, the whole construct is a pain....

def save_checkpoint(plan_id, config):
    global graph
    print("Saving checkpoint")
    state = graph.get_state(config)
    serialized_state = pickle.dumps(state)
    encoded_state = base64.b64encode(serialized_state).decode('utf-8')
    try:
        checkpoint_response = checkpoint_table.get_item(Key={'id': plan_id})
        item = checkpoint_response.get('Item')
        current_time_iso = datetime.utcnow().isoformat() + 'Z'  # ISO 8601 format for updatedAt
        current_time_epoch = int(datetime.utcnow().timestamp() * 1000)  # Epoch time for _lastChangedAt

        if item:
            current_states = item['checkpoints']
            current_states.append(encoded_state)
            current_version = item['_version'] + 1
            createdAt = item['createdAt']
        else:
            current_states = [encoded_state]
            current_version = 1
            createdAt = current_time_iso
        _lastChangedAt = current_time_epoch
        updatedAt = current_time_iso
        print("Setting owner to ", user)
        checkpoint_table.put_item(Item={
            'id': plan_id,
            'checkpoints': current_states,
            'owner': user,
            '_version': current_version,
            '_lastChangedAt': _lastChangedAt,
            '_deleted': False,
            'createdAt': createdAt,
            'updatedAt': updatedAt
        })
        print("Item successfully updated in DynamoDB")
    except ClientError as e:
        print(f"An error occurred: {e.response['Error']['Message']}")


def restore_checkpoint(plan_id, config):
    try:
        checkpoint_response = checkpoint_table.get_item(Key={'id': plan_id})
        item = checkpoint_response.get('Item')
        if item:
            current_states = item['checkpoints']
            state = current_states[-1]
            state = base64.b64decode(state)
            loaded_state = pickle.loads(state)
            graph.update_state(config=config, values=loaded_state.values, as_node="handler_interacter")
    except ClientError as e:
        print(f"An error occurred: {e.response['Error']['Message']}")


async def _run_agents(behavior, answers, plan_id, stage):
    config = {"configurable": {"thread_id": plan_id}}

    if behavior:
        for s in graph.stream({"target_behavior": behavior},
                              config=config, interrupt_after=["handler_interacter"]):
            print(s)
        save_checkpoint(plan_id, config)

    if answers:
        restore_checkpoint(plan_id, config)
        # add the answers to the state
        answers_data = json.loads(answers)
        print("Adding answers to the state: ", answers_data)
        graph.update_state(config=config, values={"handler_information": answers_data}, as_node="handler_interacter")
        for event in graph.stream(None, config=config):
            print(event)
        save_checkpoint(plan_id, config)
        save_plans(plan_id, config)
    return "done"

[AbhishekBose]

@wuodar Are you calling workflow.compile in every call to the lambda?

[elango5292]

hey @wuodar, can we talk about your experience using langraph on lambda? was there any issue in terms of scalability, latency and availability?