aisearch-docs.js

/**
 * Edge Function template for performing AI-powered search on databases created by the GitHub Action.
 * 
 * This template provides a ready-to-use Edge Function that users can copy and paste 
 * into their SQLite Cloud project to enable semantic search capabilities on the 
 * documentation database generated by the sqlite-aisearch-action GitHub Action.
 * 
 * The function combines vector embeddings with full-text search using Reciprocal 
 * Rank Fusion (RRF) to deliver accurate and contextually relevant search results.
 * 
 * See README.md for deployment instructions.
 */

//---- CONFIGURATION ----
const sqliteAIBaseUrl = "https://aiserver.vital-rhino.eks.euc1.ryujaz.sqlite.cloud";
const sqliteAIAPI = "/v1/ai/embeddings";
const topKSentences = 3;  // Number of top sentences to include in preview
const maxChars = 400;     // Maximum total characters for preview
const gap = "[...]";      // Gap indicator string
//-----------------------

const query = request.params.query;
const limit = parseInt(request.params.limit) || 5; // Number of top results to return

// Get embedding from sqlite-ai-server
const data = {"text": query };
const response = await fetch(sqliteAIBaseUrl + sqliteAIAPI, {
    method: "POST",
    headers: {
        "Content-Type": "application/json"
    },
    body: JSON.stringify(data)
});

if (!response.ok) {
    throw new Error("Failed to generate embedding. " + response.statusText);
}

const result = await response.json();
const query_embedding = result.data.embedding;

// Clean query for full-text search
const query_fts = (query.toLowerCase().match(/\b\w+\b/g) || []).join(" ") + "*";

// Vector configuration must match the embedding parameters used during database generation
await connection.sql("SELECT vector_init('chunks', 'embedding', 'type=INT8,dimension=768,distance=cosine')");
await connection.sql("SELECT vector_init('sentences', 'embedding', 'type=INT8,dimension=768,distance=cosine')");

const res = await connection.sql(
    `
        -- sqlite-vector KNN vector search results
    WITH vec_matches AS (
        SELECT
            v.rowid AS chunk_id,
            row_number() OVER (ORDER BY v.distance) AS rank_number,
            v.distance
        FROM vector_quantize_scan('chunks', 'embedding', ?, ?) AS v
    ),
    -- Full-text search results
    fts_matches AS (
        SELECT
            chunks_fts.rowid AS chunk_id,
            row_number() OVER (ORDER BY rank) AS rank_number,
            rank AS score
        FROM chunks_fts
        WHERE chunks_fts MATCH ?
        LIMIT ?
    ),
    -- combine FTS5 + vector search results with RRF
    matches AS (
        SELECT
            COALESCE(vec_matches.chunk_id, fts_matches.chunk_id) AS chunk_id,
            vec_matches.rank_number AS vec_rank,
            fts_matches.rank_number AS fts_rank,
            -- Reciprocal Rank Fusion score
            (
                COALESCE(1.0 / (60 + vec_matches.rank_number), 0.0) * 1.0 +
                COALESCE(1.0 / (60 + fts_matches.rank_number), 0.0) * 1.0
            ) AS combined_rank,
            vec_matches.distance AS vec_distance,
            fts_matches.score AS fts_score
        FROM vec_matches
            FULL OUTER JOIN fts_matches
                ON vec_matches.chunk_id = fts_matches.chunk_id
    )
        SELECT
            documents.id,
            documents.uri,
            documents.metadata,
            chunks.id AS chunk_id,
            chunks.content AS chunk_content,
            vec_rank,
            fts_rank,
            combined_rank,
            vec_distance,
            fts_score
        FROM matches
            JOIN chunks ON chunks.id = matches.chunk_id
            JOIN documents ON documents.id = chunks.document_id
    ORDER BY combined_rank DESC
    ;
    `, query_embedding, limit, query_fts, limit
);

