From be738ba9482ebe0215e09d1e20b2e061807638f0 Mon Sep 17 00:00:00 2001 From: Brandon Wolfgang Date: Thu, 22 Oct 2026 15:14:40 -0500 Subject: [PATCH] feat: implement retrieval-augmented question answering - Add answerQuestion to retrieve relevant document chunks and generate context-grounded answers - Filter retrieved chunks by a minimum similarity threshold of 0.5 - Return a fallback response when no relevant documents meet the threshold - Include source attribution with document IDs, titles, content, and similarity scores - Add POST /ask with question validation and sanitized error handling - Add unit and endpoint tests for relevance filtering, answer generation, source attribution, and failure scenarios --- src/answer.test.ts | 161 +++++++++++++++++++++++++++++++++++++++++++++ src/answer.ts | 74 +++++++++++++++++++++ src/server.test.ts | 85 +++++++++++++++++++++++- src/server.ts | 71 +++++++++++++------- 4 files changed, 366 insertions(+), 25 deletions(-) create mode 100644 src/answer.test.ts create mode 100644 src/answer.ts diff --git a/src/answer.test.ts b/src/answer.test.ts new file mode 100644 index 0000000..5ac28a0 --- /dev/null +++ b/src/answer.test.ts @@ -0,0 +1,161 @@ +import { describe, it } from "node:test"; +import assert from "node:assert/strict"; +import { answerQuestion } from "./answer.js"; +import type { SearchResult } from "./search.js"; +import { generateColumnString } from "node-pg-migrate/operations/indexes/shared"; + +const sampleResults: SearchResult[] = [ + { + documentId: "1", + title: "RAG Basics", + content: + "Retrieval augmented generation uses retrieved information to help generated a response.", + similarity: 0.88, + chunkIndex: 0, + }, +]; + +describe("answerQuestion", () => { + it("rejects an empty question", async () => { + await assert.rejects(answerQuestion(" "), { + message: "Question cannot be empty", + }); + }); + + it("returns a fallback answer when no relevant documents are found", async () => { + let generatorCalled = false; + + const result = await answerQuestion("What is RAG?", { + search: async () => [], + generate: async () => { + generatorCalled = true; + return "This should not be called."; + }, + }); + + assert.equal( + result.answer, + "I couldn't find relevant information in the available documents.", + ); + assert.deepEqual(result.sources, []); + assert.equal(generatorCalled, false); + }); + + it("generates and answer using retrieved context and returns its sources", async () => { + let receivedQuery = ""; + let receivedLimit = 0; + let receivedPrompt = ""; + + const result = await answerQuestion(" What is RAG? ", { + search: async (query, limit) => { + ((receivedQuery = query), (receivedLimit = limit ?? 0)); + return sampleResults; + }, + generate: async (prompt) => { + receivedPrompt = prompt; + return "RAG uses retrieved information to generate a response."; + }, + }); + + assert.equal(receivedQuery, "What is RAG?"); + assert.equal(receivedLimit, 5); + assert.match(receivedPrompt, /RAG Basics/); + assert.match( + receivedPrompt, + /Retrieval augmented generation uses retrieved information/, + ); + assert.match(receivedPrompt, /What is RAG\?/); + + assert.equal( + result.answer, + "RAG uses retrieved information to generate a response.", + ); + assert.deepEqual(result.sources, [ + { + documentId: "1", + title: "RAG Basics", + content: sampleResults[0].content, + similarity: 0.88, + }, + ]); + }); + + it("propagates search errors", async () => { + await assert.rejects( + answerQuestion("What is RAG?", { + search: async () => { + throw new Error("Search failed"); + }, + generate: async () => "Should not be called.", + }), + { message: "Search failed" }, + ); + }); + + it("propagates generation errors", async () => { + await assert.rejects( + answerQuestion("What is RAG?", { + search: async () => sampleResults, + generate: async () => { + throw new Error("Generation failed"); + }, + }), + { message: "Generation failed" }, + ); + }); + + it("excludes results below the similarity threshold", async () => { + const result = await answerQuestion("What is the capital of France?", { + search: async () => [ + { + documentId: "1", + chunkIndex: 0, + title: "RAG Basics", + content: "Retrieval augmented generation uses retrieved information.", + similarity: 0.308, + }, + ], + generate: async () => { + assert.fail("Generator should not be called for irrelevant results"); + }, + }); + + assert.equal( + result.answer, + "I couldn't find relevant information in the available documents.", + ); + assert.deepEqual(result.sources, []); + }); + + it("uses only relevant results for context and source attribution", async () => { + let receivedPrompt = ""; + + const result = await answerQuestion("How does RAG work?", { + search: async () => [ + { + documentId: "1", + chunkIndex: 0, + title: "Relevant document", + content: "RAG retrieves information to help generate responses.", + similarity: 0.88, + }, + { + documentId: "2", + chunkIndex: 0, + title: "Irrelevant document", + content: "Paris is the capital of France.", + similarity: 0.308, + }, + ], + generate: async (prompt) => { + receivedPrompt = prompt; + return "RAG retrieves information to help generate responses."; + }, + }); + + assert.match(receivedPrompt, /Relevant document/); + assert.doesNotMatch(receivedPrompt, /Irrelevant document/); + assert.equal(result.sources.length, 1); + assert.equal(result.sources[0].documentId, "1"); + }); +}); diff --git a/src/answer.ts b/src/answer.ts new file mode 100644 index 0000000..55fa971 --- /dev/null +++ b/src/answer.ts @@ -0,0 +1,74 @@ +import { generateResponse } from "./ollama.js"; +import { searchDocuments, type SearchResult } from "./search.js"; + +export interface AnswerResult { + answer: string; + sources: { + documentId: string; + title: string; + content: string; + similarity: number; + }[]; +} + +interface AnswerDependencies { + search: (query: string, limit?: number) => Promise; + generate: (prompt: string) => Promise; +} + +const defaultDependencies: AnswerDependencies = { + search: searchDocuments, + generate: generateResponse, +}; + +export const answerQuestion = async ( + question: string, + dependencies: AnswerDependencies = defaultDependencies, +): Promise => { + const normalizedQuestion = question.trim(); + + if (!normalizedQuestion) { + throw new Error("Question cannot be empty"); + } + + const results = await dependencies.search(normalizedQuestion, 5); + const relevantResults = results.filter((result) => result.similarity >= 0.5); + + if (relevantResults.length === 0) { + return { + answer: + "I couldn't find relevant information in the available documents.", + sources: [], + }; + } + + const context = relevantResults + .map( + (result, index) => + `[Source ${index + 1}: ${result.title}]\n${result.content}`, + ) + .join("\n\n"); + + const prompt = `Answer the question using only the context provided below. +If the context does not contain the answer, say that you don't have enough information. +Do not invent facts or claim that the context supports something it does not. + +Context: +${context} + +Question: ${normalizedQuestion} + +Answer:`; + + const answer = await dependencies.generate(prompt); + + return { + answer, + sources: relevantResults.map((result) => ({ + documentId: result.documentId, + title: result.title, + content: result.content, + similarity: result.similarity, + })), + }; +}; diff --git a/src/server.test.ts b/src/server.test.ts index fd530dc..2f703f3 100644 --- a/src/server.test.ts +++ b/src/server.test.ts @@ -2,9 +2,7 @@ import assert from "node:assert"; import { after, describe, it } from "node:test"; import { buildApp } from "./server.js"; -import { setServers } from "node:dns"; -import { allowedNodeEnvironmentFlags } from "node:process"; -import { unescape } from "node:querystring"; +import type { AnswerResult } from "./answer.js"; import { error } from "node:console"; describe("API endpoints", () => { @@ -390,3 +388,84 @@ describe("API endpoints", () => { } }); }); + +describe("POST /ask", () => { + it("returns an answer with source attribution", async () => { + const expected: AnswerResult = { + answer: "RAG uses retrieved information to generate a response.", + sources: [ + { + documentId: "1", + title: "RAG Basics", + content: "Retrieved information helps generate a response", + similarity: 0.88, + }, + ], + }; + + const app = buildApp( + undefined, + undefined, + undefined, + undefined, + async (question) => { + assert.equal(question, "What is RAG?"); + return expected; + }, + ); + + const response = await app.inject({ + method: "POST", + url: "/ask", + payload: { question: "What is RAG?" }, + }); + + assert.equal(response.statusCode, 200); + assert.deepEqual(response.json(), expected); + + await app.close(); + }); + + it("rejects an invalid question", async () => { + const app = buildApp(); + + const response = await app.inject({ + method: "POST", + url: "/ask", + payload: { question: " " }, + }); + + assert.equal(response.statusCode, 400); + assert.deepEqual(response.json(), { + error: "Invalid question", + }); + + await app.close(); + }); + + it("returns 503 when answering fails", async () => { + const app = buildApp( + undefined, + undefined, + undefined, + undefined, + async () => { + throw new Error("Sensitive internal detail"); + }, + ); + + const response = await app.inject({ + method: "POST", + url: "/ask", + payload: { question: "What is RAG?" }, + }); + + assert.equal(response.statusCode, 503); + assert.deepEqual(response.json(), { + error: "Question answering failed", + }); + assert.equal(response.body.includes("Sensitive internal detail"), false); + + await app.close(); + }); +}); diff --git a/src/server.ts b/src/server.ts index 6da0b83..4b4b7d4 100644 --- a/src/server.ts +++ b/src/server.ts @@ -1,10 +1,15 @@ import Fastify, { type FastifyError } from "fastify"; import { z } from "zod"; +import { answerQuestion as answerQuestionDefault } from "./answer.js"; import { ingestDocument as ingestDocumentDefault } from "./ingestion.js"; import { checkOllamaHealth, generateResponse } from "./ollama.js"; import { searchDocuments as searchDocumentsDefault } from "./search.js"; +const askSchema = z.object({ + question: z.string().trim().min(1).max(2_000), +}); + const echoSchema = z.object({ message: z.string().min(1).max(500), }); @@ -25,6 +30,7 @@ export const buildApp = ( checkHealth = checkOllamaHealth, ingest = ingestDocumentDefault, search = searchDocumentsDefault, + answer = answerQuestionDefault, ) => { const app = Fastify({ logger: true, @@ -74,6 +80,49 @@ export const buildApp = ( }; }); + app.post("/ask", async (request, reply) => { + const result = askSchema.safeParse(request.body); + + if (!result.success) { + return reply.code(400).send({ + error: "Invalid question", + }); + } + + try { + const response = await answer(result.data.question); + return response; + } catch (error) { + request.log.error({ err: error }, "Question answering failed"); + + return reply.code(503).send({ + error: "Question answering failed", + }); + } + }); + + app.post("/documents", async (request, reply) => { + const result = documentSchema.safeParse(request.body); + + if (!result.success) { + return reply.code(400).send({ + error: "Invalid document", + }); + } + + try { + const document = await ingest(result.data); + + return reply.code(201).send(document); + } catch (error) { + request.log.error({ err: error }, "Document ingestion failed"); + + return reply.code(503).send({ + error: "Document ingestion failed", + }); + } + }); + app.post("/echo", async (request, reply) => { const result = echoSchema.safeParse(request.body); @@ -110,28 +159,6 @@ export const buildApp = ( } }); - app.post("/documents", async (request, reply) => { - const result = documentSchema.safeParse(request.body); - - if (!result.success) { - return reply.code(400).send({ - error: "Invalid document", - }); - } - - try { - const document = await ingest(result.data); - - return reply.code(201).send(document); - } catch (error) { - request.log.error({ err: error }, "Document ingestion failed"); - - return reply.code(503).send({ - error: "Document ingestion failed", - }); - } - }); - app.post("/search", async (request, reply) => { const result = searchSchema.safeParse(request.body);