story/internal/extraction/processor.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 {
		return outbox.Retryable(ErrConversationContentNotReady)
	}
	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]) + "…"
}