472 lines
16 KiB
Go
472 lines
16 KiB
Go
package service
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"log"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
"unicode/utf8"
|
||
|
||
"tcm-agent/internal/crawler"
|
||
"tcm-agent/internal/dao"
|
||
)
|
||
|
||
// ========================================================================
|
||
// CrawlService —— 药品抓取的运行编排 + 进程内定时调度
|
||
// ========================================================================
|
||
// 两个入口：
|
||
// 1. 手动触发：面板点「立即抓取」→ handler 调 TriggerCrawlTask（异步跑）
|
||
// 2. 定时触发：StartCrawlScheduler 常驻 goroutine 每分钟检查到期任务
|
||
//
|
||
// 一次运行的完整流程（runCrawl）：
|
||
// 加载任务 → 建运行日志(running) → 抓索引 → 按游标切本批 URL →
|
||
// 逐个抓详情+解析 → upsert 到目标知识库（按药名判重）→
|
||
// 更新游标/日志终态/任务摘要 → 刷新库统计
|
||
//
|
||
// 并发控制：
|
||
// - 同一任务同一时刻只允许一次运行（crawlRunning 内存标记）
|
||
// - 全局同一时刻只跑一个抓取任务（信号量容量 1）——抓取是 IO 慢活，
|
||
// 串行既保护源站也避免多任务写同一库的统计竞态
|
||
// ========================================================================
|
||
|
||
// crawl 运行状态（内存标记；重启后自然清零，DB 里 running 状态由启动时校正）
|
||
var (
|
||
crawlMu sync.Mutex
|
||
crawlRunning = map[uint]bool{}
|
||
// 全局串行信号量：容量 1，多任务同时到期时排队执行
|
||
crawlSem = make(chan struct{}, 1)
|
||
)
|
||
|
||
// CrawlTaskRunning 查询某任务是否正在运行（API 列表展示用）
|
||
func CrawlTaskRunning(id uint) bool {
|
||
crawlMu.Lock()
|
||
defer crawlMu.Unlock()
|
||
return crawlRunning[id]
|
||
}
|
||
|
||
// markCrawlRunning 尝试标记任务开始运行；已在运行返回 false
|
||
func markCrawlRunning(id uint) bool {
|
||
crawlMu.Lock()
|
||
defer crawlMu.Unlock()
|
||
if crawlRunning[id] {
|
||
return false
|
||
}
|
||
crawlRunning[id] = true
|
||
return true
|
||
}
|
||
|
||
// unmarkCrawlRunning 清除运行标记
|
||
func unmarkCrawlRunning(id uint) {
|
||
crawlMu.Lock()
|
||
defer crawlMu.Unlock()
|
||
delete(crawlRunning, id)
|
||
}
|
||
|
||
// TriggerCrawlTask 触发一次抓取（异步）
|
||
//
|
||
// 同步做完所有"可立即拒绝"的校验（任务存在/启用/未在运行/源合法），
|
||
// 通过后丢 goroutine 后台跑，立即返回给前端——抓取一批要几分钟，HTTP 不能等
|
||
func TriggerCrawlTask(taskID uint, trigger string) error {
|
||
task, err := dao.KBCrawlTaskGet(taskID)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if task.Status != 1 {
|
||
return fmt.Errorf("任务已禁用，请先启用再触发")
|
||
}
|
||
if _, ok := crawler.GetSource(task.Source); !ok {
|
||
return fmt.Errorf("抓取源 %s 不存在", task.Source)
|
||
}
|
||
if _, err := dao.KBGetLibrary(task.LibraryID); err != nil {
|
||
return fmt.Errorf("目标知识库不存在（id=%d），请先在知识库管理里创建", task.LibraryID)
|
||
}
|
||
if !markCrawlRunning(taskID) {
|
||
return fmt.Errorf("任务正在运行中，请等待本次完成")
|
||
}
|
||
go func() {
|
||
defer unmarkCrawlRunning(taskID)
|
||
// 全局串行：等别的抓取任务跑完（保护源站 + 避免统计竞态）
|
||
crawlSem <- struct{}{}
|
||
defer func() { <-crawlSem }()
|
||
// 排队等待期间任务可能被用户禁用或改配置：
|
||
// 拿到执行权后重读最新配置再跑，避免按触发时的旧快照执行
|
||
fresh, err := dao.KBCrawlTaskGet(taskID)
|
||
if err != nil {
|
||
log.Printf("[Crawl] 任务 #%d 排队后重读失败（可能已删除），跳过: %v", taskID, err)
|
||
return
|
||
}
|
||
if fresh.Status != 1 {
|
||
log.Printf("[Crawl] 任务 #%d 排队期间被禁用，跳过本次运行", taskID)
|
||
return
|
||
}
|
||
runCrawl(fresh, trigger)
|
||
}()
|
||
return nil
|
||
}
|
||
|
||
// StartCrawlScheduler 启动定时调度器（main 里调用一次）
|
||
//
|
||
// 每 60s 醒来检查所有启用的非手动任务是否到期，到期的异步触发。
|
||
// 用轮询而非精确定时器：任务量小（个位数），分钟级精度足够，逻辑最简单
|
||
func StartCrawlScheduler() {
|
||
go func() {
|
||
// 启动时先校正一遍"孤儿 running"状态：
|
||
// 上次进程崩溃/重启时正在跑的任务，DB 里 last_status 会永远停在 running，
|
||
// 面板会一直显示进行中——这里统一改成 failed 并注明原因
|
||
fixOrphanRunning()
|
||
|
||
ticker := time.NewTicker(60 * time.Second)
|
||
defer ticker.Stop()
|
||
log.Printf("[Crawl] 定时调度器已启动（每分钟检查到期任务）")
|
||
for range ticker.C {
|
||
checkDueTasks()
|
||
}
|
||
}()
|
||
}
|
||
|
||
// fixOrphanRunning 把因进程重启滞留在 running 的任务/日志校正为 failed
|
||
func fixOrphanRunning() {
|
||
tasks, err := dao.KBCrawlTaskList()
|
||
if err != nil {
|
||
return
|
||
}
|
||
now := int(time.Now().Unix())
|
||
for _, t := range tasks {
|
||
if t.LastStatus == "running" {
|
||
_ = dao.KBCrawlTaskUpdate(t.ID, map[string]any{
|
||
"last_status": "failed",
|
||
"last_message": "服务重启中断（进程退出时任务未完成）",
|
||
})
|
||
// 对应的运行中日志也补终态
|
||
if logs, err := dao.KBCrawlLogList(t.ID, 5); err == nil {
|
||
for _, lg := range logs {
|
||
if lg.Status == "running" && lg.FinishedAt == 0 {
|
||
_ = dao.KBCrawlLogFinish(lg.ID, map[string]any{
|
||
"status": "failed", "finished_at": now,
|
||
"message": "服务重启中断",
|
||
})
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// checkDueTasks 扫描启用任务，把到期的触发执行
|
||
func checkDueTasks() {
|
||
tasks, err := dao.KBCrawlTaskList()
|
||
if err != nil {
|
||
return
|
||
}
|
||
now := time.Now()
|
||
for _, t := range tasks {
|
||
if t.Status != 1 || t.ScheduleType == "" || t.ScheduleType == "manual" {
|
||
continue
|
||
}
|
||
if CrawlTaskRunning(t.ID) {
|
||
continue
|
||
}
|
||
if crawlDue(&t, now) {
|
||
log.Printf("[Crawl] 任务 #%d「%s」到期，定时触发", t.ID, t.Name)
|
||
if err := TriggerCrawlTask(t.ID, "schedule"); err != nil {
|
||
log.Printf("[Crawl] 任务 #%d 定时触发失败: %v", t.ID, err)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// crawlDue 判断任务是否到期
|
||
//
|
||
// 三种调度语义（都基于本地时区）：
|
||
// - interval：距上次运行 >= interval_hours 小时（从未运行视为立即到期）
|
||
// - daily： 今天的 run_at_hour 整点已过，且上次运行早于该时点
|
||
// - weekly： 本周 run_at_weekday 的 run_at_hour 已过，且上次运行早于该时点
|
||
func crawlDue(t *dao.KBCrawlTaskRow, now time.Time) bool {
|
||
last := time.Unix(int64(t.LastRunAt), 0)
|
||
switch t.ScheduleType {
|
||
case "interval":
|
||
hours := t.IntervalHours
|
||
if hours < 1 {
|
||
hours = 24
|
||
}
|
||
return t.LastRunAt == 0 || now.Sub(last) >= time.Duration(hours)*time.Hour
|
||
case "daily":
|
||
due := time.Date(now.Year(), now.Month(), now.Day(), t.RunAtHour, 0, 0, 0, now.Location())
|
||
return now.After(due) && last.Before(due)
|
||
case "weekly":
|
||
// 找到本周的目标星期（周日=0 与 time.Weekday 语义一致）
|
||
offset := t.RunAtWeekday - int(now.Weekday())
|
||
due := time.Date(now.Year(), now.Month(), now.Day(), t.RunAtHour, 0, 0, 0, now.Location()).
|
||
AddDate(0, 0, offset)
|
||
// 目标时点在未来（本周还没到）则不触发；已过则比对上次运行
|
||
return now.After(due) && last.Before(due)
|
||
}
|
||
return false
|
||
}
|
||
|
||
// runCrawl 执行一次抓取（同步；调用方负责放 goroutine + 运行标记）
|
||
func runCrawl(task *dao.KBCrawlTaskRow, trigger string) {
|
||
src, _ := crawler.GetSource(task.Source)
|
||
startAt := int(time.Now().Unix())
|
||
|
||
// 建运行日志 + 任务标记 running
|
||
logRow := &dao.KBCrawlLogRow{
|
||
TaskID: task.ID, TriggerType: trigger,
|
||
StartedAt: startAt, Status: "running",
|
||
}
|
||
if err := dao.KBCrawlLogCreate(logRow); err != nil {
|
||
log.Printf("[Crawl] 任务 #%d 创建运行日志失败: %v", task.ID, err)
|
||
return
|
||
}
|
||
_ = dao.KBCrawlTaskUpdate(task.ID, map[string]any{
|
||
"last_run_at": startAt, "last_status": "running", "last_message": "",
|
||
})
|
||
|
||
// 超时预算：索引抓取 ~1min + 每条详情 ~2s（含限速）+ 入库余量
|
||
itemsPerRun := task.ItemsPerRun
|
||
if itemsPerRun < 1 {
|
||
itemsPerRun = 50
|
||
}
|
||
if itemsPerRun > 500 {
|
||
itemsPerRun = 500
|
||
}
|
||
timeout := time.Duration(itemsPerRun)*3*time.Second + 3*time.Minute
|
||
ctx, cancel := context.WithTimeout(context.Background(), timeout)
|
||
defer cancel()
|
||
|
||
finish := func(status, msg string, fetched, created, updated, failed int, newOffset int) {
|
||
_ = dao.KBCrawlLogFinish(logRow.ID, map[string]any{
|
||
"status": status, "finished_at": int(time.Now().Unix()),
|
||
"total_fetched": fetched, "created_docs": created,
|
||
"updated_docs": updated, "failed_items": failed,
|
||
"message": truncateRunes(msg, 900),
|
||
})
|
||
_ = dao.KBCrawlTaskUpdate(task.ID, map[string]any{
|
||
"last_status": status, "last_message": truncateRunes(msg, 450),
|
||
"progress_offset": newOffset,
|
||
})
|
||
}
|
||
|
||
// panic 兜底：抓取解析的是外部 HTML，任何解析路径的意外 panic
|
||
// 都不能带崩整个 Agent 服务进程；恢复后把本次运行记为失败、游标原地不动
|
||
defer func() {
|
||
if r := recover(); r != nil {
|
||
log.Printf("[Crawl] 任务 #%d panic 已恢复: %v", task.ID, r)
|
||
finish("failed", fmt.Sprintf("运行异常（panic 已恢复）: %v", r), 0, 0, 0, 0, task.ProgressOffset)
|
||
}
|
||
}()
|
||
|
||
// 1. 抓索引（全量 URL 列表，游标基于它切片）
|
||
urls, err := src.FetchIndex(ctx)
|
||
if err != nil {
|
||
log.Printf("[Crawl] 任务 #%d 索引抓取失败: %v", task.ID, err)
|
||
finish("failed", "索引抓取失败: "+err.Error(), 0, 0, 0, 0, task.ProgressOffset)
|
||
return
|
||
}
|
||
|
||
// 2. 游标切片：offset 越界（源站条目变少）自动归零重来
|
||
offset := task.ProgressOffset
|
||
if offset < 0 || offset >= len(urls) {
|
||
offset = 0
|
||
}
|
||
end := offset + itemsPerRun
|
||
if end > len(urls) {
|
||
end = len(urls)
|
||
}
|
||
batch := urls[offset:end]
|
||
log.Printf("[Crawl] 任务 #%d 开始：总 %d 条，本批 [%d, %d)", task.ID, len(urls), offset, end)
|
||
|
||
// 3. 逐条抓详情 + upsert
|
||
//
|
||
// attempted 记录"实际处理完的条数"——超时中断时游标只推进到断点，
|
||
// 而不是整批跳过（否则未抓的条目要等下一整轮才会补上）
|
||
var created, updated, failed, attempted int
|
||
var failMsgs []string
|
||
timedOut := false
|
||
for _, u := range batch {
|
||
med, err := src.FetchDetail(ctx, u)
|
||
if err != nil {
|
||
// 上下文超时/取消：当前条没有处理完，不计入 attempted，
|
||
// 游标停在这条之前，下次运行从断点续抓
|
||
if ctx.Err() != nil {
|
||
timedOut = true
|
||
break
|
||
}
|
||
attempted++
|
||
failed++
|
||
if len(failMsgs) < 3 {
|
||
failMsgs = append(failMsgs, err.Error())
|
||
}
|
||
continue
|
||
}
|
||
attempted++
|
||
isNew, err := upsertMedicine(task.LibraryID, med)
|
||
if err != nil {
|
||
failed++
|
||
if len(failMsgs) < 3 {
|
||
failMsgs = append(failMsgs, med.Name+": "+err.Error())
|
||
}
|
||
continue
|
||
}
|
||
if isNew {
|
||
created++
|
||
} else {
|
||
updated++
|
||
}
|
||
}
|
||
fetched := created + updated + failed
|
||
|
||
// 4. 新游标：只按实际处理量推进；一轮抓完归零（下轮转增量更新模式）
|
||
progressed := offset + attempted
|
||
newOffset := progressed
|
||
roundDone := newOffset >= len(urls)
|
||
if roundDone {
|
||
newOffset = 0
|
||
}
|
||
|
||
// 5. 刷新目标库统计（doc_count/chunk_count 冗余字段）
|
||
_ = dao.KBUpdateLibraryStats(task.LibraryID)
|
||
|
||
// 6. 终态判定 + 摘要（超时中断视为异常：有产出算 partial，无产出算 failed）
|
||
status := "success"
|
||
if failed > 0 || timedOut {
|
||
if created+updated > 0 {
|
||
status = "partial"
|
||
} else {
|
||
status = "failed"
|
||
}
|
||
}
|
||
msg := fmt.Sprintf("本批 %d 条：新建 %d、更新 %d、失败 %d；进度 %d/%d",
|
||
fetched, created, updated, failed, progressed, len(urls))
|
||
if timedOut {
|
||
msg += fmt.Sprintf("（运行超时中断，本批只完成 %d/%d，游标停在断点下次续抓）", attempted, len(batch))
|
||
}
|
||
if roundDone {
|
||
msg += "（一轮抓取完成，游标归零转增量更新）"
|
||
}
|
||
if len(failMsgs) > 0 {
|
||
msg += "；失败示例: " + strings.Join(failMsgs, " | ")
|
||
}
|
||
log.Printf("[Crawl] 任务 #%d 完成：%s", task.ID, msg)
|
||
finish(status, msg, fetched, created, updated, failed, newOffset)
|
||
}
|
||
|
||
// ---------------------------- 药材 → 知识库映射 ----------------------------
|
||
|
||
// coreSectionOrder 核心检索段落的优先顺序（组装 chunk 正文用）
|
||
//
|
||
// 为什么选这些：临床开方最关心的字段，且都在几百字内——
|
||
// 性味归经/功效/临床应用/禁忌是处方校验的关键依据。
|
||
// 「来源/生境分布」是源站老版条目的标签（对应新版的 药用部位/产地分布），两套都收
|
||
var coreSectionOrder = []string{
|
||
"性味归经", "功效与作用", "临床应用", "使用禁忌",
|
||
"药用部位", "来源", "产地分布", "生境分布", "采收加工", "药材性状",
|
||
}
|
||
|
||
// upsertMedicine 把一味药材写入知识库（按药名判重：有则更新，无则新建）
|
||
//
|
||
// 返回 isNew：true=新建文档，false=更新已有文档
|
||
func upsertMedicine(libraryID uint, med *crawler.Medicine) (bool, error) {
|
||
content, chunks := buildMedicineContent(med)
|
||
|
||
// 文档元数据：别名/拼音/来源 URL（面板文档详情可见，便于溯源）
|
||
meta := dao.MarshalMeta(map[string]any{
|
||
"aliases": med.Aliases,
|
||
"pinyin": med.Pinyin,
|
||
"source_url": med.SourceURL,
|
||
"crawled_at": time.Now().Format("2006-01-02 15:04:05"),
|
||
})
|
||
|
||
existing, err := dao.KBFindDocByTitle(libraryID, med.Name)
|
||
if err != nil {
|
||
return false, err
|
||
}
|
||
if existing != nil {
|
||
// 已存在：刷新文档正文 + 整体替换分段（幂等，重复抓取内容不膨胀）
|
||
if err := dao.KBUpdateDocContent(existing.ID, content, meta, "crawl"); err != nil {
|
||
return false, err
|
||
}
|
||
return false, dao.KBReplaceDocChunks(existing.ID, chunks)
|
||
}
|
||
|
||
// 不存在：新建文档 + 分段
|
||
doc := &dao.KBDocRow{
|
||
LibraryID: libraryID,
|
||
Title: med.Name,
|
||
Content: content,
|
||
SourceFile: med.SourceURL,
|
||
SourceType: "crawl",
|
||
MetaJSON: meta,
|
||
}
|
||
return true, dao.KBInsertDocWithChunks(doc, chunks)
|
||
}
|
||
|
||
// buildMedicineContent 组装文档正文与检索分段
|
||
//
|
||
// 分段策略（检索质量优先）：
|
||
// - 分段 1「核心卡」：药名+别名开头（别名参与 FULLTEXT，按别名搜也能命中）
|
||
// + 核心段落（性味归经/功效/临床应用/禁忌...），这是检索主力
|
||
// - 分段 2「配伍药方」：方剂内容独立成段（可能很长，且检索意图不同）
|
||
// - 其余长段落（植物形态/药理研究/化学成分）只进文档原文，不进分段——
|
||
// 它们对临床检索噪音大于价值，文档详情页仍可完整查看
|
||
func buildMedicineContent(med *crawler.Medicine) (string, []*dao.KBChunkRow) {
|
||
// 头行：药名（别名：a、b、c；拼音：xx）
|
||
head := med.Name
|
||
var headExtra []string
|
||
if len(med.Aliases) > 0 {
|
||
headExtra = append(headExtra, "别名："+strings.Join(med.Aliases, "、"))
|
||
}
|
||
if med.Pinyin != "" {
|
||
headExtra = append(headExtra, "拼音："+med.Pinyin)
|
||
}
|
||
if len(headExtra) > 0 {
|
||
head += "（" + strings.Join(headExtra, "；") + "）"
|
||
}
|
||
|
||
// 分段 1：核心卡（按 coreSectionOrder 顺序拼装）
|
||
var core strings.Builder
|
||
core.WriteString(head)
|
||
core.WriteString("\n")
|
||
for _, label := range coreSectionOrder {
|
||
if text := med.GetSection(label); text != "" {
|
||
core.WriteString("\n【" + label + "】" + text)
|
||
}
|
||
}
|
||
chunks := []*dao.KBChunkRow{{
|
||
Title: med.Name,
|
||
Content: truncateRunes(core.String(), 1800),
|
||
MetaJSON: dao.MarshalMeta(map[string]any{
|
||
"related_questions": med.Aliases, // 别名当关联问题，面板分段列表可见
|
||
}),
|
||
}}
|
||
|
||
// 分段 2：配伍药方（有才建）
|
||
if recipes := med.GetSection("配伍药方"); recipes != "" {
|
||
chunks = append(chunks, &dao.KBChunkRow{
|
||
Title: med.Name + "·配伍药方",
|
||
Content: truncateRunes(med.Name+"的配伍药方：\n"+recipes, 1800),
|
||
})
|
||
}
|
||
|
||
// 文档原文：头行 + 全部段落按原始顺序（详情页完整可查）
|
||
var full strings.Builder
|
||
full.WriteString(head)
|
||
full.WriteString("\n")
|
||
for _, s := range med.Sections {
|
||
full.WriteString("\n【" + s.Label + "】" + s.Text + "\n")
|
||
}
|
||
full.WriteString("\n来源：" + med.SourceURL)
|
||
return full.String(), chunks
|
||
}
|
||
|
||
// truncateRunes 按 rune 截断（不破坏中文字符边界）
|
||
func truncateRunes(s string, max int) string {
|
||
if utf8.RuneCountInString(s) <= max {
|
||
return s
|
||
}
|
||
runes := []rune(s)
|
||
return string(runes[:max]) + "…"
|
||
}
|