Compare commits
No commits in common. "master-wiki" and "main" have entirely different histories.
master-wik
...
main
@ -73,9 +73,8 @@ 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 is only enabled on non-Windows targets. Windows uses the system
|
# jemalloc allocator
|
||||||
# allocator because tikv-jemalloc-sys does not reliably build for MSVC.
|
# Note: profiling is only needed for debug/analysis, disabled in release for better performance
|
||||||
[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"] }
|
||||||
|
|
||||||
|
|||||||
@ -87,11 +87,7 @@ 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_EMBEDDING_URL` | `http://222.190.139.186:59700/v1/embeddings` | 文本向量服务 |
|
||||||
| `HTKNOW_IMAGE_EMBEDDING_URL` | `http://192.168.0.46:59802/v1/embeddings/file` | 图片向量服务 |
|
| `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_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 模型
|
### AI 模型
|
||||||
| 环境变量 | 默认值 | 说明 |
|
| 环境变量 | 默认值 | 说明 |
|
||||||
| --- | --- | --- |
|
| --- | --- | --- |
|
||||||
|
|||||||
@ -1,45 +1,61 @@
|
|||||||
|
version: "3.8"
|
||||||
|
|
||||||
services:
|
services:
|
||||||
htknow:
|
htknow:
|
||||||
image: gwknow:20260729-weknora-fix
|
build:
|
||||||
container_name: gwhtknow_ceshi
|
context: .
|
||||||
|
dockerfile: Dockerfile
|
||||||
|
container_name: htknow
|
||||||
ports:
|
ports:
|
||||||
- "12010:3000"
|
- "8080:3000"
|
||||||
volumes:
|
volumes:
|
||||||
|
# Mount data directory for persistent storage
|
||||||
- ./data:/app/data
|
- ./data:/app/data
|
||||||
environment:
|
environment:
|
||||||
RUST_LOG: "info"
|
- RUST_LOG=info
|
||||||
HTKNOW_SERVER_HOST: "0.0.0.0"
|
# 服务器配置
|
||||||
HTKNOW_SERVER_PORT: "3000"
|
- HTKNOW_SERVER_HOST=0.0.0.0
|
||||||
|
- HTKNOW_SERVER_PORT=3000
|
||||||
HTKNOW_MINERU_URL: "http://192.168.0.46:10001/file_parse"
|
- HTKNOW_SERVER_UPLOAD_LIMIT_MB=500
|
||||||
HTKNOW_REQUEST_TIMEOUT_SECS: "600"
|
- HTKNOW_SERVER_PROCESS_INTERVAL_SECS=10
|
||||||
HTKNOW_MINERU_MAX_PAGES: "0"
|
- HTKNOW_BUILD_KNOWLEDGE_GRAPH=false
|
||||||
HTKNOW_OFFICE_CONVERT_URL: "http://192.168.2.12:8003/convert"
|
- "HTKNOW_LANCEDB_COMPACT_CRON=0 0 3 * * *"
|
||||||
HTKNOW_EMBEDDING_URL: "http://192.168.2.12:59700/v1/embeddings"
|
# 外部服务地址
|
||||||
HTKNOW_RERANK_URL: "http://192.168.2.12:59600/v1/rerank"
|
- HTKNOW_MINERU_URL=http://192.168.0.46:10001/file_parse
|
||||||
HTKNOW_CUSTOM_PARSE_URL: "http://192.168.2.12:59085/split_result"
|
- HTKNOW_REQUEST_TIMEOUT_SECS=600
|
||||||
HTKNOW_IMAGE_PARSE_URL: "http://192.168.2.12:59085/split_image"
|
- HTKNOW_MINERU_MAX_PAGES=0
|
||||||
|
- HTKNOW_OFFICE_CONVERT_URL=http://192.168.0.46:8003/convert
|
||||||
HTKNOW_WEKNORA_URL: "http://192.168.2.12:59098"
|
- HTKNOW_EMBEDDING_URL=http://192.168.0.46:9700/v1/embeddings
|
||||||
HTKNOW_WEKNORA_API_KEY: "sk-R7YwleKiYxVlSyp_BPE6Giz5t8rCOQpRrcDBBZ6Bljc"
|
- HTKNOW_RERANK_URL=http://192.168.0.46:9600/v1/rerank
|
||||||
HTKNOW_WEKNORA_KB_MAP: '{"*":"f9241739-537d-4678-8e42-99ce65f45b25"}'
|
# AI 模型配置
|
||||||
|
- HTKNOW_EMBEDDING_MODEL=bge-m3
|
||||||
HTKNOW_EMBEDDING_MODEL: "bge-m3"
|
- HTKNOW_EMBEDDING_DIM=1024
|
||||||
HTKNOW_EMBEDDING_DIM: "1024"
|
- HTKNOW_EMBEDDING_BATCH_SIZE=8
|
||||||
HTKNOW_EMBEDDING_BATCH_SIZE: "8"
|
- HTKNOW_RERANK_MODEL=bge-rerank
|
||||||
HTKNOW_RERANK_MODEL: "bge-rerank"
|
- HTKNOW_RERANK_THRESHOLD=0.1
|
||||||
HTKNOW_RERANK_THRESHOLD: "0.1"
|
# 数据库配置
|
||||||
|
- DATABASE_URL=sqlite://data/app.sqlite
|
||||||
DATABASE_URL: "sqlite://data/app.sqlite"
|
- HTKNOW_DB_MAX_CONNECTIONS=10
|
||||||
HTKNOW_DB_MAX_CONNECTIONS: "10"
|
- HTKNOW_DB_BUSY_TIMEOUT_MS=5000
|
||||||
HTKNOW_DB_BUSY_TIMEOUT_MS: "5000"
|
# 存储路径配置
|
||||||
|
- HTKNOW_DATA_DIR=data
|
||||||
HTKNOW_DATA_DIR: "data"
|
- HTKNOW_LANCEDB_PATH=data/lancedb_data
|
||||||
HTKNOW_LANCEDB_PATH: "data/lancedb_data"
|
- HTKNOW_TEMP_PATH=data/temp
|
||||||
HTKNOW_TEMP_PATH: "data/temp"
|
- HTKNOW_IMAGES_PATH=data/images
|
||||||
HTKNOW_IMAGES_PATH: "data/images"
|
# 搜索配置
|
||||||
HTKNOW_SEARCH_LIMIT: "10"
|
- HTKNOW_SEARCH_LIMIT=10
|
||||||
HTKNOW_TANTIVY_INDEX_PATH: "data/tantivy_index"
|
- HTKNOW_TANTIVY_INDEX_PATH=data/tantivy_index
|
||||||
HTKNOW_TANTIVY_MEMORY_MB: "50"
|
- HTKNOW_TANTIVY_MEMORY_MB=50
|
||||||
|
# LLM 配置(可选,用于知识图谱构建)
|
||||||
|
# - LLM_API_URL=https://api.openai.com/v1/chat/completions
|
||||||
|
# - LLM_API_KEY=your-api-key
|
||||||
|
# - LLM_MODEL=gpt-3.5-turbo
|
||||||
restart: unless-stopped
|
restart: unless-stopped
|
||||||
|
# Uncomment if you need to connect to other services
|
||||||
|
# networks:
|
||||||
|
# - htknow-network
|
||||||
|
|
||||||
|
# Uncomment if you need a network
|
||||||
|
# networks:
|
||||||
|
# htknow-network:
|
||||||
|
# driver: bridge
|
||||||
|
|||||||
@ -1,11 +1,9 @@
|
|||||||
use sqlx::{QueryBuilder, Sqlite, SqlitePool};
|
use sqlx::{QueryBuilder, Sqlite, SqlitePool};
|
||||||
|
|
||||||
use crate::{
|
use crate::{
|
||||||
AuthUser,
|
AuthUser, api::{
|
||||||
api::{
|
error::{ApiError, ApiResult}, knowledge_base
|
||||||
error::{ApiError, ApiResult},
|
}
|
||||||
knowledge_base,
|
|
||||||
},
|
|
||||||
};
|
};
|
||||||
|
|
||||||
/// 要求当前用户为 admin,否则返回 BadRequest。
|
/// 要求当前用户为 admin,否则返回 BadRequest。
|
||||||
|
|||||||
@ -1,9 +1,7 @@
|
|||||||
use std::{fmt, num::ParseIntError, result};
|
use std::{fmt, num::ParseIntError, result};
|
||||||
|
|
||||||
use axum::{
|
use axum::{
|
||||||
Json,
|
Json, http::StatusCode, response::{IntoResponse, Response}
|
||||||
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 = normalize_upload_filename(field.file_name().unwrap_or("unknown"));
|
let filename = field.file_name().unwrap_or("unknown").to_string();
|
||||||
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,54 +755,6 @@ 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,
|
||||||
@ -1724,16 +1676,6 @@ 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,
|
||||||
@ -2530,13 +2472,6 @@ 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,9 +1,7 @@
|
|||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
|
|
||||||
use axum::{
|
use axum::{
|
||||||
Extension,
|
Extension, extract::{Path, Query, State}, response::Json
|
||||||
extract::{Path, Query, State},
|
|
||||||
response::Json,
|
|
||||||
};
|
};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use sqlx::SqlitePool;
|
use sqlx::SqlitePool;
|
||||||
|
|||||||
@ -5,12 +5,9 @@ use sqlx::SqlitePool;
|
|||||||
use utoipa::ToSchema;
|
use utoipa::ToSchema;
|
||||||
|
|
||||||
use crate::{
|
use crate::{
|
||||||
AuthUser,
|
AuthUser, api::{
|
||||||
api::{
|
common, error::{ApiError, ApiResult}
|
||||||
common,
|
}, search::{SearchEngine, tantivy_engine::ForceMergeStats}
|
||||||
error::{ApiError, ApiResult},
|
|
||||||
},
|
|
||||||
search::{SearchEngine, tantivy_engine::ForceMergeStats},
|
|
||||||
};
|
};
|
||||||
|
|
||||||
/// 进程内存占用
|
/// 进程内存占用
|
||||||
@ -353,17 +350,14 @@ fn force_merge_index_stats(index: &str, stats: ForceMergeStats) -> TantivyForceM
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(all(feature = "profiling", not(windows)))]
|
#[cfg(feature = "profiling")]
|
||||||
async fn heap_profile_impl() -> ApiResult<Response> {
|
async fn heap_profile_impl() -> ApiResult<Response> {
|
||||||
use std::{
|
use std::{
|
||||||
ffi::CString,
|
ffi::CString, io::Seek, time::{SystemTime, UNIX_EPOCH}
|
||||||
io::Seek,
|
|
||||||
time::{SystemTime, UNIX_EPOCH},
|
|
||||||
};
|
};
|
||||||
|
|
||||||
use axum::{
|
use axum::{
|
||||||
body::Body,
|
body::Body, http::{HeaderValue, header}
|
||||||
http::{HeaderValue, header},
|
|
||||||
};
|
};
|
||||||
use tempfile::NamedTempFile;
|
use tempfile::NamedTempFile;
|
||||||
use tikv_jemalloc_ctl::raw;
|
use tikv_jemalloc_ctl::raw;
|
||||||
@ -413,22 +407,19 @@ async fn heap_profile_impl() -> ApiResult<Response> {
|
|||||||
Ok(response)
|
Ok(response)
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(any(not(feature = "profiling"), windows))]
|
#[cfg(not(feature = "profiling"))]
|
||||||
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(all(feature = "profiling", not(windows)))]
|
#[cfg(feature = "profiling")]
|
||||||
async fn heap_profile_pdf_impl() -> ApiResult<Response> {
|
async fn heap_profile_pdf_impl() -> ApiResult<Response> {
|
||||||
use std::{
|
use std::{
|
||||||
ffi::CString,
|
ffi::CString, process::Stdio, time::{SystemTime, UNIX_EPOCH}
|
||||||
process::Stdio,
|
|
||||||
time::{SystemTime, UNIX_EPOCH},
|
|
||||||
};
|
};
|
||||||
|
|
||||||
use axum::{
|
use axum::{
|
||||||
body::Body,
|
body::Body, http::{HeaderValue, header}
|
||||||
http::{HeaderValue, header},
|
|
||||||
};
|
};
|
||||||
use tempfile::NamedTempFile;
|
use tempfile::NamedTempFile;
|
||||||
use tikv_jemalloc_ctl::raw;
|
use tikv_jemalloc_ctl::raw;
|
||||||
@ -521,7 +512,7 @@ async fn heap_profile_pdf_impl() -> ApiResult<Response> {
|
|||||||
Ok(response)
|
Ok(response)
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(any(not(feature = "profiling"), windows))]
|
#[cfg(not(feature = "profiling"))]
|
||||||
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,8 +4,7 @@
|
|||||||
//! 7Z 和 RAR 格式暂不支持。
|
//! 7Z 和 RAR 格式暂不支持。
|
||||||
|
|
||||||
use std::{
|
use std::{
|
||||||
io::{Read, Write},
|
io::{Read, Write}, path::Path
|
||||||
path::Path,
|
|
||||||
};
|
};
|
||||||
|
|
||||||
use encoding_rs::GB18030;
|
use encoding_rs::GB18030;
|
||||||
|
|||||||
@ -413,8 +413,7 @@ 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,7 +131,6 @@ 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(());
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -162,46 +161,6 @@ 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,9 +1,5 @@
|
|||||||
use axum::{
|
use axum::{
|
||||||
Router,
|
Router, body::Body, http::{StatusCode, Uri, header}, response::{IntoResponse, Response}, routing::get
|
||||||
body::Body,
|
|
||||||
http::{StatusCode, Uri, header},
|
|
||||||
response::{IntoResponse, Response},
|
|
||||||
routing::get,
|
|
||||||
};
|
};
|
||||||
use rust_embed::Embed;
|
use rust_embed::Embed;
|
||||||
|
|
||||||
|
|||||||
@ -23,8 +23,7 @@ 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,
|
pool: &SqlitePool, file_id: i64, image_filename: &str, description: &str, raw_response: &str, source: &str, jpg_filename: Option<&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)
|
||||||
@ -69,9 +68,7 @@ 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(|| {
|
.with_context(|| format!("failed to delete image_description for file_id={} filename={}", file_id, image_filename))?;
|
||||||
format!("failed to delete image_description for file_id={} filename={}", file_id, image_filename)
|
|
||||||
})?;
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -97,16 +94,13 @@ struct FileImageDescriptionRow {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// 批量查询多个文件的图片描述,返回 (file_id, image_filename) -> description 映射。
|
/// 批量查询多个文件的图片描述,返回 (file_id, image_filename) -> description 映射。
|
||||||
pub async fn list_by_file_ids(
|
pub async fn list_by_file_ids(pool: &SqlitePool, file_ids: &[i64]) -> Result<HashMap<(i64, String), (String, Option<String>)>> {
|
||||||
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> = QueryBuilder::new(
|
let mut query_builder: QueryBuilder<'_, Sqlite> =
|
||||||
"SELECT file_id, image_filename, description, jpg_filename FROM image_descriptions WHERE file_id IN (",
|
QueryBuilder::new("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);
|
||||||
@ -129,16 +123,7 @@ 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(
|
save(pool, dst_file_id, &row.image_filename, &row.description, &row.raw_response, &row.source, row.jpg_filename.as_deref()).await?;
|
||||||
pool,
|
|
||||||
dst_file_id,
|
|
||||||
&row.image_filename,
|
|
||||||
&row.description,
|
|
||||||
&row.raw_response,
|
|
||||||
&row.source,
|
|
||||||
row.jpg_filename.as_deref(),
|
|
||||||
)
|
|
||||||
.await?;
|
|
||||||
}
|
}
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|||||||
@ -1,8 +1,5 @@
|
|||||||
use axum::{
|
use axum::{
|
||||||
body::Body,
|
body::Body, http::{Request, StatusCode}, middleware::Next, response::{IntoResponse, Response}
|
||||||
http::{Request, StatusCode},
|
|
||||||
middleware::Next,
|
|
||||||
response::{IntoResponse, Response},
|
|
||||||
};
|
};
|
||||||
use base64::{Engine, engine::general_purpose::STANDARD};
|
use base64::{Engine, engine::general_purpose::STANDARD};
|
||||||
|
|
||||||
@ -22,7 +19,6 @@ 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,9 +1,7 @@
|
|||||||
use std::{env, str::FromStr};
|
use std::{env, str::FromStr};
|
||||||
|
|
||||||
use log4rs::{
|
use log4rs::{
|
||||||
append::console::ConsoleAppender,
|
append::console::ConsoleAppender, config::{Appender, Config, Logger, Root}, encode::pattern::PatternEncoder
|
||||||
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(all(not(windows), debug_assertions))]
|
#[cfg(debug_assertions)]
|
||||||
#[global_allocator]
|
#[global_allocator]
|
||||||
static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
|
static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
|
||||||
|
|
||||||
#[cfg(all(not(windows), debug_assertions, feature = "profiling"))]
|
#[cfg(all(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, weknora_sync};
|
use htknow::{api, auth, config, db, frontend, log4rs, processor, search};
|
||||||
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,7 +27,6 @@ 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,9 +5,7 @@
|
|||||||
use anyhow::Result;
|
use anyhow::Result;
|
||||||
use log::info;
|
use log::info;
|
||||||
use lopdf::{
|
use lopdf::{
|
||||||
Document, Object, ObjectId, Stream,
|
Document, Object, ObjectId, Stream, content::{Content, Operation}, dictionary
|
||||||
content::{Content, Operation},
|
|
||||||
dictionary,
|
|
||||||
};
|
};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
|
|||||||
@ -179,16 +179,8 @@ 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) = image_description::save(
|
if let Err(err) =
|
||||||
pool,
|
image_description::save(pool, file_id, name, &resp.description, &resp.raw_response, source, resp.jpg_filename.as_deref()).await
|
||||||
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);
|
||||||
}
|
}
|
||||||
@ -2054,7 +2046,12 @@ impl FileProcessor {
|
|||||||
if output.status.success() {
|
if output.status.success() {
|
||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
anyhow::bail!("{} exited with {}: {}", program, output.status, String::from_utf8_lossy(&output.stderr).trim());
|
anyhow::bail!(
|
||||||
|
"{} exited with {}: {}",
|
||||||
|
program,
|
||||||
|
output.status,
|
||||||
|
String::from_utf8_lossy(&output.stderr).trim()
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn process_uploaded_image_file(
|
async fn process_uploaded_image_file(
|
||||||
@ -2168,10 +2165,7 @@ 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!(
|
warn!("Keeping one MinerU image block for uploaded image {}, dropped {} duplicate blocks", file.filename, removed);
|
||||||
"Keeping one MinerU image block for uploaded image {}, dropped {} duplicate blocks",
|
|
||||||
file.filename, removed
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
@ -2185,10 +2179,7 @@ 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!(
|
warn!("Keeping one MinerU image block for uploaded image {}, dropped {} duplicate blocks", file.filename, removed);
|
||||||
"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? {
|
||||||
@ -2304,7 +2295,11 @@ 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() { None } else { Some(vec![caption]) },
|
image_caption: if caption.is_empty() {
|
||||||
|
None
|
||||||
|
} else {
|
||||||
|
Some(vec![caption])
|
||||||
|
},
|
||||||
table_body: None,
|
table_body: None,
|
||||||
table_caption: None,
|
table_caption: None,
|
||||||
});
|
});
|
||||||
@ -2952,10 +2947,6 @@ 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(())
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -3575,9 +3566,6 @@ 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(());
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -3652,9 +3640,6 @@ 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,48 +29,8 @@ 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 {
|
||||||
@ -1247,8 +1207,37 @@ impl SearchEngine {
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
// 两个召回分支的原始分数量纲不同,按排名进行 RRF 融合。
|
// 合并结果:使用 HashMap 按 id 去重,保留最高分数,同时去除内容为空的结果
|
||||||
let merged_results = reciprocal_rank_fusion([tantivy_results, lancedb_results]);
|
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));
|
||||||
|
|
||||||
info!("Merged results count: {}", merged_results.len());
|
info!("Merged results count: {}", merged_results.len());
|
||||||
|
|
||||||
@ -1377,7 +1366,29 @@ impl SearchEngine {
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
let merged_results = reciprocal_rank_fusion([tantivy_results, lancedb_results]);
|
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));
|
||||||
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() {
|
||||||
@ -1956,33 +1967,3 @@ 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,25 +1,14 @@
|
|||||||
use std::{
|
use std::{
|
||||||
collections::{HashMap, HashSet},
|
collections::{HashMap, HashSet}, path::Path, sync::{
|
||||||
path::Path,
|
Arc, mpsc::{self, Sender}
|
||||||
sync::{
|
}, thread, time::{Duration, Instant, SystemTime, UNIX_EPOCH}
|
||||||
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,
|
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 _}
|
||||||
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;
|
||||||
|
|||||||
@ -1,265 +0,0 @@
|
|||||||
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,7 +1,6 @@
|
|||||||
mod common;
|
mod common;
|
||||||
use std::{
|
use std::{
|
||||||
fs,
|
fs, time::{Duration, Instant}
|
||||||
time::{Duration, Instant},
|
|
||||||
};
|
};
|
||||||
|
|
||||||
use axum::http::StatusCode;
|
use axum::http::StatusCode;
|
||||||
|
|||||||
@ -1,19 +1,14 @@
|
|||||||
#![allow(dead_code, unused_imports)]
|
#![allow(dead_code, unused_imports)]
|
||||||
|
|
||||||
use std::{
|
use std::{
|
||||||
fs,
|
fs, path::PathBuf, sync::{
|
||||||
path::PathBuf,
|
OnceLock, atomic::{AtomicUsize, Ordering}
|
||||||
sync::{
|
}, time::{SystemTime, UNIX_EPOCH}
|
||||||
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,
|
body::Body, http::{Request, Response, StatusCode, header}
|
||||||
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