AI Agent的长期记忆与知识图谱融合：构建持续学习的智能系统

关键词

AI Agent、长期记忆、知识图谱、记忆整合、知识表示、持续学习、神经符号AI

摘要

本文深入探讨AI Agent的长期记忆与知识图谱融合的理论基础、架构设计、实现机制及实际应用。我们从第一性原理出发，分析记忆系统的本质需求，构建记忆-知识融合的数学框架，并详细介绍实现这一融合的多种技术路径。通过结合神经符号AI的最新进展，我们展示如何构建具有持续学习能力、可解释性和泛化能力的智能系统。文章还包括完整的代码实现、案例研究和未来发展方向，为研究人员和实践者提供全面的技术指南。

---

1. 概念基础

核心概念

在深入探讨AI Agent的长期记忆与知识图谱融合之前，我们首先需要明确几个核心概念的定义和内涵：

1. AI Agent：指能够感知环境、做出决策并执行行动以实现特定目标的智能系统。AI Agent可以是软件实体（如聊天机器人、推荐系统）或物理实体（如机器人、自动驾驶汽车）。

2. 长期记忆：在AI语境下，指智能系统能够持久存储并在需要时检索的信息、经验和知识。与短期记忆（工作记忆）不同，长期记忆具有大容量、持久性和可检索性等特点。

3. 知识图谱：一种结构化的知识表示形式，以图的形式组织信息，其中节点表示实体，边表示实体之间的关系。知识图谱能够有效表示世界知识，并支持复杂的推理和查询。

4. 记忆整合：指将新获取的信息与已有记忆系统地结合的过程，包括编码、存储、检索和更新等环节。

5. 神经符号AI：结合神经网络（连接主义）和符号系统（符号主义）优点的AI研究范式，旨在同时利用神经网络的学习能力和符号系统的推理能力。

问题背景

人工智能的发展历程中，构建具有人类级智能的系统一直是终极目标之一。人类智能的一个关键特征是能够通过经验积累知识，并在不同情境中灵活应用这些知识。然而，传统的AI系统往往缺乏有效的长期记忆机制，导致以下问题：

1. 灾难性遗忘：神经网络在学习新任务时往往会忘记之前学到的知识，这一现象被称为灾难性遗忘。

2. 知识碎片化：不同的AI模型和系统通常拥有各自独立的知识库，缺乏有效的知识共享和整合机制。

3. 有限的泛化能力：缺乏结构化知识表示的系统难以将学到的知识迁移到新的、未见过的情境中。

4. 可解释性不足：黑盒模型的决策过程难以理解和追踪，限制了AI系统在关键领域的应用。

知识图谱作为一种结构化的知识表示方法，能够有效组织和推理世界知识，但其构建和更新过程往往需要大量人工干预，且缺乏从经验中自动学习的能力。另一方面，现代深度学习系统虽然具有强大的学习能力，但其知识表示往往是隐式的、分布式的，难以进行精确的推理和解释。

因此，将AI Agent的长期记忆与知识图谱融合，结合两者的优势，成为构建更智能、更灵活的AI系统的关键研究方向。

问题描述

我们需要解决的核心问题可以概括为：如何设计一个AI Agent的记忆系统，使其能够：

1. 持久存储：将Agent的经验和知识以结构化的方式长期保存。

2. 高效检索：在需要时能够快速、准确地检索相关信息。

3. 持续学习：能够自动整合新信息，更新已有知识，并避免灾难性遗忘。

4. 灵活推理：能够基于存储的知识进行复杂推理，解决新问题。

5. 可解释性：能够解释其知识和决策过程，增强系统的可信度。

为了解决这些问题，我们需要探索长期记忆与知识图谱的融合机制，设计合适的架构和算法，并验证其在实际应用中的有效性。

问题解决

解决这一问题的关键思路是构建一个混合记忆系统，结合神经和符号方法的优势：

1. 神经记忆组件：负责处理非结构化数据，学习模式和规律，并将其编码为分布式表示。

2. 符号知识图谱：负责存储结构化知识，支持精确查询和逻辑推理。

3. 融合机制：连接神经和符号组件，实现双向信息转换和协同工作。

4. 记忆管理：包括记忆编码、存储、检索、更新和遗忘等机制，确保记忆系统的效率和有效性。

在接下来的章节中，我们将详细探讨这些组件的理论基础、设计原理和实现方法。

边界与外延

在深入探讨之前，我们需要明确本文的研究边界和相关概念的外延：

