Lengthy Context Isn’t Free — I Constructed a Secure Immediate-Pruning Layer That Makes LLM Methods Work

I’ve labored on, dialog state tends to develop shortly over time. It’s widespread to resend massive parts of the historical past on every flip—together with older software outputs, repeated RAG retrievals, and context that’s not related. As this accumulates, prompts can develop into considerably bigger, which can improve inference price and latency, and in some instances have an effect on reasoning efficiency.

I constructed a deterministic pipeline that prunes this redundant state earlier than the immediate ever reaches the mannequin. The model I applied avoids LLM calls, embeddings, and exterior dependencies. Relying strictly on customary library elements ensures each pruning choice stays totally deterministic and reproducible.

State monitoring executes in three distinct passes: Expired Context Elimination, Duplicate Context Elimination, and Dependency Restoration. The third move is what helps make the primary two safer in observe. It ensures that nothing a later message is dependent upon is unintentionally eliminated.

Whereas constructing it, I bumped into two bugs that modified the design. My first benchmark corpus used a hard and fast variety of duplicates and rancid software calls, which made discount percentages shrink as conversations grew. That didn’t mirror what I’d anticipate from real-world habits.

My dependency restoration logic additionally went fully untested at first as a result of my artificial knowledge by no means created a case the place a required message was truly eliminated. Each points are lined right here, together with how fixing them modified the outcomes.

After correcting the pipeline, I benchmarked it throughout three workloads: plain chat, a RAG assistant, and a tool-heavy agent. Every was examined at 5 dialog sizes, for a complete of 15 configurations on two completely different machines. Throughout these runs, all labeled required info had been preserved. The system additionally reached a secure mounted level after a single move, which implies pruning an already pruned immediate produces no additional adjustments.

Token discount is dependent upon the workload. It’s about 2 to 4 p.c for plain chat, 27 to 32 p.c for a RAG assistant, and 33 to 34 p.c for a tool-heavy agent. Even at 2,000 turns and 131,000 tokens, preprocessing stayed below 50 milliseconds.

Full code, all 35 checks, and the uncooked terminal output are included beneath so you may run the pipeline your self.

Full code: https://github.com/Emmimal/prompt-pruning-layer/

Each immediate I ship retains getting heavier

I stored seeing the very same sample pop up throughout each long-running agent I constructed. A dialog begins out completely clear. Fifty turns later, it’s an absolute mess.

By that fiftieth flip, the immediate payload you’re delivery out on each single request contains the system immediate, your complete chat historical past, 4 software outputs (two of that are fully stale as a result of the software ran twice), six retrieved chunks (three are near-duplicates as a result of the consumer circled again to an previous subject), a SQL consequence from twenty turns in the past that’s simply sitting there, and a single consumer choice acknowledged as soon as and by no means introduced up once more.

None of these things is technically mistaken. Each single piece made sense the precise second it was injected. The true difficulty is that nothing ever will get cleared out. The immediate turns into an append-only log of each historic occasion, and also you’re dumping the entire thing onto the mannequin, each single flip, indefinitely.

That’s not only a storage downside. Reasoning efficiency measurably degrades as enter size grows, even when the added content material is irrelevant to the duty [1], and fashions are worse at utilizing data buried in the midst of an extended context than data close to the perimeters [2].

The speedy knee-jerk response whenever you see this bloating is simply fundamental truncation. Slice the final N messages and drop every little thing else. It’s actually one line of code. The catch is it silently breaks your dialog chains in methods you gained’t even notice till it blows up on an actual consumer:

Flip 3:   Consumer: My most popular output format is CSV.
Flip 4:   Assistant: Bought it.
...
Flip 47:  Consumer: Export the outcomes.

In case your window is capped on the final 20 turns, flip 3 disappears by flip 47. The assistant loses the context that the consumer requested for a CSV. That knowledge was not stale; it was a tough dependency, and easy positional truncation can not differentiate between previous, expendable context and previous context {that a} later flip nonetheless depends on.

That’s the precise design constraint this mission addresses. Any mechanism that removes redundant state should distinguish between these two classes. Recency alone is inadequate.

