288 lines
14 KiB
Go
288 lines
14 KiB
Go
|
package service
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"encoding/json"
|
||
|
|
"errors"
|
||
|
|
"fmt"
|
||
|
|
"os"
|
||
|
|
"strings"
|
||
|
|
"sync"
|
||
|
|
"testing"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/accesspolicy"
|
||
|
|
"github.com/Tencent/WeKnora/internal/application/repository"
|
||
|
|
werrors "github.com/Tencent/WeKnora/internal/errors"
|
||
|
|
"github.com/Tencent/WeKnora/internal/models/chat"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
"github.com/google/uuid"
|
||
|
|
"github.com/stretchr/testify/require"
|
||
|
|
"gorm.io/driver/postgres"
|
||
|
|
"gorm.io/driver/sqlite"
|
||
|
|
"gorm.io/gorm"
|
||
|
|
)
|
||
|
|
|
||
|
|
type classificationTestModel struct {
|
||
|
|
*autoTagChatModel
|
||
|
|
call func(context.Context, []chat.Message, *chat.ChatOptions) (*types.ChatResponse, error)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m classificationTestModel) Chat(ctx context.Context, messages []chat.Message, options *chat.ChatOptions) (*types.ChatResponse, error) {
|
||
|
|
return m.call(ctx, messages, options)
|
||
|
|
}
|
||
|
|
|
||
|
|
func classificationFixture(t *testing.T) (*KnowledgeClassificationService, context.Context, *autoTagChatModel) {
|
||
|
|
t.Helper()
|
||
|
|
var db *gorm.DB
|
||
|
|
var e error
|
||
|
|
if dsn := os.Getenv("CLASSIFICATION_TEST_POSTGRES_DSN"); dsn != "" {
|
||
|
|
admin, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
|
||
|
|
require.NoError(t, err)
|
||
|
|
schema := "classification_test_" + strings.ReplaceAll(uuid.NewString(), "-", "")
|
||
|
|
require.NoError(t, admin.Exec("CREATE SCHEMA "+schema).Error)
|
||
|
|
t.Cleanup(func() {
|
||
|
|
require.NoError(t, admin.Exec("DROP SCHEMA "+schema+" CASCADE").Error)
|
||
|
|
raw, err := admin.DB()
|
||
|
|
require.NoError(t, err)
|
||
|
|
raw.Close()
|
||
|
|
})
|
||
|
|
db, e = gorm.Open(postgres.Open(dsn+" search_path="+schema), &gorm.Config{})
|
||
|
|
} else {
|
||
|
|
db, e = gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||
|
|
}
|
||
|
|
require.NoError(t, e)
|
||
|
|
raw, e := db.DB()
|
||
|
|
require.NoError(t, e)
|
||
|
|
if db.Dialector.Name() == "sqlite" {
|
||
|
|
raw.SetMaxOpenConns(1)
|
||
|
|
} else {
|
||
|
|
raw.SetMaxOpenConns(12)
|
||
|
|
}
|
||
|
|
t.Cleanup(func() { raw.Close() })
|
||
|
|
require.NoError(t, db.AutoMigrate(&types.User{}, &types.Tenant{}, &types.TenantMember{}, &types.KnowledgeBase{}, &types.Knowledge{}, &types.Chunk{}, &types.KnowledgeTag{}, &types.KnowledgeTagRelation{}, &types.AuditLog{}, &KnowledgeClassificationJob{}, &KnowledgeClassificationItem{}))
|
||
|
|
require.NoError(t, db.Exec("CREATE UNIQUE INDEX idx_classification_active ON knowledge_classification_jobs(tenant_id,owner_id,knowledge_base_id) WHERE status='running'").Error)
|
||
|
|
require.NoError(t, db.Create(&types.Tenant{ID: 1, Name: "test"}).Error)
|
||
|
|
require.NoError(t, db.Create(&types.User{ID: "alice", TenantID: 1, Email: "alice@example.test", Username: "Alice", IsActive: true}).Error)
|
||
|
|
require.NoError(t, db.Create(&types.TenantMember{TenantID: 1, UserID: "alice", Role: types.TenantRoleContributor, Status: types.TenantMemberStatusActive}).Error)
|
||
|
|
kb := &types.KnowledgeBase{ID: "kb", TenantID: 1, CreatorID: "alice", Type: "document", SummaryModelID: "model"}
|
||
|
|
require.NoError(t, db.Create(kb).Error)
|
||
|
|
require.NoError(t, db.Create(&types.Knowledge{ID: "doc", TenantID: 1, KnowledgeBaseID: "kb", Title: "product document", ParseStatus: types.ParseStatusCompleted}).Error)
|
||
|
|
require.NoError(t, db.Create(&types.Chunk{ID: "chunk", TenantID: 1, KnowledgeBaseID: "kb", KnowledgeID: "doc", Content: "Product education: " + strings.Repeat("x", 100000), ChunkType: types.ChunkTypeText, IndexStatus: "ready"}).Error)
|
||
|
|
for _, tag := range []*types.KnowledgeTag{{ID: "new", TenantID: 1, KnowledgeBaseID: "kb", Name: "产品 / 使用说明 / 完整分类路径"}, {ID: "old", TenantID: 1, KnowledgeBaseID: "kb", Name: "现有分类"}} {
|
||
|
|
require.NoError(t, db.Create(tag).Error)
|
||
|
|
}
|
||
|
|
model := &autoTagChatModel{response: `{"matches":[{"index":1,"confidence":0.95,"reason":"资料介绍了产品的使用方法，适合归入使用说明。"}]}`}
|
||
|
|
s := &KnowledgeClassificationService{db: db, models: &stubModelService{chatModel: model}, bases: versionCopyKBService{bases: map[string]*types.KnowledgeBase{"kb": kb}}}
|
||
|
|
ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(1))
|
||
|
|
ctx = context.WithValue(ctx, types.UserIDContextKey, "alice")
|
||
|
|
ctx = context.WithValue(ctx, types.TenantRoleContextKey, types.TenantRoleContributor)
|
||
|
|
return s, ctx, model
|
||
|
|
}
|
||
|
|
func TestKnowledgeClassificationRequiresExplicitApplyAndPreservesTags(t *testing.T) {
|
||
|
|
s, ctx, model := classificationFixture(t)
|
||
|
|
repo := repository.NewKnowledgeRepository(s.db)
|
||
|
|
require.NoError(t, repo.SetKnowledgeTags(ctx, "doc", []string{"old"}))
|
||
|
|
job, e := s.Create(ctx, "kb", "", []string{"doc"})
|
||
|
|
require.NoError(t, e)
|
||
|
|
_, e = s.Create(ctx, "kb", "", []string{"doc"})
|
||
|
|
require.ErrorContains(t, e, "已有分类任务")
|
||
|
|
require.NoError(t, s.RunNext(context.Background()))
|
||
|
|
jobs, e := s.List(ctx, "kb")
|
||
|
|
require.NoError(t, e)
|
||
|
|
require.Len(t, jobs, 1)
|
||
|
|
require.Equal(t, "completed", jobs[0].Status)
|
||
|
|
require.Equal(t, 1, jobs[0].Processed)
|
||
|
|
item := jobs[0].Items[0]
|
||
|
|
require.Equal(t, "suggested", item.Status)
|
||
|
|
var suggestions []ClassificationSuggestion
|
||
|
|
require.NoError(t, json.Unmarshal(item.Suggestions, &suggestions))
|
||
|
|
require.Equal(t, "产品 / 使用说明 / 完整分类路径", suggestions[0].Name)
|
||
|
|
tags, e := repo.GetKnowledgeTags(ctx, []string{"doc"})
|
||
|
|
require.NoError(t, e)
|
||
|
|
require.Len(t, tags["doc"], 1)
|
||
|
|
require.Less(t, len(model.messages[1].Content), 5000, "must not read whole large chunk into model")
|
||
|
|
require.NoError(t, s.Apply(ctx, "kb", job.ID, []string{item.ID}))
|
||
|
|
require.NoError(t, s.Apply(ctx, "kb", job.ID, []string{item.ID}))
|
||
|
|
tags, e = repo.GetKnowledgeTags(ctx, []string{"doc"})
|
||
|
|
require.NoError(t, e)
|
||
|
|
require.Len(t, tags["doc"], 2)
|
||
|
|
var count int64
|
||
|
|
require.NoError(t, s.db.Model(&types.AuditLog{}).Where("action=?", "knowledge.classification_applied").Count(&count).Error)
|
||
|
|
require.Positive(t, count)
|
||
|
|
}
|
||
|
|
func TestKnowledgeClassificationRejectsStaleContentCategoryAndOtherUser(t *testing.T) {
|
||
|
|
for _, change := range []string{"content", "category", "user", "permission"} {
|
||
|
|
t.Run(change, func(t *testing.T) {
|
||
|
|
s, ctx, _ := classificationFixture(t)
|
||
|
|
job, e := s.Create(ctx, "kb", "model", []string{"doc"})
|
||
|
|
require.NoError(t, e)
|
||
|
|
require.NoError(t, s.RunNext(context.Background()))
|
||
|
|
jobs, e := s.List(ctx, "kb")
|
||
|
|
require.NoError(t, e)
|
||
|
|
switch change {
|
||
|
|
case "content":
|
||
|
|
require.NoError(t, s.db.Model(&types.Chunk{}).Where("id='chunk'").Updates(map[string]any{"content": "new text", "content_revision": 1}).Error)
|
||
|
|
case "category":
|
||
|
|
require.NoError(t, s.db.Model(&types.KnowledgeTag{}).Where("id='new'").Update("name", "changed").Error)
|
||
|
|
case "user":
|
||
|
|
ctx = context.WithValue(ctx, types.UserIDContextKey, "bob")
|
||
|
|
ctx = context.WithValue(ctx, types.TenantRoleContextKey, types.TenantRoleAdmin)
|
||
|
|
case "permission":
|
||
|
|
ctx = accesspolicy.WithSnapshot(ctx, &accesspolicy.Snapshot{Restricted: true, Permissions: map[string]bool{"knowledge.read": true}})
|
||
|
|
}
|
||
|
|
require.Error(t, s.Apply(ctx, "kb", job.ID, []string{jobs[0].Items[0].ID}))
|
||
|
|
tags, e := repository.NewKnowledgeRepository(s.db).GetKnowledgeTags(ctx, []string{"doc"})
|
||
|
|
require.NoError(t, e)
|
||
|
|
require.Empty(t, tags["doc"])
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
func TestKnowledgeClassificationWorkerRevocationFailureAndRetry(t *testing.T) {
|
||
|
|
s, ctx, model := classificationFixture(t)
|
||
|
|
job, e := s.Create(ctx, "kb", "model", []string{"doc"})
|
||
|
|
require.NoError(t, e)
|
||
|
|
require.NoError(t, s.db.Model(&types.User{}).Where("id='alice'").Update("is_active", false).Error)
|
||
|
|
require.NoError(t, s.RunNext(context.Background()))
|
||
|
|
jobs, e := s.List(ctx, "kb")
|
||
|
|
require.NoError(t, e)
|
||
|
|
require.Equal(t, "completed_with_errors", jobs[0].Status)
|
||
|
|
require.Empty(t, model.messages)
|
||
|
|
require.NoError(t, s.db.Model(&types.User{}).Where("id='alice'").Update("is_active", true).Error)
|
||
|
|
require.NoError(t, s.Retry(ctx, "kb", job.ID))
|
||
|
|
require.NoError(t, s.RunNext(context.Background()))
|
||
|
|
jobs, e = s.List(ctx, "kb")
|
||
|
|
require.NoError(t, e)
|
||
|
|
require.Equal(t, "completed", jobs[0].Status)
|
||
|
|
require.Equal(t, 1, jobs[0].Processed)
|
||
|
|
require.Zero(t, jobs[0].Failed)
|
||
|
|
}
|
||
|
|
func TestKnowledgeClassificationRejectsInventedTagsAndResumesExpiredLease(t *testing.T) {
|
||
|
|
s, ctx, model := classificationFixture(t)
|
||
|
|
model.response = `{"matches":[{"index":999,"confidence":1,"reason":"invented"},{"index":1,"confidence":0.1,"reason":"uncertain"}]}`
|
||
|
|
job, e := s.Create(ctx, "kb", "model", []string{"doc"})
|
||
|
|
require.NoError(t, e)
|
||
|
|
require.NoError(t, s.db.Model(&KnowledgeClassificationItem{}).Where("job_id=?", job.ID).Updates(map[string]any{"status": "running", "lease_token": "old-worker", "lease_until": time.Now().Add(-time.Minute)}).Error)
|
||
|
|
require.NoError(t, s.RunNext(context.Background()))
|
||
|
|
jobs, e := s.List(ctx, "kb")
|
||
|
|
require.NoError(t, e)
|
||
|
|
require.Equal(t, "no_match", jobs[0].Items[0].Status)
|
||
|
|
require.Error(t, s.Apply(ctx, "kb", job.ID, []string{jobs[0].Items[0].ID}))
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestKnowledgeClassificationOldLeaseCannotOverwriteRecoveredResult(t *testing.T) {
|
||
|
|
s, ctx, _ := classificationFixture(t)
|
||
|
|
job, err := s.Create(ctx, "kb", "model", []string{"doc"})
|
||
|
|
require.NoError(t, err)
|
||
|
|
old := job.Items[0]
|
||
|
|
old.LeaseToken = "expired-worker"
|
||
|
|
require.NoError(t, s.db.Model(&KnowledgeClassificationItem{}).Where("id=?", old.ID).Updates(map[string]any{"status": "running", "lease_token": old.LeaseToken, "lease_until": time.Now().Add(-time.Minute)}).Error)
|
||
|
|
require.NoError(t, s.RunNext(context.Background()))
|
||
|
|
before, err := s.List(ctx, "kb")
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.NoError(t, s.finish(context.Background(), *job, old, "failed", "late failure", "", nil))
|
||
|
|
after, err := s.List(ctx, "kb")
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, before, after)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestKnowledgeClassificationConcurrentCompletionCounts(t *testing.T) {
|
||
|
|
s, ctx, _ := classificationFixture(t)
|
||
|
|
job, err := s.Create(ctx, "kb", "model", []string{"doc"})
|
||
|
|
require.NoError(t, err)
|
||
|
|
const total = 8
|
||
|
|
items := make([]KnowledgeClassificationItem, total)
|
||
|
|
for i := range items {
|
||
|
|
items[i] = KnowledgeClassificationItem{ID: fmt.Sprintf("parallel-%d", i), JobID: job.ID, KnowledgeID: fmt.Sprintf("parallel-doc-%d", i), Status: "running", LeaseToken: fmt.Sprintf("lease-%d", i), Suggestions: types.JSON(`[]`)}
|
||
|
|
}
|
||
|
|
require.NoError(t, s.db.Where("job_id=?", job.ID).Delete(&KnowledgeClassificationItem{}).Error)
|
||
|
|
require.NoError(t, s.db.Create(&items).Error)
|
||
|
|
require.NoError(t, s.db.Model(job).Update("total", total).Error)
|
||
|
|
var wg sync.WaitGroup
|
||
|
|
failures := make(chan error, total)
|
||
|
|
for _, item := range items {
|
||
|
|
wg.Add(1)
|
||
|
|
go func(item KnowledgeClassificationItem) {
|
||
|
|
defer wg.Done()
|
||
|
|
failures <- s.finish(context.Background(), *job, item, "no_match", "", "", nil)
|
||
|
|
}(item)
|
||
|
|
}
|
||
|
|
wg.Wait()
|
||
|
|
close(failures)
|
||
|
|
for err := range failures {
|
||
|
|
require.NoError(t, err)
|
||
|
|
}
|
||
|
|
require.NoError(t, s.db.First(job, "id=?", job.ID).Error)
|
||
|
|
require.Equal(t, total, job.Processed)
|
||
|
|
require.Equal(t, "completed", job.Status)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestKnowledgeClassificationMasksUntrustedErrors(t *testing.T) {
|
||
|
|
const secret = "sk-provider-private-value response-body password=secret"
|
||
|
|
for _, source := range []error{errors.New(secret), werrors.NewBadRequestError(secret).WithDetails(secret), werrors.NewForbiddenError(secret)} {
|
||
|
|
t.Run(fmt.Sprintf("%T-%s", source, classificationErrorCode(source)), func(t *testing.T) {
|
||
|
|
s, ctx, model := classificationFixture(t)
|
||
|
|
s.models = &stubModelService{chatModel: classificationTestModel{autoTagChatModel: model, call: func(_ context.Context, _ []chat.Message, options *chat.ChatOptions) (*types.ChatResponse, error) {
|
||
|
|
require.True(t, json.Valid(options.Format), "JSON response mode required")
|
||
|
|
return nil, source
|
||
|
|
}}}
|
||
|
|
_, err := s.Create(ctx, "kb", "model", []string{"doc"})
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.NoError(t, s.RunNext(context.Background()))
|
||
|
|
jobs, err := s.List(ctx, "kb")
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, "failed", jobs[0].Items[0].Status)
|
||
|
|
body, err := json.Marshal(jobs)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.NotContains(t, string(body), secret)
|
||
|
|
require.NotContains(t, KnowledgeClassificationPublicError(source).Error(), secret)
|
||
|
|
require.Nil(t, KnowledgeClassificationPublicError(source).Details)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
require.Equal(t, "已知的安全说明", KnowledgeClassificationPublicError(classificationConflict("已知的安全说明")).Message)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestKnowledgeClassificationRevokedDuringModelCall(t *testing.T) {
|
||
|
|
for _, revoke := range []string{"membership", "permission", "hidden"} {
|
||
|
|
t.Run(revoke, func(t *testing.T) {
|
||
|
|
s, ctx, model := classificationFixture(t)
|
||
|
|
revoked := false
|
||
|
|
s.snapshot = func(context.Context) (*accesspolicy.Snapshot, error) {
|
||
|
|
snapshot := &accesspolicy.Snapshot{Restricted: true, Permissions: map[string]bool{"knowledge.write": !revoked || revoke != "permission"}}
|
||
|
|
if revoked && revoke == "hidden" {
|
||
|
|
snapshot.Hidden = map[string]bool{"knowledge:kb": true}
|
||
|
|
}
|
||
|
|
return snapshot, nil
|
||
|
|
}
|
||
|
|
s.models = &stubModelService{chatModel: classificationTestModel{autoTagChatModel: model, call: func(_ context.Context, _ []chat.Message, _ *chat.ChatOptions) (*types.ChatResponse, error) {
|
||
|
|
revoked = true
|
||
|
|
if revoke == "membership" {
|
||
|
|
require.NoError(t, s.db.Model(&types.TenantMember{}).Where("user_id='alice'").Update("status", "inactive").Error)
|
||
|
|
}
|
||
|
|
return &types.ChatResponse{Content: model.response}, nil
|
||
|
|
}}}
|
||
|
|
_, err := s.Create(ctx, "kb", "model", []string{"doc"})
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.NoError(t, s.RunNext(context.Background()))
|
||
|
|
jobs, err := s.List(ctx, "kb")
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, "failed", jobs[0].Items[0].Status)
|
||
|
|
require.JSONEq(t, `[]`, string(jobs[0].Items[0].Suggestions))
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestKnowledgeClassificationRetryBlockedByAnotherActiveJob(t *testing.T) {
|
||
|
|
s, ctx, _ := classificationFixture(t)
|
||
|
|
job, err := s.Create(ctx, "kb", "model", []string{"doc"})
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.NoError(t, s.db.Model(&KnowledgeClassificationItem{}).Where("job_id=?", job.ID).Update("status", "failed").Error)
|
||
|
|
require.NoError(t, s.db.Model(job).Updates(map[string]any{"status": "completed_with_errors", "processed": 1, "failed": 1}).Error)
|
||
|
|
_, err = s.Create(ctx, "kb", "model", []string{"doc"})
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.ErrorContains(t, s.Retry(ctx, "kb", job.ID), "已有分类任务")
|
||
|
|
var item KnowledgeClassificationItem
|
||
|
|
require.NoError(t, s.db.First(&item, "job_id=?", job.ID).Error)
|
||
|
|
require.Equal(t, "failed", item.Status)
|
||
|
|
}
|