From 1ce10dc1f24190a98e0a94341c2335b01ee262ab Mon Sep 17 00:00:00 2001 From: Defeng Date: Tue, 28 Jul 2026 10:15:27 +0800 Subject: [PATCH] feat: synchronize processed files to WeKnora --- Cargo.toml | 5 +- src/api/common.rs | 8 +- src/api/error.rs | 4 +- src/api/file.rs | 67 +- src/api/graph.rs | 4 +- src/api/system.rs | 31 +- src/archive.rs | 3 +- src/config.rs | 5 +- src/db.rs | 41 + src/frontend.rs | 6 +- src/image_description.rs | 27 +- src/lib.rs | 6 +- src/log4rs.rs | 4 +- src/main.rs | 7 +- src/pdf_highlight.rs | 4 +- src/processor.rs | 8281 +++++++++++++++++----------------- src/search/mod.rs | 127 +- src/search/tantivy_engine.rs | 21 +- src/weknora_sync.rs | 265 ++ tests/api_concurrency.rs | 3 +- tests/common/mod.rs | 13 +- 21 files changed, 4700 insertions(+), 4232 deletions(-) create mode 100644 src/weknora_sync.rs diff --git a/Cargo.toml b/Cargo.toml index f990c57..46e4e79 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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"] } diff --git a/src/api/common.rs b/src/api/common.rs index 03c6a92..795b347 100644 --- a/src/api/common.rs +++ b/src/api/common.rs @@ -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。 diff --git a/src/api/error.rs b/src/api/error.rs index 5dc2682..77ccc39 100644 --- a/src/api/error.rs +++ b/src/api/error.rs @@ -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; diff --git a/src/api/file.rs b/src/api/file.rs index 7166e79..da5a167 100644 --- a/src/api/file.rs +++ b/src/api/file.rs @@ -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::>>() 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 } diff --git a/src/api/graph.rs b/src/api/graph.rs index 64a8527..801dadc 100644 --- a/src/api/graph.rs +++ b/src/api/graph.rs @@ -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; diff --git a/src/api/system.rs b/src/api/system.rs index 2d47735..d0a1fba 100644 --- a/src/api/system.rs +++ b/src/api/system.rs @@ -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 { 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 { Ok(response) } -#[cfg(not(feature = "profiling"))] +#[cfg(any(not(feature = "profiling"), windows))] async fn heap_profile_impl() -> ApiResult { 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 { 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 { Ok(response) } -#[cfg(not(feature = "profiling"))] +#[cfg(any(not(feature = "profiling"), windows))] async fn heap_profile_pdf_impl() -> ApiResult { Err(ApiError::BadRequest("Heap profiling is only available with 'profiling' feature enabled.".to_string())) } diff --git a/src/archive.rs b/src/archive.rs index 9948a68..75c6f31 100644 --- a/src/archive.rs +++ b/src/archive.rs @@ -4,7 +4,8 @@ //! 7Z 和 RAR 格式暂不支持。 use std::{ - io::{Read, Write}, path::Path + io::{Read, Write}, + path::Path, }; use encoding_rs::GB18030; diff --git a/src/config.rs b/src/config.rs index bdfb89b..7186127 100644 --- a/src/config.rs +++ b/src/config.rs @@ -380,7 +380,7 @@ fn env_or_parse(key: &str, default: T) -> T { /// 配置加载器 trait #[allow(dead_code)] -pub trait ConfigLoader: Send+Sync { +pub trait ConfigLoader: Send + Sync { /// 加载配置 fn load(&self) -> anyhow::Result; } @@ -413,7 +413,8 @@ impl EtcdConfigLoader { /// 启动 watch 监听配置变化 pub async fn watch(&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 变化 diff --git a/src/db.rs b/src/db.rs index e042087..de2f183 100644 --- a/src/db.rs +++ b/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(()) } diff --git a/src/frontend.rs b/src/frontend.rs index e703493..140b52c 100644 --- a/src/frontend.rs +++ b/src/frontend.rs @@ -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; diff --git a/src/image_description.rs b/src/image_description.rs index baf9030..6d093a1 100644 --- a/src/image_description.rs +++ b/src/image_description.rs @@ -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)>> { +pub async fn list_by_file_ids( + pool: &SqlitePool, file_ids: &[i64], +) -> Result)>> { 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(()) } diff --git a/src/lib.rs b/src/lib.rs index 12b486e..ed8dbe6 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -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)] diff --git a/src/log4rs.rs b/src/log4rs.rs index 7dcc814..baa773b 100644 --- a/src/log4rs.rs +++ b/src/log4rs.rs @@ -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, Vec<(String, log::LevelFilter)>) { diff --git a/src/main.rs b/src/main.rs index 805dbe3..d64d8c0 100644 --- a/src/main.rs +++ b/src/main.rs @@ -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()); diff --git a/src/pdf_highlight.rs b/src/pdf_highlight.rs index f52e0cf..abb2b3a 100644 --- a/src/pdf_highlight.rs +++ b/src/pdf_highlight.rs @@ -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}; diff --git a/src/processor.rs b/src/processor.rs index dece906..f7d9d48 100644 --- a/src/processor.rs +++ b/src/processor.rs @@ -1,197 +1,205 @@ -use std::{ - collections::{HashMap, HashSet}, - future::Future, - path::Path, - sync::{ - Arc, - atomic::{AtomicBool, AtomicU64, Ordering}, - }, - time::{Duration, Instant, SystemTime, UNIX_EPOCH}, -}; - -use aho_corasick::{AhoCorasick, MatchKind}; -use base64::{Engine, engine::general_purpose::STANDARD}; -use futures::stream::StreamExt; -use log::{debug, error, info, warn}; -use lopdf::Document; -use once_cell::sync::Lazy; -use reqwest::multipart; -use serde::{Deserialize, Deserializer, Serialize}; -use sqlx::{QueryBuilder, Sqlite, SqlitePool}; -use tokio::{ - fs, - io::{AsyncBufReadExt, AsyncReadExt}, - process::Command, - time, -}; - -use crate::{ - api::{ - FILE_COLS_NO_CONTENT, File, collect_image_paths_for_files, collect_image_raw_paths_for_files, - effective_parse_file_id, find_reusable_parsed_file, remove_image_files, resolve_image_storage_path, - update_file_custom_image_meta, - }, - archive, config, - graph::{graph_manager::KnowledgeGraph, llm_extractor::LLMGraphExtractor}, - image_description, image_parse, - search::{self, SearchEngine, content_looks_like_image_reference, tantivy_engine}, -}; - -/// 将本地文件构造为流式 multipart Part，避免大文件全量读入内存。 -async fn stream_file_part(path: &str, filename: &str, mime_type: &str) -> anyhow::Result { - let file = fs::File::open(path).await?; - let len = file.metadata().await?.len(); - Ok(reqwest::multipart::Part::stream_with_length(file, len).file_name(filename.to_string()).mime_str(mime_type)?) -} - -/// 将 HTTP 响应体流式写入文件，并返回写入的字节数。 -async fn stream_response_to_file(response: &mut reqwest::Response, path: &std::path::Path) -> anyhow::Result { - let mut file = fs::File::create(path).await?; - let mut total: u64 = 0; - while let Some(chunk) = response.chunk().await? { - total += chunk.len() as u64; - tokio::io::AsyncWriteExt::write_all(&mut file, &chunk).await?; - } - Ok(total) -} - -/// 判断文件是否以指定字节开头。 -async fn file_starts_with(path: &std::path::Path, prefix: &[u8]) -> bool { - let mut file = match fs::File::open(path).await { - Ok(f) => f, - Err(_) => return false, - }; - let mut buf = vec![0u8; prefix.len()]; - match file.read(&mut buf).await { - Ok(n) if n == prefix.len() => buf == prefix, - _ => false, - } -} - -/// 读取文件开头若干字节并转换为文本摘要，用于错误日志。 -async fn summarize_file_start(path: &std::path::Path) -> String { - match fs::read(path).await { - Ok(bytes) => summarize_http_body(&bytes), - Err(e) => format!("", e), - } -} - -/// 复用的外部服务 HTTP client（连接池），超时取自启动时配置。 -static SERVICES_HTTP_CLIENT: Lazy = Lazy::new(|| { - let timeout = Duration::from_secs(config::get().services.request_timeout_secs); - reqwest::Client::builder().timeout(timeout).build().expect("build services http client") -}); - -#[derive(Debug, Deserialize, Serialize)] -struct Result { - #[serde(default)] - content_list: String, - #[serde(default)] - images: HashMap, -} - -#[allow(clippy::type_complexity)] -type RawImageJobs = (Vec<(String, String)>, HashMap, Vec); - -/// 并发调用外部图片文本化接口，保存结果并返回 filename -> description 映射。 -async fn parse_images_to_descriptions( - pool: &SqlitePool, file_id: i64, images: &HashMap, source: &str, original_filename: Option<&str>, -) -> anyhow::Result> { - let cfg = config::get(); - let Some(parse_url) = cfg.services.image_parse_url.as_deref() else { - return Ok(HashMap::new()); - }; - if images.is_empty() || parse_url.is_empty() { - return Ok(HashMap::new()); - } - let concurrency = cfg.services.image_parse_concurrency.max(1); - - let jobs: Vec<(String, String)> = images.iter().map(|(k, v)| (k.clone(), v.clone())).collect(); - let original_filename = original_filename.map(str::to_owned); - let mut stream = futures::stream::iter(jobs.into_iter().map(|(filename, base64)| { - let source = source.to_string(); - let pool = pool.clone(); - let original_filename = original_filename.clone(); - async move { - match image_parse::parse_image_base64(&filename, &base64, None, original_filename.as_deref()).await { - Ok(resp) => { - if let Err(err) = image_description::save( - &pool, - file_id, - &filename, - &resp.description, - &resp.raw_response, - &source, - resp.jpg_filename.as_deref(), - ) - .await - { - warn!("Failed to save image description for {}: {}", filename, err); - } - Some((filename, resp.description)) - } - Err(err) => { - warn!("Image parse failed for {}: {}", filename, err); - None - } - } - } - })) - .buffer_unordered(concurrency); - - let mut map = HashMap::new(); - while let Some(result) = stream.next().await { - if let Some((filename, description)) = result { - map.insert(filename, description); - } - } - Ok(map) -} - -/// 从数据库加载图片描述；对未命中的图片尝试从已保存的文件补全。 -async fn load_or_parse_image_descriptions( - pool: &SqlitePool, file_id: i64, image_names: &[String], source: &str, -) -> anyhow::Result> { - let mut desc_map = image_description::list_by_file(pool, file_id).await?; - let stale_error_descriptions: Vec = desc_map - .iter() - .filter(|(_, description)| image_parse::is_error_response(description)) - .map(|(filename, _)| filename.clone()) - .collect(); - for filename in stale_error_descriptions { - desc_map.remove(&filename); - if let Err(err) = image_description::delete(pool, file_id, &filename).await { - warn!("Failed to remove stale image parse error for {}: {}", filename, err); - } - } - if config::get().services.image_parse_url.is_none() { - return Ok(desc_map); - } - - for name in image_names { - if desc_map.contains_key(name) { - continue; - } - let Some(path_str) = resolve_image_storage_path(name) else { continue }; - let path = Path::new(&path_str); - if !path.exists() { - continue; - } - match image_parse::parse_image_file(path, name, None).await { - Ok(resp) => { - if let Err(err) = - image_description::save(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); - } - desc_map.insert(name.clone(), resp.description); - } - Err(err) => warn!("Failed to parse image {} from disk: {}", name, err), - } - } - Ok(desc_map) -} - +use std::{ + collections::{HashMap, HashSet}, + future::Future, + path::Path, + sync::{ + Arc, + atomic::{AtomicBool, AtomicU64, Ordering}, + }, + time::{Duration, Instant, SystemTime, UNIX_EPOCH}, +}; + +use aho_corasick::{AhoCorasick, MatchKind}; +use base64::{Engine, engine::general_purpose::STANDARD}; +use futures::stream::StreamExt; +use log::{debug, error, info, warn}; +use lopdf::Document; +use once_cell::sync::Lazy; +use reqwest::multipart; +use serde::{Deserialize, Deserializer, Serialize}; +use sqlx::{QueryBuilder, Sqlite, SqlitePool}; +use tokio::{ + fs, + io::{AsyncBufReadExt, AsyncReadExt}, + process::Command, + time, +}; + +use crate::{ + api::{ + FILE_COLS_NO_CONTENT, File, collect_image_paths_for_files, collect_image_raw_paths_for_files, + effective_parse_file_id, find_reusable_parsed_file, remove_image_files, resolve_image_storage_path, + update_file_custom_image_meta, + }, + archive, config, + graph::{graph_manager::KnowledgeGraph, llm_extractor::LLMGraphExtractor}, + image_description, image_parse, + search::{self, SearchEngine, content_looks_like_image_reference, tantivy_engine}, +}; + +/// 将本地文件构造为流式 multipart Part，避免大文件全量读入内存。 +async fn stream_file_part(path: &str, filename: &str, mime_type: &str) -> anyhow::Result { + let file = fs::File::open(path).await?; + let len = file.metadata().await?.len(); + Ok(reqwest::multipart::Part::stream_with_length(file, len).file_name(filename.to_string()).mime_str(mime_type)?) +} + +/// 将 HTTP 响应体流式写入文件，并返回写入的字节数。 +async fn stream_response_to_file(response: &mut reqwest::Response, path: &std::path::Path) -> anyhow::Result { + let mut file = fs::File::create(path).await?; + let mut total: u64 = 0; + while let Some(chunk) = response.chunk().await? { + total += chunk.len() as u64; + tokio::io::AsyncWriteExt::write_all(&mut file, &chunk).await?; + } + Ok(total) +} + +/// 判断文件是否以指定字节开头。 +async fn file_starts_with(path: &std::path::Path, prefix: &[u8]) -> bool { + let mut file = match fs::File::open(path).await { + Ok(f) => f, + Err(_) => return false, + }; + let mut buf = vec![0u8; prefix.len()]; + match file.read(&mut buf).await { + Ok(n) if n == prefix.len() => buf == prefix, + _ => false, + } +} + +/// 读取文件开头若干字节并转换为文本摘要，用于错误日志。 +async fn summarize_file_start(path: &std::path::Path) -> String { + match fs::read(path).await { + Ok(bytes) => summarize_http_body(&bytes), + Err(e) => format!("", e), + } +} + +/// 复用的外部服务 HTTP client（连接池），超时取自启动时配置。 +static SERVICES_HTTP_CLIENT: Lazy = Lazy::new(|| { + let timeout = Duration::from_secs(config::get().services.request_timeout_secs); + reqwest::Client::builder().timeout(timeout).build().expect("build services http client") +}); + +#[derive(Debug, Deserialize, Serialize)] +struct Result { + #[serde(default)] + content_list: String, + #[serde(default)] + images: HashMap, +} + +#[allow(clippy::type_complexity)] +type RawImageJobs = (Vec<(String, String)>, HashMap, Vec); + +/// 并发调用外部图片文本化接口，保存结果并返回 filename -> description 映射。 +async fn parse_images_to_descriptions( + pool: &SqlitePool, file_id: i64, images: &HashMap, source: &str, original_filename: Option<&str>, +) -> anyhow::Result> { + let cfg = config::get(); + let Some(parse_url) = cfg.services.image_parse_url.as_deref() else { + return Ok(HashMap::new()); + }; + if images.is_empty() || parse_url.is_empty() { + return Ok(HashMap::new()); + } + let concurrency = cfg.services.image_parse_concurrency.max(1); + + let jobs: Vec<(String, String)> = images.iter().map(|(k, v)| (k.clone(), v.clone())).collect(); + let original_filename = original_filename.map(str::to_owned); + let mut stream = futures::stream::iter(jobs.into_iter().map(|(filename, base64)| { + let source = source.to_string(); + let pool = pool.clone(); + let original_filename = original_filename.clone(); + async move { + match image_parse::parse_image_base64(&filename, &base64, None, original_filename.as_deref()).await { + Ok(resp) => { + if let Err(err) = image_description::save( + &pool, + file_id, + &filename, + &resp.description, + &resp.raw_response, + &source, + resp.jpg_filename.as_deref(), + ) + .await + { + warn!("Failed to save image description for {}: {}", filename, err); + } + Some((filename, resp.description)) + } + Err(err) => { + warn!("Image parse failed for {}: {}", filename, err); + None + } + } + } + })) + .buffer_unordered(concurrency); + + let mut map = HashMap::new(); + while let Some(result) = stream.next().await { + if let Some((filename, description)) = result { + map.insert(filename, description); + } + } + Ok(map) +} + +/// 从数据库加载图片描述；对未命中的图片尝试从已保存的文件补全。 +async fn load_or_parse_image_descriptions( + pool: &SqlitePool, file_id: i64, image_names: &[String], source: &str, +) -> anyhow::Result> { + let mut desc_map = image_description::list_by_file(pool, file_id).await?; + let stale_error_descriptions: Vec = desc_map + .iter() + .filter(|(_, description)| image_parse::is_error_response(description)) + .map(|(filename, _)| filename.clone()) + .collect(); + for filename in stale_error_descriptions { + desc_map.remove(&filename); + if let Err(err) = image_description::delete(pool, file_id, &filename).await { + warn!("Failed to remove stale image parse error for {}: {}", filename, err); + } + } + if config::get().services.image_parse_url.is_none() { + return Ok(desc_map); + } + + for name in image_names { + if desc_map.contains_key(name) { + continue; + } + let Some(path_str) = resolve_image_storage_path(name) else { continue }; + let path = Path::new(&path_str); + if !path.exists() { + continue; + } + match image_parse::parse_image_file(path, name, None).await { + Ok(resp) => { + if let Err(err) = image_description::save( + 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); + } + desc_map.insert(name.clone(), resp.description); + } + Err(err) => warn!("Failed to parse image {} from disk: {}", name, err), + } + } + Ok(desc_map) +} + fn enrich_content_list_with_image_descriptions(content_list: &mut [ContentItem], desc_map: &HashMap) { for item in content_list.iter_mut() { if item.typ == "image" @@ -203,12 +211,12 @@ fn enrich_content_list_with_image_descriptions(content_list: &mut [ContentItem], .or_else(|| img_name.strip_prefix("images/").and_then(|name| desc_map.get(name))) .filter(|s| !s.trim().is_empty()) { - let caption = description.clone(); - let captions = item.image_caption.get_or_insert_with(Vec::new); - if !captions.iter().any(|existing| existing == &caption) { - captions.push(caption); - } - } + let caption = description.clone(); + let captions = item.image_caption.get_or_insert_with(Vec::new); + if !captions.iter().any(|existing| existing == &caption) { + captions.push(caption); + } + } } } @@ -234,1764 +242,1764 @@ fn retain_one_image_block_for_uploaded_image(content_list: &mut Vec }); removed } - -async fn claim_pending_file_by_id(pool: &SqlitePool, file_id: i64) -> anyhow::Result> { - let sql = format!( - "UPDATE files - SET status = 2, parse_run_id = lower(hex(randomblob(16))), updated_at = strftime('%s','now') - WHERE id = ? AND status = 0 - RETURNING {FILE_COLS_NO_CONTENT}" - ); - Ok(sqlx::query_as::<_, File>(&sql).bind(file_id).fetch_optional(pool).await?) -} - -/// MinerU API 返回的结果结构 -#[derive(Debug, Deserialize)] -struct MinerUResponse { - results: MinerUResults, -} - -#[derive(Debug, Deserialize)] -#[serde(untagged)] -enum MinerUResults { - Map(HashMap), - List(Vec), -} - -#[derive(Debug, Deserialize)] -struct MinerUResultItem { - #[serde(default)] - status: String, - #[serde(default)] - error: String, - #[serde(default)] - filename: String, - #[serde(default)] - content_list: String, - #[serde(default)] - images: HashMap, -} - -#[derive(Debug, Deserialize)] -struct AnalyzePdfResponse { - code: i32, - #[serde(default)] - message: String, - data: Option, -} - -#[derive(Debug, Deserialize)] -struct AnalyzePdfData { - #[serde(default)] - content_list: Vec, -} - -#[derive(Debug, Deserialize, Serialize)] -struct AudioTranscriptionResponse { - #[serde(default)] - text: String, - #[serde(default)] - language: String, -} - -#[derive(Debug, Deserialize)] -struct CustomParseResponse { - #[serde(default)] - code: i32, - #[serde(default)] - message: String, - data: Option, -} - -#[derive(Debug, Deserialize)] -struct CustomParseData { - #[serde(default)] - slices: Vec, - #[serde(default)] - full_content: Option, - #[serde(default)] - summary: Option, - #[serde(default)] - images: Option>, - #[serde(default)] - content_list: Option>, -} - -#[derive(Debug)] -struct NormalizedCustomParseData { - slices: Vec, - full_content: Option, - summary: Option, - images: HashMap, - content_list: Option>, - image_paths: Vec, -} - -#[derive(Debug, Serialize)] -struct CustomParseReuseRequest<'a> { - pdf_contents: &'a [PdfContentRow], -} - -#[derive(Debug, Deserialize)] -struct CustomSlice { - content: String, - #[serde(default, deserialize_with = "deserialize_slice_positions")] - positions: Vec, -} -#[derive(Debug, Clone, Deserialize, Serialize)] -struct ContentItem { - #[serde(default, rename = "type")] - typ: String, - #[serde(default)] - bbox: Vec, - #[serde(default)] - page_idx: i32, - #[serde(default)] - text: Option, - #[serde(default)] - text_level: Option, - #[serde(default)] - text_format: Option, // latex - #[serde(default)] - img_path: Option, - #[serde(default)] - image_caption: Option>, - #[serde(default)] - table_body: Option, - #[serde(default)] - table_caption: Option>, -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)] -struct SlicePosition { - page_idx: i32, - bbox: [i32; 4], - #[serde(default)] - sheet_name: Option, - #[serde(default)] - row_num: Option, -} - -#[derive(Debug, Deserialize)] -struct RawSlicePosition { - page_idx: i32, - #[serde(default)] - bbox: Vec, - #[serde(default)] - sheet_name: Option, - #[serde(default)] - row_num: Option, -} - -fn deserialize_slice_positions<'de, D>(deserializer: D) -> std::result::Result, D::Error> -where - D: Deserializer<'de>, -{ - let raw_positions: Option> = Option::deserialize(deserializer)?; - let mut positions = Vec::new(); - if let Some(raw_positions) = raw_positions { - for raw in raw_positions { - if raw.bbox.len() == 4 { - positions.push(SlicePosition { - page_idx: raw.page_idx, - bbox: [raw.bbox[0], raw.bbox[1], raw.bbox[2], raw.bbox[3]], - sheet_name: raw.sheet_name, - row_num: raw.row_num, - }); - } else { - debug!("Dropping custom slice position on page {} due to invalid bbox {:?}", raw.page_idx, raw.bbox); - } - } - } - Ok(positions) -} - -#[derive(Debug, Clone)] -struct SliceWithPositions { - content: String, - positions: Vec, - is_image: bool, -} - -#[derive(Debug, Clone)] -struct Segment { - start: usize, - end: usize, - positions: Vec, -} - -type PdfContentRow = crate::pdf_content::PdfContent; - -#[derive(Debug, Clone, sqlx::FromRow)] -struct SliceRow { - id: i64, - content: String, -} - -#[derive(Debug, Clone, sqlx::FromRow)] -struct SlicePositionRecord { - slice_id: i64, - page_idx: i32, - x1: i32, - y1: i32, - x2: i32, - y2: i32, - sheet_name: Option, - row_num: Option, -} - -#[derive(Debug, Clone)] -struct ClonedSlice { - old_id: i64, - new_id: i64, - content: String, -} - -/// 文件处理器：定时从数据库读取未处理的文件并处理 -pub struct FileProcessor { - pool: SqlitePool, - search_engine: search::SearchEngine, - interval: Duration, -} - -static PARSE_PAUSED: AtomicBool = AtomicBool::new(false); -static PARSE_TIMING_RUN_SEQ: AtomicU64 = AtomicU64::new(1); - -struct ParseTimingCtx { - file_id: i64, - filename: String, - run_id: String, - pipeline: &'static str, - run_started: Instant, - step_seq: u32, -} - -impl ParseTimingCtx { - fn new(file: &File, pipeline: &'static str) -> Self { - let seq = PARSE_TIMING_RUN_SEQ.fetch_add(1, Ordering::Relaxed); - let now_ms = SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_millis()).unwrap_or_default(); - let run_id = format!("f{}-{}-{}", file.id, now_ms, seq); - let ctx = Self { - file_id: file.id, - filename: file.filename.clone(), - run_id, - pipeline, - run_started: Instant::now(), - step_seq: 0, - }; - debug!( - target: "parse_timing", - "[parse_timing] event=run_start file_id={} filename={:?} run_id={} pipeline={}", - ctx.file_id, - ctx.filename, - ctx.run_id, - ctx.pipeline - ); - ctx - } - - fn set_pipeline(&mut self, pipeline: &'static str) { - if self.pipeline == pipeline { - return; - } - self.pipeline = pipeline; - debug!( - target: "parse_timing", - "[parse_timing] event=pipeline_switch file_id={} run_id={} pipeline={}", - self.file_id, - self.run_id, - self.pipeline - ); - } - - fn step_start(&mut self, step: &'static str) -> (u32, Instant) { - self.step_seq += 1; - let seq = self.step_seq; - debug!( - target: "parse_timing", - "[parse_timing] event=step_start file_id={} run_id={} pipeline={} seq={} step={}", - self.file_id, - self.run_id, - self.pipeline, - seq, - step - ); - (seq, Instant::now()) - } - - fn step_ok(&self, step: &'static str, seq: u32, started_at: Instant) { - debug!( - target: "parse_timing", - "[parse_timing] event=step_end file_id={} run_id={} pipeline={} seq={} step={} status=ok duration_ms={}", - self.file_id, - self.run_id, - self.pipeline, - seq, - step, - started_at.elapsed().as_millis() - ); - } - - fn step_err(&self, step: &'static str, seq: u32, started_at: Instant, err: &anyhow::Error) { - let err_detail = format!("{:#}", err); - debug!( - target: "parse_timing", - "[parse_timing] event=step_end file_id={} run_id={} pipeline={} seq={} step={} status=err duration_ms={} err={:?}", - self.file_id, - self.run_id, - self.pipeline, - seq, - step, - started_at.elapsed().as_millis(), - err_detail - ); - } - - fn finish(&self, result: &anyhow::Result<()>) { - match result { - Ok(_) => { - debug!( - target: "parse_timing", - "[parse_timing] event=run_end file_id={} filename={:?} run_id={} pipeline={} status=ok total_duration_ms={} steps={}", - self.file_id, - self.filename, - self.run_id, - self.pipeline, - self.run_started.elapsed().as_millis(), - self.step_seq - ); - } - Err(err) => { - let err_detail = format!("{:#}", err); - debug!( - target: "parse_timing", - "[parse_timing] event=run_end file_id={} filename={:?} run_id={} pipeline={} status=err total_duration_ms={} steps={} err={:?}", - self.file_id, - self.filename, - self.run_id, - self.pipeline, - self.run_started.elapsed().as_millis(), - self.step_seq, - err_detail - ); - } - } - } - - async fn step(&mut self, step: &'static str, fut: Fut) -> anyhow::Result - where - Fut: Future>, - { - let (seq, started_at) = self.step_start(step); - match fut.await { - Ok(value) => { - self.step_ok(step, seq, started_at); - Ok(value) - } - Err(err) => { - self.step_err(step, seq, started_at, &err); - Err(err) - } - } - } -} - -async fn timed_step_opt(timing: Option<&mut ParseTimingCtx>, step: &'static str, fut: Fut) -> anyhow::Result -where - Fut: Future>, -{ - match timing { - Some(ctx) => ctx.step(step, fut).await, - None => fut.await, - } -} - -fn summarize_http_body(body_bytes: &[u8]) -> String { - if body_bytes.is_empty() { - return "".to_string(); - } - - let raw = String::from_utf8_lossy(body_bytes).trim().to_string(); - if raw.is_empty() { - return "".to_string(); - } - - const MAX_CHARS: usize = 800; - if raw.chars().count() > MAX_CHARS { - let preview: String = raw.chars().take(MAX_CHARS).collect(); - format!("{}...(truncated, {} bytes)", preview, body_bytes.len()) - } else { - raw - } -} - -pub fn set_parse_paused(paused: bool) { - PARSE_PAUSED.store(paused, Ordering::SeqCst); - if paused { - warn!("File parsing has been paused for index maintenance"); - } else { - info!("File parsing resumed"); - } -} - -pub fn is_parse_paused() -> bool { - PARSE_PAUSED.load(Ordering::SeqCst) -} - -impl FileProcessor { - fn artifact_key(file: &File) -> String { - let cfg = config::get(); - format!( - "{}:{}:builtin-v1:{}:{}", - file.hash, file.slice_type, cfg.slice.smart_slice_max_chars, cfg.slice.fixed_slice_overlap_chars - ) - } - - async fn ensure_parse_artifact( - &self, file: &File, full_content: &str, summary: Option<&str>, - ) -> anyhow::Result { - let key = Self::artifact_key(file); - let cfg = config::get(); - let config_hash = format!("{}:{}", cfg.slice.smart_slice_max_chars, cfg.slice.fixed_slice_overlap_chars); - sqlx::query( - "INSERT OR IGNORE INTO parse_artifacts - (artifact_key, content_hash, slice_type, parser_version, config_hash, source_file_id, full_content, summary) - VALUES (?, ?, ?, 'builtin-v1', ?, ?, ?, ?)", - ) - .bind(&key) - .bind(&file.hash) - .bind(&file.slice_type) - .bind(config_hash) - .bind(file.id) - .bind(full_content) - .bind(summary) - .execute(&self.pool) - .await?; - let artifact_id: i64 = sqlx::query_scalar("SELECT id FROM parse_artifacts WHERE artifact_key = ?") - .bind(key) - .fetch_one(&self.pool) - .await?; - let updated_source_summary = sqlx::query( - "UPDATE parse_artifacts - SET summary = ?, updated_at = strftime('%s','now') - WHERE id = ? AND source_file_id = ?", - ) - .bind(summary) - .bind(artifact_id) - .bind(file.id) - .execute(&self.pool) - .await? - .rows_affected(); - if updated_source_summary == 0 && summary.is_some() { - sqlx::query( - "UPDATE parse_artifacts - SET summary = COALESCE(NULLIF(summary, ''), ?), updated_at = strftime('%s','now') - WHERE id = ? AND (summary IS NULL OR summary = '')", - ) - .bind(summary) - .bind(artifact_id) - .execute(&self.pool) - .await?; - } - sqlx::query("UPDATE files SET artifact_id = ? WHERE id = ?") - .bind(artifact_id) - .bind(file.id) - .execute(&self.pool) - .await?; - Ok(artifact_id) - } - - async fn adopt_existing_artifact_if_needed( - &self, file: &File, artifact_id: i64, full_content: &str, summary: Option<&str>, - ) -> anyhow::Result<()> { - let source_file_id: i64 = sqlx::query_scalar("SELECT source_file_id FROM parse_artifacts WHERE id = ?") - .bind(artifact_id) - .fetch_one(&self.pool) - .await?; - if source_file_id == file.id { - return Ok(()); - } - - // 两个相同内容的文件并发完成时，artifact 唯一键只允许一个源文件。 - // 丢弃当前批次的重复持久化数据，并把搜索投影切换到共享源切片。 - self.cleanup_processing_file_data_with_retry(file.id, 3).await?; - let rows = self.fetch_slice_rows(source_file_id).await?; - let shared: Vec = - rows.into_iter().map(|row| ClonedSlice { old_id: row.id, new_id: row.id, content: row.content }).collect(); - let mut source = file.clone(); - source.id = source_file_id; - source.content = Some(full_content.to_string()); - source.summary = summary.map(str::to_string); - self.reindex_cloned_slices(file, &shared, &source).await?; - Ok(()) - } - /// 创建新的文件处理器 - /// - /// # Arguments - /// * `pool` - 数据库连接池 - /// * `search_engine` - 搜索引擎实例 - /// * `interval_secs` - 处理间隔（秒） - pub fn new(pool: SqlitePool, search_engine: search::SearchEngine, interval_secs: u64) -> Self { - Self { pool, search_engine, interval: Duration::from_secs(interval_secs) } - } - - fn services_http_client(&self) -> anyhow::Result { - // 复用全局 client，避免每次请求重建连接池（reqwest::Client 内部 Arc，clone 廉价） - Ok(SERVICES_HTTP_CLIENT.clone()) - } - - /// 启动后台处理任务 - pub fn start(self) { - let processor = Arc::new(self); - tokio::spawn(async move { - info!("File processor started with interval: {:?}", processor.interval); - - if let Err(e) = processor.reset_processing_files().await { - error!("Failed to reset in-progress files: {}", e); - } - - // 自适应轮询：队列空闲时逐步拉长检查间隔，有任务到达后立即回到基础间隔， - // 避免长期无文件时仍以固定频率空转查询数据库。上限取基础间隔的 3 倍， - // 兼顾空闲负载与新文件的响应延迟。 - let base_interval = processor.interval; - let max_idle_interval = base_interval * 3; - let mut idle_interval = base_interval; - - loop { - if is_parse_paused() { - debug!("File processor paused, sleeping for {:?}", base_interval); - time::sleep(base_interval).await; - continue; - } - // 持续处理直到没有待处理的文件 - let mut did_work = false; - loop { - if is_parse_paused() { - debug!("File processor paused while processing queue"); - break; - } - match processor.process_pending_files().await { - Ok(has_more) => { - if !has_more { - // 没有更多文件需要处理，退出内循环 - debug!("No more pending files, waiting for next check"); - break; - } - // 有更多文件，继续处理 - did_work = true; - debug!("More files pending, continuing processing"); - } - Err(e) => { - error!("Error processing files: {}", e); - break; - } - } - } - - // 根据是否处理过文件调整下次检查的等待时间 - if did_work { - idle_interval = base_interval; - } else { - idle_interval = (idle_interval * 2).min(max_idle_interval); - } - time::sleep(idle_interval).await; - } - }); - } - - /// 重置异常退出时处于“处理中”的文件状态 - async fn reset_processing_files(&self) -> anyhow::Result<()> { - let file_ids: Vec = - sqlx::query_scalar("SELECT id FROM files WHERE status = 2").fetch_all(&self.pool).await?; - - if file_ids.is_empty() { - return Ok(()); - } - - info!("Found {} in-progress files from previous run, resetting", file_ids.len()); - - for file_id in file_ids { - if let Err(e) = self.reset_processing_file_data(file_id).await { - error!("Failed to reset processing file {}: {}", file_id, e); - } - } - - Ok(()) - } - - async fn reset_processing_file_data(&self, file_id: i64) -> anyhow::Result<()> { - let image_paths = collect_image_paths_for_files(&self.pool, std::slice::from_ref(&file_id)).await?; - let mut tx = self.pool.begin().await?; - - self.delete_processing_file_data(&mut tx, file_id).await?; - sqlx::query( - "UPDATE files SET status = 0, parse_run_id = NULL, log = '', updated_at = strftime('%s','now') WHERE id = ?", - ) - .bind(file_id) - .execute(&mut *tx) - .await?; - - tx.commit().await?; - - self.delete_processing_file_content(file_id).await; - remove_image_files(image_paths).await; - - if let Err(e) = self.search_engine.delete(Some(file_id), None).await { - warn!("Failed to delete search index for file {}: {}", file_id, e); - } - - Ok(()) - } - - async fn cleanup_processing_file_data(&self, file_id: i64) -> anyhow::Result<()> { - let image_paths = collect_image_paths_for_files(&self.pool, std::slice::from_ref(&file_id)).await?; - let mut tx = self.pool.begin().await?; - self.delete_processing_file_data(&mut tx, file_id).await?; - tx.commit().await?; - - self.delete_processing_file_content(file_id).await; - remove_image_files(image_paths).await; - - if let Err(e) = self.search_engine.delete(Some(file_id), None).await { - warn!("Failed to delete search index for file {}: {}", file_id, e); - } - - Ok(()) - } - - async fn cleanup_processing_file_data_with_retry(&self, file_id: i64, max_attempts: usize) -> anyhow::Result<()> { - let max_attempts = max_attempts.max(1); - let mut attempt = 1usize; - - loop { - match self.cleanup_processing_file_data(file_id).await { - Ok(()) => return Ok(()), - Err(err) => { - if attempt >= max_attempts || !Self::is_pool_timeout_error(&err) { - return Err(err); - } - - let retry_in_ms = (attempt as u64) * 200; - warn!( - "Cleanup for file {} failed due to DB pool timeout (attempt {}/{}), retrying in {}ms", - file_id, attempt, max_attempts, retry_in_ms - ); - time::sleep(Duration::from_millis(retry_in_ms)).await; - attempt += 1; - } - } - } - } - - fn is_pool_timeout_error(err: &anyhow::Error) -> bool { - err.chain().any(|cause| { - cause.downcast_ref::().is_some_and(|sqlx_err| matches!(sqlx_err, sqlx::Error::PoolTimedOut)) - }) - } - - async fn delete_processing_file_data( - &self, tx: &mut sqlx::Transaction<'_, Sqlite>, file_id: i64, - ) -> anyhow::Result<()> { - sqlx::query("DELETE FROM entity_mentions WHERE slice_id IN (SELECT id FROM slices WHERE file_id = ?)") - .bind(file_id) - .execute(&mut **tx) - .await?; - sqlx::query("DELETE FROM slice_positions WHERE slice_id IN (SELECT id FROM slices WHERE file_id = ?)") - .bind(file_id) - .execute(&mut **tx) - .await?; - sqlx::query("DELETE FROM slices WHERE file_id = ?").bind(file_id).execute(&mut **tx).await?; - sqlx::query("UPDATE files SET summary = NULL WHERE id = ?").bind(file_id).execute(&mut **tx).await?; - Ok(()) - } - - // Keep filesystem cleanup outside the SQLite transaction so disk I/O cannot hold the writer lock. - async fn delete_processing_file_content(&self, file_id: i64) { - if let Err(e) = crate::slice_content::delete(file_id).await { - warn!("Failed to delete slice content file for processing cleanup of file {}: {}", file_id, e); - } - if let Err(e) = crate::pdf_content::delete(file_id).await { - warn!("Failed to delete PDF content file for processing cleanup of file {}: {}", file_id, e); - } - if let Err(e) = crate::file_content::delete(file_id).await { - warn!("Failed to delete content file for processing cleanup of file {}: {}", file_id, e); - } - } - - async fn persist_file_summary( - &self, file_id: i64, kb_id: Option, summary: Option<&str>, - ) -> anyhow::Result<()> { - let summary = summary.and_then(|value| { - let trimmed = value.trim(); - if trimmed.is_empty() { None } else { Some(trimmed.to_string()) } - }); - if let Some(summary) = summary { - self.search_engine.write_summary(file_id, kb_id, summary).await?; - } else { - self.search_engine.delete_summary_by_file(file_id).await?; - } - Ok(()) - } - - /// 处理所有待处理的文件 - /// 返回是否还有更多文件需要处理 - async fn process_pending_files(self: &Arc) -> anyhow::Result { - if is_parse_paused() { - return Ok(false); - } - - let cfg = config::get(); - let configured_concurrency = cfg.server.process_concurrency.max(1); - let db_safe_concurrency = (cfg.database.max_connections as usize).saturating_sub(2).max(1); - let concurrency = configured_concurrency.min(db_safe_concurrency); - if concurrency < configured_concurrency { - warn!( - "Reducing file processing concurrency from {} to {} to avoid DB pool exhaustion (max_connections={})", - configured_concurrency, concurrency, cfg.database.max_connections - ); - } - info!("Processing pending queue with dynamic claiming (concurrency: {})", concurrency); - - let mut handles = Vec::with_capacity(concurrency); - for worker_idx in 0..concurrency { - let processor = Arc::clone(self); - handles.push(tokio::spawn(async move { processor.run_pending_worker(worker_idx).await })); - } - - let mut claimed_total = 0usize; - for handle in handles { - let worker_claimed = handle.await.map_err(|e| anyhow::anyhow!("pending worker join failed: {}", e))??; - claimed_total += worker_claimed; - } - - Ok(claimed_total > 0) - } - - async fn run_pending_worker(self: Arc, worker_idx: usize) -> anyhow::Result { - let mut claimed = 0usize; - - loop { - if is_parse_paused() { - break; - } - - let Some(file) = self.claim_next_pending_file().await? else { - break; - }; - claimed += 1; - - info!("Pending worker {} claimed file {} (kb_id={:?})", worker_idx, file.id, file.kb_id); - - if let Err(e) = self.process_file_claimed(&file).await { - error!("Failed to process file {}: {}", file.id, e); - if let Err(cleanup_err) = self.cleanup_processing_file_data_with_retry(file.id, 3).await { - error!("Failed to cleanup processing data for file {}: {}", file.id, cleanup_err); - } - self.mark_file_failed(file.id, &e.to_string()).await?; - } else { - info!("Successfully processed file {}", file.id); - } - } - - Ok(claimed) - } - - /// 原子领取一个待处理文件，按知识库解析优先级排序 - async fn claim_next_pending_file(&self) -> anyhow::Result> { - let sql = format!( - "UPDATE files - SET status = 2, parse_run_id = lower(hex(randomblob(16))), updated_at = strftime('%s','now') - WHERE id = ( - SELECT id - FROM files - WHERE status = 0 - ORDER BY parse_priority DESC, created_at ASC, id ASC - LIMIT 1 - ) - AND status = 0 - RETURNING {FILE_COLS_NO_CONTENT}" - ); - let file = sqlx::query_as::<_, File>(&sql).fetch_optional(&self.pool).await?; - Ok(file) - } - - /// 所有指定文件的即时解析也必须经过同一个 compare-and-set 领取入口。 - /// 返回 `None` 表示该文件已经被后台 worker 或另一个请求领取。 - async fn claim_file_by_id(&self, file_id: i64) -> anyhow::Result> { - claim_pending_file_by_id(&self.pool, file_id).await - } - - async fn try_reuse_existing_data(&self, file: &File) -> anyhow::Result { - if !config::get().server.reuse_duplicate_files { - return Ok(false); - } - let Some(source_file) = - find_reusable_parsed_file(&self.pool, &file.hash, &file.slice_type, Some(file.id)).await? - else { - return Ok(false); - }; - - info!("Found reusable parsed file {} for new file {} (hash={})", source_file.id, file.id, file.hash); - - match self.clone_file_data(&source_file, file).await { - Ok(_) => Ok(true), - Err(err) => { - error!("Failed to reuse parsed data from file {} for file {}: {}", source_file.id, file.id, err); - if let Err(clean_err) = self.cleanup_processing_file_data_with_retry(file.id, 3).await { - warn!("Failed to cleanup file {} after reuse error: {}", file.id, clean_err); - } - sqlx::query("UPDATE files SET log = ?, updated_at = strftime('%s','now') WHERE id = ? AND status = 2") - .bind(format!("Reuse failed: {}", err)) - .bind(file.id) - .execute(&self.pool) - .await?; - Ok(false) - } - } - } - - async fn process_file_claimed(&self, file: &File) -> anyhow::Result<()> { - self.process_file_inner(file, true, false).await - } - - async fn process_file_claimed_skip_reuse(&self, file: &File) -> anyhow::Result<()> { - self.process_file_inner(file, true, true).await - } - - async fn process_file_inner(&self, file: &File, already_claimed: bool, skip_reuse: bool) -> anyhow::Result<()> { - let mut timing = ParseTimingCtx::new(file, "dispatch"); - let result = async { - if is_parse_paused() { - anyhow::bail!("parse is paused for index maintenance"); - } - info!("Processing file: {} ({})", file.filename, file.id); - if !self.ensure_file_exists(file.id, "process start").await? { - return Ok(()); - } - - let current_status: Option = timing - .step("status_check", async { - Ok(sqlx::query_scalar("SELECT status FROM files WHERE id = ?") - .bind(file.id) - .fetch_optional(&self.pool) - .await?) - }) - .await?; - if let Some(status) = current_status { - if !already_claimed && status != 0 { - info!("File {} status changed to {}, skipping processing", file.id, status); - return Ok(()); - } - if already_claimed && status != 2 { - info!("Claimed file {} status changed to {}, skipping processing", file.id, status); - return Ok(()); - } - } - - let is_storage = timing.step("storage_kb_check", self.is_storage_kb(file.kb_id)).await?; - if is_storage { - info!("Skipping parsing for storage knowledge base file {}", file.id); - timing.step("mark_storage_skipped", self.mark_file_storage_skipped(file.id)).await?; - return Ok(()); - } - - if !skip_reuse && timing.step("reuse_check", self.try_reuse_existing_data(file)).await? { - timing.set_pipeline("reuse"); - info!("File {} reused existing parsed data, skipping processing pipeline", file.id); - return Ok(()); - } - - timing - .step("set_processing_status", async { - let sql = "UPDATE files SET status = 2, updated_at = strftime('%s','now') WHERE id = ?"; - sqlx::query(sql).bind(file.id).execute(&self.pool).await?; - Ok(()) - }) - .await?; - - // 检查文件是否为 PDF 或图片 - let filename_lower = file.filename.to_lowercase(); - let is_pdf = filename_lower.ends_with(".pdf"); - let is_image = crate::search::is_image_file(&filename_lower); - let is_audio = Self::is_audio_file(&filename_lower); - - // 检查文件是否为 Office 文档 - let is_word = filename_lower.ends_with(".doc") || filename_lower.ends_with(".docx"); - let is_presentation = filename_lower.ends_with(".ppt") || filename_lower.ends_with(".pptx"); - let is_excel = filename_lower.ends_with(".xls") || filename_lower.ends_with(".xlsx"); - - let cfg = config::get(); - let custom_url = cfg.services.custom_parse_url.as_deref(); - let custom_reuse_url = cfg.services.custom_parse_reuse_url.as_deref(); - - if let Some(reuse_url) = custom_reuse_url { - info!("Custom parse reuse enabled"); - match self.process_file_with_custom_reuse_parser(file, reuse_url, Some(&mut timing)).await { - Ok(true) => { - timing.set_pipeline("custom_reuse"); - info!( - "Custom parse reuse enabled, reused existing pdf_contents for file {} via {}", - file.id, reuse_url - ); - return Ok(()); - } - Ok(false) => {} - Err(err) => { - error!("Custom parse reuse failed for file {}: {}", file.id, err); - } - } - } - - if is_excel { - timing.set_pipeline("excel"); - if !self.ensure_file_exists(file.id, "before excel processing").await? { - return Ok(()); - } - info!("Detected Excel document, parsing directly: {}", file.filename); - self.process_excel_file(file, Some(&mut timing)).await?; - return Ok(()); - } - - if is_word || is_presentation { - timing.set_pipeline("office_pdf"); - if !self.ensure_file_exists(file.id, "before office conversion").await? { - return Ok(()); - } - info!("Detected Office document, converting to PDF: {}", file.filename); - - if let Some(custom_url) = custom_url { - let stored_pdf_path = - timing.step("convert_office_to_pdf", self.convert_office_to_pdf(file)).await?; - let mut temp_file = file.clone(); - temp_file.path = stored_pdf_path.to_string_lossy().to_string(); - temp_file.filename = format!("{}.pdf", file.id); - - timing.set_pipeline("custom_parser"); - info!("Custom parse enabled, routing converted PDF for file {} to {}", file.id, custom_url); - self.process_file_with_custom_parser(&temp_file, custom_url, Some(&mut timing)).await?; - return Ok(()); - } - - self.convert_office_to_pdf_and_process(file, Some(&mut timing)).await?; - return Ok(()); - } - - // 压缩文件：不解析，直接标记为跳过 - if archive::is_archive_file(&filename_lower) { - timing.set_pipeline("archive"); - info!("Detected archive file, skipping parsing: {}", file.filename); - timing - .step("mark_archive_skipped", async { - let sql = - "UPDATE files SET status = 3, parse_run_id = NULL, log = ?, updated_at = strftime('%s','now') WHERE id = ?"; - sqlx::query(sql).bind("Archive file: not parsed").bind(file.id).execute(&self.pool).await?; - Ok(()) - }) - .await?; - return Ok(()); - } - - if let Some(custom_url) = custom_url - && is_pdf - { - timing.set_pipeline("custom_parser"); - info!("Custom parse enabled, routing file {} to {}", file.id, custom_url); - self.process_file_with_custom_parser(file, custom_url, Some(&mut timing)).await?; - return Ok(()); - } - - if is_pdf { - timing.set_pipeline("pdf"); - if !self.ensure_file_exists(file.id, "before pdf processing").await? { - return Ok(()); - } - // 处理 PDF 或图片文件 - self.process_pdf_file(file, None, false, None, Some(&mut timing)).await?; - } else if is_image { - timing.set_pipeline("image"); - if !self.ensure_file_exists(file.id, "before image embedding").await? { - return Ok(()); - } - let image_embedding = if search::embedding::image_embedding_enabled() { - match timing - .step( - "get_image_embedding", - search::embedding::get_image_embedding_from_path(&file.path, Some(&file.filename)), - ) - .await - { - Ok(embedding) => Some(Arc::new(embedding)), - Err(err) => { - // 图片向量服务不一定支持所有可上传的图片格式；解析和入库不应因此失败。 - warn!("Failed to get image embedding for {}, skipping it: {}", file.filename, err); - None - } - } - } else { - info!("Image embedding URL not configured, skipping image embedding for file {}", file.id); - None - }; - if Self::requires_mineru_image_conversion(&filename_lower) { - match Self::convert_image_for_mineru(file).await { - Ok(converted_path) => { - let mut converted_file = file.clone(); - converted_file.path = converted_path.to_string_lossy().to_string(); - converted_file.filename = format!("{}.png", file.filename); - let result = self - .process_pdf_file(&converted_file, image_embedding, true, Some(&file.filename), Some(&mut timing)) - .await; - if let Err(err) = fs::remove_file(&converted_path).await { - if err.kind() != std::io::ErrorKind::NotFound { - warn!("Failed to remove converted image {}: {}", converted_path.display(), err); - } - } - result?; - } - Err(err) => { - // 保证文件可检索；但正常情况下应优先保留 MinerU 产生的 content_list。 - warn!( - "Failed to convert image {} for MinerU, falling back to a single direct slice: {}", - file.filename, err - ); - self.process_uploaded_image_file(file, image_embedding, Some(&mut timing)).await?; - } - } - } else { - self.process_pdf_file(file, image_embedding, true, None, Some(&mut timing)).await?; - } - } else if is_audio { - timing.set_pipeline("audio"); - if !self.ensure_file_exists(file.id, "before audio processing").await? { - return Ok(()); - } - self.process_audio_file(file, Some(&mut timing)).await?; - } else { - timing.set_pipeline("text"); - if !self.ensure_file_exists(file.id, "before text processing").await? { - return Ok(()); - } - // 处理普通文本文件 - self.process_text_file(file, Some(&mut timing)).await?; - } - - Ok(()) - } - .await; - timing.finish(&result); - result - } - - async fn process_file_with_custom_parser( - &self, file: &File, custom_url: &str, mut timing: Option<&mut ParseTimingCtx>, - ) -> anyhow::Result<()> { - if !self.ensure_file_exists(file.id, "custom parse start").await? { - return Ok(()); - } - - let parse_data = - timed_step_opt(timing.as_deref_mut(), "custom_parse_api", self.call_custom_parse_api(file, custom_url)) - .await?; - let normalized = Self::normalize_custom_parse_data(parse_data)?; - - if let Some(content_list) = normalized.content_list.as_ref() { - self.insert_custom_pdf_contents(file.id, content_list, timing.as_deref_mut()).await?; - } else { - timed_step_opt(timing.as_deref_mut(), "custom_update_image_meta", async { - update_file_custom_image_meta(&self.pool, file.id, &normalized.image_paths, "custom_parse_images") - .await?; - Ok(()) - }) - .await?; - } - - if !normalized.images.is_empty() { - self.save_custom_images(&normalized.images, timing.as_deref_mut()).await?; - } - - parse_images_to_descriptions(&self.pool, file.id, &normalized.images, "custom", Some(&file.filename)).await?; - - self.save_custom_slices( - file, - &normalized.slices, - normalized.full_content.as_deref(), - normalized.summary.as_deref(), - timing, - ) - .await?; - - Ok(()) - } - - async fn process_file_with_custom_reuse_parser( - &self, file: &File, custom_reuse_url: &str, mut timing: Option<&mut ParseTimingCtx>, - ) -> anyhow::Result { - let pdf_rows = - timed_step_opt(timing.as_deref_mut(), "fetch_pdf_content_rows", self.fetch_pdf_content_rows(file.id)) - .await?; - if pdf_rows.is_empty() { - debug!("Custom parse reuse skipped for file {} because pdf_contents are empty", file.id); - return Ok(false); - } - - let payload = CustomParseReuseRequest { pdf_contents: &pdf_rows }; - let client = self.services_http_client()?; - let response = timed_step_opt(timing.as_deref_mut(), "custom_parse_reuse_api", async { - Ok(client.post(custom_reuse_url).json(&payload).send().await?) - }) - .await?; - let response_url = response.url().to_string(); - let status = response.status(); - let request_id = - response.headers().get("x-request-id").and_then(|v| v.to_str().ok()).unwrap_or("-").to_string(); - let body_bytes = response.bytes().await?; - if !status.is_success() { - return Err(anyhow::anyhow!( - "Custom parse reuse API failed (status={}, url={}, request_id={}, body_len={}): {}", - status.as_u16(), - response_url, - request_id, - body_bytes.len(), - summarize_http_body(&body_bytes) - )); - } - - let parsed: CustomParseResponse = serde_json::from_slice(&body_bytes).map_err(|e| { - anyhow::anyhow!( - "Custom parse reuse API response decode failed (status={}, url={}, body_len={}): {} - {}", - status.as_u16(), - response_url, - body_bytes.len(), - e, - summarize_http_body(&body_bytes) - ) - })?; - - if parsed.code != 200 { - let msg = if parsed.message.is_empty() { "" } else { parsed.message.as_str() }; - return Err(anyhow::anyhow!( - "Custom parse reuse API returned error code={} message={} url={}", - parsed.code, - msg, - response_url - )); - } - - let data = parsed.data.ok_or_else(|| anyhow::anyhow!("Custom parse reuse API returned empty data"))?; - if data.slices.is_empty() && data.summary.as_deref().unwrap_or("").trim().is_empty() { - return Err(anyhow::anyhow!("Custom parse reuse API returned empty slices and summary")); - } - - self.save_custom_slices(file, &data.slices, data.full_content.as_deref(), data.summary.as_deref(), timing) - .await?; - - Ok(true) - } - - fn normalize_custom_parse_data(data: CustomParseData) -> anyhow::Result { - let mut image_mapping: HashMap = HashMap::new(); - let images = data.images.unwrap_or_default(); - - for image_name in images.keys() { - Self::ensure_safe_image_name(image_name)?; - image_mapping.entry(image_name.clone()).or_insert_with(|| image_name.clone()); - } - - if let Some(content_list) = data.content_list.as_ref() { - for item in content_list { - if let Some(img_path) = item.img_path.as_deref() { - Self::ensure_safe_image_name(img_path)?; - let mapped_path = image_mapping.get(img_path).cloned().or_else(|| { - Self::basename_for_path(img_path).and_then(|basename| image_mapping.get(&basename).cloned()) - }); - let mapped_path = mapped_path.unwrap_or_else(|| img_path.to_string()); - image_mapping.entry(img_path.to_string()).or_insert(mapped_path); - } - } - } - - let image_mapping = Self::expand_image_mapping_aliases(image_mapping); - let mut referenced_image_paths = Vec::new(); - for slice in &data.slices { - for image_path in Self::extract_custom_image_refs(&slice.content) { - if Self::lookup_image_mapping(&image_mapping, &image_path).is_none() - && resolve_image_storage_path(&image_path).is_some() - && !referenced_image_paths.iter().any(|existing| existing == &image_path) - { - referenced_image_paths.push(image_path); - } - } - } - if let Some(full_content) = data.full_content.as_deref() { - for image_path in Self::extract_custom_image_refs(full_content) { - if Self::lookup_image_mapping(&image_mapping, &image_path).is_none() - && resolve_image_storage_path(&image_path).is_some() - && !referenced_image_paths.iter().any(|existing| existing == &image_path) - { - referenced_image_paths.push(image_path); - } - } - } - - let mut normalized_images = HashMap::new(); - for (image_name, image_base64) in images { - let new_name = image_mapping - .get(&image_name) - .cloned() - .or_else(|| Self::basename_for_path(&image_name).and_then(|name| image_mapping.get(&name).cloned())) - .unwrap_or_else(|| image_name.clone()); - normalized_images.insert(new_name, image_base64); - } - - let mut content_list = data.content_list; - if let Some(items) = content_list.as_mut() { - for item in items { - if let Some(img_path) = item.img_path.as_deref() - && let Some(new_path) = Self::lookup_image_mapping(&image_mapping, img_path) - { - item.img_path = Some(new_path); - } - } - } - - let slices = data - .slices - .into_iter() - .map(|mut slice| { - slice.content = Self::rewrite_custom_image_refs(&slice.content, &image_mapping); - slice - }) - .collect(); - let full_content = data.full_content.map(|content| Self::rewrite_custom_image_refs(&content, &image_mapping)); - let summary = data.summary.and_then(|summary| { - let trimmed = summary.trim(); - if trimmed.is_empty() { None } else { Some(trimmed.to_string()) } - }); - - let mut image_paths = Vec::new(); - let mut seen = HashSet::new(); - for new_path in image_mapping.values() { - let trimmed = new_path.trim(); - if !trimmed.is_empty() && seen.insert(trimmed.to_string()) { - image_paths.push(trimmed.to_string()); - } - } - for path in referenced_image_paths { - let trimmed = path.trim(); - if !trimmed.is_empty() && seen.insert(trimmed.to_string()) { - image_paths.push(trimmed.to_string()); - } - } - - Ok(NormalizedCustomParseData { - slices, - full_content, - summary, - images: normalized_images, - content_list, - image_paths, - }) - } - - fn ensure_safe_image_name(image_name: &str) -> anyhow::Result<()> { - anyhow::ensure!( - resolve_image_storage_path(image_name).is_some(), - "Custom parse returned unsafe image path: {}", - image_name - ); - Ok(()) - } - - fn lookup_image_mapping(mapping: &HashMap, image_path: &str) -> Option { - mapping - .get(image_path) - .cloned() - .or_else(|| Self::basename_for_path(image_path).and_then(|basename| mapping.get(&basename).cloned())) - } - - fn basename_for_path(path: &str) -> Option { - std::path::Path::new(path).file_name().and_then(|name| name.to_str()).map(|name| name.to_string()) - } - - fn expand_image_mapping_aliases(mut mapping: HashMap) -> HashMap { - let aliases: Vec<(String, String)> = mapping - .iter() - .filter_map(|(old_path, new_path)| { - let basename = Self::basename_for_path(old_path)?; - if basename == *old_path { None } else { Some((basename, new_path.clone())) } - }) - .collect(); - for (alias, new_path) in aliases { - mapping.entry(alias).or_insert(new_path); - } - mapping - } - - fn extract_custom_image_refs(content: &str) -> Vec { - let api_re = regex::Regex::new(r"/api/v1/knowledge/files/([^\s'\)>]+)").expect("valid image reference regex"); - let markdown_re = regex::Regex::new(r"!\[[^\]]*\]\(([^)]+)\)").expect("valid markdown image regex"); - let mut refs = Vec::new(); - for cap in api_re.captures_iter(content) { - if let Some(path) = cap.get(1).map(|m| m.as_str().trim()).filter(|path| !path.is_empty()) - && !refs.iter().any(|existing| existing == path) - { - refs.push(path.to_string()); - } - } - for cap in markdown_re.captures_iter(content) { - let Some(path) = cap.get(1).map(|m| m.as_str().trim()).filter(|path| !path.is_empty()) else { - continue; - }; - if path.starts_with("http://") || path.starts_with("https://") || path.starts_with("data:") { - continue; - } - let normalized = path.strip_prefix("/api/v1/knowledge/files/").unwrap_or(path).trim(); - if !normalized.is_empty() && !refs.iter().any(|existing| existing == normalized) { - refs.push(normalized.to_string()); - } - } - refs - } - - fn fallback_rewrite_custom_image_refs(content: &str, mapping: &HashMap) -> String { - let mut rewritten = content.to_string(); - let mut pairs: Vec<(&String, &String)> = mapping.iter().collect(); - pairs.sort_by_key(|(right, _)| std::cmp::Reverse(right.len())); - for (old_path, new_path) in pairs { - rewritten = rewritten.replace( - &format!("/api/v1/knowledge/files/{}", old_path), - &format!("/api/v1/knowledge/files/{}", new_path), - ); - rewritten = rewritten.replace(&format!("]({})", old_path), &format!("]({})", new_path)); - } - rewritten - } - - fn rewrite_custom_image_refs(content: &str, mapping: &HashMap) -> String { - if mapping.is_empty() || content.is_empty() { - return content.to_string(); - } - - let mut patterns = Vec::with_capacity(mapping.len() * 2); - let mut replacements = Vec::with_capacity(mapping.len() * 2); - for (old_path, new_path) in mapping { - patterns.push(format!("/api/v1/knowledge/files/{old_path}")); - replacements.push(format!("/api/v1/knowledge/files/{new_path}")); - patterns.push(format!("]({old_path})")); - replacements.push(format!("]({new_path})")); - } - - match AhoCorasick::builder().match_kind(MatchKind::LeftmostLongest).build(&patterns) { - Ok(ac) => ac.replace_all(content, &replacements), - Err(e) => { - warn!("Failed to build aho-corasick automaton: {e}, fallback to loop"); - Self::fallback_rewrite_custom_image_refs(content, mapping) - } - } - } - - async fn call_custom_parse_api(&self, file: &File, custom_url: &str) -> anyhow::Result { - let mime_type = mime_guess::from_path(&file.filename).first_or_octet_stream().essence_str().to_string(); - let file_part = stream_file_part(&file.path, &file.filename, &mime_type).await?; - let form = multipart::Form::new().part("file", file_part); - - let client = self.services_http_client()?; - let response = client.post(custom_url).multipart(form).send().await?; - let response_url = response.url().to_string(); - let status = response.status(); - let request_id = - response.headers().get("x-request-id").and_then(|v| v.to_str().ok()).unwrap_or("-").to_string(); - let body_bytes = response.bytes().await?; - if !status.is_success() { - return Err(anyhow::anyhow!( - "Custom parse API failed (status={}, url={}, request_id={}, body_len={}): {}", - status.as_u16(), - response_url, - request_id, - body_bytes.len(), - summarize_http_body(&body_bytes) - )); - } - - let parsed: CustomParseResponse = serde_json::from_slice(&body_bytes).map_err(|e| { - anyhow::anyhow!( - "Custom parse API response decode failed (status={}, url={}, body_len={}): {} - {}", - status.as_u16(), - response_url, - body_bytes.len(), - e, - summarize_http_body(&body_bytes) - ) - })?; - - if parsed.code != 200 { - let msg = if parsed.message.is_empty() { "" } else { parsed.message.as_str() }; - return Err(anyhow::anyhow!( - "Custom parse API returned error code={} message={} url={}", - parsed.code, - msg, - response_url - )); - } - - let data = parsed.data.ok_or_else(|| anyhow::anyhow!("Custom parse API returned empty data"))?; - if data.slices.is_empty() && data.summary.as_deref().unwrap_or("").trim().is_empty() { - return Err(anyhow::anyhow!("Custom parse API returned empty slices and summary")); - } - - Ok(data) - } - - async fn save_custom_slices( - &self, file: &File, slices: &[CustomSlice], full_content: Option<&str>, summary: Option<&str>, - timing: Option<&mut ParseTimingCtx>, - ) -> anyhow::Result<()> { - if !self.ensure_file_exists(file.id, "before writing custom slices").await? { - return Ok(()); - } - - let derived_full_content = match full_content { - Some(content) if !content.trim().is_empty() => content.to_string(), - _ => { - let mut combined = String::new(); - for (idx, slice) in slices.iter().enumerate() { - if idx > 0 { - combined.push_str("\n\n"); - } - combined.push_str(&slice.content); - } - combined - } - }; - - let wrapped: Vec = slices - .iter() - .map(|slice| { - let content = slice.content.clone(); - SliceWithPositions { - content, - positions: slice.positions.clone(), - is_image: content_looks_like_image_reference(&slice.content), - } - }) - .collect(); - let embeddings = vec![None; wrapped.len()]; - self.finish_file_processing( - file, - wrapped, - embeddings, - &derived_full_content, - "Custom parse processed successfully", - None, - summary, - timing, - ) - .await - } - - async fn insert_custom_pdf_contents( - &self, file_id: i64, content_list: &[ContentItem], timing: Option<&mut ParseTimingCtx>, - ) -> anyhow::Result<()> { - let valid_content_items: Vec = - content_list.iter().filter(|item| item.typ != "discarded").cloned().collect(); - timed_step_opt(timing, "custom_write_pdf_contents", async { - let rows = Self::content_items_to_pdf_rows(&valid_content_items); - crate::pdf_content::write(file_id, &rows).await - }) - .await - } - - fn content_items_to_pdf_rows(items: &[ContentItem]) -> Vec { - items - .iter() - .map(|item| PdfContentRow { - page_idx: item.page_idx, - bbox: if item.bbox.is_empty() { None } else { serde_json::to_string(&item.bbox).ok() }, - text: item.text.clone(), - text_level: item.text_level, - img_path: item.img_path.clone(), - table_body: item.table_body.clone(), - }) - .collect() - } - - async fn save_custom_images( - &self, images: &HashMap, timing: Option<&mut ParseTimingCtx>, - ) -> anyhow::Result<()> { - timed_step_opt(timing, "custom_save_images", async { - let cfg = config::get(); - fs::create_dir_all(&cfg.storage.images_path).await?; - info!("custom image count: {}", images.len()); - for (img_name, img_base64) in images { - let payload = match img_base64.find("base64,") { - Some(idx) => &img_base64[idx + "base64,".len()..], - None => img_base64.as_str(), - }; - let preview: String = payload.chars().take(32).collect(); - debug!("Decoding custom image {} (len={}, preview=\"{}\")", img_name, payload.len(), preview); - let bytes = STANDARD.decode(payload).map_err(|err| { - error!( - "Failed to decode custom image {} (len={}, preview=\"{}\"): {}", - img_name, - payload.len(), - preview, - err - ); - anyhow::anyhow!(err) - })?; - let Some(image_path) = resolve_image_storage_path(img_name) else { - anyhow::bail!("Custom parse returned unsafe image path: {}", img_name); - }; - if let Some(parent) = std::path::Path::new(&image_path).parent() { - fs::create_dir_all(parent).await?; - } - fs::write(image_path, bytes).await?; - } - Ok(()) - }) - .await - } - - /// 调用外部服务将 Word/PPT 文档转换为 PDF - async fn convert_office_to_pdf(&self, file: &File) -> anyhow::Result { - let cfg = config::get(); - let pdf_dir = std::path::Path::new(&cfg.storage.pdf_path); - fs::create_dir_all(pdf_dir).await?; - - let pdf_filename = format!("{}.pdf", file.id); - let stored_pdf_path = pdf_dir.join(&pdf_filename); - let temp_pdf_path = pdf_dir.join(format!(".{}.pdf.tmp", file.id)); - let mime_type = mime_guess::from_path(&file.filename).first_or_octet_stream().essence_str().to_string(); - let mut convert_url = reqwest::Url::parse(&cfg.services.office_convert_url)?; - if !convert_url.query_pairs().any(|(key, _)| key == "target_format") { - convert_url.query_pairs_mut().append_pair("target_format", "pdf"); - } - - let mut last_error: Option = None; - for attempt in 0..2 { - // 每次重试重新打开文件流，避免流式 body 不可复用。 - let file_part = stream_file_part(&file.path, &file.filename, &mime_type).await?; - let form = multipart::Form::new().part("file", file_part); - - let client = self.services_http_client()?; - let mut response = match client.post(convert_url.clone()).multipart(form).send().await { - Ok(response) => response, - Err(err) => { - last_error = Some(format!("Office convert API request failed (url={}): {}", convert_url, err)); - if attempt == 0 { - warn!( - "Office convert API request failed, retrying in 3s: {}", - last_error.as_deref().unwrap_or("unknown error") - ); - time::sleep(Duration::from_secs(3)).await; - } - continue; - } - }; - - let response_url = response.url().to_string(); - let status = response.status(); - let body_len = stream_response_to_file(&mut response, &temp_pdf_path).await?; - if !status.is_success() { - let body_text = summarize_file_start(&temp_pdf_path).await; - last_error = Some(format!( - "Office convert API failed (status={}, url={}, body_len={}): {}", - status.as_u16(), - response_url, - body_len, - body_text - )); - } else if body_len == 0 { - last_error = Some(format!("Office convert API returned empty PDF body (url={})", response_url)); - } else if !file_starts_with(&temp_pdf_path, b"%PDF-").await { - let body_text = summarize_file_start(&temp_pdf_path).await; - last_error = Some(format!( - "Office convert API returned non-PDF body (status={}, url={}, body_len={}): {}", - status.as_u16(), - response_url, - body_len, - body_text - )); - } else { - fs::rename(&temp_pdf_path, &stored_pdf_path).await?; - return Ok(stored_pdf_path); - } - - let _ = fs::remove_file(&temp_pdf_path).await; - if attempt == 0 { - warn!( - "Office convert API failed, retrying in 3s: {}", - last_error.as_deref().unwrap_or("unknown error") - ); - time::sleep(Duration::from_secs(3)).await; - } - } - - let error_msg = last_error.unwrap_or_else(|| "unknown error".to_string()); - Err(anyhow::anyhow!("Failed to convert office document to PDF: {}", error_msg)) - } - - /// 将 Word/PPT 文档转换为 PDF 并处理 - async fn convert_office_to_pdf_and_process( - &self, file: &File, mut timing: Option<&mut ParseTimingCtx>, - ) -> anyhow::Result<()> { - let stored_pdf_path = - timed_step_opt(timing.as_deref_mut(), "convert_office_to_pdf", self.convert_office_to_pdf(file)).await?; - - // 创建临时 File 结构用于处理 PDF - let mut temp_file = file.clone(); - temp_file.path = stored_pdf_path.to_string_lossy().to_string(); - temp_file.filename = format!("{}.pdf", file.id); - - // 使用 process_pdf_file 处理转换后的 PDF - self.process_pdf_file(&temp_file, None, false, Some(file.filename.as_str()), timing).await - } - - /// 处理 Excel 文件，按 sheet+行 生成切片 - async fn process_excel_file(&self, file: &File, mut timing: Option<&mut ParseTimingCtx>) -> anyhow::Result<()> { - if !self.ensure_file_exists(file.id, "excel processing start").await? { - return Ok(()); - } - info!("Processing Excel file: {}", file.filename); - - let slices = - timed_step_opt(timing.as_deref_mut(), "parse_excel", async { self.parse_excel_to_slices(file).await }) - .await?; - - if slices.is_empty() { - timed_step_opt(timing.as_deref_mut(), "finalize_empty_excel", async { - let sql = "UPDATE files SET status = 1, parse_run_id = NULL, log = ?, updated_at = strftime('%s','now') WHERE id = ?"; - sqlx::query(sql).bind("Excel processed successfully (empty)").bind(file.id).execute(&self.pool).await?; - crate::file_content::delete(file.id).await?; - Ok(()) - }) - .await?; - return Ok(()); - } - - let slice_count = slices.len(); - let full_content = slices.iter().map(|s| s.content.as_str()).collect::>().join("\n\n"); - let embeddings = vec![None; slice_count]; - self.finish_file_processing( - file, - slices, - embeddings, - &full_content, - "Excel processed successfully", - None, - None, - timing, - ) - .await - } - - /// 解析 Excel 文件，按 sheet+行 生成切片 - async fn parse_excel_to_slices(&self, file: &File) -> anyhow::Result> { - // calamine 打开/解析 Excel 是同步阻塞的 CPU+IO，放到阻塞线程池，避免阻塞 async 运行时 - let path = file.path.clone(); - let file_id = file.id; - tokio::task::spawn_blocking(move || -> anyhow::Result> { - use calamine::{Reader, open_workbook_auto}; - - let path = std::path::Path::new(&path); - let mut workbook: calamine::Sheets> = - open_workbook_auto(path).map_err(|e| anyhow::anyhow!("Failed to open Excel file: {}", e))?; - - let mut slices = Vec::new(); - let sheet_names = workbook.sheet_names().to_vec(); - - for (sheet_idx, sheet_name) in sheet_names.iter().enumerate() { - let range = match workbook.worksheet_range(sheet_name) { - Ok(r) => r, - Err(e) => { - warn!("Failed to read sheet '{}' in file {}: {}", sheet_name, file_id, e); - continue; - } - }; - - let rows: Vec<_> = range.rows().collect(); - if rows.len() < 2 { - continue; - } - - let header: Vec = rows[0] - .iter() - .enumerate() - .map(|(col_idx, cell)| { - let s = cell.to_string().trim().to_string(); - if s.is_empty() { format!("列{}", col_idx + 1) } else { s } - }) - .collect(); - - for (row_idx, row) in rows.iter().enumerate().skip(1) { - let mut lines = Vec::new(); - let mut has_data = false; - - for (col_idx, cell) in row.iter().enumerate() { - let value = cell.to_string().trim().to_string(); - if !value.is_empty() { - let header_label = - header.get(col_idx).cloned().unwrap_or_else(|| format!("列{}", col_idx + 1)); - lines.push(format!("{}: {}", header_label, value)); - has_data = true; - } - } - - if !has_data { - continue; - } - - let mut content = format!("Sheet: {}\n", sheet_name); - content.push_str(&lines.join("\n")); - - let positions = vec![SlicePosition { - page_idx: sheet_idx as i32, - bbox: [0, 0, 0, 0], - sheet_name: Some(sheet_name.clone()), - row_num: Some((row_idx + 1) as i32), - }]; - - slices.push(SliceWithPositions { content, positions, is_image: false }); - } - } - - Ok(slices) - }) - .await? - } - + +async fn claim_pending_file_by_id(pool: &SqlitePool, file_id: i64) -> anyhow::Result> { + let sql = format!( + "UPDATE files + SET status = 2, parse_run_id = lower(hex(randomblob(16))), updated_at = strftime('%s','now') + WHERE id = ? AND status = 0 + RETURNING {FILE_COLS_NO_CONTENT}" + ); + Ok(sqlx::query_as::<_, File>(&sql).bind(file_id).fetch_optional(pool).await?) +} + +/// MinerU API 返回的结果结构 +#[derive(Debug, Deserialize)] +struct MinerUResponse { + results: MinerUResults, +} + +#[derive(Debug, Deserialize)] +#[serde(untagged)] +enum MinerUResults { + Map(HashMap), + List(Vec), +} + +#[derive(Debug, Deserialize)] +struct MinerUResultItem { + #[serde(default)] + status: String, + #[serde(default)] + error: String, + #[serde(default)] + filename: String, + #[serde(default)] + content_list: String, + #[serde(default)] + images: HashMap, +} + +#[derive(Debug, Deserialize)] +struct AnalyzePdfResponse { + code: i32, + #[serde(default)] + message: String, + data: Option, +} + +#[derive(Debug, Deserialize)] +struct AnalyzePdfData { + #[serde(default)] + content_list: Vec, +} + +#[derive(Debug, Deserialize, Serialize)] +struct AudioTranscriptionResponse { + #[serde(default)] + text: String, + #[serde(default)] + language: String, +} + +#[derive(Debug, Deserialize)] +struct CustomParseResponse { + #[serde(default)] + code: i32, + #[serde(default)] + message: String, + data: Option, +} + +#[derive(Debug, Deserialize)] +struct CustomParseData { + #[serde(default)] + slices: Vec, + #[serde(default)] + full_content: Option, + #[serde(default)] + summary: Option, + #[serde(default)] + images: Option>, + #[serde(default)] + content_list: Option>, +} + +#[derive(Debug)] +struct NormalizedCustomParseData { + slices: Vec, + full_content: Option, + summary: Option, + images: HashMap, + content_list: Option>, + image_paths: Vec, +} + +#[derive(Debug, Serialize)] +struct CustomParseReuseRequest<'a> { + pdf_contents: &'a [PdfContentRow], +} + +#[derive(Debug, Deserialize)] +struct CustomSlice { + content: String, + #[serde(default, deserialize_with = "deserialize_slice_positions")] + positions: Vec, +} +#[derive(Debug, Clone, Deserialize, Serialize)] +struct ContentItem { + #[serde(default, rename = "type")] + typ: String, + #[serde(default)] + bbox: Vec, + #[serde(default)] + page_idx: i32, + #[serde(default)] + text: Option, + #[serde(default)] + text_level: Option, + #[serde(default)] + text_format: Option, // latex + #[serde(default)] + img_path: Option, + #[serde(default)] + image_caption: Option>, + #[serde(default)] + table_body: Option, + #[serde(default)] + table_caption: Option>, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)] +struct SlicePosition { + page_idx: i32, + bbox: [i32; 4], + #[serde(default)] + sheet_name: Option, + #[serde(default)] + row_num: Option, +} + +#[derive(Debug, Deserialize)] +struct RawSlicePosition { + page_idx: i32, + #[serde(default)] + bbox: Vec, + #[serde(default)] + sheet_name: Option, + #[serde(default)] + row_num: Option, +} + +fn deserialize_slice_positions<'de, D>(deserializer: D) -> std::result::Result, D::Error> +where + D: Deserializer<'de>, +{ + let raw_positions: Option> = Option::deserialize(deserializer)?; + let mut positions = Vec::new(); + if let Some(raw_positions) = raw_positions { + for raw in raw_positions { + if raw.bbox.len() == 4 { + positions.push(SlicePosition { + page_idx: raw.page_idx, + bbox: [raw.bbox[0], raw.bbox[1], raw.bbox[2], raw.bbox[3]], + sheet_name: raw.sheet_name, + row_num: raw.row_num, + }); + } else { + debug!("Dropping custom slice position on page {} due to invalid bbox {:?}", raw.page_idx, raw.bbox); + } + } + } + Ok(positions) +} + +#[derive(Debug, Clone)] +struct SliceWithPositions { + content: String, + positions: Vec, + is_image: bool, +} + +#[derive(Debug, Clone)] +struct Segment { + start: usize, + end: usize, + positions: Vec, +} + +type PdfContentRow = crate::pdf_content::PdfContent; + +#[derive(Debug, Clone, sqlx::FromRow)] +struct SliceRow { + id: i64, + content: String, +} + +#[derive(Debug, Clone, sqlx::FromRow)] +struct SlicePositionRecord { + slice_id: i64, + page_idx: i32, + x1: i32, + y1: i32, + x2: i32, + y2: i32, + sheet_name: Option, + row_num: Option, +} + +#[derive(Debug, Clone)] +struct ClonedSlice { + old_id: i64, + new_id: i64, + content: String, +} + +/// 文件处理器：定时从数据库读取未处理的文件并处理 +pub struct FileProcessor { + pool: SqlitePool, + search_engine: search::SearchEngine, + interval: Duration, +} + +static PARSE_PAUSED: AtomicBool = AtomicBool::new(false); +static PARSE_TIMING_RUN_SEQ: AtomicU64 = AtomicU64::new(1); + +struct ParseTimingCtx { + file_id: i64, + filename: String, + run_id: String, + pipeline: &'static str, + run_started: Instant, + step_seq: u32, +} + +impl ParseTimingCtx { + fn new(file: &File, pipeline: &'static str) -> Self { + let seq = PARSE_TIMING_RUN_SEQ.fetch_add(1, Ordering::Relaxed); + let now_ms = SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_millis()).unwrap_or_default(); + let run_id = format!("f{}-{}-{}", file.id, now_ms, seq); + let ctx = Self { + file_id: file.id, + filename: file.filename.clone(), + run_id, + pipeline, + run_started: Instant::now(), + step_seq: 0, + }; + debug!( + target: "parse_timing", + "[parse_timing] event=run_start file_id={} filename={:?} run_id={} pipeline={}", + ctx.file_id, + ctx.filename, + ctx.run_id, + ctx.pipeline + ); + ctx + } + + fn set_pipeline(&mut self, pipeline: &'static str) { + if self.pipeline == pipeline { + return; + } + self.pipeline = pipeline; + debug!( + target: "parse_timing", + "[parse_timing] event=pipeline_switch file_id={} run_id={} pipeline={}", + self.file_id, + self.run_id, + self.pipeline + ); + } + + fn step_start(&mut self, step: &'static str) -> (u32, Instant) { + self.step_seq += 1; + let seq = self.step_seq; + debug!( + target: "parse_timing", + "[parse_timing] event=step_start file_id={} run_id={} pipeline={} seq={} step={}", + self.file_id, + self.run_id, + self.pipeline, + seq, + step + ); + (seq, Instant::now()) + } + + fn step_ok(&self, step: &'static str, seq: u32, started_at: Instant) { + debug!( + target: "parse_timing", + "[parse_timing] event=step_end file_id={} run_id={} pipeline={} seq={} step={} status=ok duration_ms={}", + self.file_id, + self.run_id, + self.pipeline, + seq, + step, + started_at.elapsed().as_millis() + ); + } + + fn step_err(&self, step: &'static str, seq: u32, started_at: Instant, err: &anyhow::Error) { + let err_detail = format!("{:#}", err); + debug!( + target: "parse_timing", + "[parse_timing] event=step_end file_id={} run_id={} pipeline={} seq={} step={} status=err duration_ms={} err={:?}", + self.file_id, + self.run_id, + self.pipeline, + seq, + step, + started_at.elapsed().as_millis(), + err_detail + ); + } + + fn finish(&self, result: &anyhow::Result<()>) { + match result { + Ok(_) => { + debug!( + target: "parse_timing", + "[parse_timing] event=run_end file_id={} filename={:?} run_id={} pipeline={} status=ok total_duration_ms={} steps={}", + self.file_id, + self.filename, + self.run_id, + self.pipeline, + self.run_started.elapsed().as_millis(), + self.step_seq + ); + } + Err(err) => { + let err_detail = format!("{:#}", err); + debug!( + target: "parse_timing", + "[parse_timing] event=run_end file_id={} filename={:?} run_id={} pipeline={} status=err total_duration_ms={} steps={} err={:?}", + self.file_id, + self.filename, + self.run_id, + self.pipeline, + self.run_started.elapsed().as_millis(), + self.step_seq, + err_detail + ); + } + } + } + + async fn step(&mut self, step: &'static str, fut: Fut) -> anyhow::Result + where + Fut: Future>, + { + let (seq, started_at) = self.step_start(step); + match fut.await { + Ok(value) => { + self.step_ok(step, seq, started_at); + Ok(value) + } + Err(err) => { + self.step_err(step, seq, started_at, &err); + Err(err) + } + } + } +} + +async fn timed_step_opt(timing: Option<&mut ParseTimingCtx>, step: &'static str, fut: Fut) -> anyhow::Result +where + Fut: Future>, +{ + match timing { + Some(ctx) => ctx.step(step, fut).await, + None => fut.await, + } +} + +fn summarize_http_body(body_bytes: &[u8]) -> String { + if body_bytes.is_empty() { + return "".to_string(); + } + + let raw = String::from_utf8_lossy(body_bytes).trim().to_string(); + if raw.is_empty() { + return "".to_string(); + } + + const MAX_CHARS: usize = 800; + if raw.chars().count() > MAX_CHARS { + let preview: String = raw.chars().take(MAX_CHARS).collect(); + format!("{}...(truncated, {} bytes)", preview, body_bytes.len()) + } else { + raw + } +} + +pub fn set_parse_paused(paused: bool) { + PARSE_PAUSED.store(paused, Ordering::SeqCst); + if paused { + warn!("File parsing has been paused for index maintenance"); + } else { + info!("File parsing resumed"); + } +} + +pub fn is_parse_paused() -> bool { + PARSE_PAUSED.load(Ordering::SeqCst) +} + +impl FileProcessor { + fn artifact_key(file: &File) -> String { + let cfg = config::get(); + format!( + "{}:{}:builtin-v1:{}:{}", + file.hash, file.slice_type, cfg.slice.smart_slice_max_chars, cfg.slice.fixed_slice_overlap_chars + ) + } + + async fn ensure_parse_artifact( + &self, file: &File, full_content: &str, summary: Option<&str>, + ) -> anyhow::Result { + let key = Self::artifact_key(file); + let cfg = config::get(); + let config_hash = format!("{}:{}", cfg.slice.smart_slice_max_chars, cfg.slice.fixed_slice_overlap_chars); + sqlx::query( + "INSERT OR IGNORE INTO parse_artifacts + (artifact_key, content_hash, slice_type, parser_version, config_hash, source_file_id, full_content, summary) + VALUES (?, ?, ?, 'builtin-v1', ?, ?, ?, ?)", + ) + .bind(&key) + .bind(&file.hash) + .bind(&file.slice_type) + .bind(config_hash) + .bind(file.id) + .bind(full_content) + .bind(summary) + .execute(&self.pool) + .await?; + let artifact_id: i64 = sqlx::query_scalar("SELECT id FROM parse_artifacts WHERE artifact_key = ?") + .bind(key) + .fetch_one(&self.pool) + .await?; + let updated_source_summary = sqlx::query( + "UPDATE parse_artifacts + SET summary = ?, updated_at = strftime('%s','now') + WHERE id = ? AND source_file_id = ?", + ) + .bind(summary) + .bind(artifact_id) + .bind(file.id) + .execute(&self.pool) + .await? + .rows_affected(); + if updated_source_summary == 0 && summary.is_some() { + sqlx::query( + "UPDATE parse_artifacts + SET summary = COALESCE(NULLIF(summary, ''), ?), updated_at = strftime('%s','now') + WHERE id = ? AND (summary IS NULL OR summary = '')", + ) + .bind(summary) + .bind(artifact_id) + .execute(&self.pool) + .await?; + } + sqlx::query("UPDATE files SET artifact_id = ? WHERE id = ?") + .bind(artifact_id) + .bind(file.id) + .execute(&self.pool) + .await?; + Ok(artifact_id) + } + + async fn adopt_existing_artifact_if_needed( + &self, file: &File, artifact_id: i64, full_content: &str, summary: Option<&str>, + ) -> anyhow::Result<()> { + let source_file_id: i64 = sqlx::query_scalar("SELECT source_file_id FROM parse_artifacts WHERE id = ?") + .bind(artifact_id) + .fetch_one(&self.pool) + .await?; + if source_file_id == file.id { + return Ok(()); + } + + // 两个相同内容的文件并发完成时，artifact 唯一键只允许一个源文件。 + // 丢弃当前批次的重复持久化数据，并把搜索投影切换到共享源切片。 + self.cleanup_processing_file_data_with_retry(file.id, 3).await?; + let rows = self.fetch_slice_rows(source_file_id).await?; + let shared: Vec = + rows.into_iter().map(|row| ClonedSlice { old_id: row.id, new_id: row.id, content: row.content }).collect(); + let mut source = file.clone(); + source.id = source_file_id; + source.content = Some(full_content.to_string()); + source.summary = summary.map(str::to_string); + self.reindex_cloned_slices(file, &shared, &source).await?; + Ok(()) + } + /// 创建新的文件处理器 + /// + /// # Arguments + /// * `pool` - 数据库连接池 + /// * `search_engine` - 搜索引擎实例 + /// * `interval_secs` - 处理间隔（秒） + pub fn new(pool: SqlitePool, search_engine: search::SearchEngine, interval_secs: u64) -> Self { + Self { pool, search_engine, interval: Duration::from_secs(interval_secs) } + } + + fn services_http_client(&self) -> anyhow::Result { + // 复用全局 client，避免每次请求重建连接池（reqwest::Client 内部 Arc，clone 廉价） + Ok(SERVICES_HTTP_CLIENT.clone()) + } + + /// 启动后台处理任务 + pub fn start(self) { + let processor = Arc::new(self); + tokio::spawn(async move { + info!("File processor started with interval: {:?}", processor.interval); + + if let Err(e) = processor.reset_processing_files().await { + error!("Failed to reset in-progress files: {}", e); + } + + // 自适应轮询：队列空闲时逐步拉长检查间隔，有任务到达后立即回到基础间隔， + // 避免长期无文件时仍以固定频率空转查询数据库。上限取基础间隔的 3 倍， + // 兼顾空闲负载与新文件的响应延迟。 + let base_interval = processor.interval; + let max_idle_interval = base_interval * 3; + let mut idle_interval = base_interval; + + loop { + if is_parse_paused() { + debug!("File processor paused, sleeping for {:?}", base_interval); + time::sleep(base_interval).await; + continue; + } + // 持续处理直到没有待处理的文件 + let mut did_work = false; + loop { + if is_parse_paused() { + debug!("File processor paused while processing queue"); + break; + } + match processor.process_pending_files().await { + Ok(has_more) => { + if !has_more { + // 没有更多文件需要处理，退出内循环 + debug!("No more pending files, waiting for next check"); + break; + } + // 有更多文件，继续处理 + did_work = true; + debug!("More files pending, continuing processing"); + } + Err(e) => { + error!("Error processing files: {}", e); + break; + } + } + } + + // 根据是否处理过文件调整下次检查的等待时间 + if did_work { + idle_interval = base_interval; + } else { + idle_interval = (idle_interval * 2).min(max_idle_interval); + } + time::sleep(idle_interval).await; + } + }); + } + + /// 重置异常退出时处于“处理中”的文件状态 + async fn reset_processing_files(&self) -> anyhow::Result<()> { + let file_ids: Vec = + sqlx::query_scalar("SELECT id FROM files WHERE status = 2").fetch_all(&self.pool).await?; + + if file_ids.is_empty() { + return Ok(()); + } + + info!("Found {} in-progress files from previous run, resetting", file_ids.len()); + + for file_id in file_ids { + if let Err(e) = self.reset_processing_file_data(file_id).await { + error!("Failed to reset processing file {}: {}", file_id, e); + } + } + + Ok(()) + } + + async fn reset_processing_file_data(&self, file_id: i64) -> anyhow::Result<()> { + let image_paths = collect_image_paths_for_files(&self.pool, std::slice::from_ref(&file_id)).await?; + let mut tx = self.pool.begin().await?; + + self.delete_processing_file_data(&mut tx, file_id).await?; + sqlx::query( + "UPDATE files SET status = 0, parse_run_id = NULL, log = '', updated_at = strftime('%s','now') WHERE id = ?", + ) + .bind(file_id) + .execute(&mut *tx) + .await?; + + tx.commit().await?; + + self.delete_processing_file_content(file_id).await; + remove_image_files(image_paths).await; + + if let Err(e) = self.search_engine.delete(Some(file_id), None).await { + warn!("Failed to delete search index for file {}: {}", file_id, e); + } + + Ok(()) + } + + async fn cleanup_processing_file_data(&self, file_id: i64) -> anyhow::Result<()> { + let image_paths = collect_image_paths_for_files(&self.pool, std::slice::from_ref(&file_id)).await?; + let mut tx = self.pool.begin().await?; + self.delete_processing_file_data(&mut tx, file_id).await?; + tx.commit().await?; + + self.delete_processing_file_content(file_id).await; + remove_image_files(image_paths).await; + + if let Err(e) = self.search_engine.delete(Some(file_id), None).await { + warn!("Failed to delete search index for file {}: {}", file_id, e); + } + + Ok(()) + } + + async fn cleanup_processing_file_data_with_retry(&self, file_id: i64, max_attempts: usize) -> anyhow::Result<()> { + let max_attempts = max_attempts.max(1); + let mut attempt = 1usize; + + loop { + match self.cleanup_processing_file_data(file_id).await { + Ok(()) => return Ok(()), + Err(err) => { + if attempt >= max_attempts || !Self::is_pool_timeout_error(&err) { + return Err(err); + } + + let retry_in_ms = (attempt as u64) * 200; + warn!( + "Cleanup for file {} failed due to DB pool timeout (attempt {}/{}), retrying in {}ms", + file_id, attempt, max_attempts, retry_in_ms + ); + time::sleep(Duration::from_millis(retry_in_ms)).await; + attempt += 1; + } + } + } + } + + fn is_pool_timeout_error(err: &anyhow::Error) -> bool { + err.chain().any(|cause| { + cause.downcast_ref::().is_some_and(|sqlx_err| matches!(sqlx_err, sqlx::Error::PoolTimedOut)) + }) + } + + async fn delete_processing_file_data( + &self, tx: &mut sqlx::Transaction<'_, Sqlite>, file_id: i64, + ) -> anyhow::Result<()> { + sqlx::query("DELETE FROM entity_mentions WHERE slice_id IN (SELECT id FROM slices WHERE file_id = ?)") + .bind(file_id) + .execute(&mut **tx) + .await?; + sqlx::query("DELETE FROM slice_positions WHERE slice_id IN (SELECT id FROM slices WHERE file_id = ?)") + .bind(file_id) + .execute(&mut **tx) + .await?; + sqlx::query("DELETE FROM slices WHERE file_id = ?").bind(file_id).execute(&mut **tx).await?; + sqlx::query("UPDATE files SET summary = NULL WHERE id = ?").bind(file_id).execute(&mut **tx).await?; + Ok(()) + } + + // Keep filesystem cleanup outside the SQLite transaction so disk I/O cannot hold the writer lock. + async fn delete_processing_file_content(&self, file_id: i64) { + if let Err(e) = crate::slice_content::delete(file_id).await { + warn!("Failed to delete slice content file for processing cleanup of file {}: {}", file_id, e); + } + if let Err(e) = crate::pdf_content::delete(file_id).await { + warn!("Failed to delete PDF content file for processing cleanup of file {}: {}", file_id, e); + } + if let Err(e) = crate::file_content::delete(file_id).await { + warn!("Failed to delete content file for processing cleanup of file {}: {}", file_id, e); + } + } + + async fn persist_file_summary( + &self, file_id: i64, kb_id: Option, summary: Option<&str>, + ) -> anyhow::Result<()> { + let summary = summary.and_then(|value| { + let trimmed = value.trim(); + if trimmed.is_empty() { None } else { Some(trimmed.to_string()) } + }); + if let Some(summary) = summary { + self.search_engine.write_summary(file_id, kb_id, summary).await?; + } else { + self.search_engine.delete_summary_by_file(file_id).await?; + } + Ok(()) + } + + /// 处理所有待处理的文件 + /// 返回是否还有更多文件需要处理 + async fn process_pending_files(self: &Arc) -> anyhow::Result { + if is_parse_paused() { + return Ok(false); + } + + let cfg = config::get(); + let configured_concurrency = cfg.server.process_concurrency.max(1); + let db_safe_concurrency = (cfg.database.max_connections as usize).saturating_sub(2).max(1); + let concurrency = configured_concurrency.min(db_safe_concurrency); + if concurrency < configured_concurrency { + warn!( + "Reducing file processing concurrency from {} to {} to avoid DB pool exhaustion (max_connections={})", + configured_concurrency, concurrency, cfg.database.max_connections + ); + } + info!("Processing pending queue with dynamic claiming (concurrency: {})", concurrency); + + let mut handles = Vec::with_capacity(concurrency); + for worker_idx in 0..concurrency { + let processor = Arc::clone(self); + handles.push(tokio::spawn(async move { processor.run_pending_worker(worker_idx).await })); + } + + let mut claimed_total = 0usize; + for handle in handles { + let worker_claimed = handle.await.map_err(|e| anyhow::anyhow!("pending worker join failed: {}", e))??; + claimed_total += worker_claimed; + } + + Ok(claimed_total > 0) + } + + async fn run_pending_worker(self: Arc, worker_idx: usize) -> anyhow::Result { + let mut claimed = 0usize; + + loop { + if is_parse_paused() { + break; + } + + let Some(file) = self.claim_next_pending_file().await? else { + break; + }; + claimed += 1; + + info!("Pending worker {} claimed file {} (kb_id={:?})", worker_idx, file.id, file.kb_id); + + if let Err(e) = self.process_file_claimed(&file).await { + error!("Failed to process file {}: {}", file.id, e); + if let Err(cleanup_err) = self.cleanup_processing_file_data_with_retry(file.id, 3).await { + error!("Failed to cleanup processing data for file {}: {}", file.id, cleanup_err); + } + self.mark_file_failed(file.id, &e.to_string()).await?; + } else { + info!("Successfully processed file {}", file.id); + } + } + + Ok(claimed) + } + + /// 原子领取一个待处理文件，按知识库解析优先级排序 + async fn claim_next_pending_file(&self) -> anyhow::Result> { + let sql = format!( + "UPDATE files + SET status = 2, parse_run_id = lower(hex(randomblob(16))), updated_at = strftime('%s','now') + WHERE id = ( + SELECT id + FROM files + WHERE status = 0 + ORDER BY parse_priority DESC, created_at ASC, id ASC + LIMIT 1 + ) + AND status = 0 + RETURNING {FILE_COLS_NO_CONTENT}" + ); + let file = sqlx::query_as::<_, File>(&sql).fetch_optional(&self.pool).await?; + Ok(file) + } + + /// 所有指定文件的即时解析也必须经过同一个 compare-and-set 领取入口。 + /// 返回 `None` 表示该文件已经被后台 worker 或另一个请求领取。 + async fn claim_file_by_id(&self, file_id: i64) -> anyhow::Result> { + claim_pending_file_by_id(&self.pool, file_id).await + } + + async fn try_reuse_existing_data(&self, file: &File) -> anyhow::Result { + if !config::get().server.reuse_duplicate_files { + return Ok(false); + } + let Some(source_file) = + find_reusable_parsed_file(&self.pool, &file.hash, &file.slice_type, Some(file.id)).await? + else { + return Ok(false); + }; + + info!("Found reusable parsed file {} for new file {} (hash={})", source_file.id, file.id, file.hash); + + match self.clone_file_data(&source_file, file).await { + Ok(_) => Ok(true), + Err(err) => { + error!("Failed to reuse parsed data from file {} for file {}: {}", source_file.id, file.id, err); + if let Err(clean_err) = self.cleanup_processing_file_data_with_retry(file.id, 3).await { + warn!("Failed to cleanup file {} after reuse error: {}", file.id, clean_err); + } + sqlx::query("UPDATE files SET log = ?, updated_at = strftime('%s','now') WHERE id = ? AND status = 2") + .bind(format!("Reuse failed: {}", err)) + .bind(file.id) + .execute(&self.pool) + .await?; + Ok(false) + } + } + } + + async fn process_file_claimed(&self, file: &File) -> anyhow::Result<()> { + self.process_file_inner(file, true, false).await + } + + async fn process_file_claimed_skip_reuse(&self, file: &File) -> anyhow::Result<()> { + self.process_file_inner(file, true, true).await + } + + async fn process_file_inner(&self, file: &File, already_claimed: bool, skip_reuse: bool) -> anyhow::Result<()> { + let mut timing = ParseTimingCtx::new(file, "dispatch"); + let result = async { + if is_parse_paused() { + anyhow::bail!("parse is paused for index maintenance"); + } + info!("Processing file: {} ({})", file.filename, file.id); + if !self.ensure_file_exists(file.id, "process start").await? { + return Ok(()); + } + + let current_status: Option = timing + .step("status_check", async { + Ok(sqlx::query_scalar("SELECT status FROM files WHERE id = ?") + .bind(file.id) + .fetch_optional(&self.pool) + .await?) + }) + .await?; + if let Some(status) = current_status { + if !already_claimed && status != 0 { + info!("File {} status changed to {}, skipping processing", file.id, status); + return Ok(()); + } + if already_claimed && status != 2 { + info!("Claimed file {} status changed to {}, skipping processing", file.id, status); + return Ok(()); + } + } + + let is_storage = timing.step("storage_kb_check", self.is_storage_kb(file.kb_id)).await?; + if is_storage { + info!("Skipping parsing for storage knowledge base file {}", file.id); + timing.step("mark_storage_skipped", self.mark_file_storage_skipped(file.id)).await?; + return Ok(()); + } + + if !skip_reuse && timing.step("reuse_check", self.try_reuse_existing_data(file)).await? { + timing.set_pipeline("reuse"); + info!("File {} reused existing parsed data, skipping processing pipeline", file.id); + return Ok(()); + } + + timing + .step("set_processing_status", async { + let sql = "UPDATE files SET status = 2, updated_at = strftime('%s','now') WHERE id = ?"; + sqlx::query(sql).bind(file.id).execute(&self.pool).await?; + Ok(()) + }) + .await?; + + // 检查文件是否为 PDF 或图片 + let filename_lower = file.filename.to_lowercase(); + let is_pdf = filename_lower.ends_with(".pdf"); + let is_image = crate::search::is_image_file(&filename_lower); + let is_audio = Self::is_audio_file(&filename_lower); + + // 检查文件是否为 Office 文档 + let is_word = filename_lower.ends_with(".doc") || filename_lower.ends_with(".docx"); + let is_presentation = filename_lower.ends_with(".ppt") || filename_lower.ends_with(".pptx"); + let is_excel = filename_lower.ends_with(".xls") || filename_lower.ends_with(".xlsx"); + + let cfg = config::get(); + let custom_url = cfg.services.custom_parse_url.as_deref(); + let custom_reuse_url = cfg.services.custom_parse_reuse_url.as_deref(); + + if let Some(reuse_url) = custom_reuse_url { + info!("Custom parse reuse enabled"); + match self.process_file_with_custom_reuse_parser(file, reuse_url, Some(&mut timing)).await { + Ok(true) => { + timing.set_pipeline("custom_reuse"); + info!( + "Custom parse reuse enabled, reused existing pdf_contents for file {} via {}", + file.id, reuse_url + ); + return Ok(()); + } + Ok(false) => {} + Err(err) => { + error!("Custom parse reuse failed for file {}: {}", file.id, err); + } + } + } + + if is_excel { + timing.set_pipeline("excel"); + if !self.ensure_file_exists(file.id, "before excel processing").await? { + return Ok(()); + } + info!("Detected Excel document, parsing directly: {}", file.filename); + self.process_excel_file(file, Some(&mut timing)).await?; + return Ok(()); + } + + if is_word || is_presentation { + timing.set_pipeline("office_pdf"); + if !self.ensure_file_exists(file.id, "before office conversion").await? { + return Ok(()); + } + info!("Detected Office document, converting to PDF: {}", file.filename); + + if let Some(custom_url) = custom_url { + let stored_pdf_path = + timing.step("convert_office_to_pdf", self.convert_office_to_pdf(file)).await?; + let mut temp_file = file.clone(); + temp_file.path = stored_pdf_path.to_string_lossy().to_string(); + temp_file.filename = format!("{}.pdf", file.id); + + timing.set_pipeline("custom_parser"); + info!("Custom parse enabled, routing converted PDF for file {} to {}", file.id, custom_url); + self.process_file_with_custom_parser(&temp_file, custom_url, Some(&mut timing)).await?; + return Ok(()); + } + + self.convert_office_to_pdf_and_process(file, Some(&mut timing)).await?; + return Ok(()); + } + + // 压缩文件：不解析，直接标记为跳过 + if archive::is_archive_file(&filename_lower) { + timing.set_pipeline("archive"); + info!("Detected archive file, skipping parsing: {}", file.filename); + timing + .step("mark_archive_skipped", async { + let sql = + "UPDATE files SET status = 3, parse_run_id = NULL, log = ?, updated_at = strftime('%s','now') WHERE id = ?"; + sqlx::query(sql).bind("Archive file: not parsed").bind(file.id).execute(&self.pool).await?; + Ok(()) + }) + .await?; + return Ok(()); + } + + if let Some(custom_url) = custom_url + && is_pdf + { + timing.set_pipeline("custom_parser"); + info!("Custom parse enabled, routing file {} to {}", file.id, custom_url); + self.process_file_with_custom_parser(file, custom_url, Some(&mut timing)).await?; + return Ok(()); + } + + if is_pdf { + timing.set_pipeline("pdf"); + if !self.ensure_file_exists(file.id, "before pdf processing").await? { + return Ok(()); + } + // 处理 PDF 或图片文件 + self.process_pdf_file(file, None, false, None, Some(&mut timing)).await?; + } else if is_image { + timing.set_pipeline("image"); + if !self.ensure_file_exists(file.id, "before image embedding").await? { + return Ok(()); + } + let image_embedding = if search::embedding::image_embedding_enabled() { + match timing + .step( + "get_image_embedding", + search::embedding::get_image_embedding_from_path(&file.path, Some(&file.filename)), + ) + .await + { + Ok(embedding) => Some(Arc::new(embedding)), + Err(err) => { + // 图片向量服务不一定支持所有可上传的图片格式；解析和入库不应因此失败。 + warn!("Failed to get image embedding for {}, skipping it: {}", file.filename, err); + None + } + } + } else { + info!("Image embedding URL not configured, skipping image embedding for file {}", file.id); + None + }; + if Self::requires_mineru_image_conversion(&filename_lower) { + match Self::convert_image_for_mineru(file).await { + Ok(converted_path) => { + let mut converted_file = file.clone(); + converted_file.path = converted_path.to_string_lossy().to_string(); + converted_file.filename = format!("{}.png", file.filename); + let result = self + .process_pdf_file(&converted_file, image_embedding, true, Some(&file.filename), Some(&mut timing)) + .await; + if let Err(err) = fs::remove_file(&converted_path).await { + if err.kind() != std::io::ErrorKind::NotFound { + warn!("Failed to remove converted image {}: {}", converted_path.display(), err); + } + } + result?; + } + Err(err) => { + // 保证文件可检索；但正常情况下应优先保留 MinerU 产生的 content_list。 + warn!( + "Failed to convert image {} for MinerU, falling back to a single direct slice: {}", + file.filename, err + ); + self.process_uploaded_image_file(file, image_embedding, Some(&mut timing)).await?; + } + } + } else { + self.process_pdf_file(file, image_embedding, true, None, Some(&mut timing)).await?; + } + } else if is_audio { + timing.set_pipeline("audio"); + if !self.ensure_file_exists(file.id, "before audio processing").await? { + return Ok(()); + } + self.process_audio_file(file, Some(&mut timing)).await?; + } else { + timing.set_pipeline("text"); + if !self.ensure_file_exists(file.id, "before text processing").await? { + return Ok(()); + } + // 处理普通文本文件 + self.process_text_file(file, Some(&mut timing)).await?; + } + + Ok(()) + } + .await; + timing.finish(&result); + result + } + + async fn process_file_with_custom_parser( + &self, file: &File, custom_url: &str, mut timing: Option<&mut ParseTimingCtx>, + ) -> anyhow::Result<()> { + if !self.ensure_file_exists(file.id, "custom parse start").await? { + return Ok(()); + } + + let parse_data = + timed_step_opt(timing.as_deref_mut(), "custom_parse_api", self.call_custom_parse_api(file, custom_url)) + .await?; + let normalized = Self::normalize_custom_parse_data(parse_data)?; + + if let Some(content_list) = normalized.content_list.as_ref() { + self.insert_custom_pdf_contents(file.id, content_list, timing.as_deref_mut()).await?; + } else { + timed_step_opt(timing.as_deref_mut(), "custom_update_image_meta", async { + update_file_custom_image_meta(&self.pool, file.id, &normalized.image_paths, "custom_parse_images") + .await?; + Ok(()) + }) + .await?; + } + + if !normalized.images.is_empty() { + self.save_custom_images(&normalized.images, timing.as_deref_mut()).await?; + } + + parse_images_to_descriptions(&self.pool, file.id, &normalized.images, "custom", Some(&file.filename)).await?; + + self.save_custom_slices( + file, + &normalized.slices, + normalized.full_content.as_deref(), + normalized.summary.as_deref(), + timing, + ) + .await?; + + Ok(()) + } + + async fn process_file_with_custom_reuse_parser( + &self, file: &File, custom_reuse_url: &str, mut timing: Option<&mut ParseTimingCtx>, + ) -> anyhow::Result { + let pdf_rows = + timed_step_opt(timing.as_deref_mut(), "fetch_pdf_content_rows", self.fetch_pdf_content_rows(file.id)) + .await?; + if pdf_rows.is_empty() { + debug!("Custom parse reuse skipped for file {} because pdf_contents are empty", file.id); + return Ok(false); + } + + let payload = CustomParseReuseRequest { pdf_contents: &pdf_rows }; + let client = self.services_http_client()?; + let response = timed_step_opt(timing.as_deref_mut(), "custom_parse_reuse_api", async { + Ok(client.post(custom_reuse_url).json(&payload).send().await?) + }) + .await?; + let response_url = response.url().to_string(); + let status = response.status(); + let request_id = + response.headers().get("x-request-id").and_then(|v| v.to_str().ok()).unwrap_or("-").to_string(); + let body_bytes = response.bytes().await?; + if !status.is_success() { + return Err(anyhow::anyhow!( + "Custom parse reuse API failed (status={}, url={}, request_id={}, body_len={}): {}", + status.as_u16(), + response_url, + request_id, + body_bytes.len(), + summarize_http_body(&body_bytes) + )); + } + + let parsed: CustomParseResponse = serde_json::from_slice(&body_bytes).map_err(|e| { + anyhow::anyhow!( + "Custom parse reuse API response decode failed (status={}, url={}, body_len={}): {} - {}", + status.as_u16(), + response_url, + body_bytes.len(), + e, + summarize_http_body(&body_bytes) + ) + })?; + + if parsed.code != 200 { + let msg = if parsed.message.is_empty() { "" } else { parsed.message.as_str() }; + return Err(anyhow::anyhow!( + "Custom parse reuse API returned error code={} message={} url={}", + parsed.code, + msg, + response_url + )); + } + + let data = parsed.data.ok_or_else(|| anyhow::anyhow!("Custom parse reuse API returned empty data"))?; + if data.slices.is_empty() && data.summary.as_deref().unwrap_or("").trim().is_empty() { + return Err(anyhow::anyhow!("Custom parse reuse API returned empty slices and summary")); + } + + self.save_custom_slices(file, &data.slices, data.full_content.as_deref(), data.summary.as_deref(), timing) + .await?; + + Ok(true) + } + + fn normalize_custom_parse_data(data: CustomParseData) -> anyhow::Result { + let mut image_mapping: HashMap = HashMap::new(); + let images = data.images.unwrap_or_default(); + + for image_name in images.keys() { + Self::ensure_safe_image_name(image_name)?; + image_mapping.entry(image_name.clone()).or_insert_with(|| image_name.clone()); + } + + if let Some(content_list) = data.content_list.as_ref() { + for item in content_list { + if let Some(img_path) = item.img_path.as_deref() { + Self::ensure_safe_image_name(img_path)?; + let mapped_path = image_mapping.get(img_path).cloned().or_else(|| { + Self::basename_for_path(img_path).and_then(|basename| image_mapping.get(&basename).cloned()) + }); + let mapped_path = mapped_path.unwrap_or_else(|| img_path.to_string()); + image_mapping.entry(img_path.to_string()).or_insert(mapped_path); + } + } + } + + let image_mapping = Self::expand_image_mapping_aliases(image_mapping); + let mut referenced_image_paths = Vec::new(); + for slice in &data.slices { + for image_path in Self::extract_custom_image_refs(&slice.content) { + if Self::lookup_image_mapping(&image_mapping, &image_path).is_none() + && resolve_image_storage_path(&image_path).is_some() + && !referenced_image_paths.iter().any(|existing| existing == &image_path) + { + referenced_image_paths.push(image_path); + } + } + } + if let Some(full_content) = data.full_content.as_deref() { + for image_path in Self::extract_custom_image_refs(full_content) { + if Self::lookup_image_mapping(&image_mapping, &image_path).is_none() + && resolve_image_storage_path(&image_path).is_some() + && !referenced_image_paths.iter().any(|existing| existing == &image_path) + { + referenced_image_paths.push(image_path); + } + } + } + + let mut normalized_images = HashMap::new(); + for (image_name, image_base64) in images { + let new_name = image_mapping + .get(&image_name) + .cloned() + .or_else(|| Self::basename_for_path(&image_name).and_then(|name| image_mapping.get(&name).cloned())) + .unwrap_or_else(|| image_name.clone()); + normalized_images.insert(new_name, image_base64); + } + + let mut content_list = data.content_list; + if let Some(items) = content_list.as_mut() { + for item in items { + if let Some(img_path) = item.img_path.as_deref() + && let Some(new_path) = Self::lookup_image_mapping(&image_mapping, img_path) + { + item.img_path = Some(new_path); + } + } + } + + let slices = data + .slices + .into_iter() + .map(|mut slice| { + slice.content = Self::rewrite_custom_image_refs(&slice.content, &image_mapping); + slice + }) + .collect(); + let full_content = data.full_content.map(|content| Self::rewrite_custom_image_refs(&content, &image_mapping)); + let summary = data.summary.and_then(|summary| { + let trimmed = summary.trim(); + if trimmed.is_empty() { None } else { Some(trimmed.to_string()) } + }); + + let mut image_paths = Vec::new(); + let mut seen = HashSet::new(); + for new_path in image_mapping.values() { + let trimmed = new_path.trim(); + if !trimmed.is_empty() && seen.insert(trimmed.to_string()) { + image_paths.push(trimmed.to_string()); + } + } + for path in referenced_image_paths { + let trimmed = path.trim(); + if !trimmed.is_empty() && seen.insert(trimmed.to_string()) { + image_paths.push(trimmed.to_string()); + } + } + + Ok(NormalizedCustomParseData { + slices, + full_content, + summary, + images: normalized_images, + content_list, + image_paths, + }) + } + + fn ensure_safe_image_name(image_name: &str) -> anyhow::Result<()> { + anyhow::ensure!( + resolve_image_storage_path(image_name).is_some(), + "Custom parse returned unsafe image path: {}", + image_name + ); + Ok(()) + } + + fn lookup_image_mapping(mapping: &HashMap, image_path: &str) -> Option { + mapping + .get(image_path) + .cloned() + .or_else(|| Self::basename_for_path(image_path).and_then(|basename| mapping.get(&basename).cloned())) + } + + fn basename_for_path(path: &str) -> Option { + std::path::Path::new(path).file_name().and_then(|name| name.to_str()).map(|name| name.to_string()) + } + + fn expand_image_mapping_aliases(mut mapping: HashMap) -> HashMap { + let aliases: Vec<(String, String)> = mapping + .iter() + .filter_map(|(old_path, new_path)| { + let basename = Self::basename_for_path(old_path)?; + if basename == *old_path { None } else { Some((basename, new_path.clone())) } + }) + .collect(); + for (alias, new_path) in aliases { + mapping.entry(alias).or_insert(new_path); + } + mapping + } + + fn extract_custom_image_refs(content: &str) -> Vec { + let api_re = regex::Regex::new(r"/api/v1/knowledge/files/([^\s'\)>]+)").expect("valid image reference regex"); + let markdown_re = regex::Regex::new(r"!\[[^\]]*\]\(([^)]+)\)").expect("valid markdown image regex"); + let mut refs = Vec::new(); + for cap in api_re.captures_iter(content) { + if let Some(path) = cap.get(1).map(|m| m.as_str().trim()).filter(|path| !path.is_empty()) + && !refs.iter().any(|existing| existing == path) + { + refs.push(path.to_string()); + } + } + for cap in markdown_re.captures_iter(content) { + let Some(path) = cap.get(1).map(|m| m.as_str().trim()).filter(|path| !path.is_empty()) else { + continue; + }; + if path.starts_with("http://") || path.starts_with("https://") || path.starts_with("data:") { + continue; + } + let normalized = path.strip_prefix("/api/v1/knowledge/files/").unwrap_or(path).trim(); + if !normalized.is_empty() && !refs.iter().any(|existing| existing == normalized) { + refs.push(normalized.to_string()); + } + } + refs + } + + fn fallback_rewrite_custom_image_refs(content: &str, mapping: &HashMap) -> String { + let mut rewritten = content.to_string(); + let mut pairs: Vec<(&String, &String)> = mapping.iter().collect(); + pairs.sort_by_key(|(right, _)| std::cmp::Reverse(right.len())); + for (old_path, new_path) in pairs { + rewritten = rewritten.replace( + &format!("/api/v1/knowledge/files/{}", old_path), + &format!("/api/v1/knowledge/files/{}", new_path), + ); + rewritten = rewritten.replace(&format!("]({})", old_path), &format!("]({})", new_path)); + } + rewritten + } + + fn rewrite_custom_image_refs(content: &str, mapping: &HashMap) -> String { + if mapping.is_empty() || content.is_empty() { + return content.to_string(); + } + + let mut patterns = Vec::with_capacity(mapping.len() * 2); + let mut replacements = Vec::with_capacity(mapping.len() * 2); + for (old_path, new_path) in mapping { + patterns.push(format!("/api/v1/knowledge/files/{old_path}")); + replacements.push(format!("/api/v1/knowledge/files/{new_path}")); + patterns.push(format!("]({old_path})")); + replacements.push(format!("]({new_path})")); + } + + match AhoCorasick::builder().match_kind(MatchKind::LeftmostLongest).build(&patterns) { + Ok(ac) => ac.replace_all(content, &replacements), + Err(e) => { + warn!("Failed to build aho-corasick automaton: {e}, fallback to loop"); + Self::fallback_rewrite_custom_image_refs(content, mapping) + } + } + } + + async fn call_custom_parse_api(&self, file: &File, custom_url: &str) -> anyhow::Result { + let mime_type = mime_guess::from_path(&file.filename).first_or_octet_stream().essence_str().to_string(); + let file_part = stream_file_part(&file.path, &file.filename, &mime_type).await?; + let form = multipart::Form::new().part("file", file_part); + + let client = self.services_http_client()?; + let response = client.post(custom_url).multipart(form).send().await?; + let response_url = response.url().to_string(); + let status = response.status(); + let request_id = + response.headers().get("x-request-id").and_then(|v| v.to_str().ok()).unwrap_or("-").to_string(); + let body_bytes = response.bytes().await?; + if !status.is_success() { + return Err(anyhow::anyhow!( + "Custom parse API failed (status={}, url={}, request_id={}, body_len={}): {}", + status.as_u16(), + response_url, + request_id, + body_bytes.len(), + summarize_http_body(&body_bytes) + )); + } + + let parsed: CustomParseResponse = serde_json::from_slice(&body_bytes).map_err(|e| { + anyhow::anyhow!( + "Custom parse API response decode failed (status={}, url={}, body_len={}): {} - {}", + status.as_u16(), + response_url, + body_bytes.len(), + e, + summarize_http_body(&body_bytes) + ) + })?; + + if parsed.code != 200 { + let msg = if parsed.message.is_empty() { "" } else { parsed.message.as_str() }; + return Err(anyhow::anyhow!( + "Custom parse API returned error code={} message={} url={}", + parsed.code, + msg, + response_url + )); + } + + let data = parsed.data.ok_or_else(|| anyhow::anyhow!("Custom parse API returned empty data"))?; + if data.slices.is_empty() && data.summary.as_deref().unwrap_or("").trim().is_empty() { + return Err(anyhow::anyhow!("Custom parse API returned empty slices and summary")); + } + + Ok(data) + } + + async fn save_custom_slices( + &self, file: &File, slices: &[CustomSlice], full_content: Option<&str>, summary: Option<&str>, + timing: Option<&mut ParseTimingCtx>, + ) -> anyhow::Result<()> { + if !self.ensure_file_exists(file.id, "before writing custom slices").await? { + return Ok(()); + } + + let derived_full_content = match full_content { + Some(content) if !content.trim().is_empty() => content.to_string(), + _ => { + let mut combined = String::new(); + for (idx, slice) in slices.iter().enumerate() { + if idx > 0 { + combined.push_str("\n\n"); + } + combined.push_str(&slice.content); + } + combined + } + }; + + let wrapped: Vec = slices + .iter() + .map(|slice| { + let content = slice.content.clone(); + SliceWithPositions { + content, + positions: slice.positions.clone(), + is_image: content_looks_like_image_reference(&slice.content), + } + }) + .collect(); + let embeddings = vec![None; wrapped.len()]; + self.finish_file_processing( + file, + wrapped, + embeddings, + &derived_full_content, + "Custom parse processed successfully", + None, + summary, + timing, + ) + .await + } + + async fn insert_custom_pdf_contents( + &self, file_id: i64, content_list: &[ContentItem], timing: Option<&mut ParseTimingCtx>, + ) -> anyhow::Result<()> { + let valid_content_items: Vec = + content_list.iter().filter(|item| item.typ != "discarded").cloned().collect(); + timed_step_opt(timing, "custom_write_pdf_contents", async { + let rows = Self::content_items_to_pdf_rows(&valid_content_items); + crate::pdf_content::write(file_id, &rows).await + }) + .await + } + + fn content_items_to_pdf_rows(items: &[ContentItem]) -> Vec { + items + .iter() + .map(|item| PdfContentRow { + page_idx: item.page_idx, + bbox: if item.bbox.is_empty() { None } else { serde_json::to_string(&item.bbox).ok() }, + text: item.text.clone(), + text_level: item.text_level, + img_path: item.img_path.clone(), + table_body: item.table_body.clone(), + }) + .collect() + } + + async fn save_custom_images( + &self, images: &HashMap, timing: Option<&mut ParseTimingCtx>, + ) -> anyhow::Result<()> { + timed_step_opt(timing, "custom_save_images", async { + let cfg = config::get(); + fs::create_dir_all(&cfg.storage.images_path).await?; + info!("custom image count: {}", images.len()); + for (img_name, img_base64) in images { + let payload = match img_base64.find("base64,") { + Some(idx) => &img_base64[idx + "base64,".len()..], + None => img_base64.as_str(), + }; + let preview: String = payload.chars().take(32).collect(); + debug!("Decoding custom image {} (len={}, preview=\"{}\")", img_name, payload.len(), preview); + let bytes = STANDARD.decode(payload).map_err(|err| { + error!( + "Failed to decode custom image {} (len={}, preview=\"{}\"): {}", + img_name, + payload.len(), + preview, + err + ); + anyhow::anyhow!(err) + })?; + let Some(image_path) = resolve_image_storage_path(img_name) else { + anyhow::bail!("Custom parse returned unsafe image path: {}", img_name); + }; + if let Some(parent) = std::path::Path::new(&image_path).parent() { + fs::create_dir_all(parent).await?; + } + fs::write(image_path, bytes).await?; + } + Ok(()) + }) + .await + } + + /// 调用外部服务将 Word/PPT 文档转换为 PDF + async fn convert_office_to_pdf(&self, file: &File) -> anyhow::Result { + let cfg = config::get(); + let pdf_dir = std::path::Path::new(&cfg.storage.pdf_path); + fs::create_dir_all(pdf_dir).await?; + + let pdf_filename = format!("{}.pdf", file.id); + let stored_pdf_path = pdf_dir.join(&pdf_filename); + let temp_pdf_path = pdf_dir.join(format!(".{}.pdf.tmp", file.id)); + let mime_type = mime_guess::from_path(&file.filename).first_or_octet_stream().essence_str().to_string(); + let mut convert_url = reqwest::Url::parse(&cfg.services.office_convert_url)?; + if !convert_url.query_pairs().any(|(key, _)| key == "target_format") { + convert_url.query_pairs_mut().append_pair("target_format", "pdf"); + } + + let mut last_error: Option = None; + for attempt in 0..2 { + // 每次重试重新打开文件流，避免流式 body 不可复用。 + let file_part = stream_file_part(&file.path, &file.filename, &mime_type).await?; + let form = multipart::Form::new().part("file", file_part); + + let client = self.services_http_client()?; + let mut response = match client.post(convert_url.clone()).multipart(form).send().await { + Ok(response) => response, + Err(err) => { + last_error = Some(format!("Office convert API request failed (url={}): {}", convert_url, err)); + if attempt == 0 { + warn!( + "Office convert API request failed, retrying in 3s: {}", + last_error.as_deref().unwrap_or("unknown error") + ); + time::sleep(Duration::from_secs(3)).await; + } + continue; + } + }; + + let response_url = response.url().to_string(); + let status = response.status(); + let body_len = stream_response_to_file(&mut response, &temp_pdf_path).await?; + if !status.is_success() { + let body_text = summarize_file_start(&temp_pdf_path).await; + last_error = Some(format!( + "Office convert API failed (status={}, url={}, body_len={}): {}", + status.as_u16(), + response_url, + body_len, + body_text + )); + } else if body_len == 0 { + last_error = Some(format!("Office convert API returned empty PDF body (url={})", response_url)); + } else if !file_starts_with(&temp_pdf_path, b"%PDF-").await { + let body_text = summarize_file_start(&temp_pdf_path).await; + last_error = Some(format!( + "Office convert API returned non-PDF body (status={}, url={}, body_len={}): {}", + status.as_u16(), + response_url, + body_len, + body_text + )); + } else { + fs::rename(&temp_pdf_path, &stored_pdf_path).await?; + return Ok(stored_pdf_path); + } + + let _ = fs::remove_file(&temp_pdf_path).await; + if attempt == 0 { + warn!( + "Office convert API failed, retrying in 3s: {}", + last_error.as_deref().unwrap_or("unknown error") + ); + time::sleep(Duration::from_secs(3)).await; + } + } + + let error_msg = last_error.unwrap_or_else(|| "unknown error".to_string()); + Err(anyhow::anyhow!("Failed to convert office document to PDF: {}", error_msg)) + } + + /// 将 Word/PPT 文档转换为 PDF 并处理 + async fn convert_office_to_pdf_and_process( + &self, file: &File, mut timing: Option<&mut ParseTimingCtx>, + ) -> anyhow::Result<()> { + let stored_pdf_path = + timed_step_opt(timing.as_deref_mut(), "convert_office_to_pdf", self.convert_office_to_pdf(file)).await?; + + // 创建临时 File 结构用于处理 PDF + let mut temp_file = file.clone(); + temp_file.path = stored_pdf_path.to_string_lossy().to_string(); + temp_file.filename = format!("{}.pdf", file.id); + + // 使用 process_pdf_file 处理转换后的 PDF + self.process_pdf_file(&temp_file, None, false, Some(file.filename.as_str()), timing).await + } + + /// 处理 Excel 文件，按 sheet+行 生成切片 + async fn process_excel_file(&self, file: &File, mut timing: Option<&mut ParseTimingCtx>) -> anyhow::Result<()> { + if !self.ensure_file_exists(file.id, "excel processing start").await? { + return Ok(()); + } + info!("Processing Excel file: {}", file.filename); + + let slices = + timed_step_opt(timing.as_deref_mut(), "parse_excel", async { self.parse_excel_to_slices(file).await }) + .await?; + + if slices.is_empty() { + timed_step_opt(timing.as_deref_mut(), "finalize_empty_excel", async { + let sql = "UPDATE files SET status = 1, parse_run_id = NULL, log = ?, updated_at = strftime('%s','now') WHERE id = ?"; + sqlx::query(sql).bind("Excel processed successfully (empty)").bind(file.id).execute(&self.pool).await?; + crate::file_content::delete(file.id).await?; + Ok(()) + }) + .await?; + return Ok(()); + } + + let slice_count = slices.len(); + let full_content = slices.iter().map(|s| s.content.as_str()).collect::>().join("\n\n"); + let embeddings = vec![None; slice_count]; + self.finish_file_processing( + file, + slices, + embeddings, + &full_content, + "Excel processed successfully", + None, + None, + timing, + ) + .await + } + + /// 解析 Excel 文件，按 sheet+行 生成切片 + async fn parse_excel_to_slices(&self, file: &File) -> anyhow::Result> { + // calamine 打开/解析 Excel 是同步阻塞的 CPU+IO，放到阻塞线程池，避免阻塞 async 运行时 + let path = file.path.clone(); + let file_id = file.id; + tokio::task::spawn_blocking(move || -> anyhow::Result> { + use calamine::{Reader, open_workbook_auto}; + + let path = std::path::Path::new(&path); + let mut workbook: calamine::Sheets> = + open_workbook_auto(path).map_err(|e| anyhow::anyhow!("Failed to open Excel file: {}", e))?; + + let mut slices = Vec::new(); + let sheet_names = workbook.sheet_names().to_vec(); + + for (sheet_idx, sheet_name) in sheet_names.iter().enumerate() { + let range = match workbook.worksheet_range(sheet_name) { + Ok(r) => r, + Err(e) => { + warn!("Failed to read sheet '{}' in file {}: {}", sheet_name, file_id, e); + continue; + } + }; + + let rows: Vec<_> = range.rows().collect(); + if rows.len() < 2 { + continue; + } + + let header: Vec = rows[0] + .iter() + .enumerate() + .map(|(col_idx, cell)| { + let s = cell.to_string().trim().to_string(); + if s.is_empty() { format!("列{}", col_idx + 1) } else { s } + }) + .collect(); + + for (row_idx, row) in rows.iter().enumerate().skip(1) { + let mut lines = Vec::new(); + let mut has_data = false; + + for (col_idx, cell) in row.iter().enumerate() { + let value = cell.to_string().trim().to_string(); + if !value.is_empty() { + let header_label = + header.get(col_idx).cloned().unwrap_or_else(|| format!("列{}", col_idx + 1)); + lines.push(format!("{}: {}", header_label, value)); + has_data = true; + } + } + + if !has_data { + continue; + } + + let mut content = format!("Sheet: {}\n", sheet_name); + content.push_str(&lines.join("\n")); + + let positions = vec![SlicePosition { + page_idx: sheet_idx as i32, + bbox: [0, 0, 0, 0], + sheet_name: Some(sheet_name.clone()), + row_num: Some((row_idx + 1) as i32), + }]; + + slices.push(SliceWithPositions { content, positions, is_image: false }); + } + } + + Ok(slices) + }) + .await? + } + /// MinerU 不支持的图片格式需要先转为 PNG；仅在转换失败时使用轻量级入库兜底。 fn requires_mineru_image_conversion(filename_lower: &str) -> bool { filename_lower.ends_with(".svg") || filename_lower.ends_with(".ico") || filename_lower.ends_with(".avif") || filename_lower.ends_with(".bmp") - } - - /// 将 MinerU 不支持的图片临时规范化为 PNG。 - /// - /// - SVG 使用 rsvg-convert，保留矢量图的渲染结果； - /// - ICO/AVIF 使用 ImageMagick，`[0]` 选取第一帧/图层； - /// - 成功后调用方必须删除返回的临时文件。 + } + + /// 将 MinerU 不支持的图片临时规范化为 PNG。 + /// + /// - SVG 使用 rsvg-convert，保留矢量图的渲染结果； + /// - ICO/AVIF 使用 ImageMagick，`[0]` 选取第一帧/图层； + /// - 成功后调用方必须删除返回的临时文件。 async fn convert_image_for_mineru(file: &File) -> anyhow::Result { - let cfg = config::get(); - let temp_dir = std::path::Path::new(&cfg.storage.temp_path); - fs::create_dir_all(temp_dir).await?; - let output_path = temp_dir.join(format!("mineru-image-{}-{}.png", file.id, uuid::Uuid::new_v4())); - let output = output_path.to_string_lossy().to_string(); - let filename_lower = file.filename.to_ascii_lowercase(); - + let cfg = config::get(); + let temp_dir = std::path::Path::new(&cfg.storage.temp_path); + fs::create_dir_all(temp_dir).await?; + let output_path = temp_dir.join(format!("mineru-image-{}-{}.png", file.id, uuid::Uuid::new_v4())); + let output = output_path.to_string_lossy().to_string(); + let filename_lower = file.filename.to_ascii_lowercase(); + let result = if filename_lower.ends_with(".svg") { - Self::run_image_converter( - "rsvg-convert", - &[ - "--format".to_string(), - "png".to_string(), - "--background-color".to_string(), - "white".to_string(), - "--output".to_string(), - output.clone(), - file.path.clone(), - ], + Self::run_image_converter( + "rsvg-convert", + &[ + "--format".to_string(), + "png".to_string(), + "--background-color".to_string(), + "white".to_string(), + "--output".to_string(), + output.clone(), + file.path.clone(), + ], ) .await } else { @@ -2003,175 +2011,173 @@ impl FileProcessor { } else { format!("{}[0]", file.path) }; - let args = vec![ - input, - "-background".to_string(), - "white".to_string(), - "-alpha".to_string(), - "remove".to_string(), - "-alpha".to_string(), - "off".to_string(), - output.clone(), - ]; - // Debian 的 ImageMagick 6 提供 convert，7 提供 magick。优先兼容两者。 - match Self::run_image_converter("magick", &args).await { - Ok(()) => Ok(()), - Err(magick_err) => { - #[cfg(not(windows))] - { - Self::run_image_converter("convert", &args) - .await - .map_err(|convert_err| anyhow::anyhow!("magick: {}; convert: {}", magick_err, convert_err)) - } - #[cfg(windows)] - { - Err(magick_err) - } - } - } - }; - - if let Err(err) = result { - let _ = fs::remove_file(&output_path).await; - return Err(err); - } - if !fs::try_exists(&output_path).await? { - anyhow::bail!("image converter did not create {}", output_path.display()); - } - Ok(output_path) - } - - async fn run_image_converter(program: &str, args: &[String]) -> anyhow::Result<()> { - let output = Command::new(program).args(args).output().await?; - if output.status.success() { - return Ok(()); - } - anyhow::bail!( - "{} exited with {}: {}", - program, - output.status, - String::from_utf8_lossy(&output.stderr).trim() - ); - } - - async fn process_uploaded_image_file( - &self, file: &File, image_embedding: Option>>, timing: Option<&mut ParseTimingCtx>, - ) -> anyhow::Result<()> { - let mut description = String::new(); - if config::get().services.image_parse_url.is_some() { - match image_parse::parse_image_file(Path::new(&file.path), &file.filename, None).await { - Ok(response) => { - description = response.description; - if let Err(err) = image_description::save( - &self.pool, - file.id, - &file.filename, - &description, - &response.raw_response, - "upload", - response.jpg_filename.as_deref(), - ) - .await - { - warn!("Failed to save image description for {}: {}", file.filename, err); - } - } - Err(err) => { - // 图片文本化是增强能力。部分格式可能不受下游模型支持，仍应保留文件并完成索引。 - warn!("Failed to parse uploaded image {}, continuing without description: {}", file.filename, err); - } - } - } - - let content = if description.trim().is_empty() { - format!("图片文件：{}", file.filename) - } else { - format!("图片文件：{}\n{}", file.filename, description) - }; - let slices = vec![SliceWithPositions { content: content.clone(), positions: Vec::new(), is_image: true }]; - self.finish_file_processing( - file, - slices, - vec![image_embedding], - &content, - "Image processed successfully (direct format)", - None, - None, - timing, - ) - .await - } - - /// 处理 PDF 文件，调用 MinerU API - async fn process_pdf_file( - &self, file: &File, image_embedding: Option>>, is_image: bool, index_filename: Option<&str>, - mut timing: Option<&mut ParseTimingCtx>, - ) -> anyhow::Result<()> { - if !self.ensure_file_exists(file.id, "pdf processing start").await? { - return Ok(()); - } - info!("Processing PDF file: {}", file.filename); - - let existing_rows = timed_step_opt(timing.as_deref_mut(), "pdf_existing_content_check", async { - crate::pdf_content::read(file.id).await - }) - .await?; - - let mut content_list: Vec; - let mut mineru_images: HashMap = HashMap::new(); - - if !existing_rows.is_empty() - && existing_rows.iter().any(|row| row.bbox.as_deref().is_some_and(|v| !v.is_empty())) - { - info!("Found existing PDF contents for file {}, skipping MinerU API call", file.id); - + let args = vec![ + input, + "-background".to_string(), + "white".to_string(), + "-alpha".to_string(), + "remove".to_string(), + "-alpha".to_string(), + "off".to_string(), + output.clone(), + ]; + // Debian 的 ImageMagick 6 提供 convert，7 提供 magick。优先兼容两者。 + match Self::run_image_converter("magick", &args).await { + Ok(()) => Ok(()), + Err(magick_err) => { + #[cfg(not(windows))] + { + Self::run_image_converter("convert", &args) + .await + .map_err(|convert_err| anyhow::anyhow!("magick: {}; convert: {}", magick_err, convert_err)) + } + #[cfg(windows)] + { + Err(magick_err) + } + } + } + }; + + if let Err(err) = result { + let _ = fs::remove_file(&output_path).await; + return Err(err); + } + if !fs::try_exists(&output_path).await? { + anyhow::bail!("image converter did not create {}", output_path.display()); + } + Ok(output_path) + } + + async fn run_image_converter(program: &str, args: &[String]) -> anyhow::Result<()> { + let output = Command::new(program).args(args).output().await?; + if output.status.success() { + return Ok(()); + } + anyhow::bail!("{} exited with {}: {}", program, output.status, String::from_utf8_lossy(&output.stderr).trim()); + } + + async fn process_uploaded_image_file( + &self, file: &File, image_embedding: Option>>, timing: Option<&mut ParseTimingCtx>, + ) -> anyhow::Result<()> { + let mut description = String::new(); + if config::get().services.image_parse_url.is_some() { + match image_parse::parse_image_file(Path::new(&file.path), &file.filename, None).await { + Ok(response) => { + description = response.description; + if let Err(err) = image_description::save( + &self.pool, + file.id, + &file.filename, + &description, + &response.raw_response, + "upload", + response.jpg_filename.as_deref(), + ) + .await + { + warn!("Failed to save image description for {}: {}", file.filename, err); + } + } + Err(err) => { + // 图片文本化是增强能力。部分格式可能不受下游模型支持，仍应保留文件并完成索引。 + warn!("Failed to parse uploaded image {}, continuing without description: {}", file.filename, err); + } + } + } + + let content = if description.trim().is_empty() { + format!("图片文件：{}", file.filename) + } else { + format!("图片文件：{}\n{}", file.filename, description) + }; + let slices = vec![SliceWithPositions { content: content.clone(), positions: Vec::new(), is_image: true }]; + self.finish_file_processing( + file, + slices, + vec![image_embedding], + &content, + "Image processed successfully (direct format)", + None, + None, + timing, + ) + .await + } + + /// 处理 PDF 文件，调用 MinerU API + async fn process_pdf_file( + &self, file: &File, image_embedding: Option>>, is_image: bool, index_filename: Option<&str>, + mut timing: Option<&mut ParseTimingCtx>, + ) -> anyhow::Result<()> { + if !self.ensure_file_exists(file.id, "pdf processing start").await? { + return Ok(()); + } + info!("Processing PDF file: {}", file.filename); + + let existing_rows = timed_step_opt(timing.as_deref_mut(), "pdf_existing_content_check", async { + crate::pdf_content::read(file.id).await + }) + .await?; + + let mut content_list: Vec; + let mut mineru_images: HashMap = HashMap::new(); + + if !existing_rows.is_empty() + && existing_rows.iter().any(|row| row.bbox.as_deref().is_some_and(|v| !v.is_empty())) + { + info!("Found existing PDF contents for file {}, skipping MinerU API call", file.id); + content_list = timed_step_opt(timing.as_deref_mut(), "load_existing_pdf_content", async { - let content_list = existing_rows - .iter() - .map(|row| { - let typ = if row.text.is_some() { - "text".to_string() - } else if row.img_path.is_some() { - "image".to_string() - } else if row.table_body.is_some() { - "table".to_string() - } else { - "unknown".to_string() - }; - - let bbox = row - .bbox - .as_ref() - .and_then(|bbox| serde_json::from_str::>(bbox).ok()) - .unwrap_or_default(); - - ContentItem { - typ, - bbox, - page_idx: row.page_idx, - text: row.text.clone(), - text_level: row.text_level, - text_format: None, - img_path: row.img_path.clone(), - image_caption: None, - table_body: row.table_body.clone(), - table_caption: None, - } - }) - .collect(); - Ok(content_list) + let content_list = existing_rows + .iter() + .map(|row| { + let typ = if row.text.is_some() { + "text".to_string() + } else if row.img_path.is_some() { + "image".to_string() + } else if row.table_body.is_some() { + "table".to_string() + } else { + "unknown".to_string() + }; + + let bbox = row + .bbox + .as_ref() + .and_then(|bbox| serde_json::from_str::>(bbox).ok()) + .unwrap_or_default(); + + ContentItem { + typ, + bbox, + page_idx: row.page_idx, + text: row.text.clone(), + text_level: row.text_level, + text_format: None, + img_path: row.img_path.clone(), + image_caption: None, + table_body: row.table_body.clone(), + table_caption: None, + } + }) + .collect(); + Ok(content_list) }) .await?; if is_image { let removed = retain_one_image_block_for_uploaded_image(&mut content_list); 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 { - // 没有数据，调用 MinerU API - let mineru_result = self.call_mineru_api(file, is_image, timing.as_deref_mut()).await?; - mineru_images = mineru_result.images.clone(); + // 没有数据，调用 MinerU API + let mineru_result = self.call_mineru_api(file, is_image, timing.as_deref_mut()).await?; + mineru_images = mineru_result.images.clone(); content_list = timed_step_opt(timing.as_deref_mut(), "parse_mineru_content_list", async { Ok(serde_json::from_str(&mineru_result.content_list)?) }) @@ -2179,2038 +2185,2047 @@ impl FileProcessor { if is_image { let removed = retain_one_image_block_for_uploaded_image(&mut content_list); 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? { - return Ok(()); - } - - // 提取文本内容并过滤掉 discarded 项 - let valid_content_items: Vec = - content_list.iter().filter(|item| item.typ != "discarded").cloned().collect(); - - // 只有在有有效内容时才插入数据库 - if !valid_content_items.is_empty() { - timed_step_opt(timing.as_deref_mut(), "write_pdf_contents", async { - let rows = Self::content_items_to_pdf_rows(&valid_content_items); - crate::pdf_content::write(file.id, &rows).await - }) - .await?; - } - - // 保存图片到本地 - timed_step_opt(timing.as_deref_mut(), "save_mineru_images", async { - let cfg = config::get(); - fs::create_dir_all(&cfg.storage.images_path).await?; - info!("image count: {}", mineru_result.images.len()); - info!("images: {:?}", mineru_result.images.keys()); - for (img_name, img_base64) in &mineru_result.images { - // 保存图片，如果以 data:image/jpeg;base64, 开头就去掉，没有也不报错 - let base64_marker = "base64,"; - let (prefix, payload) = match img_base64.find(base64_marker) { - Some(idx) => { - (&img_base64[..idx + base64_marker.len()], &img_base64[idx + base64_marker.len()..]) - } - None => ("(raw)", img_base64.as_str()), - }; - let preview: String = payload.chars().take(32).collect(); - debug!( - "Decoding mineru image {} for file {} (prefix={}, len={}, preview=\"{}\")", - img_name, - file.id, - prefix, - payload.len(), - preview - ); - let bytes = STANDARD.decode(payload).map_err(|err| { - error!( - "Failed to decode mineru image {} for file {} (prefix={}, len={}, preview=\"{}\"): {}", - img_name, - file.id, - prefix, - payload.len(), - preview, - err - ); - anyhow::anyhow!(err) - })?; - fs::write(format!("{}/{}", cfg.storage.images_path, img_name), bytes).await?; - } - Ok(()) - }) - .await?; - } - - // 图片文本化：解析 MinerU/历史图片，持久化描述并 enrich content_list。 - let image_names: Vec = content_list - .iter() - .filter_map(|item| if item.typ == "image" { item.img_path.clone() } else { None }) - .collect(); - let mut desc_map = if existing_rows.is_empty() { - parse_images_to_descriptions( - &self.pool, - file.id, - &mineru_images, - "mineru", - (!is_image).then_some(file.filename.as_str()), - ) - .await? - } else { - load_or_parse_image_descriptions(&self.pool, file.id, &image_names, "existing").await? - }; - - // 上传的图片文件本身也需要被文本化，并生成一个图片切片。 - if is_image && config::get().services.image_parse_url.is_some() && !file.filename.is_empty() { - if !desc_map.contains_key(&file.filename) { - match image_parse::parse_image_file(Path::new(&file.path), &file.filename, None).await { - Ok(resp) => { - if let Err(err) = image_description::save( - &self.pool, - file.id, - &file.filename, - &resp.description, - &resp.raw_response, - "upload", - resp.jpg_filename.as_deref(), - ) - .await - { - warn!("Failed to save upload image description for {}: {}", file.filename, err); - } - desc_map.insert(file.filename.clone(), resp.description); - } - Err(err) => warn!("Failed to parse upload image {}: {}", file.filename, err), - } - } + return Ok(()); + } + + // 提取文本内容并过滤掉 discarded 项 + let valid_content_items: Vec = + content_list.iter().filter(|item| item.typ != "discarded").cloned().collect(); + + // 只有在有有效内容时才插入数据库 + if !valid_content_items.is_empty() { + timed_step_opt(timing.as_deref_mut(), "write_pdf_contents", async { + let rows = Self::content_items_to_pdf_rows(&valid_content_items); + crate::pdf_content::write(file.id, &rows).await + }) + .await?; + } + + // 保存图片到本地 + timed_step_opt(timing.as_deref_mut(), "save_mineru_images", async { + let cfg = config::get(); + fs::create_dir_all(&cfg.storage.images_path).await?; + info!("image count: {}", mineru_result.images.len()); + info!("images: {:?}", mineru_result.images.keys()); + for (img_name, img_base64) in &mineru_result.images { + // 保存图片，如果以 data:image/jpeg;base64, 开头就去掉，没有也不报错 + let base64_marker = "base64,"; + let (prefix, payload) = match img_base64.find(base64_marker) { + Some(idx) => { + (&img_base64[..idx + base64_marker.len()], &img_base64[idx + base64_marker.len()..]) + } + None => ("(raw)", img_base64.as_str()), + }; + let preview: String = payload.chars().take(32).collect(); + debug!( + "Decoding mineru image {} for file {} (prefix={}, len={}, preview=\"{}\")", + img_name, + file.id, + prefix, + payload.len(), + preview + ); + let bytes = STANDARD.decode(payload).map_err(|err| { + error!( + "Failed to decode mineru image {} for file {} (prefix={}, len={}, preview=\"{}\"): {}", + img_name, + file.id, + prefix, + payload.len(), + preview, + err + ); + anyhow::anyhow!(err) + })?; + fs::write(format!("{}/{}", cfg.storage.images_path, img_name), bytes).await?; + } + Ok(()) + }) + .await?; + } + + // 图片文本化：解析 MinerU/历史图片，持久化描述并 enrich content_list。 + let image_names: Vec = content_list + .iter() + .filter_map(|item| if item.typ == "image" { item.img_path.clone() } else { None }) + .collect(); + let mut desc_map = if existing_rows.is_empty() { + parse_images_to_descriptions( + &self.pool, + file.id, + &mineru_images, + "mineru", + (!is_image).then_some(file.filename.as_str()), + ) + .await? + } else { + load_or_parse_image_descriptions(&self.pool, file.id, &image_names, "existing").await? + }; + + // 上传的图片文件本身也需要被文本化，并生成一个图片切片。 + if is_image && config::get().services.image_parse_url.is_some() && !file.filename.is_empty() { + if !desc_map.contains_key(&file.filename) { + match image_parse::parse_image_file(Path::new(&file.path), &file.filename, None).await { + Ok(resp) => { + if let Err(err) = image_description::save( + &self.pool, + file.id, + &file.filename, + &resp.description, + &resp.raw_response, + "upload", + resp.jpg_filename.as_deref(), + ) + .await + { + warn!("Failed to save upload image description for {}: {}", file.filename, err); + } + desc_map.insert(file.filename.clone(), resp.description); + } + Err(err) => warn!("Failed to parse upload image {}: {}", file.filename, err), + } + } // MinerU 已经提取到图片时，不能再把上传原图追加一次；否则一个 BMP 会同时 // 产生 MinerU 的派生 JPG 引用和原始 BMP 引用，造成重复切片。 if !content_list.iter().any(|item| item.typ == "image") { - let caption = desc_map.get(&file.filename).cloned().unwrap_or_default(); - content_list.push(ContentItem { - typ: "image".to_string(), - bbox: vec![], - page_idx: 0, - text: None, - text_level: None, - text_format: None, - img_path: Some(file.filename.clone()), - image_caption: if caption.is_empty() { - None - } else { - Some(vec![caption]) - }, - table_body: None, - table_caption: None, - }); - } - } - - enrich_content_list_with_image_descriptions(&mut content_list, &desc_map); - - // 构建全文与切片均为 CPU 密集操作，合并放到阻塞线程池执行，避免阻塞异步运行时。 - // content_list 在此之后不再使用，直接移入闭包。 - let slice_type = file.slice_type.clone(); - let (full_content, slices) = timed_step_opt(timing.as_deref_mut(), "slice_build", async { - tokio::task::spawn_blocking(move || -> anyhow::Result<(String, Vec)> { - let (full_content, full_segments) = Self::build_full_content_and_segments(&content_list); - let slices = if slice_type == "smart" || slice_type.is_empty() { - // 智能切片：使用 content_list - Self::smart_slice_content_with_positions(&content_list)? - } else if slice_type == "fixed" { - Self::fixed_slice_content_with_positions(&full_content, &full_segments)? - } else { - Self::slice_content(&full_content, &slice_type)? - .into_iter() - .map(|content| { - let is_image = content_looks_like_image_reference(&content); - SliceWithPositions { content, positions: vec![], is_image } - }) - .collect() - }; - Ok((full_content, slices)) - }) - .await - .map_err(|e| anyhow::anyhow!("slice task failed: {e}"))? - }) - .await?; - - if !self.ensure_file_exists(file.id, "before writing slices").await? { - return Ok(()); - } - - let slice_count = slices.len(); - let embeddings = vec![image_embedding.clone(); slice_count]; - self.finish_file_processing( - file, - slices, - embeddings, - &full_content, - "PDF processed successfully", - index_filename, - None, - timing, - ) - .await - } - - async fn call_mineru_api( - &self, file: &File, is_image: bool, mut timing: Option<&mut ParseTimingCtx>, - ) -> anyhow::Result { - let cfg = config::get(); - - if !is_image && cfg.services.mineru_max_pages > 0 { - // lopdf 解析整个 PDF 是同步阻塞操作（仅用于获取页数），放到阻塞线程池 - let path = file.path.clone(); - let page_count = - tokio::task::spawn_blocking(move || Document::load(&path).map(|doc| doc.get_pages().len())).await?; - match page_count { - Ok(total_pages) => { - if total_pages > cfg.services.mineru_max_pages { - return timed_step_opt(timing.as_deref_mut(), "mineru_api_in_ranges", async { - self.call_mineru_api_in_ranges(file, total_pages, cfg.services.mineru_max_pages).await - }) - .await; - } - } - Err(e) => { - warn!( - "Failed to read PDF page count for {}: {}, falling back to single MinerU request", - file.filename, e - ); - } - } - } - - timed_step_opt(timing, "mineru_api", async { - self.call_mineru_api_with_path(&file.path, &file.filename, is_image, None, None).await - }) - .await - } - - async fn call_mineru_api_in_ranges( - &self, file: &File, total_pages: usize, max_pages: usize, - ) -> anyhow::Result { - let total_parts = total_pages.div_ceil(max_pages); - info!( - "PDF {} has {} pages, calling MinerU in {} ranges (max {} pages per range)", - file.filename, total_pages, total_parts, max_pages - ); - - let mut merged_items: Vec = Vec::new(); - let mut merged_images: HashMap = HashMap::new(); - - for (part_index, range_start) in (0..total_pages).step_by(max_pages).enumerate() { - let range_end = std::cmp::min(range_start + max_pages, total_pages) - 1; - let chunk_filename = format!("{}_part_{}.pdf", file.filename, part_index + 1); - - let chunk_result = self - .call_mineru_api_with_path(&file.path, &chunk_filename, false, Some(range_start), Some(range_end)) - .await?; - let mut chunk_items: Vec = serde_json::from_str(&chunk_result.content_list)?; - - let image_prefix = format!("part{}_", part_index + 1); - let min_page_idx = chunk_items.iter().map(|item| item.page_idx).min(); - let needs_offset = min_page_idx.is_some_and(|min| min < range_start as i32); - for item in &mut chunk_items { - if needs_offset { - item.page_idx += range_start as i32; - } - if let Some(img_path) = item.img_path.as_deref() { - item.img_path = Some(Self::prefix_image_path(img_path, &image_prefix)); - } - } - - for (img_name, img_base64) in chunk_result.images { - let prefixed = Self::prefix_image_path(&img_name, &image_prefix); - merged_images.insert(prefixed, img_base64); - } - - merged_items.extend(chunk_items); - } - - let content_list = serde_json::to_string(&merged_items)?; - Ok(Result { content_list, images: merged_images }) - } - - fn prefix_image_path(img_path: &str, prefix: &str) -> String { - match img_path.rsplit_once('/') { - Some((dir, name)) => format!("{}/{}{}", dir, prefix, name), - None => format!("{}{}", prefix, img_path), - } - } - - async fn call_mineru_api_with_path( - &self, file_path: &str, filename: &str, is_image: bool, start_page: Option, end_page: Option, - ) -> anyhow::Result { - let cfg = config::get(); - let mineru_url = cfg.services.mineru_url.trim_end_matches('/'); - - let client = self.services_http_client()?; - if mineru_url.ends_with("/file_parse") { - // 构建 multipart form - let mime_type = if is_image { - mime_guess::from_path(filename).first_or_octet_stream().essence_str().to_string() - } else { - "application/pdf".to_string() - }; - - let file_part = stream_file_part(file_path, filename, &mime_type).await?; - let mut form = multipart::Form::new() - .text("return_middle_json", "false") - .text("return_model_output", "false") - .text("return_md", "false") - .text("return_images", "true") - .text("parse_method", "auto") - .text("lang_list", "ch") - .text("output_dir", "") - .text("server_url", "string") - .text("return_content_list", "true") - .text("backend", "pipeline") - .text("table_enable", "true") - .text("response_format_zip", "false") - .text("formula_enable", "true") - .part("files", file_part); - let start_page_id = start_page.unwrap_or(0); - let end_page_id = end_page.map(|v| v.to_string()).unwrap_or_else(|| "99999".to_string()); - form = form.text("start_page_id", start_page_id.to_string()).text("end_page_id", end_page_id); - - // 调用 MinerU API - let response = client.post(mineru_url).multipart(form).send().await?; - let status = response.status(); - let body_bytes = response.bytes().await?; - if !status.is_success() { - let error_text = String::from_utf8_lossy(&body_bytes); - return Err(anyhow::anyhow!("MinerU API failed: {}", error_text)); - } - - let mineru_response: MinerUResponse = serde_json::from_slice(&body_bytes).map_err(|e| { - let body_text = String::from_utf8_lossy(&body_bytes); - anyhow::anyhow!("MinerU API response decode failed: {} - {}", e, body_text) - })?; - let mineru_result = match mineru_response.results { - MinerUResults::Map(results) => { - results.into_values().next().ok_or_else(|| anyhow::anyhow!("MinerU API returned empty results"))? - } - MinerUResults::List(results) => { - if let Some(item) = results.iter().find(|item| item.status == "error") { - let mut error_msg = item.error.clone(); - if error_msg.is_empty() { - error_msg = "MinerU API returned error result".to_string(); - } - if !item.filename.is_empty() { - error_msg = format!("{} (file: {})", error_msg, item.filename); - } - return Err(anyhow::anyhow!("MinerU API failed: {}", error_msg)); - } - let first = results - .into_iter() - .next() - .ok_or_else(|| anyhow::anyhow!("MinerU API returned empty results"))?; - Result { content_list: first.content_list, images: first.images } - } - }; - return Ok(mineru_result); - } - - let mut url = reqwest::Url::parse(mineru_url)?; - { - let mut query = url.query_pairs_mut(); - query.append_pair("backend", "pipeline"); - query.append_pair("method", "auto"); - query.append_pair("lang", "ch"); - query.append_pair("start_page", &start_page.unwrap_or(0).to_string()); - if let Some(end_page) = end_page { - query.append_pair("end_page", &end_page.to_string()); - } - query.append_pair("formula_enable", "true"); - query.append_pair("table_enable", "true"); - } - - let mime_type = mime_guess::from_path(filename).first_or_octet_stream().essence_str().to_string(); - let file_part = stream_file_part(file_path, filename, &mime_type).await?; - let form = multipart::Form::new().part("file", file_part); - - let response = client.post(url.clone()).multipart(form).send().await?; - if !response.status().is_success() { - let error_text = response.text().await?; - return Err(anyhow::anyhow!("MinerU API failed: {}", error_text)); - } - - let analyze_response: AnalyzePdfResponse = response.json().await?; - if analyze_response.code != 200 { - return Err(anyhow::anyhow!("MinerU API failed: {}", analyze_response.message)); - } - - let Some(data) = analyze_response.data else { - return Err(anyhow::anyhow!("MinerU API returned empty data")); - }; - - let host = url.host_str().ok_or_else(|| anyhow::anyhow!("MinerU API URL missing host"))?; - let origin = if let Some(port) = url.port() { - format!("{}://{}:{}", url.scheme(), host, port) - } else { - format!("{}://{}", url.scheme(), host) - }; - - let mut content_list = data.content_list; - - // 去重收集需要下载的图片（filename -> url） - let mut to_download: HashMap = HashMap::new(); - for item in &content_list { - let Some(img_path) = item.img_path.as_deref() else { continue }; - let filename = std::path::Path::new(img_path) - .file_name() - .and_then(|name| name.to_str()) - .unwrap_or(img_path) - .to_string(); - to_download - .entry(filename) - .or_insert_with(|| format!("{}/output/{}", origin, img_path.trim_start_matches('/'))); - } - - // 并发下载图片（限制并发数，避免打爆 MinerU），各图片之间无依赖。 - use futures::stream::{StreamExt, TryStreamExt}; - let images: HashMap = - futures::stream::iter(to_download.into_iter().map(|(filename, image_url)| { - let client = client.clone(); - async move { - let image_response = client.get(&image_url).send().await?; - if !image_response.status().is_success() { - let error_text = image_response.text().await?; - return Err(anyhow::anyhow!("MinerU image download failed: {}", error_text)); - } - let image_bytes = image_response.bytes().await?; - let image_mime = mime_guess::from_path(&filename).first_or_octet_stream().essence_str().to_string(); - let image_base64 = STANDARD.encode(image_bytes); - anyhow::Ok((filename, format!("data:{};base64,{}", image_mime, image_base64))) - } - })) - .buffer_unordered(8) - .try_collect() - .await?; - - // 回填 img_path 为去掉目录的文件名 - for item in &mut content_list { - if let Some(img_path) = item.img_path.as_deref() { - let filename = std::path::Path::new(img_path) - .file_name() - .and_then(|name| name.to_str()) - .unwrap_or(img_path) - .to_string(); - item.img_path = Some(filename); - } - } - - let content_list_json = serde_json::to_string(&content_list)?; - Ok(Result { content_list: content_list_json, images }) - } - - /// 流式读取文本文件，边读边分片，避免一次性把整个文件载入内存。 - async fn stream_text_file(path: &str, slice_type: &str) -> anyhow::Result<(String, Vec)> { - use tokio::io::BufReader; - - let file = tokio::fs::File::open(path).await?; - let mut reader = BufReader::new(file); - let mut full_content = String::new(); - let mut slices = Vec::new(); - - match slice_type { - "paragraph" => { - let mut current = String::new(); - loop { - let mut line = String::new(); - let n = reader.read_line(&mut line).await?; - if n == 0 { - break; - } - - let (content, term) = if line.ends_with("\r\n") { - let c = line.drain(..line.len() - 2).collect::(); - (c, "\r\n") - } else if line.ends_with('\n') { - let c = line.drain(..line.len() - 1).collect::(); - (c, "\n") - } else { - (line, "") - }; - - full_content.push_str(&content); - full_content.push_str(term); - - if content.trim().is_empty() { - if !current.trim().is_empty() { - slices.push(current.trim().to_string()); - current.clear(); - } - } else { - if !current.is_empty() { - current.push_str(term); - } - current.push_str(&content); - } - } - if !current.trim().is_empty() { - slices.push(current.trim().to_string()); - } - } - "sentence" => { - const SENTENCE_END: [char; 6] = ['。', '.', '?', '!', '？', '！']; - let mut buf = String::new(); - loop { - let mut line = String::new(); - let n = reader.read_line(&mut line).await?; - if n == 0 { - break; - } - - let (content, term) = if line.ends_with("\r\n") { - let c = line.drain(..line.len() - 2).collect::(); - (c, "\r\n") - } else if line.ends_with('\n') { - let c = line.drain(..line.len() - 1).collect::(); - (c, "\n") - } else { - (line, "") - }; - - full_content.push_str(&content); - full_content.push_str(term); - - for ch in content.chars().chain(term.chars()) { - if SENTENCE_END.contains(&ch) { - if !buf.trim().is_empty() { - slices.push(buf.trim().to_string()); - } - buf.clear(); - } else { - buf.push(ch); - } - } - } - if !buf.trim().is_empty() { - slices.push(buf.trim().to_string()); - } - } - _ => { - let cfg = config::get(); - let chunk_size = cfg.slice.smart_slice_max_chars; - let overlap = cfg.slice.fixed_slice_overlap_chars; - let mut buf: Vec = Vec::new(); - - loop { - let mut line = String::new(); - let n = reader.read_line(&mut line).await?; - if n == 0 { - break; - } - - let (content, term) = if line.ends_with("\r\n") { - let c = line.drain(..line.len() - 2).collect::(); - (c, "\r\n") - } else if line.ends_with('\n') { - let c = line.drain(..line.len() - 1).collect::(); - (c, "\n") - } else { - (line, "") - }; - - full_content.push_str(&content); - full_content.push_str(term); - - for ch in content.chars().chain(term.chars()) { - buf.push(ch); - if buf.len() >= chunk_size { - let slice: String = buf[..chunk_size].iter().collect(); - slices.push(slice); - if overlap >= chunk_size || overlap == 0 { - buf.clear(); - } else { - let keep = buf.len() - overlap; - buf.drain(0..keep); - } - } - } - } - - while !buf.is_empty() { - let end = buf.len().min(chunk_size); - let slice: String = buf[..end].iter().collect(); - slices.push(slice); - if end == buf.len() { - break; - } - if overlap >= end || overlap == 0 { - break; - } - let keep = buf.len() - overlap; - buf.drain(0..keep); - } - } - } - - Ok((full_content, slices)) - } - - /// 处理普通文本文件 - async fn process_text_file(&self, file: &File, mut timing: Option<&mut ParseTimingCtx>) -> anyhow::Result<()> { - if !self.ensure_file_exists(file.id, "before reading text").await? { - return Ok(()); - } - let (content, slices) = timed_step_opt(timing.as_deref_mut(), "read_text_file", async { - Self::stream_text_file(&file.path, &file.slice_type).await - }) - .await?; - self.process_plain_text_content(file, content, slices, "Processing completed successfully", timing).await - } - - async fn process_audio_file(&self, file: &File, mut timing: Option<&mut ParseTimingCtx>) -> anyhow::Result<()> { - info!("Processing audio file: {}", file.filename); - - if !self.ensure_file_exists(file.id, "before reading audio").await? { - return Ok(()); - } - let mime_type = mime_guess::from_path(&file.filename).first_or_octet_stream().essence_str().to_string(); - let file_part = timed_step_opt( - timing.as_deref_mut(), - "open_audio_file", - stream_file_part(&file.path, &file.filename, &mime_type), - ) - .await?; - - let form = multipart::Form::new().part("file", file_part); - - let client = self.services_http_client()?; - let cfg = config::get(); - let mut req_builder = client.post(&cfg.services.audio_transcription_url).multipart(form); - if let Some(key) = &cfg.services.audio_transcription_key - && !key.is_empty() - { - req_builder = req_builder.header("Authorization", format!("Bearer {}", key)); - } - - let response = - timed_step_opt(timing.as_deref_mut(), "audio_transcription_api", async { Ok(req_builder.send().await?) }) - .await?; - if !response.status().is_success() { - let error_text = response.text().await?; - return Err(anyhow::anyhow!("Audio transcription API failed: {}", error_text)); - } - - let AudioTranscriptionResponse { text, language } = response.json().await?; - let log_message = if language.is_empty() { - "Audio processed successfully".to_string() - } else { - format!("Audio processed successfully (language: {})", language) - }; - - let text = timed_step_opt(timing.as_deref_mut(), "audio_transcription", async { Ok(text) }).await?; - let slices = timed_step_opt(timing.as_deref_mut(), "slice_build", async { - let content_owned = text.clone(); - let slice_type = file.slice_type.clone(); - tokio::task::spawn_blocking(move || Self::slice_content(&content_owned, &slice_type)) - .await - .map_err(|e| anyhow::anyhow!("slice_content task failed: {e}"))? - }) - .await?; - - self.process_plain_text_content(file, text, slices, &log_message, timing).await - } - - async fn process_plain_text_content( - &self, file: &File, content: String, slices: Vec, log_message: &str, - timing: Option<&mut ParseTimingCtx>, - ) -> anyhow::Result<()> { - let wrapped: Vec = slices - .into_iter() - .map(|content| { - let is_image = content_looks_like_image_reference(&content); - SliceWithPositions { content, positions: vec![], is_image } - }) - .collect(); - let embeddings = vec![None; wrapped.len()]; - self.finish_file_processing(file, wrapped, embeddings, &content, log_message, None, None, timing).await - } - - /// 标记文件处理失败 - async fn mark_file_failed(&self, file_id: i64, error_msg: &str) -> anyhow::Result<()> { - let sql = "UPDATE files SET status = -1, parse_run_id = NULL, log = ?, updated_at = strftime('%s','now') WHERE id = ?"; - sqlx::query(sql).bind(error_msg).bind(file_id).execute(&self.pool).await?; - Ok(()) - } - - async fn mark_file_storage_skipped(&self, file_id: i64) -> anyhow::Result<()> { - let sql = - "UPDATE files SET status = 3, parse_run_id = NULL, log = ?, updated_at = strftime('%s','now') WHERE id = ?"; - sqlx::query(sql).bind("Storage mode: not parsed").bind(file_id).execute(&self.pool).await?; - Ok(()) - } - - #[allow(clippy::too_many_arguments)] - async fn finish_file_processing( - &self, file: &File, slices: Vec, embeddings: Vec>>>, - full_content: &str, log_message: &str, index_name_override: Option<&str>, summary: Option<&str>, - mut timing: Option<&mut ParseTimingCtx>, - ) -> anyhow::Result<()> { - if !self.ensure_file_exists(file.id, "before writing slices").await? { - return Ok(()); - } - let slice_count = slices.len(); - anyhow::ensure!( - embeddings.len() == slice_count, - "embeddings count {} does not match slice count {}", - embeddings.len(), - slice_count - ); - - let parse_run_id: String = sqlx::query_scalar( - "SELECT parse_run_id FROM files WHERE id = ? AND status = 2 AND parse_run_id IS NOT NULL", - ) - .bind(file.id) - .fetch_one(&self.pool) - .await?; - - let (search_docs, search_embeddings) = timed_step_opt(timing.as_deref_mut(), "insert_slices", async { - let persisted = self.insert_slices_and_positions(file.id, &parse_run_id, slices).await?; - let mut docs = Vec::with_capacity(persisted.len()); - let mut embs = Vec::with_capacity(persisted.len()); - for (idx, (id, content, is_image)) in persisted.into_iter().enumerate() { - docs.push(tantivy_engine::Document::new(id, file.id, file.kb_id, content).with_is_image(is_image)); - embs.push(embeddings.get(idx).cloned().flatten()); - } - Ok((docs, embs)) - }) - .await?; - - if !search_docs.is_empty() { - timed_step_opt(timing.as_deref_mut(), "write_search_batch", async { - self.search_engine.write_batch(search_docs, search_embeddings).await?; - Ok(()) - }) - .await?; - } - - if !self.ensure_file_exists(file.id, "before writing full index").await? { - return Ok(()); - } - let index_name = index_name_override.unwrap_or(&file.filename); - let index_full_content = format!("{}\n\n{}", index_name, full_content); - timed_step_opt(timing.as_deref_mut(), "write_full_index", async { - self.search_engine - .write_full(tantivy_engine::Document::new(file.id, file.id, file.kb_id, index_full_content)) - .await?; - Ok(()) - }) - .await?; - - if !self.ensure_file_exists(file.id, "before updating status").await? { - return Ok(()); - } - timed_step_opt(timing.as_deref_mut(), "publish_parse_artifact", async { - let artifact_id = self.ensure_parse_artifact(file, full_content, summary).await?; - self.adopt_existing_artifact_if_needed(file, artifact_id, full_content, summary).await?; - Ok(()) - }) - .await?; - let normalized_summary = summary.and_then(|value| { - let trimmed = value.trim(); - if trimmed.is_empty() { None } else { Some(trimmed.to_string()) } - }); - timed_step_opt(timing.as_deref_mut(), "write_summary_index", async { - self.persist_file_summary(file.id, file.kb_id, normalized_summary.as_deref()).await - }) - .await?; - timed_step_opt(timing.as_deref_mut(), "finalize_file_status", async { - let sql = "UPDATE files SET status = 1, parse_run_id = NULL, log = ?, summary = ?, updated_at = strftime('%s','now') \ - WHERE id = ? AND status = 2 AND parse_run_id = ?"; - let updated = sqlx::query(sql) - .bind(log_message) - .bind(normalized_summary.as_deref()) - .bind(file.id) - .bind(&parse_run_id) - .execute(&self.pool) - .await? - .rows_affected(); - anyhow::ensure!(updated == 1, "parse run for file {} is no longer current", file.id); - crate::file_content::write(file.id, full_content).await?; - Ok(()) - }) - .await?; - - info!("File {} processed successfully with {} slices", file.id, slice_count); - - self.search_engine.reload_readers()?; - - timed_step_opt(timing, "build_knowledge_graph", async { - self.maybe_build_knowledge_graph(file).await; - Ok(()) - }) - .await?; - - Ok(()) - } - - /// 返回每个 slice 的 (id, content, is_image) 供后续构建搜索文档使用。 - async fn insert_slices_and_positions( - &self, file_id: i64, parse_run_id: &str, slices: Vec, - ) -> anyhow::Result> { - if slices.is_empty() { - return Ok(Vec::new()); - } - // 每批最多 500 条 slice 一个事务，避免长时间持有 SQLite 写锁阻塞其他写操作 - const SLICE_TX_BATCH: usize = 500; - let mut persisted: Vec<(i64, String, bool)> = Vec::with_capacity(slices.len()); - for (chunk_idx, slice_chunk) in slices.chunks(SLICE_TX_BATCH).enumerate() { - let mut tx = self.pool.begin().await?; - let mut chunk_persisted: Vec<(i64, String, bool)> = Vec::with_capacity(slice_chunk.len()); - let mut position_rows: Vec<(i64, SlicePosition)> = Vec::new(); - let binds_per_row = 4_usize; - let max_vars = 999_usize; - let batch_size = std::cmp::max(1, max_vars / binds_per_row); - for (insert_chunk_idx, insert_chunk) in slice_chunk.chunks(batch_size).enumerate() { - let insert_offset = insert_chunk_idx * batch_size; - let mut slice_sql = - QueryBuilder::::new("INSERT INTO slices (file_id, parse_run_id, ordinal, is_image) "); - slice_sql.push_values(insert_chunk.iter().enumerate(), |mut b, (idx, slice)| { - let ordinal = chunk_idx * SLICE_TX_BATCH + insert_offset + idx; - b.push_bind(file_id) - .push_bind(parse_run_id) - .push_bind(ordinal as i64) - .push_bind(i64::from(slice.is_image)); - }); - slice_sql.push( - " ON CONFLICT(file_id, parse_run_id, ordinal) WHERE parse_run_id IS NOT NULL AND ordinal IS NOT NULL \ - DO UPDATE SET updated_at = strftime('%s','now') RETURNING id", - ); - - let inserted_ids: Vec<(i64,)> = slice_sql.build_query_as().fetch_all(&mut *tx).await?; - anyhow::ensure!( - inserted_ids.len() == insert_chunk.len(), - "inserted slice row count mismatch: expected {}, got {}", - insert_chunk.len(), - inserted_ids.len() - ); - - for (slice, (id,)) in insert_chunk.iter().zip(inserted_ids) { - for position in &slice.positions { - position_rows.push((id, position.clone())); - } - chunk_persisted.push((id, slice.content.clone(), slice.is_image)); - } - } - if !position_rows.is_empty() { - let binds_per_row = 8_usize; - let max_vars = 999_usize; - let batch_size = std::cmp::max(1, max_vars / binds_per_row); - for chunk in position_rows.chunks(batch_size) { - let mut pos_sql = QueryBuilder::::new( - "insert into slice_positions(slice_id, page_idx, x1, y1, x2, y2, sheet_name, row_num) ", - ); - pos_sql.push_values(chunk.iter(), |mut b, (slice_id, position)| { - b.push_bind(slice_id) - .push_bind(position.page_idx) - .push_bind(position.bbox[0]) - .push_bind(position.bbox[1]) - .push_bind(position.bbox[2]) - .push_bind(position.bbox[3]) - .push_bind(&position.sheet_name) - .push_bind(position.row_num); - }); - pos_sql.build().execute(&mut *tx).await?; - } - } - tx.commit().await?; - crate::slice_content::upsert_many( - file_id, - &chunk_persisted.iter().map(|(id, content, _)| (*id, content.clone())).collect::>(), - ) - .await?; - persisted.extend(chunk_persisted); - } - Ok(persisted) - } - - async fn ensure_file_exists(&self, file_id: i64, stage: &str) -> anyhow::Result { - let exists = sqlx::query_scalar::<_, i64>("SELECT 1 FROM files WHERE id = ? LIMIT 1") - .bind(file_id) - .fetch_optional(&self.pool) - .await? - .is_some(); - if !exists { - info!("File {} no longer exists during {}, skipping further processing", file_id, stage); - } - Ok(exists) - } - - async fn is_storage_kb(&self, kb_id: Option) -> anyhow::Result { - let Some(kb_id) = kb_id else { return Ok(false) }; - let kb_type: Option = sqlx::query_scalar("SELECT kb_type FROM knowledge_bases WHERE id = ?") - .bind(kb_id) - .fetch_optional(&self.pool) - .await?; - Ok(matches!(kb_type.as_deref(), Some("storage"))) - } - - fn is_audio_file(filename_lower: &str) -> bool { - filename_lower.ends_with(".wav") - || filename_lower.ends_with(".mp3") - || filename_lower.ends_with(".m4a") - || filename_lower.ends_with(".aac") - || filename_lower.ends_with(".flac") - || filename_lower.ends_with(".ogg") - || filename_lower.ends_with(".opus") - || filename_lower.ends_with(".wma") - || filename_lower.ends_with(".amr") - || filename_lower.ends_with(".aiff") - || filename_lower.ends_with(".aif") - || filename_lower.ends_with(".alac") - || filename_lower.ends_with(".webm") - } - - /// 根据 slice_type 对内容进行分片 - fn slice_content(content: &str, slice_type: &str) -> anyhow::Result> { - match slice_type { - "paragraph" => { - // 按段落分片（以双换行符分隔） - Ok(content.split("\n\n").map(|s| s.trim()).filter(|s| !s.is_empty()).map(|s| s.to_string()).collect()) - } - "sentence" => { - // 按句子分片（简单实现：以句号、问号、感叹号分隔） - Ok(content - .split(['。', '.', '?', '!', '？', '！']) - .map(|s| s.trim()) - .filter(|s| !s.is_empty()) - .map(|s| s.to_string()) - .collect()) - } - _ => { - // 固定长度分片（每8000字符，重叠100字符） - let cfg = config::get(); - let chunk_size = cfg.slice.smart_slice_max_chars; - let overlap = cfg.slice.fixed_slice_overlap_chars; - - let chars: Vec = content.chars().collect(); - let mut slices = Vec::new(); - let mut start = 0; - - while start < chars.len() { - let end = std::cmp::min(start + chunk_size, chars.len()); - let slice: String = chars[start..end].iter().collect(); - slices.push(slice); - - if end >= chars.len() { - break; - } - - // 下一个切片从 end - overlap 开始，以实现重叠 - start = end.saturating_sub(overlap); - if start >= end { - break; - } - } - - Ok(slices) - } - } - } - - fn build_full_content_and_segments(content_list: &[ContentItem]) -> (String, Vec) { - let mut full_content = String::new(); - let mut segments = Vec::new(); - let mut current_len = 0usize; - - for pdf_content in content_list { - if pdf_content.typ == "discarded" { - continue; - } - let mut item_content = String::new(); - if pdf_content.typ == "text" { - if let Some(text) = &pdf_content.text { - if let Some(lv) = pdf_content.text_level { - item_content.push_str("#".repeat(lv as usize).as_str()); - item_content.push(' '); - } - item_content.push_str(text); - } - } else if pdf_content.typ == "image" { - if let Some(img_name) = &pdf_content.img_path { - item_content.push_str(&format!("![{}](/api/v1/knowledge/files/{})", img_name, img_name)); - } - if let Some(captions) = &pdf_content.image_caption { - for caption in captions { - item_content.push_str(&caption.to_string()); - } - } - } else if pdf_content.typ == "table" { - if let Some(table_caption) = &pdf_content.table_caption { - for caption in table_caption { - item_content.push_str(&caption.to_string()); - } - } - if let Some(table_body) = &pdf_content.table_body { - item_content.push_str(table_body); - } - } - - if item_content.is_empty() { - continue; - } - - let start = current_len; - full_content.push_str(&item_content); - current_len += item_content.chars().count(); - let end = current_len; - - let positions = Self::positions_from_item(pdf_content); - if !positions.is_empty() { - segments.push(Segment { start, end, positions }); - } - - full_content.push_str("\n\n"); - current_len += 2; - } - - (full_content, segments) - } - - fn fixed_slice_content_with_positions( - content: &str, segments: &[Segment], - ) -> anyhow::Result> { - let cfg = config::get(); - let chunk_size = cfg.slice.smart_slice_max_chars; - let overlap = cfg.slice.fixed_slice_overlap_chars; - - let chars: Vec = content.chars().collect(); - let mut slices = Vec::new(); - let mut start = 0; - - while start < chars.len() { - let end = std::cmp::min(start + chunk_size, chars.len()); - let slice: String = chars[start..end].iter().collect(); - let positions = Self::positions_for_range(segments, start, end); - slices.push(SliceWithPositions { - content: slice.clone(), - positions, - is_image: content_looks_like_image_reference(&slice), - }); - - if end >= chars.len() { - break; - } - - start = end.saturating_sub(overlap); - if start >= end { - break; - } - } - - Ok(slices) - } - - fn smart_slice_content_with_positions(content_list: &[ContentItem]) -> anyhow::Result> { - let cfg = config::get(); - let max_chars = cfg.slice.smart_slice_max_chars; - - let mut slices = Vec::new(); - let mut current_slice = String::new(); - let mut current_segments: Vec = Vec::new(); - let mut current_len = 0usize; - let mut current_header = String::new(); // 当前所在的标题 - let mut current_header_positions: Vec = Vec::new(); - - for item in content_list { - if item.typ == "discarded" { - continue; - } - - let mut item_content = String::new(); - let item_positions = Self::positions_from_item(item); - - // 处理文本内容 - if item.typ == "text" { - if let Some(text) = &item.text { - // 如果是标题，更新当前标题 - if let Some(level) = item.text_level { - // 这是一个标题 - let header_text = format!("{} {}", "#".repeat(level as usize), text); - - // 如果当前有累积的内容，先保存 - if !current_slice.trim().is_empty() { - Self::flush_slice_with_positions( - &mut slices, - &mut current_slice, - &mut current_segments, - max_chars, - ); - current_len = 0; - } - - // 更新当前标题 - current_header = header_text; - current_header_positions = item_positions.clone(); - // 开始新的切片，包含标题 - current_slice.clear(); - current_segments.clear(); - Self::append_segment( - &mut current_slice, - &mut current_len, - &mut current_segments, - ¤t_header, - current_header_positions.clone(), - ); - Self::append_separator(&mut current_slice, &mut current_len); - continue; - } else { - // 普通文本 - item_content.push_str(text); - } - } - } else if item.typ == "image" { - if let Some(img_name) = &item.img_path { - item_content.push_str(&format!("![{}](/api/v1/knowledge/files/{})", img_name, img_name)); - } - if let Some(captions) = &item.image_caption { - for caption in captions { - item_content.push_str(&caption.to_string()); - } - } - } else if item.typ == "table" { - if let Some(table_caption) = &item.table_caption { - for caption in table_caption { - item_content.push_str(&caption.to_string()); - } - } - if let Some(table_body) = &item.table_body { - item_content.push_str( - &table_body - .replace(" colspan=\"1\"", "") - .replace(" rowspan=\"1\"", "") - .replace(" colspan='1'", "") - .replace(" rowspan='1'", "") - .replace(" colspan=1", "") - .replace(" rowspan=1", ""), - ); - } - } - - if item_content.is_empty() { - continue; - } - - // 检查加入这个内容后是否超过字数限制 - let test_len = if current_slice.is_empty() { - // 如果当前切片为空，但有标题，先加上标题 - if !current_header.is_empty() { - current_header.chars().count() + 2 + item_content.chars().count() - } else { - item_content.chars().count() - } - } else { - current_len + item_content.chars().count() + 2 - }; - - if test_len > max_chars { - // 超过限制，保存当前切片 - if !current_slice.trim().is_empty() { - Self::flush_slice_with_positions(&mut slices, &mut current_slice, &mut current_segments, max_chars); - current_len = 0; - } - - // 开始新的切片 - current_slice.clear(); - current_segments.clear(); - if !current_header.is_empty() { - Self::append_segment( - &mut current_slice, - &mut current_len, - &mut current_segments, - ¤t_header, - current_header_positions.clone(), - ); - Self::append_separator(&mut current_slice, &mut current_len); - } - Self::append_segment( - &mut current_slice, - &mut current_len, - &mut current_segments, - &item_content, - item_positions.clone(), - ); - Self::append_separator(&mut current_slice, &mut current_len); - } else { - // 没有超过限制，继续累积 - if current_slice.is_empty() { - current_slice.clear(); - current_segments.clear(); - current_len = 0; - if !current_header.is_empty() { - Self::append_segment( - &mut current_slice, - &mut current_len, - &mut current_segments, - ¤t_header, - current_header_positions.clone(), - ); - Self::append_separator(&mut current_slice, &mut current_len); - } - Self::append_segment( - &mut current_slice, - &mut current_len, - &mut current_segments, - &item_content, - item_positions.clone(), - ); - } else { - Self::append_segment( - &mut current_slice, - &mut current_len, - &mut current_segments, - &item_content, - item_positions.clone(), - ); - Self::append_separator(&mut current_slice, &mut current_len); - } - } - } - - // 保存最后的切片 - if !current_slice.trim().is_empty() { - Self::flush_slice_with_positions(&mut slices, &mut current_slice, &mut current_segments, max_chars); - } - - Ok(slices) - } - - fn flush_slice_with_positions( - slices: &mut Vec, current_slice: &mut String, segments: &mut Vec, max_chars: usize, - ) { - if current_slice.trim().is_empty() { - return; - } - let content = std::mem::take(current_slice); - let segment_data = std::mem::take(segments); - let mut new_slices = Self::split_slice_with_positions(&content, &segment_data, max_chars); - slices.append(&mut new_slices); - } - - fn split_slice_with_positions(content: &str, segments: &[Segment], max_chars: usize) -> Vec { - let char_count = content.chars().count(); - if char_count == 0 { - return Vec::new(); - } - - let ranges = if char_count <= max_chars { - vec![(0, char_count)] - } else { - let sentence_ranges = Self::sentence_ranges(content); - if sentence_ranges.is_empty() { - vec![(0, char_count)] - } else { - let mut ranges = Vec::new(); - let mut slice_start = sentence_ranges[0].0; - let mut slice_end = sentence_ranges[0].1; - for (start, end) in sentence_ranges.iter().skip(1) { - if end - slice_start > max_chars && slice_end > slice_start { - ranges.push((slice_start, slice_end)); - slice_start = *start; - } - slice_end = *end; - } - ranges.push((slice_start, slice_end)); - ranges - } - }; - - let chars: Vec = content.chars().collect(); - ranges - .into_iter() - .filter_map(|(start, end)| { - if start >= end || end > chars.len() { - return None; - } - let slice: String = chars[start..end].iter().collect(); - let positions = Self::positions_for_range(segments, start, end); - Some(SliceWithPositions { - content: slice.clone(), - positions, - is_image: content_looks_like_image_reference(&slice), - }) - }) - .collect() - } - - fn sentence_ranges(content: &str) -> Vec<(usize, usize)> { - let mut ranges = Vec::new(); - let mut start = 0usize; - for (idx, ch) in content.chars().enumerate() { - if matches!(ch, '。' | '.' | '?' | '!' | '？' | '！') { - let end = idx + 1; - if end > start { - ranges.push((start, end)); - } - start = end; - } - } - let total = content.chars().count(); - if start < total { - ranges.push((start, total)); - } - ranges - } - - fn append_segment( - current_slice: &mut String, current_len: &mut usize, segments: &mut Vec, text: &str, - positions: Vec, - ) { - if text.is_empty() { - return; - } - let start = *current_len; - current_slice.push_str(text); - *current_len += text.chars().count(); - let end = *current_len; - if !positions.is_empty() { - segments.push(Segment { start, end, positions }); - } - } - - fn append_separator(current_slice: &mut String, current_len: &mut usize) { - current_slice.push_str("\n\n"); - *current_len += 2; - } - - fn positions_from_item(item: &ContentItem) -> Vec { - if item.bbox.len() == 4 { - let bbox = [item.bbox[0], item.bbox[1], item.bbox[2], item.bbox[3]]; - vec![SlicePosition { page_idx: item.page_idx, bbox, sheet_name: None, row_num: None }] - } else { - Vec::new() - } - } - - fn positions_for_range(segments: &[Segment], start: usize, end: usize) -> Vec { - let mut set: HashSet = HashSet::new(); - for segment in segments { - if segment.end > start && segment.start < end { - for position in &segment.positions { - set.insert(position.clone()); - } - } - } - set.into_iter().collect() - } - - async fn clone_file_data(&self, source: &File, target: &File) -> anyhow::Result<()> { - if source.id == target.id { - anyhow::bail!("Source and target file ids are identical"); - } - if source.status != 1 { - anyhow::bail!("Source file {} is not processed", source.id); - } - - // 新结构只复制搜索投影，SQLite 中的切片/PDF 解析结果由 artifact 共享。 - if config::get().server.reuse_duplicate_files { - let source_file_id = effective_parse_file_id(&self.pool, source.id).await?; - let mut source_for_reindex = source.clone(); - let full_content = if let Some(artifact_id) = source.artifact_id { - sqlx::query_scalar::<_, Option>("SELECT full_content FROM parse_artifacts WHERE id = ?") - .bind(artifact_id) - .fetch_optional(&self.pool) - .await? - .flatten() - .unwrap_or_default() - } else { - crate::file_content::read(source_file_id).await?.unwrap_or_default() - }; - let summary = if let Some(artifact_id) = source.artifact_id { - sqlx::query_scalar::<_, Option>("SELECT summary FROM parse_artifacts WHERE id = ?") - .bind(artifact_id) - .fetch_optional(&self.pool) - .await? - .flatten() - } else { - source.summary.clone() - }; - source_for_reindex.content = Some(full_content.clone()); - source_for_reindex.summary = summary.clone(); - let artifact_id = - self.ensure_parse_artifact(&source_for_reindex, &full_content, summary.as_deref()).await?; - let rows = self.fetch_slice_rows(source_file_id).await?; - let shared_slices: Vec = rows - .into_iter() - .map(|row| ClonedSlice { old_id: row.id, new_id: row.id, content: row.content }) - .collect(); - - sqlx::query( - "UPDATE files SET status = 2, artifact_id = ?, log = ?, updated_at = strftime('%s','now') WHERE id = ?", - ) - .bind(artifact_id) - .bind(format!("Reusing shared parse artifact {}", artifact_id)) - .bind(target.id) - .execute(&self.pool) - .await?; - let indexed_content = self.reindex_cloned_slices(target, &shared_slices, &source_for_reindex).await?; - sqlx::query( - "UPDATE files SET status = 1, parse_run_id = NULL, artifact_id = ?, log = ?, summary = ?, \ - updated_at = strftime('%s','now') WHERE id = ?", - ) - .bind(artifact_id) - .bind(format!("Reused shared parse artifact {} from file {}", artifact_id, source_file_id)) - .bind(summary.as_deref()) - .bind(target.id) - .execute(&self.pool) - .await?; - crate::file_content::write(target.id, &indexed_content).await?; - let mut updated_file = target.clone(); - updated_file.artifact_id = Some(artifact_id); - updated_file.content = Some(indexed_content); - updated_file.summary = summary.clone(); - self.maybe_build_knowledge_graph(&updated_file).await; - return Ok(()); - } - - let reuse_log = format!("Reusing parsed data from file {}", source.id); - sqlx::query("UPDATE files SET status = 2, log = ?, updated_at = strftime('%s','now') WHERE id = ?") - .bind(&reuse_log) - .bind(target.id) - .execute(&self.pool) - .await?; - - self.cleanup_processing_file_data_with_retry(target.id, 3).await?; - - let pdf_rows = self.fetch_pdf_content_rows(source.id).await?; - let mut slice_rows = self.fetch_slice_rows(source.id).await?; - let slice_ids: Vec = slice_rows.iter().map(|row| row.id).collect(); - let slice_positions = self.fetch_slice_position_rows(&slice_ids).await?; - let (image_jobs, image_mapping) = self.prepare_image_jobs(&pdf_rows, source.id); - let meta_image_paths = if pdf_rows.is_empty() { - collect_image_raw_paths_for_files(&self.pool, &[source.id]).await? - } else { - Vec::new() - }; - let (meta_image_jobs, meta_image_mapping, target_meta_image_paths) = - Self::prepare_raw_image_jobs(&meta_image_paths, source.id); - - // 兼容历史数据：复用时剥离旧的 f{file_id}_ 图片前缀，并同步重写内容中的引用。 - let mut combined_image_mapping = image_mapping.clone(); - combined_image_mapping.extend(meta_image_mapping.clone()); - if !combined_image_mapping.is_empty() { - for row in &mut slice_rows { - row.content = Self::rewrite_custom_image_refs(&row.content, &combined_image_mapping); - } - } - let mut source_for_reindex = source.clone(); - if !combined_image_mapping.is_empty() - && let Some(content) = crate::file_content::read(source.id).await? - { - source_for_reindex.content = Some(Self::rewrite_custom_image_refs(&content, &combined_image_mapping)); - } - - let mut tx = self.pool.begin().await?; - self.insert_pdf_rows(&mut tx, target.id, &pdf_rows, &image_mapping).await?; - let cloned_slices = self.insert_slice_rows(&mut tx, target.id, &slice_rows).await?; - self.insert_slice_positions(&mut tx, &cloned_slices, &slice_positions).await?; - tx.commit().await?; - - self.copy_image_files(&image_jobs).await?; - self.copy_image_files(&meta_image_jobs).await?; - self.copy_converted_pdf(source.id, target.id).await?; - - if pdf_rows.is_empty() { - update_file_custom_image_meta(&self.pool, target.id, &target_meta_image_paths, "reuse_custom_images") - .await?; - } - - let full_content = self.reindex_cloned_slices(target, &cloned_slices, &source_for_reindex).await?; - - self.search_engine.reload_readers()?; - - let final_log = format!("Reused parsed data from file {}", source.id); - sqlx::query( - "UPDATE files SET status = 1, parse_run_id = NULL, log = ?, summary = ?, updated_at = strftime('%s','now') WHERE id = ?", - ) - .bind(&final_log) - .bind(source.summary.as_deref()) - .bind(target.id) - .execute(&self.pool) - .await?; - crate::file_content::write(target.id, &full_content).await?; - - let mut updated_file = target.clone(); - updated_file.content = Some(full_content.clone()); - updated_file.summary = source.summary.clone(); - self.maybe_build_knowledge_graph(&updated_file).await; - - Ok(()) - } - - async fn fetch_pdf_content_rows(&self, file_id: i64) -> anyhow::Result> { - crate::pdf_content::read(file_id).await - } - - async fn fetch_slice_rows(&self, file_id: i64) -> anyhow::Result> { - let ids: Vec = sqlx::query_scalar("SELECT id FROM slices WHERE file_id = ? ORDER BY id") - .bind(file_id) - .fetch_all(&self.pool) - .await?; - let contents = crate::slice_content::read_all(file_id).await?; - Ok(ids.into_iter().map(|id| SliceRow { id, content: contents.get(&id).cloned().unwrap_or_default() }).collect()) - } - - async fn fetch_slice_position_rows(&self, slice_ids: &[i64]) -> anyhow::Result> { - if slice_ids.is_empty() { - return Ok(Vec::new()); - } - let mut all_rows = Vec::new(); - let chunk_size = 400; - for chunk in slice_ids.chunks(chunk_size) { - let mut qb = QueryBuilder::::new( - "SELECT slice_id, page_idx, x1, y1, x2, y2, sheet_name, row_num FROM slice_positions WHERE slice_id IN (", - ); - let mut separated = qb.separated(", "); - for slice_id in chunk { - separated.push_bind(slice_id); - } - qb.push(") ORDER BY slice_id, id"); - let rows = qb.build_query_as::().fetch_all(&self.pool).await?; - all_rows.extend(rows); - } - Ok(all_rows) - } - - fn prepare_image_jobs( - &self, rows: &[PdfContentRow], source_id: i64, - ) -> (Vec<(String, String)>, HashMap) { - let mut jobs = Vec::new(); - let mut mapping = HashMap::new(); - for row in rows { - if let Some(path) = &row.img_path { - if mapping.contains_key(path) { - continue; - } - let new_path = Self::remove_image_id_prefix(path, source_id); - if new_path != *path { - jobs.push((path.clone(), new_path.clone())); - mapping.insert(path.clone(), new_path); - } - } - } - (jobs, mapping) - } - - fn prepare_raw_image_jobs(paths: &[String], source_id: i64) -> RawImageJobs { - let mut jobs = Vec::new(); - let mut mapping = HashMap::new(); - let mut target_paths = Vec::new(); - let mut seen = HashSet::new(); - for path in paths { - let trimmed = path.trim(); - if trimmed.is_empty() || !seen.insert(trimmed.to_string()) { - continue; - } - let new_path = Self::remove_image_id_prefix(trimmed, source_id); - if new_path != trimmed { - mapping.insert(trimmed.to_string(), new_path.clone()); - jobs.push((trimmed.to_string(), new_path.clone())); - } - target_paths.push(new_path); - } - (jobs, mapping, target_paths) - } - - async fn insert_pdf_rows( - &self, _tx: &mut sqlx::Transaction<'_, Sqlite>, target_file_id: i64, rows: &[PdfContentRow], - image_mapping: &HashMap, - ) -> anyhow::Result<()> { - let rows = rows - .iter() - .map(|row| { - let new_img_path = row - .img_path - .as_ref() - .and_then(|path| image_mapping.get(path)) - .cloned() - .or_else(|| row.img_path.clone()); - PdfContentRow { img_path: new_img_path, ..row.clone() } - }) - .collect::>(); - crate::pdf_content::write(target_file_id, &rows).await - } - - async fn insert_slice_rows( - &self, tx: &mut sqlx::Transaction<'_, Sqlite>, target_file_id: i64, rows: &[SliceRow], - ) -> anyhow::Result> { - if rows.is_empty() { - return Ok(Vec::new()); - } - - let binds_per_row = 2_usize; - let max_vars = 999_usize; - let batch_size = std::cmp::max(1, max_vars / binds_per_row); - let mut cloned = Vec::with_capacity(rows.len()); - for chunk in rows.chunks(batch_size) { - let mut qb = QueryBuilder::::new("INSERT INTO slices (file_id) "); - qb.push_values(chunk.iter(), |mut b, _row| { - b.push_bind(target_file_id); - }); - qb.push(" RETURNING id"); - - let inserted_ids: Vec<(i64,)> = qb.build_query_as().fetch_all(&mut **tx).await?; - anyhow::ensure!( - inserted_ids.len() == chunk.len(), - "inserted cloned slice row count mismatch: expected {}, got {}", - chunk.len(), - inserted_ids.len() - ); - - for (row, (new_id,)) in chunk.iter().zip(inserted_ids) { - cloned.push(ClonedSlice { old_id: row.id, new_id, content: row.content.clone() }); - } - } - - crate::slice_content::upsert_many( - target_file_id, - &cloned.iter().map(|row| (row.new_id, row.content.clone())).collect::>(), - ) - .await?; - Ok(cloned) - } - - async fn insert_slice_positions( - &self, tx: &mut sqlx::Transaction<'_, Sqlite>, cloned_slices: &[ClonedSlice], positions: &[SlicePositionRecord], - ) -> anyhow::Result<()> { - if positions.is_empty() || cloned_slices.is_empty() { - return Ok(()); - } - let id_map: HashMap = cloned_slices.iter().map(|slice| (slice.old_id, slice.new_id)).collect(); - let filtered: Vec<&SlicePositionRecord> = - positions.iter().filter(|row| id_map.contains_key(&row.slice_id)).collect(); - if filtered.is_empty() { - return Ok(()); - } - let binds_per_row = 8; - let batch_size = std::cmp::max(1, 999 / binds_per_row); - for chunk in filtered.chunks(batch_size) { - let mut qb = QueryBuilder::::new( - "INSERT INTO slice_positions(slice_id, page_idx, x1, y1, x2, y2, sheet_name, row_num) ", - ); - qb.push_values(chunk.iter(), |mut b, row| { - let new_id = id_map[&row.slice_id]; - b.push_bind(new_id) - .push_bind(row.page_idx) - .push_bind(row.x1) - .push_bind(row.y1) - .push_bind(row.x2) - .push_bind(row.y2) - .push_bind(&row.sheet_name) - .push_bind(row.row_num); - }); - qb.build().execute(&mut **tx).await?; - } - Ok(()) - } - - async fn copy_image_files(&self, jobs: &[(String, String)]) -> anyhow::Result<()> { - for (old_path, new_path) in jobs { - let Some(src_abs) = resolve_image_storage_path(old_path) else { - warn!("Unable to resolve source image path {}, skipping reuse copy", old_path); - continue; - }; - let Some(dst_abs) = resolve_image_storage_path(new_path) else { - warn!("Unable to resolve destination image path {}, skipping reuse copy", new_path); - continue; - }; - if let Some(parent) = std::path::Path::new(&dst_abs).parent() { - fs::create_dir_all(parent).await?; - } - fs::copy(&src_abs, &dst_abs).await?; - } - Ok(()) - } - - async fn copy_converted_pdf(&self, source_id: i64, target_id: i64) -> anyhow::Result<()> { - let cfg = config::get(); - let pdf_dir = std::path::Path::new(&cfg.storage.pdf_path); - let src_pdf = pdf_dir.join(format!("{}.pdf", source_id)); - match fs::try_exists(&src_pdf).await { - Ok(true) => {} - Ok(false) => return Ok(()), - Err(err) => { - warn!("Failed to check converted PDF existence for file {}: {}, skipping copy", source_id, err); - return Ok(()); - } - } - fs::create_dir_all(pdf_dir).await?; - let dst_pdf = pdf_dir.join(format!("{}.pdf", target_id)); - fs::copy(&src_pdf, &dst_pdf).await?; - Ok(()) - } - - async fn reindex_cloned_slices( - &self, target: &File, cloned_slices: &[ClonedSlice], source: &File, - ) -> anyhow::Result { - let mut search_docs = Vec::new(); - for slice in cloned_slices { - let is_image = content_looks_like_image_reference(&slice.content); - search_docs.push( - tantivy_engine::Document::new(slice.new_id, target.id, target.kb_id, slice.content.clone()) - .with_is_image(is_image), - ); - } - - if !search_docs.is_empty() { - let filename_lower = target.filename.to_lowercase(); - let is_image_file = crate::search::is_image_file(&filename_lower); - let embeddings = if is_image_file && search::embedding::image_embedding_enabled() { - let embedding = - search::embedding::get_image_embedding_from_path(&target.path, Some(&target.filename)).await?; - let arc_embedding = Arc::new(embedding); - (0..search_docs.len()).map(|_| Some(Arc::clone(&arc_embedding))).collect() - } else { - if is_image_file { - info!("Image embedding URL not configured, skipping image embedding for reused file {}", target.id); - } - vec![None; search_docs.len()] - }; - self.search_engine.write_batch(search_docs, embeddings).await?; - } - - let full_content = if let Some(content) = source.content.clone() { - content - } else if cloned_slices.is_empty() { - String::new() - } else { - cloned_slices - .iter() - .map(|slice| slice.content.as_str()) - .filter(|content| !content.trim().is_empty()) - .map(|content| content.to_string()) - .collect::>() - .join("\n\n") - }; - - let index_full_content = if full_content.is_empty() { - target.filename.clone() - } else { - format!("{}\n\n{}", target.filename, full_content) - }; - self.search_engine - .write_full(tantivy_engine::Document::new(target.id, target.id, target.kb_id, index_full_content)) - .await?; - self.persist_file_summary(target.id, target.kb_id, source.summary.as_deref()).await?; - - Ok(full_content) - } - - fn remove_image_id_prefix(original: &str, source_id: i64) -> String { - let path = std::path::Path::new(original); - let filename = path.file_name().and_then(|name| name.to_str()).unwrap_or(original); - let prefix = format!("f{}_", source_id); - let stripped = if filename.starts_with(&prefix) { - filename.trim_start_matches(&prefix).to_string() - } else { - filename.to_string() - }; - if let Some(parent) = path.parent() { parent.join(stripped).to_string_lossy().to_string() } else { stripped } - } - - async fn maybe_build_knowledge_graph(&self, file: &File) { - if !config::get().server.build_knowledge_graph { - debug!("Knowledge graph building disabled by config, skipping file {}", file.id); - return; - } - - if let Err(e) = self.build_knowledge_graph(file).await { - error!("Failed to build knowledge graph for file {}: {}", file.id, e); - // 不影响主流程，仅记录错误 - } - } - - /// 构建知识图谱（完全由LLM生成） - async fn build_knowledge_graph(&self, file: &File) -> anyhow::Result<()> { - info!("Building knowledge graph for file {}", file.id); - - // 1. 初始化LLM图谱提取器 - let llm_extractor = LLMGraphExtractor::from_env(); - - if !llm_extractor.is_enabled() { - warn!("LLM not enabled, skipping knowledge graph building for file {}", file.id); - return Ok(()); - } - - info!("Using LLM to generate knowledge graph for file {}", file.id); - - // 2. 获取文件的所有切片 - let source_file_id = effective_parse_file_id(&self.pool, file.id).await?; - let slices = self - .fetch_slice_rows(source_file_id) - .await? - .into_iter() - .map(|row| (row.id, row.content)) - .collect::>(); - - if slices.is_empty() { - debug!("No slices found for file {}, skipping graph building", file.id); - return Ok(()); - } - - // 3. 合并所有切片内容（限制长度避免超出LLM上下文） - let mut combined_content = String::new(); - let max_content_length = 8000; // 限制总长度 - - for (_, content) in &slices { - if combined_content.len() + content.len() > max_content_length { - break; - } - combined_content.push_str(content); - combined_content.push_str("\n\n"); - } - - if combined_content.trim().is_empty() { - debug!("No content to process for file {}", file.id); - return Ok(()); - } - - // 4. 调用LLM提取知识图谱 - let context = format!("文件名: {}", file.filename); - - let (mut entities, mut relations) = - match llm_extractor.extract_knowledge_graph(&combined_content, &context).await { - Ok(result) => result, - Err(e) => { - error!("LLM knowledge graph extraction failed for file {}: {}", file.id, e); - return Err(e); - } - }; - - info!("LLM extracted {} entities and {} relations from file {}", entities.len(), relations.len(), file.id); - - // 5. 为实体和关系添加文件信息 - for entity in &mut entities { - entity.file_id = Some(file.id); - entity.kb_id = file.kb_id; - } - - for relation in &mut relations { - relation.file_id = Some(file.id); - } - - // 6. 增量直写知识图谱（避免把整个 KB 的图加载到内存） - KnowledgeGraph::incremental_update_direct(self.pool.clone(), file.kb_id, entities, relations).await?; - - // 7. 保存图快照 - KnowledgeGraph::save_snapshot_direct(&self.pool, file.kb_id).await?; - - info!("Knowledge graph updated successfully for file {} (LLM-generated)", file.id); - - Ok(()) - } -} - -pub async fn process_file_immediate(pool: SqlitePool, search_engine: SearchEngine, file_id: i64) -> anyhow::Result<()> { - if is_parse_paused() { - anyhow::bail!("parse is paused for index maintenance"); - } - let processor = FileProcessor::new(pool, search_engine, 0); - let Some(file) = processor.claim_file_by_id(file_id).await? else { - info!("File {} is not pending; immediate parse claim skipped", file_id); - return Ok(()); - }; - let result = processor.process_file_claimed(&file).await; - if let Err(err) = &result { - if let Err(cleanup_err) = processor.cleanup_processing_file_data_with_retry(file_id, 3).await { - error!("Failed to cleanup immediate parse data for file {}: {}", file_id, cleanup_err); - } - processor.mark_file_failed(file_id, &err.to_string()).await?; - } - result -} - -pub async fn process_file_immediate_skip_reuse( - pool: SqlitePool, search_engine: SearchEngine, file_id: i64, -) -> anyhow::Result<()> { - if is_parse_paused() { - anyhow::bail!("parse is paused for index maintenance"); - } - let processor = FileProcessor::new(pool, search_engine, 0); - let Some(file) = processor.claim_file_by_id(file_id).await? else { - info!("File {} is not pending; immediate parse claim skipped", file_id); - return Ok(()); - }; - let result = processor.process_file_claimed_skip_reuse(&file).await; - if let Err(err) = &result { - if let Err(cleanup_err) = processor.cleanup_processing_file_data_with_retry(file_id, 3).await { - error!("Failed to cleanup immediate parse data for file {}: {}", file_id, cleanup_err); - } - processor.mark_file_failed(file_id, &err.to_string()).await?; - } - result -} - -pub async fn try_reuse_file_with_file( - pool: SqlitePool, search_engine: SearchEngine, file: File, -) -> anyhow::Result { - if is_parse_paused() { - return Ok(false); - } - let processor = FileProcessor::new(pool, search_engine, 0); - let Some(claimed_file) = processor.claim_file_by_id(file.id).await? else { - return Ok(false); - }; - let result = processor.try_reuse_existing_data(&claimed_file).await; - match result { - Ok(true) => Ok(true), - Ok(false) => { - sqlx::query( - "UPDATE files SET status = 0, parse_run_id = NULL, updated_at = strftime('%s','now') \ - WHERE id = ? AND status = 2", - ) - .bind(file.id) - .execute(&processor.pool) - .await?; - Ok(false) - } - Err(err) => { - sqlx::query( - "UPDATE files SET status = 0, parse_run_id = NULL, log = ?, updated_at = strftime('%s','now') \ - WHERE id = ? AND status = 2", - ) - .bind(format!("Reuse failed: {}", err)) - .bind(file.id) - .execute(&processor.pool) - .await?; - Err(err) - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[tokio::test] - async fn immediate_parse_claim_is_atomic() -> anyhow::Result<()> { - let temp = tempfile::NamedTempFile::new()?; - let url = format!("sqlite://{}", temp.path().display()); - let pool = sqlx::sqlite::SqlitePoolOptions::new().max_connections(4).connect(&url).await?; - sqlx::query( - "CREATE TABLE files ( - id INTEGER PRIMARY KEY, - user_id TEXT NOT NULL, user_name TEXT NOT NULL, hash TEXT NOT NULL, - filename TEXT NOT NULL, path TEXT NOT NULL, size INTEGER NOT NULL, - tags TEXT NOT NULL, status INTEGER NOT NULL, log TEXT NOT NULL, - slice_type TEXT NOT NULL, kb_id INTEGER, is_public INTEGER NOT NULL, - meta TEXT, summary TEXT, parse_priority INTEGER NOT NULL, parse_run_id TEXT, artifact_id INTEGER, - created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL - )", - ) - .execute(&pool) - .await?; - sqlx::query( - "INSERT INTO files VALUES (1, 'u', 'n', 'h', 'f.txt', '/tmp/f', 1, '', 0, '', 'text', NULL, 0, NULL, NULL, 50, NULL, NULL, 1, 1)", - ) - .execute(&pool) - .await?; - - let (first, second) = tokio::join!(claim_pending_file_by_id(&pool, 1), claim_pending_file_by_id(&pool, 1)); - let claimed = usize::from(first?.is_some()) + usize::from(second?.is_some()); - assert_eq!(claimed, 1); - - let (status, run_id): (i32, Option) = - sqlx::query_as("SELECT status, parse_run_id FROM files WHERE id = 1").fetch_one(&pool).await?; - assert_eq!(status, 2); - assert!(run_id.is_some()); - Ok(()) - } - - fn image_content_item(img_path: &str) -> ContentItem { - ContentItem { - typ: "image".to_string(), - bbox: vec![1, 2, 3, 4], - page_idx: 0, - text: None, - text_level: None, - text_format: None, - img_path: Some(img_path.to_string()), - image_caption: None, - table_body: None, - table_caption: None, - } - } - - #[test] - fn custom_parse_normalization_preserves_image_names() -> anyhow::Result<()> { - let mut images = HashMap::new(); - images.insert("img.png".to_string(), "raw-base64".to_string()); - let data = CustomParseData { - slices: vec![CustomSlice { - content: "see ![x](/api/v1/knowledge/files/img.png)".to_string(), - positions: Vec::new(), - }], - full_content: Some("full /api/v1/knowledge/files/img.png".to_string()), - summary: None, - images: Some(images), - content_list: Some(vec![image_content_item("images/img.png")]), - }; - - let normalized = FileProcessor::normalize_custom_parse_data(data)?; - - assert!(normalized.images.contains_key("img.png")); - assert_eq!(normalized.content_list.as_ref().unwrap()[0].img_path.as_deref(), Some("img.png")); - assert!(normalized.slices[0].content.contains("/api/v1/knowledge/files/img.png")); - assert!(normalized.full_content.as_ref().unwrap().contains("/api/v1/knowledge/files/img.png")); - assert_eq!(normalized.image_paths, vec!["img.png".to_string()]); - Ok(()) - } - - #[test] - fn custom_parse_normalization_preserves_unmapped_slice_refs() -> anyhow::Result<()> { - let data = CustomParseData { - slices: vec![CustomSlice { - content: "legacy ![x](/api/v1/knowledge/files/legacy.png)".to_string(), - positions: Vec::new(), - }], - full_content: None, - summary: None, - images: None, - content_list: None, - }; - - let normalized = FileProcessor::normalize_custom_parse_data(data)?; - - assert!(normalized.images.is_empty()); - assert!(normalized.slices[0].content.contains("/api/v1/knowledge/files/legacy.png")); - assert_eq!(normalized.image_paths, vec!["legacy.png".to_string()]); - Ok(()) - } - - #[test] - fn reuse_removes_legacy_image_id_prefix() { - assert_eq!(FileProcessor::remove_image_id_prefix("f42_img.png", 42), "img.png"); - assert_eq!(FileProcessor::remove_image_id_prefix("images/f42_img.png", 42), "images/img.png"); - assert_eq!(FileProcessor::remove_image_id_prefix("img.png", 42), "img.png"); - assert_eq!(FileProcessor::remove_image_id_prefix("f7_img.png", 42), "f7_img.png"); - } - - #[test] + let caption = desc_map.get(&file.filename).cloned().unwrap_or_default(); + content_list.push(ContentItem { + typ: "image".to_string(), + bbox: vec![], + page_idx: 0, + text: None, + text_level: None, + text_format: None, + img_path: Some(file.filename.clone()), + image_caption: if caption.is_empty() { None } else { Some(vec![caption]) }, + table_body: None, + table_caption: None, + }); + } + } + + enrich_content_list_with_image_descriptions(&mut content_list, &desc_map); + + // 构建全文与切片均为 CPU 密集操作，合并放到阻塞线程池执行，避免阻塞异步运行时。 + // content_list 在此之后不再使用，直接移入闭包。 + let slice_type = file.slice_type.clone(); + let (full_content, slices) = timed_step_opt(timing.as_deref_mut(), "slice_build", async { + tokio::task::spawn_blocking(move || -> anyhow::Result<(String, Vec)> { + let (full_content, full_segments) = Self::build_full_content_and_segments(&content_list); + let slices = if slice_type == "smart" || slice_type.is_empty() { + // 智能切片：使用 content_list + Self::smart_slice_content_with_positions(&content_list)? + } else if slice_type == "fixed" { + Self::fixed_slice_content_with_positions(&full_content, &full_segments)? + } else { + Self::slice_content(&full_content, &slice_type)? + .into_iter() + .map(|content| { + let is_image = content_looks_like_image_reference(&content); + SliceWithPositions { content, positions: vec![], is_image } + }) + .collect() + }; + Ok((full_content, slices)) + }) + .await + .map_err(|e| anyhow::anyhow!("slice task failed: {e}"))? + }) + .await?; + + if !self.ensure_file_exists(file.id, "before writing slices").await? { + return Ok(()); + } + + let slice_count = slices.len(); + let embeddings = vec![image_embedding.clone(); slice_count]; + self.finish_file_processing( + file, + slices, + embeddings, + &full_content, + "PDF processed successfully", + index_filename, + None, + timing, + ) + .await + } + + async fn call_mineru_api( + &self, file: &File, is_image: bool, mut timing: Option<&mut ParseTimingCtx>, + ) -> anyhow::Result { + let cfg = config::get(); + + if !is_image && cfg.services.mineru_max_pages > 0 { + // lopdf 解析整个 PDF 是同步阻塞操作（仅用于获取页数），放到阻塞线程池 + let path = file.path.clone(); + let page_count = + tokio::task::spawn_blocking(move || Document::load(&path).map(|doc| doc.get_pages().len())).await?; + match page_count { + Ok(total_pages) => { + if total_pages > cfg.services.mineru_max_pages { + return timed_step_opt(timing.as_deref_mut(), "mineru_api_in_ranges", async { + self.call_mineru_api_in_ranges(file, total_pages, cfg.services.mineru_max_pages).await + }) + .await; + } + } + Err(e) => { + warn!( + "Failed to read PDF page count for {}: {}, falling back to single MinerU request", + file.filename, e + ); + } + } + } + + timed_step_opt(timing, "mineru_api", async { + self.call_mineru_api_with_path(&file.path, &file.filename, is_image, None, None).await + }) + .await + } + + async fn call_mineru_api_in_ranges( + &self, file: &File, total_pages: usize, max_pages: usize, + ) -> anyhow::Result { + let total_parts = total_pages.div_ceil(max_pages); + info!( + "PDF {} has {} pages, calling MinerU in {} ranges (max {} pages per range)", + file.filename, total_pages, total_parts, max_pages + ); + + let mut merged_items: Vec = Vec::new(); + let mut merged_images: HashMap = HashMap::new(); + + for (part_index, range_start) in (0..total_pages).step_by(max_pages).enumerate() { + let range_end = std::cmp::min(range_start + max_pages, total_pages) - 1; + let chunk_filename = format!("{}_part_{}.pdf", file.filename, part_index + 1); + + let chunk_result = self + .call_mineru_api_with_path(&file.path, &chunk_filename, false, Some(range_start), Some(range_end)) + .await?; + let mut chunk_items: Vec = serde_json::from_str(&chunk_result.content_list)?; + + let image_prefix = format!("part{}_", part_index + 1); + let min_page_idx = chunk_items.iter().map(|item| item.page_idx).min(); + let needs_offset = min_page_idx.is_some_and(|min| min < range_start as i32); + for item in &mut chunk_items { + if needs_offset { + item.page_idx += range_start as i32; + } + if let Some(img_path) = item.img_path.as_deref() { + item.img_path = Some(Self::prefix_image_path(img_path, &image_prefix)); + } + } + + for (img_name, img_base64) in chunk_result.images { + let prefixed = Self::prefix_image_path(&img_name, &image_prefix); + merged_images.insert(prefixed, img_base64); + } + + merged_items.extend(chunk_items); + } + + let content_list = serde_json::to_string(&merged_items)?; + Ok(Result { content_list, images: merged_images }) + } + + fn prefix_image_path(img_path: &str, prefix: &str) -> String { + match img_path.rsplit_once('/') { + Some((dir, name)) => format!("{}/{}{}", dir, prefix, name), + None => format!("{}{}", prefix, img_path), + } + } + + async fn call_mineru_api_with_path( + &self, file_path: &str, filename: &str, is_image: bool, start_page: Option, end_page: Option, + ) -> anyhow::Result { + let cfg = config::get(); + let mineru_url = cfg.services.mineru_url.trim_end_matches('/'); + + let client = self.services_http_client()?; + if mineru_url.ends_with("/file_parse") { + // 构建 multipart form + let mime_type = if is_image { + mime_guess::from_path(filename).first_or_octet_stream().essence_str().to_string() + } else { + "application/pdf".to_string() + }; + + let file_part = stream_file_part(file_path, filename, &mime_type).await?; + let mut form = multipart::Form::new() + .text("return_middle_json", "false") + .text("return_model_output", "false") + .text("return_md", "false") + .text("return_images", "true") + .text("parse_method", "auto") + .text("lang_list", "ch") + .text("output_dir", "") + .text("server_url", "string") + .text("return_content_list", "true") + .text("backend", "pipeline") + .text("table_enable", "true") + .text("response_format_zip", "false") + .text("formula_enable", "true") + .part("files", file_part); + let start_page_id = start_page.unwrap_or(0); + let end_page_id = end_page.map(|v| v.to_string()).unwrap_or_else(|| "99999".to_string()); + form = form.text("start_page_id", start_page_id.to_string()).text("end_page_id", end_page_id); + + // 调用 MinerU API + let response = client.post(mineru_url).multipart(form).send().await?; + let status = response.status(); + let body_bytes = response.bytes().await?; + if !status.is_success() { + let error_text = String::from_utf8_lossy(&body_bytes); + return Err(anyhow::anyhow!("MinerU API failed: {}", error_text)); + } + + let mineru_response: MinerUResponse = serde_json::from_slice(&body_bytes).map_err(|e| { + let body_text = String::from_utf8_lossy(&body_bytes); + anyhow::anyhow!("MinerU API response decode failed: {} - {}", e, body_text) + })?; + let mineru_result = match mineru_response.results { + MinerUResults::Map(results) => { + results.into_values().next().ok_or_else(|| anyhow::anyhow!("MinerU API returned empty results"))? + } + MinerUResults::List(results) => { + if let Some(item) = results.iter().find(|item| item.status == "error") { + let mut error_msg = item.error.clone(); + if error_msg.is_empty() { + error_msg = "MinerU API returned error result".to_string(); + } + if !item.filename.is_empty() { + error_msg = format!("{} (file: {})", error_msg, item.filename); + } + return Err(anyhow::anyhow!("MinerU API failed: {}", error_msg)); + } + let first = results + .into_iter() + .next() + .ok_or_else(|| anyhow::anyhow!("MinerU API returned empty results"))?; + Result { content_list: first.content_list, images: first.images } + } + }; + return Ok(mineru_result); + } + + let mut url = reqwest::Url::parse(mineru_url)?; + { + let mut query = url.query_pairs_mut(); + query.append_pair("backend", "pipeline"); + query.append_pair("method", "auto"); + query.append_pair("lang", "ch"); + query.append_pair("start_page", &start_page.unwrap_or(0).to_string()); + if let Some(end_page) = end_page { + query.append_pair("end_page", &end_page.to_string()); + } + query.append_pair("formula_enable", "true"); + query.append_pair("table_enable", "true"); + } + + let mime_type = mime_guess::from_path(filename).first_or_octet_stream().essence_str().to_string(); + let file_part = stream_file_part(file_path, filename, &mime_type).await?; + let form = multipart::Form::new().part("file", file_part); + + let response = client.post(url.clone()).multipart(form).send().await?; + if !response.status().is_success() { + let error_text = response.text().await?; + return Err(anyhow::anyhow!("MinerU API failed: {}", error_text)); + } + + let analyze_response: AnalyzePdfResponse = response.json().await?; + if analyze_response.code != 200 { + return Err(anyhow::anyhow!("MinerU API failed: {}", analyze_response.message)); + } + + let Some(data) = analyze_response.data else { + return Err(anyhow::anyhow!("MinerU API returned empty data")); + }; + + let host = url.host_str().ok_or_else(|| anyhow::anyhow!("MinerU API URL missing host"))?; + let origin = if let Some(port) = url.port() { + format!("{}://{}:{}", url.scheme(), host, port) + } else { + format!("{}://{}", url.scheme(), host) + }; + + let mut content_list = data.content_list; + + // 去重收集需要下载的图片（filename -> url） + let mut to_download: HashMap = HashMap::new(); + for item in &content_list { + let Some(img_path) = item.img_path.as_deref() else { continue }; + let filename = std::path::Path::new(img_path) + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or(img_path) + .to_string(); + to_download + .entry(filename) + .or_insert_with(|| format!("{}/output/{}", origin, img_path.trim_start_matches('/'))); + } + + // 并发下载图片（限制并发数，避免打爆 MinerU），各图片之间无依赖。 + use futures::stream::{StreamExt, TryStreamExt}; + let images: HashMap = + futures::stream::iter(to_download.into_iter().map(|(filename, image_url)| { + let client = client.clone(); + async move { + let image_response = client.get(&image_url).send().await?; + if !image_response.status().is_success() { + let error_text = image_response.text().await?; + return Err(anyhow::anyhow!("MinerU image download failed: {}", error_text)); + } + let image_bytes = image_response.bytes().await?; + let image_mime = mime_guess::from_path(&filename).first_or_octet_stream().essence_str().to_string(); + let image_base64 = STANDARD.encode(image_bytes); + anyhow::Ok((filename, format!("data:{};base64,{}", image_mime, image_base64))) + } + })) + .buffer_unordered(8) + .try_collect() + .await?; + + // 回填 img_path 为去掉目录的文件名 + for item in &mut content_list { + if let Some(img_path) = item.img_path.as_deref() { + let filename = std::path::Path::new(img_path) + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or(img_path) + .to_string(); + item.img_path = Some(filename); + } + } + + let content_list_json = serde_json::to_string(&content_list)?; + Ok(Result { content_list: content_list_json, images }) + } + + /// 流式读取文本文件，边读边分片，避免一次性把整个文件载入内存。 + async fn stream_text_file(path: &str, slice_type: &str) -> anyhow::Result<(String, Vec)> { + use tokio::io::BufReader; + + let file = tokio::fs::File::open(path).await?; + let mut reader = BufReader::new(file); + let mut full_content = String::new(); + let mut slices = Vec::new(); + + match slice_type { + "paragraph" => { + let mut current = String::new(); + loop { + let mut line = String::new(); + let n = reader.read_line(&mut line).await?; + if n == 0 { + break; + } + + let (content, term) = if line.ends_with("\r\n") { + let c = line.drain(..line.len() - 2).collect::(); + (c, "\r\n") + } else if line.ends_with('\n') { + let c = line.drain(..line.len() - 1).collect::(); + (c, "\n") + } else { + (line, "") + }; + + full_content.push_str(&content); + full_content.push_str(term); + + if content.trim().is_empty() { + if !current.trim().is_empty() { + slices.push(current.trim().to_string()); + current.clear(); + } + } else { + if !current.is_empty() { + current.push_str(term); + } + current.push_str(&content); + } + } + if !current.trim().is_empty() { + slices.push(current.trim().to_string()); + } + } + "sentence" => { + const SENTENCE_END: [char; 6] = ['。', '.', '?', '!', '？', '！']; + let mut buf = String::new(); + loop { + let mut line = String::new(); + let n = reader.read_line(&mut line).await?; + if n == 0 { + break; + } + + let (content, term) = if line.ends_with("\r\n") { + let c = line.drain(..line.len() - 2).collect::(); + (c, "\r\n") + } else if line.ends_with('\n') { + let c = line.drain(..line.len() - 1).collect::(); + (c, "\n") + } else { + (line, "") + }; + + full_content.push_str(&content); + full_content.push_str(term); + + for ch in content.chars().chain(term.chars()) { + if SENTENCE_END.contains(&ch) { + if !buf.trim().is_empty() { + slices.push(buf.trim().to_string()); + } + buf.clear(); + } else { + buf.push(ch); + } + } + } + if !buf.trim().is_empty() { + slices.push(buf.trim().to_string()); + } + } + _ => { + let cfg = config::get(); + let chunk_size = cfg.slice.smart_slice_max_chars; + let overlap = cfg.slice.fixed_slice_overlap_chars; + let mut buf: Vec = Vec::new(); + + loop { + let mut line = String::new(); + let n = reader.read_line(&mut line).await?; + if n == 0 { + break; + } + + let (content, term) = if line.ends_with("\r\n") { + let c = line.drain(..line.len() - 2).collect::(); + (c, "\r\n") + } else if line.ends_with('\n') { + let c = line.drain(..line.len() - 1).collect::(); + (c, "\n") + } else { + (line, "") + }; + + full_content.push_str(&content); + full_content.push_str(term); + + for ch in content.chars().chain(term.chars()) { + buf.push(ch); + if buf.len() >= chunk_size { + let slice: String = buf[..chunk_size].iter().collect(); + slices.push(slice); + if overlap >= chunk_size || overlap == 0 { + buf.clear(); + } else { + let keep = buf.len() - overlap; + buf.drain(0..keep); + } + } + } + } + + while !buf.is_empty() { + let end = buf.len().min(chunk_size); + let slice: String = buf[..end].iter().collect(); + slices.push(slice); + if end == buf.len() { + break; + } + if overlap >= end || overlap == 0 { + break; + } + let keep = buf.len() - overlap; + buf.drain(0..keep); + } + } + } + + Ok((full_content, slices)) + } + + /// 处理普通文本文件 + async fn process_text_file(&self, file: &File, mut timing: Option<&mut ParseTimingCtx>) -> anyhow::Result<()> { + if !self.ensure_file_exists(file.id, "before reading text").await? { + return Ok(()); + } + let (content, slices) = timed_step_opt(timing.as_deref_mut(), "read_text_file", async { + Self::stream_text_file(&file.path, &file.slice_type).await + }) + .await?; + self.process_plain_text_content(file, content, slices, "Processing completed successfully", timing).await + } + + async fn process_audio_file(&self, file: &File, mut timing: Option<&mut ParseTimingCtx>) -> anyhow::Result<()> { + info!("Processing audio file: {}", file.filename); + + if !self.ensure_file_exists(file.id, "before reading audio").await? { + return Ok(()); + } + let mime_type = mime_guess::from_path(&file.filename).first_or_octet_stream().essence_str().to_string(); + let file_part = timed_step_opt( + timing.as_deref_mut(), + "open_audio_file", + stream_file_part(&file.path, &file.filename, &mime_type), + ) + .await?; + + let form = multipart::Form::new().part("file", file_part); + + let client = self.services_http_client()?; + let cfg = config::get(); + let mut req_builder = client.post(&cfg.services.audio_transcription_url).multipart(form); + if let Some(key) = &cfg.services.audio_transcription_key + && !key.is_empty() + { + req_builder = req_builder.header("Authorization", format!("Bearer {}", key)); + } + + let response = + timed_step_opt(timing.as_deref_mut(), "audio_transcription_api", async { Ok(req_builder.send().await?) }) + .await?; + if !response.status().is_success() { + let error_text = response.text().await?; + return Err(anyhow::anyhow!("Audio transcription API failed: {}", error_text)); + } + + let AudioTranscriptionResponse { text, language } = response.json().await?; + let log_message = if language.is_empty() { + "Audio processed successfully".to_string() + } else { + format!("Audio processed successfully (language: {})", language) + }; + + let text = timed_step_opt(timing.as_deref_mut(), "audio_transcription", async { Ok(text) }).await?; + let slices = timed_step_opt(timing.as_deref_mut(), "slice_build", async { + let content_owned = text.clone(); + let slice_type = file.slice_type.clone(); + tokio::task::spawn_blocking(move || Self::slice_content(&content_owned, &slice_type)) + .await + .map_err(|e| anyhow::anyhow!("slice_content task failed: {e}"))? + }) + .await?; + + self.process_plain_text_content(file, text, slices, &log_message, timing).await + } + + async fn process_plain_text_content( + &self, file: &File, content: String, slices: Vec, log_message: &str, + timing: Option<&mut ParseTimingCtx>, + ) -> anyhow::Result<()> { + let wrapped: Vec = slices + .into_iter() + .map(|content| { + let is_image = content_looks_like_image_reference(&content); + SliceWithPositions { content, positions: vec![], is_image } + }) + .collect(); + let embeddings = vec![None; wrapped.len()]; + self.finish_file_processing(file, wrapped, embeddings, &content, log_message, None, None, timing).await + } + + /// 标记文件处理失败 + async fn mark_file_failed(&self, file_id: i64, error_msg: &str) -> anyhow::Result<()> { + let sql = "UPDATE files SET status = -1, parse_run_id = NULL, log = ?, updated_at = strftime('%s','now') WHERE id = ?"; + sqlx::query(sql).bind(error_msg).bind(file_id).execute(&self.pool).await?; + Ok(()) + } + + async fn mark_file_storage_skipped(&self, file_id: i64) -> anyhow::Result<()> { + let sql = + "UPDATE files SET status = 3, parse_run_id = NULL, log = ?, updated_at = strftime('%s','now') WHERE id = ?"; + sqlx::query(sql).bind("Storage mode: not parsed").bind(file_id).execute(&self.pool).await?; + Ok(()) + } + + #[allow(clippy::too_many_arguments)] + async fn finish_file_processing( + &self, file: &File, slices: Vec, embeddings: Vec>>>, + full_content: &str, log_message: &str, index_name_override: Option<&str>, summary: Option<&str>, + mut timing: Option<&mut ParseTimingCtx>, + ) -> anyhow::Result<()> { + if !self.ensure_file_exists(file.id, "before writing slices").await? { + return Ok(()); + } + let slice_count = slices.len(); + anyhow::ensure!( + embeddings.len() == slice_count, + "embeddings count {} does not match slice count {}", + embeddings.len(), + slice_count + ); + + let parse_run_id: String = sqlx::query_scalar( + "SELECT parse_run_id FROM files WHERE id = ? AND status = 2 AND parse_run_id IS NOT NULL", + ) + .bind(file.id) + .fetch_one(&self.pool) + .await?; + + let (search_docs, search_embeddings) = timed_step_opt(timing.as_deref_mut(), "insert_slices", async { + let persisted = self.insert_slices_and_positions(file.id, &parse_run_id, slices).await?; + let mut docs = Vec::with_capacity(persisted.len()); + let mut embs = Vec::with_capacity(persisted.len()); + for (idx, (id, content, is_image)) in persisted.into_iter().enumerate() { + docs.push(tantivy_engine::Document::new(id, file.id, file.kb_id, content).with_is_image(is_image)); + embs.push(embeddings.get(idx).cloned().flatten()); + } + Ok((docs, embs)) + }) + .await?; + + if !search_docs.is_empty() { + timed_step_opt(timing.as_deref_mut(), "write_search_batch", async { + self.search_engine.write_batch(search_docs, search_embeddings).await?; + Ok(()) + }) + .await?; + } + + if !self.ensure_file_exists(file.id, "before writing full index").await? { + return Ok(()); + } + let index_name = index_name_override.unwrap_or(&file.filename); + let index_full_content = format!("{}\n\n{}", index_name, full_content); + timed_step_opt(timing.as_deref_mut(), "write_full_index", async { + self.search_engine + .write_full(tantivy_engine::Document::new(file.id, file.id, file.kb_id, index_full_content)) + .await?; + Ok(()) + }) + .await?; + + if !self.ensure_file_exists(file.id, "before updating status").await? { + return Ok(()); + } + timed_step_opt(timing.as_deref_mut(), "publish_parse_artifact", async { + let artifact_id = self.ensure_parse_artifact(file, full_content, summary).await?; + self.adopt_existing_artifact_if_needed(file, artifact_id, full_content, summary).await?; + Ok(()) + }) + .await?; + let normalized_summary = summary.and_then(|value| { + let trimmed = value.trim(); + if trimmed.is_empty() { None } else { Some(trimmed.to_string()) } + }); + timed_step_opt(timing.as_deref_mut(), "write_summary_index", async { + self.persist_file_summary(file.id, file.kb_id, normalized_summary.as_deref()).await + }) + .await?; + timed_step_opt(timing.as_deref_mut(), "finalize_file_status", async { + let sql = "UPDATE files SET status = 1, parse_run_id = NULL, log = ?, summary = ?, updated_at = strftime('%s','now') \ + WHERE id = ? AND status = 2 AND parse_run_id = ?"; + let updated = sqlx::query(sql) + .bind(log_message) + .bind(normalized_summary.as_deref()) + .bind(file.id) + .bind(&parse_run_id) + .execute(&self.pool) + .await? + .rows_affected(); + anyhow::ensure!(updated == 1, "parse run for file {} is no longer current", file.id); + crate::file_content::write(file.id, full_content).await?; + Ok(()) + }) + .await?; + + info!("File {} processed successfully with {} slices", file.id, slice_count); + + self.search_engine.reload_readers()?; + + timed_step_opt(timing, "build_knowledge_graph", async { + self.maybe_build_knowledge_graph(file).await; + Ok(()) + }) + .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(()) + } + + /// 返回每个 slice 的 (id, content, is_image) 供后续构建搜索文档使用。 + async fn insert_slices_and_positions( + &self, file_id: i64, parse_run_id: &str, slices: Vec, + ) -> anyhow::Result> { + if slices.is_empty() { + return Ok(Vec::new()); + } + // 每批最多 500 条 slice 一个事务，避免长时间持有 SQLite 写锁阻塞其他写操作 + const SLICE_TX_BATCH: usize = 500; + let mut persisted: Vec<(i64, String, bool)> = Vec::with_capacity(slices.len()); + for (chunk_idx, slice_chunk) in slices.chunks(SLICE_TX_BATCH).enumerate() { + let mut tx = self.pool.begin().await?; + let mut chunk_persisted: Vec<(i64, String, bool)> = Vec::with_capacity(slice_chunk.len()); + let mut position_rows: Vec<(i64, SlicePosition)> = Vec::new(); + let binds_per_row = 4_usize; + let max_vars = 999_usize; + let batch_size = std::cmp::max(1, max_vars / binds_per_row); + for (insert_chunk_idx, insert_chunk) in slice_chunk.chunks(batch_size).enumerate() { + let insert_offset = insert_chunk_idx * batch_size; + let mut slice_sql = + QueryBuilder::::new("INSERT INTO slices (file_id, parse_run_id, ordinal, is_image) "); + slice_sql.push_values(insert_chunk.iter().enumerate(), |mut b, (idx, slice)| { + let ordinal = chunk_idx * SLICE_TX_BATCH + insert_offset + idx; + b.push_bind(file_id) + .push_bind(parse_run_id) + .push_bind(ordinal as i64) + .push_bind(i64::from(slice.is_image)); + }); + slice_sql.push( + " ON CONFLICT(file_id, parse_run_id, ordinal) WHERE parse_run_id IS NOT NULL AND ordinal IS NOT NULL \ + DO UPDATE SET updated_at = strftime('%s','now') RETURNING id", + ); + + let inserted_ids: Vec<(i64,)> = slice_sql.build_query_as().fetch_all(&mut *tx).await?; + anyhow::ensure!( + inserted_ids.len() == insert_chunk.len(), + "inserted slice row count mismatch: expected {}, got {}", + insert_chunk.len(), + inserted_ids.len() + ); + + for (slice, (id,)) in insert_chunk.iter().zip(inserted_ids) { + for position in &slice.positions { + position_rows.push((id, position.clone())); + } + chunk_persisted.push((id, slice.content.clone(), slice.is_image)); + } + } + if !position_rows.is_empty() { + let binds_per_row = 8_usize; + let max_vars = 999_usize; + let batch_size = std::cmp::max(1, max_vars / binds_per_row); + for chunk in position_rows.chunks(batch_size) { + let mut pos_sql = QueryBuilder::::new( + "insert into slice_positions(slice_id, page_idx, x1, y1, x2, y2, sheet_name, row_num) ", + ); + pos_sql.push_values(chunk.iter(), |mut b, (slice_id, position)| { + b.push_bind(slice_id) + .push_bind(position.page_idx) + .push_bind(position.bbox[0]) + .push_bind(position.bbox[1]) + .push_bind(position.bbox[2]) + .push_bind(position.bbox[3]) + .push_bind(&position.sheet_name) + .push_bind(position.row_num); + }); + pos_sql.build().execute(&mut *tx).await?; + } + } + tx.commit().await?; + crate::slice_content::upsert_many( + file_id, + &chunk_persisted.iter().map(|(id, content, _)| (*id, content.clone())).collect::>(), + ) + .await?; + persisted.extend(chunk_persisted); + } + Ok(persisted) + } + + async fn ensure_file_exists(&self, file_id: i64, stage: &str) -> anyhow::Result { + let exists = sqlx::query_scalar::<_, i64>("SELECT 1 FROM files WHERE id = ? LIMIT 1") + .bind(file_id) + .fetch_optional(&self.pool) + .await? + .is_some(); + if !exists { + info!("File {} no longer exists during {}, skipping further processing", file_id, stage); + } + Ok(exists) + } + + async fn is_storage_kb(&self, kb_id: Option) -> anyhow::Result { + let Some(kb_id) = kb_id else { return Ok(false) }; + let kb_type: Option = sqlx::query_scalar("SELECT kb_type FROM knowledge_bases WHERE id = ?") + .bind(kb_id) + .fetch_optional(&self.pool) + .await?; + Ok(matches!(kb_type.as_deref(), Some("storage"))) + } + + fn is_audio_file(filename_lower: &str) -> bool { + filename_lower.ends_with(".wav") + || filename_lower.ends_with(".mp3") + || filename_lower.ends_with(".m4a") + || filename_lower.ends_with(".aac") + || filename_lower.ends_with(".flac") + || filename_lower.ends_with(".ogg") + || filename_lower.ends_with(".opus") + || filename_lower.ends_with(".wma") + || filename_lower.ends_with(".amr") + || filename_lower.ends_with(".aiff") + || filename_lower.ends_with(".aif") + || filename_lower.ends_with(".alac") + || filename_lower.ends_with(".webm") + } + + /// 根据 slice_type 对内容进行分片 + fn slice_content(content: &str, slice_type: &str) -> anyhow::Result> { + match slice_type { + "paragraph" => { + // 按段落分片（以双换行符分隔） + Ok(content.split("\n\n").map(|s| s.trim()).filter(|s| !s.is_empty()).map(|s| s.to_string()).collect()) + } + "sentence" => { + // 按句子分片（简单实现：以句号、问号、感叹号分隔） + Ok(content + .split(['。', '.', '?', '!', '？', '！']) + .map(|s| s.trim()) + .filter(|s| !s.is_empty()) + .map(|s| s.to_string()) + .collect()) + } + _ => { + // 固定长度分片（每8000字符，重叠100字符） + let cfg = config::get(); + let chunk_size = cfg.slice.smart_slice_max_chars; + let overlap = cfg.slice.fixed_slice_overlap_chars; + + let chars: Vec = content.chars().collect(); + let mut slices = Vec::new(); + let mut start = 0; + + while start < chars.len() { + let end = std::cmp::min(start + chunk_size, chars.len()); + let slice: String = chars[start..end].iter().collect(); + slices.push(slice); + + if end >= chars.len() { + break; + } + + // 下一个切片从 end - overlap 开始，以实现重叠 + start = end.saturating_sub(overlap); + if start >= end { + break; + } + } + + Ok(slices) + } + } + } + + fn build_full_content_and_segments(content_list: &[ContentItem]) -> (String, Vec) { + let mut full_content = String::new(); + let mut segments = Vec::new(); + let mut current_len = 0usize; + + for pdf_content in content_list { + if pdf_content.typ == "discarded" { + continue; + } + let mut item_content = String::new(); + if pdf_content.typ == "text" { + if let Some(text) = &pdf_content.text { + if let Some(lv) = pdf_content.text_level { + item_content.push_str("#".repeat(lv as usize).as_str()); + item_content.push(' '); + } + item_content.push_str(text); + } + } else if pdf_content.typ == "image" { + if let Some(img_name) = &pdf_content.img_path { + item_content.push_str(&format!("![{}](/api/v1/knowledge/files/{})", img_name, img_name)); + } + if let Some(captions) = &pdf_content.image_caption { + for caption in captions { + item_content.push_str(&caption.to_string()); + } + } + } else if pdf_content.typ == "table" { + if let Some(table_caption) = &pdf_content.table_caption { + for caption in table_caption { + item_content.push_str(&caption.to_string()); + } + } + if let Some(table_body) = &pdf_content.table_body { + item_content.push_str(table_body); + } + } + + if item_content.is_empty() { + continue; + } + + let start = current_len; + full_content.push_str(&item_content); + current_len += item_content.chars().count(); + let end = current_len; + + let positions = Self::positions_from_item(pdf_content); + if !positions.is_empty() { + segments.push(Segment { start, end, positions }); + } + + full_content.push_str("\n\n"); + current_len += 2; + } + + (full_content, segments) + } + + fn fixed_slice_content_with_positions( + content: &str, segments: &[Segment], + ) -> anyhow::Result> { + let cfg = config::get(); + let chunk_size = cfg.slice.smart_slice_max_chars; + let overlap = cfg.slice.fixed_slice_overlap_chars; + + let chars: Vec = content.chars().collect(); + let mut slices = Vec::new(); + let mut start = 0; + + while start < chars.len() { + let end = std::cmp::min(start + chunk_size, chars.len()); + let slice: String = chars[start..end].iter().collect(); + let positions = Self::positions_for_range(segments, start, end); + slices.push(SliceWithPositions { + content: slice.clone(), + positions, + is_image: content_looks_like_image_reference(&slice), + }); + + if end >= chars.len() { + break; + } + + start = end.saturating_sub(overlap); + if start >= end { + break; + } + } + + Ok(slices) + } + + fn smart_slice_content_with_positions(content_list: &[ContentItem]) -> anyhow::Result> { + let cfg = config::get(); + let max_chars = cfg.slice.smart_slice_max_chars; + + let mut slices = Vec::new(); + let mut current_slice = String::new(); + let mut current_segments: Vec = Vec::new(); + let mut current_len = 0usize; + let mut current_header = String::new(); // 当前所在的标题 + let mut current_header_positions: Vec = Vec::new(); + + for item in content_list { + if item.typ == "discarded" { + continue; + } + + let mut item_content = String::new(); + let item_positions = Self::positions_from_item(item); + + // 处理文本内容 + if item.typ == "text" { + if let Some(text) = &item.text { + // 如果是标题，更新当前标题 + if let Some(level) = item.text_level { + // 这是一个标题 + let header_text = format!("{} {}", "#".repeat(level as usize), text); + + // 如果当前有累积的内容，先保存 + if !current_slice.trim().is_empty() { + Self::flush_slice_with_positions( + &mut slices, + &mut current_slice, + &mut current_segments, + max_chars, + ); + current_len = 0; + } + + // 更新当前标题 + current_header = header_text; + current_header_positions = item_positions.clone(); + // 开始新的切片，包含标题 + current_slice.clear(); + current_segments.clear(); + Self::append_segment( + &mut current_slice, + &mut current_len, + &mut current_segments, + ¤t_header, + current_header_positions.clone(), + ); + Self::append_separator(&mut current_slice, &mut current_len); + continue; + } else { + // 普通文本 + item_content.push_str(text); + } + } + } else if item.typ == "image" { + if let Some(img_name) = &item.img_path { + item_content.push_str(&format!("![{}](/api/v1/knowledge/files/{})", img_name, img_name)); + } + if let Some(captions) = &item.image_caption { + for caption in captions { + item_content.push_str(&caption.to_string()); + } + } + } else if item.typ == "table" { + if let Some(table_caption) = &item.table_caption { + for caption in table_caption { + item_content.push_str(&caption.to_string()); + } + } + if let Some(table_body) = &item.table_body { + item_content.push_str( + &table_body + .replace(" colspan=\"1\"", "") + .replace(" rowspan=\"1\"", "") + .replace(" colspan='1'", "") + .replace(" rowspan='1'", "") + .replace(" colspan=1", "") + .replace(" rowspan=1", ""), + ); + } + } + + if item_content.is_empty() { + continue; + } + + // 检查加入这个内容后是否超过字数限制 + let test_len = if current_slice.is_empty() { + // 如果当前切片为空，但有标题，先加上标题 + if !current_header.is_empty() { + current_header.chars().count() + 2 + item_content.chars().count() + } else { + item_content.chars().count() + } + } else { + current_len + item_content.chars().count() + 2 + }; + + if test_len > max_chars { + // 超过限制，保存当前切片 + if !current_slice.trim().is_empty() { + Self::flush_slice_with_positions(&mut slices, &mut current_slice, &mut current_segments, max_chars); + current_len = 0; + } + + // 开始新的切片 + current_slice.clear(); + current_segments.clear(); + if !current_header.is_empty() { + Self::append_segment( + &mut current_slice, + &mut current_len, + &mut current_segments, + ¤t_header, + current_header_positions.clone(), + ); + Self::append_separator(&mut current_slice, &mut current_len); + } + Self::append_segment( + &mut current_slice, + &mut current_len, + &mut current_segments, + &item_content, + item_positions.clone(), + ); + Self::append_separator(&mut current_slice, &mut current_len); + } else { + // 没有超过限制，继续累积 + if current_slice.is_empty() { + current_slice.clear(); + current_segments.clear(); + current_len = 0; + if !current_header.is_empty() { + Self::append_segment( + &mut current_slice, + &mut current_len, + &mut current_segments, + ¤t_header, + current_header_positions.clone(), + ); + Self::append_separator(&mut current_slice, &mut current_len); + } + Self::append_segment( + &mut current_slice, + &mut current_len, + &mut current_segments, + &item_content, + item_positions.clone(), + ); + } else { + Self::append_segment( + &mut current_slice, + &mut current_len, + &mut current_segments, + &item_content, + item_positions.clone(), + ); + Self::append_separator(&mut current_slice, &mut current_len); + } + } + } + + // 保存最后的切片 + if !current_slice.trim().is_empty() { + Self::flush_slice_with_positions(&mut slices, &mut current_slice, &mut current_segments, max_chars); + } + + Ok(slices) + } + + fn flush_slice_with_positions( + slices: &mut Vec, current_slice: &mut String, segments: &mut Vec, max_chars: usize, + ) { + if current_slice.trim().is_empty() { + return; + } + let content = std::mem::take(current_slice); + let segment_data = std::mem::take(segments); + let mut new_slices = Self::split_slice_with_positions(&content, &segment_data, max_chars); + slices.append(&mut new_slices); + } + + fn split_slice_with_positions(content: &str, segments: &[Segment], max_chars: usize) -> Vec { + let char_count = content.chars().count(); + if char_count == 0 { + return Vec::new(); + } + + let ranges = if char_count <= max_chars { + vec![(0, char_count)] + } else { + let sentence_ranges = Self::sentence_ranges(content); + if sentence_ranges.is_empty() { + vec![(0, char_count)] + } else { + let mut ranges = Vec::new(); + let mut slice_start = sentence_ranges[0].0; + let mut slice_end = sentence_ranges[0].1; + for (start, end) in sentence_ranges.iter().skip(1) { + if end - slice_start > max_chars && slice_end > slice_start { + ranges.push((slice_start, slice_end)); + slice_start = *start; + } + slice_end = *end; + } + ranges.push((slice_start, slice_end)); + ranges + } + }; + + let chars: Vec = content.chars().collect(); + ranges + .into_iter() + .filter_map(|(start, end)| { + if start >= end || end > chars.len() { + return None; + } + let slice: String = chars[start..end].iter().collect(); + let positions = Self::positions_for_range(segments, start, end); + Some(SliceWithPositions { + content: slice.clone(), + positions, + is_image: content_looks_like_image_reference(&slice), + }) + }) + .collect() + } + + fn sentence_ranges(content: &str) -> Vec<(usize, usize)> { + let mut ranges = Vec::new(); + let mut start = 0usize; + for (idx, ch) in content.chars().enumerate() { + if matches!(ch, '。' | '.' | '?' | '!' | '？' | '！') { + let end = idx + 1; + if end > start { + ranges.push((start, end)); + } + start = end; + } + } + let total = content.chars().count(); + if start < total { + ranges.push((start, total)); + } + ranges + } + + fn append_segment( + current_slice: &mut String, current_len: &mut usize, segments: &mut Vec, text: &str, + positions: Vec, + ) { + if text.is_empty() { + return; + } + let start = *current_len; + current_slice.push_str(text); + *current_len += text.chars().count(); + let end = *current_len; + if !positions.is_empty() { + segments.push(Segment { start, end, positions }); + } + } + + fn append_separator(current_slice: &mut String, current_len: &mut usize) { + current_slice.push_str("\n\n"); + *current_len += 2; + } + + fn positions_from_item(item: &ContentItem) -> Vec { + if item.bbox.len() == 4 { + let bbox = [item.bbox[0], item.bbox[1], item.bbox[2], item.bbox[3]]; + vec![SlicePosition { page_idx: item.page_idx, bbox, sheet_name: None, row_num: None }] + } else { + Vec::new() + } + } + + fn positions_for_range(segments: &[Segment], start: usize, end: usize) -> Vec { + let mut set: HashSet = HashSet::new(); + for segment in segments { + if segment.end > start && segment.start < end { + for position in &segment.positions { + set.insert(position.clone()); + } + } + } + set.into_iter().collect() + } + + async fn clone_file_data(&self, source: &File, target: &File) -> anyhow::Result<()> { + if source.id == target.id { + anyhow::bail!("Source and target file ids are identical"); + } + if source.status != 1 { + anyhow::bail!("Source file {} is not processed", source.id); + } + + // 新结构只复制搜索投影，SQLite 中的切片/PDF 解析结果由 artifact 共享。 + if config::get().server.reuse_duplicate_files { + let source_file_id = effective_parse_file_id(&self.pool, source.id).await?; + let mut source_for_reindex = source.clone(); + let full_content = if let Some(artifact_id) = source.artifact_id { + sqlx::query_scalar::<_, Option>("SELECT full_content FROM parse_artifacts WHERE id = ?") + .bind(artifact_id) + .fetch_optional(&self.pool) + .await? + .flatten() + .unwrap_or_default() + } else { + crate::file_content::read(source_file_id).await?.unwrap_or_default() + }; + let summary = if let Some(artifact_id) = source.artifact_id { + sqlx::query_scalar::<_, Option>("SELECT summary FROM parse_artifacts WHERE id = ?") + .bind(artifact_id) + .fetch_optional(&self.pool) + .await? + .flatten() + } else { + source.summary.clone() + }; + source_for_reindex.content = Some(full_content.clone()); + source_for_reindex.summary = summary.clone(); + let artifact_id = + self.ensure_parse_artifact(&source_for_reindex, &full_content, summary.as_deref()).await?; + let rows = self.fetch_slice_rows(source_file_id).await?; + let shared_slices: Vec = rows + .into_iter() + .map(|row| ClonedSlice { old_id: row.id, new_id: row.id, content: row.content }) + .collect(); + + sqlx::query( + "UPDATE files SET status = 2, artifact_id = ?, log = ?, updated_at = strftime('%s','now') WHERE id = ?", + ) + .bind(artifact_id) + .bind(format!("Reusing shared parse artifact {}", artifact_id)) + .bind(target.id) + .execute(&self.pool) + .await?; + let indexed_content = self.reindex_cloned_slices(target, &shared_slices, &source_for_reindex).await?; + sqlx::query( + "UPDATE files SET status = 1, parse_run_id = NULL, artifact_id = ?, log = ?, summary = ?, \ + updated_at = strftime('%s','now') WHERE id = ?", + ) + .bind(artifact_id) + .bind(format!("Reused shared parse artifact {} from file {}", artifact_id, source_file_id)) + .bind(summary.as_deref()) + .bind(target.id) + .execute(&self.pool) + .await?; + crate::file_content::write(target.id, &indexed_content).await?; + let mut updated_file = target.clone(); + updated_file.artifact_id = Some(artifact_id); + updated_file.content = Some(indexed_content); + updated_file.summary = summary.clone(); + 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(()); + } + + let reuse_log = format!("Reusing parsed data from file {}", source.id); + sqlx::query("UPDATE files SET status = 2, log = ?, updated_at = strftime('%s','now') WHERE id = ?") + .bind(&reuse_log) + .bind(target.id) + .execute(&self.pool) + .await?; + + self.cleanup_processing_file_data_with_retry(target.id, 3).await?; + + let pdf_rows = self.fetch_pdf_content_rows(source.id).await?; + let mut slice_rows = self.fetch_slice_rows(source.id).await?; + let slice_ids: Vec = slice_rows.iter().map(|row| row.id).collect(); + let slice_positions = self.fetch_slice_position_rows(&slice_ids).await?; + let (image_jobs, image_mapping) = self.prepare_image_jobs(&pdf_rows, source.id); + let meta_image_paths = if pdf_rows.is_empty() { + collect_image_raw_paths_for_files(&self.pool, &[source.id]).await? + } else { + Vec::new() + }; + let (meta_image_jobs, meta_image_mapping, target_meta_image_paths) = + Self::prepare_raw_image_jobs(&meta_image_paths, source.id); + + // 兼容历史数据：复用时剥离旧的 f{file_id}_ 图片前缀，并同步重写内容中的引用。 + let mut combined_image_mapping = image_mapping.clone(); + combined_image_mapping.extend(meta_image_mapping.clone()); + if !combined_image_mapping.is_empty() { + for row in &mut slice_rows { + row.content = Self::rewrite_custom_image_refs(&row.content, &combined_image_mapping); + } + } + let mut source_for_reindex = source.clone(); + if !combined_image_mapping.is_empty() + && let Some(content) = crate::file_content::read(source.id).await? + { + source_for_reindex.content = Some(Self::rewrite_custom_image_refs(&content, &combined_image_mapping)); + } + + let mut tx = self.pool.begin().await?; + self.insert_pdf_rows(&mut tx, target.id, &pdf_rows, &image_mapping).await?; + let cloned_slices = self.insert_slice_rows(&mut tx, target.id, &slice_rows).await?; + self.insert_slice_positions(&mut tx, &cloned_slices, &slice_positions).await?; + tx.commit().await?; + + self.copy_image_files(&image_jobs).await?; + self.copy_image_files(&meta_image_jobs).await?; + self.copy_converted_pdf(source.id, target.id).await?; + + if pdf_rows.is_empty() { + update_file_custom_image_meta(&self.pool, target.id, &target_meta_image_paths, "reuse_custom_images") + .await?; + } + + let full_content = self.reindex_cloned_slices(target, &cloned_slices, &source_for_reindex).await?; + + self.search_engine.reload_readers()?; + + let final_log = format!("Reused parsed data from file {}", source.id); + sqlx::query( + "UPDATE files SET status = 1, parse_run_id = NULL, log = ?, summary = ?, updated_at = strftime('%s','now') WHERE id = ?", + ) + .bind(&final_log) + .bind(source.summary.as_deref()) + .bind(target.id) + .execute(&self.pool) + .await?; + crate::file_content::write(target.id, &full_content).await?; + + let mut updated_file = target.clone(); + updated_file.content = Some(full_content.clone()); + updated_file.summary = source.summary.clone(); + 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(()) + } + + async fn fetch_pdf_content_rows(&self, file_id: i64) -> anyhow::Result> { + crate::pdf_content::read(file_id).await + } + + async fn fetch_slice_rows(&self, file_id: i64) -> anyhow::Result> { + let ids: Vec = sqlx::query_scalar("SELECT id FROM slices WHERE file_id = ? ORDER BY id") + .bind(file_id) + .fetch_all(&self.pool) + .await?; + let contents = crate::slice_content::read_all(file_id).await?; + Ok(ids.into_iter().map(|id| SliceRow { id, content: contents.get(&id).cloned().unwrap_or_default() }).collect()) + } + + async fn fetch_slice_position_rows(&self, slice_ids: &[i64]) -> anyhow::Result> { + if slice_ids.is_empty() { + return Ok(Vec::new()); + } + let mut all_rows = Vec::new(); + let chunk_size = 400; + for chunk in slice_ids.chunks(chunk_size) { + let mut qb = QueryBuilder::::new( + "SELECT slice_id, page_idx, x1, y1, x2, y2, sheet_name, row_num FROM slice_positions WHERE slice_id IN (", + ); + let mut separated = qb.separated(", "); + for slice_id in chunk { + separated.push_bind(slice_id); + } + qb.push(") ORDER BY slice_id, id"); + let rows = qb.build_query_as::().fetch_all(&self.pool).await?; + all_rows.extend(rows); + } + Ok(all_rows) + } + + fn prepare_image_jobs( + &self, rows: &[PdfContentRow], source_id: i64, + ) -> (Vec<(String, String)>, HashMap) { + let mut jobs = Vec::new(); + let mut mapping = HashMap::new(); + for row in rows { + if let Some(path) = &row.img_path { + if mapping.contains_key(path) { + continue; + } + let new_path = Self::remove_image_id_prefix(path, source_id); + if new_path != *path { + jobs.push((path.clone(), new_path.clone())); + mapping.insert(path.clone(), new_path); + } + } + } + (jobs, mapping) + } + + fn prepare_raw_image_jobs(paths: &[String], source_id: i64) -> RawImageJobs { + let mut jobs = Vec::new(); + let mut mapping = HashMap::new(); + let mut target_paths = Vec::new(); + let mut seen = HashSet::new(); + for path in paths { + let trimmed = path.trim(); + if trimmed.is_empty() || !seen.insert(trimmed.to_string()) { + continue; + } + let new_path = Self::remove_image_id_prefix(trimmed, source_id); + if new_path != trimmed { + mapping.insert(trimmed.to_string(), new_path.clone()); + jobs.push((trimmed.to_string(), new_path.clone())); + } + target_paths.push(new_path); + } + (jobs, mapping, target_paths) + } + + async fn insert_pdf_rows( + &self, _tx: &mut sqlx::Transaction<'_, Sqlite>, target_file_id: i64, rows: &[PdfContentRow], + image_mapping: &HashMap, + ) -> anyhow::Result<()> { + let rows = rows + .iter() + .map(|row| { + let new_img_path = row + .img_path + .as_ref() + .and_then(|path| image_mapping.get(path)) + .cloned() + .or_else(|| row.img_path.clone()); + PdfContentRow { img_path: new_img_path, ..row.clone() } + }) + .collect::>(); + crate::pdf_content::write(target_file_id, &rows).await + } + + async fn insert_slice_rows( + &self, tx: &mut sqlx::Transaction<'_, Sqlite>, target_file_id: i64, rows: &[SliceRow], + ) -> anyhow::Result> { + if rows.is_empty() { + return Ok(Vec::new()); + } + + let binds_per_row = 2_usize; + let max_vars = 999_usize; + let batch_size = std::cmp::max(1, max_vars / binds_per_row); + let mut cloned = Vec::with_capacity(rows.len()); + for chunk in rows.chunks(batch_size) { + let mut qb = QueryBuilder::::new("INSERT INTO slices (file_id) "); + qb.push_values(chunk.iter(), |mut b, _row| { + b.push_bind(target_file_id); + }); + qb.push(" RETURNING id"); + + let inserted_ids: Vec<(i64,)> = qb.build_query_as().fetch_all(&mut **tx).await?; + anyhow::ensure!( + inserted_ids.len() == chunk.len(), + "inserted cloned slice row count mismatch: expected {}, got {}", + chunk.len(), + inserted_ids.len() + ); + + for (row, (new_id,)) in chunk.iter().zip(inserted_ids) { + cloned.push(ClonedSlice { old_id: row.id, new_id, content: row.content.clone() }); + } + } + + crate::slice_content::upsert_many( + target_file_id, + &cloned.iter().map(|row| (row.new_id, row.content.clone())).collect::>(), + ) + .await?; + Ok(cloned) + } + + async fn insert_slice_positions( + &self, tx: &mut sqlx::Transaction<'_, Sqlite>, cloned_slices: &[ClonedSlice], positions: &[SlicePositionRecord], + ) -> anyhow::Result<()> { + if positions.is_empty() || cloned_slices.is_empty() { + return Ok(()); + } + let id_map: HashMap = cloned_slices.iter().map(|slice| (slice.old_id, slice.new_id)).collect(); + let filtered: Vec<&SlicePositionRecord> = + positions.iter().filter(|row| id_map.contains_key(&row.slice_id)).collect(); + if filtered.is_empty() { + return Ok(()); + } + let binds_per_row = 8; + let batch_size = std::cmp::max(1, 999 / binds_per_row); + for chunk in filtered.chunks(batch_size) { + let mut qb = QueryBuilder::::new( + "INSERT INTO slice_positions(slice_id, page_idx, x1, y1, x2, y2, sheet_name, row_num) ", + ); + qb.push_values(chunk.iter(), |mut b, row| { + let new_id = id_map[&row.slice_id]; + b.push_bind(new_id) + .push_bind(row.page_idx) + .push_bind(row.x1) + .push_bind(row.y1) + .push_bind(row.x2) + .push_bind(row.y2) + .push_bind(&row.sheet_name) + .push_bind(row.row_num); + }); + qb.build().execute(&mut **tx).await?; + } + Ok(()) + } + + async fn copy_image_files(&self, jobs: &[(String, String)]) -> anyhow::Result<()> { + for (old_path, new_path) in jobs { + let Some(src_abs) = resolve_image_storage_path(old_path) else { + warn!("Unable to resolve source image path {}, skipping reuse copy", old_path); + continue; + }; + let Some(dst_abs) = resolve_image_storage_path(new_path) else { + warn!("Unable to resolve destination image path {}, skipping reuse copy", new_path); + continue; + }; + if let Some(parent) = std::path::Path::new(&dst_abs).parent() { + fs::create_dir_all(parent).await?; + } + fs::copy(&src_abs, &dst_abs).await?; + } + Ok(()) + } + + async fn copy_converted_pdf(&self, source_id: i64, target_id: i64) -> anyhow::Result<()> { + let cfg = config::get(); + let pdf_dir = std::path::Path::new(&cfg.storage.pdf_path); + let src_pdf = pdf_dir.join(format!("{}.pdf", source_id)); + match fs::try_exists(&src_pdf).await { + Ok(true) => {} + Ok(false) => return Ok(()), + Err(err) => { + warn!("Failed to check converted PDF existence for file {}: {}, skipping copy", source_id, err); + return Ok(()); + } + } + fs::create_dir_all(pdf_dir).await?; + let dst_pdf = pdf_dir.join(format!("{}.pdf", target_id)); + fs::copy(&src_pdf, &dst_pdf).await?; + Ok(()) + } + + async fn reindex_cloned_slices( + &self, target: &File, cloned_slices: &[ClonedSlice], source: &File, + ) -> anyhow::Result { + let mut search_docs = Vec::new(); + for slice in cloned_slices { + let is_image = content_looks_like_image_reference(&slice.content); + search_docs.push( + tantivy_engine::Document::new(slice.new_id, target.id, target.kb_id, slice.content.clone()) + .with_is_image(is_image), + ); + } + + if !search_docs.is_empty() { + let filename_lower = target.filename.to_lowercase(); + let is_image_file = crate::search::is_image_file(&filename_lower); + let embeddings = if is_image_file && search::embedding::image_embedding_enabled() { + let embedding = + search::embedding::get_image_embedding_from_path(&target.path, Some(&target.filename)).await?; + let arc_embedding = Arc::new(embedding); + (0..search_docs.len()).map(|_| Some(Arc::clone(&arc_embedding))).collect() + } else { + if is_image_file { + info!("Image embedding URL not configured, skipping image embedding for reused file {}", target.id); + } + vec![None; search_docs.len()] + }; + self.search_engine.write_batch(search_docs, embeddings).await?; + } + + let full_content = if let Some(content) = source.content.clone() { + content + } else if cloned_slices.is_empty() { + String::new() + } else { + cloned_slices + .iter() + .map(|slice| slice.content.as_str()) + .filter(|content| !content.trim().is_empty()) + .map(|content| content.to_string()) + .collect::>() + .join("\n\n") + }; + + let index_full_content = if full_content.is_empty() { + target.filename.clone() + } else { + format!("{}\n\n{}", target.filename, full_content) + }; + self.search_engine + .write_full(tantivy_engine::Document::new(target.id, target.id, target.kb_id, index_full_content)) + .await?; + self.persist_file_summary(target.id, target.kb_id, source.summary.as_deref()).await?; + + Ok(full_content) + } + + fn remove_image_id_prefix(original: &str, source_id: i64) -> String { + let path = std::path::Path::new(original); + let filename = path.file_name().and_then(|name| name.to_str()).unwrap_or(original); + let prefix = format!("f{}_", source_id); + let stripped = if filename.starts_with(&prefix) { + filename.trim_start_matches(&prefix).to_string() + } else { + filename.to_string() + }; + if let Some(parent) = path.parent() { parent.join(stripped).to_string_lossy().to_string() } else { stripped } + } + + async fn maybe_build_knowledge_graph(&self, file: &File) { + if !config::get().server.build_knowledge_graph { + debug!("Knowledge graph building disabled by config, skipping file {}", file.id); + return; + } + + if let Err(e) = self.build_knowledge_graph(file).await { + error!("Failed to build knowledge graph for file {}: {}", file.id, e); + // 不影响主流程，仅记录错误 + } + } + + /// 构建知识图谱（完全由LLM生成） + async fn build_knowledge_graph(&self, file: &File) -> anyhow::Result<()> { + info!("Building knowledge graph for file {}", file.id); + + // 1. 初始化LLM图谱提取器 + let llm_extractor = LLMGraphExtractor::from_env(); + + if !llm_extractor.is_enabled() { + warn!("LLM not enabled, skipping knowledge graph building for file {}", file.id); + return Ok(()); + } + + info!("Using LLM to generate knowledge graph for file {}", file.id); + + // 2. 获取文件的所有切片 + let source_file_id = effective_parse_file_id(&self.pool, file.id).await?; + let slices = self + .fetch_slice_rows(source_file_id) + .await? + .into_iter() + .map(|row| (row.id, row.content)) + .collect::>(); + + if slices.is_empty() { + debug!("No slices found for file {}, skipping graph building", file.id); + return Ok(()); + } + + // 3. 合并所有切片内容（限制长度避免超出LLM上下文） + let mut combined_content = String::new(); + let max_content_length = 8000; // 限制总长度 + + for (_, content) in &slices { + if combined_content.len() + content.len() > max_content_length { + break; + } + combined_content.push_str(content); + combined_content.push_str("\n\n"); + } + + if combined_content.trim().is_empty() { + debug!("No content to process for file {}", file.id); + return Ok(()); + } + + // 4. 调用LLM提取知识图谱 + let context = format!("文件名: {}", file.filename); + + let (mut entities, mut relations) = + match llm_extractor.extract_knowledge_graph(&combined_content, &context).await { + Ok(result) => result, + Err(e) => { + error!("LLM knowledge graph extraction failed for file {}: {}", file.id, e); + return Err(e); + } + }; + + info!("LLM extracted {} entities and {} relations from file {}", entities.len(), relations.len(), file.id); + + // 5. 为实体和关系添加文件信息 + for entity in &mut entities { + entity.file_id = Some(file.id); + entity.kb_id = file.kb_id; + } + + for relation in &mut relations { + relation.file_id = Some(file.id); + } + + // 6. 增量直写知识图谱（避免把整个 KB 的图加载到内存） + KnowledgeGraph::incremental_update_direct(self.pool.clone(), file.kb_id, entities, relations).await?; + + // 7. 保存图快照 + KnowledgeGraph::save_snapshot_direct(&self.pool, file.kb_id).await?; + + info!("Knowledge graph updated successfully for file {} (LLM-generated)", file.id); + + Ok(()) + } +} + +pub async fn process_file_immediate(pool: SqlitePool, search_engine: SearchEngine, file_id: i64) -> anyhow::Result<()> { + if is_parse_paused() { + anyhow::bail!("parse is paused for index maintenance"); + } + let processor = FileProcessor::new(pool, search_engine, 0); + let Some(file) = processor.claim_file_by_id(file_id).await? else { + info!("File {} is not pending; immediate parse claim skipped", file_id); + return Ok(()); + }; + let result = processor.process_file_claimed(&file).await; + if let Err(err) = &result { + if let Err(cleanup_err) = processor.cleanup_processing_file_data_with_retry(file_id, 3).await { + error!("Failed to cleanup immediate parse data for file {}: {}", file_id, cleanup_err); + } + processor.mark_file_failed(file_id, &err.to_string()).await?; + } + result +} + +pub async fn process_file_immediate_skip_reuse( + pool: SqlitePool, search_engine: SearchEngine, file_id: i64, +) -> anyhow::Result<()> { + if is_parse_paused() { + anyhow::bail!("parse is paused for index maintenance"); + } + let processor = FileProcessor::new(pool, search_engine, 0); + let Some(file) = processor.claim_file_by_id(file_id).await? else { + info!("File {} is not pending; immediate parse claim skipped", file_id); + return Ok(()); + }; + let result = processor.process_file_claimed_skip_reuse(&file).await; + if let Err(err) = &result { + if let Err(cleanup_err) = processor.cleanup_processing_file_data_with_retry(file_id, 3).await { + error!("Failed to cleanup immediate parse data for file {}: {}", file_id, cleanup_err); + } + processor.mark_file_failed(file_id, &err.to_string()).await?; + } + result +} + +pub async fn try_reuse_file_with_file( + pool: SqlitePool, search_engine: SearchEngine, file: File, +) -> anyhow::Result { + if is_parse_paused() { + return Ok(false); + } + let processor = FileProcessor::new(pool, search_engine, 0); + let Some(claimed_file) = processor.claim_file_by_id(file.id).await? else { + return Ok(false); + }; + let result = processor.try_reuse_existing_data(&claimed_file).await; + match result { + Ok(true) => Ok(true), + Ok(false) => { + sqlx::query( + "UPDATE files SET status = 0, parse_run_id = NULL, updated_at = strftime('%s','now') \ + WHERE id = ? AND status = 2", + ) + .bind(file.id) + .execute(&processor.pool) + .await?; + Ok(false) + } + Err(err) => { + sqlx::query( + "UPDATE files SET status = 0, parse_run_id = NULL, log = ?, updated_at = strftime('%s','now') \ + WHERE id = ? AND status = 2", + ) + .bind(format!("Reuse failed: {}", err)) + .bind(file.id) + .execute(&processor.pool) + .await?; + Err(err) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn immediate_parse_claim_is_atomic() -> anyhow::Result<()> { + let temp = tempfile::NamedTempFile::new()?; + let url = format!("sqlite://{}", temp.path().display()); + let pool = sqlx::sqlite::SqlitePoolOptions::new().max_connections(4).connect(&url).await?; + sqlx::query( + "CREATE TABLE files ( + id INTEGER PRIMARY KEY, + user_id TEXT NOT NULL, user_name TEXT NOT NULL, hash TEXT NOT NULL, + filename TEXT NOT NULL, path TEXT NOT NULL, size INTEGER NOT NULL, + tags TEXT NOT NULL, status INTEGER NOT NULL, log TEXT NOT NULL, + slice_type TEXT NOT NULL, kb_id INTEGER, is_public INTEGER NOT NULL, + meta TEXT, summary TEXT, parse_priority INTEGER NOT NULL, parse_run_id TEXT, artifact_id INTEGER, + created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL + )", + ) + .execute(&pool) + .await?; + sqlx::query( + "INSERT INTO files VALUES (1, 'u', 'n', 'h', 'f.txt', '/tmp/f', 1, '', 0, '', 'text', NULL, 0, NULL, NULL, 50, NULL, NULL, 1, 1)", + ) + .execute(&pool) + .await?; + + let (first, second) = tokio::join!(claim_pending_file_by_id(&pool, 1), claim_pending_file_by_id(&pool, 1)); + let claimed = usize::from(first?.is_some()) + usize::from(second?.is_some()); + assert_eq!(claimed, 1); + + let (status, run_id): (i32, Option) = + sqlx::query_as("SELECT status, parse_run_id FROM files WHERE id = 1").fetch_one(&pool).await?; + assert_eq!(status, 2); + assert!(run_id.is_some()); + Ok(()) + } + + fn image_content_item(img_path: &str) -> ContentItem { + ContentItem { + typ: "image".to_string(), + bbox: vec![1, 2, 3, 4], + page_idx: 0, + text: None, + text_level: None, + text_format: None, + img_path: Some(img_path.to_string()), + image_caption: None, + table_body: None, + table_caption: None, + } + } + + #[test] + fn custom_parse_normalization_preserves_image_names() -> anyhow::Result<()> { + let mut images = HashMap::new(); + images.insert("img.png".to_string(), "raw-base64".to_string()); + let data = CustomParseData { + slices: vec![CustomSlice { + content: "see ![x](/api/v1/knowledge/files/img.png)".to_string(), + positions: Vec::new(), + }], + full_content: Some("full /api/v1/knowledge/files/img.png".to_string()), + summary: None, + images: Some(images), + content_list: Some(vec![image_content_item("images/img.png")]), + }; + + let normalized = FileProcessor::normalize_custom_parse_data(data)?; + + assert!(normalized.images.contains_key("img.png")); + assert_eq!(normalized.content_list.as_ref().unwrap()[0].img_path.as_deref(), Some("img.png")); + assert!(normalized.slices[0].content.contains("/api/v1/knowledge/files/img.png")); + assert!(normalized.full_content.as_ref().unwrap().contains("/api/v1/knowledge/files/img.png")); + assert_eq!(normalized.image_paths, vec!["img.png".to_string()]); + Ok(()) + } + + #[test] + fn custom_parse_normalization_preserves_unmapped_slice_refs() -> anyhow::Result<()> { + let data = CustomParseData { + slices: vec![CustomSlice { + content: "legacy ![x](/api/v1/knowledge/files/legacy.png)".to_string(), + positions: Vec::new(), + }], + full_content: None, + summary: None, + images: None, + content_list: None, + }; + + let normalized = FileProcessor::normalize_custom_parse_data(data)?; + + assert!(normalized.images.is_empty()); + assert!(normalized.slices[0].content.contains("/api/v1/knowledge/files/legacy.png")); + assert_eq!(normalized.image_paths, vec!["legacy.png".to_string()]); + Ok(()) + } + + #[test] + fn reuse_removes_legacy_image_id_prefix() { + assert_eq!(FileProcessor::remove_image_id_prefix("f42_img.png", 42), "img.png"); + assert_eq!(FileProcessor::remove_image_id_prefix("images/f42_img.png", 42), "images/img.png"); + assert_eq!(FileProcessor::remove_image_id_prefix("img.png", 42), "img.png"); + assert_eq!(FileProcessor::remove_image_id_prefix("f7_img.png", 42), "f7_img.png"); + } + + #[test] fn unsupported_mineru_image_formats_require_conversion() { for filename in ["icon.ico", "diagram.svg", "photo.avif", "photo.bmp"] { - assert!(crate::search::is_image_file(filename)); - assert!(FileProcessor::requires_mineru_image_conversion(filename)); - } - assert!(!FileProcessor::requires_mineru_image_conversion("photo.png")); - } - - #[test] + assert!(crate::search::is_image_file(filename)); + assert!(FileProcessor::requires_mineru_image_conversion(filename)); + } + assert!(!FileProcessor::requires_mineru_image_conversion("photo.png")); + } + + #[test] fn image_descriptions_are_only_appended_once() { - let mut items = vec![image_content_item("upload.gif")]; - let descriptions = HashMap::from([("upload.gif".to_string(), "animated image".to_string())]); - - enrich_content_list_with_image_descriptions(&mut items, &descriptions); - enrich_content_list_with_image_descriptions(&mut items, &descriptions); - - let captions = items[0].image_caption.as_ref().expect("image caption should be present"); - assert_eq!(captions.len(), 1); + let mut items = vec![image_content_item("upload.gif")]; + let descriptions = HashMap::from([("upload.gif".to_string(), "animated image".to_string())]); + + enrich_content_list_with_image_descriptions(&mut items, &descriptions); + enrich_content_list_with_image_descriptions(&mut items, &descriptions); + + let captions = items[0].image_caption.as_ref().expect("image caption should be present"); + assert_eq!(captions.len(), 1); assert_eq!(captions[0], "animated image"); } diff --git a/src/search/mod.rs b/src/search/mod.rs index ee2899c..f444f4c 100644 --- a/src/search/mod.rs +++ b/src/search/mod.rs @@ -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 = Lazy::new(Client::new); +/// Fuse ranked recall lists without comparing incompatible raw score scales. +fn reciprocal_rank_fusion(result_lists: impl IntoIterator>) -> Vec { + let mut fused: HashMap = 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 = 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 = 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 = 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 = 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![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![10, 11]); + assert!(fused[0].score > fused[1].score); + } +} diff --git a/src/search/tantivy_engine.rs b/src/search/tantivy_engine.rs index 0031f4b..7f28f5b 100644 --- a/src/search/tantivy_engine.rs +++ b/src/search/tantivy_engine.rs @@ -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, + index_writer: &mut tantivy::IndexWriter, schema: &Schema, docs: impl IntoIterator, ) -> tantivy::Result { let mut count = 0_usize; for doc in docs { diff --git a/src/weknora_sync.rs b/src/weknora_sync.rs new file mode 100644 index 0000000..9d91c89 --- /dev/null +++ b/src/weknora_sync.rs @@ -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 = + Lazy::new(|| Client::builder().timeout(Duration::from_secs(120)).build().expect("build WeKnora sync HTTP client")); + +#[derive(Debug, Serialize)] +struct ExternalChunk { + external_slice_id: String, + index: usize, + content: String, +} + +#[derive(Debug, Serialize)] +struct ExternalImport { + source: &'static str, + external_file_id: String, + source_version: String, + title: String, + filename: String, + file_type: String, + file_size: i64, + file_hash: String, + chunks: Vec, +} + +#[derive(Debug, FromRow)] +struct SyncEvent { + event_id: String, + event_type: String, + file_id: i64, + htknow_kb_id: i64, + attempts: i64, +} + +pub async fn enqueue_file_upsert(pool: &SqlitePool, file_id: i64) -> Result<()> { + if !sync_enabled() { + return Ok(()); + } + let kb_id: i64 = sqlx::query_scalar("SELECT kb_id FROM files WHERE id = ?") + .bind(file_id) + .fetch_optional(pool) + .await? + .flatten() + .ok_or_else(|| anyhow!("file {} has no knowledge base", file_id))?; + enqueue(pool, "upsert", file_id, kb_id).await +} + +pub async fn enqueue_file_delete(pool: &SqlitePool, file_id: i64, htknow_kb_id: i64) -> Result<()> { + if !sync_enabled() { + return Ok(()); + } + enqueue(pool, "delete", file_id, htknow_kb_id).await +} + +async fn enqueue(pool: &SqlitePool, event_type: &str, file_id: i64, htknow_kb_id: i64) -> Result<()> { + let event_id = uuid::Uuid::new_v4().to_string(); + sqlx::query( + "INSERT INTO weknora_sync_outbox + (event_id, event_type, file_id, htknow_kb_id, attempts, next_attempt_at, last_error, created_at, updated_at) + VALUES (?, ?, ?, ?, 0, strftime('%s','now'), '', strftime('%s','now'), strftime('%s','now')) + ON CONFLICT(file_id) DO UPDATE SET + event_id = excluded.event_id, + event_type = excluded.event_type, + htknow_kb_id = excluded.htknow_kb_id, + attempts = 0, + next_attempt_at = excluded.next_attempt_at, + last_error = '', + updated_at = excluded.updated_at", + ) + .bind(event_id) + .bind(event_type) + .bind(file_id) + .bind(htknow_kb_id) + .execute(pool) + .await?; + Ok(()) +} + +pub fn start_worker(pool: SqlitePool) { + if !sync_enabled() { + info!("WeKnora synchronization is disabled"); + return; + } + tokio::spawn(async move { + loop { + if let Err(err) = process_due_events(&pool).await { + warn!("WeKnora outbox worker failed: {}", err); + } + tokio::time::sleep(Duration::from_secs(2)).await; + } + }); +} + +async fn process_due_events(pool: &SqlitePool) -> Result<()> { + let events: Vec = sqlx::query_as( + "SELECT event_id, event_type, file_id, htknow_kb_id, attempts + FROM weknora_sync_outbox + WHERE next_attempt_at <= strftime('%s','now') + ORDER BY created_at, file_id + LIMIT 10", + ) + .fetch_all(pool) + .await?; + for event in events { + let result = match event.event_type.as_str() { + "upsert" => sync_file(pool, event.file_id).await, + "delete" => delete_file(event.file_id, event.htknow_kb_id).await, + other => Err(anyhow!("unsupported outbox event type {}", other)), + }; + match result { + Ok(()) => { + sqlx::query("DELETE FROM weknora_sync_outbox WHERE file_id = ? AND event_id = ?") + .bind(event.file_id) + .bind(&event.event_id) + .execute(pool) + .await?; + info!("WeKnora {} accepted for file {}", event.event_type, event.file_id); + } + Err(err) => { + let attempts = event.attempts + 1; + let delay = (2_i64.saturating_pow(attempts.min(10) as u32)).min(3600); + sqlx::query( + "UPDATE weknora_sync_outbox + SET attempts = ?, next_attempt_at = strftime('%s','now') + ?, last_error = ?, updated_at = strftime('%s','now') + WHERE file_id = ? AND event_id = ?", + ) + .bind(attempts) + .bind(delay) + .bind(err.to_string()) + .bind(event.file_id) + .bind(&event.event_id) + .execute(pool) + .await?; + warn!( + "WeKnora {} attempt {} failed for file {}; retry in {}s: {}", + event.event_type, attempts, event.file_id, delay, err + ); + } + } + } + Ok(()) +} + +fn sync_enabled() -> bool { + std::env::var("HTKNOW_WEKNORA_URL").ok().is_some_and(|v| !v.trim().is_empty()) +} + +fn mapped_weknora_kb_id(htknow_kb_id: i64) -> Result { + let kb_map_raw = std::env::var("HTKNOW_WEKNORA_KB_MAP").context("HTKNOW_WEKNORA_KB_MAP is required")?; + let kb_map: HashMap = + serde_json::from_str(&kb_map_raw).context("parse HTKNOW_WEKNORA_KB_MAP JSON")?; + kb_map + .get(&htknow_kb_id.to_string()) + .or_else(|| kb_map.get("*")) + .cloned() + .ok_or_else(|| anyhow!("no WeKnora KB mapping for htknow KB {} and no '*' default mapping", htknow_kb_id)) +} + +async fn sync_file(pool: &SqlitePool, file_id: i64) -> Result<()> { + let file: File = sqlx::query_as(&format!("SELECT {} FROM files WHERE id = ?", FILE_COLS_NO_CONTENT)) + .bind(file_id) + .fetch_one(pool) + .await + .context("load parsed file")?; + if file.status != 1 { + return Err(anyhow!("file status is {}, expected completed", file.status)); + } + let htknow_kb_id = file.kb_id.ok_or_else(|| anyhow!("file has no knowledge base"))?; + let weknora_kb_id = mapped_weknora_kb_id(htknow_kb_id)?; + + let source_id = crate::api::effective_parse_file_id(pool, file_id).await?; + let ids: Vec = sqlx::query_scalar("SELECT id FROM slices WHERE file_id = ? ORDER BY ordinal, id") + .bind(source_id) + .fetch_all(pool) + .await?; + let contents = crate::slice_content::read_all(source_id).await?; + let chunks = ids + .into_iter() + .enumerate() + .filter_map(|(index, id)| { + let content = contents.get(&id)?.trim(); + if content.is_empty() { + return None; + } + Some(ExternalChunk { external_slice_id: id.to_string(), index, content: content.to_string() }) + }) + .collect::>(); + if chunks.is_empty() { + return Err(anyhow!("parsed file contains no text chunks")); + } + + let mut version_hasher = Sha256::new(); + version_hasher.update(file.filename.as_bytes()); + version_hasher.update(file.hash.as_bytes()); + for chunk in &chunks { + version_hasher.update(chunk.external_slice_id.as_bytes()); + version_hasher.update([0]); + version_hasher.update(chunk.content.as_bytes()); + version_hasher.update([0xff]); + } + let source_version = hex::encode(version_hasher.finalize()); + let file_type = + std::path::Path::new(&file.filename).extension().and_then(|v| v.to_str()).unwrap_or("markdown").to_lowercase(); + let payload = ExternalImport { + source: "htknow", + external_file_id: file.id.to_string(), + source_version, + title: file.filename.clone(), + filename: file.filename, + file_type, + file_size: file.size, + file_hash: file.hash, + chunks, + }; + let base = std::env::var("HTKNOW_WEKNORA_URL")?; + let url = + format!("{}/api/v1/knowledge-bases/{}/knowledge/external-chunks", base.trim_end_matches('/'), weknora_kb_id); + let mut request = CLIENT.post(url).json(&payload); + if let Ok(key) = std::env::var("HTKNOW_WEKNORA_API_KEY") + && !key.trim().is_empty() + { + request = request.header("X-API-Key", key); + } + let response = request.send().await.context("call WeKnora external chunk import")?; + let status = response.status(); + if !status.is_success() { + let body = response.text().await.unwrap_or_default(); + return Err(anyhow!("WeKnora returned {}: {}", status, body)); + } + Ok(()) +} + +async fn delete_file(file_id: i64, htknow_kb_id: i64) -> Result<()> { + let weknora_kb_id = mapped_weknora_kb_id(htknow_kb_id)?; + let base = std::env::var("HTKNOW_WEKNORA_URL")?; + let url = format!( + "{}/api/v1/knowledge-bases/{}/knowledge/external-chunks/htknow/{}", + base.trim_end_matches('/'), + weknora_kb_id, + file_id + ); + let mut request = CLIENT.delete(url); + if let Ok(key) = std::env::var("HTKNOW_WEKNORA_API_KEY") + && !key.trim().is_empty() + { + request = request.header("X-API-Key", key); + } + let response = request.send().await.context("call WeKnora external knowledge delete")?; + let status = response.status(); + if !status.is_success() { + let body = response.text().await.unwrap_or_default(); + return Err(anyhow!("WeKnora returned {}: {}", status, body)); + } + Ok(()) +} diff --git a/tests/api_concurrency.rs b/tests/api_concurrency.rs index ac404d5..bd749af 100644 --- a/tests/api_concurrency.rs +++ b/tests/api_concurrency.rs @@ -1,6 +1,7 @@ mod common; use std::{ - fs, time::{Duration, Instant} + fs, + time::{Duration, Instant}, }; use axum::http::StatusCode; diff --git a/tests/common/mod.rs b/tests/common/mod.rs index cd21534..a0a9f9c 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -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;