Borrowing an concept from working programs

Right here is the reframe I constructed the remainder of this round: an working system is consistently deciding which pages keep resident in RAM and which get evicted. A protracted-running LLM dialog has the very same downside, besides nothing performs the position of the reminiscence supervisor. Context simply accumulates perpetually as a result of no course of owns the job of deciding what stops incomes its place within the immediate.

The Immediate Pruner on this article is that lacking piece.

Every bit of dialog state (a consumer flip, an assistant flip, a software output, a retrieved chunk) is a Message object with a task, a flip quantity, and a little bit of bookkeeping metadata:

@dataclass
class Message:
    id: str
    position: str
    content material: str
    flip: int
    tool_call_key: Non-compulsory[str] = None
    expires_after_turn: Non-compulsory[int] = None
    defines_keys: checklist = area(default_factory=checklist)

Three passes run over that checklist so as. Each is a pure perform: messages in, a filtered checklist out. No mannequin is concerned wherever within the pruning step itself.

Why there’s no mannequin contained in the pruner itself

I might have used an embedding mannequin to attain message relevance, which might have made this fuzzier and quite a bit simpler to put in writing. I didn’t, and it isn’t for the sake of purity.

As soon as a pruning choice is dependent upon a mannequin’s judgment, you lose the power to purpose about what a dialog will appear to be on the following flip. The identical enter not ensures the identical output. A pruning layer meant to make a manufacturing system extra predictable shouldn’t be the least predictable a part of it. All the things right here runs on dataclasses, regex, and dict lookups, which is strictly the toolset this downside wants.

Cross 1: Expired Context Elimination

If a software will get referred to as greater than as soon as below the identical key (the identical search question, the identical SQL lookup, the identical file learn), solely the most recent consequence continues to be reliable. All the things earlier below that key’s expired.

def _pass1_expired_context_elimination(self, messages):
    last_occurrence = {}
    for m in messages:
        if m.tool_call_key:
            last_occurrence[m.tool_call_key] = m.id

    stored, eliminated = [], []
    for m in messages:
        if m.tool_call_key and last_occurrence[m.tool_call_key] != m.id:
            eliminated.append(m)
        else:
            stored.append(m)
    return stored, eliminated

Cross 2: Duplicate Context Elimination

Retrieval pipelines continually pull up an identical or near-duplicate passages, particularly when a consumer loops again to an earlier subject. This move normalizes the whitespace and casing, retains solely the primary prevalence, and drops each duplicate after it.

def _pass2_duplicate_context_elimination(self, messages):
    seen, stored, eliminated = {}, [], []
    for m in messages:
        if m.position == ROLE_RETRIEVED_DOC:
            norm = " ".be part of(m.content material.decrease().cut up())
            if norm in seen:
                eliminated.append(m)
                proceed
            seen[norm] = m.id
        stored.append(m)
    return stored, eliminated

Cross 3: Dependency Restoration (and the bug that made me construct it correctly)

This move is the explanation the primary two are protected to run in any respect. If Cross 1 or Cross 2 drops a message that occurs to be the one place a still-referenced reality acquired outlined, this move catches it and places it again.

The mechanism is deliberately easy: a message marks itself with a literal DEFINE: tag, and a later message references it utilizing REF:. If a REF survives into the ultimate stored set however its matching DEFINE acquired dropped upstream, Dependency Restoration restores that lacking message to the checklist.

def _pass3_dependency_restoration(self, all_messages, kept_messages, removed_messages):
    kept_ids = {m.id for m in kept_messages}
    by_id = {m.id: m for m in all_messages}

    key_definer = {}
    for m in all_messages:
        for key in m.defines_keys:
            key_definer[key] = m.id

    referenced_keys = set()
    for m in kept_messages:
        referenced_keys.replace(m.references())

    restored = []
    for key in referenced_keys:
        definer_id = key_definer.get(key)
        if definer_id and definer_id not in kept_ids:
            restored_msg = by_id[definer_id]
            kept_messages.append(restored_msg)
            kept_ids.add(definer_id)
            restored.append(restored_msg)

    kept_messages.kind(key=lambda m: (m.flip, m.id))
    return kept_messages, restored

