用 Go 构建上下文引擎

AI 原生的内存管理,从零开始。

用 Go 构建上下文引擎
AI模型价格对比 | AI工具导航 | ONNX模型库 | Vibe Coding教程 | PLC在线仿真器 | Tripo 3D | Meshy AI | ElevenLabs | KlingAI | ArtSpace | Phot.AI | InVideo

LLM 本质上是无状态的。如果没有健壮的内存系统,它们基本上就像金鱼——聪明,但一旦上下文窗口发生偏移,就容易产生幻觉。

在构建 AI 代理时,"内存"不仅仅是一个数据库;它是一个高性能的编排层,决定代理保留什么、遗忘什么、优先处理什么。

在这篇文章中,我们将完整地走一遍一个可运行的原生、零依赖内存引擎的实现,它完全基于 Go 标准库构建。每个函数都会覆盖到,每个设计决策都会解释清楚,每一行代码都有交代。实现代码可以从这里下载。

None

1、开始之前:导入与数据模型

我们从导入开始。这里的每个包都属于 Go 标准库——没有需要安装的外部依赖,没有 go.mod 技巧,也没有构建管线的意外。

import (
	"container/heap"
	"container/list"
	"encoding/gob"
	"fmt"
	"math"
	"os"
	"strings"
)

container/heap 为我们提供了优先队列,container/list 为滑动窗口提供了双向链表,encoding/gob 处理二进制序列化,math 提供计算余弦相似度所需的平方根,os 管理文件 I/O,strings 驱动我们的词袋(bag-of-words)向量化器。就这些——六个包,合在一起覆盖了 AI 代理内存的完整生命周期。

这个系统中的一切都围绕一个结构体展开:MemoryItem

type MemoryItem struct {
	ID       int
	Content  string
	Priority int
	Vector   []float64
	index    int
}

ID 是一个简单的标识符。Content 保存内存的实际文本——对话轮次、系统指令、关键笔记。Priority 是一个整数,决定这条内存有多重要;数值越高,意味着该项应被堆更积极地保留和浮现。Vector 是内容的数值表示,我们将用词袋方法填充它,以便代理执行语义搜索。

index 字段是让 Go 堆接口新手感到惊讶的地方:container/heap 需要它来跟踪每个项在底层切片中的位置,从而支持高效的重新排序,以及在插入后某项优先级动态变化时的更新。

2、优先级引擎:使用 container/heap 的最大堆

为了确保代理在上下文拥挤时不会丢失关键指令,我们使用最大堆。通过实现 heap.Interface,我们获得了 O(log n) 的插入和删除复杂度,这意味着即使存储了数千条记忆,检索最重要的项也始终很快。

我们将 PriorityQueue 类型定义为指向 MemoryItem 的指针切片,然后实现接口要求的五个方法。

type PriorityQueue []*MemoryItem

func (pq PriorityQueue) Len() int { return len(pq) }

Len 很直接——它只是返回底层切片的长度。堆用它来判断队列何时为空,并管理其内部树结构。

func (pq PriorityQueue) Less(i, j int) bool {
	return pq[i].Priority > pq[j].Priority // Max-heap logic
}

Less 是最大堆行为的来源。通过比较 pq[i].Priority > pq[j].Priority,我们确保高优先级的项会向队列前端冒泡。如果反转比较符号为 <,我们得到的就是最小堆——在不同场景下很有用,但对我们的代理来说,我们想让最重要的记忆排在顶部。

func (pq PriorityQueue) Swap(i, j int) {
	pq[i], pq[j] = pq[j], pq[i]
	pq[i].index = i
	pq[j].index = j
}

Swap 身兼两职。它在堆操作期间对切片中的项重新排序,同时更新每个项的 index 字段,使堆始终知道所有东西的位置。这对正确性至关重要——如果不更新 index,堆在任意一次重新排序后就会丢失项的踪迹。

func (pq *PriorityQueue) Push(x interface{}) {
	n := len(*pq)
	item := x.(*MemoryItem)
	item.index = n
	*pq = append(*pq, item)
}

