diff --git a/common/pool.go b/common/pool.go new file mode 100644 index 0000000..6b0df19 --- /dev/null +++ b/common/pool.go @@ -0,0 +1,31 @@ +package common + +import ( + "context" + + "rag-local/kb/consts" + + "github.com/gogf/gf/v2/frame/g" + "github.com/gogf/gf/v2/os/grpool" +) + +// 各并行点协程池。池内任务只做读查询与 LLM/Embedding 调用（纯 IO）， +// SQLite 写一律收敛回主 goroutine 串行执行（business.db 无 WAL，并发写会 database is locked）。 +// 防死锁：被等待的池（AnnotationDatasetPool / ChatRetrievePool）任务内不得再等待任何池， +// 等待链单向「主 → A池 → B池」；池无 Wait 方法，等待用调用方的 sync.WaitGroup。 +var ( + KgExtractPool = newPool(consts.KgExtractPoolSize, "kg_extract") + AnnotationClausePool = newPool(consts.AnnotationClausePoolSize, "annotation_clause") + AnnotationDatasetPool = newPool(consts.AnnotationDatasetPoolSize, "annotation_dataset") + ChatPool = newPool(consts.ChatPoolSize, "chat") + ChatRetrievePool = newPool(consts.ChatRetrievePoolSize, "chat_retrieve") +) + +// newPool 从 config pool. 读取池大小（<1 或缺失时回退默认值 def） +func newPool(def int, key string) *grpool.Pool { + size := g.Cfg().MustGet(context.Background(), "pool."+key, def).Int() + if size < 1 { + size = def + } + return grpool.New(size) +} diff --git a/config.yml b/config.yml index 04ef6ae..0d18bb5 100644 --- a/config.yml +++ b/config.yml @@ -28,3 +28,11 @@ vector: chat: timeout: 600 # 对话模型API请求超时时间(秒) max_retries: 3 # 请求失败最大重试次数 + +# 协程池配置（各并行点并发度，缺失或 <1 时回退代码内默认值） +pool: + kg_extract: 4 # 知识图谱：逐 chunk LLM 抽取 + annotation_clause: 4 # 合同标注：逐条款 (recall+judge) + annotation_dataset: 8 # 条款内：逐 dataset 召回 + chat: 4 # 问答：retrieve 与 GraphEnhance 并行 + chat_retrieve: 4 # 检索：vec 与 fts 并行 diff --git a/data/business.db b/data/business.db index 03e17f7..98c8824 100644 Binary files a/data/business.db and b/data/business.db differ diff --git a/docs/实现方案.md b/docs/实现方案.md index d1c0830..92794fe 100644 --- a/docs/实现方案.md +++ b/docs/实现方案.md @@ -131,6 +131,7 @@ rag-local/ │ ├── text_parser.go # txt / md │ ├── tokenizer.go # gse 中文分词（写索引/检索共用） │ ├── util.go # RandomToken（随机文件名/令牌）、TokenFingerprint（SHA-256 指纹） +│ ├── pool.go # grpool 协程池单例（5 池，config pool 段配置大小，见 §7.9） │ └── parser_test.go # 解析器单元测试 ├── kb/ # 业务模块（knowledge base，对应 video-factory 的 shortdrama） │ ├── consts/ @@ -143,7 +144,7 @@ rag-local/ │ │ ├── dto/ # 11 个文件（chunk/contract 各含多组 Req/Res；含 g.Meta 路由定义） │ │ └── domain/ # 领域对象（RAG 检索结果、引用来源、流式事件等） │ ├── dao/ # 13 个文件，与实体表一一对应（chunk_dao 兼管 vec0/fts5 虚拟表） -│ ├── service/ # 12 个文件，与 dao 对应（文档流水线、RAG 问答、合同标注编排归入对应 service） +│ ├── service/ # 12 个文件，与 dao 对应（文档流水线、RAG 问答、合同标注编排归入对应 service；grpool 协程池单例在 common/pool.go，见 §7.9） │ ├── controller/ # 11 个文件，与 service 对应（合同标注路由在 contract_controller） ├── ui-src/ # Vue 3 前端工程 │ ├── package.json @@ -464,7 +465,7 @@ CREATE INDEX IF NOT EXISTS idx_chat_message_conversation ON chat_message(convers 7. **SQLite 并发**：3 个库文件各自独立连接；写操作集中在任务轮询 goroutine 与用户操作。**当前未启用 WAL / busy_timeout**（代码无 PRAGMA），并发写 business.db 偶发 `database is locked (5)`，缓解措施与实际缺口见 §5.6 / §14.3。 8. **文档状态机 6 态**：解析 → 向量生成 → 图谱构建分三段推进（`0 待处理/1 解析中/2 向量生成中/3 图谱构建中/4 已完成/5 失败`）。图谱抽取失败**不阻断**完成：文档仍置 4，`error_msg` 记录「知识图谱未构建」原因，前端以黄色标签提示（不再出现"显示已完成但图谱没建完"的假象）。 9. **合同标注宁滥毋缺**：标注业务的召回策略与问答相反——问答要精（topK=5 + 重排门槛 max(最高分×50%, 6)），标注宁滥毋缺（漏标比多标严重）。召回放宽（每数据集向量+FTS 各 15 条）、不做重排门槛、全部候选交 LLM 判定后保留（含 0 分），见 §7.8。 -10. **轮询任务并发模型**：`StartParsePoller` 与 `StartAnnotationPoller` 各自**单 goroutine 串行**消费（3 秒间隔），不并发处理多个任务，避免 SQLite 写冲突；任务粒度（kb_parse_task / kb_contract_task）+ 子粒度（clause）断点续跑。 +10. **轮询任务并发模型**：`StartParsePoller` 与 `StartAnnotationPoller` 各自**gtimer 单例定时器串行**消费（5 秒间隔，job 未结束不重入），不并发处理多个任务，避免 SQLite 写冲突；任务粒度（kb_parse_task / kb_contract_task）+ 子粒度（clause）断点续跑。任务内部热点（kg 逐 chunk 抽取、标注逐条款、多数据集召回、问答双路检索+图增强）用 **grpool 协程池并行化**，池大小 config.yml `pool` 段配置——并行段只做读查询与 LLM/Embedding 调用，SQLite 写全部收敛回主 goroutine 串行（见 §7.9）。 --- @@ -744,6 +745,7 @@ score = Σ(1 / (60 + rank)) // RRF 分区间过窄无区分度，仅用于候 ``` - `chat_service.go` 中直接编排 eino 组件：`ChatService.Ask` = HybridRetriever.Retrieve → 组装引用列表 + 系统提示词 → OpenAIChatModel.Stream 流式生成（组件已实现 eino 接口，可随时迁入 graph/chain 拓扑；eino v0.9.13 的 graph 节点类型约束与"中间取引用"需求不匹配，故直接编排） +- **并行化（§7.9）**：`Ask` 中 `retrieve`（common.ChatPool）与 `GraphEnhance`（纯读）并行执行，汇合后再组装提示词；`HybridRetriever.Retrieve` 内向量段与 FTS 段（common.ChatRetrievePool）并行，RRF 融合、LLM 重排仍由主 goroutine 串行 - `message_service.go`：会话解析/创建、用户消息落库、历史消息组装（最近 10 轮）、助手消息 + citations JSON 落库 - Prompt 模板： ``` @@ -773,7 +775,7 @@ POST /document/upload │ 插入 kb_document(status=0) + kb_parse_task(status=0) │ （上传不校验向量模型——数据集在 Save 时已强制绑定，见下） ▼ -StartParsePoller（main.go 启动，3 秒轮询，单 goroutine 串行） +StartParsePoller（main.go 启动，gtimer 单例 5 秒轮询，job 未结束不重入） │ 取 status=0 任务 → 置 status=1（解析中） │ 1. 校验 embedding 配置 → 构建 OpenAIEmbedder（无配置 → 任务失败「数据集未绑定向量模型」； │ 数据集 Save 强制绑定向量模型「数据集必须绑定向量模型，请先选择向量模型」，此处为防御性校验） @@ -787,8 +789,9 @@ StartParsePoller（main.go 启动，3 秒轮询，单 goroutine 串行） │ c. 语义分块：eino semantic splitter（MinChunkSize=chunkSize/2） │ d. 递归兜底：eino recursive（KeepTypeEnd） │ 5. 批量 Embedding → InsertAll（chunk+vec0+FTS5 同事务，按 EmbedBatchSize=16 分批） - │ 6. 知识图谱抽取（document → status=3 图谱构建中；逐 chunk LLM 抽取，失败不阻断， - │ error_msg 记录「知识图谱未构建」原因，前端黄色标签提示） + │ 6. 知识图谱抽取（document → status=3 图谱构建中；逐 chunk LLM 抽取并发执行 + │ （kg_extract 池，见 §7.9），失败不阻断，error_msg 记录「知识图谱未构建」原因， + │ 前端黄色标签提示） │ 7. 完成：更新 document.status=4（已完成）、chunk_count；任务 status=2 │ 失败 → document.status=5 + error_msg，前端可重试 ``` @@ -800,7 +803,8 @@ StartParsePoller（main.go 启动，3 秒轮询，单 goroutine 串行） **构建（LLM 抽取，挂在解析流水线向量化之后、完成之前——对应文档状态机第 3 态「图谱构建中」）**： ``` -解析 → 分块 → 向量化 ──▶ 批量抽取（每 chunk 一次 LLM 调用，JSON 输出） +解析 → 分块 → 向量化 ──▶ 批量抽取（每 chunk 一次 LLM 调用，JSON 输出；经 kg_extract 池并发， + 池内只做 LLM 调用，upsert/insert 收敛回主 goroutine 串行，见 §7.9） ▼ {entities:[{name, type}], relations:[{head, relation, tail}]} ▼ @@ -836,12 +840,13 @@ StartParsePoller（main.go 启动，3 秒轮询，单 goroutine 串行） **与问答检索的本质差异——宁滥毋缺**：问答策略 `HybridRetriever`（topK=5 截断 + 重排门槛 max(最高分×50%, 6)）不适用于标注——topK 截断会漏标、6 分门槛会杀光「段落对段落」的弱相关匹配。标注业务**漏标比多标严重**：召回放宽（`AnnoRecallTopK=15`/数据集，向量+FTS 各 15 条）、RRF 融合截 60 候选、不做重排门槛、全部候选交 LLM 判定后保留（含 0 分）。 -**流水线**（`annotation_service.go`，`StartAnnotationPoller` 单 goroutine 3 秒轮询）： +**流水线**（`annotation_service.go`，`StartAnnotationPoller` gtimer 单例 5 秒轮询）： ``` 上传合同 → 任务落库(pending) → poller 取任务 → ParseFile 解析文本 → 正则切分条款 -→ 逐条款：多 dataset 各召回(Vec 15 + FTS 15) → RRF 融合截 60 候选 -→ LLM 一次调用判定(0-10 分 + 理由, JSON) → 全保留按分降序落库 → 更新进度 +→ 逐条款并发（annotation_clause 池）：多 dataset 各召回(Vec 15 + FTS 15，annotation_dataset 池并行) +→ RRF 融合截 60 候选 → LLM 一次调用判定(0-10 分 + 理由, JSON) +→ 主 goroutine 串行落库（清旧标→插 mark→置状态→更新进度） → 全部条款完成 → 任务 done（断点：重启后按 clause 状态续跑，已 done 条款跳过） ``` @@ -853,6 +858,30 @@ StartParsePoller（main.go 启动，3 秒轮询，单 goroutine 串行） - **断点续跑**：任务按 `status IN (0,1)` 领取，clause 粒度续跑（已 done 不重复产生 mark；重跑任务先 DeleteByClause 幂等重建） - **导出**：`AnnotatedHTML` 生成自包含 HTML（条款 + 内嵌标注，score ≥8 绿 / ≥5 蓝 / 其余灰，打印按钮 `window.print()`）；controller 直接写响应体（中间件检测已写入则不包装 JSON），前端原生 fetch + 手动 Authorization 获取 blob +### 7.9 并行化设计（grpool 协程池） + +**背景**：四个串行热点用 GoFrame `grpool` 并行化，池大小在 config.yml `pool` 段配置（缺失或 <1 回退代码内默认值，见 `kb/consts/consts.go`）： + +| 池 | 配置键 | 默认 | 作用点 | +|---|---|---|---| +| 知识图谱抽取 | `pool.kg_extract` | 4 | 文档解析第 6 步：逐 chunk LLM 抽取 | +| 合同标注-条款 | `pool.annotation_clause` | 4 | 标注流水线：逐条款 (recall+judge) | +| 合同标注-数据集 | `pool.annotation_dataset` | 8 | 条款内：逐 dataset 召回（向量化+vec+fts） | +| 问答编排 | `pool.chat` | 4 | Ask：retrieve 与 GraphEnhance 并行 | +| 问答检索 | `pool.chat_retrieve` | 4 | HybridRetriever：vec 段与 fts 段并行 | + +**两条铁律**： + +1. **写收敛**：business.db 无 WAL / busy_timeout（§14.3），多 goroutine 并发写会 `database is locked (5)`（已实测）。因此并发段只做**读查询与 LLM/Embedding 调用（纯 IO）**，所有 SQLite 写（状态更新、进度、落库）一律收敛回主 goroutine 串行执行——锁定风险归零，并发只赢在 IO 等待上。 +2. **死锁防护**：池内任务若等待同一池的任务会饿死（worker 全部阻塞在等待上）。被等待的池其任务必须**无嵌套等待**，等待链单向「主 → A池 → B池」（如 主→common.AnnotationClausePool→common.AnnotationDatasetPool、主→common.ChatPool→common.ChatRetrievePool）。`grpool` 无 Wait 方法，等待一律用调用方 `sync.WaitGroup`；任务结果经 buffered channel 回主 goroutine。 + +**实现形态**（`kb/service/pool.go`，包 init 从 `g.Cfg().MustGet("pool.", 默认值)` 读取；`grpool.New(limit)` 建池，`Pool.AddWithRecover` 提交并防 panic）： + +- 知识图谱（§7.7）：`ExtractDocument` 拆 `callExtract`（池内 LLM+JSON 解析，不写库）/ `saveExtract`（主 goroutine 串行 upsert/insert），「发一批、收一批」循环 +- 合同标注（§7.8）：`processOne` 条款循环提交「召回+判定」入池，主 goroutine 收结果串行清旧标/插 mark/置状态/更新进度；`recallCandidates` 内逐 dataset 并行召回（`recallOneDataset`），RRF 融合回主 goroutine +- 问答（§7.3/§7.4）：`Ask` 中 `retrieve`（ChatPool，内部再并行 vec/fts）与 `GraphEnhance`（纯读，主 goroutine 直接跑）并行；`Retrieve` 内 vec 段与 fts 段并行（`vecRetrieve`/`ftsRetrieve`），RRF 合并、LLM 重排保持串行 +- 任务级仍由轮询器单 goroutine 串行消费（决策 10），并行只发生在任务内部 + --- ## 8. 鉴权设计（单用户 + 启动令牌） @@ -1073,7 +1102,7 @@ FTS5 召回依赖 gse 分词质量；专有名词（人名/产品名）可能切 ### 14.3 SQLite 写并发 -解析任务轮询与用户操作可能并发写 business.db，偶发 `database is locked (5)`（**已实测，未根治**）。当前缓解：两个 poller 各自单 goroutine 串行消费 + 写操作集中到 service/dao 单事务（参照 video-factory 单机场景）。**改进方向（未实现）**：连接串初始化时执行 `PRAGMA journal_mode=WAL` 与 `PRAGMA busy_timeout=5000`（modernc 驱动支持），或 config.yml `database` 段配置 `busy_timeout` 后重试机制。 +解析任务轮询与用户操作可能并发写 business.db，偶发 `database is locked (5)`（**已实测，未根治**）。当前缓解（三层）：① 两个 poller 各自单 goroutine 串行消费任务；② 任务内并行段（grpool 协程池，§7.9）只做读查询与 LLM/Embedding 调用，**SQLite 写一律收敛回主 goroutine 串行执行**——并发不会引入新的写竞争；③ 写操作集中到 service/dao 单事务（参照 video-factory 单机场景）。**改进方向（未实现）**：连接串初始化时执行 `PRAGMA journal_mode=WAL` 与 `PRAGMA busy_timeout=5000`（modernc 驱动支持），或 config.yml `database` 段配置 `busy_timeout` 后重试机制。 ### 14.4 切换 embedding 模型 @@ -1090,5 +1119,5 @@ FTS5 召回依赖 gse 分词质量；专有名词（人名/产品名）可能切 | AI 层 | 自研 ReAct agent（openai 直连） | **Eino** 组件化编排（Indexer/Retriever/Graph） | | 鉴权 | 多用户 + 角色 + JWT（user/login_log 表） | **单用户访问令牌**：每次启动重新生成（内存持有不落库）打印，登录换取 JWT，无用户表 | | 前端 | HTML + 少量 Vue | 纯 Vue 3 SPA（构建产物 Go 托管） | -| 异步任务 | 视频生成轮询 | 文档解析/向量化 + 合同条款标注 双轮询（各单 goroutine 串行） | +| 异步任务 | 视频生成轮询 | 文档解析/向量化 + 合同条款标注 双轮询（任务级单 goroutine 串行，任务内热点用 grpool 协程池并行，见 §7.9） | | 其余 | 分层/路由/鉴权/缓存/部署模式 | 完全对齐 | diff --git a/kb/consts/consts.go b/kb/consts/consts.go index 967dc8a..a8434d7 100644 --- a/kb/consts/consts.go +++ b/kb/consts/consts.go @@ -20,7 +20,7 @@ const ( SettingsKeyChunkSize = "chunk_default_size" SettingsKeyChunkOverlap = "chunk_default_overlap" - ParsePollIntervalSeconds = 3 // 解析任务轮询间隔 + ParsePollIntervalSeconds = 5 // 解析/标注任务轮询间隔 EmbedBatchSize = 16 // 单次向量化请求的文本批量 @@ -28,4 +28,11 @@ const ( AnnoRecallTopK = 15 // 每数据集向量+FTS 各召回数（标注宁滥毋缺，放宽召回） AnnoMaxCandidates = 60 // 多数据集融合后的候选上限（喂给 LLM 判定） AnnoMaxClauseChars = 2000 // 合同条款全文上限（超长截断，控制 prompt） + + // 协程池默认大小（config.yml pool 段缺失时兜底） + KgExtractPoolSize = 4 // 知识图谱：逐 chunk LLM 抽取 + AnnotationClausePoolSize = 4 // 合同标注：逐条款 (recall+judge) + AnnotationDatasetPoolSize = 8 // 条款内：逐 dataset 召回 + ChatPoolSize = 4 // 问答：retrieve 与 GraphEnhance 并行 + ChatRetrievePoolSize = 4 // 检索：vec 与 fts 并行 ) diff --git a/kb/dao/kg_entity_dao.go b/kb/dao/kg_entity_dao.go index 3035d47..b68e692 100644 --- a/kb/dao/kg_entity_dao.go +++ b/kb/dao/kg_entity_dao.go @@ -87,16 +87,41 @@ func (d *kgEntityDao) ListNames(ctx context.Context, datasetId int64) ([]string, return names, nil } -// Upsert 按 (dataset_id, name) 去重，已存在则更新类型与来源 -func (d *kgEntityDao) Upsert(ctx context.Context, datasetId, chunkId int64, name, entityType string) error { +// KgEntityItem 批量写入条目 +type KgEntityItem struct { + Name string + EntityType string +} + +// UpsertBatch 按 (dataset_id, name) 去重批量写入，已存在则更新类型与来源（单条 SQL 多行 VALUES） +func (d *kgEntityDao) UpsertBatch(ctx context.Context, datasetId, chunkId int64, items []KgEntityItem) error { + if len(items) == 0 { + return nil + } now := gtime.Now().Format("Y-m-d H:i:s") - _, err := g.DB(consts.DbGroupDefault).Exec(ctx, `INSERT INTO `+consts.TableNameKgEntity+` - (dataset_id, name, entity_type, chunk_id, created_at, updated_at) - VALUES (?,?,?,?,?,?) - ON CONFLICT(dataset_id, name) DO UPDATE SET - entity_type=excluded.entity_type, chunk_id=excluded.chunk_id, updated_at=excluded.updated_at`, - datasetId, name, entityType, chunkId, now, now) - return err + // SQLite 变量数上限 999，按 100 行/条分片 + for start := 0; start < len(items); start += 100 { + end := min(start+100, len(items)) + var sb strings.Builder + sb.WriteString("INSERT INTO " + consts.TableNameKgEntity + ` + (dataset_id, name, entity_type, chunk_id, created_at, updated_at) + VALUES `) + args := make([]any, 0, (end-start)*6) + for i := start; i < end; i++ { + if i > start { + sb.WriteString(",") + } + sb.WriteString("(?,?,?,?,?,?)") + args = append(args, datasetId, items[i].Name, items[i].EntityType, chunkId, now, now) + } + sb.WriteString(` + ON CONFLICT(dataset_id, name) DO UPDATE SET + entity_type=excluded.entity_type, chunk_id=excluded.chunk_id, updated_at=excluded.updated_at`) + if _, err := g.DB(consts.DbGroupDefault).Exec(ctx, sb.String(), args...); err != nil { + return err + } + } + return nil } func (d *kgEntityDao) DeleteByChunkIds(ctx context.Context, chunkIds []int64) error { diff --git a/kb/dao/kg_relation_dao.go b/kb/dao/kg_relation_dao.go index 3265bb4..82f2af1 100644 --- a/kb/dao/kg_relation_dao.go +++ b/kb/dao/kg_relation_dao.go @@ -71,16 +71,39 @@ func (d *kgRelationDao) List(ctx context.Context, datasetId int64, page, pageSiz return list, total, err } -func (d *kgRelationDao) Insert(ctx context.Context, datasetId, chunkId int64, head, relation, tail string) error { - _, err := g.DB(consts.DbGroupDefault).Model(consts.TableNameKgRelation).Ctx(ctx).Data(g.Map{ - "dataset_id": datasetId, - "head": head, - "relation": relation, - "tail": tail, - "chunk_id": chunkId, - "created_at": gtime.Now().Format("Y-m-d H:i:s"), - }).Insert() - return err +// KgRelationItem 批量写入条目（三元组） +type KgRelationItem struct { + Head string + Relation string + Tail string +} + +// InsertBatch 批量写入三元组（单条 SQL 多行 VALUES） +func (d *kgRelationDao) InsertBatch(ctx context.Context, datasetId, chunkId int64, items []KgRelationItem) error { + if len(items) == 0 { + return nil + } + now := gtime.Now().Format("Y-m-d H:i:s") + // SQLite 变量数上限 999，按 100 行/条分片 + for start := 0; start < len(items); start += 100 { + end := min(start+100, len(items)) + var sb strings.Builder + sb.WriteString("INSERT INTO " + consts.TableNameKgRelation + ` + (dataset_id, head, relation, tail, chunk_id, created_at) + VALUES `) + args := make([]any, 0, (end-start)*6) + for i := start; i < end; i++ { + if i > start { + sb.WriteString(",") + } + sb.WriteString("(?,?,?,?,?,?)") + args = append(args, datasetId, items[i].Head, items[i].Relation, items[i].Tail, chunkId, now) + } + if _, err := g.DB(consts.DbGroupDefault).Exec(ctx, sb.String(), args...); err != nil { + return err + } + } + return nil } // Neighbors 一跳邻居：head 或 tail 命中实体名的三元组（实体链接用） diff --git a/kb/service/annotation_service.go b/kb/service/annotation_service.go index e7f8e8c..09dd80c 100644 --- a/kb/service/annotation_service.go +++ b/kb/service/annotation_service.go @@ -11,6 +11,7 @@ import ( "sort" "strconv" "strings" + "sync" "time" "rag-local/common" @@ -22,6 +23,7 @@ import ( "github.com/cloudwego/eino/schema" "github.com/gogf/gf/v2/errors/gerror" "github.com/gogf/gf/v2/frame/g" + "github.com/gogf/gf/v2/os/gtimer" ) var AnnotationService = &annotationService{} @@ -53,19 +55,12 @@ type annoCandidate struct { RrfScore float64 } -// StartAnnotationPoller 启动标注任务轮询：单 goroutine 串行消费 +// StartAnnotationPoller 启动标注任务轮询：gtimer 单例定时器串行消费（job 未结束不重入） func (s *annotationService) StartAnnotationPoller(ctx context.Context) { - go func() { - g.Log().Info(ctx, "annotation task poller started") - for { - select { - case <-ctx.Done(): - return - case <-time.After(consts.ParsePollIntervalSeconds * time.Second): - s.processOne(ctx) - } - } - }() + g.Log().Info(ctx, "annotation task poller started") + gtimer.AddSingleton(ctx, consts.ParsePollIntervalSeconds*time.Second, func(ctx context.Context) { + s.processOne(ctx) + }) } // SplitClauses 按行首标记切分条款；无结构时整篇作为单条 @@ -193,7 +188,16 @@ func (s *annotationService) processOne(ctx context.Context) { } } - failed := 0 + // 并行标注：召回+判定（纯读与 LLM 调用）提交 common.AnnotationClausePool， + // 落库（SQLite 写）收敛回主 goroutine 串行，避免无 WAL 下的 database is locked。 + type clauseJobOut struct { + clauseId int64 + marks []*entity.ContractMark + noCands bool + err error + } + ch := make(chan clauseJobOut, len(clauses)) + var wg sync.WaitGroup for _, cl := range clauses { if cl.Status == consts.TaskStatusDone { continue @@ -202,37 +206,59 @@ func (s *annotationService) processOne(ctx context.Context) { g.Log().Errorf(ctx, "mark clause running failed: %v", err) continue } - cands, err := s.recallCandidates(ctx, cl, dsIds, dsNames, embedders) - if err != nil { + wg.Add(1) + if err := common.AnnotationClausePool.AddWithRecover(ctx, func(ctx context.Context) { + defer wg.Done() + out := clauseJobOut{clauseId: cl.Id} + cands, err := s.recallCandidates(ctx, cl, dsIds, dsNames, embedders) + if err != nil { + out.err = err + ch <- out + return + } + if len(cands) == 0 { + out.noCands = true + ch <- out + return + } + out.marks, out.err = s.judgeClause(ctx, chatModel, cl, cands) + ch <- out + }, func(ctx context.Context, e error) { + defer wg.Done() + ch <- clauseJobOut{clauseId: cl.Id, err: e} + }); err != nil { + wg.Done() + g.Log().Errorf(ctx, "submit clause %d failed: %v", cl.Id, err) + } + } + go func() { wg.Wait(); close(ch) }() + + failed := 0 + for out := range ch { + if out.err != nil { failed++ - _ = dao.ContractClause.UpdateStatus(ctx, cl.Id, consts.TaskStatusFailed, err.Error()) + _ = dao.ContractClause.UpdateStatus(ctx, out.clauseId, consts.TaskStatusFailed, out.err.Error()) continue } - if len(cands) == 0 { + if out.noCands { // 无候选视为完成（无标注），避免卡住进度 - _ = dao.ContractClause.UpdateStatus(ctx, cl.Id, consts.TaskStatusDone, "") + _ = dao.ContractClause.UpdateStatus(ctx, out.clauseId, consts.TaskStatusDone, "") s.updateProgress(ctx, task.Id) continue } - marks, err := s.judgeClause(ctx, chatModel, cl, cands) - if err != nil { - failed++ - _ = dao.ContractClause.UpdateStatus(ctx, cl.Id, consts.TaskStatusFailed, err.Error()) - continue - } // 幂等：重跑前清旧标注，避免断点续跑产生重复 mark - if err := dao.ContractMark.DeleteByClause(ctx, cl.Id); err != nil { + if err := dao.ContractMark.DeleteByClause(ctx, out.clauseId); err != nil { g.Log().Warningf(ctx, "clear old marks failed: %v", err) } - for _, m := range marks { - m.ClauseId = cl.Id + for _, m := range out.marks { + m.ClauseId = out.clauseId } - if err := dao.ContractMark.InsertAll(ctx, marks); err != nil { + if err := dao.ContractMark.InsertAll(ctx, out.marks); err != nil { failed++ - _ = dao.ContractClause.UpdateStatus(ctx, cl.Id, consts.TaskStatusFailed, err.Error()) + _ = dao.ContractClause.UpdateStatus(ctx, out.clauseId, consts.TaskStatusFailed, err.Error()) continue } - _ = dao.ContractClause.UpdateStatus(ctx, cl.Id, consts.TaskStatusDone, "") + _ = dao.ContractClause.UpdateStatus(ctx, out.clauseId, consts.TaskStatusDone, "") s.updateProgress(ctx, task.Id) } @@ -262,39 +288,47 @@ func (s *annotationService) updateProgress(ctx context.Context, taskId int64) { } } -// recallCandidates 多数据集召回：每数据集向量+FTS 各取 AnnoRecallTopK，全局 RRF 融合截断 +// annoRecallHit 单数据集召回结果（排名用于 RRF 融合） +type annoRecallHit struct { + ChunkId int64 + DsId int64 + LawTitle string + Rank float64 +} + +// recallCandidates 多数据集召回：每数据集向量+FTS 各取 AnnoRecallTopK，全局 RRF 融合截断。 +// 各数据集召回（向量化+vec+fts，纯读）提交 common.AnnotationDatasetPool 并行，融合回主 goroutine 串行。 func (s *annotationService) recallCandidates(ctx context.Context, clause *entity.ContractClause, dsIds []int64, dsNames map[int64]string, embedders map[int64]*OpenAIEmbedder) ([]annoCandidate, error) { - merged := make(map[int64]*annoCandidate) clauseText := clause.Title + " " + clause.Content ftsText := clause.Title + " " + clause.Content if rs := []rune(ftsText); len(rs) > 200 { ftsText = string(rs[:200]) } ftsQuery := common.TokenizeQuery(ftsText) + + ch := make(chan []annoRecallHit, len(dsIds)) + var wg sync.WaitGroup for _, dsId := range dsIds { - if emb := embedders[dsId]; emb != nil { - vecs, err := emb.EmbedStrings(ctx, []string{clauseText}) - if err != nil { - g.Log().Warningf(ctx, "clause embed failed (dataset %d): %v", dsId, err) - } else if len(vecs) > 0 { - hits, err := dao.Chunk.VecSearch(ctx, dsId, domain.VecJsonF64(vecs[0]), consts.AnnoRecallTopK) - if err != nil { - g.Log().Warningf(ctx, "vec search failed (dataset %d): %v", dsId, err) - } else { - for i, h := range hits { - s.mergeHit(merged, h.ChunkId, dsId, dsNames[dsId], float64(i)) - } - } - } + wg.Add(1) + if err := common.AnnotationDatasetPool.AddWithRecover(ctx, func(ctx context.Context) { + defer wg.Done() + ch <- s.recallOneDataset(ctx, dsId, dsNames[dsId], clauseText, ftsQuery, embedders[dsId]) + }, func(ctx context.Context, e error) { + defer wg.Done() + ch <- nil + g.Log().Warningf(ctx, "dataset %d recall failed: %v", dsId, e) + }); err != nil { + wg.Done() + g.Log().Warningf(ctx, "submit dataset %d recall failed: %v", dsId, err) } - hits, err := dao.Chunk.FtsSearch(ctx, dsId, ftsQuery, consts.AnnoRecallTopK) - if err != nil { - g.Log().Warningf(ctx, "fts search failed (dataset %d): %v", dsId, err) - continue - } - for i, h := range hits { - s.mergeHit(merged, h.ChunkId, dsId, dsNames[dsId], float64(i)) + } + go func() { wg.Wait(); close(ch) }() + + merged := make(map[int64]*annoCandidate) + for hits := range ch { + for _, h := range hits { + s.mergeHit(merged, h.ChunkId, h.DsId, h.LawTitle, h.Rank) } } cands := make([]annoCandidate, 0, len(merged)) @@ -313,6 +347,35 @@ func (s *annotationService) recallCandidates(ctx context.Context, clause *entity return cands, nil } +// recallOneDataset 单数据集召回：向量检索 + FTS 检索（纯读，供池内并发调用） +func (s *annotationService) recallOneDataset(ctx context.Context, dsId int64, lawTitle, clauseText, ftsQuery string, emb *OpenAIEmbedder) []annoRecallHit { + var hits []annoRecallHit + if emb != nil { + vecs, err := emb.EmbedStrings(ctx, []string{clauseText}) + if err != nil { + g.Log().Warningf(ctx, "clause embed failed (dataset %d): %v", dsId, err) + } else if len(vecs) > 0 { + res, err := dao.Chunk.VecSearch(ctx, dsId, domain.VecJsonF64(vecs[0]), consts.AnnoRecallTopK) + if err != nil { + g.Log().Warningf(ctx, "vec search failed (dataset %d): %v", dsId, err) + } else { + for i, h := range res { + hits = append(hits, annoRecallHit{ChunkId: h.ChunkId, DsId: dsId, LawTitle: lawTitle, Rank: float64(i)}) + } + } + } + } + res, err := dao.Chunk.FtsSearch(ctx, dsId, ftsQuery, consts.AnnoRecallTopK) + if err != nil { + g.Log().Warningf(ctx, "fts search failed (dataset %d): %v", dsId, err) + return hits + } + for i, h := range res { + hits = append(hits, annoRecallHit{ChunkId: h.ChunkId, DsId: dsId, LawTitle: lawTitle, Rank: float64(i)}) + } + return hits +} + func (s *annotationService) mergeHit(merged map[int64]*annoCandidate, chunkId, dsId int64, lawTitle string, rank float64) { c := merged[chunkId] if c == nil { diff --git a/kb/service/chat_service.go b/kb/service/chat_service.go index b2410b9..aba6bb2 100644 --- a/kb/service/chat_service.go +++ b/kb/service/chat_service.go @@ -13,6 +13,7 @@ import ( "sort" "strconv" "strings" + "sync" "time" "rag-local/common" @@ -241,30 +242,40 @@ func (r *HybridRetriever) Retrieve(ctx context.Context, query string, opts ...er scores := make(map[int64]float64) srcs := make(map[int64][]string) + // 向量段与全文段互不依赖，提交 common.ChatRetrievePool 并行执行，RRF 合并回主 goroutine 串行 + ch := make(chan []retrieveHit, 2) + var wg sync.WaitGroup if r.embedder != nil { - vecs, err := r.embedder.EmbedStrings(ctx, []string{query}) - if err != nil { - g.Log().Warningf(ctx, "query embed failed: %v", err) - } else if len(vecs) > 0 { - hits, err := dao.Chunk.VecSearch(ctx, r.datasetId, domain.VecJsonF64(vecs[0]), consts.VectorTopK) - if err != nil { - g.Log().Warningf(ctx, "vec search failed: %v", err) - } else { - for i, h := range hits { - scores[h.ChunkId] += 1 / (float64(consts.RrfK) + float64(i) + 1) - srcs[h.ChunkId] = append(srcs[h.ChunkId], "vector") - } - } + wg.Add(1) + if err := common.ChatRetrievePool.AddWithRecover(ctx, func(ctx context.Context) { + defer wg.Done() + ch <- r.vecRetrieve(ctx, query) + }, func(ctx context.Context, e error) { + defer wg.Done() + ch <- nil + g.Log().Warningf(ctx, "vec retrieve failed: %v", e) + }); err != nil { + wg.Done() + g.Log().Warningf(ctx, "submit vec retrieve failed: %v", err) } } - - ftsHits, err := dao.Chunk.FtsSearch(ctx, r.datasetId, common.TokenizeQuery(query), consts.FtsTopK) - if err != nil { - g.Log().Warningf(ctx, "fts search failed: %v", err) - } else { - for i, h := range ftsHits { - scores[h.ChunkId] += 1 / (float64(consts.RrfK) + float64(i) + 1) - srcs[h.ChunkId] = append(srcs[h.ChunkId], "fts") + wg.Add(1) + if err := common.ChatRetrievePool.AddWithRecover(ctx, func(ctx context.Context) { + defer wg.Done() + ch <- r.ftsRetrieve(ctx, query) + }, func(ctx context.Context, e error) { + defer wg.Done() + ch <- nil + g.Log().Warningf(ctx, "fts retrieve failed: %v", e) + }); err != nil { + wg.Done() + g.Log().Warningf(ctx, "submit fts retrieve failed: %v", err) + } + go func() { wg.Wait(); close(ch) }() + for hits := range ch { + for _, h := range hits { + scores[h.chunkId] += 1 / (float64(consts.RrfK) + float64(h.rank) + 1) + srcs[h.chunkId] = append(srcs[h.chunkId], h.source) } } @@ -334,6 +345,52 @@ type scoredChunk struct { sources []string } +// retrieveHit 单路检索命中（rank 用于 RRF 融合） +type retrieveHit struct { + chunkId int64 + rank int + source string +} + +// vecRetrieve 向量段检索（纯读，供池内并发调用） +func (r *HybridRetriever) vecRetrieve(ctx context.Context, query string) []retrieveHit { + var hits []retrieveHit + if r.embedder == nil { + return hits + } + vecs, err := r.embedder.EmbedStrings(ctx, []string{query}) + if err != nil { + g.Log().Warningf(ctx, "query embed failed: %v", err) + return hits + } + if len(vecs) == 0 { + return hits + } + res, err := dao.Chunk.VecSearch(ctx, r.datasetId, domain.VecJsonF64(vecs[0]), consts.VectorTopK) + if err != nil { + g.Log().Warningf(ctx, "vec search failed: %v", err) + return hits + } + for i, h := range res { + hits = append(hits, retrieveHit{chunkId: h.ChunkId, rank: i, source: "vector"}) + } + return hits +} + +// ftsRetrieve 全文检索段（纯读，供池内并发调用） +func (r *HybridRetriever) ftsRetrieve(ctx context.Context, query string) []retrieveHit { + res, err := dao.Chunk.FtsSearch(ctx, r.datasetId, common.TokenizeQuery(query), consts.FtsTopK) + if err != nil { + g.Log().Warningf(ctx, "fts search failed: %v", err) + return nil + } + hits := make([]retrieveHit, 0, len(res)) + for i, h := range res { + hits = append(hits, retrieveHit{chunkId: h.ChunkId, rank: i, source: "fts"}) + } + return hits +} + // rerankByLLM 用默认对话模型对候选分块打分（0-10，JSON 输出），返回 chunk_id → 相关性分； // 任何失败（调用/解析/空结果）返回错误，由调用方回退 RRF 排序 func (r *HybridRetriever) rerankByLLM(ctx context.Context, query string, items []scoredChunk) (map[int64]float64, error) { @@ -394,19 +451,40 @@ const MaxHistoryRounds = 10 // Ask RAG 问答工作流：混合检索 → 组装提示（含引用编号）→ 对话模型流式生成。 // history 需已包含最新一条用户问题；onCitations 在检索完成后先于流式输出回调；onDelta 接收增量文本，均可为 nil。 func (s *chatService) Ask(ctx context.Context, datasetId int64, question string, history []*schema.Message, onCitations func([]domain.Citation), onDelta func(string)) (string, []domain.Citation, error) { - docs, err := s.retrieve(ctx, datasetId, question) - if err != nil { - return "", nil, err + // 检索（较重，内部再并行 vec/fts）与图增强互不依赖，并行执行；检索放 common.ChatPool，图增强主 goroutine 直接跑 + type askOut struct { + docs []*schema.Document + err error } - citations := buildCitations(docs) - if onCitations != nil { - onCitations(citations) + ch := make(chan askOut, 1) + var wg sync.WaitGroup + wg.Add(1) + if err := common.ChatPool.AddWithRecover(ctx, func(ctx context.Context) { + defer wg.Done() + docs, err := s.retrieve(ctx, datasetId, question) + ch <- askOut{docs: docs, err: err} + }, func(ctx context.Context, e error) { + defer wg.Done() + ch <- askOut{err: e} + }); err != nil { + wg.Done() + return "", nil, err } triples, err := KgRelationService.GraphEnhance(ctx, datasetId, question) if err != nil { g.Log().Warningf(ctx, "graph enhance failed: %v", err) } + wg.Wait() + out := <-ch + if out.err != nil { + return "", nil, out.err + } + docs := out.docs + citations := buildCitations(docs) + if onCitations != nil { + onCitations(citations) + } defaultChatModel, err := dao.ModelConfig.GetDefault(ctx, consts.ModelTypeChat) if err != nil { diff --git a/kb/service/kg_entity_service.go b/kb/service/kg_entity_service.go index 82f26f3..7713c28 100644 --- a/kb/service/kg_entity_service.go +++ b/kb/service/kg_entity_service.go @@ -4,7 +4,9 @@ import ( "context" "encoding/json" "strings" + "sync" + "rag-local/common" "rag-local/kb/consts" "rag-local/kb/dao" "rag-local/kb/model/entity" @@ -36,6 +38,7 @@ type kgExtractResult struct { // ExtractDocument 对文档全部新分块做 LLM 抽取（挂在解析流水线分块落库之后）。 // 每个分块一次调用，任何失败只记日志，不阻断解析流水线；返回未成功抽取的分块数，供调用方标记图谱未构建。 +// 并行：LLM 调用（纯 IO）提交 common.KgExtractPool 并发执行，落库（SQLite 写）收敛回本 goroutine 串行。 func (s *kgEntityService) ExtractDocument(ctx context.Context, datasetId, documentId int64) (int, error) { chunks, _, err := dao.Chunk.ListByDocument(ctx, documentId, 1, 100000) if err != nil { @@ -45,11 +48,42 @@ func (s *kgEntityService) ExtractDocument(ctx context.Context, datasetId, docume if err != nil { return 0, err } - failed := 0 + + type extractOut struct { + chunkId int64 + entities []kgEntityItem + relations []kgRelationItem + err error + } + ch := make(chan extractOut, len(chunks)) + var wg sync.WaitGroup for _, c := range chunks { - if err := s.extractChunk(ctx, model, datasetId, c); err != nil { + wg.Add(1) + if err := common.KgExtractPool.AddWithRecover(ctx, func(ctx context.Context) { + defer wg.Done() + out := extractOut{chunkId: c.Id} + out.entities, out.relations, out.err = s.callExtract(ctx, model, c) + ch <- out + }, func(ctx context.Context, e error) { + defer wg.Done() + ch <- extractOut{chunkId: c.Id, err: e} + }); err != nil { + wg.Done() + g.Log().Warningf(ctx, "kg extract chunk %d submit failed: %v", c.Id, err) + } + } + go func() { wg.Wait(); close(ch) }() + + failed := 0 + for out := range ch { + if out.err != nil { failed++ - g.Log().Warningf(ctx, "kg extract chunk %d failed: %v", c.Id, err) + g.Log().Warningf(ctx, "kg extract chunk %d failed: %v", out.chunkId, out.err) + continue + } + if err := s.saveExtract(ctx, datasetId, out.chunkId, out.entities, out.relations); err != nil { + failed++ + g.Log().Warningf(ctx, "kg extract chunk %d failed: %v", out.chunkId, err) } } return failed, nil @@ -66,34 +100,54 @@ func (s *kgEntityService) buildModel(ctx context.Context) (*OpenAIChatModel, err return BuildChatModel(ctx, defaultChatModel) } -func (s *kgEntityService) extractChunk(ctx context.Context, model *OpenAIChatModel, datasetId int64, chunk *entity.Chunk) error { +// callExtract 池内执行：LLM 抽取 + JSON 解析（纯 IO，不做任何写库） +func (s *kgEntityService) callExtract(ctx context.Context, model *OpenAIChatModel, chunk *entity.Chunk) ([]kgEntityItem, []kgRelationItem, error) { msgs := []*schema.Message{ {Role: schema.System, Content: "你是知识抽取助手。从文档片段中抽取实体（人名、组织、地名、产品等专有名词）及实体间的关系（动词或介词短语）。只输出 JSON，不要 markdown 代码块或任何解释，格式：{\"entities\":[{\"name\":\"实体名\",\"type\":\"类型\"}],\"relations\":[{\"head\":\"主体\",\"relation\":\"关系\",\"tail\":\"客体\"}]}"}, {Role: schema.User, Content: "文档片段：\n" + chunk.Content}, } resp, err := model.Generate(ctx, msgs) if err != nil { - return err + return nil, nil, err } data, err := parseKgJSON(resp.Content) if err != nil { - return err + return nil, nil, err } + entities := make([]kgEntityItem, 0, len(data.Entities)) for _, e := range data.Entities { - name := strings.TrimSpace(e.Name) - if name == "" { - continue - } - if err := dao.KgEntity.Upsert(ctx, datasetId, chunk.Id, name, strings.TrimSpace(e.Type)); err != nil { - return err + if name := strings.TrimSpace(e.Name); name != "" { + entities = append(entities, kgEntityItem{Name: name, Type: strings.TrimSpace(e.Type)}) } } + relations := make([]kgRelationItem, 0, len(data.Relations)) for _, r := range data.Relations { head, relation, tail := strings.TrimSpace(r.Head), strings.TrimSpace(r.Relation), strings.TrimSpace(r.Tail) if head == "" || relation == "" || tail == "" || head == tail { continue } - if err := dao.KgRelation.Insert(ctx, datasetId, chunk.Id, head, relation, tail); err != nil { + relations = append(relations, kgRelationItem{Head: head, Relation: relation, Tail: tail}) + } + return entities, relations, nil +} + +// saveExtract 主 goroutine 串行落库：实体与关系各一条批量 SQL（多行 VALUES） +func (s *kgEntityService) saveExtract(ctx context.Context, datasetId, chunkId int64, entities []kgEntityItem, relations []kgRelationItem) error { + if len(entities) > 0 { + items := make([]dao.KgEntityItem, 0, len(entities)) + for _, e := range entities { + items = append(items, dao.KgEntityItem{Name: e.Name, EntityType: e.Type}) + } + if err := dao.KgEntity.UpsertBatch(ctx, datasetId, chunkId, items); err != nil { + return err + } + } + if len(relations) > 0 { + items := make([]dao.KgRelationItem, 0, len(relations)) + for _, r := range relations { + items = append(items, dao.KgRelationItem{Head: r.Head, Relation: r.Relation, Tail: r.Tail}) + } + if err := dao.KgRelation.InsertBatch(ctx, datasetId, chunkId, items); err != nil { return err } } diff --git a/kb/service/parse_task_service.go b/kb/service/parse_task_service.go index 18bdc0e..258dadd 100644 --- a/kb/service/parse_task_service.go +++ b/kb/service/parse_task_service.go @@ -13,25 +13,19 @@ import ( "github.com/gogf/gf/v2/errors/gerror" "github.com/gogf/gf/v2/frame/g" + "github.com/gogf/gf/v2/os/gtimer" ) var ParseTaskService = &parseTaskService{} type parseTaskService struct{} -// StartParsePoller 启动任务轮询：单 goroutine 串行消费待处理任务（与 video-factory StartVideoPoller 同模式） +// StartParsePoller 启动任务轮询：gtimer 单例定时器串行消费待处理任务（job 未结束不重入，与视频工厂同模式） func (s *parseTaskService) StartParsePoller(ctx context.Context) { - go func() { - g.Log().Info(ctx, "parse task poller started") - for { - select { - case <-ctx.Done(): - return - case <-time.After(consts.ParsePollIntervalSeconds * time.Second): - s.processOne(ctx) - } - } - }() + g.Log().Info(ctx, "parse task poller started") + gtimer.AddSingleton(ctx, consts.ParsePollIntervalSeconds*time.Second, func(ctx context.Context) { + s.processOne(ctx) + }) } // processOne 处理一个待处理任务：解析 → 分块 → 落库（向量化在 M3 接入）