Agent-System/internal/ctxmgr/case_indexing.go

package ctxmgr

import (
	"context"
	"crypto/sha256"
	"encoding/json"
	"errors"
	"fmt"
	"regexp"
	"sort"
	"strings"
	"time"

	"chat-agent-system/internal/docgov"
)

var regPattern = regexp.MustCompile(`(?:GB[/\s-T]*\d+(?:\.\d+)?(?:-\d{4})?|YY[/\s-T]*\d+(?:\.\d+)?(?:-\d{4})?|ISO[/\s-T]*\d+|IEC[/\s-T]*\d+|MDR|FDA|NMPA)`)

// IndexApprovedDocument 将终审批准的交付物按二级章节切片、向量化并持久化到知识记忆切片库
func (m *ContextManager) IndexApprovedDocument(ctx context.Context, doc *docgov.DocumentRecord, markdownBody string) ([]int64, error) {
	if doc == nil || strings.TrimSpace(markdownBody) == "" {
		return nil, errors.New("document and markdown body required")
	}

	// 1. 按二级章节切片
	sections := splitMarkdownSections(markdownBody)
	if len(sections) == 0 {
		sections = []struct {
			Title   string
			Content string
		}{
			{
				Title:   doc.Title,
				Content: markdownBody,
			},
		}
	}

	textsToEmbed := make([]string, len(sections))
	for i, sec := range sections {
		// 结合文档主标题、器械类型与章节内容增强特征表征
		textsToEmbed[i] = fmt.Sprintf("%s %s %s\n%s", doc.Title, doc.TaskKey, sec.Title, sec.Content)
	}

	// 2. 生成嵌入向量
	emb := m.embedder
	if emb == nil {
		emb = NewDeterministicEmbedder(1024)
	}
	vecs, err := emb.Embed(ctx, textsToEmbed)
	if err != nil {
		return nil, fmt.Errorf("generate embeddings failed: %w", err)
	}

	var insertedIDs []int64

	// 3. 落库与更新内存缓存
	for idx, sec := range sections {
		var vec []float32
		if idx < len(vecs) {
			vec = vecs[idx]
		}
		chunkHash := fmt.Sprintf("%x", sha256.Sum256([]byte(fmt.Sprintf("%s:%s:%s", doc.DocID, sec.Title, sec.Content))))

		// 提取法规标准引用
		regs := extractRegulations(sec.Content)

		meta := map[string]any{
			"task_key":       doc.TaskKey,
			"artifact_title": doc.Title,
			"section_title":  sec.Title,
			"regulations":    regs,
			"author_name":    doc.AuthorName,
			"reviewer_name":  doc.ReviewerName,
			"approver_name":  doc.ApproverName,
			"stage":          doc.Stage,
			"approved_at":    doc.ApprovedAt,
		}

		if m.pool != nil {
			metaJSON, _ := json.Marshal(meta)
			vecLit := VectorLiteral(vec)

			var chunkID int64
			err := m.pool.QueryRow(ctx, `
				INSERT INTO context_memory_chunks (domain, project_id, user_id, reference_id, title, content, chunk_hash, embedding, importance_weight, metadata, created_at)
				VALUES ($1, $2, $3, $4, $5, $6, $7, $8::vector, 1.00, $9, now())
				ON CONFLICT (chunk_hash) DO UPDATE SET
					title = EXCLUDED.title,
					content = EXCLUDED.content,
					embedding = EXCLUDED.embedding,
					metadata = EXCLUDED.metadata
				RETURNING id
			`, "approved_case", doc.ProjectID, doc.AuthorID, doc.DocID, sec.Title, sec.Content, chunkHash, vecLit, metaJSON).Scan(&chunkID)
			if err == nil {
				insertedIDs = append(insertedIDs, chunkID)
			} else {
				fmt.Printf("[IndexApprovedDocument] DB insert failed: %v\n", err)
			}
		} else {
			m.mu.Lock()
			chunkID := m.nextChunkID
			m.nextChunkID++
			item := MemoryChunkItem{
				ID:               chunkID,
				Domain:           "approved_case",
				ProjectID:        doc.ProjectID,
				UserID:           doc.AuthorID,
				ReferenceID:      doc.DocID,
				Title:            sec.Title,
				Content:          sec.Content,
				ChunkHash:        chunkHash,
				Embedding:        vec,
				ImportanceWeight: 1.00,
				Metadata:         meta,
				CreatedAt:        time.Now(),
			}

			// 检查是否已有相同 chunkHash
			updated := false
			for i, existing := range m.memoryChunks {
				if existing.ChunkHash == chunkHash {
					m.memoryChunks[i] = item
					updated = true
					insertedIDs = append(insertedIDs, existing.ID)
					break
				}
			}
			if !updated {
				m.memoryChunks = append(m.memoryChunks, item)
				insertedIDs = append(insertedIDs, chunkID)
			}
			m.mu.Unlock()
		}
	}

	return insertedIDs, nil
}