Right here is the bug. I ran my first full benchmark throughout all three workloads and 5 sizes, and each single row printed “Restored (deps): 0.” Each one. My first response was that this seemed nice—an ideal security report. It wasn’t. It meant Cross 3 had by no means truly restored a single message on any run at any dimension.

I went again into my corpus generator to search out out why, and the reply was embarrassing as soon as I noticed it. My artificial conversations solely ever connected DEFINE markers to plain consumer messages, whereas Cross 1 and Cross 2 solely ever take away software outputs and retrieved paperwork. The 2 classes by no means overlapped. Dependency Restoration was sitting within the codebase, totally written, however fully untested by my very own benchmark as a result of nothing I generated ever gave it a purpose to fireside.

The repair was to let some software outputs additionally outline a dependency, the identical method an actual “get consumer settings” software name would possibly floor a reality the dialog is dependent upon later, though that precise software name might get outmoded and marked expired by Cross 1. As soon as I added that, the numbers modified instantly: 2 restorations on the smallest tool-agent dialog, climbing to 127 on the largest. Required info stayed at 100% preserved your complete time, however now that quantity truly meant one thing, as a result of the move being examined was able to failing and didn’t.

I wish to be direct in regards to the limitation that’s nonetheless right here even after the repair: dependency detection is literal identifier matching, not understanding. It catches a precise REF matching a precise DEFINE. It is not going to catch a consumer paraphrasing, asking “what format did I point out earlier” with no matching tag. A semantic dependency resolver would want an embedding mannequin or an LLM name, and that’s explicitly exterior what this deterministic pipeline does. I’d quite ship a narrower assure I can show than a broader one I can’t.

Design objectives, and why each exists

Deterministic. Similar enter at all times yields the identical output. Preserving the mannequin out of the loop eliminates run-to-run variance and any danger of hallucinating what ought to stay within the immediate.

Dependency-safe. It by no means silently drops a reality {that a} later flip nonetheless wants. Positional truncation fully lacks this property, which makes it non-negotiable right here. A pruner saving 40 p.c of tokens that often breaks a dialog is a worse trade-off than one saving 4 p.c that by no means breaks something.

Idempotent. Operating the pruner a second time does nothing. If that’s not true, you can’t safely re-prune on each single flip of a rising dialog with out risking compounding drift.

Light-weight. The pruning step ought to by no means develop into the precise bottleneck it was constructed to remove.

The benchmark: three workloads, and a mistake I virtually shipped

My first model of the artificial corpus generator picked a hard and fast variety of duplicate passages and repeated software calls (six software calls and eight duplicates) no matter how lengthy the dialog was. I ran it at 5 sizes and the discount proportion went down because the dialog acquired longer: 9.9 p.c at 50 turns, dropping to 0.3 p.c at 2,000 turns. That runs backward from precise manufacturing visitors, and it isn’t defensible. If the quantity of waste in a benchmark is only a fixed I hand-picked, anybody studying it’s proper to ask whether or not I constructed the benchmark to show the algorithm works.

So I threw that generator out and rebuilt it round an specific workload mannequin as a substitute: retrieval-per-turn, retrieval overlap charge, software name charge, and power repetition charge. All of those had been mounted earlier than operating a single benchmark, primarily based on what appeared like believable manufacturing habits, quite than adjusted afterward to hit a quantity I favored. Three workloads got here out of that:

Regular chat. No retrieval, occasional software calls, largely atypical back-and-forth.

RAG assistant. Retrieves paperwork on each flip, with an actual probability any given passage overlaps one thing retrieved just lately, as a result of customers revisit matters and retrieval re-surfaces the identical chunks.

Device agent. Frequent calls throughout 5 software varieties (search, SQL, calculator, filesystem, net fetch), excessive repetition charge, modeling one thing that re-plans and re-queries continually.

Each artificial corpus additionally ships with floor reality: each message some later message is dependent upon will get labeled required up entrance. So “did pruning maintain every little thing it wanted to” is a test in opposition to recognized labels, not a guess.

