跳到主要内容
登录

github/dsh-companion-memory

公开

DeepSeek Harness 本地聊天记忆：按来源保存聊天原文，关键词 + 向量检索，可选大模型重排与检索摘要。源码公开，仅限非商用

浏览文件
main

分支

host.js

193 行 · 10.6 KB历史下载
1import { homedir } from 'node:os';
2import { join, resolve } from 'node:path';
3import { existsSync, readFileSync, writeFileSync, renameSync } from 'node:fs';
4import { randomUUID, randomBytes } from 'node:crypto';
5import { RetrievalMemory, SUMMARY_BATCH } from './store.js';
6import { recall } from './recall.js';
7import { normalizeSettings } from './config.js';
8import { detectLocalEmbedding, createEmbedder } from './models.js';
9import { compileRules, resolveOrigin, viewOf, normalizeOrigin } from './origin.js';
10import { turnProjection, PROJECTION_KEY, textOf, isUser, setSpeakTools, setHearTools } from './projection.js';
11import { registerManagement } from './management.js';
12import { registerTools } from './tools.js';
13
14const USER = 'local-user';
15/** After DSH's own runtime facts (orders 100–210), so they stay first. */
16const CONTEXT_ORDER = 300;
17
18// Fallback only, for hosts without runtime prompt context. Uses the published wire shape.
19function evidenceMessage(text, hits) {
20 return Object.freeze({ id: randomUUID(), role: 'user',
21 content: Object.freeze([Object.freeze({ type: 'text', text })]),
22 source: Object.freeze({ kind: 'dsh-companion-memory', records: hits.map(hit => hit.record_id) }) });
23}
24
25function writePrivate(path, text) {
26 const temporary = path + '.' + randomUUID() + '.tmp';
27 writeFileSync(temporary, text, { mode: 0o600 });
28 renameSync(temporary, path);
29}
30/** Local token for bridges calling the HTTP API; created once, readable only by this OS user. */
31function loadApiToken(dir) {
32 const path = join(dir, 'api-token');
33 const existing = existsSync(path) ? readFileSync(path, 'utf8').trim() : '';
34 if (/^[a-f0-9]{64}$/.test(existing)) return existing;
35 const token = randomBytes(32).toString('hex');
36 writePrivate(path, token + '\n');
37 return token;
38}
39
40export function applyMemory(ctx, config = {}) {
41 const dir = resolve(config.dir || join(process.env.DSH_PROFILE_DIR || process.env.DSH_HOME || join(homedir(), '.dsh'), 'data', 'companion-memory'));
42 const settingsPath = join(dir, 'settings.json');
43 const stored = existsSync(settingsPath) ? JSON.parse(readFileSync(settingsPath, 'utf8')) : {};
44 let settings = normalizeSettings({ ...config, ...stored });
45 let rules = compileRules(settings.originRules);
46 setSpeakTools(settings.speakTools); setHearTools(settings.hearTools);
47 const namespace = String(config.namespace || 'dsh');
48 if (!namespace || namespace.length > 240) throw new Error('无效 namespace');
49 const warn = code => { try { ctx.logger?.warn?.('dsh-companion-memory: ' + code); } catch {} };
50 const memory = new RetrievalMemory({ dir, ...settings, logger: warn });
51 const apiToken = loadApiToken(dir);
52 let disposed = false;
53 let operationAbort = new AbortController();
54 const jobs = new Set();
55 const schedule = job => { jobs.add(job); job.catch(() => warn('index-unavailable')).finally(() => jobs.delete(job)); };
56 // Summaries first, so their vectors are computed in the same pass. New turns wait until a full
57 // batch of SUMMARY_BATCH has accumulated: one request per ten records, not one per turn.
58 const indexBatch = () => {
59 const signal = operationAbort.signal;
60 const summarize = memory.summarizer && memory.summariesPending() >= SUMMARY_BATCH
61 ? memory.summarize({ signal }).catch(() => warn('summary-unavailable')) : Promise.resolve();
62 schedule(summarize.then(() => memory.index({ limit: 32, signal })));
63 };
64
65 const originOf = sessionId => resolveOrigin(sessionId, { rules, ownerIds: settings.ownerIds,
66 unknownOrigin: settings.unknownOrigin, binding: id => memory.binding(id) });
67 const access = origin => ({ namespace, userId: USER, scope: origin });
68 const view = origin => ({ ...access(origin), ...viewOf(origin, settings) });
69
70 /** Store completed turns not seen before; returns the source keys of the session's recent tail. */
71 const importTail = session => {
72 const state = ctx.sessionProjections.stateOf(session, PROJECTION_KEY);
73 if (!state) return new Set();
74 const origin = originOf(session.id);
75 const keys = new Set();
76 for (const row of state.completed) {
77 const sourceKey = `${session.id}:${row.seq}`;
78 keys.add(sourceKey);
79 if (!settings.autoRecord || memory.known(access(origin), sourceKey)) continue;
80 const user = memory.stripBoilerplate(row.user, { count: true });
81 const heard = row.heard ? memory.freshLines(access(origin), row.heard) : '';
82 const said = [row.assistant, row.spoken].filter(Boolean).join('\n');
83 if (!heard && !said) continue;
84 memory.put({ ...access(origin), sourceKey, createdAt: row.time,
85 content: [user && `用户：${user}`, heard && `收到的消息：\n${heard}`, said && `助手：${said}`].filter(Boolean).join('\n'),
86 source: { platform: 'deepseek-harness', session: String(session.id), turn: row.turn, seq: row.seq } });
87 }
88 return keys;
89 };
90
91 // Retrieval result per agent, read by the runtime-context provider at request admission.
92 const evidence = new Map();
93 const agentKey = agent => String(agent?.id ?? agent?.session?.id ?? '');
94 const remember = (key, value) => { evidence.delete(key); evidence.set(key, value); if (evidence.size > 256) evidence.delete(evidence.keys().next().value); };
95 let contextReady = false;
96 ctx.inject?.(['systemPrompt'], scoped => {
97 scoped.effect(() => {
98 // Runtime context is logged after the cached history and only when its text changes,
99 // so identical evidence across the steps of one turn costs nothing and keeps the KV prefix.
100 const dispose = scoped.systemPrompt.context({ name: 'dsh-companion-memory', order: CONTEXT_ORDER, interpolate: false,
101 text: context => (settings.autoRetrieve && evidence.get(agentKey(context?.agent))?.text) || '' });
102 contextReady = true;
103 return () => { contextReady = false; dispose?.(); };
104 });
105 });
106
107 ctx.sessionProjections.register(turnProjection);
108 ctx.on('session/event', (session, event) => {
109 if (disposed || !settings.autoRecord || event.type !== 'turn/end' || event.data.reason?.kind !== 'completed') return;
110 try { importTail(session); indexBatch(); } catch { warn('record-unavailable'); }
111 });
112 ctx.on('agent/pre-step', async ({ agent, signal }, next) => {
113 const decision = await next();
114 if (disposed || decision.kind === 'reject' || signal?.aborted || !agent?.session) return decision;
115 // Only a newly admitted user message triggers work; tool-continuation steps keep this turn's evidence.
116 const latest = decision.messages.findLast(message => isUser(message) && textOf(message));
117 const key = agentKey(agent);
118 if (!latest || (latest.id !== undefined && evidence.get(key)?.messageId === latest.id)) return decision;
119 let recent = new Set();
120 try { recent = importTail(agent.session); indexBatch(); } catch { warn('record-unavailable'); }
121 if (!settings.autoRetrieve) return decision;
122 try {
123 const origin = originOf(agent.session.id);
124 const query = memory.stripBoilerplate(textOf(latest)).slice(-8000);
125 const joinedSignal = signal ? AbortSignal.any([signal, operationAbort.signal]) : operationAbort.signal;
126 // Turns still in this session's tail are already in the model's context; do not repeat them.
127 const { hits, text } = query ? await recall(memory, query, view(origin),
128 { topK: settings.topK, maxChars: settings.maxContextChars, signal: joinedSignal, exclude: recent, here: origin })
129 : { hits: [], text: '' };
130 remember(key, { messageId: latest.id, text });
131 if (contextReady || !text) return decision;
132 return { ...decision, messages: [...decision.messages, evidenceMessage(text, hits)] };
133 } catch { signal?.throwIfAborted(); warn('retrieval-unavailable'); return decision; }
134 });
135
136 const service = {
137 memory, dir, namespace, apiToken, originOf, access, view,
138 settings: () => settings, signal: () => operationAbort.signal,
139 collections: () => memory.collections(namespace, USER),
140 normalizeOrigin: value => normalizeOrigin(value, settings.ownerIds),
141 /** Store text under an origin and index it in the background. */
142 note(origin, { sourceKey = 'note:' + randomUUID(), text, createdAt = Date.now(), source = {} }) {
143 const result = memory.put({ ...access(origin), sourceKey, createdAt, content: text, source });
144 indexBatch(); return result;
145 },
146 save(input) {
147 const next = normalizeSettings(input);
148 writePrivate(settingsPath, JSON.stringify(next, null, 2) + '\n');
149 operationAbort.abort(); operationAbort = new AbortController();
150 memory.configure(next); settings = next;
151 rules = compileRules(next.originRules); setSpeakTools(next.speakTools); setHearTools(next.hearTools); evidence.clear();
152 // Origin and sharing changes apply to future turns; existing records keep their scope.
153 indexBatch(); return settings;
154 },
155 async index() {
156 let indexed = 0;
157 const signal = operationAbort.signal;
158 // Backfill summaries for everything first; a failing model leaves the rest for next time.
159 try {
160 while (!disposed) {
161 const batch = await memory.summarize({ limit: 50, signal });
162 if (batch.disabled || batch.busy || batch.interrupted || !batch.remaining) break;
163 }
164 } catch { warn('summary-unavailable'); }
165 while (!disposed) {
166 const batch = await memory.index({ limit: 32, signal }); indexed += batch.indexed;
167 if (batch.disabled || batch.busy || batch.interrupted || !batch.remaining) return { ...batch, indexed };
168 }
169 return { indexed };
170 },
171 };
172 // Zero configuration: if no embedding model is configured, look for one in a local Ollama
173 // (loopback only, nothing downloaded or sent elsewhere) and backfill vectors in the background.
174 let localEmbedding = null;
175 const useLocalEmbedding = async () => {
176 localEmbedding = null;
177 if (disposed || !settings.autoEmbedding || settings.embedding.enabled) return;
178 const found = await detectLocalEmbedding({ fetcher: config.fetcher });
179 if (!found || disposed || !settings.autoEmbedding || settings.embedding.enabled) return;
180 localEmbedding = found;
181 memory.configure({ ...settings, embedder: createEmbedder(found, { fetcher: config.fetcher }) });
182 await service.index();
183 };
184 service.embeddingStatus = () => settings.embedding.enabled ? `${settings.embedding.model}（手动配置）`
185 : localEmbedding ? `${localEmbedding.model}（自动使用本机 Ollama）` : '未启用（只用关键词检索）';
186 const save = service.save;
187 service.save = input => { const next = save(input); schedule(useLocalEmbedding()); return next; };
188 schedule(useLocalEmbedding());
189 registerManagement(ctx, service);
190 registerTools(ctx, service);
191 ctx.effect(() => () => { disposed = true; operationAbort.abort(); memory.close(); });
192 return service;
193}