178 lines
10 KiB
Go
178 lines
10 KiB
Go
package extraction
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"strings"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"data-terminal-cloud/internal/outbox"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
)
|
|
|
|
func TestProcessorExtractionIdempotencyAndSensitiveContent(t *testing.T) {
|
|
databaseURL := os.Getenv("TEST_DATABASE_URL")
|
|
if databaseURL == "" {
|
|
t.Skip("TEST_DATABASE_URL is not set")
|
|
}
|
|
ctx := context.Background()
|
|
pool, err := pgxpool.New(ctx, databaseURL)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer pool.Close()
|
|
if _, err := NewProcessor(nil); err == nil {
|
|
t.Fatal("nil database pool was accepted")
|
|
}
|
|
processor, err := NewProcessor(pool)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
consentSnapshotID, consentSnapshotPublicID := uuid.NewString(), uuid.NewString()
|
|
contextID, contextPublicID := uuid.NewString(), uuid.NewString()
|
|
conversationID, conversationPublicID := uuid.NewString(), uuid.NewString()
|
|
emptyConversationID, emptyConversationPublicID := uuid.NewString(), uuid.NewString()
|
|
defer func() {
|
|
_, _ = pool.Exec(ctx, `DELETE FROM conversation_content_assessment WHERE conversation_id=$1::uuid`, conversationID)
|
|
_, _ = pool.Exec(ctx, `DELETE FROM fact_candidate WHERE story_id IN (SELECT id FROM story_draft WHERE conversation_id=$1::uuid)`, conversationID)
|
|
_, _ = pool.Exec(ctx, `DELETE FROM story_source WHERE story_revision_id IN (SELECT sr.id FROM story_revision sr JOIN story_draft sd ON sd.id=sr.story_id WHERE sd.conversation_id=$1::uuid)`, conversationID)
|
|
_, _ = pool.Exec(ctx, `UPDATE story_draft SET current_revision_id=NULL WHERE conversation_id=$1::uuid`, conversationID)
|
|
_, _ = pool.Exec(ctx, `DELETE FROM story_revision WHERE story_id IN (SELECT id FROM story_draft WHERE conversation_id=$1::uuid)`, conversationID)
|
|
_, _ = pool.Exec(ctx, `DELETE FROM story_draft WHERE conversation_id=$1::uuid`, conversationID)
|
|
_, _ = pool.Exec(ctx, `DELETE FROM conversation_summary WHERE conversation_id=$1::uuid`, conversationID)
|
|
_, _ = pool.Exec(ctx, `DELETE FROM conversation_turn WHERE conversation_id=$1::uuid`, conversationID)
|
|
_, _ = pool.Exec(ctx, `DELETE FROM conversation WHERE id=$1::uuid`, conversationID)
|
|
_, _ = pool.Exec(ctx, `DELETE FROM conversation WHERE id=$1::uuid`, emptyConversationID)
|
|
_, _ = pool.Exec(ctx, `DELETE FROM context_snapshot WHERE id=$1::uuid`, contextID)
|
|
_, _ = pool.Exec(ctx, `DELETE FROM consent_snapshot WHERE id=$1::uuid`, consentSnapshotID)
|
|
}()
|
|
if _, err := pool.Exec(ctx, `
|
|
INSERT INTO consent_snapshot(id,public_id,tenant_id,subject_id,snapshot_version,effective_consents,source_consent_ids)
|
|
VALUES($1::uuid,$2::uuid,'10000000-0000-0000-0000-000000000001','10000000-0000-0000-0000-000000000401',1,
|
|
'{"story_extraction":true,"transcript_storage":true}','[]')`, consentSnapshotID, consentSnapshotPublicID); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := pool.Exec(ctx, `
|
|
INSERT INTO context_snapshot(id,public_id,tenant_id,subject_id,profile_version,consent_snapshot_id,memory_versions,prompt_bundle_version,engine_mode,content_digest,sanitized_summary)
|
|
VALUES($1::uuid,$2::uuid,'10000000-0000-0000-0000-000000000001','10000000-0000-0000-0000-000000000401',1,
|
|
$3::uuid,'[]','extraction-integration-v1','s2s_hybrid',$4,'合成提取上下文')`,
|
|
contextID, contextPublicID, consentSnapshotID, "d"+strings.Repeat("0", 63)); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := pool.Exec(ctx, `
|
|
INSERT INTO conversation(id,public_id,tenant_id,household_id,subject_id,device_id,consent_snapshot_id,context_snapshot_id,
|
|
provider_task_ref,engine_mode,retention_mode,status,allocation_expires_at)
|
|
VALUES($1::uuid,$2::uuid,'10000000-0000-0000-0000-000000000001','10000000-0000-0000-0000-000000000201',
|
|
'10000000-0000-0000-0000-000000000401','10000000-0000-0000-0000-000000000501',
|
|
$3::uuid,$4::uuid,$5,'s2s_hybrid','transcript_only','ACTIVE',now()+interval '5 minutes')`,
|
|
conversationID, conversationPublicID, consentSnapshotID, contextID, "extract-task-"+conversationID); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := pool.Exec(ctx, `
|
|
INSERT INTO conversation(id,public_id,tenant_id,household_id,subject_id,device_id,consent_snapshot_id,context_snapshot_id,
|
|
provider_task_ref,engine_mode,retention_mode,status,allocation_expires_at,started_at,ended_at,stop_reason)
|
|
VALUES($1::uuid,$2::uuid,'10000000-0000-0000-0000-000000000001','10000000-0000-0000-0000-000000000201',
|
|
'10000000-0000-0000-0000-000000000401','10000000-0000-0000-0000-000000000501',
|
|
$3::uuid,$4::uuid,$5,'s2s_hybrid','transcript_only','COMPLETED',now()+interval '5 minutes',now()-interval '1 minute',now(),'user_hangup')`,
|
|
emptyConversationID, emptyConversationPublicID, consentSnapshotID, contextID, "extract-empty-task-"+emptyConversationID); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := processor.Handle(ctx, outbox.Event{AggregateID: emptyConversationID, Payload: []byte(`{"generation_version":1}`)}); err != nil {
|
|
t.Fatalf("terminal conversation without subject turns should not retry: %v", err)
|
|
}
|
|
if _, err := pool.Exec(ctx, `
|
|
INSERT INTO conversation_turn(tenant_id,subject_id,conversation_id,turn_key,sequence_no,speaker,start_at,end_at,final_text,source_type)
|
|
VALUES('10000000-0000-0000-0000-000000000001','10000000-0000-0000-0000-000000000401',$1::uuid,$2,1,'subject',now()-interval '2 minutes',now()-interval '110 seconds',$3,'provider'),
|
|
('10000000-0000-0000-0000-000000000001','10000000-0000-0000-0000-000000000401',$1::uuid,$4,2,'agent',now()-interval '100 seconds',now()-interval '90 seconds',$5,'provider')`,
|
|
conversationID, "extract-subject-"+conversationID, "我年轻时在河边小学教书。", "extract-agent-"+conversationID, "这是很珍贵的经历。"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
event := outbox.Event{AggregateID: conversationID, Payload: []byte(`{"generation_version":2}`)}
|
|
if err := processor.Handle(ctx, event); !errors.Is(err, ErrConversationContentNotReady) || !outbox.IsRetryable(err) {
|
|
t.Fatalf("active conversation extraction error=%v", err)
|
|
}
|
|
if _, err := pool.Exec(ctx, `UPDATE conversation SET status='COMPLETED',started_at=now()-interval '3 minutes',ended_at=now(),stop_reason='user_hangup' WHERE id=$1::uuid`, conversationID); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var modelCalls atomic.Int32
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
|
modelCalls.Add(1)
|
|
updateContext, cancel := context.WithTimeout(context.Background(), time.Second)
|
|
defer cancel()
|
|
if _, updateErr := pool.Exec(updateContext, `UPDATE conversation SET updated_at=updated_at WHERE id=$1::uuid`, conversationID); updateErr != nil {
|
|
t.Errorf("model call was made while the conversation row was transaction-locked: %v", updateErr)
|
|
http.Error(w, "database lock", http.StatusServiceUnavailable)
|
|
return
|
|
}
|
|
_, _ = io.WriteString(w, `{"choices":[{"message":{"content":"{\"meaningful\":true,\"story_worthy\":true,\"decision_reason\":\"包含具体的教书经历\",\"summary\":\"长者谈到在河边小学教书\",\"keywords\":[\"河边小学\",\"教书经历\"],\"open_topic\":\"下次可以聊聊当年的学生\",\"story_title\":\"河边小学的回忆\",\"story_content\":\"我年轻时在河边小学教书。\",\"fact_content\":\"年轻时曾在河边小学教书\",\"confidence\":0.88}"}}]}`)
|
|
}))
|
|
defer server.Close()
|
|
arkGenerator, err := NewArkGenerator(server.URL+"/api/v3", "integration-key", "doubao-integration", 2*time.Second)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
modelProcessor, err := NewProcessor(pool, arkGenerator)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := modelProcessor.Handle(ctx, event); err != nil {
|
|
t.Fatalf("extract completed conversation: %v", err)
|
|
}
|
|
if err := modelProcessor.Handle(ctx, event); err != nil {
|
|
t.Fatalf("replay extraction: %v", err)
|
|
}
|
|
var summaries, stories, revisions, sources, facts int
|
|
if err := pool.QueryRow(ctx, `SELECT
|
|
(SELECT count(*) FROM conversation_summary WHERE conversation_id=$1::uuid),
|
|
(SELECT count(*) FROM story_draft WHERE conversation_id=$1::uuid),
|
|
(SELECT count(*) FROM story_revision WHERE story_id IN (SELECT id FROM story_draft WHERE conversation_id=$1::uuid)),
|
|
(SELECT count(*) FROM story_source WHERE story_revision_id IN (SELECT sr.id FROM story_revision sr JOIN story_draft sd ON sd.id=sr.story_id WHERE sd.conversation_id=$1::uuid)),
|
|
(SELECT count(*) FROM fact_candidate WHERE story_id IN (SELECT id FROM story_draft WHERE conversation_id=$1::uuid))`, conversationID).
|
|
Scan(&summaries, &stories, &revisions, &sources, &facts); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if summaries != 2 || stories != 1 || revisions != 1 || sources != 1 || facts != 1 {
|
|
t.Fatalf("normal extraction counts summaries=%d stories=%d revisions=%d sources=%d facts=%d", summaries, stories, revisions, sources, facts)
|
|
}
|
|
if _, err := pool.Exec(ctx, `
|
|
INSERT INTO conversation_turn(tenant_id,subject_id,conversation_id,turn_key,sequence_no,speaker,start_at,end_at,final_text,source_type)
|
|
VALUES('10000000-0000-0000-0000-000000000001','10000000-0000-0000-0000-000000000401',$1::uuid,$2,3,'subject',now()-interval '20 seconds',now()-interval '10 seconds',$3,'provider')`,
|
|
conversationID, "extract-sensitive-"+conversationID, "我的银行卡号是6222021234567890123，短信验证码是482931。"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
event.Payload = []byte(`{"generation_version":3}`)
|
|
if err := modelProcessor.Handle(ctx, event); err != nil {
|
|
t.Fatalf("extract sensitive conversation: %v", err)
|
|
}
|
|
if modelCalls.Load() != 2 {
|
|
t.Fatalf("model calls=%d want=2; sensitive content must not leave the local safety boundary", modelCalls.Load())
|
|
}
|
|
var sensitivePrompt, sensitiveSummary string
|
|
if err := pool.QueryRow(ctx, `SELECT prompt_version,content FROM conversation_summary WHERE conversation_id=$1::uuid AND summary_type='conversation_summary' AND generation_version=3`, conversationID).Scan(&sensitivePrompt, &sensitiveSummary); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if !strings.HasPrefix(sensitivePrompt, "summary-sensitive-filter-v1:") || strings.Contains(sensitiveSummary, "银行卡") {
|
|
t.Fatalf("sensitive summary prompt=%q content=%q", sensitivePrompt, sensitiveSummary)
|
|
}
|
|
if err := pool.QueryRow(ctx, `SELECT count(*) FROM story_draft WHERE conversation_id=$1::uuid`, conversationID).Scan(&stories); err != nil || stories != 1 {
|
|
t.Fatalf("sensitive extraction created story count=%d err=%v", stories, err)
|
|
}
|
|
}
|
|
|
|
func TestTruncateRunes(t *testing.T) {
|
|
if got := truncateRunes("一二三四", 3); got != "一二三…" {
|
|
t.Fatalf("truncateRunes=%q", got)
|
|
}
|
|
if got := truncateRunes("一二", 3); got != "一二" {
|
|
t.Fatalf("short truncateRunes=%q", got)
|
|
}
|
|
}
|