Right here is the entire output. All 15 configurations, not a slice of it:

Each single row says Information stored: Sure and Idempotent: Sure. Not most rows. All fifteen.

The sample by workload is smart when you take a look at the place the waste truly comes from. Regular chat barely retrieves something and infrequently repeats a software name, so there’s virtually nothing for Cross 1 or Cross 2 to catch; it stays round 4 p.c regardless of how lengthy the dialog runs. The RAG assistant retrieves each flip with actual overlap, so Duplicate Context Elimination carries a lot of the weight, touchdown round 32 p.c. The Device agent combines each issues (frequent software repetition and retrieval overlap) and hits the very best discount at 33 to 34 p.c.

Totally different workloads accumulate completely different sorts of waste. The pruner responds on to no matter waste is sitting in entrance of it, quite than producing a flat quantity that might recommend the benchmark was reverse-engineered to hit a goal.

If I solely acquired to maintain one consequence from this entire benchmark, it’s the security property: 15 out of 15 configurations preserved 100% of required info. Zero lacking dependencies, throughout three structurally completely different workloads and a 40x vary in dialog size. Truncation by place can not provide that. For me, this was crucial sign when evaluating whether or not the strategy is perhaps usable in manufacturing.

That quantity can be computed, not hand-counted. The benchmark script itself tallies how most of the 15 (workload, dimension) pairs preserved each required reality and what number of reached the idempotent mounted level, printing an combination abstract block after the 15 detailed runs:

============================================================
SUMMARY
============================================================
Configurations run:               15 (3 workloads x 5 sizes)
Required info preserved:         15/15
Reached mounted level (idempotent): 15/15

Workload            Token discount vary    Information preserved   Idempotent
Regular chat          1.9-4.1%                YES               YES
RAG assistant        27.0-32.5%               YES               YES
Device agent           33.1-33.7%               YES               YES

I wrote this test as a result of I caught myself hand-counting desk rows for an early draft. That metric belongs within the script output, not my very own eyes squinting at a terminal log. If a take a look at run ever hits something lower than 15/15, it means I broke the pruner and have a regression to search out, not a typo to edit within the submit.

I left one particular metric out of the ultimate numbers: tokens eliminated per millisecond of execution overhead. The code computes it—it peaks at round 4,000 tokens per millisecond on small tool-agent runs and lands between 900 and 1,700 tokens at bigger scales. It stays within the codebase as an inner area as a result of it helps observe scaling prices, nevertheless it belongs exterior the primary desk. Readers can not act on it the way in which they will with uncooked token depend, discount proportion, or millisecond overhead. Three direct metrics exhibiting a transparent trade-off are higher than a fourth that acts as a novelty.

The idempotence result’s the half I favored monitoring essentially the most. Proving prune(prune(x)) == prune(x) means the pipeline hits a secure mounted level on the primary move. Operating it once more on an already-pruned immediate adjustments nothing:

That guidelines out oscillation. It additionally guidelines out cumulative shrinkage throughout turns in the event you re-pruner on each single message of a rising dialog, which is strictly how this runs in manufacturing.

Reproducing it on a second machine

I ran the total benchmark on a Linux container operating Python 3.12.3, then once more on Home windows 11 in PyCharm utilizing Python 3.12 in a separate venv. Each token depend and message depend matched precisely throughout each machines. Solely the millisecond timings moved, which is customary for various {hardware}, and even these stayed below 50 milliseconds for the biggest, most tool-heavy dialog on each setups.

One factor I seen whereas evaluating the 2 runs and wish to be straight about: prompt-build time (the time to serialize the ultimate message checklist right into a string) often got here out slower after pruning than earlier than on the Regular Chat workload. Not by a lot—below a millisecond—however the course was backward.

My learn is that Regular Chat solely removes 2 to 4 p.c of messages, so the earlier than and after checklist sizes are practically an identical. At sub-two-millisecond operations, system jitter and rubbish assortment pauses simply swamp the precise sign. Including a warm-up name and utilizing the median of 30 runs largely stabilized the metric, however the anomaly nonetheless pops up on that workload. It’s noise at a scale the place the delta is smaller than the measurement error, so I left it uncooked quite than massaging the information.

