Advanced RAG: Dynamic Chunk/Document Retrieval (with LlamaCloud)

Key Takeaways

hey everyone uh Jerry from llama index here and today we'll talk about an advanced brag technique that allows you to do Dynamic retrieval or in other words figuring out whether to retrieve specific top K chunks from a document Corpus or returning the entire document um different types of retrieval might be more suitable for different types of questions for instance if you're trying to ask more uh simple questions like you know where does uh the specific person do action Y then you might want to do more uh specific like top K retrieval which is standard in an a rag pipeline however if you're trying to actually ask something that requires the content of an entire document like how did Tesla do in this financial year you might for instance want to return the entire 10K document into the context window this is increasingly more feasible with both longer context Windows as well as lower cost and latency and actually running through large amounts of context and ideally you want to unify this into a single query interface so we'll show you how to build a dynamic retrieval mechanism through that can do both chunk level retrieval and file level retrieval with llama cloud and workflows so we can go into this diagram to take a look at what this agentic architecture would look like to help you do uh Dynamic retrieval first um we the the core is at this thing that we call a router query workflow and it's a simple router prompt which takes it in an input and then depending on the properties of that input decides to Route it to either a chunk level retriever which returns you know specific chunks from documents or a document retriever which returns entire documents from your knowledge Corpus um this uh these retrieval endpoints are indexed through the Llama Cloud interface and so llama Cloud provides you out of the box ways of doing both chunk level and document level retrieval at the end regardless of what retrieval endpoint you pick you take in all the context and feed it to an llm to synthesize a response the outer layer um actually we can wrap this with some sort of agentic reasoning architecture where um instead of directly feeding the query to a router you can actually feed this to some sort of Chain of Thought uh layer first which will iterate through um the different subsections of this question and decide to call this router prompt as a tool and so by wrapping it in this overall you know multi-step reasoning layer then you also get the properties of being able to break down more complex questions into simpler ones so let's walk through and and see how this works first things first we will be loading in um five apple and Tesla 10ks um and then uploading them into L Cloud uh these are basically you know the 10K documents for apple and Tesla um through from uh 2019 to 20123 um and you can see an example of what this data looks like right here um 10K is are a pretty popular data source for building rag demos um and so you can see there's a lot of text there's a lot of tables as well and llama cloud has outof thebox capabilities thanks to our parsing technology to parse that text into uh or parse the document into wall formatted text and tables in markdown format in llama Cloud we can see that you know we've already created an index containing these apple and Tesla 10ks um so just to give you an example um all you have to do is drag and drop these files and create a new index and this is what you get so by clicking into one of these um files you can see the raw document you can see the parse document which is basically uh markdown right and you can see what the chunks look like in this case we specified page level chunking so that every page is a separate chunk um and here for instance is what the table of contents looks like and if you go to the overview page um you'll see this is how we actually plug in this index into the rest ofar jch workflow it gives you back retrieval endpoints that you can call in order to get back relevant context over your data so now we will go in and Define the core router workflow um and so um maybe taking a step back we'll actually do an integration we'll showcase an integration with uh llama Trace which is a collaboration that we have done with arise Phoenix where once we run these workflows you can actually log all these traces to a hosted instance um that's uh that's um like hosted arise Phoenix it's a pretty comprehensive observability and evaluation tool that allows you to capture all these different logs and traces from the workflows and help you see them in a really nice display so we'll show you that in just a bit but first let's actually Define the workflow so in this section we are going to Define um the sections that you need to actually build a router workflow um this box right here just shows you some core prompts um and and also pidan classes that you're going to need one item here is the router prompt where you tell the llm to basically decide uh Which choice to pick given some information about these choices and given the input task and so we are going to use this router prompt to decide whether to pick chunk level retrieval or file level retrieval depending on the question the format string right here tells you how you should format the outputs and basically what you want here is to basically uh tell the router prompt to format your choices in some sort of structured representation here we see that we have the answer classes um and so the answer is defined as a pantic base model um and when you pick a choice it has you you output both the choice index as as well as a reason for it here you can output um multiple choices instead of just a single choice and so uh when you do a router you can select uh multiple choices in this list and then the router output parser basically is a object or a class that you apply at the end of the llm call and so after the llm returns some output you can use the output parser to basically parse that string output into this identic object so in this case we actually are not using um function calling so most LMS these days or many of them support function calling where you actually get back a structured output through the LM API but this is actually a more general form where instead of you know having to uh use the tool interface of the LM directly you basically are prompting the LM to basically output stuff in the schema um of course this doesn't necessarily guarantee that it will output stuff in the schema but most models should obey this format these days and then we're going to use this parser object to then convert the llm response into a structured output you can also do this with function calling if you want um this is just showing you what the current notebook does so once we have that those core components we can define a router query workflow um and this is going to be the first workflow that we Define and it's going to consist of the following steps um given a question we are going to uh choose basically uh what um kind of underlying query engine whether it's file level or chunk level retrieval to call based on uh based on the input um and then we are going to query each one um of the choices depending on which one was selected and then we are going to synthesize a final response based on the results from the queries above similar to how we've defined workflows previously a workflow is basically defined as a python class plus a set of steps and every step passes events from um you know or listens to input events and emits output events so you can pass in whatever initialization parameters that you want here we just pass in the different choices uh or the query engines that we want to specify and so the query engines by the way um once we'll we'll show this once we actually initialize this later on the query engines are basically just full rag pipelines defined over the different retri that we Define in this case the retrievers are again the file and chunk level retrievals um the first step here is we are going to given the start event which contains the input task pick the right query engine so we are going to use the router prompts um uh and plus the output parsing um or structured output parsing to basically uh tell us you know what exactly is the selected Choice um what choices are selected by the L given those choices uh we are going to then uh call the next event which will basically help uh decide you know the query engine that that we want to to query and so this um just returns a choose query engines event which takes in um the output and the query and then in the next step we are going to actually call the relevant selected query engines um you could choose to break this down into separate events actually so that depending on you know the number of query engines you you decide to call you actually emit one or more events here we just call all the query engines in a single step and then the last step here is you know given that we call all the query engines try to Sy synthesize the response and so you know given um kind of the the input um we are going to call the llm one more time or use some sort of summarizer prompt to basically get back the final response so this is the overall workflow and notice that it does depend on this query engines piece which requires us to actually pass in the specific end points that uh or the specific choices that we want the router to pick between and as mentioned we want the router to be able to pick between chunk level retrieval and file level retrieval so we are going to define a llama Cloud retriever over these um over these different documents or Define two different LL Cloud retrievers and then PL plug this into the to the router um so we are going to integrate with that index that we defined um you know for for um the the like apple and Tesla 10ks and then notice that we defined two different types of retrievers so here this is indexa as retriever uh with retrieval mode files via content and then here there's another Retriever with retrieval mode uh chunks this gives you back the two different retrieval interfaces one will return the most relevant chunks given the input and the other will return the most relevant documents we'll plug those both in so a retriever takes in an input and Returns the most relevant chunks um by plugging this into a retriever query engine then this takes in an input and synthesizes the answer uh the retrieve chunks using an llm um and actually gives you back the final answer so by plugging in a retriever into this the retriever query engine you get back an end to end pipeline so we have these two query engines and then the last step is to actually Define you know some descriptions around these query engines because we're passing both as choices into this router the router needs descriptions of what these uh query engines actually do in order to differentiate between the two and prompt the all to decide which one to pick so we are going to define the two different descriptions of these query engines here we have uh the document uh which we description um which basically says you know let's synthesize an answer to your question by feeding in an entire relevant document as context and so it's best used for higher level summarization options um and so try not to use it if the answer can be found in a specific chunk of a given document the chunk level retriever will say synthesizing the answer to a question by feeding in a relevant chunk as context and so it's best suited for questions are more pointed in nature you know you were basically trying to guide the L to help it really understand uh what types of questions uh are better suited for the document Retriever and which questions are better better suited for the chunk level retriever um better prompting probably would include actual few shot examples of questions and answers and also maybe negative examples as well in either case once we defined both the query engine as well as the tool descriptions we plug this in along with the descriptions into this router query workflow that we defined above so we could stop there this already G gives us a router workflow which can dynamically decide whether to query file level or chunk level retrieval but uh we are going to go One Step Beyond and you know the one thing here is that this doesn't really have the capabilities of doing for instance Chain of Thought processing what if the input is actually somewhat complex and you want to be able to decide whether to break this uh task into more simpler steps before deciding uh which you know retrievers to call for each step for instance if you ask you know what is the revenue of both apple and Tesla you probably want to break that into two different questions what is the uh revenue of apple and also what is the revenue of Tesla and then ask those independently against this router system so we'll wrap it with this kind of agentic workflow and so we're actually going to define a router workflow that or an outer workflow that we plug in the subw workflow or the router uh into this this agentic outer workflow will be able to take in a question and then basically do some sort of function calling to dynamically break down that question into individual substeps and keep iterating until all the sub questions are answered so this is very similar to our implementation of what a function calling agent is or what a react agent is and if you want additional references you should definitely check out the the notebooks where we show you how to basically build a react or function call engagement from scratch um and it's very similar to that and it will consist of the following steps first is given an input we are going to append that message to the chat history and so we are going to feed that chat history into the llm along with the set of tools the tool there's a single tool in this case and that's basically just the uh underlying router workflow and the LM will decide uh whether to call that tool or whether it's done and also what is the input to that tool by deciding what is the input to that tool because you know most of these LMS have some sort of function calling capabilities um it will implicitly be able to decide whether to break this more complex question down into simpler calls because it might call the the tool with a more simple input and then the next step is to actually call the tool and run it so we run the router workflow um and you know we we'll we'll basically dispatch the calls to to the tool and call the tool and then we finally gather all the results from the tool calls um and then we finally feed it into an LM for the final response it's basically a while loop that keeps running over uh the the set of tools until the llm implicitly decides that it doesn't need to call any more tools and it can figure out the final answer so you can see how we Define the workflow right here we plug in the rag workflow um you know basically this is the underlying workflow that acts as a tool for this outer agent workflow and then we also um yeah we we plug in this rag workflow um as a tool so we wrap it with what we call a function tool and this basically is an abstraction that allows it to be fed into an agent this workflow has the following steps so the step uh the first step is prepare chat which basically takes in the input and adds it to the trat history and emits an input event the next step is we actually do um we call the llm with tool calling um so we plug in uh the rag workflow tool we use one of the functions that we have within LOM index uh which is a bit lower level on the LM side and it's called a chat with tools and this basically allows you to um not just you know get pass in a general chat completion input and get back a response you also plug in the set of tools that you want and then you can actually parse the output tool calls and so by calling aat with tools with this tool and then calling get tool calls from the response you can directly parse the set of tool calls that it was um it was able to call and this this these tool calls are basically calls to the underlying router if the there are tool calls um you know then we actually emit gather tools event so we process these tool calls and continue to um as in we actually run the tools um and get back the response and then we keep iterating from there if the LM decides that you know there's no more tool calls to be done so that generally means that I can answer this question without needing to call a router anymore then it'll just return the final answer and then this entire workflow will stop so that's basically how it concludes execution if there are tool tool calls to be run then it'll call dispatch calls and then for every tool call it'll call uh what we call a tool call event um so this part you know uh you can see that instead of just sending one event from dispatch calls we send n events where n is the number of tool calls and every tool call event will be routed to a specific step which is basically called call tool and then it will be responsible for calling that specific tool with that set of inputs so for instance if you decide to um for instance what is the revenue of apple and Tesla um the llm might actually decide especially those llms that support parallel function calling like the open AI models um to do two tool tool calls at once one to call the the underlying um query Engine with like what is the revenue of apple and the other what is the revenue of Tesla if it decides to actually call two tool calls at once you're going to be emitting two events right here and then each of those two events will get P to call tool with you know what is revenue of apple and what is revenue of Tesla we will run each of these workflows independently get back an output and then emit like a tool call event result so our responsibility after we call all these tools is to then gather all these tool calls um and then uh the high level idea of gather is to wait for all the tool calls to complete and then finally um append the final response to the trat history so gather is a step that basically uh has a call called collect events and collect events basically allows this step to wait um until all the underlying tool calls have completed collect events will return none if uh you know you haven't reached the count that's specified here um and you know this count because the llm right uh when it Returns the number of uh the the set of tools it wants to call gives you uh that number already and so you're basically waiting until you actually finish calling all these tools before you can then proceed so if there are no if it's it hasn't filled up the counter you'll return none and then if it has then you'll continue to the end for every tool event in tool events we're going to add the the tool call and then also the output of calling the tool to the chat history and then go back to the input event and remember the input event just goes back into chat and then we are going to restart this tool calling process um and so you know let's say we answered the the the questions what is the revenue of apple and what is the revenue of Tesla we added those to the chat history we go back to the chat step we call the llm again you know the llm has full visibility over the entire chat history uh the original question is again what was the revenue of both apple and Tesla if the Alum sees that both of those questions have already been answered in the chat history it will probably decide to not call any more tools and then it will respond uh it will return a stop event so what I just described to you is actually the Second Step Beyond just this initial router layer on top of both file and chunk level retrieval and it's actually optional if you don't want to implement it but it basically is a just generic a a simple agent architecture that does Chain of Thought processing um that you can basically plug in whatever workflow you want and then this thing will always be able to call that workflow as a tool by doing Chain of Thought over it so let's define this entire thing you know we have this router output agent workflow which is this entire agent workflow and then we can visualize what it looks like right here um this is just an illustration of the outer workflow um where again given a start event shows prepare chat uh calling the tools dispatching the tool calls actually executing the tools and Gathering the results and then you know basically iterates until it is done and finally let us uh run some of some of these workflows so uh we are going to try out some of these example queries like tell me the revenue for apple and Tesla in 20121 and if we just go through the logged outputs um you can see it's actually basically doing exactly what I said uh first It produced input event then it is going to uh you know uh call it called the lb and basically looked at you know what are the underlying tool calls I need to execute and you can lcly see what calls that it decided to execute here with the outputs it calls the function query workflow with the message Apple Revenue 2021 and also uh the uh query workflow with message Tesla Revenue 2021 and it's emitting both of these at the same time so the LM actually decided to you know do parallel function calling after seeing the input um next because we call the underlying router you know the router again decides you know whether or not we want to answer this question using uh file level retrieval or chunk level retrieval and choice one is document level retrieval or file level and choice two is chunk level for both of these questions it said it pick choice two and here the question Apple Revenue 2021 is pointed in nature so the llm you know probably based on prior knowledge figure this is something that probably only requires a specific lookup from the document so you don't need the entire document um and so using a relevant trunk as context as more appropr R same thing for Tesla Revenue 2021 then it will gather the responses synthesize the answers um and then finally it will call a gather you know running step gather and give you back the final response um Apple's total net sales and then Tesla's total revenue uh the formatting is a little messed up mostly because of uh kind of like the italics and stuff but basically it gives you back the right answer there's another section which uh is tell me the Tailwinds for apple and Tesla in 2021 um and then here you know it actually picks Choice One or document level retrieval um because the llm based on prior knowledge uh decided that like you know generally it probably needs to understand the entire document uh to um to uh to actually get enough context to answer the question um this part is maybe a little bit more ambiguous so it's understandable if the um decided to choose either chunk level or file level and then the last thing I want to show you is just you know something that's very clear where you need the entire document as context um if you ask a question like how is Apple doing generally in 2019 um this is pretty clear that you basically should probably look at the entirety of Apple 2019 uh as a document and you can see here the selected choice is Choice One or document level retrieval this question requires as a general summary of the entire document and is best suited for higher level summarization options and synthesizing an answer by feeding in the entire document this context and then you can see what this uh response is so now we can take a look at how these traces are actually logged in llama trace and I won't go through all the details but basically you can see that these two workflows are actually logged right here um you know we asked tell me the revenue for apple and Tesla and then also how is Apple doing generally and you can basically see that we log the entire call stack right here um you can see every single step through the workflow including both the outer agent workflow as well as the router and you can see the inputs and outputs to every step so definitely check it out especially if you sign up um and basically all your workflow steps will get logged to this tool great that's it for today um there's some more sections over here but this is basically more advanced uh and you can basically do it as an exercise um to to the viewer uh and so yeah stay tuned for more content um there will be more advanced rag sections along with additional sections on report generation and so thanks and see you next time