Log In to continue learning
Nvidia’s NeMo Agent Toolkit: Making Agents Reliable
Intermediate
1h20m
Topics
Agents
In this video, you'll expand your NAT workflow to serve as a ReAct agent. And you'll add Python tools for analyzing NOAA climate data. You'll register your functions, define input schemas with Pydantic, and you'll add descriptions so your ReAct agent can determine which tools to use. You'll turn your agent from a simple chatbot into one that takes intelligent data-driven actions. Alright, let's code. In this lesson we're going to build on what we learned previously. This time we're going to build a climate analysis assistant. And we're going to use multiple tools with that assistant. We're going to take plain Python functions and we're going to turn them into NAT tools. And then we're going to register those tools to make them available to our agents. We'll package all that up as a Python package, and we'll also examine some other concepts such as how a ReAct agent uses tools. What is a ReAct agent? It's a common agentic pattern. It stands for reasoning and acting. And it's basically a loop. An LLM will reason about what to do next. It'll formulate one or more actions, it'll call those actions, and then it'll reason some more. And it'll continue that loop until it's finished. This ReAct Agent has access to three tools, and you'll notice that each of these tools is written in a completely different framework. When it runs, it'll decide which of these tools to use, and it'll iterate on tool use until it has a final answer. The data set we'll be working off of is real climate data from NOAA. And it includes information from multiple countries like the United States, France, Japan, Brazil. We'll introduce some plain Python functions that are able to analyze this data, and we'll see how to wire these functions into an agentic workflow. Python functions are the foundation of NAT tools. The top function is a vanilla Python function to calculate temperature statistics. It optionally takes a country and then it returns the statistics. If we want to use this as a NeMo Agent Toolkit tool, we simply register it as a function, and then we wrap it in some boilerplate, such as an async method and an async wrapper. Ultimately, it calls the calculate_statistics function, and then we yield some FunctionInfo to let the NeMo Agent Toolkit know how this function runs. It's a bit of boilerplate, but it's not too much. And now we're ready to wire it into our agentic workflow. The NeMo Agent Toolkit CLI can generate most of the boilerplate you need to wrap functions. But we're going to look at this here and then in a lot more detail on the slides. Every tool starts with a configuration class. This is the class that takes information in our YAML file and makes it available to the tool. Every tool also has an input schema. This is a Pydantic class that tells the large language model or the agentic workflow how to call the tool. And finally, we have to decorate the function with a registration decorator. and this wires the tool into the NeMo Agent Toolkit. We will look at this in a lot more depth in the notebook. To package our workflows, we just use a standard Python project, a pyproject.toml file. There's nothing NeMo Agent Toolkit specific here. We give the project a name, we specify entry points, and then we install it using pip or uv pip install. Once we have all of our tools registered in the NeMo Agent Toolkit, we simply have to create a configuration file. In this configuration file, you can see we have a list of functions, and the top one is calculate_statistics. This is the tool we were just discussing that we wrapped to make it available to the NeMo Agent Toolkit. There's a couple other tools as well. filter_by_country and create_visualization. These are similar functions that we've wrapped. You can see that our react_agent has gained some powers. It has a tool_names property, and we have passed the three functions that we have defined above to the react_agent, now is ready to reason and act using these tools. Of course, the configuration file is not just about running the agent. While it provides tools, it also provides for telemetry, such as open telemetry tracing. and running evaluations against our agent to make sure it doesn't regress over time. The NeMo Agent Toolkit project is just a standard Python project. You can use the NAT CLI to generate a lot of the code that we've looked at. For instance, this command, nat workflow create climate_assistant will create the following project structure. This is just a standard Python project that is ready to go with NAT. All right, let's add some intelligence to our agents with tools. First, let's look at some standalone climate tools. These are not agentic tools, these are just Python functions. We're going to import a series of Python functions here. that can load climate data, they can calculate climate statistics, they can filter climate data by country, find extreme years, create visualizations, and list the countries of the climate data. Let's take a look at what this climate data looks like. So we have 1210 temperature records ranging from 1950 to 2025 and spanning 10 countries. And now let's look at some of the functions. So the documentation for calculate_statistics says it calculates basic statistics from temperature data. It takes a data frame with temperature data and an optional country to filter it by and it'll return a JSON string with the statistics. So let's call this function and see how it works. In this case we've called calculate_statistics and we've passed it the df, but not a country. Remember country is optional. And let's print the global statistics and we get some data output here. Now let's look at a different function, create_visualization. And this creates climate data visualizations and saves it to a file as an image. And we can see the arguments that it takes here and what it returns. Let's test this function out and see what it does. Here we see that calling create_visualization and we pass it a plot_type and a place to save will output a plot of climate visualization data. Each of these tools returns JSON. They have a single responsibility, but they are not plugged into our agentic workflow. Let's go ahead and do that. Now we're going to register a single tool into our NeMo Agent Toolkit workflow. When you register a tool with the NeMo Agent Toolkit, there are three things that you need to consider. The first is the input to the tool. So NeMo Agent Toolkit will use pydantic to let you define what the input to the function looks like. So we're going to import BaseModel and Field from pydantic, and our tool is calculating statistics, so we will call this input definition CalculateStatsInput. and calculate stats takes a country that is optional. And so in our case, we will give it a description, country name to filter by, leave empty for global statistics. And now that we have our inputs defined, we can move on. Next, we'll set up our configuration. NeMo Agent Toolkit is config based and we've seen YAML files already. How do the properties from the YAML get into our tool or our LLM or our workflow? it's through a config object. So in our case, we're going to define a CalculateStatisticsConfig. We'll give it a name, and this is going to become the type that we set in our YAML file. and we can give it a description. And in our case, we're just going to pass, which means there's no properties on this class, which means that our calculate statistics tool does not take any configuration. But if it was going to take any configuration properties from the YAML, we would define them here. And now we're going to register the tool itself. We register with a decorator, register_function, and we pass into that decorator the config_type, which is that config that we created earlier. This is a bit of boilerplate and the NAT CLI has generate commands that will generate some of it for you. We're not going to talk in detail about all the aspects of this, but we'll see here that we have a wrapper function that wraps the Python function that we're going to expose as a tool. It takes two properties that we'll talk about in greater depth later on. the configuration which will be injected in. And if we had properties on our configuration object that we defined earlier, then this would be an argument that contains those properties pulled from the YAML. And then a builder which we can use to grab other aspects of our workflow such as an LLM or a retriever or any other thing that we've defined in our YAML. In our case, we're not going to be using either of these. We are going to load our climate data as a data frame. And then as we wrap our tool, we will pass that data frame in. Here's the actual function that we're wrapping, calculate_statistics that we looked at earlier. We pass in the data frame that we loaded a few lines above. And then you see that we are also accepting a country parameter. This is optional of course, as we defined earlier, and we will pass that country in as well to calculate statistics and then we'll return the result. So we've wrapped the Python function, but we still need to describe it to the NeMo Agent Toolkit. And so what we will do from this is we will yield a FunctionInfo object. and the function that we'll pass to that FunctionInfo object is the _wrapper that we just looked at. And the input_schema that we talked about earlier that defines what the input looks like to our tool, and a description, and this description will help the agentic workflow know whether it should call this tool or not. So a bit of boilerplate code, but again, this would be generated for you. A lot of it would be generated for you by the NeMo Agent Toolkit CLI. If you wanted to see and run the full register.py file, you could run this cell here and we would be able to see what the Python file looks like. Now let's look at what the configuration looks like. We have this configuration already saved to disk, so we're going to load it into our cell and take a look. And here's the config file that we're going to use for this NAT workflow. Some of it will already be familiar to you. the climate_llm that we looked at in the previous notebook. But we have a new section and this is called functions. And these are the tools that we can make available to our agents. and the first and only function we have listed here is calculate_statistics. And we give it a type and we looked at this type earlier simple_tool_demo/calculate_statistics. This is how the NeMo Agent Toolkit knows how to find this tool in our Python functions. and we can override the description here for the tool. Our workflow is similar but different than before. So you'll recall before the type was just a simple LLM call, and now we're telling it that the workflow is a react_agent. A ReAct agent reasons and then acts in a loop, and this is a built-in agent to the NeMo Agent Toolkit. So we're going to say run a ReAct agent and the ReAct agent is going to have the following tools available. and we pull the tool from the functions listed above. So we have calculate_statistics here and we have passed that in as a tool to our react_agent. It also takes an LLM, which we defined earlier, the climate_llm. And we're going to set verbose to true, and this way we'll be able to see detailed outputs for the tool use. All right, let's go ahead and install and run our ReAct Agent workflow. So the tools that we looked at earlier do need to be installed into our virtual environment. And so I'm going to go ahead and install them here. And now that those are installed, we can run our workflow. So you're familiar with nat run. We'll pass it the config file that we looked at earlier, and an input, What is the global temperature trend per decade? And we can see it reasoning through. the agentic input, what is the global temperature trend per decade? and the agent thinks about this and decides that it's going to call calculate_statistics. So it calls calculate_statistics and it gets an error, no data found for country None. So this is probably room for us to improve our agent later on. realizes it should call with an empty string. So it calls calculate_statistics again with an empty country, and it gets back the results that we looked at earlier when we ran calculate_statistics. The agent concludes it knows the final answer and it gives us the answer here. The global temperature trend per decade is 0.241 degrees Celsius per decade. Let's uninstall our simple tool demo package, just so we'll be ready to run other packages and other notebooks. So let's talk through what we did. To register a Python function as a tool, we need a configuration class with a unique name. We need an async wrapper function. we need to yield a FunctionInfo with an LLM description and then we can reference it in our config file as package_name/tool_name. You can use NAT workflow create to create all the scaffolding you've done by hand in this notebook. And if you want more information on that, you can run NAT workflow create dash dash help. Here's a configuration file that has all of the tools that we looked at earlier registered in. So we have calculate_statistics, which we've previously seen, but now we've wired in the other tools as well. list_countries, filter_by_country, find_extreme_years, create_visualization. These are just sub properties of the functions top level property. And then we'll pass them into our ReAct agent here. Previously we just had calculate_statistics listed as a tool. But now we've listed all of the tools and made them available to our ReAct agent. We've also made two other small changes. Since we have many tools, sometimes the ReAct agent needs to iterate several times. We've told it that it can iterate up to five times. And then also, sometimes it will call a tool incorrectly and have to retry. and we've told it that it can retry up to two times. All right, now that we have all of our tools registered, we are ready to run more complicated queries against our agent. First, let's ask about global warming trends. We're going to run NAT run and we'll pass it that long configuration file we just looked at with the input of what is the global temperature trend per decade. And we can see the ReAct Agent reasoning down through the tools that we provided it, decides to call calculate_statistics and it gets us the global temperature trend per decade. Now let's ask for country analysis. In this case, we're going to ask about a specific country, specifically France's climate data. Let's see if it'll use multiple tools during this run. We can see it filtering by country. and we get the results back from the tool of specific country data. And then we see it calculating statistics, and this time it's passing a specific country, which is France, and it gets data back from that. And now we see the LLM take the provided JSON that it got from the various tools and put it together into a final answer, which it returns to us. Let's test the visualization tools next. In this case, we're going to ask for a visualization showing which countries have the highest warming trends. And we can see that it's chosen to call create_visualization. And it's let us know that it has saved that to a file and that file is ready for us to view. Let's take a look at what it rendered. And here we can see it gave us a rendering of the top five countries by warming trend. Finally, let's give it a really difficult question that's definitely going to require multiple tools and see how it handles it. In this case, we want it to compare the temperature trends of Canada and Brazil to answer the question which one is warming faster and then also create a visualization of global trends. And we can see it reasoning down through calculating the statistics per country, first for Canada, then for Brazil. We can see it having to retry a few times. But we finally get to the point where it's able to save a PNG image of the global trends and give us a final response giving us the trend per decade. And let's look at that output. And here we can see a plot of temperature over year. And now we'll do some cleanup. So let's delete the images that we've created. And let's uninstall the workflow that we created just so it doesn't conflict with later notebooks in the course. In this notebook, we've seen how we can take standalone Python functions and wrap them as tools to be used by NeMo Agent Toolkit agents.