1. 研究边界：

   • 本文主要关注软件AI Agent的记忆系统，不包括具身Agent的感知-运动循环。
   • 我们重点讨论陈述性知识（事实和概念）和程序性知识（技能和策略）的记忆，较少涉及感知记忆。
   • 我们假设Agent具有基本的感知和行动能力，重点关注这些能力与记忆系统的交互。

2. 相关概念外延：

   • 长期记忆与工作记忆的关系：我们将讨论如何将工作记忆中的信息转化为长期记忆。
   • 记忆与注意力的关系：我们将探讨注意力机制在记忆检索和更新中的作用。
   • 知识图谱与其他知识表示方法的关系：我们将比较知识图谱与逻辑表示、语义网络等方法的优缺点。

概念结构与核心要素组成

为了构建AI Agent的长期记忆与知识图谱融合系统，我们需要理解以下核心要素及其结构关系：

1. 记忆类型层次：

   • 感觉记忆：短暂保存原始感知数据。
   • 工作记忆：临时保存和处理当前任务相关信息。
   • 长期记忆：持久保存大量信息，包括陈述性记忆和程序性记忆。

2. 知识表示组件：

   • 实体：表示现实世界或抽象概念中的对象。
   • 关系：表示实体之间的联系。
   • 属性：表示实体或关系的特征。
   • 规则：表示知识之间的逻辑关系和约束。

3. 记忆过程：

   • 编码：将感知信息转换为记忆存储格式。
   • 存储：将编码后的信息保存到记忆系统中。
   • 检索：根据当前需求从记忆中获取相关信息。
   • 巩固：将临时记忆转化为长期稳定记忆。
   • 再巩固：检索记忆时对其进行更新和强化。
   • 遗忘：移除不必要或错误的记忆。

4. 系统组件：

   • 感知接口：从环境中获取信息。
   • 记忆编码器：将感知信息转换为记忆表示。
   • 记忆存储：保存编码后的记忆。
   • 记忆检索器：根据查询检索相关记忆。
   • 推理引擎：基于记忆进行推理和决策。
   • 动作接口：执行决策并影响环境。

在接下来的章节中，我们将详细探讨这些要素的理论基础、设计原理和实现方法，并构建一个完整的融合系统。

---

2. 理论框架

第一性原理推导

为了构建AI Agent的长期记忆与知识图谱融合系统，我们首先从第一性原理出发，分析记忆系统的基本需求和约束：

1. 信息存储的基本原理：

   • 任何记忆系统都需要解决三个基本问题：存储什么信息、如何组织信息、如何检索信息。
   • 信息论告诉我们，有效记忆需要权衡信息的完整性、存储效率和检索速度。

2. 记忆系统的热力学约束：

   • 兰道尔原理指出，擦除一比特信息至少需要消耗kTln2的能量，其中k是玻尔兹曼常数，T是温度。
   • 这意味着记忆系统必须有选择地存储信息，避免不必要的能量消耗。

3. 学习与记忆的关系：

   • 从学习理论的角度，记忆可以看作是学习的结果，而学习可以看作是记忆的更新过程。
   • 统计学习理论告诉我们，有效的记忆系统需要在拟合数据和泛化能力之间取得平衡。

4. 知识表示的基本要求：

   • 表示定理指出，任何可计算函数都可以用适当的表示形式来表示。
   • 对于AI系统，知识表示需要支持高效推理、学习和更新。

基于这些第一性原理，我们可以推导出长期记忆与知识图谱融合系统的基本设计原则：

1. 层次性：记忆系统应该具有层次结构，从原始感知数据到抽象概念知识。
2. 选择性：系统应该有选择地存储和保留重要信息，忘记无关或错误信息。
3. 关联性：信息应该以关联的方式组织，支持基于上下文的检索和推理。
4. 适应性：系统应该能够根据经验自动调整其结构和内容。
5. 可解释性：系统的知识和决策过程应该是可追踪和可解释的。

数学形式化

为了更精确地描述长期记忆与知识图谱融合系统，我们引入以下数学形式化：

记忆系统的基本定义

我们将记忆系统定义为一个五元组：
$$\mathcal{M}=(\mathcal{E},\mathcal{R},\mathcal{A},\mathcal{K},\mathcal{O})$$

其中：

• $\mathcal{E}$ 是实体集合，每个实体$e\in \mathcal{E}$代表一个概念或对象。
• $\mathcal{R}$ 是关系集合，每个关系$r\in \mathcal{R}$代表实体之间的一种联系。
• $\mathcal{A}$ 是属性集合，每个属性$a\in \mathcal{A}$代表实体或关系的一个特征。
• $\mathcal{K}$ 是知识集合，每个知识单元$k\in \mathcal{K}$是一个三元组$(e_1,r,e_2)$或属性值对$(e,a,v)$。
• $\mathcal{O}$ 是操作集合，包括编码、存储、检索、更新等操作。

