feat: add query reformulation and LLM reranking to search pipeline

server/lib/query-reformulation.ts

@@ -1,22 +1,22 @@
 import { generateText } from "ai";
 import { chatModel } from "./mesh-provider";
 
-const REFORMULATION_PROMPT = `You are a query optimizer for semantic search in a documentation database.
-Your task is to transform user questions into optimized search queries.
+const REFORMULATION_PROMPT = `You are a query optimizer for semantic vector search in a documentation database.
+Your task is to transform user questions into dense, meaning-rich sentences optimized for embedding similarity.
 
 Rules:
-1. Extract the main concepts and keywords from the question
-2. Remove filler words and question structure
-3. Keep technical terms intact
-4. Output a concise search query (max 10 words)
+1. Rephrase the question as a declarative statement that captures the full semantic meaning
+2. Keep all technical terms intact
+3. Include relevant synonyms or related terms inline to broaden recall
+4. Output a single concise sentence (max 20 words)
 5. Output ONLY the reformulated query, nothing else
-6. Keep the same language as the input
+6. Always output in English regardless of input language (embeddings work best in English)
 
 Examples:
-- "O que é uma section no deco.cx?" → "section deco.cx definição conceito"
-- "How do I create a new page?" → "create new page tutorial"
-- "Como faço para editar uma seção?" → "editar seção como fazer"
-- "What are the best practices for loaders?" → "loaders best practices guidelines"`;
+- "O que é uma section no deco.cx?" → "definition and concept of sections as UI building blocks in deco.cx"
+- "How do I create a new page?" → "guide to creating and setting up a new page in deco.cx"
+- "Como faço para editar uma seção?" → "how to edit and modify sections in deco.cx pages"
+- "What are the best practices for loaders?" → "best practices and guidelines for using loaders and data fetching in deco.cx"`;
 
 export async function reformulateQuery(query: string): Promise<string> {
   try {

server/lib/reranker.ts

@@ -0,0 +1,62 @@
+import { generateObject } from "ai";
+import { z } from "zod";
+import { chatModel } from "./mesh-provider";
+
+interface RankCandidate {
+  content: string;
+  title: string;
+  [key: string]: unknown;
+}
+
+const RERANK_PROMPT = `You are a relevance judge. Given a search query and a list of document excerpts, score each document's relevance to the query.
+
+Score each document from 0 to 10:
+- 0: Completely irrelevant
+- 1-3: Tangentially related but doesn't answer the query
+- 4-6: Partially relevant, contains some useful information
+- 7-9: Highly relevant, directly addresses the query
+- 10: Perfect match, exactly what the user is looking for
+
+Be strict. Most documents should score below 5 unless they truly address the query.`;
+
+export async function rerankResults<T extends RankCandidate>(
+  query: string,
+  candidates: T[],
+): Promise<T[]> {
+  if (candidates.length === 0) return candidates;
+
+  try {
+    const docsText = candidates
+      .map((doc, i) => `[${i}] ${doc.title}\n${doc.content.slice(0, 300)}`)
+      .join("\n\n");
+
+    const { object } = await generateObject({
+      model: chatModel(),
+      schema: z.object({
+        scores: z
+          .array(z.number().min(0).max(10))
+          .describe("Relevance score for each document, in the same order"),
+      }),
+      system: RERANK_PROMPT,
+      prompt: `Query: ${query}\n\nDocuments:\n${docsText}`,
+      temperature: 0,
+    });
+
+    const { scores } = object;
+
+    if (scores.length !== candidates.length) {
+      console.warn(
+        `Reranker returned ${scores.length} scores for ${candidates.length} docs, using original order`,
+      );
+      return candidates;
+    }
+
+    const indexed = candidates.map((doc, i) => ({ doc, score: scores[i] }));
+    indexed.sort((a, b) => b.score - a.score);
+
+    return indexed.map((item) => item.doc);
+  } catch (error) {
+    console.error("Reranking failed, using original order:", error);
+    return candidates;
+  }
+}

server/tools/assistant.ts

@@ -3,6 +3,8 @@ import { createClient } from "@supabase/supabase-js";
 import { embed, generateText } from "ai";
 import { z } from "zod";
 import { chatModel, embeddingModel } from "../lib/mesh-provider";
+import { reformulateQuery } from "../lib/query-reformulation";
+import { rerankResults } from "../lib/reranker";
 
 const supabase = createClient(
   process.env.SUPABASE_URL!,
@@ -31,16 +33,21 @@ export const assistantTool = createTool({
   execute: async ({ context }) => {
     const { question, language } = context;
 
+    const RETRIEVAL_POOL_SIZE = 20;
+    const TOP_K = 5;
+
+    const optimizedQuery = await reformulateQuery(question);
+
     const { embedding } = await embed({
       model: embeddingModel(),
-      value: question,
+      value: optimizedQuery,
     });
 
     const filter = language ? { language } : {};
     const { data: docs, error } = await supabase.rpc("hybrid_search", {
       query_text: question,
       query_embedding: embedding,
-      match_count: 5,
+      match_count: RETRIEVAL_POOL_SIZE,
       rrf_k: 60,
       semantic_weight: 0.5,
       filter_metadata: filter,
@@ -50,8 +57,17 @@ export const assistantTool = createTool({
       throw new Error(`Search failed: ${error.message}`);
     }
 
-    const docsContext = (docs || [])
-      .map((doc: any, i: number) => `[${i + 1}] ${doc.metadata?.title || "Doc"}\n${doc.content}`)
+    const candidates = (docs || []).map((doc: any) => ({
+      content: doc.content,
+      title: doc.metadata?.title || "Untitled",
+      source: doc.metadata?.source || "",
+    }));
+
+    const reranked = await rerankResults(question, candidates);
+    const topDocs = reranked.slice(0, TOP_K);
+
+    const docsContext = topDocs
+      .map((doc, i) => `[${i + 1}] ${doc.title}\n${doc.content}`)
       .join("\n\n---\n\n");
 
     const { text } = await generateText({
@@ -64,9 +80,9 @@ export const assistantTool = createTool({
 
     const answer = text || "Sorry, I couldn't generate an answer.";
 
-    const sources = (docs || []).map((doc: any) => ({
-      title: doc.metadata?.title || "Untitled",
-      source: doc.metadata?.source || "",
+    const sources = topDocs.map((doc) => ({
+      title: doc.title,
+      source: doc.source,
     }));
 
     return { answer, sources };

server/tools/search-docs.ts

@@ -3,6 +3,8 @@ import { createClient } from "@supabase/supabase-js";
 import { embed } from "ai";
 import { z } from "zod";
 import { embeddingModel } from "../lib/mesh-provider";
+import { reformulateQuery } from "../lib/query-reformulation";
+import { rerankResults } from "../lib/reranker";
 
 const supabase = createClient(
   process.env.SUPABASE_URL!,
@@ -11,7 +13,7 @@ const supabase = createClient(
 
 export const searchDocsTool = createTool({
   id: "SEARCH_DOCS",
-  description: "Search the deco.cx documentation using hybrid search (semantic + full-text). Returns relevant documentation chunks based on the query.",
+  description: "Search the deco.cx documentation. Returns relevant results based on the query.",
   inputSchema: z.object({
     query: z.string().describe("The search query in natural language"),
     language: z.enum(["en", "pt"]).optional().describe("Filter by language (optional)"),
@@ -32,17 +34,21 @@ export const searchDocsTool = createTool({
   execute: async ({ context }) => {
     const { query, language, limit = 5, semanticWeight = 0.5 } = context;
 
+    const RETRIEVAL_POOL_SIZE = 20;
+
+    const optimizedQuery = await reformulateQuery(query);
+
     const { embedding } = await embed({
       model: embeddingModel(),
-      value: query,
+      value: optimizedQuery,
     });
 
     const filter = language ? { language } : {};
 
     const { data, error } = await supabase.rpc("hybrid_search", {
       query_text: query,
       query_embedding: embedding,
-      match_count: limit,
+      match_count: RETRIEVAL_POOL_SIZE,
       rrf_k: 60,
       semantic_weight: semanticWeight,
       filter_metadata: filter,
@@ -52,7 +58,7 @@ export const searchDocsTool = createTool({
       throw new Error(`Search failed: ${error.message}`);
     }
 
-    const results = (data || []).map((doc: any) => ({
+    const candidates = (data || []).map((doc: any) => ({
       content: doc.content,
       title: doc.metadata?.title || "Untitled",
       source: doc.metadata?.source || "",
@@ -62,6 +68,8 @@ export const searchDocsTool = createTool({
       hybridScore: doc.hybrid_score,
     }));
 
-    return { results };
+    const reranked = await rerankResults(query, candidates);
+
+    return { results: reranked.slice(0, limit) };
   },
 });