feat(ctxmgr): 增加已批准受控文档自动切片向量化与标杆案例语义检索
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
WeiCHEN
committed by
weichen
weichen
co-authored by
Cursor
Cursor
parent
db3dc282ba
commit
fb8232639b
@@ -51,7 +51,7 @@ func NewService(cfg *config.Config) *Service {
|
||||
ctxMgr := ctxmgr.NewContextManager(pool, casStore, ctxmgr.DefaultTimeDecayConfig())
|
||||
docGov := docgov.NewGovernanceEngine(pool)
|
||||
|
||||
return &Service{
|
||||
svc := &Service{
|
||||
cfg: cfg,
|
||||
kb: kb.NewClient(cfg),
|
||||
squad: squad.NewService(cfg),
|
||||
@@ -59,6 +59,31 @@ func NewService(cfg *config.Config) *Service {
|
||||
cas: casStore,
|
||||
docGov: docGov,
|
||||
}
|
||||
|
||||
// 注册终审批准回调：自动切片向量化沉淀为往期标杆案例
|
||||
docGov.SetOnApprovedHook(func(ctx context.Context, doc *docgov.DocumentRecord) {
|
||||
if ctxMgr == nil || doc == nil {
|
||||
return
|
||||
}
|
||||
var contentPath string
|
||||
if doc.WorkDir != "" {
|
||||
if filepath.IsAbs(doc.WorkDir) {
|
||||
contentPath = filepath.Join(doc.WorkDir, "output.md")
|
||||
} else {
|
||||
contentPath = filepath.Join(cfg.DataDir, doc.WorkDir, "output.md")
|
||||
if _, err := os.Stat(contentPath); os.IsNotExist(err) {
|
||||
contentPath = filepath.Join(doc.WorkDir, "output.md")
|
||||
}
|
||||
}
|
||||
}
|
||||
if contentPath != "" {
|
||||
if b, err := os.ReadFile(contentPath); err == nil && len(b) > 0 {
|
||||
_, _ = ctxMgr.IndexApprovedDocument(ctx, doc, string(b))
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
return svc
|
||||
}
|
||||
|
||||
func (s *Service) ContextManager() *ctxmgr.ContextManager {
|
||||
|
||||
@@ -0,0 +1,364 @@
|
||||
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,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
var textsToEmbed []string
|
||||
for _, sec := range sections {
|
||||
textsToEmbed = append(textsToEmbed, fmt.Sprintf("%s\n%s", 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 {
|
||||
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.50
|
||||
}
|
||||
|
||||
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, nil
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
package ctxmgr_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"chat-agent-system/internal/ctxmgr"
|
||||
"chat-agent-system/internal/docgov"
|
||||
)
|
||||
|
||||
func TestContextManager_ApprovedCaseIndexingAndSemanticSearch(t *testing.T) {
|
||||
mgr := ctxmgr.NewContextManager(nil, nil, ctxmgr.DefaultTimeDecayConfig())
|
||||
ctx := context.Background()
|
||||
|
||||
// 1. 模拟一个已被终审批准的交付物
|
||||
approvedTime := time.Now()
|
||||
doc := &docgov.DocumentRecord{
|
||||
DocID: "doc_proj101_feasibility_reg",
|
||||
ProjectID: 101,
|
||||
TaskKey: "feasibility_reg",
|
||||
Title: "注册路径可行性评估报告",
|
||||
Stage: "V0.5_feasibility",
|
||||
CurrentState: docgov.StateApproved,
|
||||
IsLocked: true,
|
||||
ApprovedAt: &approvedTime,
|
||||
AuthorName: "张工",
|
||||
}
|
||||
|
||||
markdownBody := `# 注册路径可行性评估报告
|
||||
|
||||
## 概述与器械分类
|
||||
超声软组织切割止血手术设备，主要用于外科手术中软组织的切割与血管闭合，按国家药监局三类医疗器械管理。
|
||||
|
||||
## 关键性能与型检预估
|
||||
设备需满足 GB 9706.1-2020 电气安全标准，以及 YY 0644 超声手术设备专用安全要求。型式检验周期预估 4-6 个月。
|
||||
|
||||
## 注册申报周期与成本策略
|
||||
整体 CDMO 研发与型检预算约为 180 万元人民币，临床评价采用同品种对比路径。`
|
||||
|
||||
// 2. 切片并向量化入库
|
||||
chunkIDs, err := mgr.IndexApprovedDocument(ctx, doc, markdownBody)
|
||||
if err != nil {
|
||||
t.Fatalf("IndexApprovedDocument failed: %v", err)
|
||||
}
|
||||
if len(chunkIDs) == 0 {
|
||||
t.Fatalf("expected indexed chunks, got 0")
|
||||
}
|
||||
|
||||
// 3. 语义搜索相似标杆案例
|
||||
cases, err := mgr.SearchSimilarCases(ctx, ctxmgr.CaseSearchOptions{
|
||||
Query: "超声刀设备需要测哪些电气安全和专用标准？",
|
||||
Limit: 3,
|
||||
MinScore: 0.50,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("SearchSimilarCases failed: %v", err)
|
||||
}
|
||||
if len(cases) == 0 {
|
||||
t.Fatalf("expected similar cases found, got 0")
|
||||
}
|
||||
|
||||
foundStandard := false
|
||||
for _, c := range cases {
|
||||
if c.DocID == doc.DocID && c.TaskKey == "feasibility_reg" {
|
||||
foundStandard = true
|
||||
t.Logf("✓ 成功命中往期标杆案例: 标题=%s, 章节=%s, 相似度=%.2f, 提要=%s",
|
||||
c.ArtifactTitle, c.SectionTitle, c.SimilarityScore, c.SummarySnippet)
|
||||
break
|
||||
}
|
||||
}
|
||||
if !foundStandard {
|
||||
t.Fatalf("expected to match indexed document doc_proj101_feasibility_reg")
|
||||
}
|
||||
}
|
||||
@@ -18,6 +18,7 @@ type ContextManager struct {
|
||||
cas *storage.CAS
|
||||
reflector *PreferenceReflector
|
||||
decayCfg TimeDecayConfig
|
||||
embedder VectorEmbedder
|
||||
|
||||
// 离线降级内存缓存
|
||||
mu sync.RWMutex
|
||||
@@ -26,8 +27,10 @@ type ContextManager struct {
|
||||
messages map[string][]ChatMessageRecord // session_id -> messages
|
||||
sessions map[string]*ChatSessionSummary // session_id -> summary
|
||||
activities map[int64][]ActivityItem // project_id -> activities
|
||||
memoryChunks []MemoryChunkItem // 语义记忆切片缓存
|
||||
nextActivityID int64
|
||||
nextMessageID int64
|
||||
nextChunkID int64
|
||||
}
|
||||
|
||||
// NewContextManager 创建上下文管理器
|
||||
@@ -40,13 +43,23 @@ func NewContextManager(pool *pgxpool.Pool, cas *storage.CAS, decayCfg TimeDecayC
|
||||
cas: cas,
|
||||
reflector: NewPreferenceReflector(),
|
||||
decayCfg: decayCfg,
|
||||
embedder: NewDeterministicEmbedder(1024),
|
||||
users: make(map[int64]*UserPreference),
|
||||
projects: make(map[int64]*ProjectCollectiveContext),
|
||||
messages: make(map[string][]ChatMessageRecord),
|
||||
sessions: make(map[string]*ChatSessionSummary),
|
||||
activities: make(map[int64][]ActivityItem),
|
||||
memoryChunks: make([]MemoryChunkItem, 0),
|
||||
nextActivityID: 1,
|
||||
nextMessageID: 1,
|
||||
nextChunkID: 1,
|
||||
}
|
||||
}
|
||||
|
||||
// SetEmbedder 设置自定义向量嵌入器 (如生产环境 bge-m3)
|
||||
func (m *ContextManager) SetEmbedder(emb VectorEmbedder) {
|
||||
if emb != nil {
|
||||
m.embedder = emb
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -25,11 +25,19 @@ type GovernanceEngine struct {
|
||||
pool *pgxpool.Pool
|
||||
sm *StateMachine
|
||||
|
||||
mu sync.RWMutex
|
||||
memDocs map[string]*DocumentRecord
|
||||
memComms map[string][]*DocumentComment
|
||||
memProps map[string]*ChangeProposal
|
||||
memAudits map[string][]*AuditLog
|
||||
mu sync.RWMutex
|
||||
memDocs map[string]*DocumentRecord
|
||||
memComms map[string][]*DocumentComment
|
||||
memProps map[string]*ChangeProposal
|
||||
memAudits map[string][]*AuditLog
|
||||
onApprovedHook func(ctx context.Context, doc *DocumentRecord)
|
||||
}
|
||||
|
||||
// SetOnApprovedHook 注册文档终审批准时的回调钩子
|
||||
func (e *GovernanceEngine) SetOnApprovedHook(hook func(ctx context.Context, doc *DocumentRecord)) {
|
||||
e.mu.Lock()
|
||||
defer e.mu.Unlock()
|
||||
e.onApprovedHook = hook
|
||||
}
|
||||
|
||||
// NewGovernanceEngine 初始化治理引擎
|
||||
@@ -738,6 +746,14 @@ func (e *GovernanceEngine) ApproverApprove(ctx context.Context, docID string, ac
|
||||
|
||||
e.memAudits[docID] = append(e.memAudits[docID], audit)
|
||||
e.syncDocAndAuditToDB(ctx, doc, audit)
|
||||
|
||||
if e.onApprovedHook != nil {
|
||||
hook := e.onApprovedHook
|
||||
e.mu.Unlock()
|
||||
hook(ctx, doc)
|
||||
e.mu.Lock()
|
||||
}
|
||||
|
||||
return doc, nil
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user