知识图谱的图论表示

我们可以将知识图谱表示为一个有向标记图：
$$\mathcal{G}=(V,E,L_V,L_E)$$

其中：

• $V$ 是顶点集合，对应实体$\mathcal{E}$。
• $E\subseteq V\times V$ 是有向边集合，对应关系实例。
• $L_V:V\to {\Sigma }_V$ 是顶点标记函数，为每个实体分配类型或标签。
• $L_E:E\to {\Sigma }_E$ 是边标记函数，为每条边分配关系类型。

记忆编码的神经表示

我们使用分布式向量表示实体和关系：
$${\mathbf{e}}_i\in {\mathbb{R}}^d,{\mathbf{r}}_j\in {\mathbb{R}}^d$$

其中$d$是嵌入维度。

对于知识三元组$(h,r,t)$，我们定义一个评分函数$f(h,r,t)$来衡量其合理性：
$$f(h,r,t)=g(\mathbf{h},\mathbf{r},\mathbf{t})$$

常见的评分函数包括：

1. TransE：
   $$f(h,r,t)=-\Vert \mathbf{h}+\mathbf{r}-\mathbf{t}\Vert$$

2. DistMult：
   $$f(h,r,t)={\mathbf{h}}^T\cdot \text{diag}(\mathbf{r})\cdot \mathbf{t}$$

3. ComplEx：
   $$f(h,r,t)=\text{Re}({\mathbf{h}}^T\cdot \text{diag}(\mathbf{r})\cdot \overline{\mathbf{t}})$$

其中$\overline{\mathbf{t}}$是$\mathbf{t}$的复共轭。

记忆检索的数学模型

记忆检索可以形式化为一个优化问题：给定查询$q$，找到最相关的记忆项$k^{\ast }$：
$$k^{\ast }=\arg \underset{k\in \mathcal{K}}{\max }\text{sim}(q,k)$$

其中$\text{sim}(q,k)$是查询$q$和记忆项$k$之间的相似度函数。

对于基于内容的检索，我们可以使用向量空间模型：
$$\text{sim}(q,k)=\frac{\mathbf{q}\cdot \mathbf{k}}{\Vert \mathbf{q}\Vert \Vert \mathbf{k}\Vert }$$

其中$\mathbf{q}$和$\mathbf{k}$分别是查询和记忆项的向量表示。

记忆更新的动态模型

我们将记忆更新建模为一个随机过程，使用贝叶斯框架：
$$P(k|D)\propto P(D|k)P(k)$$

其中：

• $P(k)$是记忆项$k$的先验概率。
• $P(D|k)$是给定记忆项$k$时观测数据$D$的似然。
• $P(k|D)$是更新后记忆项$k$的后验概率。

对于持续学习场景，我们可以使用在线学习方法更新记忆：
$$P_{t+1}(k)=\alpha P_t(k)+(1-\alpha )P(D_t|k)$$

其中$\alpha$是学习率，控制新旧记忆的权重。

理论局限性

虽然我们的数学框架提供了一个形式化的基础，但它也存在一些局限性：

1. 表示能力的限制：

   • 目前的知识表示方法难以完全捕捉人类知识的丰富性和灵活性。
   • 分布式向量表示虽然有效，但缺乏符号表示的精确性和可解释性。

2. 计算复杂性：

   • 大规模知识图谱的推理和检索问题本质上是计算密集型的。
   • 随着知识规模的增长，保持推理和检索的效率变得越来越困难。

3. 记忆更新的稳定性-可塑性困境：

   • 记忆系统需要在保持已有知识稳定性的同时，能够灵活地学习新知识，这是一个固有矛盾。
   • 目前的理论还不能完全解决这一困境。

4. 上下文和情境的建模：

   • 知识和记忆往往高度依赖于上下文和情境，但目前的理论框架难以充分建模这种依赖性。
   • 如何表示和利用情境信息仍然是一个开放问题。

竞争范式分析

在AI记忆和知识表示领域，存在多种竞争范式，每种都有其优势和劣势：

1. 符号主义范式：

   • 优势：精确、可解释、支持逻辑推理。
   • 劣势：知识获取困难、难以处理不确定性、扩展性有限。

2. 连接主义范式：

   • 优势：学习能力强、能够处理噪声和不确定性、适合模式识别。
   • 劣势：可解释性差、缺乏推理能力、容易灾难性遗忘。

3. 贝叶斯范式：

   • 优势：能够处理不确定性、提供概率推理框架、支持决策理论。
   • 劣势：计算复杂度高、先验概率难以确定、难以处理大规模知识。

