relay-knowledge 1.1.5

Graph-database-based knowledge graph project.
Documentation
use std::{thread, time::Duration};

use rusqlite::{Connection, params};

use crate::storage::{GraphSearchRequest, StorageError};

const GRAPH_BM25_QUERY_RETRY_DELAYS_MS: [u64; 3] = [5, 15, 45];

pub(super) struct RawBm25Row {
    pub(super) document_id: String,
    pub(super) document_kind: String,
    pub(super) evidence_id: String,
    pub(super) parent_evidence_id: Option<String>,
    pub(super) modality: String,
    pub(super) source_scope: String,
    pub(super) source_path: Option<String>,
    pub(super) entity_labels: Vec<String>,
    pub(super) content: String,
    pub(super) rank: f64,
}

pub(super) fn bm25_candidate_rows(
    connection: &Connection,
    request: &GraphSearchRequest,
    match_query: &str,
) -> Result<Vec<RawBm25Row>, StorageError> {
    for delay_ms in GRAPH_BM25_QUERY_RETRY_DELAYS_MS {
        match bm25_candidate_rows_once(connection, request, match_query) {
            Ok(rows) => return Ok(rows),
            Err(error) if graph_bm25_query_error_is_retryable(&error) => {
                thread::sleep(Duration::from_millis(delay_ms));
            }
            Err(error) => return Err(error),
        }
    }

    bm25_candidate_rows_once(connection, request, match_query)
}

fn bm25_candidate_rows_once(
    connection: &Connection,
    request: &GraphSearchRequest,
    match_query: &str,
) -> Result<Vec<RawBm25Row>, StorageError> {
    let mut statement = connection.prepare(
        "
        SELECT
            graph_bm25.document_id,
            graph_bm25.document_kind,
            graph_bm25.evidence_id,
            graph_bm25.parent_evidence_id,
            graph_bm25.modality,
            graph_bm25.source_scope,
            graph_bm25.source_path,
            graph_bm25.entity_labels,
            graph_bm25.content,
            bm25(graph_bm25) AS rank
        FROM graph_bm25
        LEFT JOIN evidence e
          ON graph_bm25.document_kind = 'evidence'
         AND e.id = graph_bm25.evidence_id
        WHERE graph_bm25 MATCH ?1
          AND (?2 IS NULL OR graph_bm25.source_scope = ?2)
          AND graph_bm25.created_graph_version <= ?3
          AND (
              graph_bm25.document_kind != 'evidence'
              OR e.status IN ('accepted', 'proposed')
          )
        ORDER BY rank ASC, graph_bm25.document_id ASC
        LIMIT ?4
        ",
    )?;
    let rows = statement.query_map(
        params![
            match_query,
            request.source_scope.as_deref(),
            request.graph_version.get(),
            request.limit.saturating_mul(4).max(request.limit)
        ],
        |row| {
            Ok(RawBm25Row {
                document_id: row.get(0)?,
                document_kind: row.get(1)?,
                evidence_id: row.get(2)?,
                parent_evidence_id: row.get(3)?,
                modality: row.get(4)?,
                source_scope: row.get(5)?,
                source_path: row.get(6)?,
                entity_labels: super::split_labels(row.get(7)?),
                content: row.get(8)?,
                rank: row.get(9)?,
            })
        },
    )?;
    rows.collect::<Result<Vec<_>, _>>()
        .map_err(StorageError::from)
}

pub(super) fn graph_bm25_query_error_is_retryable(error: &StorageError) -> bool {
    match error {
        StorageError::Sqlite(error) => {
            graph_bm25_query_error_message_is_retryable(&error.to_string())
        }
        _ => false,
    }
}

pub(super) fn graph_bm25_query_error_message_is_retryable(message: &str) -> bool {
    super::graph_bm25_transient_error_message(message)
}