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

@@ -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");
+  });
+});

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<SearchResult[]>;
+  generate: (prompt: string) => Promise<string>;
+}
+
+const defaultDependencies: AnswerDependencies = {
+  search: searchDocuments,
+  generate: generateResponse,
+};
+
+export const answerQuestion = async (
+  question: string,
+  dependencies: AnswerDependencies = defaultDependencies,
+): Promise<AnswerResult> => {
+  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,
+    })),
+  };
+};

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();
+  });
+});

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);