NeMo Relay LangChain Integration

1
import asyncio
2
3
import nemo_relay
4
from langchain.agents import create_agent
5
from langchain_core.tools import tool
6
from nemo_relay.integrations.langchain import NemoRelayCallbackHandler, NemoRelayMiddleware
7
8
@tool
9
def get_weather(location: str) -> str:
10
    """Get the current weather for a location."""
11
    return f"The weather in {location} is sunny and 72 degrees."
12
13
agent = create_agent(
14
    model="nvidia:nvidia/nemotron-3-nano-30b-a3b",
15
    tools=[get_weather],
16
    middleware=[NemoRelayMiddleware()],
17
    system_prompt="Use tools when they are relevant. Keep the final answer brief.",
18
)
19
20
input_payload = {
21
    "messages": [
22
        {
23
            "role": "user",
24
            "content": "What is the weather in San Francisco?",
25
        }
26
    ]
27
}
28
29
with nemo_relay.scope.scope("langchain-request", nemo_relay.ScopeType.Agent):
30
    result = asyncio.run(
31
        agent.ainvoke(input_payload, config={"callbacks": [NemoRelayCallbackHandler()]})
32
    )
33
34
final_message = result["messages"][-1]
35
print(f"Final response: {final_message.content}")