""" ContextEngine — AgentFrame 核心引擎 ==================================== 整合四层上下文保持 + LLM + Embedding: L1 MetaCog 认知层: 任务分解 / 信息缺口 / 检索指令 L2 LandmarkRouter 路由层: landmark 摘要检索 top-k L3 AbsorbedMLA 存储层: 吸收式 MLA 压缩 (35.6x) 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 .quad import QuadLayerAgent # 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 表示这段对话整体重要性。" ) 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) # 外部能力 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=m.n_layers * 0 or 576) # 会话状态 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 self.agent.pager.place(kv, self.now) # 建 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) self.agent.chunk_meta[cid] = { "tags": tags, "text": text, "source": source or "", "created_at": self.now, } self._chunk_counter += 1 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: 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) # 2. 路由层: landmark 检索 q_vec = self.embedder.embed(query) selection = self.agent.router.route(q_vec, q_vec, self.agent.summaries) # 3. 物理层: 访问更新 (遗忘曲线刷新) 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)) 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: ctx_texts = [] for cid in selection.chunk_ids: meta = self.agent.chunk_meta.get(cid, {}) if meta.get("text"): ctx_texts.append(f"[知识块 {cid}] {meta['text']}") ctx_block = "\n".join(ctx_texts) if ctx_texts else "(无检索到相关知识块)" 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)}) # ============ 工具增强版 ask ============ def ask_with_hands(self, query: str, thinking: dict = None) -> AnswerResult: """带工具循环的查询 (Agent 可自行执行代码验证)""" result = self.ask(query, chat=False) ctx_texts = [] for cid in result.retrieved: meta = self.agent.chunk_meta.get(cid[0], {}) if meta.get("text"): ctx_texts.append(f"[知识块 {cid[0]}] {meta['text']}") ctx_block = "\n".join(ctx_texts) if ctx_texts else "(无检索到相关知识块)" 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] 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 = 35.6x (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()} self.agent.summaries = {int(k): 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