4. 进化范式：

   • 优势：能够通过进化优化知识结构、适合解决复杂优化问题。
   • 劣势：收敛速度慢、缺乏理论保证、可解释性差。

我们提出的长期记忆与知识图谱融合范式试图结合符号主义和连接主义的优势，同时缓解它们的劣势。通过使用知识图谱提供结构化、可解释的知识表示，同时使用神经网络提供学习和模式识别能力，我们希望构建一个更加灵活、强大的记忆系统。

---

3. 架构设计

系统分解

基于我们的理论框架，我们将AI Agent的长期记忆与知识图谱融合系统分解为以下核心组件：

1. 感知与编码模块：

   • 感知接口：从环境获取多模态信息。
   • 特征提取器：从原始感知数据中提取有意义的特征。
   • 记忆编码器：将特征转换为记忆存储格式。

2. 工作记忆模块：

   • 临时存储：保存当前任务相关的信息。
   • 注意力机制：选择性地关注相关信息。
   • 信息处理：对工作记忆中的信息进行基本处理。

3. 长期记忆模块：

   • 陈述性记忆：存储事实和概念知识。
   • 程序性记忆：存储技能和策略。
   • 情景记忆：存储个人经历和事件。

4. 知识图谱引擎：

   • 图存储：高效存储实体、关系和属性。
   • 图查询：支持复杂的图查询和遍历。
   • 图推理：执行基于规则的逻辑推理。

5. 神经记忆组件：

   • 记忆网络：使用神经网络实现可微分的记忆操作。
   • 嵌入层：将符号知识转换为分布式向量表示。
   • 神经推理：执行基于向量的近似推理。

6. 记忆控制器：

   • 检索控制器：决定何时检索、检索什么。
   • 更新控制器：决定何时更新、如何更新记忆。
   • 遗忘控制器：决定何时遗忘、遗忘什么。

7. 推理与决策模块：

   • 混合推理器：结合符号和神经推理方法。
   • 决策制定器：基于记忆和推理结果做出决策。
   • 解释生成器：生成决策过程的解释。

8. 学习模块：

   • 知识获取：从数据和经验中学习新知识。
   • 结构学习：学习和优化记忆系统的结构。
   • 参数学习：优化记忆系统的参数。

这些组件通过精心设计的接口相互连接，形成一个完整的记忆系统。接下来，我们将详细描述这些组件之间的交互关系。

组件交互模型

我们的系统采用分层架构，组件之间的交互遵循以下原则：

1. 自底向上的信息流：感知信息从底层向上流动，经过处理和编码后存储到长期记忆中。
2. 自顶向下的控制流：高层认知过程（如推理、决策）向下发送控制信号，调节低层处理。
3. 双向交互：相邻层之间存在双向连接，允许信息和控制信号在两个方向上流动。
4. 反馈循环：系统包含多个反馈循环，用于学习和优化。

下面是组件交互的详细描述：

1. 感知与编码流程：

   • 感知接口从环境获取原始数据。
   • 特征提取器处理这些数据，提取有意义的特征。
   • 记忆编码器将特征转换为适合存储的格式。
   • 编码后的信息首先进入工作记忆。

2. 工作记忆处理：

   • 注意力机制选择性地关注工作记忆中的相关信息。
   • 工作记忆与长期记忆进行交互，检索相关的已有知识。
   • 信息在工作记忆中被处理和整合，形成更高级的表示。

3. 长期记忆交互：

   • 记忆控制器决定何时将工作记忆中的信息转移到长期记忆。
   • 知识图谱引擎和神经记忆组件协同工作，存储和检索结构化和非结构化信息。
   • 记忆更新和遗忘过程由记忆控制器调节。

4. 推理与决策：

   • 推理与决策模块从长期记忆和工作记忆中检索相关信息。
   • 混合推理器结合符号和神经方法进行推理。
   • 决策制定器基于推理结果做出决策。
   • 解释生成器创建决策过程的解释。

5. 学习与适应：

   • 学习模块监控系统的性能和环境反馈。
   • 知识获取组件从数据和经验中提取新知识。
   • 结构学习和参数学习组件优化记忆系统的结构和参数。

可视化表示

为了更直观地展示系统架构，我们使用Mermaid创建以下图表：

系统总体架构图

知识图谱与神经记忆融合示意图

记忆处理流程图

设计模式应用

在我们的架构设计中，我们应用了以下软件设计模式：

1. 分层模式：系统被组织成层次结构，每一层都有明确的职责，并且只与相邻层交互。这提高了系统的可维护性和可扩展性。

