1"""

2 AgentOS v1.14.1 — 长期记忆巩固系统 (Memory Consolidation)。

3

4 受 Letta/MemGPT 三层记忆体系启发，在虚拟内存分页器之上增加主动记忆巩固层。

5 核心机制:

6- Reflection: 定期分析对话历史，提取关键事实、模式、用户偏好

7- Consolidation: 将 Reflection 结果写入长期向量存储

8- Retrieval: 智能检索历史记忆，注入后续对话上下文

9

10 与 memory/pager.py 的关系:

11- pager.py: 被动分页（上下文窗口溢出时 page_out / keyword search page_in）

12- consolidation.py: 主动巩固（定期分析 → 提取 → 向量化存储）

13"""

14

15 from __future__ import annotations

16

17 import asyncio

18 import json

19 import time

20 import uuid

21 from dataclasses import dataclass, field

22 from enum import Enum

23 from typing import (

24 Any, Callable, Dict, List, Optional, Set, Tuple, Union,

25)

26

27

28# ── Memory Data Models ──────────────────────

29

30

31 class MemoryType(str, Enum):

32 """记忆类型。"""

33 FACT = "fact" # 事实性信息

34 PREFERENCE = "preference" # 用户偏好

35 PATTERN = "pattern" # 行为模式

36 DECISION = "decision" # 决策记录

37 LESSON = "lesson" # 经验教训

38 CONTEXT = "context" # 上下文摘要

39

40

41 class MemoryImportance(str, Enum):

42 """记忆重要性。"""

43 LOW = "low"

44 MEDIUM = "medium"

45 HIGH = "high"

46 CRITICAL = "critical"

47

48

49@dataclass

50 class MemoryFragment:

51 """记忆片段 — 从对话中提取的原子事实。"""

52

53 memory_id: str = field(default_factory=lambda: f"mem-{uuid.uuid4().hex[:12]}")

54 memory_type: MemoryType = MemoryType.FACT

55 content: str = ""

56 importance: MemoryImportance = MemoryImportance.MEDIUM

57 source_messages: List[int] = field(default_factory=list) # 源自哪些消息

58 tags: List[str] = field(default_factory=list)

59 confidence: float = 1.0 # 0.0~1.0 置信度

60 created_at: float = field(default_factory=time.time)

61 last_accessed: float = field(default_factory=time.time)

62 access_count: int = 0

63 embedding: Optional[List[float]] = None # 向量嵌入（惰性计算）

64 metadata: Dict[str, Any] = field(default_factory=dict)

65

66 def to_dict(self) -> dict:

67 return {

68 "memory_id": self.memory_id,

69 "memory_type": self.memory_type.value,

70 "content": self.content,

71 "importance": self.importance.value,

72 "source_messages": self.source_messages,

73 "tags": self.tags,

74 "confidence": self.confidence,

75 "created_at": self.created_at,

76 "last_accessed": self.last_accessed,

77 "access_count": self.access_count,

78 "metadata": self.metadata,

79 }

80

81 @classmethod

82 def from_dict(cls, d: dict) -> "MemoryFragment":

83 return cls(

84 memory_id=d.get("memory_id", ""),

85 memory_type=MemoryType(d.get("memory_type", "fact")),

86 content=d.get("content", ""),

87 importance=MemoryImportance(d.get("importance", "medium")),

88 source_messages=d.get("source_messages", []),

89 tags=d.get("tags", []),

90 confidence=d.get("confidence", 1.0),

91 created_at=d.get("created_at", time.time()),

92 last_accessed=d.get("last_accessed", time.time()),

93 access_count=d.get("access_count", 0),

94 metadata=d.get("metadata", {}),

95 )

96

97 def touch(self) -> None:

98 """更新访问时间。"""

99 self.last_accessed = time.time()

100 self.access_count += 1

101

102

103@dataclass

104 class ReflectionResult:

105 """一次 Reflection 的输出。"""

106