Push 将新项追加到切片末尾,并在 index 字段中记录其位置。然后堆执行向上冒泡以恢复不变量。注意这里的指针接收者——Push 会修改切片头本身,所以它需要 *PriorityQueue

func (pq *PriorityQueue) Pop() interface{} {
	old := *pq
	n := len(old)
	item := old[n-1]
	old[n-1] = nil
	item.index = -1
	*pq = old[0 : n-1]
	return item
}

Pop 是最棘手的部分。堆重组之后,根节点(最高优先级项)位于切片末尾。Pop 移除该元素,将已移除的槽位设为 nil 以防止内存泄漏(Go 的 GC 无法回收仍被切片引用的对象),将项的 index 重置为 -1 以标记其不再位于堆中,然后截断切片。这种清理纪律使得堆在长时间运行的会话中(项不断被压入和弹出)保持健康。

3、滑动窗口:使用 container/list 的时间上下文

对于短期上下文,我们需要一个尊重事件时间顺序的结构。container/list 提供双向链表,非常适合 FIFO 队列,因为从头部移除和向尾部插入都是 O(1) 操作。

type SlidingWindow struct {
	queue    *list.List
	capacity int
}

func NewSlidingWindow(capacity int) *SlidingWindow {
	return &SlidingWindow{queue: list.New(), capacity: capacity}
}

SlidingWindow 结构体包装了一个 *list.List 和一个 capacity 整数。构造函数 NewSlidingWindow 初始化两者,给我们一个具有固定大小限制、开箱即用的窗口。

func (sw *SlidingWindow) Add(item *MemoryItem) {
	if sw.queue.Len() >= sw.capacity {
		sw.queue.Remove(sw.queue.Front())
	}
	sw.queue.PushBack(item)
}

Add 方法是滑动窗口的引擎。如果队列已达到容量,它会在将新项追加到尾部之前移除最前面的元素——也就是最旧的那个。这就是驱逐策略:先进先出。RemovePushBack 的常数时间保证意味着窗口可以处理高吞吐的事件流而没有任何性能退化,这正是代理处理快速对话流时所需要的。

4、向量化与相似度:语义搜索

要让代理"聪明",我们需要语义搜索。由于我们不使用外部向量数据库,我们仅用 mathstrings 从零实现一个词袋向量化器和余弦相似度。

func TextToVector(text string) []float64 {
	vocab := []string{"database", "connection", "fix", "leak", "system", "user"}
	vec := make([]float64, len(vocab))
	words := strings.Fields(strings.ToLower(text))
	for _, word := range words {
		for i, v := range vocab {
			if strings.Contains(word, v) {
				vec[i]++
			}
		}
	}
	return vec
}

TextToVector 首先定义一个固定词表——对我们的演示很重要的六个词项。然后它创建一个零填充向量,每个词表项占一个槽位。输入文本被转成小写,并使用 strings.Fields 拆分成单词,它能干净地处理空白分隔。

对于每个单词,我们使用 strings.Contains 检查它是否包含任何词表项作为子串。这是一个刻意的设计选择:这意味着"connection"能匹配"database connection",甚至能匹配"connections"(因为"connection"是"connections"的子串),免费获得了一点模糊匹配。每次匹配都会递增对应的向量索引,构建出文本的基于频率的表示。

func CosineSimilarity(a, b []float64) float64 {
	var dot, normA, normB float64
	for i := range a {
		dot += a[i] * b[i]
		normA += a[i] * a[i]
		normB += b[i] * b[i]
	}
	if normA == 0 || normB == 0 {
		return 0
	}
	return dot / (math.Sqrt(normA) * math.Sqrt(normB))
}

CosineSimilarity 计算两个向量的点积,再除以它们欧几里得范数的乘积。顶部的零范数保护处理了一个向量完全为空的边界情况——返回 0 而不是除以零。结果是一个介于 0 和 1 之间的值(对于像我们这样的非负向量),其中 1 表示向量方向完全相同,0 表示它们不共享任何维度。这使代理能够针对自然语言查询找到最相关的记忆项,而无需外部嵌入模型。

5、持久化:使用 encoding/gob 的二进制序列化

