Skip to main content

retainkit

CI PyPI Python License: MIT

Context and memory policies for LLM agents, scored by the evidence that survives the token budget.

An agent that has been talking to someone for thirty sessions cannot put all of it in the prompt. Something is dropped on every turn, and what is dropped decides which questions the agent can still answer. That decision is usually made by a default in a framework: keep the last N messages, summarise the rest.

retainkit measures the decision. Give it a conversation, a question, the turns that support the answer, and a budget; it reports whether those turns are still in front of the model, how much of the budget went on them, and how the answer gets worse the further back the evidence sits.

1,527 questions over ten long conversations

LoCoMo records ten conversations running 19 to 32 sessions each. Every question names the turns that answer it. An average conversation is 23,056 tokens; the evidence for the average question is 84 of them, so 0.4% of the transcript carries the answer and the other 99.6% is the search problem.

At 2,048 tokens, 8.9% of the conversation:

Policy Sees the question Recall 95% CI Whole turns
oracle (keeps only the evidence) yes 100.0% 100.0–100.0% 100.0%
fact_memory yes 62.7% 60.2–65.1% 3.9%
retrieval yes 59.5% 56.8–62.0% 59.5%
session_summary no 13.4% 11.7–15.2% 2.0%
head_tail no 8.2% 6.9–9.6% 8.2%
recency no 7.7% 6.4–9.0% 7.7%

A sliding window at this budget answers 7.7% of what it could. Reading the question first takes that to 59.5% for the same money. Put the other way: a window needs 16,384 tokens to reach the 50% that retrieval reaches at 1,024.

LoCoMo results

The right panel is where the money is. Recency holds every question whose evidence is in the session in progress and 62.3% of the ones a session or two back. Three sessions back it holds none of them, while retrieval is still at 72.1%. Long-term memory is a different problem from the window problem, and the window scores zero on it.

Full numbers, including every budget from 512 to 32,768 and a split by question type: examples/locomo/results/locomo.md.

Two results that do not flatter the tooling

Multi-hop questions are not solved by retrieval. At 2,048 tokens the best policy holds the evidence for 76.8% of single-hop questions and 16.2% of multi-hop ones. A multi-hop question needs several turns from several sessions at once, and one query against a store tends to bring back one of them. Anyone reporting a single memory-recall number is averaging over this gap.

fact_memory wins on a technicality worth knowing about. It stores sentences rather than turns, so it fits more of the conversation into the same budget and scores 62.7% against retrieval's 59.5%. But the whole evidence turn survives only 3.9% of the time; the rest of its wins hand the model one sentence out of several. Whichever column you take as the real one, take the same column for every policy.

What is being measured

Whether a named span survived, and nothing else. A model handed the evidence can still answer wrongly, so every number here is a ceiling on what any model behind the policy could do, not a score for the model. The point of a ceiling is that it is cheap: no API calls, no judge, no variance between runs, and a policy that loses the evidence is not worth benchmarking end to end.

The same convention rules out one thing: an abstractive summariser rewrites what it keeps, and a rewritten sentence cannot be traced to a turn. The summary policy here is extractive so that survival stays answerable. Token counts come from a four-characters-per-token estimate, which moves absolute numbers around by up to 22 points for a window policy, so pass your own tokenizer if the budget has to be exact. It does not move which policies win, and there is a test holding it to that.

Install

pip install retainkit

No runtime dependencies. pip install "retainkit[viz]" adds matplotlib for the figures.

Use

from retainkit import Case, Probe, Recency, Retrieval, Transcript, Turn, run, sweep_table

case = Case(
    transcript=Transcript(
        id="guest-42",
        turns=(
            Turn(id="s1:1", session=1, speaker="guest", text="I never eat shellfish."),
            Turn(id="s1:2", session=1, speaker="agent", text="Noted, thank you."),
            Turn(id="s9:7", session=9, speaker="guest", text="Book me somewhere near the water."),
        ),
    ),
    probes=(Probe(id="q1", question="What can this guest not eat?", evidence=("s1:1",)),),
)

print(sweep_table(run([case], (Recency(), Retrieval()), budgets=(64, 512))))

Your own conversations go in as JSON and out as a report:

retainkit eval conversations.json --budget 2048 --budget 8192 --out results/

The file format is in src/retainkit/io.py. retainkit demo runs a seeded fixture and needs no data at all.

The policies

Policy Keeps Reads the question
recency The newest turns that fit. no
head_tail The opening of the conversation and the newest turns. no
session_summary The live session, plus the lead sentences of each closed one. no
retrieval The turns that rank highest against the question, put back in order. yes
fact_memory The live session, plus sentences pulled from a long-term store. yes
oracle The evidence and nothing else. Not runnable in production. yes

