NeMo Relay LangGraph Integration

1
from typing_extensions import TypedDict
2
3
import nemo_relay
4
from langgraph.graph import END, START, StateGraph
5
from nemo_relay.integrations.langgraph import NemoRelayCallbackHandler
6
7
class State(TypedDict):
8
    value: int
9
10
def increment(state: State) -> State:
11
    return {"value": state["value"] + 1}
12
13
builder = StateGraph(State)
14
builder.add_node("increment", increment)
15
builder.add_edge(START, "increment")
16
builder.add_edge("increment", END)
17
18
graph = builder.compile()
19
20
with nemo_relay.scope.scope("langgraph-request", nemo_relay.ScopeType.Agent):
21
    result = graph.invoke(
22
        {"value": 1},
23
        config={"callbacks": [NemoRelayCallbackHandler()]},
24
    )
25
26
print(result)

1
from langchain.agents import create_agent
2
from langchain_core.runnables import RunnableConfig
3
from nemo_relay.integrations.langgraph import NemoRelayMiddleware
4
5
agent = create_agent(
6
    model="nvidia:nvidia/nemotron-3-nano-30b-a3b",
7
    tools=[],
8
    middleware=[NemoRelayMiddleware()],
9
)
10
11
def agent_node(state: dict, config: RunnableConfig) -> dict:
12
    return agent.invoke({"messages": state["messages"]}, config=config)