一个重启就忘掉一切的代理毫无用处。我们使用 encoding/gob 将内存状态直接序列化到磁盘,确保代理"醒来"时上下文完好无损。

type AgentMemory struct {
	LongTerm []*MemoryItem
}

AgentMemory 是长期优先队列的简单包装器。它持有一个指向 MemoryItem 的指针切片,这正是堆产生的结构。这个结构体就是我们编码并写入磁盘的对象。

func SaveMemory(filename string, mem *AgentMemory) error {
	file, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer file.Close()
	return gob.NewEncoder(file).Encode(mem)
}

SaveMemory 创建一个新文件(或截断现有文件),延迟关闭以确保即使编码失败文件句柄也能被释放,然后使用 gob.NewEncoder 以 Go 的原生二进制格式写入整个结构体。在这种用例下,gob 比 JSON 明显更快、更紧凑,因为它将类型信息与数据一起编码,无需在序列化输出中包含字段名。错误返回模式是地道的 Go 风格:调用者可以决定一次失败的保存是致命的还是可以重试的。

6、执行:main()

main 函数编排一切。它初始化堆和滑动窗口,摄入一组带自动向量化的示例记忆,演示滑动窗口的驱逐行为,使用自然语言查询执行语义搜索,最后将状态持久化到磁盘。

我们先建立数据结构。

pq := &PriorityQueue{}
heap.Init(pq)
sw := NewSlidingWindow(3)

在空优先队列上调用 heap.Init 以建立堆不变量。滑动窗口以容量 3 创建,意味着它在任何时刻最多容纳三个项。

接下来,我们定义示例记忆,并同时将它们摄入两个层。

rawItems := []struct {
	content  string
	priority int
}{
	{"System: You are a helpful assistant.", 10},
	{"User: How do I fix the database connection?", 5},
	{"CRITICAL: Fix database connection leak.", 9},
}

for _, ri := range rawItems {
	item := &MemoryItem{
		Content:  ri.content,
		Priority: ri.priority,
		Vector:   TextToVector(ri.content),
	}
	heap.Push(pq, item)
	sw.Add(item)
}

每个原始项都被转换成一个 MemoryItemContentPriority 直接来自示例数据,Vector 通过调用 TextToVector 填充内容来获得——这就是自动向量化步骤。然后我们把项推入优先队列并添加到滑动窗口。注意,同一个 MemoryItem 指针在两个结构之间共享,所以对其中一个的更新会反映在另一个上。

现在让我们看看滑动窗口的实际运行。

fmt.Println("--- Sliding Window (capacity 3) ---")
for e := sw.queue.Front(); e != nil; e = e.Next() {
	fmt.Printf("  Window holds: '%s'\n", e.Value.(*MemoryItem).Content)
}

我们从链表的前端向后迭代,打印每个项的内容。此时,窗口按插入顺序持有全部三个初始项。

for i := 0; i < 2; i++ {
	sw.Add(&MemoryItem{
		ID:       i + 1,
		Content:  fmt.Sprintf("Stream item %d: recent event", i+1),
		Priority: 1,
		Vector:   TextToVector("recent event"),
	})
}

然后我们再添加两个低优先级(1)和通用向量的项。因为窗口容量是 3,前两个项("System: You are a helpful assistant." 和 "User: How do I fix the database connection?")会随着新项的到来被驱逐。只有最近的三个存活下来。

fmt.Println("  After 2 more stream items added:")
for e := sw.queue.Front(); e != nil; e = e.Next() {
	fmt.Printf("  Window holds: '%s'\n", e.Value.(*MemoryItem).Content)
}

输出证实了驱逐:窗口现在持有高优先级的 CRITICAL 项,加上两个新的流式项。

现在到了语义搜索部分——代理真正"思考"它记住了什么的环节。

query := "database connection fix"
queryVec := TextToVector(query)

var bestMatch *MemoryItem
maxSim := -1.0

fmt.Printf("Querying for: '%s'\n", query)
for _, item := range *pq {
	sim := CosineSimilarity(item.Vector, queryVec)
	fmt.Printf("Comparing with: '%s' (Sim: %.2f)\n", item.Content, sim)
	if sim > maxSim {
		maxSim = sim
		bestMatch = item
	}
}

