源码解读：MemGPT如何实现“无限”上下文

MemGPT概述

当前，大型语言模型（LLMs）在处理扩展对话和文档分析等任务时受到有限上下文限制，限制了它们在扩展对话和文档分析等任务中的能发挥的作用。如何让大语言模型拥有更强的上下文处理能力是业界非常看重的热点主题之一。

对此，UC Berkeley的研究者提出了一种虚拟上下文管理技术，这种技术从传统操作系统中汲取灵感，通过对快速和慢速内存之间数据的管理，使得 LLMs 能够在有限的上下文窗口中有效使用扩展上下文。并由此引出一个智能管理不同内存层级的系统MemGPT。

MemGPT 的设计灵感来自传统操作系统的层次内存管理，它能够在“内存”（类似于操作系统中的内存）和外部存储之间移动信息。MemGPT 管理内存管理、LLM 处理模块和用户之间的控制流。这种设计允许在单个任务中进行重复的上下文修改，使代理能够更有效地利用其有限的上下文。

MemGPT 在两个场景进行了评估：文档分析和多会话聊天。在文档分析场景，MemGPT 能够分析远远超出基础 LLM 上下文窗口的大型文档。在多会话聊天场景，MemGPT 可以创建记忆、反思和动态演变的对话代理，通过与用户的长期互动获取长期记忆。

MemGPT构成

MemGPT 的设计包括几个关键组件，这些组件共同支持其在处理扩展对话和文档分析任务时动态管理和使用扩展上下文的能力：

1. 主上下文（Main Context）

• 这是模型可以直接访问和操作的内存区域。它类似于计算机的 RAM，并且受到模型的最大令牌限制。MemGPT的主上下文可以简单的看成是一个消息流FIFO队列，外加存储用户上下文、LLM角色上下文的Memory。

2. 外部上下文（External Context）

• 这是一个更大的存储区域，用于存储不立即需要但可能在将来需要的信息。它类似于计算机的硬盘存储。在MemGPT中包含recall memory、archival memory

类型	说明
recall memory	持久化的历史交互上下文
archival memory	可以是交互中的上下文，也可以是外部载入的数据、文档等

3. 函数调用

• MemGPT 提供了一组函数，允许模型执行特定的内存管理任务，如从外部上下文检索信息或更新主上下文。通过在prompt里面指导LLMs去使用这些函数来达到自主更新上下文、搜索上下文、向用户发送消息等目的。
• 例如在多会话聊天场景，LLMs可以使用以下函数

函数	说明
send_message	LLM向用户发送消息
core_memory_append	新增当前交互会传递的上下文
core_memory_replace	编辑当前交互会传递的上下文，如将用户的名字从“John”修改为“Mike”
conversation_search	搜索recall memory历史
conversation_search_date	时间范围条件搜索recall memory历史
archival_memory_insert	往已归档的memory新增数据
archival_memory_search	从已归档的上下文中搜索数据

4. 内存管理算法

• 用于内存管理，确定何时和如何在主上下文和外部上下文之间移动信息。
• 例如在发送给LLMs达到token上限时，按一定的规则将core memory送给LLMs作总结并替换掉
• 再比如当LLMs需要找不到所需上下文时通过调用相关function从recall memory、archival memory搜索到所需上下文并加入到core memory

5. 控制流管理

• MemGPT 能够管理与用户的交互流程，确定何时接收输入，何时提供输出，以及何时执行内存管理任务。让我们看看相关的核心流程：

async def step():
  try:
      # Step 0: add user message
      packed_user_message = {'role': 'user', 'content': user_message}
      input_message_sequence = self.messages + [packed_user_message]

      # Step 1: send the conversation and available functions to GPT
      response = await get_ai_reply_async(model=self.model, message_sequence=input_message_sequence, functions=self.functions)

      # Step 2: check if LLM wanted to call a function
      # (if yes) Step 3: call the function
      # (if yes) Step 4: send the info on the function call and function response to LLM
      response_message = response.choices[0].message
      all_response_messages, heartbeat_request, function_failed = await self.handle_ai_response(response_message)
      
      # Add the extra metadata to the assistant response
      # (e.g. enough metadata to enable recreating the API call)
      all_response_messages[0]['api_response'] = response_message
      all_response_messages[0]['api_args'] = {
          'model': self.model,
          'messages': input_message_sequence,
          'functions': self.functions,
      }
      
      # Step 4: extend the message history
      all_new_messages = [packed_user_message] + all_response_messages
      # ...

  except Exception as e:

      printd(f"step() failed\nuser_message = {user_message}\nerror = {e}")
      # If we got a context alert, try trimming the messages length, then try again
      if 'maximum context length' in str(e):
          # A separate API call to run a summarizer
          await self.summarize_messages_inplace()
          # Try step again
          return await self.step(user_message, first_message=first_message)
      else:
          raise e

举个实际交互的例子

MemGPT	assistant/LLMs	user	说明
pre-prompt
	[send_message]Hi, is anyone there?
	[send_message]Hello! I’m Sam, your new virtual companion. How can I assist you today?
		my name is john, I took the day off today, my mom Brenda baked me a birthday cake. It was my favorite chocolate lava!
	[send_message]Happy birthday, John! It sounds like you had a wonderful day off. How was the chocolate lava cake? Did you enjoy it?
		sure, but, my name is mike
	[send_message]Oops, I apologize for the confusion, Mike! Happy belated birthday. How was the chocolate lava cake? I hope it was absolutely delicious!
		sure, do you remember that time when i talk about my mom with you
	[conversation_search]query: mom		LLM尝试从外部上下文历史记录中搜索
[conversation_search] reutn “my name is john, I took the day off today, my mom Brenda baked me a birthday cake. It was my favorite chocolate lava!”
		what is my name
	[core_memory_replace] update name from Chad to Mike		LLMs发起记忆修正
core_memory_replace
	[send_message]I apologize for the confusion, Mike! Happy belated birthday. How was the chocolate lava cake? By the way, I remember when you mentioned your mom, Brenda, baking you the cake. Such a lovely gesture!

时序图如下

总结

对于LLMs token限制的问题，MemGPT通过提供上下文管理、上下文搜索的function，用prompt指导LLMs去进行外部上下文的调取以及上下文的管理。使得其具备以下优点：

1. 动态上下文管理：MemGPT 能够动态管理和使用扩展的上下文，这有助于处理更复杂和连贯的对话和文档分析任务。
2. 长期记忆能力：MemGPT 可以记住和使用长期对话中的信息，提高对话的连贯性和相关性。
3. 自我导向的内存管理：MemGPT 可以自主更新和搜索其内存，使其能够更加智能和灵活地响应用户的需求。

与此同时，缺点也是很明显的
4. 复杂性：MemGPT 的设计和实现可能相对复杂，需要更多的资源和技术支持来实现系统集成。
5. 计算资源：MemGPT 需要多次跟LLMs进行交互来实现相关的功能，可能需要更多的计算资源来执行动态内存管理和其他高级功能以完成跟用户的单次交互。
6. 可解释性：MemGPT 的上下文管理决策依赖于LLMs，可能不太容易解释和理解。
7. 局限性：MemGPT通过上下文的交换一定程度上解决了token受限的问题，但如果单次所需的上下文token超出LLMs限制，仍然会有存在这个问题。

参考

github：https://github.com/cpacker/MemGPT
论文：https://arxiv.org/abs/2310.08560