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

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

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<number[]>;
+  query: (sql: string, values: unknown[]) => Promise<{ rows: SearchResult[] }>;
+};
+
+const defaultDependencies: SearchDependencies = {
+  embed: generateEmbedding,
+  query: (sql, values) => pool.query<SearchResult>(sql, values),
+};
+
+export const searchDocuments = async (
+  query: string,
+  limit = 5,
+  dependencies: SearchDependencies = defaultDependencies,
+): Promise<SearchResult[]> => {
+  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;
+};

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