Skip to content

[MRG] add hybrid memory for long-term and fast retrieval memory#311

Merged
YuanmingLeee merged 5 commits into
MLSysOps:mainfrom
leeeizhang:lei/hybrid_memory
Jul 14, 2025
Merged

[MRG] add hybrid memory for long-term and fast retrieval memory#311
YuanmingLeee merged 5 commits into
MLSysOps:mainfrom
leeeizhang:lei/hybrid_memory

Conversation

@leeeizhang
Copy link
Copy Markdown
Collaborator

@leeeizhang leeeizhang commented Jul 13, 2025

Closes #310

This PR introduces the HybridMemory that combines a slow, persistent memory backend (e.g., Mem0) with a fast, high-recall vector store (e.g., LanceDBMemory). It enables:

  • Dual-memory: supports querying either the fast store alone or both fast and slow stores for more comprehensive recall.
  • Memory consolidation (distill slow-memory into fast-memory):
    • last_n_consolidate: distill the most recent N entries from slow memory into fast memory.
    • topk_consolidate: distill top-K entries based on metadata ranking from slow memory into fast memory.

Usages:

from mle.utils.memory import Mem0, LanceDBMemory, HybridMemory

chat_sessions: list[dict[str, str]] = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "assistant", "content": "What is your favorite city?"},
    {"role": "user", "content": "The capital of France, Paris."},
    {"role": "assistant", "content": "Sure, I will remember for these"},
]
metadata = {"HITL-score": 1.0, "Eval-score": 99.1}

slow_mem = Mem0()
fast_mem = LanceDBMemory("/path/to/mle/project")

slow_mem.reset()
fast_mem.reset()

hybrid_mem = HybridMemory(slow_mem, fast_mem)

hybrid_mem.add(chat_sessions, metadata=metadata, prompt="summary all information")
hybrid_mem.last_n_consolidate(n=1)
hybrid_mem.top_k_consolidate(k=1, metadata_key="HITL-score")

print(hybrid_mem.query("what is the user's favorite city?", n_results=1)[0][0]["text"])

Example outputs:

Query: what is the user's favorite city?
Favorite city is Paris, the capital of France

Before submitting this PR, please make sure you have:

  • confirmed all checks still pass OR confirm CI build passes.
  • verified that any code or assets from external sources are properly credited in comments and/or in
    the credit file.

@dosubot dosubot Bot added size:L This PR changes 100-499 lines, ignoring generated files. enhancement New feature or request labels Jul 13, 2025
@leeeizhang leeeizhang changed the title [WIP] add hybrid memory [WIP] add hybrid memory for long-term and fast retrieval memory Jul 13, 2025
@leeeizhang leeeizhang changed the title [WIP] add hybrid memory for long-term and fast retrieval memory [MRG] add hybrid memory for long-term and fast retrieval memory Jul 14, 2025
Copy link
Copy Markdown
Contributor

@YuanmingLeee YuanmingLeee left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lgtm

@YuanmingLeee YuanmingLeee merged commit 8c2ec57 into MLSysOps:main Jul 14, 2025
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request size:L This PR changes 100-499 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants