PydanticAI - Building a Research Agent

Okay. So in the last video, we looked at an introduction to Pydantic AI. Since then, I've been playing with it quite a bit. I gotta say it's really nice for building things. What I wanted to do was create a couple of follow-up videos.

I've got another one on RAG coming, but what I wanted to do in this one was build a research agent. So, I guess a lot of people have done things like Perplexity and similar models. This is not exactly like that, but it's going to show you many of the elements that you can use for such applications and how Pydentic.ai makes those tasks really easy. It's a similar setup to last time. I'm basically using a simple notebook, so everyone can load this and use it straight away without having to set up their own environment, etc. Once you install the different components, you need to restart the notebook. I mentioned that last time. Another thing I have set up here is two different search engines. I have configured DuckDuckGo Search, which is the one I would probably normally show more often.

The challenge with that at the moment, though, is definitely that they seem to be rate-limiting searches. So, if you've got a proxy system, you can use that with DuckDuckGo, and that works quite nicely. The other alternative is to use Tavilli, which I've got here. I'm just showing you those two quickly: DuckDuckGo search is pretty easy to do. There's a nice unofficial Python package for this.

It supports both synchronous and asynchronous calls. In this particular agent, we want to use asynchronous calls. It has a few different types of searches. The same is true for Tavilli. If you set this up with the client being asynchronous, then everything will run in an asynchronous manner here.

And you can see that you can perform a basic search where we pass in the maximum number of results. Then we can observe that we receive three results back. Alternatively, they offer what is called a contextual search, which will also go through and provide more detailed responses. What I recommend is that you pick one of these options. The cool thing about Tavili, even though it's a paid service, is that you can obtain an API that allows for a thousand free calls per month.

So you can certainly get started for nothing. You don't even need to put in a credit card or anything like that. They're quite reasonable in that way. So you've got a choice: you can either use the DuckDuckGo search engine or the Tavily search engine.

I'm going to go with the Tavili search option, just so we don't encounter any rate limits or other issues here. All right, what I want to show you is how to build an agent where we're going to include a number of dependencies. You'll see that we're bringing in various components from data classes and some Pydantic libraries as well.

What we're going to do is create a search data class, which I will primarily use for maximum results. This class will be passed through the agent to facilitate the search tool. Additionally, we will introduce a research dependency, but I will come back to that later; I will not be using it at the start.

I'll use this later on and I'll explain why. The other big thing is we're going to use a result type. So this is basically a data class—a pedantic data class, not actually Pedantic AI, just the original Pedantic. In this, we're going to specify what we want in the final result. For the final result, I want to have a research title.

I want to have a main section for research, followed by some summaries at the end. Each of these will be represented as strings. I would like the output to be in Markdown format, although we may encounter some challenges with the heading titles later on. For example, the main section will state: "This is the main section that provides the answers to the query and research for the bullets." The summaries will be presented as a set of bullet points that summarize the answers to the query.

All right, next up, I'm going to set up the answers for the query here. Next, I'm going to set up the agent itself. In this case, I'm using OpenAI's GPT-4, but you could experiment with different models, right?

This is the cool thing: it's literally to take this from GPT-4 to GPT-4 Mini or to Gemini, etc. We can just swap this out, and it will take care of the rest on the backend for us. While they don't have Anthropic yet, it's supposedly on the way. So, hopefully, we'll get that in the not-too-distant future. All right.

So you're going to see here, in my setting up the agent, I am passing in these research dependencies and the result type. Let's focus on the result type first. As I mentioned, this will basically define what the structured output is going to be. We also have our system prompt. Here, we have "Your helpful research assistant."

You're an expert in research. If you were given a question, you would write strong keywords to conduct three to five expert searches. The agent itself will decide how many searches to perform, with each assigned a query number. I want to demonstrate this process by passing the query number through, allowing us to see which queries were chosen and other relevant details.

And then you're going to combine the results. Now, I can certainly elaborate more on this system prompt, right? This is a very simple system prompt here. I could, for example, if I wanted a research aspect to be academic, I could say to focus on academic topics. If I wanted it to be more commercial, such as checking prices or something similar, I could load all of that.

