114 lines
3.5 KiB
TypeScript
114 lines
3.5 KiB
TypeScript
'use strict';
|
|
|
|
import { Router, type Request, type Response, type NextFunction } from 'express';
|
|
import logger from '../logger.js';
|
|
import { embedTexts, search, rerank } from '../services/retrievalService.js';
|
|
import type { ScoredChunk, RankedChunk, TextChunk } from '../services/retrievalService.js';
|
|
import { generate } from '../services/generationService.js';
|
|
import type { GenerationResult } from '../services/generationService.js';
|
|
import { computeGroundedness } from '../services/scoringService.js';
|
|
import * as notebookStore from '../stores/notebookStore.js';
|
|
import type { SourceGroup } from '../stores/notebookStore.js';
|
|
|
|
const TOP_K_SEARCH = 20;
|
|
const TOP_K_RERANK = 5;
|
|
|
|
const router = Router();
|
|
|
|
interface QueryBody {
|
|
notebookId?: string;
|
|
question?: string;
|
|
}
|
|
|
|
interface Citation {
|
|
sourceIndex: number;
|
|
sourceId: string | null;
|
|
name: string | null;
|
|
chunkTexts: string[];
|
|
}
|
|
|
|
interface QueryResponse {
|
|
answer: string;
|
|
citations: Citation[];
|
|
groundednessScore: number | null;
|
|
followUpQuestions: string[];
|
|
}
|
|
|
|
router.post('/', async (req: Request<unknown, unknown, QueryBody>, res: Response, next: NextFunction) => {
|
|
try {
|
|
const { notebookId, question } = req.body;
|
|
if (!notebookId || !question) {
|
|
res.status(400).json({ error: 'notebookId and question are required' });
|
|
return;
|
|
}
|
|
|
|
logger.info({ notebookId, question }, 'query received');
|
|
|
|
const [queryEmbedding]: number[][] = await embedTexts([question], 'query');
|
|
|
|
const searchResults: ScoredChunk[] = search(queryEmbedding, notebookId, TOP_K_SEARCH);
|
|
if (searchResults.length === 0) {
|
|
res.json({
|
|
answer: 'No sources found for this notebook. Upload some documents first.',
|
|
citations: [],
|
|
groundednessScore: null,
|
|
followUpQuestions: [],
|
|
});
|
|
return;
|
|
}
|
|
|
|
const candidateChunks: TextChunk[] = searchResults.map((r) => r.chunk);
|
|
const reranked: RankedChunk[] = await rerank(question, candidateChunks);
|
|
const topChunks: TextChunk[] = reranked.slice(0, TOP_K_RERANK).map((r) => r.chunk);
|
|
|
|
logger.debug(
|
|
{
|
|
searchHits: searchResults.length,
|
|
rerankedTop: topChunks.length,
|
|
},
|
|
'retrieval complete'
|
|
);
|
|
|
|
const sourceGroups: SourceGroup[] = notebookStore.buildSourceGroups(notebookId, topChunks);
|
|
|
|
const { answer, citedSourceIndices, followUpQuestions }: GenerationResult = await generate(question, sourceGroups);
|
|
|
|
const citedChunkIds: string[] = [];
|
|
for (const idx of citedSourceIndices) {
|
|
const group: SourceGroup | undefined = sourceGroups.find((g) => g.docIndex === idx);
|
|
if (group) {
|
|
citedChunkIds.push(...group.chunks.map((c) => c.id));
|
|
}
|
|
}
|
|
|
|
const groundednessScore: number = await computeGroundedness(answer, citedChunkIds);
|
|
|
|
logger.info({ groundednessScore: groundednessScore.toFixed(3) }, 'query complete');
|
|
|
|
const citations: Citation[] = citedSourceIndices.map((idx) => {
|
|
const group: SourceGroup | undefined = sourceGroups.find((g) => g.docIndex === idx);
|
|
return group
|
|
? {
|
|
sourceIndex: idx,
|
|
sourceId: group.sourceId,
|
|
name: group.name,
|
|
chunkTexts: group.chunks.map((c) => c.text),
|
|
}
|
|
: { sourceIndex: idx, sourceId: null, name: null, chunkTexts: [] as string[] };
|
|
});
|
|
|
|
const response: QueryResponse = {
|
|
answer,
|
|
citations,
|
|
groundednessScore,
|
|
followUpQuestions,
|
|
};
|
|
|
|
res.json(response);
|
|
} catch (err) {
|
|
next(err);
|
|
}
|
|
});
|
|
|
|
export default router;
|