1
This commit is contained in:
@@ -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)
|
||||
}
|
||||
@@ -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 并行
|
||||
|
||||
Binary file not shown.
+40
-11
@@ -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.<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 agent(openai 直连) | **Eino** 组件化编排(Indexer/Retriever/Graph) |
|
||||
| 鉴权 | 多用户 + 角色 + JWT(user/login_log 表) | **单用户访问令牌**:每次启动重新生成(内存持有不落库)打印,登录换取 JWT,无用户表 |
|
||||
| 前端 | HTML + 少量 Vue | 纯 Vue 3 SPA(构建产物 Go 托管) |
|
||||
| 异步任务 | 视频生成轮询 | 文档解析/向量化 + 合同条款标注 双轮询(各单 goroutine 串行) |
|
||||
| 异步任务 | 视频生成轮询 | 文档解析/向量化 + 合同条款标注 双轮询(任务级单 goroutine 串行,任务内热点用 grpool 协程池并行,见 §7.9) |
|
||||
| 其余 | 分层/路由/鉴权/缓存/部署模式 | 完全对齐 |
|
||||
|
||||
+8
-1
@@ -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
@@ -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
@@ -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 命中实体名的三元组(实体链接用)
|
||||
|
||||
@@ -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
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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 接入)
|
||||
|
||||
Reference in New Issue
Block a user