All those kinds of things are included in this system prompt. This information will be used to determine the searches and how they are filtered to provide our responses. Next, we need to create the search agent tool. If you look back, we have this search agent tool, and the main parameter we will pass in here is the maximum number of results.

So, while creating the tool, we essentially have a decorator that will be at `search_agent.tool`. In this case, I have commented out the DuckDuckGo search option, but you could include that one. Just don't include both, as it would probably cause confusion. Now, the question is, how will it know what this tool is?

We've got a docstring for that, right? So this is going to handle the search for a keyword query. The query will be the keywords to search for here. We have the query that we pass in, which will come from the model, and we also have the query number that we refer to up here, which will come in as well.

These are decided by the LLM. We then have the search data class, which really defines the dependencies for this. I’m going to constrain it so that for each query, you’ll see later that I will set a maximum result limit. I will pass that in at query time. You can imagine that if we had some kind of UI with a dropdown asking, "How many results do you want for each of these queries?" we could pass that value in there.

Now we're going to show what the search query is and the number associated with it. These two are obviously coming from the LLM. Then, Tavilli is going to perform a search. It will execute a "get search context" and pass in the query.

It's going to pass in the maximum results that it got from the dependencies that we've put in there. The same kind of thing applies for DuckDuckGo here. All right. Now, when setting up the dependencies, you can see that I want to have the date, which I'm not going to use at the start, but we will use it later on. I'm going to show you how you would inject that into the system prompt.

But we can basically create a date string and pass that into the search data class. Actually, where I'm going to end up using it is not here. I will use it in the system prompt. I found that to be an easier way to do it, and it also demonstrates how to update the system prompt on the fly.

That's why I've put it in there. You could place it here and then append it to the query in a very specific way if you prefer. One of the advantages of the method I'm going to show you is that if it doesn't need to be appended, it simply won't be injected. The large language model (LLM) can make that decision by itself. All right, we have these dependencies, with the key one being the maximum results.

All right, so I've basically set these up, and now I'm just going to inject my query into this. I've got my search agent, and I'm just going to say, "Can you give me a very detailed bio on Sam Altman?" We can see that, in this case, it decided to execute four queries as it went through this. Now, you notice that each of the queries is not waiting for the previous query to complete, right?

It's all asynchronous in this case. We can see that the queries it decided on—remember, it could choose between three to five queries—and for each query, it's still getting three results back. So we're getting quite a lot of tokens as we look at this. If we look here, we can see that it decided to go for the search queries: "Sam Altman biography," "early life," "Sam Altman career," "OpenAI," "Y Combinator," "Sam Altman achievements," "contributions," "personal life," and "philanthropy." All right.

If we come in here and look at the data, we can see that, sure enough, this data has come back. It has gathered all that information and sent it back to the LLM to format it into a response for us. Our response needs to have a research title.

It needs to have the research main, and it should have research bullets included as well. So let's review these as we go through. I'm going to append the Markdown heading simply because I got frustrated with trying to get the right prompt to actually include that. It may be that some kind of filtering is occurring; I'm not sure why it wouldn't automatically add that.

But anyway, you can see that we've got our research title there. If we look at the main research, we've got that as well. Additionally, if we examine our research bullet points, we have those too. Now we have a structured output that includes the following: `result.data.research title`, `result.data.research main`, and `result.data.research bullets`. If we combine these into Markdown and print it out, we can see that we are generating a comprehensive Markdown response. Each time you run it, the output will vary. Sometimes, you may even see subheadings included, which we could have explicitly requested.

We could have said, "Okay, break it up into Subheading One, Research Main Part One, Subheading Two, Research Main Part Two," etc. We can structure this however we want it to turn out. That's one of the cool things about this. Now, if we come in here and actually look at the result—the whole thing that got sent back—and examine what happened throughout this process, we can see that we created our agent. Our agent had a system prompt.

We had a user, a user prompt, and then we had the different queries. Essentially, Query One was sent to the search tool, followed by Query Two, Query Three, Query Four, and so on. One of the interesting aspects of this approach is that if I ask for a bio of Sam Altman without specifying "detailed" or similar terms, it will likely perform just three query searches. In fact, during most of my testing, when I requested very detailed information, it actually returned five queries.

