This commit is contained in:
2026-08-06 14:25:03 +08:00
parent 9afde8e2d6
commit 410c87e02f
11 changed files with 448 additions and 136 deletions
+31
View File
@@ -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.<key> 读取池大小(<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)
}
+8
View File
@@ -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 并行
BIN
View File
Binary file not shown.
+40 -11
View File
@@ -131,6 +131,7 @@ rag-local/
│ ├── text_parser.go # txt / md
│ ├── tokenizer.go # gse 中文分词(写索引/检索共用)
│ ├── util.go # RandomToken(随机文件名/令牌)、TokenFingerprintSHA-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 问答、合同标注编排归入对应 servicegrpool 协程池单例在 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 时已强制绑定,见下)
StartParsePollermain.go 启动,3 秒轮询,单 goroutine 串行
StartParsePollermain.go 启动,gtimer 单例 5 秒轮询,job 未结束不重入
│ 取 status=0 任务 → 置 status=1(解析中)
│ 1. 校验 embedding 配置 → 构建 OpenAIEmbedder(无配置 → 任务失败「数据集未绑定向量模型」;
│ 数据集 Save 强制绑定向量模型「数据集必须绑定向量模型,请先选择向量模型」,此处为防御性校验)
@@ -787,8 +789,9 @@ StartParsePollermain.go 启动,3 秒轮询,单 goroutine 串行)
│ c. 语义分块:eino semantic splitterMinChunkSize=chunkSize/2
│ d. 递归兜底:eino recursiveKeepTypeEnd
│ 5. 批量 Embedding → InsertAllchunk+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 @@ StartParsePollermain.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 @@ StartParsePollermain.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 15annotation_dataset 池并行)
RRF 融合截 60 候选 → LLM 一次调用判定(0-10 分 + 理由, JSON)
→ 主 goroutine 串行落库(清旧标→插 mark→置状态→更新进度)
→ 全部条款完成 → 任务 done(断点:重启后按 clause 状态续跑,已 done 条款跳过)
```
@@ -853,6 +858,30 @@ StartParsePollermain.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 | Askretrieve 与 GraphEnhance 并行 |
| 问答检索 | `pool.chat_retrieve` | 4 | HybridRetrievervec 段与 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.<key>", 默认值)` 读取;`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 agentopenai 直连) | **Eino** 组件化编排(Indexer/Retriever/Graph |
| 鉴权 | 多用户 + 角色 + JWTuser/login_log 表) | **单用户访问令牌**:每次启动重新生成(内存持有不落库)打印,登录换取 JWT,无用户表 |
| 前端 | HTML + 少量 Vue | 纯 Vue 3 SPA(构建产物 Go 托管) |
| 异步任务 | 视频生成轮询 | 文档解析/向量化 + 合同条款标注 双轮询(单 goroutine 串行) |
| 异步任务 | 视频生成轮询 | 文档解析/向量化 + 合同条款标注 双轮询(任务级单 goroutine 串行,任务内热点用 grpool 协程池并行,见 §7.9 |
| 其余 | 分层/路由/鉴权/缓存/部署模式 | 完全对齐 |
+8 -1
View File
@@ -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 并行
)
+34 -9
View File
@@ -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 {
+33 -10
View File
@@ -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 命中实体名的三元组(实体链接用)
+116 -53
View File
@@ -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 {
+105 -27
View File
@@ -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 {
+67 -13
View File
@@ -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
}
}
+6 -12
View File
@@ -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 接入)