What this benchmark intentionally doesn’t measure

This benchmark isolates token discount and pruning overhead as a result of these are the metrics the pipeline truly controls. Finish-to-end LLM latency is a totally separate variable. It is dependent upon supplier structure, batching, regional caching, and community situations that this mission can not see. Making an attempt to transform a token discount proportion immediately right into a latency delta means inventing an arbitrary conversion fixed.

The baseline actuality is easy: slicing 30 to 34 p.c of enter tokens means the mannequin does much less work per name. Usually, inference price and latency have a tendency to extend with immediate dimension [4], making this a helpful price lever. However a real latency quantity requires a dwell validation move in opposition to your particular supplier. Publishing a generic latency determine right here would imply making claims about infrastructure I don’t management, quite than evaluating the pruning layer itself.

How this suits into an actual agent loop

The pruner lives in precisely one spot: proper after the dialog historical past is pulled collectively for a flip, and simply earlier than it will get serialized into the ultimate immediate string.

from prompt_pruning import PromptPruner, PromptBuilder

pruner = PromptPruner()
builder = PromptBuilder()

def handle_turn(conversation_state, new_user_message):
    conversation_state.append(new_user_message)

    pruned_messages, report = pruner.prune(conversation_state)
    immediate = builder.construct(pruned_messages)
    response = call_llm(immediate)

    conversation_state.append(make_assistant_message(response))
    return response

As a result of the pipeline is idempotent, calling prune() each flip of a rising dialog is protected. Operating the pruner ten occasions on a historical past pruned 9 occasions yields the very same consequence as a clear run from scratch. This makes “run it each flip” a protected default, eliminating the necessity to observe state or purpose about earlier passes.

The one integration choice left is how REF and DEFINE tags get connected to messages. Right here they’re literal markers contained in the message content material, which is the only mechanism for a prototype. A manufacturing system would doubtless connect them as structured metadata on the message object so the tags by no means leak into the uncooked textual content the mannequin reads. Cross 3’s logic stays the identical both method. An upstream course of nonetheless has to find out what counts as a dependency value tagging, as a result of Cross 3 can solely restore what it’s explicitly instructed to trace.

What this doesn’t cowl

Dependency detection is literal, not semantic. If a reference is paraphrased and lacks an identical tag, the script will miss it.

These workloads are additionally fully artificial. I selected the three parameter units primarily based on believable manufacturing habits, not actual telemetry. When you have manufacturing logs, regenerating these three workload classes from precise utilization is the apparent subsequent step. The numbers will shift relying on what your precise visitors seems to be like.

This pipeline omits semantic compression, embeddings, and LLM-scored pruning. These are legitimate, various approaches, however avoiding them retains this implementation totally deterministic and dependency-free. LLMLingua is a chief instance of the discovered various; it makes use of a small language mannequin to attain and drop tokens, reaching a lot increased compression ratios than this script [3]. Selecting between them is a direct trade-off: you alternate determinism and zero-dependency execution for tighter compression.

The token counts are additionally approximations. The script makes use of a whitespace and punctuation-boundary heuristic as a substitute of a manufacturing subword tokenizer like tiktoken. As a result of the heuristic runs constantly earlier than and after pruning, the relative discount percentages stay correct, even when absolutely the numbers don’t completely match an official tokenizer.

Lastly, there is no such thing as a direct latency measurement for the infrastructure causes detailed earlier.

The place I’d take this subsequent

Two extensions make sense right here quite than increasing the present three passes.

The primary is a hybrid strategy. Hold these three deterministic passes as a quick, protected first stage, then hand the output to an embedding-aware or LLM-scored compression software like LLMLingua. This catches semantic redundancy that literal identifier matching misses, similar to two passages saying the identical factor in several phrases.

Operating the deterministic passes first preserves the protection assure. If the discovered stage misbehaves, the dialog drops again to the deterministic baseline as a substitute of failing unpredictably. Architecturally, the discovered move solely operates on the output of Cross 3. A mannequin bug within the compression step would possibly shrink a immediate too aggressively, nevertheless it can not reintroduce a dependency failure that the deterministic passes already cleared.