2. 模型-视图-控制器（MVC）模式：我们的架构将数据模型（知识图谱和神经记忆）、视图（感知和动作接口）和控制器（记忆控制器和推理决策模块）分离，提高了系统的模块化程度。

3. 事件驱动模式：系统的许多组件通过事件进行通信，使得组件之间的耦合更加松散，系统更加灵活。

4. 发布-订阅模式：记忆系统中的更新可以被多个组件订阅，实现了高效的一对多通信。

5. 策略模式：我们的记忆检索和更新模块可以使用不同的策略，使得系统可以根据不同的应用场景灵活调整。

6. 装饰器模式：我们可以使用装饰器模式为基础记忆组件添加额外的功能，如缓存、日志、安全检查等。

通过应用这些设计模式，我们的架构不仅满足了功能需求，还具有良好的可维护性、可扩展性和灵活性。

---

4. 实现机制

算法复杂度分析

在设计和实现长期记忆与知识图谱融合系统时，我们需要仔细分析各个算法的时间和空间复杂度，以确保系统的效率和可扩展性。

知识图谱嵌入算法复杂度

知识图谱嵌入是我们系统的核心组件之一，其复杂度分析如下：

1. TransE算法：

   • 时间复杂度：$O(n_e+n_r+n_t\cdot k\cdot d)$，其中$n_e$是实体数，$n_r$是关系数，$n_t$是三元组数，$k$是训练轮数，$d$是嵌入维度。
   • 空间复杂度：$O((n_e+n_r)\cdot d)$，用于存储实体和关系的嵌入向量。

2. DistMult算法：

   • 时间复杂度：与TransE相同，$O(n_e+n_r+n_t\cdot k\cdot d)$。
   • 空间复杂度：与TransE相同，$O((n_e+n_r)\cdot d)$。

3. ComplEx算法：

   • 时间复杂度：$O(n_e+n_r+n_t\cdot k\cdot d)$，但每个嵌入操作涉及复数运算，实际计算量大约是实数嵌入的2倍。
   • 空间复杂度：$O((n_e+n_r)\cdot 2d)$，因为需要存储复数的实部和虚部。

记忆检索算法复杂度

记忆检索是系统的另一个关键操作，其复杂度分析如下：

1. 基于图遍历的检索：

   • 时间复杂度：$O(b^d)$，其中$b$是分支因子，$d$是遍历深度。对于大图，这可能会非常高。
   • 空间复杂度：$O(b^d)$，用于存储遍历队列。

2. 基于向量相似度的检索：

   • 暴力搜索的时间复杂度：$O(n\cdot d)$，其中$n$是向量数量，$d$是维度。
   • 使用近似最近邻算法（如FAISS）的时间复杂度：$O(\log n\cdot d)$。
   • 空间复杂度：$O(n\cdot d)$，用于存储向量索引。

3. 混合检索：

   • 时间复杂度：取决于具体的混合策略，但通常是图遍历和向量检索复杂度的组合。
   • 空间复杂度：需要同时存储图结构和向量索引。

记忆更新算法复杂度

记忆更新是系统持续学习的关键，其复杂度分析如下：

1. 知识图谱更新：

   • 添加新实体/关系的时间复杂度：$O(1)$（使用适当的索引结构）。
   • 添加新三元组的时间复杂度：$O(1)$。
   • 更新推理规则的时间复杂度：取决于规则的复杂性，可能是$O(r\cdot n_t)$，其中$r$是规则数。

2. 神经记忆更新：

   • 在线学习的时间复杂度：$O(d)$每次更新，其中$d$是嵌入维度。
   • 批量学习的时间复杂度：$O(b\cdot k\cdot d)$，其中$b$是批量大小，$k$是训练步数。
   • 记忆网络的写操作时间复杂度：$O(d)$。

通过分析这些复杂度，我们可以做出更明智的设计决策，例如选择合适的算法、优化数据结构、使用近似方法等，以确保系统在实际应用中的效率和可扩展性。

优化代码实现

在本节中，我们提供一个简化但功能完整的AI Agent长期记忆与知识图谱融合系统的Python实现。我们将重点关注核心组件的实现，并提供必要的优化。

首先，我们需要安装一些必要的库：

# 环境安装
!pip install numpy torch networkx transformers faiss-cpu

接下来，我们实现核心组件：

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import networkx as nx
from transformers import BertTokenizer, BertModel
import faiss
from typing import List, Tuple, Dict, Any, Optional
import random

# 设置随机种子以确保结果可复现
random.seed(42)
np.random.seed(42)
torch.manual_seed(42)

