john
277 lines
9.1 KiB
JavaScript
277 lines
9.1 KiB
JavaScript
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,
|
|
llmProviderKeyId: row.llm_provider_key_id || '',
|
|
llmModel: row.llm_model || '',
|
|
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,
|
|
llm_provider_key_id TEXT NOT NULL DEFAULT '',
|
|
llm_model TEXT NOT NULL DEFAULT '',
|
|
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 taskColumns = new Set(db.prepare('PRAGMA table_info(research_tasks)').all().map((row) => row.name));
|
|
if (!taskColumns.has('llm_provider_key_id')) {
|
|
db.exec("ALTER TABLE research_tasks ADD COLUMN llm_provider_key_id TEXT NOT NULL DEFAULT ''");
|
|
}
|
|
if (!taskColumns.has('llm_model')) {
|
|
db.exec("ALTER TABLE research_tasks ADD COLUMN llm_model TEXT NOT NULL DEFAULT ''");
|
|
}
|
|
|
|
const createTaskStmt = db.prepare(`
|
|
INSERT INTO research_tasks
|
|
(id, user_id, question, depth, llm_provider_key_id, llm_model, 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,
|
|
llmProviderKeyId = '',
|
|
llmModel = '',
|
|
}) {
|
|
const timestamp = now();
|
|
createTaskStmt.run(
|
|
id,
|
|
userId,
|
|
question,
|
|
depth,
|
|
llmProviderKeyId,
|
|
llmModel,
|
|
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();
|
|
},
|
|
};
|
|
}
|