From da134a1689ac033dd7df4a84f7a7175e0f534af9 Mon Sep 17 00:00:00 2001 From: Brandon Wolfgang Date: Thu, 22 Oct 2026 11:10:10 -0500 Subject: [PATCH] feat: implement semantic document search - Generate query embeddings using Ollama and normalize search input - Retrieve the most similar chunks from completed documents using pgvector cosine distance - Support configurable result limits from 1 to 20, defaulting to 5 - Add tests for query validation, vector search, result limits, and embedding errors - Inject embedding and database query dependencies for isolated testing --- src/search.test.ts | 59 ++++++++++++++++++++++++++++++++++++++++++++++ src/search.ts | 56 +++++++++++++++++++++++++++++++++++++++++++ src/server.ts | 1 + 3 files changed, 116 insertions(+) create mode 100644 src/search.test.ts create mode 100644 src/search.ts diff --git a/src/search.test.ts b/src/search.test.ts new file mode 100644 index 0000000..34ae0a8 --- /dev/null +++ b/src/search.test.ts @@ -0,0 +1,59 @@ +import assert from "node:assert/strict"; +import { describe, it } from "node:test"; + +import { searchDocuments } from "./search.js"; + +describe("searchDocuments", () => { + it("rejects an empty query", async () => { + await assert.rejects( + searchDocuments(" "), + /Search query cannot be empty/, + ); + }); + + it("embeds the query and returns matching chunks", async () => { + let embeddedText = ""; + let executedSql = ""; + let queryValues: unknown[] = []; + + const expectedResults = [ + { + documentId: "1", + title: "RAG smoke test", + content: "Relevant knowledge base content.", + chunkIndex: 0, + similarity: 0.92, + }, + ]; + + const results = await searchDocuments(" test query ", 3, { + embed: async (text) => { + embeddedText = text; + return [0.1, 0.2, 0.3]; + }, + query: async (sql, values) => { + executedSql = sql; + queryValues = values; + return { rows: expectedResults }; + }, + }); + + assert.equal(embeddedText, "test query"); + assert.match(executedSql, /ORDER BY dc\.embedding <=> \$1::vector/); + assert.match(executedSql, /d\.status = 'completed'/); + assert.deepEqual(queryValues, ["[0.1,0.2,0.3]", 3]); + assert.deepEqual(results, expectedResults); + }); + + it("propagates embedding errors", async () => { + await assert.rejects( + searchDocuments("test query", 5, { + embed: async () => { + throw new Error("Embedding unavailable"); + }, + query: async () => ({ rows: [] }), + }), + /Embedding unavailable/, + ); + }); +}); diff --git a/src/search.ts b/src/search.ts new file mode 100644 index 0000000..6b8cf9b --- /dev/null +++ b/src/search.ts @@ -0,0 +1,56 @@ +import { pool } from "./db.js"; +import { generateEmbedding } from "./ollama.js"; + +export interface SearchResult { + documentId: string; + title: string; + content: string; + chunkIndex: number; + similarity: number; +} + +type SearchDependencies = { + embed: (text: string) => Promise; + query: (sql: string, values: unknown[]) => Promise<{ rows: SearchResult[] }>; +}; + +const defaultDependencies: SearchDependencies = { + embed: generateEmbedding, + query: (sql, values) => pool.query(sql, values), +}; + +export const searchDocuments = async ( + query: string, + limit = 5, + dependencies: SearchDependencies = defaultDependencies, +): Promise => { + const normalizedQuery = query.trim(); + + if (!normalizedQuery) { + throw new Error("Search query cannot be empty"); + } + + if (!Number.isInteger(limit) || limit < 1 || limit > 20) { + throw new Error("Search limit must be an integer between 1 and 20"); + } + + const embedding = await dependencies.embed(normalizedQuery); + const vector = `[${embedding.join(",")}]`; + + const result = await dependencies.query( + `SELECT + d.id::text AS "documentId", + d.title, + dc.content, + dc.chunk_index AS "chunkIndex", + (1 - (dc.embedding <=> $1::vector))::float8 AS similarity + FROM document_chunks dc + JOIN documents d ON d.id = dc.document_id + WHERE d.status = 'completed' + ORDER BY dc.embedding <=> $1::vector + LIMIT $2`, + [vector, limit], + ); + + return result.rows; +}; diff --git a/src/server.ts b/src/server.ts index 8c473d7..00601c6 100644 --- a/src/server.ts +++ b/src/server.ts @@ -3,6 +3,7 @@ import { z } from "zod"; import { ingestDocument as ingestDocumentDefault } from "./ingestion.js"; import { checkOllamaHealth, generateResponse } from "./ollama.js"; +import { searchDocuments as searchDocumentsDefault } from "./search.js"; const echoSchema = z.object({ message: z.string().min(1).max(500),