kefu/im/backend/internal/app/ai_agent.go

package app

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"math/rand"
	"strings"
	"sync"
	"time"
)

const (
	aiHistoryMessages = 20
	aiReplyMaxRunes   = 120
	aiJobMaxAttempts  = 2
	// A burst of incoming messages collapses into one reply, but the reader
	// should not wait longer than this for it.
	aiReplyMaxDelaySeconds = 15
)

type aiAgentBinding struct {
	UserID          int64
	ModelID         int64
	Persona         string
	DailyReplyLimit int
	Nickname        string
	Age             int
	City            string
	Bio             string
}

type aiReplyJob struct {
	ID             int64
	ConversationID int64
	AgentUserID    int64
	Attempts       int
}

// Called right after a message is committed, from every send path. It must stay
// cheap and must never fail the send: a missing reply is better than a failed
// message.
// aiGeneratedKey marks the context in which the worker writes its own reply.
// The loop it prevents is "a reply triggers another reply", which is a property
// of the message, not of the account: a person operating a managed test account
// is having a real conversation and deserves an answer.
type aiGeneratedKey struct{}

func markAIGenerated(ctx context.Context) context.Context {
	return context.WithValue(ctx, aiGeneratedKey{}, true)
}

func isAIGenerated(ctx context.Context) bool {
	value, _ := ctx.Value(aiGeneratedKey{}).(bool)
	return value
}

func (a *App) enqueueAIReply(ctx context.Context, conversationID, senderID int64, members []int64, messageID int64) {
	if !a.configBool(ctx, "ai.enabled", false) {
		return
	}
	if isAIGenerated(ctx) {
		return // The worker's own reply must never trigger the next one.
	}
	peerID := int64(0)
	for _, member := range members {
		if member != senderID {
			if peerID != 0 {
				return // Group conversations have no single managed recipient.
			}
			peerID = member
		}
	}
	if peerID == 0 {
		return
	}
	var agentUserID int64
	var testBatch string
	// Note there is no condition on the sender: whoever typed this, the reply is
	// owed by the recipient. Runaway agent-to-agent chatter is prevented above,
	// where the worker's own writes are recognised and skipped, and bounded by
	// the per-agent daily reply limits.
	err := a.db.QueryRowContext(ctx, `SELECT g.user_id,u.test_batch
		FROM ai_agents g JOIN users u ON u.id=g.user_id
		WHERE g.user_id=? AND g.enabled=1 AND u.is_test=1 AND u.status=1 AND u.deleted_at IS NULL`,
		peerID).Scan(&agentUserID, &testBatch)
	if err != nil || agentUserID == 0 {
		return
	}
	if !aiBatchAllowed(a.configPlain(ctx, "ai.allow_batches", ""), testBatch) {
		return
	}
	delay := a.aiReplyDelay(ctx)
	if _, err = a.db.ExecContext(ctx, `INSERT INTO ai_reply_jobs(conversation_id,agent_user_id,trigger_message_id,run_after)
		VALUES(?,?,?,DATE_ADD(NOW(3),INTERVAL ? MICROSECOND))
		ON DUPLICATE KEY UPDATE trigger_message_id=VALUES(trigger_message_id),
			run_after=LEAST(DATE_ADD(created_at,INTERVAL ? SECOND),VALUES(run_after))`,
		conversationID, agentUserID, messageID, delay.Microseconds(), aiReplyMaxDelaySeconds); err != nil {
		log.Printf("ai: 入队失败 conversation=%d: %v", conversationID, err)
	}
}

// An empty allow list means every test account; otherwise the account's batch
// has to be named explicitly.
func aiBatchAllowed(allowed, batch string) bool {
	allowed = strings.TrimSpace(allowed)
	if allowed == "" {
		return true
	}
	for _, item := range strings.Split(allowed, ",") {
		if strings.TrimSpace(item) == strings.TrimSpace(batch) && batch != "" {
			return true
		}
	}
	return false
}

func (a *App) aiReplyDelay(ctx context.Context) time.Duration {
	minDelay := a.configInt(ctx, "ai.reply_delay_min_ms", 2000)
	maxDelay := a.configInt(ctx, "ai.reply_delay_max_ms", 6000)
	if minDelay < 0 {
		minDelay = 0
	}
	if maxDelay <= minDelay {
		maxDelay = minDelay
	}
	spread := maxDelay - minDelay
	if spread > 0 {
		minDelay += rand.Intn(spread)
	}
	return time.Duration(minDelay) * time.Millisecond
}