// SearchSimilarCases 检索相似标杆案例
func (m *ContextManager) SearchSimilarCases(ctx context.Context, opts CaseSearchOptions) ([]CaseRecommendationItem, error) {
	if strings.TrimSpace(opts.Query) == "" {
		return nil, nil
	}
	if opts.Limit <= 0 {
		opts.Limit = 3
	}
	if opts.MinScore <= 0 {
		opts.MinScore = 0.30
	}

	emb := m.embedder
	if emb == nil {
		emb = NewDeterministicEmbedder(1024)
	}
	qVecs, err := emb.Embed(ctx, []string{opts.Query})
	if err != nil || len(qVecs) == 0 {
		return nil, fmt.Errorf("embed query failed: %w", err)
	}
	queryVec := qVecs[0]

	if m.pool != nil {
		vecLit := VectorLiteral(queryVec)
		query := `
			SELECT c.id, c.project_id, COALESCE(p.name, ''), c.reference_id, c.title, c.content, c.metadata,
			       1 - (c.embedding <=> $1::vector) as similarity_score
			FROM context_memory_chunks c
			LEFT JOIN projects p ON c.project_id = p.id
			WHERE c.domain = 'approved_case'
			  AND ($2 = '' OR c.reference_id != $2)
			  AND ($3 = '' OR c.metadata->>'task_key' = $3)
			ORDER BY c.embedding <=> $1::vector ASC
			LIMIT $4
		`
		rows, err := m.pool.Query(ctx, query, vecLit, opts.ExcludeDoc, opts.TaskKey, opts.Limit)
		if err != nil {
			return nil, fmt.Errorf("search similar cases failed: %w", err)
		}
		defer rows.Close()

		var results []CaseRecommendationItem
		for rows.Next() {
			var id, projectID int64
			var projectName, refID, title, content string
			var metaJSON []byte
			var score float64

			if err := rows.Scan(&id, &projectID, &projectName, &refID, &title, &content, &metaJSON, &score); err == nil {
				if score < opts.MinScore {
					continue
				}
				var meta map[string]any
				_ = json.Unmarshal(metaJSON, &meta)

				item := buildRecommendationItem(projectID, projectName, refID, title, content, score, meta)
				results = append(results, item)
			}
		}
		return results, rows.Err()
	}

	m.mu.RLock()
	defer m.mu.RUnlock()

	type scoredItem struct {
		chunk MemoryChunkItem
		score float64
	}

	var scored []scoredItem
	for _, c := range m.memoryChunks {
		if c.Domain != "approved_case" {
			continue
		}
		if opts.ExcludeDoc != "" && c.ReferenceID == opts.ExcludeDoc {
			continue
		}
		if opts.TaskKey != "" {
			if tk, ok := c.Metadata["task_key"].(string); ok && tk != opts.TaskKey {
				continue
			}
		}

		score := cosineSimilarity(queryVec, c.Embedding)
		if score >= opts.MinScore {
			scored = append(scored, scoredItem{chunk: c, score: score})
		}
	}

	sort.Slice(scored, func(i, j int) bool {
		return scored[i].score > scored[j].score
	})

	var results []CaseRecommendationItem
	for i := 0; i < len(scored) && i < opts.Limit; i++ {
		c := scored[i].chunk
		projectName := ""
		if proj, ok := m.projects[c.ProjectID]; ok {
			projectName = proj.ProjectName
		}
		item := buildRecommendationItem(c.ProjectID, projectName, c.ReferenceID, c.Title, c.Content, scored[i].score, c.Metadata)
		results = append(results, item)
	}

	return results, nil
}

