DEV Community

Siddhi Prajapati
Siddhi Prajapati

Posted on

IMMUNE: Deterministic Memory Attribution for AI Agents Using Delta Debugging

#ai #python #machinelearning #opensource

TL;DR

I built IMMUNE because I hit a wall: my AI agent kept giving wrong answers and I had no idea which stored memory was causing it. IMMUNE applies delta debugging — a classical algorithm from software testing — to pinpoint the exact memory responsible. No LLM judge, no guessing. It is a drop-in for LangGraph and quarantines bad memories at the Redis index level. This post explains how it works, shows the code, and is honest about what it cannot do yet.

The Problem: Memory Attribution in AI Agents

I ran into this while building an agent that accumulated memories across sessions. At some point it started giving answers that were just wrong — confidently wrong. The hard part was not fixing it. The hard part was figuring out which memory caused it. With dozens of stored memories, there is no obvious way to trace a wrong answer back to its source.

Today's approaches fail in predictable ways:

  • LLM-as-judge: Non-deterministic, expensive, and gives explanations that are not real causal chains
  • Attention-weight analysis: Requires white-box model access and does not generalize across providers
  • Manual review: Does not scale past a few dozen memories
  • Embedding similarity: Finds topically related memories, not causally responsible ones

What we need is a method that isolates the minimal set of memories that actually causes a specific wrong answer — deterministically, without requiring model internals.

The Approach: Applying Delta Debugging to Memory Attribution

Core Insight

The idea came from reading about delta debugging in software testing. Zeller's 2002 paper describes ddmin: given a binary predicate (does this input cause the failure?), find the minimal failing subset in O(n log n) calls. The problem maps directly. Instead of minimizing a sequence of code changes, I minimize a set of retrieved memories. The predicate is the same: does this subset reproduce the wrong answer?

Algorithm: ddmin Applied to Memory Attribution 