我们使用同一个 TextToVector 函数将自然语言查询"database connection fix"转换为向量,然后遍历优先队列中的每个项。

对于每个项,我们计算其向量与查询向量之间的余弦相似度。我们跟踪最佳匹配及其相似度分数。输出显示,系统指令得分 0.00(与查询没有重叠),而用户问题和 CRITICAL 笔记都得分 0.87——最高的相似度。先遇到的那个作为 bestMatch 胜出。

fmt.Printf("\nBest match found: '%s' (Sim: %.2f)\n", bestMatch.Content, maxSim)

这会打印最佳匹配,证实代理成功基于语义相似度而非精确关键词匹配检索到了最相关的记忆。

最后,我们将整个内存状态持久化到磁盘。

mem := &AgentMemory{LongTerm: *pq}
SaveMemory("agent_memory.gob", mem)
fmt.Println("Memory state saved to agent_memory.gob")

我们把优先队列包装进一个 AgentMemory 结构体并调用 SaveMemory,它将一切序列化到 agent_memory.gob。在后续运行时,这个文件可以被加载回来,恢复代理的完整内存状态。

下面是运行该程序时你应该看到的完整输出:

PS C:\Users\User\ai-memory> go run .\ai_memory_experiment.go
--- Sliding Window (capacity 3) ---
  Window holds: 'System: You are a helpful assistant.'
  Window holds: 'User: How do I fix the database connection?'
  Window holds: 'CRITICAL: Fix database connection leak.'
  After 2 more stream items added:
  Window holds: 'CRITICAL: Fix database connection leak.'
  Window holds: 'Stream item 1: recent event'
  Window holds: 'Stream item 2: recent event'

Querying for: 'database connection fix'
Comparing with: 'System: You are a helpful assistant.' (Sim: 0.00)
Comparing with: 'User: How do I fix the database connection?' (Sim: 0.87)
Comparing with: 'CRITICAL: Fix database connection leak.' (Sim: 0.87)

Best match found: 'User: How do I fix the database connection?' (Sim: 0.87)
Memory state saved to agent_memory.gob
PS C:\Users\User\ai-memory>

7、用 MCP 为代理赋能

这个实现不仅仅是一次编码练习——它是模型上下文协议(MCP)服务器的基础。通过把这个内存引擎包装进一个 MCP 接口,你可以把代理的长期记忆暴露给任何 LLM(比如 Claude 或本地的 Ollama 模型)。LLM 可以像调用工具一样查询你的 Go 原生内存存储,实际上给了它一个存在于模型瞬时上下文窗口之外的、持久的、可搜索的大脑。

这种方法的美妙之处在于每个组件都是模块化、可替换的。把词袋向量化器换成真正的嵌入模型,为更长的对话增大滑动窗口容量,或者添加一个 LoadMemory 函数在启动时从磁盘恢复状态。因为一切都建立在标准库之上,整个系统可以编译成零运行时依赖的单一静态二进制——非常适合嵌入代理框架、部署到边缘设备,或作为轻量级 MCP 服务器与您喜爱的 LLM 一起运行。

8、关键要点

  • 零依赖:使用的每个包都属于 Go 标准库。没有 go get,没有 vendor 目录,没有供应链担忧。
  • 双内存层:滑动窗口以 O(1) 驱逐处理短期上下文,而最大堆以 O(log n) 操作保存长期优先级。
  • 无需外部服务的语义搜索:词袋向量化和余弦相似度让你在不调用嵌入 API 的情况下获得检索增强生成能力。
  • 开箱即用的持久化encoding/gob 将整个代理状态序列化到单个文件,使得跨重启恢复对话变得轻而易举。
  • MCP 就绪架构:模块化设计意味着这个内存引擎可以被包装成一个 MCP 工具,并插入任何支持该协议的代理框架。

原文链接: The Memory-Efficient AI Agent: Building a Context Engine in Go

汇智网翻译整理,转载请标明出处