107 fragments: List[MemoryFragment] = field(default_factory=list)

108 summary: str = "" # 会话级摘要

109 contradictions: List[Tuple[str, str]] = field(default_factory=list) # 新旧矛盾

110 deprecated_ids: List[str] = field(default_factory=list) # 需淘汰的旧记忆

111 user_profile_update: Dict[str, Any] = field(default_factory=dict)

112 timestamp: float = field(default_factory=time.time)

113

114 @property

115 def total_fragments(self) -> int:

116 return len(self.fragments)

117

118 @property

119 def has_insights(self) -> bool:

120 return bool(self.fragments or self.summary or self.user_profile_update)

121

122

123# ── Vector Store Interface ──────────────────

124

125

126 class VectorStoreBackend:

127 """向量存储后端抽象。

128

129 支持多种后端: 内存、FAISS、Chroma、Pinecone 等。

130 """

131

132 async def add(

133 self,

134 fragments: List[MemoryFragment],

135 embeddings: List[List[float]],

136 ) -> List[str]:

137 """批量添加记忆片段（含嵌入向量）。返回 memory_ids。"""

138 raise NotImplementedError

139

140 async def search(

141 self,

142 query_embedding: List[float],

143 top_k: int = 10,

144 filter_types: Optional[List[MemoryType]] = None,

145 min_importance: MemoryImportance = MemoryImportance.LOW,

146 ) -> List[Tuple[MemoryFragment, float]]:

147 """向量相似度搜索。返回 (fragment, score)。"""

148 raise NotImplementedError

149

150 async def delete(self, memory_ids: List[str]) -> int:

151 """删除指定记忆。返回删除数量。"""

152 raise NotImplementedError

153

154 async def count(self) -> int:

155 """记忆总数。"""

156 raise NotImplementedError

157

158

159 class InMemoryVectorStore(VectorStoreBackend):

160 """内存向量存储（开发/测试用）。"""

161

162 def __init__(self):

163 self._fragments: Dict[str, MemoryFragment] = {}

164 self._embeddings: Dict[str, List[float]] = {}

165

166 async def add(

167 self,

168 fragments: List[MemoryFragment],

169 embeddings: List[List[float]],

170 ) -> List[str]:

171 ids = []

172 for frag, emb in zip(fragments, embeddings):

173 self._fragments[frag.memory_id] = frag

174 self._embeddings[frag.memory_id] = emb

175 ids.append(frag.memory_id)

176 return ids

177

178 async def search(

179 self,

180 query_embedding: List[float],

181 top_k: int = 10,

182 filter_types: Optional[List[MemoryType]] = None,

183 min_importance: MemoryImportance = MemoryImportance.LOW,

184 ) -> List[Tuple[MemoryFragment, float]]:

185 results = []

186 importance_rank = {

187 MemoryImportance.LOW: 0,

188 MemoryImportance.MEDIUM: 1,

189 MemoryImportance.HIGH: 2,

190 MemoryImportance.CRITICAL: 3,

191 }

192 min_rank = importance_rank[min_importance]

193

194 for mid, emb in self._embeddings.items():

195 frag = self._fragments[mid]

196 # Type filter

197 if filter_types and frag.memory_type not in filter_types:

198 continue

199 # Importance filter

200 if importance_rank[frag.importance] < min_rank:

201 continue

202 # Cosine similarity

203 score = self._cosine_similarity(query_embedding, emb)

204 results.append((frag, score))

205

206 results.sort(key=lambda x: x[1], reverse=True)

207 return results[:top_k]

208

209 async def delete(self, memory_ids: List[str]) -> int:

210 count = 0

211 for mid in memory_ids:

212 if mid in self._fragments:

213 del self._fragments[mid]

214 self._embeddings.pop(mid, None)

215 count += 1

216 return count

217

218 async def count(self) -> int:

219 return len(self._fragments)

220

221 @staticmethod

222 def _cosine_similarity(a: List[float], b: List[float]) -> float:

223 dot = sum(x * y for x, y in zip(a, b))

224 norm_a = sum(x * x for x in a) ** 0.5

225 norm_b = sum(x * x for x in b) ** 0.5

226 if norm_a == 0 or norm_b == 0:

227 return 0.0

228 return dot / (norm_a * norm_b)

229

230

231# ── Embedding Provider ──────────────────────

232

233

234 class EmbeddingProvider:

235 """嵌入向量生成器抽象。"""

236

237 async def embed(self, texts: List[str]) -> List[List[float]]:

238 """批量生成嵌入向量。"""

239 raise NotImplementedError

240

241 async def embed_single(self, text: str) -> List[float]:

242 """单条文本嵌入。"""

243 results = await self.embed([text])

244 return results[0]

245

246

247 class SimpleHashEmbedding(EmbeddingProvider):

248 """简单位次嵌入（开发/测试用，非语义向量）。

249

250 用字符 n-gram 哈希作为伪嵌入，提供基本的相似度。

251 生产环境应替换为 OpenAI/Cohere 等真实嵌入模型。

252 """

253

254 def __init__(self, dim: int = 128):

255 self.dim = dim

256

257 async def embed(self, texts: List[str]) -> List[List[float]]:

258 results = []

259 for text in texts:

260 vec = [0.0] * self.dim

261 # Character 3-gram hashing

262 for i in range(len(text) - 2):

263 gram = text[i:i+3]

264 h = hash(gram) % self.dim

265 vec[h] += 1.0

266 # L2 normalize

267 norm = sum(v * v for v in vec) ** 0.5

268 if norm > 0:

269 vec = [v / norm for v in vec]

270 results.append(vec)

271 return results

272

273

274# ── Reflection Engine ───────────────────────

275

276

277 class ReflectionConfig:

278 """Reflection 触发配置。"""

279

280 def __init__(

281 self,

282 min_messages_since_last: int = 10,

283 min_seconds_since_last: float = 300.0, # 5 分钟

284 max_conversation_turns: int = 50,

285 auto_reflect: bool = True,

286 ):

287 self.min_messages_since_last = min_messages_since_last

288 self.min_seconds_since_last = min_seconds_since_last

289 self.max_conversation_turns = max_conversation_turns

290 self.auto_reflect = auto_reflect

291

292

293 class ReflectionEngine:

294 """记忆反思引擎。

295

296 定期分析对话历史，提取:

297 - Facts: 用户提到的具体信息

298 - Preferences: 用户偏好与习惯

299 - Patterns: 反复出现的行为模式

300 - Lessons: 从错误中学到的经验

301

302 Usage:

303 engine = ReflectionEngine(llm_reflect_fn, vector_store, embedding_provider)

304 # 在 agent loop 中定期调用

305 should_reflect = engine.should_reflect(message_count)

306 if should_reflect:

307 result = await engine.reflect(messages_history)

308 """

309

310 def __init__(

311 self,

312 llm_reflect_fn: Optional[Callable[[List[dict], str], Any]] = None,

313 vector_store: Optional[Any] = None,

314 embedding_provider: Optional[Any] = None,

315 config: Optional[ReflectionConfig] = None,

316 ):

317 """

318 Args:

319 llm_reflect_fn: LLM 调用函数，签名 (messages, prompt) -> reflection_text

320 vector_store: 向量存储后端

321 embedding_provider: 嵌入向量生成器

322 config: 触发配置

323 """

324 self._llm_reflect = llm_reflect_fn

325 self._vector_store = vector_store

326 self._embedding_provider = embedding_provider

327 self.config = config or ReflectionConfig()

328 self._last_reflection_time: float = 0.0

329 self._message_count_since_reflection: int = 0

330 self._reflection_count: int = 0

331

332 def should_reflect(self, current_message_count: int) -> bool:

333 """判断是否应该触发 Reflection。"""

334 if not self.config.auto_reflect:

335 return False

336 if self._reflection_count == 0 and current_message_count >= 5:

