348 lines
13 KiB
Go
348 lines
13 KiB
Go
package extraction
|
||
|
||
import (
|
||
"context"
|
||
"crypto/sha256"
|
||
"encoding/hex"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"strings"
|
||
"unicode/utf8"
|
||
|
||
"data-terminal-cloud/internal/contentpolicy"
|
||
"data-terminal-cloud/internal/outbox"
|
||
|
||
"github.com/jackc/pgx/v5"
|
||
"github.com/jackc/pgx/v5/pgxpool"
|
||
)
|
||
|
||
var ErrConversationContentNotReady = errors.New("conversation content is not ready for extraction")
|
||
var ErrConversationContentChanged = errors.New("conversation content changed while extraction was running")
|
||
|
||
type Processor struct {
|
||
pool *pgxpool.Pool
|
||
generator Generator
|
||
safety contentpolicy.Classifier
|
||
}
|
||
|
||
type turn struct {
|
||
id, publicID, speaker, text string
|
||
interrupted bool
|
||
}
|
||
|
||
type extractionSnapshot struct {
|
||
exists bool
|
||
tenantID, subjectID, status string
|
||
storyExtractionAllowed bool
|
||
turns []turn
|
||
maxSequence int64
|
||
digest string
|
||
}
|
||
|
||
type queryer interface {
|
||
QueryRow(context.Context, string, ...any) pgx.Row
|
||
Query(context.Context, string, ...any) (pgx.Rows, error)
|
||
}
|
||
|
||
func NewProcessor(pool *pgxpool.Pool, generators ...Generator) (*Processor, error) {
|
||
if pool == nil {
|
||
return nil, errors.New("database pool is required")
|
||
}
|
||
var generator Generator = DeterministicGenerator{}
|
||
if len(generators) > 1 || (len(generators) == 1 && generators[0] == nil) {
|
||
return nil, errors.New("one valid extraction generator is required")
|
||
}
|
||
if len(generators) == 1 {
|
||
generator = generators[0]
|
||
}
|
||
return &Processor{pool: pool, generator: generator, safety: contentpolicy.LocalClassifier{}}, nil
|
||
}
|
||
|
||
func NewProcessorWithSafety(pool *pgxpool.Pool, generator Generator, safety contentpolicy.Classifier) (*Processor, error) {
|
||
if generator == nil || safety == nil {
|
||
return nil, errors.New("valid extraction generator and content safety classifier are required")
|
||
}
|
||
processor, err := NewProcessor(pool, generator)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
processor.safety = safety
|
||
return processor, nil
|
||
}
|
||
|
||
func (p *Processor) Handle(ctx context.Context, event outbox.Event) error {
|
||
snapshot, err := loadExtractionSnapshot(ctx, p.pool, event.AggregateID, false)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if !snapshot.exists {
|
||
return nil
|
||
}
|
||
if snapshot.status != "COMPLETED" && snapshot.status != "FAILED" {
|
||
return outbox.Retryable(ErrConversationContentNotReady)
|
||
}
|
||
if !snapshot.storyExtractionAllowed {
|
||
return nil
|
||
}
|
||
generationVersion := requestedGenerationVersion(event.Payload)
|
||
if generationVersion < snapshot.maxSequence {
|
||
return nil
|
||
}
|
||
sourceTurns := subjectTurns(snapshot.turns)
|
||
if len(sourceTurns) == 0 {
|
||
// Terminal conversations cannot receive more subject turns. Treat the
|
||
// extraction request as complete so permanently empty conversations do
|
||
// not cycle through retry and dead-letter queues forever.
|
||
return nil
|
||
}
|
||
content := joinedConversationContent(snapshot.turns)
|
||
generated, err := p.generateSafely(ctx, content, snapshot.tenantID)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
// The model call above is deliberately outside a database transaction. The
|
||
// conversation row is locked only for the short revalidation and write phase.
|
||
tx, err := p.pool.BeginTx(ctx, pgx.TxOptions{})
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer tx.Rollback(context.Background())
|
||
current, err := loadExtractionSnapshot(ctx, tx, event.AggregateID, true)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if !current.exists {
|
||
return tx.Commit(ctx)
|
||
}
|
||
if !current.storyExtractionAllowed {
|
||
return tx.Commit(ctx)
|
||
}
|
||
if current.status != "COMPLETED" && current.status != "FAILED" {
|
||
return outbox.Retryable(ErrConversationContentNotReady)
|
||
}
|
||
if current.digest != snapshot.digest {
|
||
return outbox.Retryable(ErrConversationContentChanged)
|
||
}
|
||
return storeGeneration(ctx, tx, event.AggregateID, current, generated, generationVersion)
|
||
}
|
||
|
||
func (p *Processor) generateSafely(ctx context.Context, content string, tenantIDs ...string) (GenerationResult, error) {
|
||
tenantID := ""
|
||
if len(tenantIDs) > 0 {
|
||
tenantID = tenantIDs[0]
|
||
}
|
||
decision, err := p.safety.Classify(ctx, content)
|
||
if err != nil {
|
||
return GenerationResult{}, err
|
||
}
|
||
if decision.AllowOrdinaryStory {
|
||
generated, generateErr := p.generator.Generate(ctx, GenerationInput{TenantID: tenantID, Content: content})
|
||
if generateErr != nil {
|
||
return GenerationResult{}, generateErr
|
||
}
|
||
outputDecision, safetyErr := p.safety.Classify(ctx, generatedSafetyContent(generated))
|
||
if safetyErr != nil {
|
||
return GenerationResult{}, safetyErr
|
||
}
|
||
if !outputDecision.AllowOrdinaryStory {
|
||
return sensitiveGeneration(outputDecision), nil
|
||
}
|
||
return generated, nil
|
||
}
|
||
return sensitiveGeneration(decision), nil
|
||
}
|
||
|
||
func generatedSafetyContent(result GenerationResult) string {
|
||
return strings.Join([]string{result.Summary, strings.Join(result.Keywords, "、"), result.OpenTopic, result.StoryTitle, result.StoryContent, result.FactContent}, "\n")
|
||
}
|
||
|
||
func loadExtractionSnapshot(ctx context.Context, q queryer, conversationID string, lock bool) (extractionSnapshot, error) {
|
||
lockClause := ""
|
||
if lock {
|
||
lockClause = " FOR UPDATE OF c"
|
||
}
|
||
var snapshot extractionSnapshot
|
||
var consentsJSON []byte
|
||
err := q.QueryRow(ctx, `
|
||
SELECT c.tenant_id::text,c.subject_id::text,c.status,cs.effective_consents
|
||
FROM conversation c JOIN consent_snapshot cs ON cs.id=c.consent_snapshot_id AND cs.tenant_id=c.tenant_id
|
||
JOIN subject_profile sp ON sp.id=c.subject_id AND sp.tenant_id=c.tenant_id AND sp.status IN ('active','suspended')
|
||
WHERE c.id=$1::uuid`+lockClause, conversationID).Scan(&snapshot.tenantID, &snapshot.subjectID, &snapshot.status, &consentsJSON)
|
||
if errors.Is(err, pgx.ErrNoRows) {
|
||
return extractionSnapshot{}, nil
|
||
}
|
||
if err != nil {
|
||
return extractionSnapshot{}, fmt.Errorf("read extraction conversation: %w", err)
|
||
}
|
||
snapshot.exists = true
|
||
consents := map[string]bool{}
|
||
if err = json.Unmarshal(consentsJSON, &consents); err != nil {
|
||
return extractionSnapshot{}, fmt.Errorf("decode extraction consents: %w", err)
|
||
}
|
||
snapshot.storyExtractionAllowed = consents["story_extraction"]
|
||
rows, err := q.Query(ctx, `
|
||
SELECT id::text,public_id::text,speaker,final_text,interrupted,sequence_no
|
||
FROM conversation_turn WHERE conversation_id=$1::uuid ORDER BY sequence_no,id`, conversationID)
|
||
if err != nil {
|
||
return extractionSnapshot{}, fmt.Errorf("read extraction turns: %w", err)
|
||
}
|
||
defer rows.Close()
|
||
for rows.Next() {
|
||
var item turn
|
||
var sequence int64
|
||
if err = rows.Scan(&item.id, &item.publicID, &item.speaker, &item.text, &item.interrupted, &sequence); err != nil {
|
||
return extractionSnapshot{}, err
|
||
}
|
||
if sequence > snapshot.maxSequence {
|
||
snapshot.maxSequence = sequence
|
||
}
|
||
snapshot.turns = append(snapshot.turns, item)
|
||
}
|
||
if err = rows.Err(); err != nil {
|
||
return extractionSnapshot{}, err
|
||
}
|
||
snapshot.digest = snapshotDigest(snapshot)
|
||
return snapshot, nil
|
||
}
|
||
|
||
func requestedGenerationVersion(payload json.RawMessage) int64 {
|
||
var value struct {
|
||
GenerationVersion int64 `json:"generation_version"`
|
||
}
|
||
if json.Unmarshal(payload, &value) == nil && value.GenerationVersion > 0 {
|
||
return value.GenerationVersion
|
||
}
|
||
return 1
|
||
}
|
||
|
||
func snapshotDigest(snapshot extractionSnapshot) string {
|
||
type digestTurn struct {
|
||
ID string `json:"id"`
|
||
PublicID string `json:"public_id"`
|
||
Speaker string `json:"speaker"`
|
||
Text string `json:"text"`
|
||
Interrupted bool `json:"interrupted"`
|
||
}
|
||
digestTurns := make([]digestTurn, 0, len(snapshot.turns))
|
||
for _, item := range snapshot.turns {
|
||
digestTurns = append(digestTurns, digestTurn{item.id, item.publicID, item.speaker, item.text, item.interrupted})
|
||
}
|
||
digestInput := struct {
|
||
Exists bool
|
||
TenantID, SubjectID, Status string
|
||
Allowed bool
|
||
Turns []digestTurn
|
||
}{snapshot.exists, snapshot.tenantID, snapshot.subjectID, snapshot.status, snapshot.storyExtractionAllowed, digestTurns}
|
||
encoded, _ := json.Marshal(digestInput)
|
||
digest := sha256.Sum256(encoded)
|
||
return hex.EncodeToString(digest[:])
|
||
}
|
||
|
||
func subjectTurns(turns []turn) []turn {
|
||
items := make([]turn, 0, len(turns))
|
||
for _, item := range turns {
|
||
if item.speaker == "subject" {
|
||
items = append(items, item)
|
||
}
|
||
}
|
||
return items
|
||
}
|
||
|
||
func joinedConversationContent(turns []turn) string {
|
||
parts := make([]string, 0, len(turns))
|
||
for _, item := range turns {
|
||
label := "豆包"
|
||
if item.speaker == "subject" {
|
||
label = "长者"
|
||
}
|
||
parts = append(parts, label+"："+item.text)
|
||
}
|
||
return strings.Join(parts, "\n")
|
||
}
|
||
|
||
func storeGeneration(ctx context.Context, tx pgx.Tx, conversationID string, snapshot extractionSnapshot, generated GenerationResult, generationVersion int64) error {
|
||
sourceTurns := subjectTurns(snapshot.turns)
|
||
sourcePublicIDs := make([]string, 0, len(sourceTurns))
|
||
for _, item := range sourceTurns {
|
||
sourcePublicIDs = append(sourcePublicIDs, item.publicID)
|
||
}
|
||
sourceJSON, _ := json.Marshal(sourcePublicIDs)
|
||
if _, err := tx.Exec(ctx, `
|
||
INSERT INTO conversation_content_assessment (tenant_id,subject_id,conversation_id,generation_version,meaningful,story_worthy,decision_reason,model_ref,prompt_version,confidence)
|
||
VALUES ($1::uuid,$2::uuid,$3::uuid,$4,$5,$6,$7,$8,$9,$10)
|
||
ON CONFLICT (conversation_id,generation_version) DO NOTHING`, snapshot.tenantID, snapshot.subjectID, conversationID, generationVersion, generated.Meaningful, generated.StoryWorthy, generated.DecisionReason, generated.ModelRef, generated.PromptVersion, generated.Confidence); err != nil {
|
||
return fmt.Errorf("store conversation content assessment: %w", err)
|
||
}
|
||
if !generated.Meaningful {
|
||
return tx.Commit(ctx)
|
||
}
|
||
if _, err := tx.Exec(ctx, `
|
||
INSERT INTO conversation_summary (tenant_id,subject_id,conversation_id,summary_type,content,source_turn_ids,model_ref,prompt_version,generation_version,confidence,status)
|
||
VALUES ($1::uuid,$2::uuid,$3::uuid,'conversation_summary',$4,$5::jsonb,$6,$7,$8,$9,'READY')
|
||
ON CONFLICT (conversation_id,summary_type,generation_version) DO NOTHING`, snapshot.tenantID, snapshot.subjectID, conversationID, summaryWithKeywords(generated.Summary, generated.Keywords), sourceJSON, generated.ModelRef, generated.PromptVersion, generationVersion, generated.Confidence); err != nil {
|
||
return fmt.Errorf("store conversation summary: %w", err)
|
||
}
|
||
if _, err := tx.Exec(ctx, `
|
||
INSERT INTO conversation_summary (tenant_id,subject_id,conversation_id,summary_type,content,source_turn_ids,model_ref,prompt_version,generation_version,confidence,status)
|
||
VALUES ($1::uuid,$2::uuid,$3::uuid,'open_topic',$4,$5::jsonb,$6,$7,$8,$9,'READY')
|
||
ON CONFLICT (conversation_id,summary_type,generation_version) DO NOTHING`, snapshot.tenantID, snapshot.subjectID, conversationID, generated.OpenTopic, sourceJSON, generated.ModelRef, generated.PromptVersion, generationVersion, generated.Confidence); err != nil {
|
||
return fmt.Errorf("store open topic: %w", err)
|
||
}
|
||
if !generated.AllowStory {
|
||
return tx.Commit(ctx)
|
||
}
|
||
var storyID string
|
||
err := tx.QueryRow(ctx, `
|
||
INSERT INTO story_draft (tenant_id,subject_id,conversation_id,title,status,confidence,sensitivity,generation_version)
|
||
VALUES ($1::uuid,$2::uuid,$3::uuid,$4,'CANDIDATE',$5,'normal',$6)
|
||
ON CONFLICT (conversation_id,generation_version) DO UPDATE SET updated_at=story_draft.updated_at
|
||
RETURNING id::text`, snapshot.tenantID, snapshot.subjectID, conversationID, generated.StoryTitle, generated.Confidence, generationVersion).Scan(&storyID)
|
||
if err != nil {
|
||
return fmt.Errorf("store story candidate: %w", err)
|
||
}
|
||
var revisionID string
|
||
err = tx.QueryRow(ctx, `
|
||
INSERT INTO story_revision (story_id,revision_no,content,author_type,model_ref,prompt_version)
|
||
VALUES ($1::uuid,1,$2,'model',$3,$4)
|
||
ON CONFLICT (story_id,revision_no) DO UPDATE SET content=story_revision.content
|
||
RETURNING id::text`, storyID, generated.StoryContent, generated.ModelRef, generated.PromptVersion).Scan(&revisionID)
|
||
if err != nil {
|
||
return fmt.Errorf("store story revision: %w", err)
|
||
}
|
||
if _, err = tx.Exec(ctx, `UPDATE story_draft SET current_revision_id=$1::uuid WHERE id=$2::uuid AND current_revision_id IS NULL`, revisionID, storyID); err != nil {
|
||
return err
|
||
}
|
||
for _, item := range sourceTurns {
|
||
if _, err = tx.Exec(ctx, `
|
||
INSERT INTO story_source (story_revision_id,conversation_turn_id,source_start_offset,source_end_offset,source_role)
|
||
VALUES ($1::uuid,$2::uuid,0,$3,'subject') ON CONFLICT DO NOTHING`, revisionID, item.id, utf8.RuneCountInString(item.text)); err != nil {
|
||
return fmt.Errorf("store story source: %w", err)
|
||
}
|
||
}
|
||
if _, err = tx.Exec(ctx, `
|
||
INSERT INTO fact_candidate (tenant_id,subject_id,story_id,fact_type,content,statement_basis,time_range,confidence,sensitivity,source_turn_ids,status,generation_version)
|
||
SELECT $1::uuid,$2::uuid,$3::uuid,'life_event',$4,'subject_statement','{}'::jsonb,$5,'normal',$6::jsonb,'CANDIDATE',$7
|
||
WHERE NOT EXISTS (SELECT 1 FROM fact_candidate WHERE story_id=$3::uuid AND generation_version=$7)`, snapshot.tenantID, snapshot.subjectID, storyID, generated.FactContent, generated.Confidence, sourceJSON, generationVersion); err != nil {
|
||
return fmt.Errorf("store fact candidate: %w", err)
|
||
}
|
||
return tx.Commit(ctx)
|
||
}
|
||
|
||
func summaryWithKeywords(summary string, keywords []string) string {
|
||
keywords = normalizedKeywords(keywords)
|
||
if len(keywords) == 0 {
|
||
return summary
|
||
}
|
||
return summary + "\n关键词：" + strings.Join(keywords, "、")
|
||
}
|
||
|
||
func truncateRunes(value string, limit int) string {
|
||
characters := []rune(value)
|
||
if len(characters) <= limit {
|
||
return value
|
||
}
|
||
return string(characters[:limit]) + "…"
|
||
}
|