feat: implement local deep search engine
This commit is contained in:
john
@@ -0,0 +1,187 @@
|
||||
import http from 'node:http';
|
||||
import path from 'node:path';
|
||||
import { pathToFileURL } from 'node:url';
|
||||
import { createDeepSearchEngine } from './deep-search-engine.mjs';
|
||||
import { createDeepSearchStore } from './deep-search-store.mjs';
|
||||
|
||||
function sendJson(res, status, body) {
|
||||
const payload = JSON.stringify(body);
|
||||
res.writeHead(status, {
|
||||
'content-type': 'application/json; charset=utf-8',
|
||||
'content-length': Buffer.byteLength(payload),
|
||||
'cache-control': 'no-store',
|
||||
});
|
||||
res.end(payload);
|
||||
}
|
||||
|
||||
async function readJson(req, { maxBytes = 1_000_000 } = {}) {
|
||||
const chunks = [];
|
||||
let size = 0;
|
||||
for await (const chunk of req) {
|
||||
size += chunk.length;
|
||||
if (size > maxBytes) {
|
||||
const error = new Error('Request body is too large');
|
||||
error.status = 413;
|
||||
throw error;
|
||||
}
|
||||
chunks.push(chunk);
|
||||
}
|
||||
if (!chunks.length) return {};
|
||||
try {
|
||||
return JSON.parse(Buffer.concat(chunks).toString('utf8'));
|
||||
} catch {
|
||||
const error = new Error('Request body must be valid JSON');
|
||||
error.status = 400;
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
function publicTask(task, { includeContent = false, includeEvents = false } = {}) {
|
||||
if (!task) return null;
|
||||
return {
|
||||
...task,
|
||||
sources: task.sources?.map((source) => ({
|
||||
...source,
|
||||
...(!includeContent ? { content: undefined } : {}),
|
||||
})),
|
||||
...(!includeEvents ? { events: undefined } : {}),
|
||||
};
|
||||
}
|
||||
|
||||
function matchTaskPath(pathname, suffix = '') {
|
||||
const match = pathname.match(new RegExp(`^/v1/research/([^/]+)${suffix}$`));
|
||||
return match ? decodeURIComponent(match[1]) : null;
|
||||
}
|
||||
|
||||
function canAccessTask(task, requestUserId) {
|
||||
return !requestUserId || !task?.userId || task.userId === requestUserId;
|
||||
}
|
||||
|
||||
export function createDeepSearchHttpServer({
|
||||
engine,
|
||||
secret = process.env.TKMIND_DEEP_SEARCH_SECRET || '',
|
||||
logger = console,
|
||||
} = {}) {
|
||||
if (!engine) throw new Error('Deep Search engine is required');
|
||||
return http.createServer(async (req, res) => {
|
||||
const url = new URL(req.url ?? '/', 'http://deep-search.local');
|
||||
if (url.pathname === '/health' && req.method === 'GET') {
|
||||
return sendJson(res, 200, engine.getHealth());
|
||||
}
|
||||
if (secret && req.headers['x-secret-key'] !== secret) {
|
||||
return sendJson(res, 401, { code: 'UNAUTHORIZED', message: 'Invalid service secret' });
|
||||
}
|
||||
const requestUserId = String(req.headers['x-user-id'] ?? '').trim();
|
||||
try {
|
||||
if (url.pathname === '/v1/search' && req.method === 'POST') {
|
||||
const input = await readJson(req);
|
||||
const results = await engine.search(input);
|
||||
return sendJson(res, 200, { query: input.query, results, service: 'tkmind-deep-search' });
|
||||
}
|
||||
if (url.pathname === '/v1/research' && req.method === 'POST') {
|
||||
const input = await readJson(req);
|
||||
const task = engine.start({
|
||||
...input,
|
||||
userId: requestUserId || input.userId || null,
|
||||
});
|
||||
return sendJson(res, 202, {
|
||||
task_id: task.taskId,
|
||||
status: task.status,
|
||||
service: task.service,
|
||||
});
|
||||
}
|
||||
if (url.pathname === '/v1/research' && req.method === 'GET') {
|
||||
const tasks = engine.listTasks(url.searchParams.get('limit'))
|
||||
.filter((task) => canAccessTask(task, requestUserId));
|
||||
return sendJson(res, 200, { tasks });
|
||||
}
|
||||
const eventTaskId = matchTaskPath(url.pathname, '/events');
|
||||
if (eventTaskId && req.method === 'GET') {
|
||||
const task = engine.getTask(eventTaskId);
|
||||
if (!task) return sendJson(res, 404, { code: 'NOT_FOUND', message: 'Research task not found' });
|
||||
if (!canAccessTask(task, requestUserId)) {
|
||||
return sendJson(res, 403, { code: 'FORBIDDEN', message: 'Research task belongs to another user' });
|
||||
}
|
||||
res.writeHead(200, {
|
||||
'content-type': 'text/event-stream',
|
||||
'cache-control': 'no-cache, no-transform',
|
||||
connection: 'keep-alive',
|
||||
});
|
||||
for (const event of task.events ?? []) {
|
||||
res.write(`id: ${event.seq}\nevent: ${event.type}\ndata: ${JSON.stringify(event.payload)}\n\n`);
|
||||
}
|
||||
const unsubscribe = engine.subscribe(eventTaskId, (event) => {
|
||||
res.write(`event: ${event.type}\ndata: ${JSON.stringify(event.payload)}\n\n`);
|
||||
if (['completed', 'failed', 'cancelled'].includes(event.type)) res.end();
|
||||
});
|
||||
const heartbeat = setInterval(() => res.write(': heartbeat\n\n'), 15_000);
|
||||
const cleanup = () => {
|
||||
clearInterval(heartbeat);
|
||||
unsubscribe();
|
||||
};
|
||||
req.once('close', cleanup);
|
||||
res.once('close', cleanup);
|
||||
return;
|
||||
}
|
||||
const taskId = matchTaskPath(url.pathname);
|
||||
if (taskId && req.method === 'GET') {
|
||||
const task = engine.getTask(taskId);
|
||||
if (!task) return sendJson(res, 404, { code: 'NOT_FOUND', message: 'Research task not found' });
|
||||
if (!canAccessTask(task, requestUserId)) {
|
||||
return sendJson(res, 403, { code: 'FORBIDDEN', message: 'Research task belongs to another user' });
|
||||
}
|
||||
return sendJson(res, 200, publicTask(task, {
|
||||
includeContent: url.searchParams.get('include_content') === '1',
|
||||
includeEvents: url.searchParams.get('include_events') === '1',
|
||||
}));
|
||||
}
|
||||
if (taskId && req.method === 'DELETE') {
|
||||
const existing = engine.getTask(taskId);
|
||||
if (existing && !canAccessTask(existing, requestUserId)) {
|
||||
return sendJson(res, 403, { code: 'FORBIDDEN', message: 'Research task belongs to another user' });
|
||||
}
|
||||
const task = engine.cancel(taskId);
|
||||
if (!task) return sendJson(res, 404, { code: 'NOT_FOUND', message: 'Research task not found' });
|
||||
return sendJson(res, 202, publicTask(task));
|
||||
}
|
||||
return sendJson(res, 404, { code: 'NOT_FOUND', message: 'Route not found' });
|
||||
} catch (error) {
|
||||
logger.warn?.('[deep-search] request failed', error);
|
||||
return sendJson(res, Number(error?.status) || 400, {
|
||||
code: 'REQUEST_FAILED',
|
||||
message: String(error?.message ?? error),
|
||||
});
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
export function startDeepSearchServer({
|
||||
host = process.env.TKMIND_DEEP_SEARCH_HOST || '127.0.0.1',
|
||||
port = Number(process.env.TKMIND_DEEP_SEARCH_PORT || 20100),
|
||||
databasePath = process.env.TKMIND_DEEP_SEARCH_DB
|
||||
|| path.join(process.cwd(), '.deep-search', 'research.sqlite'),
|
||||
logger = console,
|
||||
} = {}) {
|
||||
const store = createDeepSearchStore({ databasePath });
|
||||
const engine = createDeepSearchEngine({ store });
|
||||
const server = createDeepSearchHttpServer({ engine, logger });
|
||||
server.listen(port, host, () => {
|
||||
logger.info?.(`[deep-search] listening on http://${host}:${port}`);
|
||||
});
|
||||
const close = async () => {
|
||||
await new Promise((resolve) => server.close(resolve));
|
||||
store.close();
|
||||
};
|
||||
return { server, engine, store, close };
|
||||
}
|
||||
|
||||
const invokedPath = process.argv[1] ? pathToFileURL(path.resolve(process.argv[1])).href : '';
|
||||
if (import.meta.url === invokedPath) {
|
||||
const runtime = startDeepSearchServer();
|
||||
const shutdown = async () => {
|
||||
await runtime.close();
|
||||
process.exit(0);
|
||||
};
|
||||
process.once('SIGINT', shutdown);
|
||||
process.once('SIGTERM', shutdown);
|
||||
}
|
||||
Reference in New Issue
Block a user