feat: implement local deep search engine
This commit is contained in:
john
@@ -0,0 +1,249 @@
|
||||
import fs from 'node:fs';
|
||||
import path from 'node:path';
|
||||
import { DatabaseSync } from 'node:sqlite';
|
||||
|
||||
const TASK_COLUMNS = new Set([
|
||||
'status',
|
||||
'progress',
|
||||
'phase',
|
||||
'plan_json',
|
||||
'report',
|
||||
'error',
|
||||
'updated_at',
|
||||
'completed_at',
|
||||
]);
|
||||
|
||||
function parseJson(value, fallback) {
|
||||
if (!value) return fallback;
|
||||
try {
|
||||
return JSON.parse(value);
|
||||
} catch {
|
||||
return fallback;
|
||||
}
|
||||
}
|
||||
|
||||
function mapTask(row) {
|
||||
if (!row) return null;
|
||||
return {
|
||||
id: row.id,
|
||||
userId: row.user_id || null,
|
||||
question: row.question,
|
||||
depth: row.depth,
|
||||
status: row.status,
|
||||
progress: row.progress,
|
||||
phase: row.phase,
|
||||
plan: parseJson(row.plan_json, []),
|
||||
report: row.report || '',
|
||||
error: row.error || null,
|
||||
createdAt: row.created_at,
|
||||
updatedAt: row.updated_at,
|
||||
completedAt: row.completed_at || null,
|
||||
};
|
||||
}
|
||||
|
||||
function mapSource(row) {
|
||||
return {
|
||||
id: row.id,
|
||||
taskId: row.task_id,
|
||||
url: row.url,
|
||||
title: row.title,
|
||||
snippet: row.snippet,
|
||||
content: row.content,
|
||||
provider: row.provider,
|
||||
score: row.score,
|
||||
query: row.query,
|
||||
createdAt: row.created_at,
|
||||
};
|
||||
}
|
||||
|
||||
export function createDeepSearchStore({
|
||||
databasePath = process.env.TKMIND_DEEP_SEARCH_DB || path.join(process.cwd(), '.deep-search', 'research.sqlite'),
|
||||
now = () => Date.now(),
|
||||
} = {}) {
|
||||
if (databasePath !== ':memory:') {
|
||||
fs.mkdirSync(path.dirname(path.resolve(databasePath)), { recursive: true });
|
||||
}
|
||||
const db = new DatabaseSync(databasePath);
|
||||
db.exec(`
|
||||
PRAGMA journal_mode = WAL;
|
||||
PRAGMA foreign_keys = ON;
|
||||
CREATE TABLE IF NOT EXISTS research_tasks (
|
||||
id TEXT PRIMARY KEY,
|
||||
user_id TEXT,
|
||||
question TEXT NOT NULL,
|
||||
depth TEXT NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
progress INTEGER NOT NULL DEFAULT 0,
|
||||
phase TEXT NOT NULL DEFAULT 'queued',
|
||||
plan_json TEXT NOT NULL DEFAULT '[]',
|
||||
report TEXT NOT NULL DEFAULT '',
|
||||
error TEXT,
|
||||
created_at INTEGER NOT NULL,
|
||||
updated_at INTEGER NOT NULL,
|
||||
completed_at INTEGER
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS research_sources (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
task_id TEXT NOT NULL REFERENCES research_tasks(id) ON DELETE CASCADE,
|
||||
url TEXT NOT NULL,
|
||||
title TEXT NOT NULL DEFAULT '',
|
||||
snippet TEXT NOT NULL DEFAULT '',
|
||||
content TEXT NOT NULL DEFAULT '',
|
||||
provider TEXT NOT NULL DEFAULT 'unknown',
|
||||
score REAL NOT NULL DEFAULT 0,
|
||||
query TEXT NOT NULL DEFAULT '',
|
||||
created_at INTEGER NOT NULL,
|
||||
UNIQUE(task_id, url)
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS research_events (
|
||||
seq INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
task_id TEXT NOT NULL REFERENCES research_tasks(id) ON DELETE CASCADE,
|
||||
type TEXT NOT NULL,
|
||||
payload_json TEXT NOT NULL DEFAULT '{}',
|
||||
created_at INTEGER NOT NULL
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS research_memories (
|
||||
id TEXT PRIMARY KEY,
|
||||
user_id TEXT,
|
||||
topic TEXT NOT NULL,
|
||||
summary TEXT NOT NULL,
|
||||
source_task_id TEXT NOT NULL REFERENCES research_tasks(id) ON DELETE CASCADE,
|
||||
created_at INTEGER NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_research_tasks_status ON research_tasks(status, updated_at DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_research_sources_task ON research_sources(task_id, score DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_research_events_task ON research_events(task_id, seq);
|
||||
`);
|
||||
|
||||
const createTaskStmt = db.prepare(`
|
||||
INSERT INTO research_tasks
|
||||
(id, user_id, question, depth, status, progress, phase, created_at, updated_at)
|
||||
VALUES (?, ?, ?, ?, 'queued', 0, 'queued', ?, ?)
|
||||
`);
|
||||
const getTaskStmt = db.prepare('SELECT * FROM research_tasks WHERE id = ?');
|
||||
const listTasksStmt = db.prepare('SELECT * FROM research_tasks ORDER BY created_at DESC LIMIT ?');
|
||||
const sourcesStmt = db.prepare('SELECT * FROM research_sources WHERE task_id = ? ORDER BY score DESC, id ASC');
|
||||
const eventsStmt = db.prepare('SELECT * FROM research_events WHERE task_id = ? ORDER BY seq ASC');
|
||||
const addSourceStmt = db.prepare(`
|
||||
INSERT INTO research_sources
|
||||
(task_id, url, title, snippet, content, provider, score, query, created_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(task_id, url) DO UPDATE SET
|
||||
title = excluded.title,
|
||||
snippet = excluded.snippet,
|
||||
content = CASE
|
||||
WHEN length(excluded.content) > length(research_sources.content) THEN excluded.content
|
||||
ELSE research_sources.content
|
||||
END,
|
||||
provider = excluded.provider,
|
||||
score = MAX(research_sources.score, excluded.score),
|
||||
query = excluded.query
|
||||
`);
|
||||
const appendEventStmt = db.prepare(`
|
||||
INSERT INTO research_events (task_id, type, payload_json, created_at)
|
||||
VALUES (?, ?, ?, ?)
|
||||
`);
|
||||
const saveMemoryStmt = db.prepare(`
|
||||
INSERT INTO research_memories (id, user_id, topic, summary, source_task_id, created_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(id) DO UPDATE SET summary = excluded.summary
|
||||
`);
|
||||
|
||||
function getTask(id, { includeSources = true, includeEvents = true } = {}) {
|
||||
const task = mapTask(getTaskStmt.get(String(id)));
|
||||
if (!task) return null;
|
||||
if (includeSources) task.sources = sourcesStmt.all(task.id).map(mapSource);
|
||||
if (includeEvents) {
|
||||
task.events = eventsStmt.all(task.id).map((event) => ({
|
||||
seq: event.seq,
|
||||
type: event.type,
|
||||
payload: parseJson(event.payload_json, {}),
|
||||
createdAt: event.created_at,
|
||||
}));
|
||||
}
|
||||
return task;
|
||||
}
|
||||
|
||||
return {
|
||||
createTask({ id, userId = null, question, depth }) {
|
||||
const timestamp = now();
|
||||
createTaskStmt.run(id, userId, question, depth, timestamp, timestamp);
|
||||
appendEventStmt.run(id, 'queued', JSON.stringify({ depth }), timestamp);
|
||||
return getTask(id);
|
||||
},
|
||||
|
||||
updateTask(id, patch = {}) {
|
||||
const normalized = { ...patch, updated_at: patch.updatedAt ?? now() };
|
||||
if (Array.isArray(normalized.plan)) {
|
||||
normalized.plan_json = JSON.stringify(normalized.plan);
|
||||
delete normalized.plan;
|
||||
}
|
||||
if (Object.hasOwn(normalized, 'completedAt')) {
|
||||
normalized.completed_at = normalized.completedAt;
|
||||
delete normalized.completedAt;
|
||||
}
|
||||
if (Object.hasOwn(normalized, 'updatedAt')) delete normalized.updatedAt;
|
||||
const entries = Object.entries(normalized).filter(([key]) => TASK_COLUMNS.has(key));
|
||||
if (entries.length) {
|
||||
const assignments = entries.map(([key]) => `${key} = ?`).join(', ');
|
||||
db.prepare(`UPDATE research_tasks SET ${assignments} WHERE id = ?`)
|
||||
.run(...entries.map(([, value]) => value), String(id));
|
||||
}
|
||||
return getTask(id);
|
||||
},
|
||||
|
||||
addSource(taskId, source = {}) {
|
||||
addSourceStmt.run(
|
||||
String(taskId),
|
||||
String(source.url ?? ''),
|
||||
String(source.title ?? '').slice(0, 1000),
|
||||
String(source.snippet ?? '').slice(0, 4000),
|
||||
String(source.content ?? '').slice(0, 100_000),
|
||||
String(source.provider ?? 'unknown').slice(0, 100),
|
||||
Number(source.score ?? 0) || 0,
|
||||
String(source.query ?? '').slice(0, 1000),
|
||||
now(),
|
||||
);
|
||||
},
|
||||
|
||||
appendEvent(taskId, type, payload = {}) {
|
||||
appendEventStmt.run(String(taskId), String(type), JSON.stringify(payload), now());
|
||||
},
|
||||
|
||||
saveMemory({ id, userId = null, topic, summary, sourceTaskId }) {
|
||||
saveMemoryStmt.run(id, userId, topic, summary, sourceTaskId, now());
|
||||
},
|
||||
|
||||
getTask,
|
||||
|
||||
listTasks(limit = 50) {
|
||||
const bounded = Math.max(1, Math.min(200, Number(limit) || 50));
|
||||
return listTasksStmt.all(bounded).map(mapTask);
|
||||
},
|
||||
|
||||
listMemories({ userId = null, limit = 50 } = {}) {
|
||||
const bounded = Math.max(1, Math.min(200, Number(limit) || 50));
|
||||
const rows = userId
|
||||
? db.prepare('SELECT * FROM research_memories WHERE user_id = ? ORDER BY created_at DESC LIMIT ?')
|
||||
.all(String(userId), bounded)
|
||||
: db.prepare('SELECT * FROM research_memories ORDER BY created_at DESC LIMIT ?').all(bounded);
|
||||
return rows.map((row) => ({
|
||||
id: row.id,
|
||||
userId: row.user_id || null,
|
||||
topic: row.topic,
|
||||
summary: row.summary,
|
||||
sourceTaskId: row.source_task_id,
|
||||
createdAt: row.created_at,
|
||||
}));
|
||||
},
|
||||
|
||||
getStats() {
|
||||
const rows = db.prepare('SELECT status, COUNT(*) AS count FROM research_tasks GROUP BY status').all();
|
||||
return Object.fromEntries(rows.map((row) => [row.status, Number(row.count)]));
|
||||
},
|
||||
|
||||
close() {
|
||||
db.close();
|
||||
},
|
||||
};
|
||||
}
|
||||
Reference in New Issue
Block a user