A policy is anything with a name, a query_aware flag and a select method, so your own goes straight into the same sweep:

class KeepEverySecondTurn:
    name = "every_second"
    query_aware = False

    def select(self, transcript, probe, budget, tokenizer):
        from retainkit.policies import fill, package
        from retainkit.policies.base import turn_fragments

        kept, _ = fill(turn_fragments(transcript)[::2], budget, tokenizer)
        return package(self.name, budget, kept, tokenizer)

Reproducing

make install
make test
make demo                      # the seeded fixture, written to results/
python examples/locomo/fetch_data.py
python examples/locomo/run_locomo.py

LoCoMo is fetched at run time and is not redistributed here; it carries the terms of its own release. The fixture results in results/ are rebuilt by CI on every push and the build fails if a committed number no longer matches the code, so nothing quoted above can drift away from what the package produces.

Citing

@software{mohammadi_retainkit,
  author  = {Mohammadi, Hadi},
  title   = {retainkit: context and memory policies for LLM agents,
             scored by the evidence that survives the token budget},
  year    = {2026},
  url     = {https://github.com/mohammadi-hadi/retainkit},
  license = {MIT}
}

LoCoMo is by Maharana et al., Evaluating Very Long-Term Conversational Memory of LLM Agents, ACL 2024. Cite it if you use these results.

judgekit audits an LLM judge, arenakit audits a pairwise leaderboard, trajectory-judge scores agent trajectories. More at mohammadi.cv.

Metadata

Release files for retainkit 0.1.0

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for retainkit 0.1.0
File Size Uploaded Provenance Actions
retainkit-0.1.0.tar.gz 162.4 kB With provenance Details for retainkit-0.1.0.tar.gz

Built distribution (wheel)

Table of built distributions (wheels) for retainkit 0.1.0
File Interpreter ABI Platform Actions
retainkit-0.1.0-py3-none-any.whl Python 3 none any Details for retainkit-0.1.0-py3-none-any.whl

Total release size: 188.0 kB

Release files / retainkit-0.1.0.tar.gz

Download URL retainkit-0.1.0.tar.gz
Size 162.4 kB
Tags Source
SHA-256 checksum
How to use checksums
c372edc64144e5aa1ba4ffa922d9ae501f5dd4514572967b8bf1f471193ee099
BLAKE2b-256 checksum
How to use checksums
b4921421b39e01910ba316068adc49f4a9d188571ca2eabeeb5e03f5254dd0f3
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
Yes
Uploaded via twine/7.0.0 CPython/3.13.14

Provenance

Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.

PyPI Publish Attestation

PyPI verified that this artifact, at this checksum, originated from the publisher listed below.

Signed by GitHub Actions, verified by PyPI on Sep 5, 2026.

Transparency log

Release files / retainkit-0.1.0-py3-none-any.whl

Download URL retainkit-0.1.0-py3-none-any.whl
Size 25.6 kB
Tags Python 3
SHA-256 checksum
How to use checksums
6592553f3ee733606b87c08206ca1502adba47a5763125946a9761e011b2c04c
BLAKE2b-256 checksum
How to use checksums
cd1998ef558bd56268a137bee8bc6f28f72dd7f81637cdfb14a977b66dbad652
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
Yes
Uploaded via twine/7.0.0 CPython/3.13.14

Provenance

Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.

PyPI Publish Attestation

PyPI verified that this artifact, at this checksum, originated from the publisher listed below.

Signed by GitHub Actions, verified by PyPI on Sep 5, 2026.

Transparency log

Security for retainkit 0.1.0

Provenance

Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.

Full provenance

Both files in this release have been uploaded with provenance metadata.

Select from 2 files

Attestation for retainkit-0.1.0.tar.gz

PyPI Publish Attestation

PyPI verified that this artifact, at this checksum, originated from the publisher listed below.

Signed by GitHub Actions, verified by PyPI on Sep 5, 2026.

Transparency log

A note on package security

PyPI is a repository for distributing packages, not a service for auditing them. While third-party security researchers and our staff do their best to detect and remove malicious packages, PyPI cannot guarantee the quality, security or suitability of any given package.

To learn more about protecting yourself from malicious actors, visit our package security documentation.

Release history Release notifications | RSS feed

This release

0.1.0 This release

2 release files Full provenance

Anthropic, PBC Visionary sponsor Bloomberg Visionary sponsor Hudson River Trading Visionary sponsor Meta Visionary sponsor NVIDIA Visionary sponsor Microsoft Sustainability sponsor Depot Continuous Integration AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page