class KnowledgeGraph:
    """知识图谱实现"""
    
    def __init__(self):
        self.graph = nx.DiGraph()
        self.entity_to_id = {}
        self.relation_to_id = {}
        self.id_to_entity = {}
        self.id_to_relation = {}
        
    def add_entity(self, entity: str) -> int:
        """添加实体到知识图谱"""
        if entity not in self.entity_to_id:
            entity_id = len(self.entity_to_id)
            self.entity_to_id[entity] = entity_id
            self.id_to_entity[entity_id] = entity
            self.graph.add_node(entity_id, name=entity)
        return self.entity_to_id[entity]
    
    def add_relation(self, relation: str) -> int:
        """添加关系类型到知识图谱"""
        if relation not in self.relation_to_id:
            relation_id = len(self.relation_to_id)
            self.relation_to_id[relation] = relation_id
            self.id_to_relation[relation_id] = relation
        return self.relation_to_id[relation]
    
    def add_triple(self, head: str, relation: str, tail: str) -> None:
        """添加三元组到知识图谱"""
        head_id = self.add_entity(head)
        tail_id = self.add_entity(tail)
        relation_id = self.add_relation(relation)
        
        self.graph.add_edge(head_id, tail_id, relation=relation_id)
    
    def get_triples(self) -> List[Tuple[int, int, int]]:
        """获取所有三元组"""
        triples = []
        for u, v, data in self.graph.edges(data=True):
            triples.append((u, data['relation'], v))
        return triples
    
    def get_neighbors(self, entity: str) -> List[Tuple[str, str]]:
        """获取实体的邻居"""
        if entity not in self.entity_to_id:
            return []
        
        entity_id = self.entity_to_id[entity]
        neighbors = []
        
        # 出边
        for neighbor_id in self.graph.successors(entity_id):
            relation_id = self.graph[entity_id][neighbor_id]['relation']
            relation = self.id_to_relation[relation_id]
            neighbor = self.id_to_entity[neighbor_id]
            neighbors.append((relation, neighbor))
        
        # 入边
        for neighbor_id in self.graph.predecessors(entity_id):
            relation_id = self.graph[neighbor_id][entity_id]['relation']
            relation = self.id_to_relation[relation_id]
            neighbor = self.id_to_entity[neighbor_id]
            neighbors.append((f"~{relation}", neighbor))
        
        return neighbors
    
    def query(self, head: Optional[str] = None, relation: Optional[str] = None, 
              tail: Optional[str] = None) -> List[Tuple[str, str, str]]:
        """查询知识图谱"""
        results = []
        
        # 构建查询模式
        if head is not None and relation is not None and tail is None:
            # (h, r, ?)
            if head in self.entity_to_id and relation in self.relation_to_id:
                head_id = self.entity_to_id[head]
                relation_id = self.relation_to_id[relation]
                
                for neighbor_id in self.graph.successors(head_id):
                    if self.graph[head_id][neighbor_id]['relation'] == relation_id:
                        results.append((head, relation, self.id_to_entity[neighbor_id]))
        
        elif head is None and relation is not None and tail is not None:
            # (?, r, t)
            if relation in self.relation_to_id and tail in self.entity_to_id:
                tail_id = self.entity_to_id[tail]
                relation_id = self.relation_to_id[relation]
                
                for neighbor_id in self.graph.predecessors(tail_id):
                    if self.graph[neighbor_id][tail_id]['relation'] == relation_id:
                        results.append((self.id_to_entity[neighbor_id], relation, tail))
        
        elif head is not None and relation is None and tail is not None:
            # (h, ?, t)
            if head in self.entity_to_id and tail in self.entity_to_id:
                head_id = self.entity_to_id[head]
                tail_id = self.entity_to_id[tail]
                
                if self.graph.has_edge(head_id, tail_id):
                    relation_id = self.graph[head_id][tail_id]['relation']
                    results.append((head, self.id_to_relation[relation_id], tail))
        
        else:
            # 其他情况，遍历所有三元组
            for h, r, t in self.get_triples():
                head_str = self.id_to_entity[h]
                relation_str = self.id_to_relation[r]
                tail_str = self.id_to_entity[t]
                
                if (head is None or head_str == head) and \
                   (relation is None or relation_str == relation) and \
                   (tail is None or tail_str == tail):
                    results.append((head_str, relation_str, tail_str))
        
        return results