// The worker follows the same ticker shape as processDueAccountClosures.
func (a *App) runAIReplyWorker(ctx context.Context) {
	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			a.processDueAIReplies(ctx)
		}
	}
}

func (a *App) processDueAIReplies(ctx context.Context) {
	if !a.configBool(ctx, "ai.enabled", false) {
		return
	}
	concurrency := a.configInt(ctx, "ai.max_concurrency", 4)
	if concurrency < 1 || concurrency > 64 {
		concurrency = 4
	}
	jobs, err := a.claimAIReplyJobs(ctx, concurrency)
	if err != nil || len(jobs) == 0 {
		return
	}
	var wait sync.WaitGroup
	for _, job := range jobs {
		wait.Add(1)
		go func(job aiReplyJob) {
			defer wait.Done()
			if replyErr := a.processAIReplyJob(ctx, job); replyErr != nil {
				a.failAIReplyJob(ctx, job, replyErr)
				return
			}
			_, _ = a.db.ExecContext(ctx, `UPDATE ai_reply_jobs SET status=2,last_error='' WHERE id=?`, job.ID)
		}(job)
	}
	wait.Wait()
}

// Claiming stamps a token with a single UPDATE and reads the rows back: each
// row can only leave status 0 once, so two instances never take the same job.
// SELECT ... FOR UPDATE SKIP LOCKED would be the obvious alternative, but it
// does not parse on MySQL 5.7, which some deployments still run.
func (a *App) claimAIReplyJobs(ctx context.Context, limit int) ([]aiReplyJob, error) {
	token := randomToken()[:32]
	result, err := a.db.ExecContext(ctx, `UPDATE ai_reply_jobs SET status=1,attempts=attempts+1,claim_token=?
		WHERE status=0 AND run_after<=NOW(3) ORDER BY run_after LIMIT ?`, token, limit)
	if err != nil {
		return nil, err
	}
	if affected, _ := result.RowsAffected(); affected == 0 {
		return nil, nil
	}
	rows, err := a.db.QueryContext(ctx, `SELECT id,conversation_id,agent_user_id,attempts FROM ai_reply_jobs WHERE claim_token=? AND status=1`, token)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	jobs := []aiReplyJob{}
	for rows.Next() {
		var job aiReplyJob
		if rows.Scan(&job.ID, &job.ConversationID, &job.AgentUserID, &job.Attempts) == nil {
			jobs = append(jobs, job)
		}
	}
	return jobs, rows.Err()
}

func (a *App) failAIReplyJob(ctx context.Context, job aiReplyJob, cause error) {
	message := cause.Error()
	if len([]rune(message)) > 500 {
		message = string([]rune(message)[:500])
	}
	if job.Attempts >= aiJobMaxAttempts {
		_, _ = a.db.ExecContext(ctx, `UPDATE ai_reply_jobs SET status=3,last_error=? WHERE id=?`, message, job.ID)
		return
	}
	// Returning the job to the queue can collide with a newer pending job for
	// the same conversation; that newer job supersedes this one.
	if _, err := a.db.ExecContext(ctx, `UPDATE ai_reply_jobs SET status=0,claim_token='',last_error=?,run_after=DATE_ADD(NOW(3),INTERVAL 5 SECOND) WHERE id=?`, message, job.ID); err != nil {
		_, _ = a.db.ExecContext(ctx, `UPDATE ai_reply_jobs SET status=3,last_error=? WHERE id=?`, message, job.ID)
	}
}

