Context Engineering for AI

In this demo you’ll apply one of the context engineering techniques from the previous section, context pruning, using a LangGraph agent.

Context pruning simply means removing irrelevant information from a context. The way it’ll be implemented in this demo is similar to how you’d do it for many of the other processes. Note that every model behaves differently. The way a context could degrade over time with GPT-4 may be different from Gemini-2.5. For the best outcome, always assess a model’s behavior for your particular use case.

Open the context-pruning.ipynb file from the starter project in Visual Studio Code or any other editor that can open Python notebooks. cd into the starter project directory in the Terminal, and run uv venv to create a virtual environment.

The prerequisites for this notebook are in the first cell. Copy the pip command to install the dependencies. At the top right corner of VS code, select the .venv interpreter as the Python interpreter for your notebook.

Then, get an OpenAI API key for at least $5. Without topping up your credits, you won’t be able to use the OpenAI models. If you already have some performant models stored locally or have credits for some other model in the cloud, you can easily swap the OpenAI models here with those. Just make sure that the model you use for the tool calling is capable and compatible with the LangChain agent framework.

Back in the terminal, set the OpenAI API key to the OPENAI_API_KEY environment variable:

export OPENAI_API_KEY='insert-open-ai-api-key'

Or execute the code in the second cell to set the key for the current session. Make sure to use this exact variable name, since that’s what the OpenAI package will look for by default.

All set. Here’s what’s happening in the remaining cells.

In the next code cell, you prepare a list of extra data from the URLs. You’ll use these to provide extra data to your RAG which forms part of your AI agent. Run the cell to use the WebBaseLoader to retrieve all the data in the given URLs.

In the next cell, you break all the data from the previous cell into smaller manageable chunks, before storing them in memory in the next cell. You create a retriever to have access to this data for use by an LLM. Run the cell.

A vector store is a kind of storage that stores data in a format models understand. Likewise the retriever is a class that’s able to access, search and return such data. Execute this cell to embed the data and create the retriever.

In the next cell is a quick test to see whether the retriever is set up properly. It does so by creating a retriever tool based on the retriever object from the previous cell, together with a name for the tool and a prompt. Run it and monitor the output. It prints a portion of the response based on the prompt. The name for the retriever tool is key to helping the LLM use the right tool for the right job, always give it a good descriptive name.

Now, go to the next cell. Here, you bind the retriever tool with the gpt-4.1-mini LLM. This means this model can be used to call the retrieve_blog_posts tool when the agent needs information from those blog posts. Execute the cell.

With all the housekeeping stuff out of the way, it’s time to set up the agent in the next cell. Go to the next cell.

Here, the state is effectively the context window that contains everything for that context. You define an llm_call function which decides whether to call a tool or provide an answer. You create a prompt for your agent to act as a research assistant.

Then comes the actual prompt for pruning the context.

The should_continue function decides whether to end the cycle or keep getting more information from the LLM. It quits when the LLM doesn’t request for a tool call, suggesting that the LLM has enough information to return an answer. Otherwise, it makes a tool call.

To help you understand what’s happening, run the cell to see a graph of the flow within the various agent’s components.

Finally, in the last cell, you provide the prompt to the AI agent that sets the agent off to work. Run the cell. After a while, you’ll see the output for the prompt. To better appreciate how good this response is, use the same prompt but in a basic RAG, and compare the output. This is cleaner.

That’s all for this demo. Use the notebook in the final project directory as a reference as you practice. Continue to see a few other techniques used in context engineering.