diff --git a/internal/materialindex/pipeline.go b/internal/materialindex/pipeline.go new file mode 100644 index 00000000..bc9e1943 --- /dev/null +++ b/internal/materialindex/pipeline.go @@ -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*$`) + +// 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 +} diff --git a/internal/materialindex/pipeline_test.go b/internal/materialindex/pipeline_test.go new file mode 100644 index 00000000..da15f194 --- /dev/null +++ b/internal/materialindex/pipeline_test.go @@ -0,0 +1,34 @@ +package materialindex_test + +import ( + "strings" + "testing" + + "chat-agent-system/internal/materialindex" +) + +func TestBuildTOCAndFormatRAG(t *testing.T) { + md := "\n## 需求概述\n正文A\n\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) + } +}