func buildRecommendationItem(projectID int64, projectName, refID, title, content string, score float64, meta map[string]any) CaseRecommendationItem {
	taskKey, _ := meta["task_key"].(string)
	artTitle, _ := meta["artifact_title"].(string)
	if artTitle == "" {
		artTitle = title
	}
	approvedAt, _ := meta["approved_at"].(string)

	var regs []string
	if rawRegs, ok := meta["regulations"].([]any); ok {
		for _, r := range rawRegs {
			if rs, ok := r.(string); ok {
				regs = append(regs, rs)
			}
		}
	} else if rawRegsStr, ok := meta["regulations"].([]string); ok {
		regs = rawRegsStr
	}

	snippet := strings.TrimSpace(content)
	if len([]rune(snippet)) > 200 {
		snippet = string([]rune(snippet)[:200]) + "..."
	}

	return CaseRecommendationItem{
		DocID:           refID,
		ProjectID:       projectID,
		ProjectName:     projectName,
		TaskKey:         taskKey,
		ArtifactTitle:   artTitle,
		SectionTitle:    title,
		SummarySnippet:  snippet,
		KeyRegulations:  regs,
		SimilarityScore: score,
		ApprovedAt:      approvedAt,
	}
}

// splitMarkdownSections 按二级标题 "## " 切分文档正文
func splitMarkdownSections(md string) []struct {
	Title   string
	Content string
} {
	lines := strings.Split(md, "\n")
	var sections []struct {
		Title   string
		Content string
	}

	var currentTitle string
	var currentLines []string

	for _, line := range lines {
		trimmed := strings.TrimSpace(line)
		if strings.HasPrefix(trimmed, "## ") {
			if currentTitle != "" && len(currentLines) > 0 {
				sections = append(sections, struct {
					Title   string
					Content string
				}{
					Title:   currentTitle,
					Content: strings.TrimSpace(strings.Join(currentLines, "\n")),
				})
			}
			currentTitle = trimmed
			currentLines = nil
		} else if currentTitle != "" {
			currentLines = append(currentLines, line)
		}
	}

	if currentTitle != "" && len(currentLines) > 0 {
		sections = append(sections, struct {
			Title   string
			Content string
		}{
			Title:   currentTitle,
			Content: strings.TrimSpace(strings.Join(currentLines, "\n")),
		})
	}

	return sections
}

// extractRegulations 抽取正文中提及的标准法规编号
func extractRegulations(text string) []string {
	matches := regPattern.FindAllString(text, -1)
	if len(matches) == 0 {
		return nil
	}
	seen := make(map[string]bool)
	var unique []string
	for _, m := range matches {
		m = strings.TrimSpace(m)
		if !seen[m] {
			seen[m] = true
			unique = append(unique, m)
		}
	}
	return unique
}

// cosineSimilarity 计算两个归一化向量的余弦相似度
func cosineSimilarity(a, b []float32) float64 {
	if len(a) == 0 || len(b) == 0 || len(a) != len(b) {
		return 0.0
	}
	var dot float64
	for i := 0; i < len(a); i++ {
		dot += float64(a[i] * b[i])
	}
	// 映射到 [0, 1] 范围
	if dot < 0 {
		dot = 0
	}
	if dot > 1 {
		dot = 1
	}
	return dot
}