feat(materialindex): parse pipeline with chunk embed and RAG search
EnsureParsed waits or lazily triggers MinerU, indexes oversized markdown into material_chunks, and supports reranked retrieval. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
weichen
@@ -0,0 +1,654 @@
|
||||
package materialindex
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"chat-agent-system/internal/config"
|
||||
"chat-agent-system/internal/kb"
|
||||
"chat-agent-system/internal/mineru"
|
||||
"chat-agent-system/internal/squad"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
const (
|
||||
engineMinerU = "mineru"
|
||||
statePending = "pending"
|
||||
stateRunning = "running"
|
||||
stateDone = "done"
|
||||
stateFailed = "failed"
|
||||
maxAttempts = 3
|
||||
defaultWait = 90 * time.Second
|
||||
chunkRuneLimit = 1200
|
||||
)
|
||||
|
||||
var pageAnchorRe = regexp.MustCompile(`(?m)^<!--\s*page\s+(\d+)\s*-->\s*$`)
|
||||
|
||||
// ParseRecord is a row from asset_parses.
|
||||
type ParseRecord struct {
|
||||
ID int64
|
||||
BlobSHA256 string
|
||||
FileName string
|
||||
Engine string
|
||||
State string
|
||||
TaskID string
|
||||
Markdown string
|
||||
ContentList json.RawMessage
|
||||
TotalPages int
|
||||
ErrorMessage string
|
||||
Attempts int
|
||||
}
|
||||
|
||||
// ChunkHit is a retrieved material chunk.
|
||||
type ChunkHit struct {
|
||||
FileName string
|
||||
Title string
|
||||
Content string
|
||||
PageStart int
|
||||
PageEnd int
|
||||
Score float64
|
||||
}
|
||||
|
||||
// Pipeline owns parse cache + chunk index.
|
||||
type Pipeline struct {
|
||||
pool *pgxpool.Pool
|
||||
client *mineru.Client
|
||||
embedder *kb.Embedder
|
||||
reranker *kb.Reranker
|
||||
waitLimit time.Duration
|
||||
fullMax int
|
||||
ragTopK int
|
||||
|
||||
inflightMu sync.Mutex
|
||||
inflight map[string]chan struct{}
|
||||
}
|
||||
|
||||
// NewPipeline builds a material parse/index pipeline. client may be nil/disabled.
|
||||
func NewPipeline(cfg *config.Config, pool *pgxpool.Pool) *Pipeline {
|
||||
var client *mineru.Client
|
||||
var embedder *kb.Embedder
|
||||
var reranker *kb.Reranker
|
||||
wait := defaultWait
|
||||
fullMax := 60 * 1024
|
||||
topK := 8
|
||||
if cfg != nil {
|
||||
client = mineru.NewClient(mineru.Config{
|
||||
Endpoint: cfg.MineruEndpointURL,
|
||||
APIKey: cfg.MineruAPIKey,
|
||||
ModelVersion: cfg.MineruModelVersion,
|
||||
Language: cfg.MineruLanguage,
|
||||
HTTPTimeout: cfg.MineruTimeout,
|
||||
})
|
||||
if cfg.MineruTimeout > 0 {
|
||||
wait = cfg.MineruTimeout
|
||||
}
|
||||
if cfg.MaterialFulltextMaxBytes > 0 {
|
||||
fullMax = cfg.MaterialFulltextMaxBytes
|
||||
}
|
||||
if cfg.MaterialRAGTopK > 0 {
|
||||
topK = cfg.MaterialRAGTopK
|
||||
}
|
||||
if cfg.EmbeddingEndpointURL != "" {
|
||||
embedder = kb.NewEmbedder(cfg.EmbeddingEndpointURL, cfg.EmbeddingAPIKey, cfg.EmbeddingModel)
|
||||
}
|
||||
if cfg.RerankerEndpointURL != "" {
|
||||
reranker = kb.NewReranker(cfg.RerankerEndpointURL, cfg.RerankerAPIKey, cfg.RerankerModel)
|
||||
}
|
||||
}
|
||||
return &Pipeline{
|
||||
pool: pool,
|
||||
client: client,
|
||||
embedder: embedder,
|
||||
reranker: reranker,
|
||||
waitLimit: wait,
|
||||
fullMax: fullMax,
|
||||
ragTopK: topK,
|
||||
inflight: make(map[string]chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// Enabled reports whether MinerU credentials are present.
|
||||
func (p *Pipeline) Enabled() bool {
|
||||
return p != nil && p.client != nil && p.client.Enabled()
|
||||
}
|
||||
|
||||
// FulltextMaxBytes returns the hybrid injection threshold.
|
||||
func (p *Pipeline) FulltextMaxBytes() int {
|
||||
if p == nil || p.fullMax <= 0 {
|
||||
return 60 * 1024
|
||||
}
|
||||
return p.fullMax
|
||||
}
|
||||
|
||||
// RAGTopK returns retrieval top-k.
|
||||
func (p *Pipeline) RAGTopK() int {
|
||||
if p == nil || p.ragTopK <= 0 {
|
||||
return 8
|
||||
}
|
||||
return p.ragTopK
|
||||
}
|
||||
|
||||
// GetParseState returns parse state for a blob (empty string if unknown).
|
||||
func (p *Pipeline) GetParseState(ctx context.Context, blobSHA string) string {
|
||||
rec, err := p.getRecord(ctx, blobSHA)
|
||||
if err != nil || rec == nil {
|
||||
return ""
|
||||
}
|
||||
return rec.State
|
||||
}
|
||||
|
||||
// GetParseStates batch-loads parse states keyed by blob sha.
|
||||
func (p *Pipeline) GetParseStates(ctx context.Context, blobSHAs []string) map[string]string {
|
||||
out := make(map[string]string, len(blobSHAs))
|
||||
if p == nil || p.pool == nil || len(blobSHAs) == 0 {
|
||||
return out
|
||||
}
|
||||
rows, err := p.pool.Query(ctx, `
|
||||
SELECT blob_sha256, state FROM asset_parses
|
||||
WHERE engine = $1 AND blob_sha256 = ANY($2)
|
||||
`, engineMinerU, blobSHAs)
|
||||
if err != nil {
|
||||
return out
|
||||
}
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
var sha, state string
|
||||
if rows.Scan(&sha, &state) == nil {
|
||||
out[sha] = state
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// EnsureParsed waits for or triggers MinerU parse. On miss/disable/fail returns
|
||||
// fallback text from squad.ExtractTextWithAnchors (may be empty).
|
||||
func (p *Pipeline) EnsureParsed(ctx context.Context, fileName string, blobSHA string, data []byte) (markdown string, fromMinerU bool, err error) {
|
||||
fallback := func() (string, bool, error) {
|
||||
txt, e := squad.ExtractTextWithAnchors(fileName, data)
|
||||
return strings.TrimSpace(txt), false, e
|
||||
}
|
||||
if p == nil {
|
||||
return fallback()
|
||||
}
|
||||
if blobSHA == "" && len(data) > 0 {
|
||||
sum := sha256.Sum256(data)
|
||||
blobSHA = hex.EncodeToString(sum[:])
|
||||
}
|
||||
if blobSHA == "" {
|
||||
return fallback()
|
||||
}
|
||||
|
||||
rec, _ := p.getRecord(ctx, blobSHA)
|
||||
if rec != nil && rec.State == stateDone && strings.TrimSpace(rec.Markdown) != "" {
|
||||
return rec.Markdown, true, nil
|
||||
}
|
||||
if !p.Enabled() || !mineru.SupportsExt(fileName) {
|
||||
return fallback()
|
||||
}
|
||||
if rec != nil && rec.State == stateFailed && rec.Attempts >= maxAttempts {
|
||||
return fallback()
|
||||
}
|
||||
|
||||
waitCtx, cancel := context.WithTimeout(ctx, p.waitLimit)
|
||||
defer cancel()
|
||||
|
||||
if rec != nil && (rec.State == statePending || rec.State == stateRunning) {
|
||||
if done, md := p.waitUntilDone(waitCtx, blobSHA); done {
|
||||
return md, true, nil
|
||||
}
|
||||
return fallback()
|
||||
}
|
||||
|
||||
if err := p.triggerParse(waitCtx, fileName, blobSHA, data); err != nil {
|
||||
slog.Warn("materialindex parse failed", "file", fileName, "sha", blobSHA[:min(12, len(blobSHA))], "error", err)
|
||||
return fallback()
|
||||
}
|
||||
rec, _ = p.getRecord(ctx, blobSHA)
|
||||
if rec != nil && rec.State == stateDone && strings.TrimSpace(rec.Markdown) != "" {
|
||||
return rec.Markdown, true, nil
|
||||
}
|
||||
return fallback()
|
||||
}
|
||||
|
||||
// EnqueueAsync starts background parse without blocking the caller.
|
||||
func (p *Pipeline) EnqueueAsync(fileName string, blobSHA string, data []byte, projectID int64) {
|
||||
if p == nil || !p.Enabled() || !mineru.SupportsExt(fileName) || blobSHA == "" || len(data) == 0 {
|
||||
return
|
||||
}
|
||||
go func() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), p.waitLimit+2*time.Minute)
|
||||
defer cancel()
|
||||
md, ok, err := p.EnsureParsed(ctx, fileName, blobSHA, data)
|
||||
if err != nil || !ok || strings.TrimSpace(md) == "" {
|
||||
return
|
||||
}
|
||||
if utf8.RuneCountInString(md) > 0 && len(md) > p.FulltextMaxBytes() {
|
||||
if idxErr := p.IndexChunks(ctx, projectID, blobSHA, fileName, md); idxErr != nil {
|
||||
slog.Warn("materialindex index chunks failed", "sha", blobSHA[:min(12, len(blobSHA))], "error", idxErr)
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func (p *Pipeline) triggerParse(ctx context.Context, fileName, blobSHA string, data []byte) error {
|
||||
key := blobSHA
|
||||
p.inflightMu.Lock()
|
||||
if ch, ok := p.inflight[key]; ok {
|
||||
p.inflightMu.Unlock()
|
||||
select {
|
||||
case <-ch:
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
done := make(chan struct{})
|
||||
p.inflight[key] = done
|
||||
p.inflightMu.Unlock()
|
||||
defer func() {
|
||||
close(done)
|
||||
p.inflightMu.Lock()
|
||||
delete(p.inflight, key)
|
||||
p.inflightMu.Unlock()
|
||||
}()
|
||||
|
||||
if err := p.upsertPending(ctx, blobSHA, fileName); err != nil {
|
||||
return err
|
||||
}
|
||||
_ = p.setState(ctx, blobSHA, stateRunning, "", "", "", nil, 0)
|
||||
|
||||
result, err := p.client.ParseFile(ctx, fileName, data)
|
||||
if err != nil {
|
||||
msg := err.Error()
|
||||
var retry *mineru.RetryableError
|
||||
attempts := 1
|
||||
if rec, _ := p.getRecord(ctx, blobSHA); rec != nil {
|
||||
attempts = rec.Attempts + 1
|
||||
}
|
||||
_ = p.bumpAttempts(ctx, blobSHA, attempts, msg)
|
||||
if errors.As(err, &retry) && attempts < maxAttempts {
|
||||
return err
|
||||
}
|
||||
_ = p.setState(ctx, blobSHA, stateFailed, "", "", msg, nil, 0)
|
||||
return err
|
||||
}
|
||||
return p.setState(ctx, blobSHA, stateDone, result.BatchID, result.Markdown, "", result.ContentList, result.TotalPages)
|
||||
}
|
||||
|
||||
func (p *Pipeline) waitUntilDone(ctx context.Context, blobSHA string) (bool, string) {
|
||||
ticker := time.NewTicker(2 * time.Second)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
rec, _ := p.getRecord(ctx, blobSHA)
|
||||
if rec != nil {
|
||||
if rec.State == stateDone && strings.TrimSpace(rec.Markdown) != "" {
|
||||
return true, rec.Markdown
|
||||
}
|
||||
if rec.State == stateFailed {
|
||||
return false, ""
|
||||
}
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return false, ""
|
||||
case <-ticker.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (p *Pipeline) getRecord(ctx context.Context, blobSHA string) (*ParseRecord, error) {
|
||||
if p.pool == nil {
|
||||
return nil, nil
|
||||
}
|
||||
row := p.pool.QueryRow(ctx, `
|
||||
SELECT id, blob_sha256, file_name, engine, state, task_id, markdown, content_list, total_pages, error_message, attempts
|
||||
FROM asset_parses WHERE blob_sha256 = $1 AND engine = $2
|
||||
`, blobSHA, engineMinerU)
|
||||
var rec ParseRecord
|
||||
var contentList []byte
|
||||
err := row.Scan(&rec.ID, &rec.BlobSHA256, &rec.FileName, &rec.Engine, &rec.State, &rec.TaskID,
|
||||
&rec.Markdown, &contentList, &rec.TotalPages, &rec.ErrorMessage, &rec.Attempts)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return nil, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rec.ContentList = contentList
|
||||
return &rec, nil
|
||||
}
|
||||
|
||||
func (p *Pipeline) upsertPending(ctx context.Context, blobSHA, fileName string) error {
|
||||
if p.pool == nil {
|
||||
return nil
|
||||
}
|
||||
_, err := p.pool.Exec(ctx, `
|
||||
INSERT INTO asset_parses (blob_sha256, file_name, engine, state, attempts, updated_at)
|
||||
VALUES ($1, $2, $3, $4, 0, now())
|
||||
ON CONFLICT (blob_sha256, engine) DO UPDATE SET
|
||||
file_name = EXCLUDED.file_name,
|
||||
updated_at = now()
|
||||
WHERE asset_parses.state NOT IN ('done', 'running')
|
||||
`, blobSHA, fileName, engineMinerU, statePending)
|
||||
return err
|
||||
}
|
||||
|
||||
func (p *Pipeline) setState(ctx context.Context, blobSHA, state, taskID, markdown, errMsg string, contentList json.RawMessage, pages int) error {
|
||||
if p.pool == nil {
|
||||
return nil
|
||||
}
|
||||
cl := []byte("[]")
|
||||
if len(contentList) > 0 {
|
||||
cl = contentList
|
||||
}
|
||||
_, err := p.pool.Exec(ctx, `
|
||||
UPDATE asset_parses SET
|
||||
state = $2,
|
||||
task_id = CASE WHEN $3 = '' THEN task_id ELSE $3 END,
|
||||
markdown = CASE WHEN $4 = '' THEN markdown ELSE $4 END,
|
||||
content_list = CASE WHEN $5::jsonb = '[]'::jsonb THEN content_list ELSE $5::jsonb END,
|
||||
total_pages = CASE WHEN $6 = 0 THEN total_pages ELSE $6 END,
|
||||
error_message = $7,
|
||||
updated_at = now()
|
||||
WHERE blob_sha256 = $1 AND engine = $8
|
||||
`, blobSHA, state, taskID, markdown, cl, pages, errMsg, engineMinerU)
|
||||
return err
|
||||
}
|
||||
|
||||
func (p *Pipeline) bumpAttempts(ctx context.Context, blobSHA string, attempts int, errMsg string) error {
|
||||
if p.pool == nil {
|
||||
return nil
|
||||
}
|
||||
_, err := p.pool.Exec(ctx, `
|
||||
UPDATE asset_parses SET attempts = $2, error_message = $3, updated_at = now()
|
||||
WHERE blob_sha256 = $1 AND engine = $4
|
||||
`, blobSHA, attempts, errMsg, engineMinerU)
|
||||
return err
|
||||
}
|
||||
|
||||
// IndexChunks splits markdown and stores embeddings into material_chunks.
|
||||
func (p *Pipeline) IndexChunks(ctx context.Context, projectID int64, blobSHA, fileName, markdown string) error {
|
||||
if p == nil || p.pool == nil || p.embedder == nil {
|
||||
return nil
|
||||
}
|
||||
sections := splitMaterialSections(markdown)
|
||||
if len(sections) == 0 {
|
||||
return nil
|
||||
}
|
||||
texts := make([]string, len(sections))
|
||||
for i, s := range sections {
|
||||
texts[i] = fmt.Sprintf("%s %s\n%s", fileName, s.Title, s.Content)
|
||||
}
|
||||
vecs, err := p.embedder.Embed(ctx, texts)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for i, sec := range sections {
|
||||
var vec []float32
|
||||
if i < len(vecs) {
|
||||
vec = vecs[i]
|
||||
}
|
||||
hash := fmt.Sprintf("%x", sha256.Sum256([]byte(fmt.Sprintf("%s:%d:%s:%s", blobSHA, i, sec.Title, sec.Content))))
|
||||
meta, _ := json.Marshal(map[string]any{
|
||||
"file_name": fileName,
|
||||
"page_start": sec.PageStart,
|
||||
"page_end": sec.PageEnd,
|
||||
})
|
||||
vecLit := kb.VectorLiteral(vec)
|
||||
_, err := p.pool.Exec(ctx, `
|
||||
INSERT INTO material_chunks (project_id, blob_sha256, file_name, chunk_index, title, content, page_start, page_end, chunk_hash, embedding, metadata)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10::vector, $11)
|
||||
ON CONFLICT (chunk_hash) DO UPDATE SET
|
||||
title = EXCLUDED.title,
|
||||
content = EXCLUDED.content,
|
||||
embedding = EXCLUDED.embedding,
|
||||
metadata = EXCLUDED.metadata
|
||||
`, projectID, blobSHA, fileName, i, sec.Title, sec.Content, sec.PageStart, sec.PageEnd, hash, vecLit, meta)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Search retrieves top chunks for a project+optional blob, with optional rerank.
|
||||
func (p *Pipeline) Search(ctx context.Context, projectID int64, blobSHA, query string, limit int) ([]ChunkHit, error) {
|
||||
if p == nil || p.pool == nil || p.embedder == nil || strings.TrimSpace(query) == "" {
|
||||
return nil, nil
|
||||
}
|
||||
if limit <= 0 {
|
||||
limit = p.RAGTopK()
|
||||
}
|
||||
qVecs, err := p.embedder.Embed(ctx, []string{query})
|
||||
if err != nil || len(qVecs) == 0 {
|
||||
return nil, err
|
||||
}
|
||||
vecLit := kb.VectorLiteral(qVecs[0])
|
||||
fetch := limit * 3
|
||||
if fetch < limit {
|
||||
fetch = limit
|
||||
}
|
||||
rows, err := p.pool.Query(ctx, `
|
||||
SELECT file_name, title, content, page_start, page_end,
|
||||
1 - (embedding <=> $1::vector) AS score
|
||||
FROM material_chunks
|
||||
WHERE project_id = $2
|
||||
AND ($3 = '' OR blob_sha256 = $3)
|
||||
ORDER BY embedding <=> $1::vector ASC
|
||||
LIMIT $4
|
||||
`, vecLit, projectID, blobSHA, fetch)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var hits []ChunkHit
|
||||
var docs []string
|
||||
for rows.Next() {
|
||||
var h ChunkHit
|
||||
if err := rows.Scan(&h.FileName, &h.Title, &h.Content, &h.PageStart, &h.PageEnd, &h.Score); err != nil {
|
||||
continue
|
||||
}
|
||||
hits = append(hits, h)
|
||||
docs = append(docs, h.Title+"\n"+h.Content)
|
||||
}
|
||||
if p.reranker != nil && len(docs) > 0 {
|
||||
scores, rerr := p.reranker.Rerank(ctx, query, docs)
|
||||
if rerr == nil && len(scores) == len(hits) {
|
||||
for i := range hits {
|
||||
hits[i].Score = scores[i]
|
||||
}
|
||||
sortHitsByScore(hits)
|
||||
}
|
||||
}
|
||||
if len(hits) > limit {
|
||||
hits = hits[:limit]
|
||||
}
|
||||
return hits, nil
|
||||
}
|
||||
|
||||
// BuildTOC extracts heading outline with page hints from markdown.
|
||||
func BuildTOC(fileName, markdown string) string {
|
||||
var sb strings.Builder
|
||||
sb.WriteString(fmt.Sprintf("# 材料目录: %s\n\n", filepath.Base(fileName)))
|
||||
sb.WriteString("本文档体积较大，已改为目录 + 按任务检索片段注入。请结合 `.rag.md` 引用出处。\n\n")
|
||||
currentPage := 1
|
||||
count := 0
|
||||
for _, line := range strings.Split(markdown, "\n") {
|
||||
trim := strings.TrimSpace(line)
|
||||
if m := pageAnchorRe.FindStringSubmatch(trim); len(m) == 2 {
|
||||
fmt.Sscanf(m[1], "%d", ¤tPage)
|
||||
continue
|
||||
}
|
||||
if strings.HasPrefix(trim, "#") {
|
||||
level := 0
|
||||
for level < len(trim) && trim[level] == '#' {
|
||||
level++
|
||||
}
|
||||
title := strings.TrimSpace(trim[level:])
|
||||
if title == "" {
|
||||
continue
|
||||
}
|
||||
sb.WriteString(fmt.Sprintf("- [p.%d] %s %s\n", currentPage, strings.Repeat("#", min(level, 3)), title))
|
||||
count++
|
||||
if count >= 80 {
|
||||
sb.WriteString("\n…（目录已截断）\n")
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
if count == 0 {
|
||||
sb.WriteString("- （未识别到标题；请直接查看检索片段）\n")
|
||||
}
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
// FormatRAGMarkdown formats retrieved hits for workspace injection.
|
||||
func FormatRAGMarkdown(fileName string, hits []ChunkHit) string {
|
||||
var sb strings.Builder
|
||||
sb.WriteString(fmt.Sprintf("# 材料检索片段: %s\n\n", filepath.Base(fileName)))
|
||||
if len(hits) == 0 {
|
||||
sb.WriteString("（未命中相关片段；可缩小问题范围或上传更小文档）\n")
|
||||
return sb.String()
|
||||
}
|
||||
for i, h := range hits {
|
||||
title := h.Title
|
||||
if title == "" {
|
||||
title = "片段"
|
||||
}
|
||||
page := h.PageStart
|
||||
if page <= 0 {
|
||||
page = 1
|
||||
}
|
||||
cite := fmt.Sprintf("〔材料: %s §%s p.%d〕", filepath.Base(fileName), title, page)
|
||||
sb.WriteString(fmt.Sprintf("## %d. %s\n\n", i+1, title))
|
||||
sb.WriteString(fmt.Sprintf("- 出处: %s\n", cite))
|
||||
if h.PageEnd > h.PageStart && h.PageEnd > 0 {
|
||||
sb.WriteString(fmt.Sprintf("- 页码范围: p.%d–p.%d\n", h.PageStart, h.PageEnd))
|
||||
}
|
||||
sb.WriteString("\n")
|
||||
sb.WriteString(strings.TrimSpace(h.Content))
|
||||
sb.WriteString("\n\n---\n\n")
|
||||
}
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
type section struct {
|
||||
Title string
|
||||
Content string
|
||||
PageStart int
|
||||
PageEnd int
|
||||
}
|
||||
|
||||
func splitMaterialSections(md string) []section {
|
||||
lines := strings.Split(md, "\n")
|
||||
var sections []section
|
||||
var cur section
|
||||
cur.PageStart = 1
|
||||
cur.PageEnd = 1
|
||||
currentPage := 1
|
||||
flush := func() {
|
||||
cur.Content = strings.TrimSpace(cur.Content)
|
||||
if cur.Content == "" && cur.Title == "" {
|
||||
return
|
||||
}
|
||||
if cur.Title == "" {
|
||||
cur.Title = "正文"
|
||||
}
|
||||
cur.PageEnd = currentPage
|
||||
// further split oversized sections
|
||||
parts := splitByRuneLimit(cur, chunkRuneLimit)
|
||||
sections = append(sections, parts...)
|
||||
cur = section{PageStart: currentPage, PageEnd: currentPage}
|
||||
}
|
||||
for _, line := range lines {
|
||||
trim := strings.TrimSpace(line)
|
||||
if m := pageAnchorRe.FindStringSubmatch(trim); len(m) == 2 {
|
||||
fmt.Sscanf(m[1], "%d", ¤tPage)
|
||||
if cur.PageStart == 0 {
|
||||
cur.PageStart = currentPage
|
||||
}
|
||||
cur.Content += line + "\n"
|
||||
continue
|
||||
}
|
||||
if strings.HasPrefix(trim, "## ") {
|
||||
flush()
|
||||
cur.Title = strings.TrimSpace(strings.TrimPrefix(trim, "## "))
|
||||
cur.PageStart = currentPage
|
||||
continue
|
||||
}
|
||||
cur.Content += line + "\n"
|
||||
}
|
||||
flush()
|
||||
return sections
|
||||
}
|
||||
|
||||
func splitByRuneLimit(sec section, limit int) []section {
|
||||
if utf8.RuneCountInString(sec.Content) <= limit {
|
||||
return []section{sec}
|
||||
}
|
||||
paras := strings.Split(sec.Content, "\n\n")
|
||||
var out []section
|
||||
var buf strings.Builder
|
||||
part := 1
|
||||
flush := func() {
|
||||
if buf.Len() == 0 {
|
||||
return
|
||||
}
|
||||
title := sec.Title
|
||||
if part > 1 {
|
||||
title = fmt.Sprintf("%s (%d)", sec.Title, part)
|
||||
}
|
||||
out = append(out, section{
|
||||
Title: title,
|
||||
Content: strings.TrimSpace(buf.String()),
|
||||
PageStart: sec.PageStart,
|
||||
PageEnd: sec.PageEnd,
|
||||
})
|
||||
buf.Reset()
|
||||
part++
|
||||
}
|
||||
for _, p := range paras {
|
||||
if utf8.RuneCountInString(buf.String())+utf8.RuneCountInString(p) > limit && buf.Len() > 0 {
|
||||
flush()
|
||||
}
|
||||
if buf.Len() > 0 {
|
||||
buf.WriteString("\n\n")
|
||||
}
|
||||
buf.WriteString(p)
|
||||
}
|
||||
flush()
|
||||
if len(out) == 0 {
|
||||
return []section{sec}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func sortHitsByScore(hits []ChunkHit) {
|
||||
for i := 0; i < len(hits); i++ {
|
||||
for j := i + 1; j < len(hits); j++ {
|
||||
if hits[j].Score > hits[i].Score {
|
||||
hits[i], hits[j] = hits[j], hits[i]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func min(a, b int) int {
|
||||
if a < b {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
package materialindex_test
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"chat-agent-system/internal/materialindex"
|
||||
)
|
||||
|
||||
func TestBuildTOCAndFormatRAG(t *testing.T) {
|
||||
md := "<!-- page 1 -->\n## 需求概述\n正文A\n<!-- page 2 -->\n## 技术参数\n正文B\n"
|
||||
toc := materialindex.BuildTOC("spec.pdf", md)
|
||||
if !strings.Contains(toc, "[p.1]") || !strings.Contains(toc, "需求概述") {
|
||||
t.Fatalf("toc=%s", toc)
|
||||
}
|
||||
rag := materialindex.FormatRAGMarkdown("spec.pdf", []materialindex.ChunkHit{{
|
||||
FileName: "spec.pdf",
|
||||
Title: "技术参数",
|
||||
Content: "关键参数 220V",
|
||||
PageStart: 2,
|
||||
PageEnd: 2,
|
||||
Score: 0.9,
|
||||
}})
|
||||
if !strings.Contains(rag, "〔材料: spec.pdf §技术参数 p.2〕") {
|
||||
t.Fatalf("rag=%s", rag)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatRAG_Empty(t *testing.T) {
|
||||
out := materialindex.FormatRAGMarkdown("x.pdf", nil)
|
||||
if !strings.Contains(out, "未命中") {
|
||||
t.Fatalf("out=%s", out)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user