// The results from the query may contain multiple resulting chunks per document.
// We want to return one result per document, so we will group by document id and take
// the top-ranked chunk as a snippet.
const seenDocuments = new Set();
const topResults = res
    .filter(item => !seenDocuments.has(item.id) && seenDocuments.add(item.id))
    .slice(0, limit);

// ----- Fetch top sentences for each top result -----
for (const result of topResults) {
    result.sentences = await connection.sql(
        `WITH vec_matches AS (
            SELECT
                v.rowid AS sentence_id,
                row_number() OVER (ORDER BY v.distance) AS rank_number,
                v.distance
            FROM vector_quantize_scan_stream('sentences', 'embedding', ?) AS v
                JOIN sentences ON sentences.rowid = v.rowid
            WHERE sentences.chunk_id = ?
            LIMIT ?
        )
        SELECT
            sentence_id,
            -- Extract sentence directly from document content
            COALESCE(
                substr(chunks.content, sentences.start_offset + 1, sentences.end_offset - sentences.start_offset),
                ""
            ) AS content,
            sentences.start_offset AS sentence_start_offset,
            sentences.end_offset AS sentence_end_offset,
            rank_number,
            distance
        FROM vec_matches
            JOIN sentences ON sentences.rowid = vec_matches.sentence_id
            JOIN chunks ON chunks.id = sentences.chunk_id
        ORDER BY rank_number ASC
        ;
        `, query_embedding, result.chunk_id, topKSentences
    );
}

// ----- Build snippets from sentences -----
for (const item of topResults) {
    const topSentences = item.sentences ? item.sentences.slice(0, topKSentences) : [];
    let snippet = "";
    
    if (topSentences.length === 0) {
        // Fallback: no sentences, return truncated chunk content
        const chunkContent = item.chunk_content || "";
        snippet = chunkContent.substring(0, maxChars);
    } else {
        // Sort by start_offset to maintain document order
        topSentences.sort((a, b) => {
            const offsetA = a.sentence_start_offset !== null ? a.sentence_start_offset : -1;
            const offsetB = b.sentence_start_offset !== null ? b.sentence_start_offset : -1;
            return offsetA - offsetB;
        });
        
        const previewParts = [];
        let totalChars = 0;
        let prevEndOffset = null;
        
        for (const sentence of topSentences) {
            const sentenceText = sentence.content;
                        
            // Check for gap between sentences
            if (prevEndOffset !== null && sentence.sentence_start_offset !== null) {
                const gapSize = sentence.sentence_start_offset - prevEndOffset;
                if (gapSize > 10) {
                    previewParts.push(gap);
                    totalChars += gap.length;
                }
            }
            
            previewParts.push(sentenceText);
            totalChars += sentenceText.length;
            prevEndOffset = sentence.sentence_end_offset;
        }
        
        const preview = previewParts.join(" ");
        snippet = preview.length > maxChars ? preview.substring(0, maxChars - 3) + "..." : preview;
    }
    
    item.snippet = snippet;
}

// ----- URLs for results -----
// Customize this section based on how URLs should be constructed for your documents.
// This example uses 'base_url' from metadata and 'slug' if available, otherwise derives from URI.
// ----------------------------
const finalResults = topResults
    .map(item => {
        const metadata = JSON.parse(item.metadata);
        const baseUrl = metadata.base_url;
        const slug = metadata.extracted?.slug;
        const uri = item.uri;
        
        let fullUrl;
        if (slug) {
            fullUrl = `${baseUrl}${slug}`;
        } else {
            const uriWithoutExtension = uri
                .toLowerCase()
                .replace(/\.(mdx?|md)$/i, '');
            fullUrl = `${baseUrl}${uriWithoutExtension}`;
        }
        
        return {
            id: item.id,
            url: fullUrl,
            title: metadata.extracted?.title || metadata.generated?.title,
            snippet: item.snippet
        };
    });

return {
    data: {
        /**
         * @type {Array<{id: number, url: string, title: string, snippet: string}>}
         * The search results with constructed URLs, titles, and snippets.
         */
        search: finalResults
    }
}