Files
RAG-Lambdas/query.js
2026-08-12 19:41:32 -04:00

90 lines
2.7 KiB
JavaScript

import { Pinecone } from "@pinecone-database/pinecone";
import { embedTexts, rerank } from "./lib/voyage.js";
import { generate } from "./lib/anthropic.js";
import { computeGroundedness } from "./lib/scoring.js";
const pc = new Pinecone({ apiKey: process.env.PINECONE_API_KEY });
const TOP_K_SEARCH = 20;
const TOP_K_RERANK = 5;
export const handler = async (event) => {
const { question } = JSON.parse(event.body || "{}");
if (!question || typeof question !== "string") {
return {
statusCode: 400,
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ error: "question is required" }),
};
}
const [queryEmbedding] = await embedTexts([question], "query");
const index = pc.index(process.env.PINECONE_INDEX_NAME);
const queryResponse = await index.query({
vector: queryEmbedding,
topK: TOP_K_SEARCH,
includeMetadata: true,
});
const candidates = (queryResponse.matches || [])
.filter((m) => m.metadata && m.metadata.text)
.map((m) => ({
id: m.id,
text: m.metadata.text,
source: m.metadata.source || m.id,
chunkIndex: m.metadata.chunkIndex,
}));
if (candidates.length === 0) {
return {
statusCode: 200,
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
answer: "No sources found. Ingest documents first.",
citations: [],
groundednessScore: null,
followUpQuestions: [],
contextUsed: 0,
}),
};
}
const ranked = await rerank(question, candidates);
const topChunks = ranked.slice(0, TOP_K_RERANK).map((r) => r.document);
const sources = topChunks.map((chunk, i) => ({
docIndex: i + 1,
name: chunk.source,
text: chunk.text,
id: chunk.id,
}));
const { answer, citedSourceIndices, followUpQuestions } = await generate(question, sources);
const citedChunks = citedSourceIndices
.map((idx) => sources.find((s) => s.docIndex === idx))
.filter(Boolean);
const groundednessScore = await computeGroundedness(
answer,
citedChunks.map((c) => c.text)
);
const citations = citedChunks.map((c) => ({
sourceIndex: c.docIndex,
source: c.name,
chunkText: c.text,
}));
return {
statusCode: 200,
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
answer,
citations,
groundednessScore,
followUpQuestions,
contextUsed: topChunks.length,
}),
};
};