sundynix-agentix/sundynix-mcp-go/internal/rag/ingest_large.go

package rag

import (
	"context"
	"fmt"
	"log"
	"strings"
	"sync"

	"github.com/sundynix/sundynix-shared/contract"
)

// embedAll 并发分批把 chunks 向量化，保序返回（vecs[i] 对应 chunks[i]）。
// 串行版几十万字上千块 = 上百次顺序 HTTP，慢且阻塞；这里限并发 embedConcurrency 跑，
// 各批写入预分配切片的对应区间保持顺序，任一批失败即整体失败（上层会清理重入）。
func (e *Engine) embedAll(ctx context.Context, chunks []string, emit func(contract.IngestEvent)) ([][]float32, error) {
	vecs := make([][]float32, len(chunks))
	sem := make(chan struct{}, embedConcurrency)
	var wg sync.WaitGroup
	var mu sync.Mutex
	var firstErr error
	done := 0

	for start := 0; start < len(chunks); start += embedBatch {
		end := min(start+embedBatch, len(chunks))
		wg.Add(1)
		sem <- struct{}{}
		go func(start, end int) {
			defer wg.Done()
			defer func() { <-sem }()
			mu.Lock()
			stop := firstErr != nil
			mu.Unlock()
			if stop {
				return // 已有批失败，不再发起新请求
			}
			bv, err := e.embed().Embed(ctx, chunks[start:end])
			mu.Lock()
			defer mu.Unlock()
			if err != nil {
				if firstErr == nil {
					firstErr = err
				}
				return
			}
			copy(vecs[start:end], bv)
			done += end - start
			emit(contract.IngestEvent{Stage: "向量化", Done: done, Total: len(chunks)})
		}(start, end)
	}
	wg.Wait()
	if firstErr != nil {
		return nil, firstErr
	}
	return vecs, nil
}

// packGraphWindows 把已切好的语义块合并成抽取窗口（窗内拼接，每窗 ~target 字）。
// 复用语义块边界，避免再切一遍；超出 maxWindows 截断并回报 truncated=true（调用方告警）。
func packGraphWindows(chunks []string, target, maxWindows int) (windows []string, truncated bool) {
	var cur strings.Builder
	curLen := 0
	flush := func() {
		if curLen > 0 {
			windows = append(windows, cur.String())
			cur.Reset()
			curLen = 0
		}
	}
	for _, c := range chunks {
		cl := runeLen(c)
		if curLen > 0 && curLen+cl > target {
			flush()
			if len(windows) >= maxWindows {
				return windows, true
			}
		}
		if curLen > 0 {
			cur.WriteByte('\n')
			curLen++
		}
		cur.WriteString(c)
		curLen += cl
	}
	flush()
	return windows, false
}

// extractGraphWindowed 把全文按窗口并发抽三元组并 MERGE 进 Neo4j。
// 取代"整篇喂 LLM"——几十万字也能全覆盖；实体按 kb+name 在 Neo4j 天然去重，
// 这里再做一次内存去重减少重复 MERGE。单窗失败只丢该窗，不影响其余。
func (e *Engine) extractGraphWindowed(ctx context.Context, kb, fileID string, chunks []string, emit func(contract.IngestEvent)) {
	windows, truncated := packGraphWindows(chunks, graphWindowRunes(), graphMaxWindows())
	if len(windows) == 0 {
		return
	}
	if truncated {
		log.Printf("[rag] 文档过大，图谱仅抽取前 %d 窗（其余略过）kb=%s", len(windows), kb)
		emit(contract.IngestEvent{Stage: "抽实体", Msg: fmt.Sprintf("文档较大，图谱抽取前 %d 段", len(windows))})
	}
	emit(contract.IngestEvent{Stage: "抽实体", Done: 0, Total: len(windows), Msg: fmt.Sprintf("LLM 分 %d 段抽取知识三元组…", len(windows))})

	chat := e.graphChatClient()
	sem := make(chan struct{}, graphConcurrency())
	var wg sync.WaitGroup
	var mu sync.Mutex
	var all []Triple
	doneWin := 0

	for _, w := range windows {
		wg.Add(1)
		sem <- struct{}{}
		go func(w string) {
			defer wg.Done()
			defer func() { <-sem }()
			ts, err := extractTriples(ctx, chat, w)
			mu.Lock()
			defer mu.Unlock()
			doneWin++
			if err != nil {
				log.Printf("[rag] 窗口三元组抽取失败（忽略该窗）: %v", err)
			} else {
				all = append(all, ts...)
			}
			emit(contract.IngestEvent{Stage: "抽实体", Done: doneWin, Total: len(windows), Msg: fmt.Sprintf("图谱抽取 %d/%d 段", doneWin, len(windows))})
		}(w)
	}
	wg.Wait()

	all = dedupeTriples(all)
	if len(all) == 0 {
		return
	}
	// 实时回流给 UI（边出现边渲染图谱）。
	tv := make([]contract.TripleView, len(all))
	for i, t := range all {
		tv[i] = contract.TripleView{S: t.S, P: t.P, O: t.O}
	}
	emit(contract.IngestEvent{Stage: "抽实体", Total: len(all), Triples: tv, Msg: fmt.Sprintf("抽出 %d 条知识三元组", len(all))})
	emit(contract.IngestEvent{Stage: "写Neo4j", Total: len(all), Msg: fmt.Sprintf("%d 条三元组写入图谱", len(all))})
	if n, gerr := e.graph.store(ctx, kb, fileID, all); gerr != nil {
		log.Printf("[rag] 写 Neo4j 失败（图谱降级）: %v", gerr)
	} else {
		log.Printf("[rag] 图谱: 写入 %d 条三元组到 kb=%s（%d 窗）", n, kb, len(windows))
	}
}

// dedupeTriples 去掉空项与重复的 (s,p,o)，减少跨窗重复 MERGE。
func dedupeTriples(ts []Triple) []Triple {
	seen := make(map[string]bool, len(ts))
	out := ts[:0]
	for _, t := range ts {
		if t.S == "" || t.P == "" || t.O == "" {
			continue
		}
		k := t.S + "\x00" + t.P + "\x00" + t.O
		if !seen[k] {
			seen[k] = true
			out = append(out, t)
		}
	}
	return out
}