feat: synchronize processed files to WeKnora
This commit is contained in:
parent
fefc5cf80a
commit
1ce10dc1f2
@ -73,8 +73,9 @@ encoding_rs = "0.8"
|
|||||||
etcd-client = { version = "0.14", optional = true }
|
etcd-client = { version = "0.14", optional = true }
|
||||||
serde_with = "3.16.1"
|
serde_with = "3.16.1"
|
||||||
sysinfo = "0.29"
|
sysinfo = "0.29"
|
||||||
# jemalloc allocator
|
# jemalloc is only enabled on non-Windows targets. Windows uses the system
|
||||||
# Note: profiling is only needed for debug/analysis, disabled in release for better performance
|
# allocator because tikv-jemalloc-sys does not reliably build for MSVC.
|
||||||
|
[target.'cfg(not(windows))'.dependencies]
|
||||||
tikv-jemallocator = { version = "0.6" }
|
tikv-jemallocator = { version = "0.6" }
|
||||||
tikv-jemalloc-ctl = { version = "0.6", optional = true, features = ["use_std"] }
|
tikv-jemalloc-ctl = { version = "0.6", optional = true, features = ["use_std"] }
|
||||||
|
|
||||||
|
|||||||
@ -1,9 +1,11 @@
|
|||||||
use sqlx::{QueryBuilder, Sqlite, SqlitePool};
|
use sqlx::{QueryBuilder, Sqlite, SqlitePool};
|
||||||
|
|
||||||
use crate::{
|
use crate::{
|
||||||
AuthUser, api::{
|
AuthUser,
|
||||||
error::{ApiError, ApiResult}, knowledge_base
|
api::{
|
||||||
}
|
error::{ApiError, ApiResult},
|
||||||
|
knowledge_base,
|
||||||
|
},
|
||||||
};
|
};
|
||||||
|
|
||||||
/// 要求当前用户为 admin,否则返回 BadRequest。
|
/// 要求当前用户为 admin,否则返回 BadRequest。
|
||||||
|
|||||||
@ -1,7 +1,9 @@
|
|||||||
use std::{fmt, num::ParseIntError, result};
|
use std::{fmt, num::ParseIntError, result};
|
||||||
|
|
||||||
use axum::{
|
use axum::{
|
||||||
Json, http::StatusCode, response::{IntoResponse, Response}
|
Json,
|
||||||
|
http::StatusCode,
|
||||||
|
response::{IntoResponse, Response},
|
||||||
};
|
};
|
||||||
use serde::Serialize;
|
use serde::Serialize;
|
||||||
use sqlx;
|
use sqlx;
|
||||||
|
|||||||
@ -456,7 +456,7 @@ pub async fn upload(
|
|||||||
|
|
||||||
match field_name.as_deref() {
|
match field_name.as_deref() {
|
||||||
Some("file") => {
|
Some("file") => {
|
||||||
let filename = field.file_name().unwrap_or("unknown").to_string();
|
let filename = normalize_upload_filename(field.file_name().unwrap_or("unknown"));
|
||||||
debug!("Uploading file: {}", filename);
|
debug!("Uploading file: {}", filename);
|
||||||
let tempname = uuid::Uuid::new_v4().to_string();
|
let tempname = uuid::Uuid::new_v4().to_string();
|
||||||
let filepath = format!("{}/{}", dir, tempname);
|
let filepath = format!("{}/{}", dir, tempname);
|
||||||
@ -755,6 +755,54 @@ pub async fn upload(
|
|||||||
Ok(Json(uploaded_files))
|
Ok(Json(uploaded_files))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn normalize_upload_filename(raw: &str) -> String {
|
||||||
|
if raw.chars().any(is_cjk) {
|
||||||
|
return raw.to_string();
|
||||||
|
}
|
||||||
|
let Some(bytes) = raw.chars().map(|ch| u8::try_from(ch as u32).ok()).collect::<Option<Vec<_>>>() else {
|
||||||
|
return raw.to_string();
|
||||||
|
};
|
||||||
|
if let Ok(decoded) = std::str::from_utf8(&bytes)
|
||||||
|
&& decoded.chars().any(is_cjk)
|
||||||
|
{
|
||||||
|
return decoded.to_string();
|
||||||
|
}
|
||||||
|
let high_bytes = bytes.iter().filter(|byte| **byte >= 0x80).count();
|
||||||
|
if high_bytes >= 4 {
|
||||||
|
let (decoded, _, had_errors) = encoding_rs::GBK.decode(&bytes);
|
||||||
|
if !had_errors && decoded.chars().any(is_cjk) {
|
||||||
|
return decoded.into_owned();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
raw.to_string()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn is_cjk(ch: char) -> bool {
|
||||||
|
matches!(ch as u32, 0x3400..=0x4DBF | 0x4E00..=0x9FFF | 0xF900..=0xFAFF)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod upload_filename_tests {
|
||||||
|
use super::normalize_upload_filename;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn repairs_utf8_bytes_misread_as_latin1() {
|
||||||
|
assert_eq!(normalize_upload_filename("å\u{8f}\u{91}å\u{8a}¨æ\u{9c}º.pdf"), "发动机.pdf");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn repairs_gbk_bytes_misread_as_latin1() {
|
||||||
|
assert_eq!(normalize_upload_filename("À×´ï-άÐÞÊÖ²á.pdf"), "雷达-维修手册.pdf");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn preserves_valid_names() {
|
||||||
|
assert_eq!(normalize_upload_filename("发动机.pdf"), "发动机.pdf");
|
||||||
|
assert_eq!(normalize_upload_filename("resume.pdf"), "resume.pdf");
|
||||||
|
assert_eq!(normalize_upload_filename("résumé.pdf"), "résumé.pdf");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// 获取文件详情
|
/// 获取文件详情
|
||||||
#[utoipa::path(
|
#[utoipa::path(
|
||||||
get,
|
get,
|
||||||
@ -1676,6 +1724,16 @@ async fn execute_batch_delete(
|
|||||||
step_start.elapsed().as_millis()
|
step_start.elapsed().as_millis()
|
||||||
);
|
);
|
||||||
|
|
||||||
|
// Local database and index deletion has committed. Preserve the captured
|
||||||
|
// KB IDs long enough to propagate idempotent delete events to WeKnora.
|
||||||
|
for file in &files_to_delete {
|
||||||
|
if let Some(kb_id) = file.kb_id {
|
||||||
|
if let Err(err) = crate::weknora_sync::enqueue_file_delete(pool, file.id, kb_id).await {
|
||||||
|
warn!("Failed to enqueue WeKnora delete for file {}: {}", file.id, err);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
debug!(
|
debug!(
|
||||||
"file_batch_delete total requested={} accepted={} deleted={} {}ms",
|
"file_batch_delete total requested={} accepted={} deleted={} {}ms",
|
||||||
requested,
|
requested,
|
||||||
@ -2472,6 +2530,13 @@ pub async fn update_slices(
|
|||||||
warn!("Failed to update full search index for file {}: {}", id, e);
|
warn!("Failed to update full search index for file {}: {}", id, e);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// The slice contents and htknow indexes are committed at this point.
|
||||||
|
// Notify WeKnora as well so an edited document updates its external
|
||||||
|
// knowledge record without requiring the source file to be reparsed.
|
||||||
|
if let Err(err) = crate::weknora_sync::enqueue_file_upsert(&pool, id).await {
|
||||||
|
warn!("Failed to enqueue WeKnora sync for file {}: {}", id, err);
|
||||||
|
}
|
||||||
|
|
||||||
// 8. 返回更新后的切片列表
|
// 8. 返回更新后的切片列表
|
||||||
get_slices(State(pool), Path(id)).await
|
get_slices(State(pool), Path(id)).await
|
||||||
}
|
}
|
||||||
|
|||||||
@ -1,7 +1,9 @@
|
|||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
|
|
||||||
use axum::{
|
use axum::{
|
||||||
Extension, extract::{Path, Query, State}, response::Json
|
Extension,
|
||||||
|
extract::{Path, Query, State},
|
||||||
|
response::Json,
|
||||||
};
|
};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use sqlx::SqlitePool;
|
use sqlx::SqlitePool;
|
||||||
|
|||||||
@ -5,9 +5,12 @@ use sqlx::SqlitePool;
|
|||||||
use utoipa::ToSchema;
|
use utoipa::ToSchema;
|
||||||
|
|
||||||
use crate::{
|
use crate::{
|
||||||
AuthUser, api::{
|
AuthUser,
|
||||||
common, error::{ApiError, ApiResult}
|
api::{
|
||||||
}, search::{SearchEngine, tantivy_engine::ForceMergeStats}
|
common,
|
||||||
|
error::{ApiError, ApiResult},
|
||||||
|
},
|
||||||
|
search::{SearchEngine, tantivy_engine::ForceMergeStats},
|
||||||
};
|
};
|
||||||
|
|
||||||
/// 进程内存占用
|
/// 进程内存占用
|
||||||
@ -350,14 +353,17 @@ fn force_merge_index_stats(index: &str, stats: ForceMergeStats) -> TantivyForceM
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(feature = "profiling")]
|
#[cfg(all(feature = "profiling", not(windows)))]
|
||||||
async fn heap_profile_impl() -> ApiResult<Response> {
|
async fn heap_profile_impl() -> ApiResult<Response> {
|
||||||
use std::{
|
use std::{
|
||||||
ffi::CString, io::Seek, time::{SystemTime, UNIX_EPOCH}
|
ffi::CString,
|
||||||
|
io::Seek,
|
||||||
|
time::{SystemTime, UNIX_EPOCH},
|
||||||
};
|
};
|
||||||
|
|
||||||
use axum::{
|
use axum::{
|
||||||
body::Body, http::{HeaderValue, header}
|
body::Body,
|
||||||
|
http::{HeaderValue, header},
|
||||||
};
|
};
|
||||||
use tempfile::NamedTempFile;
|
use tempfile::NamedTempFile;
|
||||||
use tikv_jemalloc_ctl::raw;
|
use tikv_jemalloc_ctl::raw;
|
||||||
@ -407,19 +413,22 @@ async fn heap_profile_impl() -> ApiResult<Response> {
|
|||||||
Ok(response)
|
Ok(response)
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(not(feature = "profiling"))]
|
#[cfg(any(not(feature = "profiling"), windows))]
|
||||||
async fn heap_profile_impl() -> ApiResult<Response> {
|
async fn heap_profile_impl() -> ApiResult<Response> {
|
||||||
Err(ApiError::BadRequest("Heap profiling is only available with 'profiling' feature enabled.".to_string()))
|
Err(ApiError::BadRequest("Heap profiling is only available with 'profiling' feature enabled.".to_string()))
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(feature = "profiling")]
|
#[cfg(all(feature = "profiling", not(windows)))]
|
||||||
async fn heap_profile_pdf_impl() -> ApiResult<Response> {
|
async fn heap_profile_pdf_impl() -> ApiResult<Response> {
|
||||||
use std::{
|
use std::{
|
||||||
ffi::CString, process::Stdio, time::{SystemTime, UNIX_EPOCH}
|
ffi::CString,
|
||||||
|
process::Stdio,
|
||||||
|
time::{SystemTime, UNIX_EPOCH},
|
||||||
};
|
};
|
||||||
|
|
||||||
use axum::{
|
use axum::{
|
||||||
body::Body, http::{HeaderValue, header}
|
body::Body,
|
||||||
|
http::{HeaderValue, header},
|
||||||
};
|
};
|
||||||
use tempfile::NamedTempFile;
|
use tempfile::NamedTempFile;
|
||||||
use tikv_jemalloc_ctl::raw;
|
use tikv_jemalloc_ctl::raw;
|
||||||
@ -512,7 +521,7 @@ async fn heap_profile_pdf_impl() -> ApiResult<Response> {
|
|||||||
Ok(response)
|
Ok(response)
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(not(feature = "profiling"))]
|
#[cfg(any(not(feature = "profiling"), windows))]
|
||||||
async fn heap_profile_pdf_impl() -> ApiResult<Response> {
|
async fn heap_profile_pdf_impl() -> ApiResult<Response> {
|
||||||
Err(ApiError::BadRequest("Heap profiling is only available with 'profiling' feature enabled.".to_string()))
|
Err(ApiError::BadRequest("Heap profiling is only available with 'profiling' feature enabled.".to_string()))
|
||||||
}
|
}
|
||||||
|
|||||||
@ -4,7 +4,8 @@
|
|||||||
//! 7Z 和 RAR 格式暂不支持。
|
//! 7Z 和 RAR 格式暂不支持。
|
||||||
|
|
||||||
use std::{
|
use std::{
|
||||||
io::{Read, Write}, path::Path
|
io::{Read, Write},
|
||||||
|
path::Path,
|
||||||
};
|
};
|
||||||
|
|
||||||
use encoding_rs::GB18030;
|
use encoding_rs::GB18030;
|
||||||
|
|||||||
@ -380,7 +380,7 @@ fn env_or_parse<T: std::str::FromStr>(key: &str, default: T) -> T {
|
|||||||
|
|
||||||
/// 配置加载器 trait
|
/// 配置加载器 trait
|
||||||
#[allow(dead_code)]
|
#[allow(dead_code)]
|
||||||
pub trait ConfigLoader: Send+Sync {
|
pub trait ConfigLoader: Send + Sync {
|
||||||
/// 加载配置
|
/// 加载配置
|
||||||
fn load(&self) -> anyhow::Result<AppConfig>;
|
fn load(&self) -> anyhow::Result<AppConfig>;
|
||||||
}
|
}
|
||||||
@ -413,7 +413,8 @@ impl EtcdConfigLoader {
|
|||||||
/// 启动 watch 监听配置变化
|
/// 启动 watch 监听配置变化
|
||||||
pub async fn watch<F>(&self, _callback: F) -> anyhow::Result<()>
|
pub async fn watch<F>(&self, _callback: F) -> anyhow::Result<()>
|
||||||
where
|
where
|
||||||
F: Fn(AppConfig)+Send+Sync+'static, {
|
F: Fn(AppConfig) + Send + Sync + 'static,
|
||||||
|
{
|
||||||
// TODO: 实现 etcd watch 逻辑
|
// TODO: 实现 etcd watch 逻辑
|
||||||
// 1. 连接 etcd
|
// 1. 连接 etcd
|
||||||
// 2. 监听 prefix 下的 key 变化
|
// 2. 监听 prefix 下的 key 变化
|
||||||
|
|||||||
41
src/db.rs
41
src/db.rs
@ -131,6 +131,7 @@ async fn run_schema_migrations(pool: &SqlitePool) -> anyhow::Result<()> {
|
|||||||
run_image_description_migration(pool).await?;
|
run_image_description_migration(pool).await?;
|
||||||
run_file_summary_migration(pool).await?;
|
run_file_summary_migration(pool).await?;
|
||||||
run_image_description_jpg_filename_migration(pool).await?;
|
run_image_description_jpg_filename_migration(pool).await?;
|
||||||
|
run_weknora_outbox_migration(pool).await?;
|
||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -161,6 +162,46 @@ async fn run_schema_migrations(pool: &SqlitePool) -> anyhow::Result<()> {
|
|||||||
run_image_description_migration(pool).await?;
|
run_image_description_migration(pool).await?;
|
||||||
run_file_summary_migration(pool).await?;
|
run_file_summary_migration(pool).await?;
|
||||||
run_image_description_jpg_filename_migration(pool).await?;
|
run_image_description_jpg_filename_migration(pool).await?;
|
||||||
|
run_weknora_outbox_migration(pool).await?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn run_weknora_outbox_migration(pool: &SqlitePool) -> anyhow::Result<()> {
|
||||||
|
const VERSION: i64 = 6;
|
||||||
|
let mut tx = pool.begin().await?;
|
||||||
|
let claimed = sqlx::query("INSERT OR IGNORE INTO schema_migrations(version, name) VALUES (?, ?)")
|
||||||
|
.bind(VERSION)
|
||||||
|
.bind("weknora_sync_outbox")
|
||||||
|
.execute(&mut *tx)
|
||||||
|
.await?
|
||||||
|
.rows_affected();
|
||||||
|
if claimed == 0 {
|
||||||
|
tx.rollback().await?;
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
sqlx::query(
|
||||||
|
"CREATE TABLE IF NOT EXISTS weknora_sync_outbox (
|
||||||
|
event_id TEXT NOT NULL,
|
||||||
|
event_type TEXT NOT NULL CHECK(event_type IN ('upsert','delete')),
|
||||||
|
file_id INTEGER PRIMARY KEY,
|
||||||
|
htknow_kb_id INTEGER NOT NULL,
|
||||||
|
attempts INTEGER NOT NULL DEFAULT 0,
|
||||||
|
next_attempt_at INTEGER NOT NULL DEFAULT (strftime('%s','now')),
|
||||||
|
last_error TEXT NOT NULL DEFAULT '',
|
||||||
|
created_at INTEGER NOT NULL DEFAULT (strftime('%s','now')),
|
||||||
|
updated_at INTEGER NOT NULL DEFAULT (strftime('%s','now'))
|
||||||
|
)",
|
||||||
|
)
|
||||||
|
.execute(&mut *tx)
|
||||||
|
.await?;
|
||||||
|
sqlx::query(
|
||||||
|
"CREATE INDEX IF NOT EXISTS idx_weknora_sync_outbox_due
|
||||||
|
ON weknora_sync_outbox(next_attempt_at, created_at)",
|
||||||
|
)
|
||||||
|
.execute(&mut *tx)
|
||||||
|
.await?;
|
||||||
|
tx.commit().await?;
|
||||||
|
info!("Applied schema migration {}: weknora_sync_outbox", VERSION);
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@ -1,5 +1,9 @@
|
|||||||
use axum::{
|
use axum::{
|
||||||
Router, body::Body, http::{StatusCode, Uri, header}, response::{IntoResponse, Response}, routing::get
|
Router,
|
||||||
|
body::Body,
|
||||||
|
http::{StatusCode, Uri, header},
|
||||||
|
response::{IntoResponse, Response},
|
||||||
|
routing::get,
|
||||||
};
|
};
|
||||||
use rust_embed::Embed;
|
use rust_embed::Embed;
|
||||||
|
|
||||||
|
|||||||
@ -23,7 +23,8 @@ pub struct ImageDescription {
|
|||||||
|
|
||||||
/// 保存或更新单条图片描述
|
/// 保存或更新单条图片描述
|
||||||
pub async fn save(
|
pub async fn save(
|
||||||
pool: &SqlitePool, file_id: i64, image_filename: &str, description: &str, raw_response: &str, source: &str, jpg_filename: Option<&str>,
|
pool: &SqlitePool, file_id: i64, image_filename: &str, description: &str, raw_response: &str, source: &str,
|
||||||
|
jpg_filename: Option<&str>,
|
||||||
) -> Result<()> {
|
) -> Result<()> {
|
||||||
sqlx::query(
|
sqlx::query(
|
||||||
"INSERT INTO image_descriptions (file_id, image_filename, description, raw_response, source, jpg_filename)
|
"INSERT INTO image_descriptions (file_id, image_filename, description, raw_response, source, jpg_filename)
|
||||||
@ -68,7 +69,9 @@ pub async fn delete(pool: &SqlitePool, file_id: i64, image_filename: &str) -> Re
|
|||||||
.bind(image_filename)
|
.bind(image_filename)
|
||||||
.execute(pool)
|
.execute(pool)
|
||||||
.await
|
.await
|
||||||
.with_context(|| format!("failed to delete image_description for file_id={} filename={}", file_id, image_filename))?;
|
.with_context(|| {
|
||||||
|
format!("failed to delete image_description for file_id={} filename={}", file_id, image_filename)
|
||||||
|
})?;
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -94,13 +97,16 @@ struct FileImageDescriptionRow {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// 批量查询多个文件的图片描述,返回 (file_id, image_filename) -> description 映射。
|
/// 批量查询多个文件的图片描述,返回 (file_id, image_filename) -> description 映射。
|
||||||
pub async fn list_by_file_ids(pool: &SqlitePool, file_ids: &[i64]) -> Result<HashMap<(i64, String), (String, Option<String>)>> {
|
pub async fn list_by_file_ids(
|
||||||
|
pool: &SqlitePool, file_ids: &[i64],
|
||||||
|
) -> Result<HashMap<(i64, String), (String, Option<String>)>> {
|
||||||
if file_ids.is_empty() {
|
if file_ids.is_empty() {
|
||||||
return Ok(HashMap::new());
|
return Ok(HashMap::new());
|
||||||
}
|
}
|
||||||
|
|
||||||
let mut query_builder: QueryBuilder<'_, Sqlite> =
|
let mut query_builder: QueryBuilder<'_, Sqlite> = QueryBuilder::new(
|
||||||
QueryBuilder::new("SELECT file_id, image_filename, description, jpg_filename FROM image_descriptions WHERE file_id IN (");
|
"SELECT file_id, image_filename, description, jpg_filename FROM image_descriptions WHERE file_id IN (",
|
||||||
|
);
|
||||||
let mut separated = query_builder.separated(", ");
|
let mut separated = query_builder.separated(", ");
|
||||||
for file_id in file_ids {
|
for file_id in file_ids {
|
||||||
separated.push_bind(file_id);
|
separated.push_bind(file_id);
|
||||||
@ -123,7 +129,16 @@ pub async fn copy_for_file(pool: &SqlitePool, src_file_id: i64, dst_file_id: i64
|
|||||||
.with_context(|| format!("failed to copy image_descriptions from file_id={}", src_file_id))?;
|
.with_context(|| format!("failed to copy image_descriptions from file_id={}", src_file_id))?;
|
||||||
|
|
||||||
for row in rows {
|
for row in rows {
|
||||||
save(pool, dst_file_id, &row.image_filename, &row.description, &row.raw_response, &row.source, row.jpg_filename.as_deref()).await?;
|
save(
|
||||||
|
pool,
|
||||||
|
dst_file_id,
|
||||||
|
&row.image_filename,
|
||||||
|
&row.description,
|
||||||
|
&row.raw_response,
|
||||||
|
&row.source,
|
||||||
|
row.jpg_filename.as_deref(),
|
||||||
|
)
|
||||||
|
.await?;
|
||||||
}
|
}
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|||||||
@ -1,5 +1,8 @@
|
|||||||
use axum::{
|
use axum::{
|
||||||
body::Body, http::{Request, StatusCode}, middleware::Next, response::{IntoResponse, Response}
|
body::Body,
|
||||||
|
http::{Request, StatusCode},
|
||||||
|
middleware::Next,
|
||||||
|
response::{IntoResponse, Response},
|
||||||
};
|
};
|
||||||
use base64::{Engine, engine::general_purpose::STANDARD};
|
use base64::{Engine, engine::general_purpose::STANDARD};
|
||||||
|
|
||||||
@ -19,6 +22,7 @@ pub mod pdf_highlight;
|
|||||||
pub mod processor;
|
pub mod processor;
|
||||||
pub mod search;
|
pub mod search;
|
||||||
pub mod slice_content;
|
pub mod slice_content;
|
||||||
|
pub mod weknora_sync;
|
||||||
|
|
||||||
/// User authentication info extracted from request headers
|
/// User authentication info extracted from request headers
|
||||||
#[derive(Clone, Debug)]
|
#[derive(Clone, Debug)]
|
||||||
|
|||||||
@ -1,7 +1,9 @@
|
|||||||
use std::{env, str::FromStr};
|
use std::{env, str::FromStr};
|
||||||
|
|
||||||
use log4rs::{
|
use log4rs::{
|
||||||
append::console::ConsoleAppender, config::{Appender, Config, Logger, Root}, encode::pattern::PatternEncoder
|
append::console::ConsoleAppender,
|
||||||
|
config::{Appender, Config, Logger, Root},
|
||||||
|
encode::pattern::PatternEncoder,
|
||||||
};
|
};
|
||||||
|
|
||||||
fn parse_rust_log(rust_log: &str) -> (Option<log::LevelFilter>, Vec<(String, log::LevelFilter)>) {
|
fn parse_rust_log(rust_log: &str) -> (Option<log::LevelFilter>, Vec<(String, log::LevelFilter)>) {
|
||||||
|
|||||||
@ -1,8 +1,8 @@
|
|||||||
#[cfg(debug_assertions)]
|
#[cfg(all(not(windows), debug_assertions))]
|
||||||
#[global_allocator]
|
#[global_allocator]
|
||||||
static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
|
static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
|
||||||
|
|
||||||
#[cfg(all(debug_assertions, feature = "profiling"))]
|
#[cfg(all(not(windows), debug_assertions, feature = "profiling"))]
|
||||||
#[unsafe(export_name = "_rjem_malloc_conf")]
|
#[unsafe(export_name = "_rjem_malloc_conf")]
|
||||||
// lg_prof_sample:0 表示每次分配都采样(最详细,但性能开销大)
|
// lg_prof_sample:0 表示每次分配都采样(最详细,但性能开销大)
|
||||||
// lg_prof_sample:10 表示每 1KB 采样一次(推荐)
|
// lg_prof_sample:10 表示每 1KB 采样一次(推荐)
|
||||||
@ -13,7 +13,7 @@ use std::{net::SocketAddr, sync::Arc};
|
|||||||
|
|
||||||
use axum::{Router, extract::DefaultBodyLimit, middleware, response::Html, routing::get};
|
use axum::{Router, extract::DefaultBodyLimit, middleware, response::Html, routing::get};
|
||||||
use chrono::Local;
|
use chrono::Local;
|
||||||
use htknow::{api, auth, config, db, frontend, log4rs, processor, search};
|
use htknow::{api, auth, config, db, frontend, log4rs, processor, search, weknora_sync};
|
||||||
use tokio::net::TcpListener;
|
use tokio::net::TcpListener;
|
||||||
use tokio_cron_scheduler::{Job, JobScheduler};
|
use tokio_cron_scheduler::{Job, JobScheduler};
|
||||||
use utoipa_swagger_ui::SwaggerUi;
|
use utoipa_swagger_ui::SwaggerUi;
|
||||||
@ -27,6 +27,7 @@ async fn main() -> anyhow::Result<()> {
|
|||||||
log::info!("Configuration loaded: server={}:{}", cfg.server.host, cfg.server.port);
|
log::info!("Configuration loaded: server={}:{}", cfg.server.host, cfg.server.port);
|
||||||
|
|
||||||
let pool = db::init().await?;
|
let pool = db::init().await?;
|
||||||
|
weknora_sync::start_worker(pool.clone());
|
||||||
log::info!("Initializing search engine...");
|
log::info!("Initializing search engine...");
|
||||||
let search_init_started_at = std::time::Instant::now();
|
let search_init_started_at = std::time::Instant::now();
|
||||||
let search_engine = search::SearchEngine::init().await.with_pool(pool.clone());
|
let search_engine = search::SearchEngine::init().await.with_pool(pool.clone());
|
||||||
|
|||||||
@ -5,7 +5,9 @@
|
|||||||
use anyhow::Result;
|
use anyhow::Result;
|
||||||
use log::info;
|
use log::info;
|
||||||
use lopdf::{
|
use lopdf::{
|
||||||
Document, Object, ObjectId, Stream, content::{Content, Operation}, dictionary
|
Document, Object, ObjectId, Stream,
|
||||||
|
content::{Content, Operation},
|
||||||
|
dictionary,
|
||||||
};
|
};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
|
|||||||
@ -179,8 +179,16 @@ async fn load_or_parse_image_descriptions(
|
|||||||
}
|
}
|
||||||
match image_parse::parse_image_file(path, name, None).await {
|
match image_parse::parse_image_file(path, name, None).await {
|
||||||
Ok(resp) => {
|
Ok(resp) => {
|
||||||
if let Err(err) =
|
if let Err(err) = image_description::save(
|
||||||
image_description::save(pool, file_id, name, &resp.description, &resp.raw_response, source, resp.jpg_filename.as_deref()).await
|
pool,
|
||||||
|
file_id,
|
||||||
|
name,
|
||||||
|
&resp.description,
|
||||||
|
&resp.raw_response,
|
||||||
|
source,
|
||||||
|
resp.jpg_filename.as_deref(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
{
|
{
|
||||||
warn!("Failed to save image description for {}: {}", name, err);
|
warn!("Failed to save image description for {}: {}", name, err);
|
||||||
}
|
}
|
||||||
@ -2046,12 +2054,7 @@ impl FileProcessor {
|
|||||||
if output.status.success() {
|
if output.status.success() {
|
||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
anyhow::bail!(
|
anyhow::bail!("{} exited with {}: {}", program, output.status, String::from_utf8_lossy(&output.stderr).trim());
|
||||||
"{} exited with {}: {}",
|
|
||||||
program,
|
|
||||||
output.status,
|
|
||||||
String::from_utf8_lossy(&output.stderr).trim()
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn process_uploaded_image_file(
|
async fn process_uploaded_image_file(
|
||||||
@ -2165,7 +2168,10 @@ impl FileProcessor {
|
|||||||
if is_image {
|
if is_image {
|
||||||
let removed = retain_one_image_block_for_uploaded_image(&mut content_list);
|
let removed = retain_one_image_block_for_uploaded_image(&mut content_list);
|
||||||
if removed > 0 {
|
if removed > 0 {
|
||||||
warn!("Keeping one MinerU image block for uploaded image {}, dropped {} duplicate blocks", file.filename, removed);
|
warn!(
|
||||||
|
"Keeping one MinerU image block for uploaded image {}, dropped {} duplicate blocks",
|
||||||
|
file.filename, removed
|
||||||
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
@ -2179,7 +2185,10 @@ impl FileProcessor {
|
|||||||
if is_image {
|
if is_image {
|
||||||
let removed = retain_one_image_block_for_uploaded_image(&mut content_list);
|
let removed = retain_one_image_block_for_uploaded_image(&mut content_list);
|
||||||
if removed > 0 {
|
if removed > 0 {
|
||||||
warn!("Keeping one MinerU image block for uploaded image {}, dropped {} duplicate blocks", file.filename, removed);
|
warn!(
|
||||||
|
"Keeping one MinerU image block for uploaded image {}, dropped {} duplicate blocks",
|
||||||
|
file.filename, removed
|
||||||
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if !self.ensure_file_exists(file.id, "before writing pdf contents").await? {
|
if !self.ensure_file_exists(file.id, "before writing pdf contents").await? {
|
||||||
@ -2295,11 +2304,7 @@ impl FileProcessor {
|
|||||||
text_level: None,
|
text_level: None,
|
||||||
text_format: None,
|
text_format: None,
|
||||||
img_path: Some(file.filename.clone()),
|
img_path: Some(file.filename.clone()),
|
||||||
image_caption: if caption.is_empty() {
|
image_caption: if caption.is_empty() { None } else { Some(vec![caption]) },
|
||||||
None
|
|
||||||
} else {
|
|
||||||
Some(vec![caption])
|
|
||||||
},
|
|
||||||
table_body: None,
|
table_body: None,
|
||||||
table_caption: None,
|
table_caption: None,
|
||||||
});
|
});
|
||||||
@ -2947,6 +2952,10 @@ impl FileProcessor {
|
|||||||
})
|
})
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
|
if let Err(err) = crate::weknora_sync::enqueue_file_upsert(&self.pool, file.id).await {
|
||||||
|
warn!("Failed to enqueue WeKnora sync for file {}: {}", file.id, err);
|
||||||
|
}
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -3566,6 +3575,9 @@ impl FileProcessor {
|
|||||||
updated_file.content = Some(indexed_content);
|
updated_file.content = Some(indexed_content);
|
||||||
updated_file.summary = summary.clone();
|
updated_file.summary = summary.clone();
|
||||||
self.maybe_build_knowledge_graph(&updated_file).await;
|
self.maybe_build_knowledge_graph(&updated_file).await;
|
||||||
|
if let Err(err) = crate::weknora_sync::enqueue_file_upsert(&self.pool, target.id).await {
|
||||||
|
warn!("Failed to enqueue WeKnora sync for file {}: {}", target.id, err);
|
||||||
|
}
|
||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -3640,6 +3652,9 @@ impl FileProcessor {
|
|||||||
updated_file.content = Some(full_content.clone());
|
updated_file.content = Some(full_content.clone());
|
||||||
updated_file.summary = source.summary.clone();
|
updated_file.summary = source.summary.clone();
|
||||||
self.maybe_build_knowledge_graph(&updated_file).await;
|
self.maybe_build_knowledge_graph(&updated_file).await;
|
||||||
|
if let Err(err) = crate::weknora_sync::enqueue_file_upsert(&self.pool, target.id).await {
|
||||||
|
warn!("Failed to enqueue WeKnora sync for file {}: {}", target.id, err);
|
||||||
|
}
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|||||||
@ -29,8 +29,48 @@ pub use tantivy_engine::{FullSearchResultItem, SearchResultItem};
|
|||||||
const FULL_SNIPPET_MAX_CHARS: usize = 400;
|
const FULL_SNIPPET_MAX_CHARS: usize = 400;
|
||||||
const MAX_QUERY_TERMS_FOR_SYNONYM_LOOKUP: usize = 100;
|
const MAX_QUERY_TERMS_FOR_SYNONYM_LOOKUP: usize = 100;
|
||||||
const DEFAULT_REBUILD_BATCH_SIZE: i64 = 100;
|
const DEFAULT_REBUILD_BATCH_SIZE: i64 = 100;
|
||||||
|
/// Reciprocal Rank Fusion constant. 60 is the commonly used default.
|
||||||
|
const RRF_K: f32 = 60.0;
|
||||||
static RERANK_HTTP_CLIENT: Lazy<Client> = Lazy::new(Client::new);
|
static RERANK_HTTP_CLIENT: Lazy<Client> = Lazy::new(Client::new);
|
||||||
|
|
||||||
|
/// Fuse ranked recall lists without comparing incompatible raw score scales.
|
||||||
|
fn reciprocal_rank_fusion(result_lists: impl IntoIterator<Item = Vec<SearchResultItem>>) -> Vec<SearchResultItem> {
|
||||||
|
let mut fused: HashMap<i64, (SearchResultItem, f32, usize)> = HashMap::new();
|
||||||
|
|
||||||
|
for results in result_lists {
|
||||||
|
for (index, mut result) in results.into_iter().filter(|result| !result.content.trim().is_empty()).enumerate() {
|
||||||
|
let rank = index + 1;
|
||||||
|
let contribution = 1.0 / (RRF_K + rank as f32);
|
||||||
|
fused
|
||||||
|
.entry(result.id)
|
||||||
|
.and_modify(|(_, score, best_rank)| {
|
||||||
|
*score += contribution;
|
||||||
|
*best_rank = (*best_rank).min(rank);
|
||||||
|
})
|
||||||
|
.or_insert_with(|| {
|
||||||
|
result.score = contribution;
|
||||||
|
(result, contribution, rank)
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut results: Vec<(SearchResultItem, usize)> = fused
|
||||||
|
.into_values()
|
||||||
|
.map(|(mut result, score, best_rank)| {
|
||||||
|
result.score = score;
|
||||||
|
(result, best_rank)
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
results.sort_by(|(a, a_rank), (b, b_rank)| {
|
||||||
|
b.score
|
||||||
|
.partial_cmp(&a.score)
|
||||||
|
.unwrap_or(std::cmp::Ordering::Equal)
|
||||||
|
.then_with(|| a_rank.cmp(b_rank))
|
||||||
|
.then_with(|| a.id.cmp(&b.id))
|
||||||
|
});
|
||||||
|
results.into_iter().map(|(result, _)| result).collect()
|
||||||
|
}
|
||||||
|
|
||||||
/// /v1/rerank 格式的请求体
|
/// /v1/rerank 格式的请求体
|
||||||
#[derive(Debug, Serialize)]
|
#[derive(Debug, Serialize)]
|
||||||
struct RerankRequest {
|
struct RerankRequest {
|
||||||
@ -1207,37 +1247,8 @@ impl SearchEngine {
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
// 合并结果:使用 HashMap 按 id 去重,保留最高分数,同时去除内容为空的结果
|
// 两个召回分支的原始分数量纲不同,按排名进行 RRF 融合。
|
||||||
let mut merged_map: HashMap<i64, SearchResultItem> = HashMap::new();
|
let merged_results = reciprocal_rank_fusion([tantivy_results, lancedb_results]);
|
||||||
|
|
||||||
for result in tantivy_results {
|
|
||||||
// 跳过空内容(包括仅有空白的情况)
|
|
||||||
if result.content.trim().is_empty() {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
merged_map.insert(result.id, result);
|
|
||||||
}
|
|
||||||
|
|
||||||
for result in lancedb_results {
|
|
||||||
// 跳过空内容(包括仅有空白的情况)
|
|
||||||
if result.content.trim().is_empty() {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
merged_map
|
|
||||||
.entry(result.id)
|
|
||||||
.and_modify(|e| {
|
|
||||||
// 如果已存在,取两者中分数较高的
|
|
||||||
if result.score > e.score {
|
|
||||||
*e = result.clone();
|
|
||||||
}
|
|
||||||
})
|
|
||||||
.or_insert(result);
|
|
||||||
}
|
|
||||||
|
|
||||||
// 转换为 Vec 并按分数降序排序
|
|
||||||
let mut merged_results: Vec<SearchResultItem> = merged_map.into_values().collect();
|
|
||||||
merged_results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal));
|
|
||||||
|
|
||||||
info!("Merged results count: {}", merged_results.len());
|
info!("Merged results count: {}", merged_results.len());
|
||||||
|
|
||||||
@ -1366,29 +1377,7 @@ impl SearchEngine {
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
let mut merged_map: HashMap<i64, SearchResultItem> = HashMap::new();
|
let merged_results = reciprocal_rank_fusion([tantivy_results, lancedb_results]);
|
||||||
for result in tantivy_results {
|
|
||||||
if result.content.trim().is_empty() {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
merged_map.insert(result.id, result);
|
|
||||||
}
|
|
||||||
for result in lancedb_results {
|
|
||||||
if result.content.trim().is_empty() {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
merged_map
|
|
||||||
.entry(result.id)
|
|
||||||
.and_modify(|e| {
|
|
||||||
if result.score > e.score {
|
|
||||||
*e = result.clone();
|
|
||||||
}
|
|
||||||
})
|
|
||||||
.or_insert(result);
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut merged_results: Vec<SearchResultItem> = merged_map.into_values().collect();
|
|
||||||
merged_results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal));
|
|
||||||
info!("Image-by-text merged results count: {}", merged_results.len());
|
info!("Image-by-text merged results count: {}", merged_results.len());
|
||||||
|
|
||||||
if merged_results.is_empty() {
|
if merged_results.is_empty() {
|
||||||
@ -1967,3 +1956,33 @@ fn reload_reader(reader: &IndexReader, label: &str) -> tantivy::Result<()> {
|
|||||||
debug!("Tantivy {} reader reload {}ms", label, start.elapsed().as_millis());
|
debug!("Tantivy {} reader reload {}ms", label, start.elapsed().as_millis());
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
fn result(id: i64, content: &str, score: f32) -> SearchResultItem {
|
||||||
|
SearchResultItem { id, file_id: id, kb_id: Some(1), content: content.to_string(), score }
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn rrf_rewards_results_recalled_by_both_branches() {
|
||||||
|
let text = vec![result(1, "text-first", 100.0), result(2, "shared", 1.0)];
|
||||||
|
let vector = vec![result(2, "shared", 0.1), result(3, "vector-second", 0.09)];
|
||||||
|
|
||||||
|
let fused = reciprocal_rank_fusion([text, vector]);
|
||||||
|
|
||||||
|
assert_eq!(fused.iter().map(|item| item.id).collect::<Vec<_>>(), vec![2, 1, 3]);
|
||||||
|
assert!((fused[0].score - (1.0 / 62.0 + 1.0 / 61.0)).abs() < f32::EPSILON);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn rrf_uses_rank_instead_of_raw_score_and_drops_empty_content() {
|
||||||
|
let text = vec![result(10, "first", 0.01), result(11, "second", 10_000.0), result(12, " ", 20_000.0)];
|
||||||
|
|
||||||
|
let fused = reciprocal_rank_fusion([text, Vec::new()]);
|
||||||
|
|
||||||
|
assert_eq!(fused.iter().map(|item| item.id).collect::<Vec<_>>(), vec![10, 11]);
|
||||||
|
assert!(fused[0].score > fused[1].score);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@ -1,14 +1,25 @@
|
|||||||
use std::{
|
use std::{
|
||||||
collections::{HashMap, HashSet}, path::Path, sync::{
|
collections::{HashMap, HashSet},
|
||||||
Arc, mpsc::{self, Sender}
|
path::Path,
|
||||||
}, thread, time::{Duration, Instant, SystemTime, UNIX_EPOCH}
|
sync::{
|
||||||
|
Arc,
|
||||||
|
mpsc::{self, Sender},
|
||||||
|
},
|
||||||
|
thread,
|
||||||
|
time::{Duration, Instant, SystemTime, UNIX_EPOCH},
|
||||||
};
|
};
|
||||||
|
|
||||||
use aho_corasick::{AhoCorasick, AhoCorasickBuilder, MatchKind};
|
use aho_corasick::{AhoCorasick, AhoCorasickBuilder, MatchKind};
|
||||||
use log::{debug, info, warn};
|
use log::{debug, info, warn};
|
||||||
use serde::Serialize;
|
use serde::Serialize;
|
||||||
use tantivy::{
|
use tantivy::{
|
||||||
Index, IndexReader, Result, TantivyDocument, TantivyError, Term, collector::TopDocs, directory::error::LockError, doc, merge_policy::LogMergePolicy, query::{BooleanQuery, BoostQuery, Occur, Query, TermQuery}, schema::{FAST, Field, INDEXED, IndexRecordOption, STORED, Schema, TextFieldIndexing, TextOptions, Value as _}
|
Index, IndexReader, Result, TantivyDocument, TantivyError, Term,
|
||||||
|
collector::TopDocs,
|
||||||
|
directory::error::LockError,
|
||||||
|
doc,
|
||||||
|
merge_policy::LogMergePolicy,
|
||||||
|
query::{BooleanQuery, BoostQuery, Occur, Query, TermQuery},
|
||||||
|
schema::{FAST, Field, INDEXED, IndexRecordOption, STORED, Schema, TextFieldIndexing, TextOptions, Value as _},
|
||||||
};
|
};
|
||||||
|
|
||||||
use super::chinese_tokenizer;
|
use super::chinese_tokenizer;
|
||||||
@ -451,7 +462,7 @@ pub async fn create_rebuild_writer(index: &Index, label: &str) -> tantivy::Resul
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn add_documents(
|
pub fn add_documents(
|
||||||
index_writer: &mut tantivy::IndexWriter, schema: &Schema, docs: impl IntoIterator<Item=Document>,
|
index_writer: &mut tantivy::IndexWriter, schema: &Schema, docs: impl IntoIterator<Item = Document>,
|
||||||
) -> tantivy::Result<usize> {
|
) -> tantivy::Result<usize> {
|
||||||
let mut count = 0_usize;
|
let mut count = 0_usize;
|
||||||
for doc in docs {
|
for doc in docs {
|
||||||
|
|||||||
265
src/weknora_sync.rs
Normal file
265
src/weknora_sync.rs
Normal file
@ -0,0 +1,265 @@
|
|||||||
|
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<Client> =
|
||||||
|
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<ExternalChunk>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[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<SyncEvent> = 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<String> {
|
||||||
|
let kb_map_raw = std::env::var("HTKNOW_WEKNORA_KB_MAP").context("HTKNOW_WEKNORA_KB_MAP is required")?;
|
||||||
|
let kb_map: HashMap<String, String> =
|
||||||
|
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<i64> = 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::<Vec<_>>();
|
||||||
|
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(())
|
||||||
|
}
|
||||||
@ -1,6 +1,7 @@
|
|||||||
mod common;
|
mod common;
|
||||||
use std::{
|
use std::{
|
||||||
fs, time::{Duration, Instant}
|
fs,
|
||||||
|
time::{Duration, Instant},
|
||||||
};
|
};
|
||||||
|
|
||||||
use axum::http::StatusCode;
|
use axum::http::StatusCode;
|
||||||
|
|||||||
@ -1,14 +1,19 @@
|
|||||||
#![allow(dead_code, unused_imports)]
|
#![allow(dead_code, unused_imports)]
|
||||||
|
|
||||||
use std::{
|
use std::{
|
||||||
fs, path::PathBuf, sync::{
|
fs,
|
||||||
OnceLock, atomic::{AtomicUsize, Ordering}
|
path::PathBuf,
|
||||||
}, time::{SystemTime, UNIX_EPOCH}
|
sync::{
|
||||||
|
OnceLock,
|
||||||
|
atomic::{AtomicUsize, Ordering},
|
||||||
|
},
|
||||||
|
time::{SystemTime, UNIX_EPOCH},
|
||||||
};
|
};
|
||||||
|
|
||||||
use axum::{Router, extract::DefaultBodyLimit, middleware};
|
use axum::{Router, extract::DefaultBodyLimit, middleware};
|
||||||
pub use axum::{
|
pub use axum::{
|
||||||
body::Body, http::{Request, Response, StatusCode, header}
|
body::Body,
|
||||||
|
http::{Request, Response, StatusCode, header},
|
||||||
};
|
};
|
||||||
use htknow::{api, auth, db, search::SearchEngine, slice_content};
|
use htknow::{api, auth, db, search::SearchEngine, slice_content};
|
||||||
pub use http_body_util::BodyExt;
|
pub use http_body_util::BodyExt;
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user