def delta_debug(candidates, question, failing_answer, n=2):
    if len(candidates) == 1:
        return candidates

    chunk_size = max(1, len(candidates) // n)
    chunks = [candidates[i*chunk_size:(i+1)*chunk_size]
              for i in range(n)]
    if len(candidates) % n:
        chunks[-1] += candidates[n*chunk_size:]

    for chunk in chunks:
        if still_fails(chunk, question, failing_answer):
            return delta_debug(chunk, question, failing_answer, 2)

    for chunk in chunks:
        complement = [m for m in candidates if m not in chunk]
        if still_fails(complement, question, failing_answer):
            return delta_debug(complement, question, failing_answer,
                               max(2, n - 1))

    if n < len(candidates):
        return delta_debug(candidates, question, failing_answer,
                           min(2*n, len(candidates)))

    return candidates
Enter fullscreen mode Exit fullscreen mode

Predicate: still_fails() 

def still_fails(memory_subset, question, failing_answer):
    context = build_context(memory_subset)
    candidate_answer = replay_query(context, question)
    return is_contradiction(candidate_answer, failing_answer)

class ContradictionDetector:
    NEGATION_PAIRS = [
        ("is", "is not"), ("can", "cannot"), ("will", "will not"),
        ("should", "should not"), ("does", "does not"),
        ("allergic", "not allergic"), ("safe", "unsafe"),
    ]

    def contradicts(self, a, b):
        a_lower, b_lower = a.lower(), b.lower()
        for pos, neg in self.NEGATION_PAIRS:
            if pos in a_lower and neg in b_lower:
                return True
            if neg in a_lower and pos in b_lower:
                return True
        sim = cosine_similarity(embed(a), embed(b))
        return sim < -0.3
Enter fullscreen mode Exit fullscreen mode

System Architecture

Drop-in Replacement for LangGraph 

# Before IMMUNE
from langgraph.store.memory import InMemoryStore
store = InMemoryStore()

# After IMMUNE — all other code unchanged
from immune import ImmuneStore
store = ImmuneStore(redis_url='redis://localhost:6379')
Enter fullscreen mode Exit fullscreen mode

Redis Quarantine Enforcement

Quarantine is enforced at the database layer, not the application layer. Even if a future developer removes the filter from the application code, the index schema guarantees quarantined memories are never returned.

FT.CREATE memory_idx ON HASH PREFIX 1 mem:
  SCHEMA
    content TEXT WEIGHT 2.0
    status TAG
    embedding VECTOR HNSW 6 TYPE FLOAT32 DIM 1536 DISTANCE_METRIC COSINE

FT.SEARCH memory_idx
  '(@status:{active})=>[KNN 10 @embedding $vec AS score]'
  PARAMS 2 vec $query_vector
  SORTBY score ASC
  DIALECT 2
Enter fullscreen mode Exit fullscreen mode

Observability 

with tracer.start_as_current_span('memory_attribution') as span:
    span.set_attribute('total_candidates', len(candidates))
    span.set_attribute('oracle_calls', oracle_call_count)
    span.set_attribute('minimal_culprits', len(result))
    span.set_attribute('quarantine_triggered', bool(result))
Enter fullscreen mode Exit fullscreen mode

MCP Tool Integration 

from fastmcp import FastMCP
mcp = FastMCP('IMMUNE Memory Safety')

@mcp.tool()
def quarantine_memory(memory_id: str, reason: str) -> dict:
    redis_client.hset(f'mem:{memory_id}', 'status', 'quarantined')
    return {'quarantined': memory_id, 'status': 'success'}

@mcp.tool()
def attribute_answer(question: str, answer: str) -> dict:
    candidates = fetch_relevant_memories(question)
    culprits = delta_debug(candidates, question, answer)
    return {'culprits': [m['id'] for m in culprits], 'count': len(culprits)}
Enter fullscreen mode Exit fullscreen mode

Results

Benchmarks

Tests run on MacBook Pro M2, Redis 7.2, Python 3.11.

Single poison, 10 memories — Oracle calls: 8 — Detected: Yes — False positives: 0
Single poison, 50 memories — Oracle calls: 22 — Detected: Yes — False positives: 0
Single poison, 100 memories — Oracle calls: 39 — Detected: Yes — False positives: 0
Dual poison, 50 memories — Oracle calls: 41 — Detected: Yes — False positives: 0
Dual poison, 100 memories — Oracle calls: 74 — Detected: Yes — False positives: 1*
Triple poison, 50 memories — Oracle calls: 89 — Detected: Yes — False positives: 0
No poison (baseline), 50 memories — Oracle calls: 0 — Detected: No flag — False positives: 0
Semantic contradiction, 50 memories — Oracle calls: 28 — Detected: Yes — False positives: 0

*One healthy memory falsely included; it was semantically adjacent to the poison, not causally related.

Performance

  • Oracle calls: O(n log n) — consistent with ddmin theory
  • Each oracle call: 200-400ms on Ollama llama3, 800ms on GPT-4o-mini
  • 100 memories, 1 culprit: ~8 seconds end-to-end
  • Quarantine write: less than 5ms (Redis HSET)
  • Retrieval overhead vs InMemoryStore: +12ms average

Tech Stack

  • Python 3.11 — core runtime
  • LangGraph — agent framework (ImmuneStore is a BaseStore drop-in)
  • Redis 7.2 + RediSearch — vector store with schema-level quarantine
  • OpenAI text-embedding-3-small (or local Ollama) — memory embeddings
  • Arize Phoenix + OpenTelemetry — attribution trace observability
  • FastMCP — MCP server for tool integration
  • pytest + pytest-asyncio — 44 tests
  • Docker Compose — local dev environment

How to Run 

git clone https://github.com/Sixxhi/AI-Hackathon-20206
cd AI-Hackathon-20206
pip install -r requirements.txt
docker-compose up -d
make demo
make test
Enter fullscreen mode Exit fullscreen mode

Honest Limitations

  • Semantic predicate gap: The ContradictionDetector uses rule-based negation pairs. It misses subtle domain-specific contradictions.
  • Predicate quality is the bottleneck: ddmin isolates the minimum failing set relative to the predicate. A weak predicate returns a superset of culprits.
  • Single-turn attribution only: Multi-turn conversation poisoning is not yet addressed.
  • No production load testing: The benchmark is synthetic.
  • No adversarial robustness testing: A sophisticated attacker who knows ddmin could craft memories to evade detection.

Future Work

  • Replace rule-based ContradictionDetector with a fine-tuned classifier
  • Extend ddmin to conversation-turn x memory-id pairs for multi-turn attribution
  • Design adversarial red-team evaluation
  • Production pilot with real latency and false positive measurements

Conclusion

Memory poisoning is a real problem in deployed agent systems that does not get enough attention. Most safety work focuses on training time. But agents in production accumulate memories that can go bad — that is a runtime problem.

IMMUNE is my attempt to address it with a concrete, deterministic approach. It is not production-ready and the limitations section says so clearly. But the core idea — applying ddmin to memory attribution — works for the cases tested, and I think it is worth building on.

Code: https://github.com/Sixxhi/AI-Hackathon-20206

Happy to discuss in the comments, especially on the predicate design and any related work I missed.

Top comments (0)