func (a *App) processAIReplyJob(ctx context.Context, job aiReplyJob) error {
	binding, err := a.loadAIAgent(ctx, job.AgentUserID)
	if err != nil {
		return err
	}
	if err = a.checkAIQuota(ctx, binding); err != nil {
		return err
	}
	history, lastSenderID, err := a.aiConversationHistory(ctx, job.ConversationID, job.AgentUserID)
	if err != nil {
		return err
	}
	if len(history) == 0 || lastSenderID == job.AgentUserID {
		// Someone already answered, or there is nothing to answer.
		return nil
	}
	model, err := a.aiModelForAgent(ctx, binding)
	if err != nil {
		return err
	}
	provider, err := newAIProvider(model, a.config.Environment == "production")
	if err != nil {
		return err
	}
	first, err := a.aiFirstReply(ctx, job.ConversationID, job.AgentUserID)
	if err != nil {
		return err
	}
	request := aiChatRequest{
		System:   aiSystemPrompt(binding, a.configBool(ctx, "ai.disclose_in_chat", true)),
		Messages: history,
	}
	callCtx, cancel := context.WithTimeout(ctx, time.Duration(model.TimeoutMS+2000)*time.Millisecond)
	defer cancel()
	started := time.Now()
	result, callErr := provider.Chat(callCtx, request)
	latency := int(time.Since(started).Milliseconds())
	a.logAICall(ctx, model.ID, "chat", job.AgentUserID, job.ConversationID, latency, result, callErr)
	if callErr != nil {
		if fallback := strings.TrimSpace(a.configPlain(ctx, "ai.fallback_text", "")); fallback != "" {
			_ = a.sendAIReply(ctx, job, fallback)
		}
		return callErr
	}
	text := sanitizeAIReply(result.Text)
	if text == "" {
		return fmt.Errorf("模型返回内容为空")
	}
	if first && a.configBool(ctx, "ai.disclose_in_chat", true) {
		text = "[AI 助手] " + text
	}
	if err = a.sendAIReply(ctx, job, text); err != nil {
		return err
	}
	_, _ = a.db.ExecContext(ctx, `UPDATE ai_agents SET reply_count=reply_count+1,last_reply_at=NOW(3) WHERE user_id=?`, job.AgentUserID)
	return nil
}

func (a *App) sendAIReply(ctx context.Context, job aiReplyJob, text string) error {
	item, members, err := a.persistMessageContext(markAIGenerated(ctx), job.ConversationID, job.AgentUserID,
		"ai-"+randomToken()[:20], 1, map[string]any{"text": text})
	if err != nil {
		return err
	}
	// The reply reaches clients through the same push the sender's own messages
	// use, so unread counts and conversation ordering need no special casing.
	a.hub.broadcast(members, map[string]any{"command": "MESSAGE_PUSH", "data": item})
	return nil
}

func (a *App) loadAIAgent(ctx context.Context, userID int64) (aiAgentBinding, error) {
	binding := aiAgentBinding{UserID: userID}
	var birthday sql.NullTime
	err := a.db.QueryRowContext(ctx, `SELECT g.model_id,g.persona,g.daily_reply_limit,p.nickname,p.birthday,p.city_name,p.bio
		FROM ai_agents g JOIN user_profiles p ON p.user_id=g.user_id JOIN users u ON u.id=g.user_id
		WHERE g.user_id=? AND g.enabled=1 AND u.is_test=1 AND u.status=1 AND u.deleted_at IS NULL`, userID).
		Scan(&binding.ModelID, &binding.Persona, &binding.DailyReplyLimit, &binding.Nickname, &birthday, &binding.City, &binding.Bio)
	if errors.Is(err, sql.ErrNoRows) {
		return binding, fmt.Errorf("账号未开启 AI 托管")
	}
	if birthday.Valid {
		binding.Age = age(birthday.Time)
	}
	return binding, err
}

func (a *App) aiModelForAgent(ctx context.Context, binding aiAgentBinding) (aiModel, error) {
	if binding.ModelID > 0 {
		model, err := a.loadAIModel(ctx, binding.ModelID)
		if err == nil && model.Status == 1 {
			return model, nil
		}
	}
	return a.defaultAIModel(ctx)
}

// Quotas are counted from the call log so a restart cannot reset them.
func (a *App) checkAIQuota(ctx context.Context, binding aiAgentBinding) error {
	if limit := a.configInt(ctx, "ai.daily_call_limit", 0); limit > 0 {
		var used int
		_ = a.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM ai_call_logs WHERE scene='chat' AND created_at>=CURRENT_DATE()`).Scan(&used)
		if used >= limit {
			return fmt.Errorf("已达到今日全局调用上限 %d", limit)
		}
	}
	limit := binding.DailyReplyLimit
	if limit <= 0 {
		limit = a.configInt(ctx, "ai.agent_daily_reply_limit", 0)
	}
	if limit > 0 {
		var used int
		_ = a.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM ai_call_logs WHERE scene='chat' AND agent_user_id=? AND status=1 AND created_at>=CURRENT_DATE()`, binding.UserID).Scan(&used)
		if used >= limit {
			return fmt.Errorf("账号已达到今日回复上限 %d", limit)
		}
	}
	return nil
}