337 return True # 首次在 5 条消息后触发

338 msg_check = (

339 self._message_count_since_reflection

340 >= self.config.min_messages_since_last

341 )

342 time_check = (

343 time.time() - self._last_reflection_time

344 >= self.config.min_seconds_since_last

345 )

346 return msg_check or time_check

347

348 async def reflect(

349 self,

350 messages: List[dict],

351 existing_fragments: Optional[List[MemoryFragment]] = None,

352 ) -> ReflectionResult:

353 """执行一次 Reflection。

354

355 Args:

356 messages: 对话历史（dict 列表，含 role/content）

357 existing_fragments: 已有的记忆片段（用于矛盾检测）

358

359 Returns:

360 ReflectionResult 含新提取的记忆片段

361 """

362 self._last_reflection_time = time.time()

363 self._reflection_count += 1

364

365 # 1. 构建 Reflection prompt

366 prompt = self._build_reflection_prompt(messages, existing_fragments)

367

368 # 2. 调用 LLM 提取记忆

369 reflection_text = await self._llm_reflect(messages, prompt)

370

371 # 3. 解析 LLM 输出

372 result = self._parse_reflection_output(reflection_text, len(messages))

373

374 # 4. 生成嵌入向量

375 if result.fragments:

376 texts = [f.content for f in result.fragments]

377 embeddings = await self._embedding_provider.embed(texts)

378 for frag, emb in zip(result.fragments, embeddings):

379 frag.embedding = emb

380

381 # 5. 存入向量库

382 if result.fragments:

383 await self._vector_store.add(result.fragments, embeddings)

384

385 # 6. 淘汰旧记忆

386 if result.deprecated_ids:

387 await self._vector_store.delete(result.deprecated_ids)

388

389 # Reset counter

390 self._message_count_since_reflection = 0

391

392 return result

393

394 def record_message(self) -> None:

395 """记录一条新消息（用于计数触发）。"""

396 self._message_count_since_reflection += 1

397

398 async def retrieve_relevant(

399 self,

400 query: str,

401 top_k: int = 5,

402 filter_types: Optional[List[MemoryType]] = None,

403 ) -> List[MemoryFragment]:

404 """检索与查询相关的记忆。

405

406 Args:

407 query: 查询文本

408 top_k: 返回数量

409 filter_types: 按类型过滤

410

411 Returns:

412 相关记忆片段列表

413 """

414 query_embedding = await self._embedding_provider.embed_single(query)

415 results = await self._vector_store.search(

416 query_embedding,

417 top_k=top_k,

418 filter_types=filter_types,

419 )

420 fragments = []

421 for frag, score in results:

422 frag.touch()

423 fragments.append(frag)

424 return fragments

425

426 def _build_reflection_prompt(

427 self,

428 messages: List[dict],

429 existing_fragments: Optional[List[MemoryFragment]] = None,

430 ) -> str:

431 """构建 Reflection prompt。"""

432 existing_str = ""

433 if existing_fragments:

434 existing_items = [

435 f"- [{f.memory_type.value}] {f.content}"

436 for f in existing_fragments[:20]

437 ]

438 existing_str = (

439 "\n\nExisting memories:\n" + "\n".join(existing_items)

440 )

441

442 return f"""You are a memory consolidation system. Analyze the conversation and extract:

443

444 1. FACTS: Specific information mentioned (names, dates, numbers, tools used, decisions made)

445 2. PREFERENCES: User preferences, likes, dislikes, habits

446 3. PATTERNS: Repeated behaviors, common workflows, recurring topics

447 4. LESSONS: What went wrong, what worked, what to avoid next time

448

449 For each extracted item, assign:

450- type: "fact" | "preference" | "pattern" | "decision" | "lesson"

451- importance: "low" | "medium" | "high" | "critical"

452- confidence: 0.0 to 1.0

453

454{existing_str}

455

456 Output JSON array only:

457[{{"type": "...", "content": "...", "importance": "...", "confidence": 0.9, "tags": ["..."]}}]

458

459 If nothing significant to extract, output empty array: []"""

