Context Engineering: Lessons Learned from Scaling CoCounsel
Jake Heller, founder of Casetext, shares a pragmatic framework for turning powerful large language models like GPT-4 into reliable, professional-grade products. He details a rigorous, evaluation-driven approach to prompt and context engineering, emphasizing iterative testing, the critical role of high-quality context, and advanced techniques like reinforcement fine-tuning and strategic model selection.
Jake Heller, founder of Casetext, provides a masterclass on the practical lessons learned in building Co-Counsel, an AI legal assistant powered by GPT-4. The journey began with the realization that GPT-4 represented a monumental leap over its predecessors, scoring in the 90th percentile on the bar exam, a task where GPT-3.5 only reached the 10th. This breakthrough demonstrated that the model could perform complex legal work at a human-professional level, leading Casetext to pivot its entire company and eventually be acquired by Thomson Reuters.
A Framework for Building AI Applications
Developing a robust AI application like Co-Counsel involves a structured, top-down approach that begins long before the first prompt is written.
-
Define the Customer Experience: Start by envisioning the ideal user experience. For Co-Counsel, this meant creating a suite of “skills” (e.g., legal research, contract review) that mirrored the capabilities a lawyer might list on a resume. This approach translates user needs directly into product features.
-
Model the World’s Best Expert: Deconstruct how the world’s best human expert would perform a given task. For a complex legal research query, a top lawyer would:
- Clarify the user’s request.
- Formulate dozens of potential search queries.
- Execute the searches and meticulously review every page of the results.
- Synthesize notes from relevant documents.
- Compile a final, comprehensive answer.
-
Architect the Workflow: This expert workflow becomes the application’s architecture. Determine if the process is linear (a fixed sequence of steps) or “agentic” (requiring loops and conditional logic, like reformulating searches if initial results are poor). This workflow is then broken down into micro-steps, where each step is handled by either traditional code or a specific, targeted LLM prompt.
The Rigorous Craft of Prompt Engineering
At the micro-step level, making a prompt work reliably is a demanding and iterative process. Heller outlines a disciplined methodology:
- Start with Evals: Begin by creating a small set of 10 evaluation cases. For each case, define the input (instructions and context) and the objectively correct output.
- Iterate Relentlessly: Your first prompt will likely fail on several evals. The key is persistent iteration. As Heller states, a good prompt engineer is someone “willing to not sleep for two weeks straight until they get it right.” Continuously refine the instructions, model parameters (like temperature), and context until you pass all 10 tests.
- Scale the Evals: Expand your test suite from 10 to 50, then 100, and eventually to 1,000 tests per prompt. These tests should not just be self-generated but should also include anticipated user inputs—including the “dumb” or unexpected queries that real users will inevitably try. A beta program is invaluable for sourcing these real-world test cases.
By the time you’re passing 999 out of a thousand tests, you have some degree of reliability and are ready to release to a wider audience.
The Critical Importance of Context
While the term “prompt engineering” is common, the reality is that most prompts consist of instructions + context. Often, when a model produces a poor output, the problem isn’t the model or the instructions, but the context it was given.
- Garbage In, Garbage Out: If your retrieval system provides irrelevant documents or your OCR process mangles the text, no amount of prompt tuning can save the output. Before blaming the prompt, meticulously review the exact context the model sees. “If it’s hard for you to read, it’ll be hard for the AI to read.”
- Constrain the Model: To prevent hallucinations and ensure answers are grounded in facts, use a powerful instruction:
“Don’t answer based on any of your previously existing knowledge, just based on what I’m giving you right here. If this thing says the sky is purple, the sky is damn purple.” This forces the model to act as a reasoning engine solely over the provided data.
Advanced Techniques and Practical Tricks
-
Optimize for Speed and Cost: For tasks that require processing many documents, performance is key. A useful trick is to design prompts that require only a single token response (e.g., a number from 1-10, “true/false”). You can even instruct the model to prepare a detailed explanation but use a
max_tokens=1setting or stop words to cut it off after the single-token answer is generated. This forces the model to “think” more deeply before providing its concise, fast response. -
Design for Easy Evaluation: Structure prompts so the model outputs the objective, easily verifiable part of its answer first (e.g., a numeric score or a classification), followed by the qualitative explanation. This dramatically simplifies the process of running automated evaluations.
-
Embrace Reinforcement Fine-Tuning (RFT): Unlike older methods that required thousands of examples and yielded minimal gains, modern RFT can deliver significant improvements with just 50-100 high-quality examples. At Casetext, the team is moving towards a system where every single micro-prompt has its own reinforcement fine-tuned model, squeezing out extra percentage points of accuracy.
-
Use the Right Model for the Job: There is no rule that an entire workflow must use a single model. Experiment with different models for different micro-steps. A complex reasoning task might require GPT-4 Turbo or Claude 3 Opus, while a simpler classification task might be handled perfectly well by a much cheaper and faster model like Gemini 1.5 Flash. This strategic selection optimizes both performance and cost.