☰
LightMem开发者指南:用工厂模式扩展自定义存储引擎与检索策略
2026/10/13 14:53:28 网站建设 项目流程

【免费下载链接】LightMem

[ICLR 2026] LightMem: Lightweight and Efficient Memory-Augmented Generation

项目地址:https://gitcode.com/gh_mirrors/li/LightMem
点击查看免费下载

LightMem 是一个轻量级、高效的记忆增强生成(Memory-Augmented Generation)框架(ICLR 2026 录用),为 LLM 与 AI Agent 提供记忆存储、检索与更新能力。它的核心设计亮点是采用**工厂模式(Factory Pattern)**将存储引擎与检索策略模块化:你只需通过配置即可在 OpenAI、Ollama、vLLM 等后端间切换,也能按统一规范扩展自己的向量存储引擎(如 Chroma、Milvus)或检索算法(如 TF-IDF)。本文带你读懂 LightMem 工厂架构,并给出扩展自定义存储引擎与检索策略的完整路径。

🧩 为什么 LightMem 需要工厂模式

长对话场景下,模型上下文越积越长,成本与延迟急剧上升。LightMem 的思路是:把历史消息压缩、分段成结构化的记忆单元,存进向量库或倒排索引,回答时再按检索策略召回相关记忆。

问题来了:向量库有 Qdrant、Chroma、Milvus 等好几家,LLM 后端有云端 API 也有本地 Ollama / vLLM。如果每个组件都用if/else硬编码,代码很快会失控。LightMem 的解法就是工厂模式:

  • 配置层(src/lightmem/configs/):用 Pydantic 定义每种组件的配置模型,model_name决定用哪个实现;
  • 工厂层(src/lightmem/factory/):每个工厂类维护一张model_name → 模块.类的映射表,通过importlib动态导入并实例化;
  • 入口层(LightMemory):根据retrieve_strategy、index_strategy等开关按需初始化组件。

📂 工厂全景图:一张目录地图

LightMem 的所有可插拔组件都集中在两个目录,结构一一对应:

组件工厂实现目录配置目录已内置实现
记忆管理器(LLM 后端)memory_manager/memory_manager/openai、deepseek、ollama、transformers、vllm、vllm_offline
向量检索器(存储引擎)embeddingretriever/embeddingretriever/qdrant
上下文检索器contextretriever/contextretriever/BM25
文本嵌入器text_embedder/text_embedder/huggingface、openai
预压缩器pre_compressor/pre_compressor/llmlingua-2、entropy_compress

其中向量检索器工厂 EmbeddingRetrieverFactory 就是扩展自定义存储引擎的"注册表":

class EmbeddingRetrieverFactory: _MODEL_MAPPING: Dict[str, str] = { "qdrant": "lightmem.factory.retriever.embeddingretriever.qdrant.Qdrant", }

⚙️ 工作原理:从名字到实例的三步跳转

以 ContextRetrieverFactory.from_config() 为例,每个工厂的实例化逻辑完全一致,只有三步:

  1. 查表:拿配置里的model_name(如"BM25")去_MODEL_MAPPING找目标类的完整路径;
  2. 动态导入:用rsplit('.', 1)拆出模块路径与类名,import_module()+getattr()拿到类;
  3. 实例化:若配置里带了configs字典,则以config=configs构造,否则直接无参实例化。

与之配合的配置侧同样"自描述"。例如 ContextRetrieverConfig 维护了一份_model_list = ["BM25"],在model_name不在列表内时直接抛出ValueError,帮你把错误拦在启动阶段。而 EmbeddingRetrieverConfig 更进一步:校验通过后会自动把configs字典转成对应的具体配置类(如 QdrantConfig),后者还内置了"host/port、url/api_key、path 三选一"的连接校验。

🚀 四步接入自定义存储引擎

假设你想把 Qdrant 换成 Chroma。按 LightMem 的规范,只需四处改动:

第 1 步:实现组件类。参照 Qdrant 的接口风格,新类接收config参数,实现insert()、search()等方法即可。

第 2 步:定义配置类。在configs/retriever/embeddingretriever/下新建一个 Pydantic 配置模型(字段、默认值、连接校验照 QdrantConfig 写)。

第 3 步:注册映射。在两处各加一行:

# 工厂侧:factory/retriever/embeddingretriever/factory.py _MODEL_MAPPING = { "qdrant": "lightmem.factory.retriever.embeddingretriever.qdrant.Qdrant", "chroma": "your.package.chroma.ChamraStore", # 新增 } # 配置侧:configs/retriever/embeddingretriever/base.py _model_list = { "qdrant": "...qdrant.QdrantConfig", "chroma": "...chroma.ChamraConfig", # 新增 }

第 4 步:用配置启用。无需改任何业务代码:

from lightmem.configs.base import BaseMemoryConfigs config = BaseMemoryConfigs( retrieve_strategy="embedding", embedding_retriever={ "model_name": "chroma", "configs": {"collection_name": "my_mem", "embedding_model_dims": 1536}, }, )

retrieve_strategy支持embedding/context/hybrid三档(见 BaseMemoryConfigs),LightMemory 初始化时会据此按需创建检索器:hybrid模式下同时启动上下文检索器与向量检索器(lightmem.py),未选中的组件完全不加载。

🧪 自定义检索策略怎么配

除了换存储引擎,LightMem 还提供正交的两个开关:

  • index_strategy(建索引用):embedding/context/hybrid,决定是否建向量索引、倒排索引或两者都建(base.py);
  • retrieve_strategy(检索用):同上三选一,决定查询走哪条路。

扩展新检索算法(如 TF-IDF)的路径与向量库相同:在 ContextRetrieverFactory._MODEL_MAPPING 注册实现类 → 在 ContextRetrieverConfig._model_list 登记名字 → 配置中指定context_retriever={"model_name": "tfidf"}。

💡 提示:仓库内还有 MCP Server(mcp/server.py)与实验脚本(experiments/locomo/),可结合 LoCoMo、LongMemEval 数据集验证你的新策略效果。

🛠 常见扩展点速查

想扩展什么改哪里参照实现
新的向量存储引擎EmbeddingRetrieverFactory + EmbeddingRetrieverConfigqdrant.py
新的检索算法ContextRetrieverFactory + ContextRetrieverConfigbm25.py
新的 LLM 后端MemoryManagerFactoryollama.py
新的嵌入模型TextEmbedderFactoryhuggingface.py
新的压缩算法PreCompressorFactoryllmlingua_2.py

📌 小结

LightMem 的工厂模式把"用什么后端"收敛成了一行配置:配置层负责校验与自描述,工厂层负责动态加载,入口层按策略懒加载。理解了_MODEL_MAPPING(工厂侧)与_model_list(配置侧)这两张注册表,你就可以用自己的向量库、检索算法或 LLM 后端无痛替换 LightMem 的任意组件,构建专属的记忆增强应用。

  • 完整架构说明:README.md
  • 记忆核心实现:src/lightmem/memory/lightmem.py
  • 组件配置基类:src/lightmem/configs/base.py
  • 运行实验:experiments/

【免费下载链接】LightMem

[ICLR 2026] LightMem: Lightweight and Efficient Memory-Augmented Generation

项目地址:https://gitcode.com/gh_mirrors/li/LightMem
点击查看免费下载

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询