460

461 def _parse_reflection_output(

462 self,

463 text: str,

464 source_msg_count: int,

465 ) -> ReflectionResult:

466 """解析 LLM 输出的 JSON。"""

467 result = ReflectionResult()

468

469 try:

470 # Extract JSON array

471 start = text.find("[")

472 end = text.rfind("]")

473 if start >= 0 and end > start:

474 json_str = text[start:end + 1]

475 items = json.loads(json_str)

476 for item in items:

477 frag = MemoryFragment(

478 memory_type=MemoryType(item.get("type", "fact")),

479 content=item.get("content", ""),

480 importance=MemoryImportance(

481 item.get("importance", "medium")

482 ),

483 confidence=float(item.get("confidence", 1.0)),

484 tags=item.get("tags", []),

485 source_messages=list(

486 range(

487 max(0, source_msg_count - 20),

488 source_msg_count,

489 )

490 ),

491 )

492 if frag.content.strip():

493 result.fragments.append(frag)

494 except (json.JSONDecodeError, KeyError, ValueError):

495 pass

496

497 return result

498

499 @property

500 def stats(self) -> Dict[str, Any]:

501 return {

502 "reflection_count": self._reflection_count,

503 "last_reflection_time": self._last_reflection_time,

504 "messages_since_last": self._message_count_since_reflection,

505 }

506

507 # ── Persistence (v1.14.9) ────────────────

508

509 def get_state(self) -> Dict[str, Any]:

510 """Export ReflectionEngine state for persistence."""

511 return {

512 "reflection_count": self._reflection_count,

513 "last_reflection_time": self._last_reflection_time,

514 "message_count_since_reflection": self._message_count_since_reflection,

515 "vector_store_fragments": {

516 mid: frag.to_dict()

517 for mid, frag in self._vector_store._fragments.items()

518 } if hasattr(self._vector_store, '_fragments') and self._vector_store else {},

519 "vector_store_embeddings": {

520 mid: list(emb) if emb else []

521 for mid, emb in (self._vector_store._embeddings.items()

522 if hasattr(self._vector_store, '_embeddings') and self._vector_store

523 else {}.items())

524 },

525 }

526

527 def restore_state(self, state: Dict[str, Any]) -> None:

528 """Restore ReflectionEngine from a persisted snapshot."""

529 self._reflection_count = state.get("reflection_count", 0)

530 self._last_reflection_time = state.get("last_reflection_time", 0.0)

531 self._message_count_since_reflection = state.get("message_count_since_reflection", 0)

532

533 if self._vector_store and hasattr(self._vector_store, '_fragments'):

534 self._vector_store._fragments.clear()

535 self._vector_store._embeddings.clear()

536 for mid, frag_data in state.get("vector_store_fragments", {}).items():

537 self._vector_store._fragments[mid] = MemoryFragment.from_dict(frag_data)

538 for mid, emb in state.get("vector_store_embeddings", {}).items():

539 self._vector_store._embeddings[mid] = emb

540

541

542# ── Memory Context Injector ─────────────────

543

544

545 class MemoryContextInjector:

546 """记忆上下文注入器。

547

548 在每次 Agent 对话开始时，自动检索相关历史记忆，

549 注入到 system prompt 或上下文中。

550

551 Usage:

552 injector = MemoryContextInjector(reflection_engine)

553 context = await injector.build_context("user query here")

554 messages.insert(0, {"role": "system", "content": context})

555 """

556

557 def __init__(

558 self,

559 reflection_engine: ReflectionEngine,

560 max_context_length: int = 2000,

561 max_fragments: int = 5,

562 ):

563 self._engine = reflection_engine

564 self.max_context_length = max_context_length

565 self.max_fragments = max_fragments

566

567 async def build_context(

568 self,

569 query: str,

570 include_types: Optional[List[MemoryType]] = None,

571 ) -> str:

572 """构建上下文注入文本。"""

573 fragments = await self._engine.retrieve_relevant(

574 query,

575 top_k=self.max_fragments,

576 filter_types=include_types,

577 )

578

579 if not fragments:

580 return ""

581

582 lines = ["[Relevant Memories]"]

583 for frag in fragments:

584 lines.append(

585 f"- [{frag.memory_type.value}] {frag.content}"

586 f" (confidence: {frag.confidence:.0%})"

587 )

588

589 context = "\n".join(lines)

590 if len(context) > self.max_context_length:

591 context = context[:self.max_context_length] + "..."

592

593 return context

594

595 async def build_condensed_context(

596 self,

597 query: str,

598 ) -> str:

599 """构建紧凑上下文（仅高重要性记忆）。"""

600 fragments = await self._engine.retrieve_relevant(

601 query,

602 top_k=self.max_fragments,

603 )

604 # Filter: only HIGH/CRITICAL

605 important = [

606 f for f in fragments

607 if f.importance in (MemoryImportance.HIGH, MemoryImportance.CRITICAL)

608 ]

609 if not important:

610 return ""

611

612 lines = ["[Key Context]"]

613 for frag in important[:3]:

614 lines.append(f"- {frag.content}")

615

616 return "\n".join(lines)

617

618

619# ── Memory Consolidation Pipeline ───────────

620

621

622 class MemoryConsolidationPipeline:

623 """记忆巩固流水线（一键集成）。

624

625 组合 ReflectionEngine + MemoryContextInjector，

626 提供开箱即用的记忆系统。

627

628 Usage:

629 pipeline = MemoryConsolidationPipeline(llm_fn)

630 # 在 agent loop 中:

631 pipeline.record_message()

632 if pipeline.should_reflect():

633 await pipeline.reflect(messages)

634 context = await pipeline.get_context(user_query)

635 """

636

637 def __init__(

638 self,

639 llm_reflect_fn: Callable,

640 vector_store: Optional[VectorStoreBackend] = None,

641 embedding_provider: Optional[EmbeddingProvider] = None,

642 config: Optional[ReflectionConfig] = None,

643 ):

644 self._vector_store = vector_store or InMemoryVectorStore()

645 self._embedding_provider = embedding_provider or SimpleHashEmbedding(128)

646 self._reflection_engine = ReflectionEngine(

647 llm_reflect_fn=llm_reflect_fn,

648 vector_store=self._vector_store,

649 embedding_provider=self._embedding_provider,

650 config=config,

651 )

652 self._injector = MemoryContextInjector(self._reflection_engine)

653

654 def record_message(self) -> None:

655 self._reflection_engine.record_message()

656

657 def should_reflect(self) -> bool:

658 return self._reflection_engine.should_reflect(

659 self._reflection_engine._message_count_since_reflection

660 )

661

662 async def reflect(self, messages: List[dict]) -> ReflectionResult:

663 return await self._reflection_engine.reflect(messages)

664

665 async def get_context(self, query: str) -> str:

666 return await self._injector.build_context(query)

667

668 async def get_condensed_context(self, query: str) -> str:

669 return await self._injector.build_condensed_context(query)

670

671 @property

672 def stats(self) -> Dict[str, Any]:

673 return {

674 "reflection": self._reflection_engine.stats,

675 "total_memories": asyncio.get_event_loop().run_until_complete(

676 self._vector_store.count()

677 ) if asyncio.get_event_loop().is_running() else 0,

678 }

679

680 # ── Persistence (v1.14.9) ────────────────

681

682 def get_state(self) -> Dict[str, Any]:

683 """Export consolidation pipeline state for persistence. Delegates to ReflectionEngine."""

684 return self._reflection_engine.get_state()

685

686 def restore_state(self, state: Dict[str, Any]) -> None:

687 """Restore consolidation pipeline from a persisted snapshot."""

688 self._reflection_engine.restore_state(state)