open-notebook/open_notebook/database/migrations/4.surrealql

REMOVE FUNCTION IF EXISTS fn::text_search;


DEFINE FUNCTION IF NOT EXISTS fn::text_search($query_text: string, $match_count: int, $sources:bool, $show_notes:bool) {
  
    let $source_title_search = 
        IF $sources {(
            SELECT id, title, 
            search::highlight('`', '`', 1) as content,
            id as parent_id,
            math::max(search::score(1)) AS relevance
            FROM source
            WHERE title @1@ $query_text
            GROUP BY id)}
        ELSE { [] };
    
    let $source_embedding_search = 
         IF $sources {(
            SELECT source.id as id, source.title as title, search::highlight('`', '`', 1) as content, source.id as parent_id, math::max(search::score(1)) AS relevance
            FROM source_embedding
            WHERE content @1@ $query_text
            GROUP BY id)}
        ELSE { [] };

    let $source_full_search = 
         IF $sources {(
            SELECT id, title, search::highlight('`', '`', 1) as content, id as parent_id, math::max(search::score(1)) AS relevance
            FROM source
            WHERE full_text @1@ $query_text
            GROUP BY id)}
        ELSE { [] };
    
    let $source_insight_search = 
         IF $sources {(
             SELECT id, insight_type + " - " + (source.title OR '') as title, search::highlight('`', '`', 1) as content, id as parent_id,  math::max(search::score(1)) AS relevance
            FROM source_insight
            WHERE content @1@ $query_text
            GROUP BY id)}
        ELSE { [] };

    let $note_title_search = 
         IF $show_notes {(
             SELECT id, title, search::highlight('`', '`', 1) as content,  id as parent_id, math::max(search::score(1)) AS relevance
            FROM note
            WHERE title @1@ $query_text
            GROUP BY id)}
        ELSE { [] };

     let $note_content_search = 
         IF $show_notes {(
             SELECT id, title, search::highlight('`', '`', 1) as content,  id as parent_id, math::max(search::score(1)) AS relevance
            FROM note
            WHERE content @1@ $query_text
            GROUP BY id)}
        ELSE { [] };

    let $source_chunk_results = array::union($source_embedding_search, $source_full_search);
    
    let $source_asset_results = array::union($source_title_search, $source_insight_search);

    let $source_results = array::union($source_chunk_results, $source_asset_results );
    let $note_results = array::union($note_title_search, $note_content_search );
    let $final_results = array::union($source_results, $note_results );

        RETURN (select id, parent_id, title, math::max(relevance) as relevance
        from $final_results where id is not None
        group by id, parent_id, title ORDER BY relevance DESC LIMIT $match_count);

};


REMOVE FUNCTION IF EXISTS fn::vector_search;

DEFINE FUNCTION IF NOT EXISTS fn::vector_search($query: array<float>, $match_count: int, $sources: bool, $show_notes: bool, $min_similarity: float) {
    let $source_embedding_search = 
        IF $sources {(
            SELECT 
                source.id as id,
                source.title as title,
                content,
                source.id as parent_id,
                vector::similarity::cosine(embedding, $query) as similarity
            FROM source_embedding 
            WHERE vector::similarity::cosine(embedding, $query) >= $min_similarity
            ORDER BY similarity DESC
            LIMIT $match_count
        )}
        ELSE { [] };

    let $source_insight_search = 
        IF $sources {(
            SELECT 
                id,
                insight_type + ' - ' + (source.title OR '') as title,
                content,
                source.id as parent_id,
                vector::similarity::cosine(embedding, $query) as similarity
            FROM source_insight
            WHERE vector::similarity::cosine(embedding, $query) >= $min_similarity
            ORDER BY similarity DESC
            LIMIT $match_count
        )}
        ELSE { [] };


    let $note_content_search = 
        IF $show_notes {(
            SELECT 
                id,
                title,
                content,
                id as parent_id,
                vector::similarity::cosine(embedding, $query) as similarity
            FROM note
            WHERE vector::similarity::cosine(embedding, $query) >= $min_similarity
            ORDER BY similarity DESC
            LIMIT $match_count
        )}
        ELSE { [] };


    let $all_results = array::union(
        array::union($source_embedding_search, $source_insight_search),
        $note_content_search
    );


    RETURN (select id, parent_id, title, math::max(similarity) as similarity,
    array::flatten(content) as matches
    from $all_results where id is not None
    group by id, parent_id, title ORDER BY similarity DESC LIMIT $match_count);

};