The second extension is closing the hole between artificial knowledge and manufacturing actuality. The following step is regenerating these similar three workload classes from precise manufacturing traces. Preserving the benchmark methodology an identical—the identical mounted parameters, ground-truth dependency labels, and 15-configuration sweep—ensures the outcomes stay immediately comparable to those printed figures whereas changing guesses with actual telemetry.

Precise utilization logs will doubtless shift the RAG and tool-agent metrics in both course. Actual-world visitors doesn’t robotically imply increased compression. For example, a manufacturing system with aggressive upstream deduplication would possibly already remove the waste this artificial mannequin assumes. That will be a extremely helpful discovering in its personal proper, quite than a failure.

A 3rd, smaller level to notice: the REF/DEFINE conference used right here is only a placeholder. A manufacturing system ought to derive these tags robotically from structured knowledge, software name arguments, session variables, or specific consumer settings, quite than counting on handbook textual content markers.

Whereas the deterministic logic in Cross 3 stays an identical both method, the precise worth of this pipeline relies upon solely on how cleanly you may generate correct dependency tags upstream. That’s an architecture-specific integration downside quite than one thing a general-purpose pruning library can remedy out of the field.

Manufacturing serving programs already deal with immediate bloat as a reminiscence administration downside on the infrastructure layer, evicting and sharing KV cache pages very similar to an working system handles bodily reminiscence [4]. You may consider this mission as working one layer above that, on the immediate building part, earlier than the request ever reaches the infrastructure.

The 2 layers don’t compete. A smaller, deduplicated immediate supplies a greater enter to a well-managed cache quite than performing as a alternative for it.

The Core Takeaway

If a vital reality is preserved, the pipeline must exhibit that retention quite than assume it primarily based on the absence of apparent errors. This precept guided the design of the system.

Lengthy-running conversations don’t at all times require a bigger, smarter mannequin to determine what to recollect. Usually, the extra sturdy resolution is a predictable, three-pass system that may programmatically show what it didn’t lose.

Sources

[1] Levy, M., Jacoby, A., & Goldberg, Y. (2024). Similar activity, extra tokens: The influence of enter size on the reasoning efficiency of enormous language fashions. In Proceedings of the 62nd Annual Assembly of the Affiliation for Computational Linguistics (Quantity 1: Lengthy Papers) (pp. 15339–15353). Affiliation for Computational Linguistics. https://doi.org/10.18653/v1/2024.acl-long.818

[2] Liu, N. F., Lin, Ok., Hewitt, J., Paranjape, A., Bevilacqua, M., Petroni, F., & Liang, P. (2024). Misplaced within the center: How language fashions use lengthy contexts. Transactions of the Affiliation for Computational Linguistics, 12, 157–173. https://doi.org/10.1162/tacl_a_00638

[3] Jiang, H., Wu, Q., Lin, C.-Y., Yang, Y., & Qiu, L. (2023). LLMLingua: Compressing prompts for accelerated inference of enormous language fashions. In Proceedings of the 2023 Convention on Empirical Strategies in Pure Language Processing (pp. 13358–13376). Affiliation for Computational Linguistics. https://doi.org/10.18653/v1/2023.emnlp-main.825

[4] Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J. E., Zhang, H., & Stoica, I. (2023). Environment friendly reminiscence administration for giant language mannequin serving with PagedAttention. In Proceedings of the twenty ninth Symposium on Working Methods Rules (SOSP ’23) (pp. 611–626). Affiliation for Computing Equipment. https://doi.org/10.1145/3600006.3613165

All code, benchmark numbers, and take a look at outcomes on this article are my very own, generated by operating the included codebase immediately and reproduced on two separate machines. No proprietary datasets, copyrighted textual content, or third-party code had been utilized in constructing or benchmarking this technique. All code, the corpus generator, the pruner, the benchmark harness, and all 35 checks, is offered within the repository linked beneath.

https://github.com/Emmimal/prompt-pruning-layer