func (a *App) aiFirstReply(ctx context.Context, conversationID, agentUserID int64) (bool, error) {
	var replied int
	err := a.db.QueryRowContext(ctx, `SELECT EXISTS(SELECT 1 FROM im_messages WHERE conversation_id=? AND sender_id=?)`, conversationID, agentUserID).Scan(&replied)
	return replied == 0, err
}

// Only the tail of the conversation is sent: enough for continuity, bounded in
// cost. Media turns become placeholders because the model gets text only.
func (a *App) aiConversationHistory(ctx context.Context, conversationID, agentUserID int64) ([]aiChatMessage, int64, error) {
	rows, err := a.db.QueryContext(ctx, `SELECT sender_id,message_type,body,recalled_at,admin_removed_at
		FROM im_messages WHERE conversation_id=? ORDER BY seq DESC LIMIT ?`, conversationID, aiHistoryMessages)
	if err != nil {
		return nil, 0, err
	}
	defer rows.Close()
	ordered := []aiChatMessage{}
	lastSenderID := int64(0)
	for rows.Next() {
		var senderID int64
		var messageType int
		var body []byte
		var recalledAt, removedAt sql.NullTime
		if rows.Scan(&senderID, &messageType, &body, &recalledAt, &removedAt) != nil {
			continue
		}
		if lastSenderID == 0 {
			lastSenderID = senderID
		}
		if recalledAt.Valid || removedAt.Valid {
			continue
		}
		text := aiMessageText(messageType, body)
		if text == "" {
			continue
		}
		role := "user"
		if senderID == agentUserID {
			role = "assistant"
		}
		ordered = append([]aiChatMessage{{Role: role, Text: text}}, ordered...)
	}
	return ordered, lastSenderID, rows.Err()
}

func aiMessageText(messageType int, body []byte) string {
	switch messageType {
	case 2:
		return "[图片]"
	case 3:
		return "[语音]"
	}
	var content struct {
		Text string `json:"text"`
	}
	if json.Unmarshal(body, &content) != nil {
		return ""
	}
	return strings.TrimSpace(content.Text)
}

func aiSystemPrompt(binding aiAgentBinding, disclose bool) string {
	persona := strings.TrimSpace(binding.Persona)
	if persona == "" {
		// Falling back to the profile keeps the replies consistent with what the
		// other person can see on the page.
		details := []string{}
		if binding.Nickname != "" {
			details = append(details, "昵称"+binding.Nickname)
		}
		if binding.Age > 0 {
			details = append(details, fmt.Sprintf("%d 岁", binding.Age))
		}
		if binding.City != "" {
			details = append(details, "在"+binding.City)
		}
		if bio := strings.TrimSpace(binding.Bio); bio != "" {
			details = append(details, "个人简介："+bio)
		}
		persona = "你是一位社交软件用户，" + strings.Join(details, "，") + "。"
	}
	rules := []string{
		"用中文口语回复，像真人聊天一样简短自然，不超过 40 个字。",
		"一次只说一件事，不要分点、不要 Markdown、不要堆叠表情。",
		"不索取或提供手机号、微信等联系方式，不谈金钱、转账、投资。",
		"不主动约线下见面，不输出链接。",
		"不知道的事就说不知道，不要编造经历细节。",
	}
	if disclose {
		rules = append(rules, "如果对方询问你是不是机器人或 AI，必须如实说明你是 AI 助手。")
	}
	return persona + "\n\n对话要求：\n- " + strings.Join(rules, "\n- ")
}

// Models drift into assistant formatting; the chat bubble wants one plain line.
func sanitizeAIReply(text string) string {
	text = strings.TrimSpace(text)
	if text == "" {
		return ""
	}
	if index := strings.Index(text, "```"); index >= 0 {
		text = text[:index]
	}
	replacer := strings.NewReplacer("**", "", "__", "", "##", "", "> ", "", "\r", "\n")
	text = replacer.Replace(text)
	lines := []string{}
	for _, line := range strings.Split(text, "\n") {
		line = strings.TrimSpace(strings.TrimLeft(strings.TrimSpace(line), "-*•0123456789. "))
		if line != "" {
			lines = append(lines, line)
		}
	}
	text = strings.TrimSpace(strings.Join(lines, " "))
	if runes := []rune(text); len(runes) > aiReplyMaxRunes {
		text = strings.TrimSpace(string(runes[:aiReplyMaxRunes]))
	}
	return text
}