class TransEModel(nn.Module):
    """TransE知识图谱嵌入模型"""
    
    def __init__(self, num_entities: int, num_relations: int, embedding_dim: int = 100, 
                 margin: float = 1.0):
        super(TransEModel, self).__init__()
        
        self.num_entities = num_entities
        self.num_relations = num_relations
        self.embedding_dim = embedding_dim
        self.margin = margin
        
        # 初始化实体和关系嵌入
        self.entity_embeddings = nn.Embedding(num_entities, embedding_dim)
        self.relation_embeddings = nn.Embedding(num_relations, embedding_dim)
        
        # 初始化权重
        nn.init.xavier_uniform_(self.entity_embeddings.weight.data)
        nn.init.xavier_uniform_(self.relation_embeddings.weight.data)
        
        # 归一化关系嵌入
        self.relation_embeddings.weight.data.div_(
            self.relation_embeddings.weight.data.norm(p=2, dim=1, keepdim=True)
        )
    
    def forward(self, head_ids: torch.Tensor, relation_ids: torch.Tensor, 
                tail_ids: torch.Tensor) -> torch.Tensor:
        """前向传播，计算三元组的分数"""
        # 获取嵌入
        head_embeds = self.entity_embeddings(head_ids)
        relation_embeds = self.relation_embeddings(relation_ids)
        tail_embeds = self.entity_embeddings(tail_ids)
        
        # 归一化实体嵌入
        head_embeds = head_embeds / head_embeds.norm(p=2, dim=1, keepdim=True)
        tail_embeds = tail_embeds / tail_embeds.norm(p=2, dim=1, keepdim=True)
        
        # 计算分数：-||h + r - t||
        scores = -torch.norm(head_embeds + relation_embeds - tail_embeds, p=2, dim=1)
        
        return scores
    
    def distance(self, head_ids: torch.Tensor, relation_ids: torch.Tensor, 
                 tail_ids: torch.Tensor) -> torch.Tensor:
        """计算距离：||h + r - t||"""
        # 获取嵌入
        head_embeds = self.entity_embeddings(head_ids)
        relation_embeds = self.relation_embeddings(relation_ids)
        tail_embeds = self.entity_embeddings(tail_ids)
        
        # 归一化实体嵌入
        head_embeds = head_embeds / head_embeds.norm(p=2, dim=1, keepdim=True)
        tail_embeds = tail_embeds / tail_embeds.norm(p=2, dim=1, keepdim=True)
        
        # 计算距离
        distances = torch.norm(head_embeds + relation_embeds - tail_embeds, p=2, dim=1)
        
        return distances


