An LLM Exercise on Retrieval: Essential of Context Engineering

The emerging discourse on context engineering has captured everyone’s attention in the AI world. This concept, and practice captures the way that engineers now structure information that they put in the context window of an LLM. Different from prompt engineering, which describes how the engineer structures each prompt to make the LLM to create a certain output. Now with the increasing capability of LLM such as using tools, thinking, etc. information being funneled into an LLM is a lot more structured, and requires a lot more system design.

Andrej Karpathy describes how context engineering is both an art and a science for creating industrial-strength LLM app [emphasis mine]. First and foremost, it’s a science because there is a structure around what information one should put into the context window such as examples, tools, chain of thought prompting, etc. This is all engineering practices that we have learned in the past three years (since 2022 when chatGPT was launched). We know that certain prompt structures work better now. And in the age of agentic AI, LLM can plan, take actions given the right tools, right functions, context engineering describes an entire practice where what kind of tools being used might determine whether an agentic system succeed or not. Second, it is an art because “of the guiding intuition around LLM psychology of people spirits.“ This second part might need further unpacking. My reading of this phrase “the guiding intuition around LLM psychology of people spirits“ means each LLM has its own quirks, and by mimicking human-generated data, it has internalized certain human behaviors. Thus, the art part of context engineering lies in the idea that LLMs are not deterministic, and have behaviors that are unpredictable, and need experience to get at. In other words, context engineering is “experiential,“ one that needs a lot of tinkering to get right. Understanding the fundamentals themselves won’t replace building, and tinkering with different LLMs in order to create an industrial-grade LLM app.

In many ways, whatever we’re doing now while creating an LLM app is one way or another context engineering.

In creating the AI exercise series, I am mindful of this emerging discourse, and would like to include as many context engineering exercises as I can. The most important exercise regarding context engineering in my opinion is retrieval, and retrieval augmented generation (RAG). RAG extends the long term memory of the LLM app, keeps the LLM up to date with the most relevant events. In AI engineering, Chip Huyen argues that before considering fine tuning an LLM, one should experiment with prompting, then adding RAG. If those two fail, fine-tuning might be the next option. RAG helps the model focus its response to a specific piece of information, thus arguably increase the specificity and accuracy of the final output.

RAG is particularly helpful in domains where LLMs have little information on such as closed source data such as medical records, law cases, etc. I created such an exercise for the legal domain which is ripe for innovation, and is a place where industrial grade LLM apps, and agentic systems might make the most impacts because the field relies so much on structured documents.

Following is the description of the problem, which can be found on my github: https://github.com/ngathan/ai-exercises/tree/main/01_llm_retrieval

Problem Statement: For domain specific problems such as answering complicated legal questions, using pre-trained LLMs as they are might not suffice. In such scenarios, retrieving the right documents, then add them to the context, and ask the same LLM to see how the accuracy might change. Once it is determined that using the right document, the system might do better than using the pretrained LLM alone, we then set up a retrieval augmented generation system, and use an LLM to answer the original question.

The dataset to be used in this exercise is housing_qa from Stanford Regulation, Evaluation and Governance Lab. You can read more about the dataset on Hugging Face, and the original paper A Reasoning-Focused Legal Retrieval Benchmark by Zheng et al.

This exercise helps you set up a simple RAG system.

Step 1: Evaluate when an LLM answer legal questions alone

Step 2: Evaluate when an LLM answer legal questions with the right document

Step 3: Evaluate when the RAG system where LLM answer legal questions with the documents retrieved using a retrieval system

At each step, there is an accuracy number associated with that particular step. Remember to calculate these accuracies along the way, and see if at the end they follow your expectations.

Hints:

Hint 1: When building the pipeline experimenting with the first 20, 50, 100, 200 data points in the dataset to see whether the pipeline works as expected Hint 2 (Advanced level): Beyond TFIDF and ngram in the problem statement, you can experiment with more advanced embedding methodology such as embedding models from OpenAI, Anthropic.
Hint 3 (Advanced level): Implement an index using FAISS to store all the statutes to speed up the cosine similarities calculation for the entire dataset

RAG is bread and butter of context engineering. Thus understanding how retrieval works is essential for AI engineers’ daily practice. This exercise is to help you visualize and concretize how RAG works in the context of legal documents. The proposed solution is one of many solutions. If you come up with better solutions (faster, higher accuracy), please send pull and merge requests.