Compare commits
2 Commits
main
...
master-wik
| Author | SHA1 | Date | |
|---|---|---|---|
| 611d999365 | |||
| 1ce10dc1f2 |
@ -73,8 +73,9 @@ encoding_rs = "0.8"
|
||||
etcd-client = { version = "0.14", optional = true }
|
||||
serde_with = "3.16.1"
|
||||
sysinfo = "0.29"
|
||||
# jemalloc allocator
|
||||
# Note: profiling is only needed for debug/analysis, disabled in release for better performance
|
||||
# jemalloc is only enabled on non-Windows targets. Windows uses the system
|
||||
# allocator because tikv-jemalloc-sys does not reliably build for MSVC.
|
||||
[target.'cfg(not(windows))'.dependencies]
|
||||
tikv-jemallocator = { version = "0.6" }
|
||||
tikv-jemalloc-ctl = { version = "0.6", optional = true, features = ["use_std"] }
|
||||
|
||||
|
||||
@ -87,7 +87,11 @@ docker run -d --name mineru-api --restart unless-stopped --ipc host -p 10001:100
|
||||
| `HTKNOW_EMBEDDING_URL` | `http://222.190.139.186:59700/v1/embeddings` | 文本向量服务 |
|
||||
| `HTKNOW_IMAGE_EMBEDDING_URL` | `http://192.168.0.46:59802/v1/embeddings/file` | 图片向量服务 |
|
||||
| `HTKNOW_RERANK_URL` | `http://222.190.139.186:59600/v1/rerank` | Rerank 服务 |
|
||||
|
||||
| `HTKNOW_CUSTOM_IMAGE_URL` | `http://192.168.0.46:59085/split_result` | 自定义图片文本化 解析 |
|
||||
| `HTKNOW_CUSTOM_IMAGE_URL` | `http://192.168.0.46:59085/split_result` | 自定义图片文本化 解析 |
|
||||
| `HTKNOW_WEKNORA_URL` | `http://172.19.152.51:8081` | weknora 后端服务 |
|
||||
| `HTKNOW_WEKNORA_API_KEY` | `sk-6U04CeFoY7JSEBKBp7fQSGYoW3kv7e5qdB14aFnEPt` | weknora api接入服务 |
|
||||
| `HTKNOW_WEKNORA_KB_MAP` | `{"*":"881ee2e9-633a-426e-b26a-b9175eec9a64"}` | weknora 知识库id和htknow映射接入服务 |
|
||||
### AI 模型
|
||||
| 环境变量 | 默认值 | 说明 |
|
||||
| --- | --- | --- |
|
||||
|
||||
@ -1,9 +1,11 @@
|
||||
use sqlx::{QueryBuilder, Sqlite, SqlitePool};
|
||||
|
||||
use crate::{
|
||||
AuthUser, api::{
|
||||
error::{ApiError, ApiResult}, knowledge_base
|
||||
}
|
||||
AuthUser,
|
||||
api::{
|
||||
error::{ApiError, ApiResult},
|
||||
knowledge_base,
|
||||
},
|
||||
};
|
||||
|
||||
/// 要求当前用户为 admin,否则返回 BadRequest。
|
||||
|
||||
@ -1,7 +1,9 @@
|
||||
use std::{fmt, num::ParseIntError, result};
|
||||
|
||||
use axum::{
|
||||
Json, http::StatusCode, response::{IntoResponse, Response}
|
||||
Json,
|
||||
http::StatusCode,
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use serde::Serialize;
|
||||
use sqlx;
|
||||
|
||||
@ -456,7 +456,7 @@ pub async fn upload(
|
||||
|
||||
match field_name.as_deref() {
|
||||
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);
|
||||
let tempname = uuid::Uuid::new_v4().to_string();
|
||||
let filepath = format!("{}/{}", dir, tempname);
|
||||
@ -755,6 +755,54 @@ pub async fn upload(
|
||||
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(
|
||||
get,
|
||||
@ -1676,6 +1724,16 @@ async fn execute_batch_delete(
|
||||
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!(
|
||||
"file_batch_delete total requested={} accepted={} deleted={} {}ms",
|
||||
requested,
|
||||
@ -2472,6 +2530,13 @@ pub async fn update_slices(
|
||||
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. 返回更新后的切片列表
|
||||
get_slices(State(pool), Path(id)).await
|
||||
}
|
||||
|
||||
@ -1,7 +1,9 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use axum::{
|
||||
Extension, extract::{Path, Query, State}, response::Json
|
||||
Extension,
|
||||
extract::{Path, Query, State},
|
||||
response::Json,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sqlx::SqlitePool;
|
||||
|
||||
@ -5,9 +5,12 @@ use sqlx::SqlitePool;
|
||||
use utoipa::ToSchema;
|
||||
|
||||
use crate::{
|
||||
AuthUser, api::{
|
||||
common, error::{ApiError, ApiResult}
|
||||
}, search::{SearchEngine, tantivy_engine::ForceMergeStats}
|
||||
AuthUser,
|
||||
api::{
|
||||
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> {
|
||||
use std::{
|
||||
ffi::CString, io::Seek, time::{SystemTime, UNIX_EPOCH}
|
||||
ffi::CString,
|
||||
io::Seek,
|
||||
time::{SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
|
||||
use axum::{
|
||||
body::Body, http::{HeaderValue, header}
|
||||
body::Body,
|
||||
http::{HeaderValue, header},
|
||||
};
|
||||
use tempfile::NamedTempFile;
|
||||
use tikv_jemalloc_ctl::raw;
|
||||
@ -407,19 +413,22 @@ async fn heap_profile_impl() -> ApiResult<Response> {
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "profiling"))]
|
||||
#[cfg(any(not(feature = "profiling"), windows))]
|
||||
async fn heap_profile_impl() -> ApiResult<Response> {
|
||||
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> {
|
||||
use std::{
|
||||
ffi::CString, process::Stdio, time::{SystemTime, UNIX_EPOCH}
|
||||
ffi::CString,
|
||||
process::Stdio,
|
||||
time::{SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
|
||||
use axum::{
|
||||
body::Body, http::{HeaderValue, header}
|
||||
body::Body,
|
||||
http::{HeaderValue, header},
|
||||
};
|
||||
use tempfile::NamedTempFile;
|
||||
use tikv_jemalloc_ctl::raw;
|
||||
@ -512,7 +521,7 @@ async fn heap_profile_pdf_impl() -> ApiResult<Response> {
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "profiling"))]
|
||||
#[cfg(any(not(feature = "profiling"), windows))]
|
||||
async fn heap_profile_pdf_impl() -> ApiResult<Response> {
|
||||
Err(ApiError::BadRequest("Heap profiling is only available with 'profiling' feature enabled.".to_string()))
|
||||
}
|
||||
|
||||
@ -4,7 +4,8 @@
|
||||
//! 7Z 和 RAR 格式暂不支持。
|
||||
|
||||
use std::{
|
||||
io::{Read, Write}, path::Path
|
||||
io::{Read, Write},
|
||||
path::Path,
|
||||
};
|
||||
|
||||
use encoding_rs::GB18030;
|
||||
|
||||
@ -380,7 +380,7 @@ fn env_or_parse<T: std::str::FromStr>(key: &str, default: T) -> T {
|
||||
|
||||
/// 配置加载器 trait
|
||||
#[allow(dead_code)]
|
||||
pub trait ConfigLoader: Send+Sync {
|
||||
pub trait ConfigLoader: Send + Sync {
|
||||
/// 加载配置
|
||||
fn load(&self) -> anyhow::Result<AppConfig>;
|
||||
}
|
||||
@ -413,7 +413,8 @@ impl EtcdConfigLoader {
|
||||
/// 启动 watch 监听配置变化
|
||||
pub async fn watch<F>(&self, _callback: F) -> anyhow::Result<()>
|
||||
where
|
||||
F: Fn(AppConfig)+Send+Sync+'static, {
|
||||
F: Fn(AppConfig) + Send + Sync + 'static,
|
||||
{
|
||||
// TODO: 实现 etcd watch 逻辑
|
||||
// 1. 连接 etcd
|
||||
// 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_file_summary_migration(pool).await?;
|
||||
run_image_description_jpg_filename_migration(pool).await?;
|
||||
run_weknora_outbox_migration(pool).await?;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
@ -161,6 +162,46 @@ async fn run_schema_migrations(pool: &SqlitePool) -> anyhow::Result<()> {
|
||||
run_image_description_migration(pool).await?;
|
||||
run_file_summary_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(())
|
||||
}
|
||||
|
||||
|
||||
@ -1,5 +1,9 @@
|
||||
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;
|
||||
|
||||
|
||||
@ -23,7 +23,8 @@ pub struct ImageDescription {
|
||||
|
||||
/// 保存或更新单条图片描述
|
||||
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<()> {
|
||||
sqlx::query(
|
||||
"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)
|
||||
.execute(pool)
|
||||
.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(())
|
||||
}
|
||||
|
||||
@ -94,13 +97,16 @@ struct FileImageDescriptionRow {
|
||||
}
|
||||
|
||||
/// 批量查询多个文件的图片描述,返回 (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() {
|
||||
return Ok(HashMap::new());
|
||||
}
|
||||
|
||||
let mut query_builder: QueryBuilder<'_, Sqlite> =
|
||||
QueryBuilder::new("SELECT file_id, image_filename, description, jpg_filename FROM image_descriptions WHERE file_id IN (");
|
||||
let mut query_builder: QueryBuilder<'_, Sqlite> = QueryBuilder::new(
|
||||
"SELECT file_id, image_filename, description, jpg_filename FROM image_descriptions WHERE file_id IN (",
|
||||
);
|
||||
let mut separated = query_builder.separated(", ");
|
||||
for file_id in file_ids {
|
||||
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))?;
|
||||
|
||||
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(())
|
||||
}
|
||||
|
||||
@ -1,5 +1,8 @@
|
||||
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};
|
||||
|
||||
@ -19,6 +22,7 @@ pub mod pdf_highlight;
|
||||
pub mod processor;
|
||||
pub mod search;
|
||||
pub mod slice_content;
|
||||
pub mod weknora_sync;
|
||||
|
||||
/// User authentication info extracted from request headers
|
||||
#[derive(Clone, Debug)]
|
||||
|
||||
@ -1,7 +1,9 @@
|
||||
use std::{env, str::FromStr};
|
||||
|
||||
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)>) {
|
||||
|
||||
@ -1,8 +1,8 @@
|
||||
#[cfg(debug_assertions)]
|
||||
#[cfg(all(not(windows), debug_assertions))]
|
||||
#[global_allocator]
|
||||
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")]
|
||||
// lg_prof_sample:0 表示每次分配都采样(最详细,但性能开销大)
|
||||
// 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 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_cron_scheduler::{Job, JobScheduler};
|
||||
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);
|
||||
|
||||
let pool = db::init().await?;
|
||||
weknora_sync::start_worker(pool.clone());
|
||||
log::info!("Initializing search engine...");
|
||||
let search_init_started_at = std::time::Instant::now();
|
||||
let search_engine = search::SearchEngine::init().await.with_pool(pool.clone());
|
||||
|
||||
@ -5,7 +5,9 @@
|
||||
use anyhow::Result;
|
||||
use log::info;
|
||||
use lopdf::{
|
||||
Document, Object, ObjectId, Stream, content::{Content, Operation}, dictionary
|
||||
Document, Object, ObjectId, Stream,
|
||||
content::{Content, Operation},
|
||||
dictionary,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
|
||||
8281
src/processor.rs
8281
src/processor.rs
File diff suppressed because it is too large
Load Diff
@ -29,8 +29,48 @@ pub use tantivy_engine::{FullSearchResultItem, SearchResultItem};
|
||||
const FULL_SNIPPET_MAX_CHARS: usize = 400;
|
||||
const MAX_QUERY_TERMS_FOR_SYNONYM_LOOKUP: usize = 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);
|
||||
|
||||
/// 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 格式的请求体
|
||||
#[derive(Debug, Serialize)]
|
||||
struct RerankRequest {
|
||||
@ -1207,37 +1247,8 @@ impl SearchEngine {
|
||||
}
|
||||
};
|
||||
|
||||
// 合并结果:使用 HashMap 按 id 去重,保留最高分数,同时去除内容为空的结果
|
||||
let mut merged_map: HashMap<i64, SearchResultItem> = HashMap::new();
|
||||
|
||||
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));
|
||||
// 两个召回分支的原始分数量纲不同,按排名进行 RRF 融合。
|
||||
let merged_results = reciprocal_rank_fusion([tantivy_results, lancedb_results]);
|
||||
|
||||
info!("Merged results count: {}", merged_results.len());
|
||||
|
||||
@ -1366,29 +1377,7 @@ impl SearchEngine {
|
||||
}
|
||||
};
|
||||
|
||||
let mut merged_map: HashMap<i64, SearchResultItem> = HashMap::new();
|
||||
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));
|
||||
let merged_results = reciprocal_rank_fusion([tantivy_results, lancedb_results]);
|
||||
info!("Image-by-text merged results count: {}", merged_results.len());
|
||||
|
||||
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());
|
||||
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::{
|
||||
collections::{HashMap, HashSet}, path::Path, sync::{
|
||||
Arc, mpsc::{self, Sender}
|
||||
}, thread, time::{Duration, Instant, SystemTime, UNIX_EPOCH}
|
||||
collections::{HashMap, HashSet},
|
||||
path::Path,
|
||||
sync::{
|
||||
Arc,
|
||||
mpsc::{self, Sender},
|
||||
},
|
||||
thread,
|
||||
time::{Duration, Instant, SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
|
||||
use aho_corasick::{AhoCorasick, AhoCorasickBuilder, MatchKind};
|
||||
use log::{debug, info, warn};
|
||||
use serde::Serialize;
|
||||
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;
|
||||
@ -451,7 +462,7 @@ pub async fn create_rebuild_writer(index: &Index, label: &str) -> tantivy::Resul
|
||||
}
|
||||
|
||||
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> {
|
||||
let mut count = 0_usize;
|
||||
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;
|
||||
use std::{
|
||||
fs, time::{Duration, Instant}
|
||||
fs,
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
use axum::http::StatusCode;
|
||||
|
||||
@ -1,14 +1,19 @@
|
||||
#![allow(dead_code, unused_imports)]
|
||||
|
||||
use std::{
|
||||
fs, path::PathBuf, sync::{
|
||||
OnceLock, atomic::{AtomicUsize, Ordering}
|
||||
}, time::{SystemTime, UNIX_EPOCH}
|
||||
fs,
|
||||
path::PathBuf,
|
||||
sync::{
|
||||
OnceLock,
|
||||
atomic::{AtomicUsize, Ordering},
|
||||
},
|
||||
time::{SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
|
||||
use axum::{Router, extract::DefaultBodyLimit, middleware};
|
||||
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};
|
||||
pub use http_body_util::BodyExt;
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user