""" ContextEngine — AgentFrame 核心引擎 ==================================== 整合四层上下文保持 + LLM + Embedding: L1 MetaCog 认知层: 任务分解 / 信息缺口 / 检索指令 L2 LandmarkRouter 路由层: landmark 摘要检索 top-k L3 AbsorbedMLA 存储层: 吸收式 MLA 压缩 (L40S 实测 28.4x) L4 KVPager 物理层: 三级换页 + Aura 遗忘曲线 对外接口 (引擎级, 与 API/CLI 解耦): ingest(text, tags) → chunk_id ask(query, chat=True) → AnswerResult stats() → 状态 dict forget(threshold) → 遗忘列表 save(path) / load(path) → 持久化 """ import json import numpy as np from ..config import AgentFrameConfig, MemoryConfig from ..llm import create_llm, LLMProvider from ..embed import create_embedding, EmbeddingProvider from ..memory.store import StateStore from ..memory.incremental import IncrementalKVStore from .quad import QuadLayerAgent from .couple import CouplePrefetcher from .tiers import TieredSummarizer # DeepSeek 工具定义 (给 Agent 的"手") RUN_PYTHON_TOOL = [{ "type": "function", "function": { "name": "run_python", "description": "执行一段 Python 代码并返回 stdout/stderr。写代码前先调用此工具验证。", "parameters": { "type": "object", "properties": { "code": {"type": "string", "description": "要执行的 Python 代码"} }, "required": ["code"] } } }] class AnswerResult: """ask() 的返回""" def __init__(self, answer: str, retrieved: list, scores: list, context_dim: int, directive=None, meta: dict = None): self.answer = answer self.retrieved = retrieved # [(chunk_id, text_preview)] self.scores = scores self.context_dim = context_dim self.directive = directive self.meta = meta or {} def to_dict(self) -> dict: return { "answer": self.answer, "retrieved": self.retrieved, "scores": self.scores, "context_dim": self.context_dim, "meta": self.meta, } MEMORY_DIRECTOR_SYSTEM = ( "你是记忆主管 (MemoryDirector)。分析一段对话, 只输出 JSON 记忆决策:\n" "{\"remember\": [\"值得长期记住的新信息(具体、可复用)\", ...], " "\"forget\": [\"已过时/无关紧要/可丢弃的信息\", ...], " "\"promote_importance\": 0-1, \"reason\": \"一句话理由\"}\n" "规则: remember 只收对未来有复用价值的事实(偏好/身份/参数/数据); " "forget 收闲聊/一次性内容/已被替代信息; 宁可少记不可乱记; " "promote_importance 表示这段对话整体重要性。" ) # ============================================================ # 工具执行安全防护 (P0: LLM 生成代码不可信) # ============================================================ # 危险命令黑名单 (子串匹配, 命中即拒绝执行) DANGEROUS_PATTERNS = [ # 破坏性文件操作 "rm -rf", "rm -fr", "rmdir /", "mkfs", "dd if=", "shred", "> /dev/sd", "mkfs.ext", "fdisk", "parted", # 系统级操作 "shutdown", "reboot", "init 0", "init 6", "poweroff", "halt", "systemctl stop", "systemctl disable", "kill -9 1", # 提权/敏感读取 "sudo ", "su -", "chmod 777 /", "chown -R", "cat /etc/shadow", "cat /etc/passwd", "cat ~/.ssh", "cat .ssh", "cat /root/.ssh", "cat /root/.openclaw", "cat /etc/ssl/private", "base64 -d /etc", "openssl genrsa", "ssh-keygen", # 反弹 shell / 外联 "nc -e", "ncat -e", "bash -i >&", "/dev/tcp/", "curl | sh", "curl|sh", "wget | sh", "wget|sh", "curl -o /etc", "wget -o /etc", "python -c 'import os; os.system", # 挂载/防火墙/网络篡改 "mount ", "umount", "iptables", "nft ", "ip link set", "tcpdump", ] MAX_TOOL_CODE_LEN = 4096 # 单次工具调用代码长度上限 class ToolExecutionError(Exception): """工具执行被安全策略拦截""" pass def _check_code_safety(code: str) -> None: """工具代码安全检查: 长度 + 危险模式黑名单""" if len(code) > MAX_TOOL_CODE_LEN: raise ToolExecutionError(f"代码过长 ({len(code)} > {MAX_TOOL_CODE_LEN} 字符), 已拦截") for pat in DANGEROUS_PATTERNS: if pat in code: raise ToolExecutionError(f"命中危险模式 '{pat}', 已拦截") class ContextEngine: """四层上下文保持引擎""" def __init__(self, config: AgentFrameConfig = None, llm: LLMProvider = None, embedder: EmbeddingProvider = None): self.config = config or AgentFrameConfig.from_env() m: MemoryConfig = self.config.memory # 四层核心 self.agent = QuadLayerAgent( n_layers=m.n_layers, quant_bits=m.quant_bits, top_k=m.top_k, seed=m.seed, reversible=m.reversible, vram_limit_mb=m.vram_limit_mb, ram_limit_mb=m.ram_limit_mb) # 外部能力 self.llm = llm or create_llm( provider=self.config.llm.provider, api_key=self.config.llm.api_key, base_url=self.config.llm.base_url, model=self.config.llm.model, fast_model=self.config.llm.fast_model, timeout=self.config.llm.timeout) self.embedder = embedder or create_embedding("hash", dim=self.agent.store.DIM) # 跨轮共现预取 (colibrì couple_prefetch 移植) self.couple = CouplePrefetcher(top_k=m.couple_k) # 增量持久化 (colibrì kv_persist 移植): 每轮 append, 不重写全量 self.incr_store = None self.incr_path = None # 查询前缀复用 (colibrì kv_prefix 移植) self._last_query = None self._last_retrieved = [] self._prefix_hits = 0 # 会话状态 self.chat_history = [] self.max_history_turns = self.config.api.max_history_turns self.now = 0.0 self._chunk_counter = 0 # ============ 摄入 ============ def ingest(self, text: str, tags: list = None, source: str = None) -> int: """摄入知识块 → 压缩进 KV 缓存""" tags = tags or [] # 文本 → 向量 (真实嵌入) vec = self.embedder.embed(text) # 四层摄取: 用向量作为 hidden states 输入 kv = self.agent.store.encode(vec) kv.heat = 0.9 # 新知识 = 热: force_hot 优先进 VRAM (修复 heat 设置无效问题) self.agent.pager.place(kv, self.now, force_hot=True) # 建 landmark 摘要 k_prime, bias = self.agent.router.build_summary(kv.latent.reshape(1, -1)) cid = kv.chunk_id self.agent.summaries[cid] = (k_prime, bias) # L0/L1/L2 分层摘要 (OpenViking 启发): 写入时生成, 读取时按需展开 tiers = TieredSummarizer.make_tiers(text) self.agent.chunk_meta[cid] = { "tags": tags, "text": text, "source": source or "", "created_at": self.now, "l0": tiers["l0"], "l1": tiers["l1"], "l2_len": tiers["l2_len"], } self._chunk_counter += 1 # 增量追加持久化: 新 chunk 立即落盘 (crash-safe) if self.incr_store is not None: self.incr_store.append( cid, kv.q4, kv.scales, kv.latent, kv.quant_bits, kv.size_bytes, self.agent.chunk_meta[cid]) return cid def ingest_batch(self, chunks: list) -> list: """批量摄入: [{"text":..., "tags":[...]}, ...]""" return [self.ingest(c["text"], c.get("tags", []), c.get("source")) for c in chunks] # ============ 查询 ============ def _exec_tool(self, code: str, timeout: int = 15) -> str: """执行 run_python 工具 (带安全防护)""" import subprocess import sys try: _check_code_safety(code) except ToolExecutionError as e: return f"[安全拦截] {e}" try: r = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=timeout, cwd="/tmp") out = (r.stdout or "") + (("\n[stderr] " + r.stderr) if r.stderr else "") if r.returncode != 0: out = f"[退出码 {r.returncode}] " + out return (out or "(无输出) 执行成功")[:800] except subprocess.TimeoutExpired: return f"[超时 >{timeout}s]" except Exception as e: return f"[执行异常] {e}" def _llm_with_tools(self, system: str, user: str, max_rounds: int = 4, thinking: dict = None, use_fast: bool = False) -> str: """带工具循环的 LLM 调用 (脑+手)""" messages = [{"role": "system", "content": system}, {"role": "user", "content": user}] for _ in range(max_rounds + 1): resp = self.llm.chat(messages, max_tokens=2000, temperature=0.8, thinking=thinking, tools=RUN_PYTHON_TOOL, use_fast=use_fast) tool_calls = resp.get("tool_calls") if not tool_calls: return resp.get("content") or "" # 执行工具并回传 messages.append({"role": "assistant", "content": resp.get("content") or "", "tool_calls": tool_calls}) for tc in tool_calls: fn = tc["function"] try: args = json.loads(fn.get("arguments", "{}")) except Exception: args = {} result = self._exec_tool(args.get("code", "")) messages.append({"role": "tool", "tool_call_id": tc.get("id", ""), "content": result}) return messages[-1].get("content", "") if messages else "" def ask(self, query: str, chat: bool = True, use_fast: bool = False, thinking: dict = None) -> AnswerResult: """查询 + 生成""" # 1. 认知层: 分解 + 检索指令 subtasks = self.agent.metacog.decompose(query) directive = self.agent.metacog.build_directive(subtasks[0], self.agent.chunk_meta) # 1.5 查询前缀复用 (colibrì kv_prefix 移植): # 新 query 与上一条共享长前缀 → 直接复用上次检索结果, 保持热块 reuse = self._prefix_reuse(query) if reuse is not None: selection = reuse else: # 2. 路由层: landmark 检索 (认知层 directive 的 required_chunks 参与加权) q_vec = self.embedder.embed(query) selection = self.agent.router.route( q_vec, q_vec, self.agent.summaries, boost_chunks=directive.required_chunks) self._last_query = query self._last_retrieved = list(selection.chunk_ids) # 2.5 跨轮共现预取: 记录本轮检索 + 预测下轮 → 提升冷块到 RAM self.couple.record(selection.chunk_ids) predicted = self.couple.predict(selection.chunk_ids) if predicted: self.agent.pager.prefetch(predicted, self.now) # 3. 物理层: 访问更新 (遗忘曲线刷新) + Top-K 高精度保护 for cid, score in zip(selection.chunk_ids, selection.scores): self.agent.pager.access(cid, abs(score), self.now) self.agent.metacog.promote(cid, abs(score)) # Top-K 保护 (colibrì): 路由命中的块标记 16bit 高精度, 防量化翻转 self.agent.store.protect_topk(cid) self.now += 1 # 4. 存储层: 聚合检索结果 gathered, hit_ids = [], [] for cid in selection.chunk_ids: kv = (self.agent.pager.vram.get(cid) or self.agent.pager.ram.get(cid) or self.agent.pager.disk.get(cid)) if kv is not None: gathered.append(kv.latent) meta = self.agent.chunk_meta.get(cid, {}) hit_ids.append((cid, meta.get("text", "")[:40])) context_dim = len(gathered[0]) if gathered else 0 # 5. 生成回答 answer = "" if chat: # 分层上下文 (OpenViking L0/L1 启发): 速览 + 概述, 比全文省 token ctx_chunks = [(cid, self.agent.chunk_meta.get(cid, {})) for cid in selection.chunk_ids] ctx_block = TieredSummarizer.build_tiered_context(ctx_chunks) system = ( "你是 AgentFrame 上下文保持系统。下面是从压缩缓存中检索到的知识块，" "回答时优先使用这些知识。如果知识不足就如实说明。" f"\n\n【检索到的知识】\n{ctx_block}") messages = [{"role": "system", "content": system}] messages.extend(self.chat_history[-self.max_history_turns * 2:]) messages.append({"role": "user", "content": query}) resp = self.llm.chat(messages, max_tokens=self.config.llm.max_tokens, temperature=self.config.llm.temperature, use_fast=use_fast) answer = resp.get("content") or "" # 更新会话历史 (老对话由压缩系统兜底) self.chat_history.append({"role": "user", "content": query}) self.chat_history.append({"role": "assistant", "content": answer}) if len(self.chat_history) > self.max_history_turns * 2: self.chat_history = self.chat_history[-self.max_history_turns * 2:] return AnswerResult( answer=answer, retrieved=hit_ids, scores=[round(float(s), 4) for s in selection.scores[:len(hit_ids)]], context_dim=context_dim, directive=directive, meta={"now": self.now, "chunks": len(self.agent.chunk_meta)}) # ============ 前缀复用 (colibrì kv_prefix 移植) ============ def _prefix_reuse(self, query: str): """ 查询前缀复用: 新 query 与上一条共享长前缀 (前 70%+ 且 ≥8 字符) → 直接复用上次检索结果 (跳过重新路由), 并保持上次热块不冷却。 返回 ChunkSelection 或 None (不复用)。 colibrì 对应: fed[] 序列前缀匹配, 跳过重复 prefill。 """ if self._last_query is None or not self._last_retrieved: return None prev, cur = self._last_query, query if len(cur) < 8 or len(prev) < 8: return None # 公共前缀长度 common = 0 for a, b in zip(prev, cur): if a != b: break common += 1 # 前缀须覆盖旧 query 的 70% 且新 query 的 50% (防过度复用) if common >= len(prev) * 0.7 and common >= len(cur) * 0.5: # 保持上次热块: 触碰一次 (刷新遗忘曲线) for cid in self._last_retrieved: kv = (self.agent.pager.vram.get(cid) or self.agent.pager.ram.get(cid) or self.agent.pager.disk.get(cid)) if kv is not None: self.agent.pager.access(cid, 0.5, self.now) self._prefix_hits += 1 from .quad import ChunkSelection return ChunkSelection( chunk_ids=self._last_retrieved, scores=[0.5] * len(self._last_retrieved), weights=[1.0 / len(self._last_retrieved)] * len(self._last_retrieved)) return None # ============ 增量持久化 (colibrì kv_persist 移植) ============ def enable_incremental(self, path: str): """开启增量持久化: 新 chunk 追加到日志, 重启可增量恢复""" self.incr_path = path self.incr_store = IncrementalKVStore(path) return self def load_incremental(self, path: str = None) -> int: """ 从增量日志恢复所有 chunk (colibrì: 重启恢复, 不用重新 prefill)。 重建 chunk_meta + summaries + pager. 返回恢复的 chunk 数。 """ path = path or self.incr_path if not path: return 0 store = IncrementalKVStore(path) if not store.exists(): return 0 recs = store.load() restored = 0 for rec in recs: cid = rec["chunk_id"] # 跳过已存在的 if cid in self.agent.chunk_meta: continue from .quad import CompressedKV kv = CompressedKV( chunk_id=cid, latent=rec["latent"], quant_bits=rec["quant_bits"], size_bytes=rec["size_bytes"], q4=rec["q4"], scales=rec["scales"], ) kv.heat = 0.5 self.agent.chunk_meta[cid] = rec["meta"] # 重建 landmark 摘要 k_prime, bias = self.agent.router.build_summary( rec["latent"].reshape(1, -1)) self.agent.summaries[cid] = (k_prime, bias) self.agent.pager.place(kv, self.now) self._chunk_counter = max(self._chunk_counter, cid + 1) restored += 1 return restored # ============ 工具增强版 ask ============ def ask_with_hands(self, query: str, thinking: dict = None) -> AnswerResult: """带工具循环的查询 (Agent 可自行执行代码验证)""" result = self.ask(query, chat=False) ctx_chunks = [(cid, self.agent.chunk_meta.get(cid, {})) for cid, _preview in result.retrieved] ctx_block = TieredSummarizer.build_tiered_context(ctx_chunks) system = ( "你是 AgentFrame 上下文保持系统。下面是检索到的知识块。" "回答时优先使用这些知识，知识不足就如实说明。" "你有 run_python 工具可以执行代码验证你的方案。" f"\n\n【检索到的知识】\n{ctx_block}") answer = self._llm_with_tools(system, query, thinking=thinking) return AnswerResult(answer=answer, retrieved=result.retrieved, scores=result.scores, context_dim=result.context_dim, meta={"hands": True, "chunks": len(self.agent.chunk_meta)}) # ============ 自主记忆管理 (MemoryDirector) ============ def _parse_memory_decision(self, raw: str) -> dict: """解析 LLM 输出的记忆决策 JSON (容错处理)""" import re m = re.search(r'\{.*\}', raw, re.DOTALL) if not m: return {"remember": [], "forget": [], "promote_importance": 0.5, "reason": ""} try: return json.loads(m.group(0)) except Exception: return {"remember": [], "forget": [], "promote_importance": 0.5, "reason": ""} def autopilot_memory(self, query: str, answer: str) -> dict: """ 模型自主记忆管理: 分析一轮对话, 自动决定记什么/忘什么. 返回决策摘要. """ conversation = f"用户: {query}\n助手: {answer[:800]}" resp = self.llm.chat( [{"role": "system", "content": MEMORY_DIRECTOR_SYSTEM}, {"role": "user", "content": f"对话:\n{conversation}"}], max_tokens=600, temperature=0.3, use_fast=False) decision = self._parse_memory_decision(resp.get("content") or "") # 执行决策 remembered = [] forgotten = [] # 1. remember → 摄入新知识块 (先去重: 与已有块相似则跳过) for item in decision.get("remember", []): if not (isinstance(item, str) and item.strip()): continue # 去重: 检查是否已有高相似度块 dup = False new_vec = self.embedder.embed(item.strip()) for cid, meta in self.agent.chunk_meta.items(): old_vec = self.embedder.embed(meta.get("text", "")) sim = float(new_vec @ old_vec) if sim > 0.8: # 余弦相似度阈值 dup = True break if dup: continue cid = self.ingest(item.strip(), tags=["auto"]) remembered.append((cid, item.strip())) # 2. forget → 按文本匹配降低重要性/删除 forget_list = decision.get("forget", []) if forget_list: for cid, meta in list(self.agent.chunk_meta.items()): text = meta.get("text", "") for f_item in forget_list: if isinstance(f_item, str) and f_item and \ (f_item in text or text in f_item): kv = (self.agent.pager.vram.get(cid) or self.agent.pager.ram.get(cid) or self.agent.pager.disk.get(cid)) if kv: kv.importance *= 0.3 # 大幅降权 kv.heat = 0.1 forgotten.append((cid, f_item)) # 3. promote_importance → 记录整体重要性 imp = float(decision.get("promote_importance", 0.5)) self._last_importance = imp return { "remembered": remembered, "forgotten": forgotten, "importance": imp, "reason": decision.get("reason", ""), } def ask_autopilot(self, query: str) -> AnswerResult: """查询 + 自主记忆管理 (一步到位)""" result = self.ask(query, chat=True) decision = self.autopilot_memory(query, result.answer) result.meta["memory_decision"] = decision return result # ============ 遗忘 / 状态 / 持久化 ============ def forget(self, threshold: float = 0.15) -> list: """按遗忘曲线清理弱记忆 (同步更新内存计数, 防泄漏)""" victims = [] for layer_name in ["disk", "ram", "vram"]: pool = getattr(self.agent.pager, layer_name) for cid, kv in list(pool.items()): strength = self.agent.pager.curve.strength(kv, self.now) if strength < threshold: victims.append(cid) del pool[cid] # 同步扣减占用计数 (vram_used/ram_used), 防漂移 if layer_name == "vram": self.agent.pager.vram_used -= kv.size_bytes elif layer_name == "ram": self.agent.pager.ram_used -= kv.size_bytes self.agent.chunk_meta.pop(cid, None) self.agent.summaries.pop(cid, None) return victims def stats(self) -> dict: return { "chunks_total": len(self.agent.chunk_meta), "layers": self.agent.pager.stats(), "bytes_per_token": round(self.agent.store.bytes_per_token(), 2), "chat_history_turns": len(self.chat_history) // 2, "now": self.now, "compression": "吸收式MLA + INT4 = 28.4x (L40S 真实推理实测)", } def save(self, path: str): """保存引擎状态""" store = StateStore(path) state = { "chunk_meta": self.agent.chunk_meta, "summaries": self.agent.summaries, "chat_history": self.chat_history, "now": self.now, "counter": self._chunk_counter, # pager 状态 "vram": self.agent.pager.vram, "ram": self.agent.pager.ram, "disk": self.agent.pager.disk, } store.save(state) def load(self, path: str): """恢复引擎状态""" store = StateStore(path) state = store.load() if not state: return False self.agent.chunk_meta = {int(k): v for k, v in state.get("chunk_meta", {}).items()} # JSON 序列化后 (k_prime, bias) 变 list, 转回 tuple self.agent.summaries = {int(k): (v[0], v[1]) if isinstance(v, list) and len(v) == 2 else v for k, v in state.get("summaries", {}).items()} self.chat_history = state.get("chat_history", []) self.now = state.get("now", 0.0) self._chunk_counter = state.get("counter", 0) # 反序列化: dict → CompressedKV from .quad import CompressedKV for layer_name in ["vram", "ram", "disk"]: raw = state.get(layer_name, {}) pool = {} for cid, d in raw.items(): if isinstance(d, CompressedKV): pool[int(cid)] = d elif isinstance(d, dict): kv = CompressedKV( chunk_id=int(d.get("chunk_id", int(cid))), latent=np.asarray(d.get("latent", [0.0]), dtype=np.float32), quant_bits=int(d.get("quant_bits", 4)), heat=float(d.get("heat", 0.5)), last_access=float(d.get("last_access", 0.0)), size_bytes=int(d.get("size_bytes", 0)), importance=float(d.get("importance", 0.5)), access_count=int(d.get("access_count", 0)), protected=bool(d.get("protected", False)), reversible=bool(d.get("reversible", False)), sector_id=int(d.get("sector_id", -1)), block_id=int(d.get("block_id", -1)), module_id=int(d.get("module_id", -1)), ) vh = d.get("value_hist") if isinstance(vh, np.ndarray): kv.value_hist = vh pool[int(cid)] = kv setattr(self.agent.pager, layer_name, pool) return True