feat: functional tests with a set of predefined questions + route for counting docs

src/lib/build-server.ts

@@ -13,11 +13,17 @@ import {
 	SimilaritySearchConfig,
 } from "./common.js";
 import { ApplicationError, EnvError, UserError } from "./errors.js";
-import { bodySchema, healthSchema, responseSchema } from "./json-schemas.js";
+import {
+	bodySchema,
+	countSchema,
+	healthSchema,
+	responseSchema,
+} from "./json-schemas.js";
 import { similaritySearchOnChunksAndSummaries } from "./similarity-search-chunks-and-summaries.js";
 import { similaritySearchOnChunksOnly } from "./similarity-search-chunks-only.js";
 import { createPrompt } from "./create-prompt.js";
 import { similaritySearchFirstSummariesThenChunks } from "./similarity-search-summaries-then-chunks.js";
+import supabase from "./supabase.js";
 
 export async function buildServer({
 	OPENAI_MODEL,
@@ -91,6 +97,34 @@ export async function buildServer({
 		},
 		{ prefix: "/health" },
 	);
+	fastify.register(
+		(app, options, next) => {
+			app.register(cors, { origin: "*" });
+			app.get(
+				"/count",
+				{
+					schema: {
+						response: countSchema,
+					},
+				},
+				async (_request, reply) => {
+					const { data, error, count } = await supabase
+						.from("processed_documents")
+						.select("*", { count: "exact", head: true });
+					console.log(data, error, count);
+					if (!error && count) {
+						reply.status(200).send({ registered_documents_count: count });
+					} else {
+						reply
+							.status(500)
+							.send({ error: "Could not count processed_documents" });
+					}
+				},
+			);
+			next();
+		},
+		{ prefix: "/processed_documents" },
+	);
 	fastify.register(
 		(app, options, next) => {
 			app.register(cors, {
@@ -183,6 +217,7 @@ export async function buildServer({
 						document_limit,
 						search_algorithm,
 						include_summary_in_response_generation,
+						generate_answer,
 					} = request.body;
 
 					app.log.info({ query });
@@ -197,6 +232,7 @@ export async function buildServer({
 					app.log.info({ document_limit });
 					app.log.info({ search_algorithm });
 					app.log.info({ include_summary_in_response_generation });
+					app.log.info({ generate_answer });
 
 					// 2. moderate content
 					// Moderate the content to comply with OpenAI T&C
@@ -308,28 +344,32 @@ export async function buildServer({
 						requestBody: request.body,
 					} as ResponseDetail;
 
-					// eslint-disable-next-line @typescript-eslint/ban-ts-comment
-					//@ts-ignore
-					const response = await fetch(
-						"https://api.openai.com/v1/chat/completions",
-						{
-							method: "POST",
-							headers: {
-								Authorization: `Bearer ${OPENAI_KEY}`,
-								"Content-Type": "application/json",
+					let answer = undefined;
+					if (generate_answer) {
+						// eslint-disable-next-line @typescript-eslint/ban-ts-comment
+						// @ts-ignore
+						const response = await fetch(
+							"https://api.openai.com/v1/chat/completions",
+							{
+								method: "POST",
+								headers: {
+									Authorization: `Bearer ${OPENAI_KEY}`,
+									"Content-Type": "application/json",
+								},
+								body: JSON.stringify(responseDetail.completionOptions),
 							},
-							body: JSON.stringify(responseDetail.completionOptions),
-						},
-					);
-
-					if (response.status !== 200) {
-						throw new ApplicationError(
-							"Failed to create completion for question",
-							{ response },
 						);
+
+						if (response.status !== 200) {
+							throw new ApplicationError(
+								"Failed to create completion for question",
+								{ response },
+							);
+						}
+						answer = await response.json();
 					}
-					const json = await response.json();
-					responseDetail.gpt = json;
+
+					responseDetail.gpt = answer;
 					responseDetail.requestBody = request.body;
 					reply.status(201).send(responseDetail);
 				},

src/lib/common.ts

@@ -85,6 +85,7 @@ export interface Body {
 	document_limit: number;
 	search_algorithm: string;
 	include_summary_in_response_generation: boolean;
+	generate_answer: boolean;
 }
 
 export interface SimilaritySearchConfig {

src/lib/json-schemas.ts

@@ -17,6 +17,7 @@ export const bodySchema = S.object()
 	.prop("document_limit", S.number().minimum(1).maximum(10).default(3))
 	.prop("min_content_length", S.number().minimum(0).maximum(10000).default(50))
 	.prop("include_summary_in_response_generation", S.boolean().default(false))
+	.prop("generate_answer", S.boolean().default(true))
 	.prop("search_algorithm", S.string().default("chunks-and-summaries"))
 	.prop(
 		"openai_model",
@@ -124,3 +125,7 @@ export const responseSchema = {
 		.prop("completionOptions", createChatCompletionRequestSchema)
 		.prop("documentMatches", S.array().items(documentMatch)),
 };
+
+export const countSchema = {
+	200: S.object().prop("registered_documents_count", S.number()),
+};

tests/api_test.ts

@@ -0,0 +1,112 @@
+import fs from "fs";
+
+enum Algorithms {
+	ChunksAndSummaries = "chunks-and-summaries",
+	ChunksOnly = "chunks-only",
+	SummariesThenChunks = "summaries-then-chunks",
+}
+
+const availableAlgorithms = [
+	{
+		temperature: 0,
+		match_threshold: 0.85,
+		num_probes: 8,
+		openai_model: "gpt-3.5-turbo-16k",
+		document_limit: 3,
+		search_algorithm: Algorithms.ChunksOnly,
+		match_count: 64,
+		include_summary_in_response_generation: false,
+		generate_answer: false,
+	},
+	{
+		temperature: 0,
+		match_threshold: 0.85,
+		num_probes: 8,
+		openai_model: "gpt-3.5-turbo-16k",
+		chunk_limit: 128,
+		summary_limit: 16,
+		document_limit: 3,
+		search_algorithm: Algorithms.ChunksAndSummaries,
+		include_summary_in_response_generation: false,
+		generate_answer: false,
+	},
+	{
+		temperature: 0,
+		match_threshold: 0.85,
+		num_probes: 8,
+		openai_model: "gpt-3.5-turbo-16k",
+		document_limit: 3,
+		search_algorithm: Algorithms.SummariesThenChunks,
+		match_count: 64,
+		include_summary_in_response_generation: false,
+		generate_answer: false,
+	},
+];
+
+interface AlgorithmCount {
+	algorithm: Algorithms;
+	count: number;
+}
+
+let counts: Array<AlgorithmCount> = availableAlgorithms.map((alg) => {
+	return { algorithm: alg.search_algorithm, count: 0 } as AlgorithmCount;
+});
+
+const testQuestions = fs.readFileSync("tests/example_questions.json", "utf-8");
+const questions = JSON.parse(testQuestions); //.slice(0, 2);
+
+for (let algIdx = 0; algIdx < availableAlgorithms.length; algIdx++) {
+	const algorithm = availableAlgorithms[algIdx];
+	console.log(
+		`Testing search algorithm "${algorithm.search_algorithm}" with configuration...`,
+	);
+	console.log(algorithm);
+
+	for (let idx = 0; idx < questions.length; idx++) {
+		const question: string = questions[idx].question;
+		const groundTruthUrl: string = questions[idx].document;
+		const data = await fetch("http://0.0.0.0:8080/vector-search", {
+			method: "POST",
+			headers: {
+				"Content-Type": "application/json",
+			},
+			body: JSON.stringify({ ...algorithm, query: question }),
+		});
+		const res = await data.json();
+
+		if (res.documentMatches) {
+			const foundDocument = res.documentMatches.filter((d: any) => {
+				return d.registered_document.source_url.includes(groundTruthUrl);
+			})[0];
+			if (foundDocument) {
+				console.log(
+					`✅ Found the expected document ${groundTruthUrl} in search results, documentId=${foundDocument.registered_document.id} with similarity=${foundDocument.similarity}.`,
+				);
+				counts[algIdx] = {
+					algorithm: algorithm.search_algorithm,
+					count: counts[algIdx].count + 1,
+				};
+			} else {
+				console.log(
+					`❌ Did not find the expected document ${groundTruthUrl} in search results.`,
+				);
+			}
+		} else {
+			console.log(
+				`❌ Did not find the expected document ${groundTruthUrl} in search results.`,
+			);
+		}
+	}
+
+	console.log(`Tests completed for ${algorithm.search_algorithm}....\n\n`);
+}
+
+console.log(`All tests completed, summary:`);
+counts.forEach((c) => {
+	console.log(
+		`Algorithm "${c.algorithm}" has match rate = ${(
+			(c.count / questions.length) *
+			100
+		).toFixed(2)}% on test set of ${questions.length} questions.`,
+	);
+});