use std::{collections::HashMap, time::Duration}; use anyhow::{Context, Result, anyhow}; use log::{info, warn}; use once_cell::sync::Lazy; use reqwest::Client; use serde::Serialize; use sha2::{Digest, Sha256}; use sqlx::{FromRow, SqlitePool}; use crate::api::{FILE_COLS_NO_CONTENT, File}; static CLIENT: Lazy = Lazy::new(|| Client::builder().timeout(Duration::from_secs(120)).build().expect("build WeKnora sync HTTP client")); #[derive(Debug, Serialize)] struct ExternalChunk { external_slice_id: String, index: usize, content: String, } #[derive(Debug, Serialize)] struct ExternalImport { source: &'static str, external_file_id: String, source_version: String, title: String, filename: String, file_type: String, file_size: i64, file_hash: String, chunks: Vec, } #[derive(Debug, FromRow)] struct SyncEvent { event_id: String, event_type: String, file_id: i64, htknow_kb_id: i64, attempts: i64, } pub async fn enqueue_file_upsert(pool: &SqlitePool, file_id: i64) -> Result<()> { if !sync_enabled() { return Ok(()); } let kb_id: i64 = sqlx::query_scalar("SELECT kb_id FROM files WHERE id = ?") .bind(file_id) .fetch_optional(pool) .await? .flatten() .ok_or_else(|| anyhow!("file {} has no knowledge base", file_id))?; enqueue(pool, "upsert", file_id, kb_id).await } pub async fn enqueue_file_delete(pool: &SqlitePool, file_id: i64, htknow_kb_id: i64) -> Result<()> { if !sync_enabled() { return Ok(()); } enqueue(pool, "delete", file_id, htknow_kb_id).await } async fn enqueue(pool: &SqlitePool, event_type: &str, file_id: i64, htknow_kb_id: i64) -> Result<()> { let event_id = uuid::Uuid::new_v4().to_string(); sqlx::query( "INSERT INTO weknora_sync_outbox (event_id, event_type, file_id, htknow_kb_id, attempts, next_attempt_at, last_error, created_at, updated_at) VALUES (?, ?, ?, ?, 0, strftime('%s','now'), '', strftime('%s','now'), strftime('%s','now')) ON CONFLICT(file_id) DO UPDATE SET event_id = excluded.event_id, event_type = excluded.event_type, htknow_kb_id = excluded.htknow_kb_id, attempts = 0, next_attempt_at = excluded.next_attempt_at, last_error = '', updated_at = excluded.updated_at", ) .bind(event_id) .bind(event_type) .bind(file_id) .bind(htknow_kb_id) .execute(pool) .await?; Ok(()) } pub fn start_worker(pool: SqlitePool) { if !sync_enabled() { info!("WeKnora synchronization is disabled"); return; } tokio::spawn(async move { loop { if let Err(err) = process_due_events(&pool).await { warn!("WeKnora outbox worker failed: {}", err); } tokio::time::sleep(Duration::from_secs(2)).await; } }); } async fn process_due_events(pool: &SqlitePool) -> Result<()> { let events: Vec = sqlx::query_as( "SELECT event_id, event_type, file_id, htknow_kb_id, attempts FROM weknora_sync_outbox WHERE next_attempt_at <= strftime('%s','now') ORDER BY created_at, file_id LIMIT 10", ) .fetch_all(pool) .await?; for event in events { let result = match event.event_type.as_str() { "upsert" => sync_file(pool, event.file_id).await, "delete" => delete_file(event.file_id, event.htknow_kb_id).await, other => Err(anyhow!("unsupported outbox event type {}", other)), }; match result { Ok(()) => { sqlx::query("DELETE FROM weknora_sync_outbox WHERE file_id = ? AND event_id = ?") .bind(event.file_id) .bind(&event.event_id) .execute(pool) .await?; info!("WeKnora {} accepted for file {}", event.event_type, event.file_id); } Err(err) => { let attempts = event.attempts + 1; let delay = (2_i64.saturating_pow(attempts.min(10) as u32)).min(3600); sqlx::query( "UPDATE weknora_sync_outbox SET attempts = ?, next_attempt_at = strftime('%s','now') + ?, last_error = ?, updated_at = strftime('%s','now') WHERE file_id = ? AND event_id = ?", ) .bind(attempts) .bind(delay) .bind(err.to_string()) .bind(event.file_id) .bind(&event.event_id) .execute(pool) .await?; warn!( "WeKnora {} attempt {} failed for file {}; retry in {}s: {}", event.event_type, attempts, event.file_id, delay, err ); } } } Ok(()) } fn sync_enabled() -> bool { std::env::var("HTKNOW_WEKNORA_URL").ok().is_some_and(|v| !v.trim().is_empty()) } fn mapped_weknora_kb_id(htknow_kb_id: i64) -> Result { let kb_map_raw = std::env::var("HTKNOW_WEKNORA_KB_MAP").context("HTKNOW_WEKNORA_KB_MAP is required")?; let kb_map: HashMap = serde_json::from_str(&kb_map_raw).context("parse HTKNOW_WEKNORA_KB_MAP JSON")?; kb_map .get(&htknow_kb_id.to_string()) .or_else(|| kb_map.get("*")) .cloned() .ok_or_else(|| anyhow!("no WeKnora KB mapping for htknow KB {} and no '*' default mapping", htknow_kb_id)) } async fn sync_file(pool: &SqlitePool, file_id: i64) -> Result<()> { let file: File = sqlx::query_as(&format!("SELECT {} FROM files WHERE id = ?", FILE_COLS_NO_CONTENT)) .bind(file_id) .fetch_one(pool) .await .context("load parsed file")?; if file.status != 1 { return Err(anyhow!("file status is {}, expected completed", file.status)); } let htknow_kb_id = file.kb_id.ok_or_else(|| anyhow!("file has no knowledge base"))?; let weknora_kb_id = mapped_weknora_kb_id(htknow_kb_id)?; let source_id = crate::api::effective_parse_file_id(pool, file_id).await?; let ids: Vec = sqlx::query_scalar("SELECT id FROM slices WHERE file_id = ? ORDER BY ordinal, id") .bind(source_id) .fetch_all(pool) .await?; let contents = crate::slice_content::read_all(source_id).await?; let chunks = ids .into_iter() .enumerate() .filter_map(|(index, id)| { let content = contents.get(&id)?.trim(); if content.is_empty() { return None; } Some(ExternalChunk { external_slice_id: id.to_string(), index, content: content.to_string() }) }) .collect::>(); if chunks.is_empty() { return Err(anyhow!("parsed file contains no text chunks")); } let mut version_hasher = Sha256::new(); version_hasher.update(file.filename.as_bytes()); version_hasher.update(file.hash.as_bytes()); for chunk in &chunks { version_hasher.update(chunk.external_slice_id.as_bytes()); version_hasher.update([0]); version_hasher.update(chunk.content.as_bytes()); version_hasher.update([0xff]); } let source_version = hex::encode(version_hasher.finalize()); let file_type = std::path::Path::new(&file.filename).extension().and_then(|v| v.to_str()).unwrap_or("markdown").to_lowercase(); let payload = ExternalImport { source: "htknow", external_file_id: file.id.to_string(), source_version, title: file.filename.clone(), filename: file.filename, file_type, file_size: file.size, file_hash: file.hash, chunks, }; let base = std::env::var("HTKNOW_WEKNORA_URL")?; let url = format!("{}/api/v1/knowledge-bases/{}/knowledge/external-chunks", base.trim_end_matches('/'), weknora_kb_id); let mut request = CLIENT.post(url).json(&payload); if let Ok(key) = std::env::var("HTKNOW_WEKNORA_API_KEY") && !key.trim().is_empty() { request = request.header("X-API-Key", key); } let response = request.send().await.context("call WeKnora external chunk import")?; let status = response.status(); if !status.is_success() { let body = response.text().await.unwrap_or_default(); return Err(anyhow!("WeKnora returned {}: {}", status, body)); } Ok(()) } async fn delete_file(file_id: i64, htknow_kb_id: i64) -> Result<()> { let weknora_kb_id = mapped_weknora_kb_id(htknow_kb_id)?; let base = std::env::var("HTKNOW_WEKNORA_URL")?; let url = format!( "{}/api/v1/knowledge-bases/{}/knowledge/external-chunks/htknow/{}", base.trim_end_matches('/'), weknora_kb_id, file_id ); let mut request = CLIENT.delete(url); if let Ok(key) = std::env::var("HTKNOW_WEKNORA_API_KEY") && !key.trim().is_empty() { request = request.header("X-API-Key", key); } let response = request.send().await.context("call WeKnora external knowledge delete")?; let status = response.status(); if !status.is_success() { let body = response.text().await.unwrap_or_default(); return Err(anyhow!("WeKnora returned {}: {}", status, body)); } Ok(()) }