class KnowledgeGraphEmbedder:
    """知识图谱嵌入器"""
    
    def __init__(self, kg: KnowledgeGraph, embedding_dim: int = 100, margin: float = 1.0,
                 learning_rate: float = 0.01, batch_size: int = 64, num_epochs: int = 100):
        self.kg = kg
        self.embedding_dim = embedding_dim
        
        # 初始化TransE模型
        self.model = TransEModel(
            num_entities=len(kg.entity_to_id),
            num_relations=len(kg.relation_to_id),
            embedding_dim=embedding_dim,
            margin=margin
        )
        
        self.optimizer = optim.Adam(self.model.parameters(), lr=learning_rate)
        self.batch_size = batch_size
        self.num_epochs = num_epochs
        
        # 检查是否有可用的GPU
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model.to(self.device)
    
    def generate_negative_samples(self, triples: List[Tuple[int, int, int]], 
                                  num_negatives: int = 1) -> List[Tuple[int, int, int, int]]:
        """生成负样本"""
        negative_samples = []
        
        for head, relation, tail in triples:
            for _ in range(num_negatives):
                # 随机决定替换头实体还是尾实体
                if random.random() < 0.5:
                    # 替换头实体
                    new_head = random.randint(0, self.model.num_entities - 1)
                    while new_head == head:
                        new_head = random.randint(0, self.model.num_entities - 1)
                    negative_samples.append((new_head, relation, tail, 0))
                else:
                    # 替换尾实体
                    new_tail = random.randint(0, self.model.num_entities - 1)
                    while new_tail == tail:
                        new_tail = random.randint(0, self.model.num_entities - 1)
                    negative_samples.append((head, relation, new_tail, 0))
            
            # 添加正样本
            negative_samples.append((head, relation, tail, 1))
        
        return negative_samples
    
    def train(self):
        """训练知识图谱嵌入模型"""
        triples = self.kg.get_triples()
        
        if not triples:
            print("知识图谱为空，无法训练嵌入模型")
            return
        
        for epoch in range(self.num_epochs):
            total_loss = 0.0
            
            # 打乱三元组
            random.shuffle(triples)
            
            # 分批处理
            for i in range(0, len(triples), self.batch_size):
                batch_triples = triples[i:i+self.batch_size]
                
                # 生成负样本
                batch_samples = self.generate_negative_samples(batch_triples)
                
                # 转换为张量
                heads = torch.tensor([s[0] for s in batch_samples], dtype=torch.long).to(self.device)
                relations = torch.tensor([s[1] for s in batch_samples], dtype=torch.long).to(self.device)
                tails = torch.tensor([s[2] for s in batch_samples], dtype=torch.long).to(self.device)
                labels = torch.tensor([s[3] for s in batch_samples], dtype=torch.float).to(self.device)
                
                # 清零梯度
                self.optimizer.zero_grad()
                
                # 计算距离
                distances = self.model.distance(heads, relations, tails)
                
                # 计算损失
                # 对于正样本，我们希望距离小；对于负样本，我们希望距离至少为margin
                loss = torch.max(
                    labels * distances + (1 - labels) * (self.model.margin - distances),
                    torch.zeros_like(distances)
                ).mean()
                
                # 反向传播和优化
                loss.backward()
                self.optimizer.step()
                
                total_loss += loss.item()
            
            # 打印每个epoch的损失
            if (epoch + 1) % 10 == 0:
                print(f"Epoch {epoch+1}/{self.num_epochs}, Loss: {total_loss / len(triples):.4f}")
    
    def get_entity_embedding(self, entity: str) -> Optional[np.ndarray]:
        """获取实体的嵌入向量"""
        if entity not in self.kg.entity_to_id:
            return None
        
        entity_id = self.kg.entity_to_id[entity]
        with torch.no_grad():
            embedding = self.model.entity_embeddings(torch.tensor([entity_id], dtype=torch.long).to(self.device))
            embedding = embedding / embedding.norm(p=2, dim=1, keepdim=True)
            return embedding.cpu().numpy()[0]
    
    def get_relation_embedding(self, relation: str) -> Optional[np.ndarray]:
        """获取关系的嵌入向量"""
        if relation not in self.kg.relation_to_id:
            return None
        
        relation_id = self.kg.relation_to_id[relation]
        with torch.no_grad():
            embedding = self.model.relation_embeddings(torch.tensor([relation_id], dtype=torch.long).to(self.device))
            embedding = embedding / embedding.norm(p=2, dim=1, keepdim=True)
            return embedding.cpu().numpy()[0]
    
    def predict_tail(self, head: str, relation: str, top_k: int = 5) -> List[Tuple[str, float]]:
        """预测给定头实体和关系的尾实体"""
        if head not in self.kg.entity_to_id or relation not in self.kg.relation_to_id:
            return []
        
        head_id = self.kg.entity_to_id[head]
        relation_id = self.kg.relation_to_id[relation]
        
        with torch.no_grad():
            # 计算所有可能的尾实体的分数
            head_tensor = torch.tensor([head_id], dtype=torch.long).to(self.device)
            relation_tensor = torch.tensor([relation_id], dtype=torch.long).to(self.device)
            all_tails = torch.arange(self.model.num_entities, dtype=torch.long).to(self.device)
            
            # 扩展头和关系张量以匹配所有尾实体
            head_tensor = head_tensor.expand_as(all_tails)
            relation_tensor = relation_tensor.expand_as(all_tails)
            
            # 计算分数
            scores = self.model(head_tensor, relation_tensor, all_tails)
            
            # 获取top-k结果
            top_scores, top_indices = torch.topk(scores, min(top_k, len(scores)))
            
            # 转换为实体名称和分数的列表
            results = []
            for score, idx in zip(top_scores.cpu().numpy(), top_indices.cpu().numpy()):
                entity = self.kg.id_to_entity[idx]
                results.append((entity, float(score)))
            
            return results


class NeuralMemory:
    """神经记忆组件"""
    
    def __init__(self, embedding_dim: int = 100):
        self.embedding_dim = embedding_dim
        self.memory_vectors = []
        self.memory_metadata = []
        self.index = None
        self.is_index_dirty = False
    
    def add(self, vector: np.ndarray, metadata: Dict[str, Any]) -> None:
        """添加记忆项"""
        if vector.shape[0] != self.embedding_dim:
            raise ValueError(f"向量维度应为{self.embedding_dim}，但得到{vector.shape[0]}")
        
        # 归一化向量
        normalized_vector = vector / np.linalg.norm(vector)
        
        self.memory_vectors.append(normalized_vector)
        self.memory_metadata.append(metadata)
        self.is_index_dirty = True
    
    def build_index(self) -> None:
        """构建FAISS索引"""
        if not self.memory_vectors:
            return
        
        # 将记忆向量转换为numpy数组
        memory_matrix = np.array(self.memory_vectors).astype('float32')