So, the cool thing here is that the agent is deciding for itself what the key keyword searches should be for this topic. I can select between three and five of those keywords. We can then see the tool returning results for each of these keywords. For example, we have four different keywords, and each of them will generate three responses. In this case, it's very clear; we have one response here, the next one is another response, and so on. It's evident that we are generating a lot of tokens in this process, which is the key takeaway.

And then now it goes back to the model. The model is using this to generate a structured response that results in our final output. The final result will include our research title and the main research components that we have identified.

And then somewhere down here are going to be the research bullets. The cool thing is that if we decide we want to write a 500-word essay or something like that, and we want it structured in a very particular way, all we have to do is go back to the result type class and flesh that out further. We could change that to be a variety of different things that fit the Pydantic schema, etc. So in the end, this is what we get back. Interestingly, it did return the hashtag, but maybe it is getting filtered out somewhere.

Anyway, we can see the cost associated with the number of tokens we've used. Here, the total tokens are 3,700. Obviously, we could change this and even provide information in the system prompt about how it should be economical. It shouldn't use more searches than necessary because they cost money, nor should it utilize more tokens than required for the same reason. All of these considerations will respond quite well when going through here. All right, next up, I want to show you something. Let's consider what happens if we take that exact same agent and ask, "Hey, what is the latest AI news?" Now, it's purely relying on its understanding that it needs to retrieve the latest news. So, it realizes that it probably needs to specify a date for the search, but it doesn't have a date provided in its system prompt or anything like that. Therefore, it's defaulting to the cutoff date, which we can see is October 2023.

And so it basically comes back with the latest AI news, October 2023. If we print that out and look at it—this is from an earlier version—we can see that we're clearly getting outdated information. There's nothing new in there. So what if we add in the date? To do that, we can simply inject it into the system prompt.

Sure! Here's the text reformatted with correct capitalization and spelling: "We are just overriding the system prompt, and we are doing that with the research dependencies. The only research dependency we had was today's date. We set that up above, where we used the date time and converted it to a date, etc. We established that above, where we got it from date time."

We converted it to a date, etc. Now we've got: You’re a helpful research assistant. You’re an expert. We have everything the same, except we say, "If you need today’s date, it is," and then we pass it in.

And that's going to be in this format up here, where we can see that we've got the year, month, and day in the string. All right, we inject that in. Now we ask, "What are the major AI news announcements in the last few days?" In this case, it determines that only three search queries are needed, right?

AI News Announcements, December 2024, Possibly Because It realizes that we're only in the first week of December. AI recently announced AI developments and the latest AI advancements. So it comes back with all our data, just like before. We get this, we get our research title, our research main, etc. If we go through and look at those, we can see what they are.

And then, if we print this out, we can see that, okay, sure enough, it has come back with some information. The bullet points are basically saying that OpenAI announced 12 days of Shipmas to eventually reveal new AI products. We can also see rumors about OpenAI launching the ChatGPT Pro plan at $200 and the text-to-video Sora model. So, it has most of those items.

And it has some other features as well. You can imagine that we could ask it to focus on announcements from specific companies—this company, that company, and so on. We could incorporate those into the system prompt as we proceed. Now, however, we have a proper response that matches our expectations for this kind of query today, rather than missing today's date and losing context.

And if we come in here and look at this, we can see that it starts with the original system prompt. However, the system prompt gets overridden to include the date in the year, month, day format. This is a reasonably simple research agent where one agent makes various calls. We could add different tools if we wanted to perform searches on different topics. For example, if it's a search about a person, we would use one search engine, whereas if it's a search about a company, we would use another.

We could include that to improve the responses we receive. This highlights the fundamentals of structuring outputs to ensure you get results that are useful. It's essential to structure the dependencies so you can inject different elements into the agent as you use it. Additionally, it’s important to set up some of the tools properly. In the next video, I will probably demonstrate something similar with a RAG system, which will likely be very comparable. We can then decide how we want to handle the context and other relevant matters, incorporating a vector store as well. As always, if you have any comments or questions, please leave them in the comments below. If you found the video useful, please click like and subscribe, and I will talk to you later.

In the next video. Bye for now!