diff --git a/server/lib/query-reformulation.ts b/server/lib/query-reformulation.ts index 85ecffb..5af2bc9 100644 --- a/server/lib/query-reformulation.ts +++ b/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 { try { diff --git a/server/lib/reranker.ts b/server/lib/reranker.ts new file mode 100644 index 0000000..189bbe2 --- /dev/null +++ b/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( + query: string, + candidates: T[], +): Promise { + 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; + } +} diff --git a/server/tools/assistant.ts b/server/tools/assistant.ts index e323962..689a3b1 100644 --- a/server/tools/assistant.ts +++ b/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 }; diff --git a/server/tools/search-docs.ts b/server/tools/search-docs.ts index ca86cb5..bf9af59 100644 --- a/server/tools/search-docs.ts +++ b/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,9 +34,13 @@ 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 } : {}; @@ -42,7 +48,7 @@ export const searchDocsTool = createTool({ 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) }; }, });