refactor(prompt): 重构提示词构建服务与数据模型

service/prompt/prompt_compose_service.go

@@ -5,171 +5,229 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"prompts-core/dao"
-	"prompts-core/model/entity"
 	"strings"
 	"time"
 
-	"prompts-core/common/util"
-	"prompts-core/consts/public"
-	promptDto "prompts-core/model/dto/prompt"
-	"prompts-core/service/gateway"
-
 	"gitea.com/red-future/common/beans"
 	"gitea.com/red-future/common/utils"
 	"github.com/gogf/gf/v2/container/gvar"
 	"github.com/gogf/gf/v2/frame/g"
+
+	"prompts-core/common/util"
+	"prompts-core/consts/public"
+	"prompts-core/dao"
+	"prompts-core/model/dto"
+	"prompts-core/model/entity"
+	"prompts-core/service/gateway"
 )
 
 // ComposeMessages 核心拼接提示词主流程
-func ComposeMessages(ctx context.Context, req *promptDto.ComposeMessagesReq) (*promptDto.ComposeMessagesRes, error) {
-	var (
-		epicycleId int64
-		taskID     string
-		history    []map[string]any
-		message    map[string]any
-		err        error
-		taskRecord *entity.ComposeTask
-	)
-	// 获取模型信息
+func ComposeMessages(ctx context.Context, req *dto.ComposeMessagesReq) (*dto.ComposeMessagesRes, error) {
 	chatModel, aiModel, err := GetModelMessage(ctx, req)
 	if err != nil {
 		return nil, err
 	}
-	// 根据构建类型进行判断处理
-	switch req.BuildType {
-	//提示词构建
-	case 1:
-		maxRetryTimes := g.Cfg().MustGet(ctx, "promptsRetry.maxRetryTimes", 3).Int()
-		//1. 获取历史会话
-		history, err = GetHistoryMessages(ctx, req.SessionId)
-		if err != nil {
-			g.Log().Errorf(ctx, "获取历史会话失败: %v，将不使用历史会话", err)
-			history = nil // 出错就用空的，不影响主流程
-		}
-		// 重试循环
-		for attempt := 0; attempt <= 0; attempt++ {
-			if attempt > 0 {
-				g.Log().Warningf(ctx, "[重试]第 %d/%d 次调用推理模型", attempt, maxRetryTimes)
-			}
-			// 2. 调用推理模型
-			taskID, err = callInferenceModel(ctx, req, chatModel, aiModel, history)
-			if err != nil {
-				g.Log().Errorf(ctx, "调用推理模型失败(第%d次): %v", attempt+1, err)
-				continue
-			}
-
-			// 3. 保存记录
-			_, err = dao.ComposeTask.Insert(ctx, &entity.ComposeTask{
-				TaskId:         taskID,
-				ModelName:      req.ModelName,
-				SkillName:      req.SkillName,
-				RequestPayload: util.MustMarshal(req),
-				Status:         public.ComposeStatusPending,
-			})
-			if err != nil {
-				g.Log().Errorf(ctx, "保存任务记录失败(第%d次): %v", attempt+1, err)
-				continue
-			}
-
-			// 4. 等待结果
-			taskRecord, err = waitForResult(ctx, taskID)
-			if err != nil {
-				g.Log().Errorf(ctx, "等待结果失败(第%d次): %v", attempt+1, err)
-				continue
-			}
-			// 校验结果
-			message = parsePromptBuild(taskRecord, chatModel)
-			if message != nil && util.IsMessageValid(message) {
-				break
-			}
-			g.Log().Warningf(ctx, "[重试] 推理结果不合法(第%d次)，准备重新请求", attempt+1)
-			message = nil
-		}
-		if message == nil {
-			return nil, errors.New("推理模型调用失败，请稍后再试")
-		}
-		//5.创建会话记录
-		epicycleId, err = dao.ComposeSession.Insert(ctx, &entity.ComposeSession{
-			SessionId:      req.SessionId,
-			RequestContent: message,
-		})
-	//节点构建
-	case 2:
-		//1. 调用推理模型
-		taskID, err = callInferenceModel(ctx, req, chatModel, aiModel, nil)
-		if err != nil {
-			return nil, err
-		}
-		//2. 保存相关记录
-		_, err = dao.ComposeTask.Insert(ctx, &entity.ComposeTask{
-			TaskId:         taskID,
-			ModelName:      req.ModelName,
-			SkillName:      req.SkillName,
-			RequestPayload: util.MustMarshal(req),
-			Status:         public.ComposeStatusPending,
-		})
-		//5. 等待结果
-		taskRecord, err := waitForResult(ctx, taskID)
-		if err != nil {
-			return nil, err
-		}
-		message = parseNodeBuild(taskRecord)
-	default:
-		epicycleId, err = dao.ComposeSession.Insert(ctx, &entity.ComposeSession{
-			SessionId: req.SessionId,
-			Remark:    req.Cause,
-		})
-		return &promptDto.ComposeMessagesRes{
-			EpicycleId: epicycleId,
-		}, nil
+	if err = validateUserForm(ctx, req, aiModel); err != nil {
+		return nil, err
 	}
-	return &promptDto.ComposeMessagesRes{
+	switch req.BuildType {
+	case public.BuildTypePrompt:
+		return handlePromptBuild(ctx, req, chatModel, aiModel) // 提示词构建
+	case public.BuildTypeNode:
+		return handleNodeBuild(ctx, req, chatModel, aiModel) // 节点构建
+	default:
+		return handleDefaultCase(ctx, req)
+	}
+}
+
+// validateUserForm 校验用户表单
+func validateUserForm(ctx context.Context, req *dto.ComposeMessagesReq, model *entity.AsynchModel) error {
+	if len(req.UserForm) == 0 {
+		return nil
+	}
+	isValid, exceedTokens, err := util.CheckUserFormWithinWindow(req.UserForm, model.TokenConfig)
+	if err != nil {
+		return fmt.Errorf("校验用户表单失败: %w", err)
+	}
+
+	if !isValid {
+		availableWindow := util.GetAvailableWindow(model.TokenConfig)
+		return fmt.Errorf("UserForm 内容超出窗口大小: 超出 %d tokens，可用窗口 %d tokens，请精简后重试",
+			exceedTokens, availableWindow)
+	}
+
+	return nil
+}
+
+// handlePromptBuild 处理提示词构建（BuildType=1）
+func handlePromptBuild(ctx context.Context, req *dto.ComposeMessagesReq, chatModel, aiModel *entity.AsynchModel) (*dto.ComposeMessagesRes, error) {
+	maxRetryTimes := g.Cfg().MustGet(ctx, "promptsRetry.maxRetryTimes", 3).Int()
+	history, err := GetHistoryMessages(ctx, req.SessionId)
+	if err != nil {
+		g.Log().Errorf(ctx, "获取历史会话失败: %v，将不使用历史会话", err)
+		history = nil
+	}
+
+	var message *dto.MultiRoundResult
+	var taskRecord *entity.ComposeTask
+	for attempt := 0; attempt <= 0; attempt++ {
+		if attempt > 0 {
+			g.Log().Warningf(ctx, "[重试]第 %d/%d 次调用推理模型", attempt, maxRetryTimes)
+		}
+
+		taskID, err := callInferenceModel(ctx, req, chatModel, aiModel, history)
+		if err != nil {
+			g.Log().Errorf(ctx, "调用推理模型失败(第%d次): %v", attempt+1, err)
+			continue
+		}
+
+		if err = saveComposeTask(ctx, taskID, req); err != nil {
+			g.Log().Errorf(ctx, "保存任务记录失败(第%d次): %v", attempt+1, err)
+			continue
+		}
+
+		taskRecord, err = waitForResult(ctx, taskID)
+		if err != nil {
+			g.Log().Errorf(ctx, "等待结果失败(第%d次): %v", attempt+1, err)
+			continue
+		}
+
+		message = parsePromptBuild(taskRecord, chatModel)
+		if message != nil {
+			break
+		}
+
+		g.Log().Warningf(ctx, "[重试] 推理结果不合法(第%d次)，准备重新请求", attempt+1)
+	}
+
+	if message == nil {
+		return nil, errors.New("推理模型调用失败，请稍后再试")
+	}
+	epicycleId, err := dao.ComposeSession.Insert(ctx, &entity.ComposeSession{
+		SessionId:      req.SessionId,
+		RequestContent: message,
+	})
+	if err != nil {
+		g.Log().Errorf(ctx, "创建会话记录失败: %v", err)
+	}
+	return &dto.ComposeMessagesRes{
 		Messages:   message,
 		EpicycleId: epicycleId,
 	}, nil
 }
 
+// handleNodeBuild 处理节点构建（BuildType=2）
+func handleNodeBuild(ctx context.Context, req *dto.ComposeMessagesReq, chatModel, aiModel *entity.AsynchModel) (*dto.ComposeMessagesRes, error) {
+	taskID, err := callInferenceModel(ctx, req, chatModel, aiModel, nil)
+	if err != nil {
+		return nil, fmt.Errorf("调用推理模型失败: %w", err)
+	}
+
+	if err := saveComposeTask(ctx, taskID, req); err != nil {
+		return nil, fmt.Errorf("保存任务记录失败: %w", err)
+	}
+
+	taskRecord, err := waitForResult(ctx, taskID)
+	if err != nil {
+		return nil, fmt.Errorf("等待结果失败: %w", err)
+	}
+
+	message := parseNodeBuild(taskRecord)
+
+	return &dto.ComposeMessagesRes{
+		Messages:   message,
+		EpicycleId: 0,
+	}, nil
+}
+
+// handleDefaultCase 处理默认情况
+func handleDefaultCase(ctx context.Context, req *dto.ComposeMessagesReq) (*dto.ComposeMessagesRes, error) {
+	epicycleId, err := dao.ComposeSession.Insert(ctx, &entity.ComposeSession{
+		SessionId: req.SessionId,
+		Remark:    req.Cause,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("创建会话记录失败: %w", err)
+	}
+
+	return &dto.ComposeMessagesRes{
+		EpicycleId: epicycleId,
+	}, nil
+}
+
+// saveComposeTask 保存组合任务
+func saveComposeTask(ctx context.Context, taskID string, req *dto.ComposeMessagesReq) error {
+	_, err := dao.ComposeTask.Insert(ctx, &entity.ComposeTask{
+		TaskId:         taskID,
+		ModelName:      req.ModelName,
+		SkillName:      req.SkillName,
+		RequestPayload: util.MustMarshal(req),
+		Status:         public.ComposeStatusPending,
+	})
+	return err
+}
+
 // GetModelMessage 获取模型信息
-func GetModelMessage(ctx context.Context, req *promptDto.ComposeMessagesReq) (*entity.AsynchModel, *entity.AsynchModel, error) {
+func GetModelMessage(ctx context.Context, req *dto.ComposeMessagesReq) (*entity.AsynchModel, *entity.AsynchModel, error) {
 	userInfo, err := utils.GetUserInfo(ctx)
 	if err != nil {
-		return nil, nil, err
+		return nil, nil, fmt.Errorf("获取用户信息失败: %w", err)
 	}
-	// 1. 获取当前用户的会话模型
-	chatModel, err := dao.Model.Get(ctx, &entity.AsynchModel{
-		SQLBaseDO:   beans.SQLBaseDO{Creator: userInfo.UserName},
-		IsChatModel: 1,
-	})
+
+	chatModel, err := getChatModel(ctx, userInfo.UserName)
 	if err != nil {
 		return nil, nil, err
 	}
-	if chatModel == nil {
-		return nil, nil, errors.New("当前没有对话模型，请添加")
-	}
-	// 2. 获取要构建的模型信息
-	aiModel, err := dao.Model.Get(ctx, &entity.AsynchModel{
-		SQLBaseDO: beans.SQLBaseDO{Creator: userInfo.UserName},
-		ModelName: req.ModelName,
-	})
+
+	aiModel, err := getAIModel(ctx, userInfo.UserName, req.ModelName)
 	if err != nil {
 		return nil, nil, err
 	}
-	if aiModel == nil {
-		return nil, nil, fmt.Errorf("需要构建的模型 %s 不存在", req.ModelName)
-	}
+
 	return chatModel, aiModel, nil
 }
 
+// getChatModel 获取聊天模型
+func getChatModel(ctx context.Context, userName string) (*entity.AsynchModel, error) {
+	chatModel, err := dao.Model.Get(ctx, &entity.AsynchModel{
+		SQLBaseDO:   beans.SQLBaseDO{Creator: userName},
+		IsChatModel: new(1),
+	})
+	if err != nil {
+		return nil, fmt.Errorf("查询聊天模型失败: %w", err)
+	}
+
+	if chatModel == nil {
+		return nil, errors.New("当前没有对话模型，请添加")
+	}
+
+	return chatModel, nil
+}
+
+// getAIModel 获取AI模型
+func getAIModel(ctx context.Context, userName, modelName string) (*entity.AsynchModel, error) {
+	aiModel, err := dao.Model.Get(ctx, &entity.AsynchModel{
+		SQLBaseDO: beans.SQLBaseDO{Creator: userName},
+		ModelName: modelName,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("查询AI模型失败: %w", err)
+	}
+
+	if aiModel == nil {
+		return nil, fmt.Errorf("需要构建的模型 %s 不存在", modelName)
+	}
+
+	return aiModel, nil
+}
+
 // callInferenceModel 调用推理模型
-func callInferenceModel(ctx context.Context, req *promptDto.ComposeMessagesReq, chatModel *entity.AsynchModel, model *entity.AsynchModel, history []map[string]any) (string, error) {
-	// 构建推理模型请求
+func callInferenceModel(ctx context.Context, req *dto.ComposeMessagesReq, chatModel *entity.AsynchModel, model *entity.AsynchModel, history []map[string]any) (string, error) {
 	taskReq, err := buildInferenceRequest(ctx, req, chatModel, model, history)
 	if err != nil {
 		return "", fmt.Errorf("构建推理请求失败: %w", err)
 	}
 
-	// 创建网关任务
 	taskID, err := gateway.CreateGatewayTask(ctx, taskReq)
 	if err != nil {
 		return "", fmt.Errorf("创建网关任务失败: %w", err)
@@ -186,96 +244,131 @@ func callInferenceModel(ctx context.Context, req *promptDto.ComposeMessagesReq,
 func waitForResult(ctx context.Context, taskID string) (*entity.ComposeTask, error) {
 	timeout := time.Duration(g.Cfg().MustGet(ctx, "task.waitTimeoutSeconds", 300).Int()) * time.Second
 	pollInterval := time.Duration(g.Cfg().MustGet(ctx, "task.pollIntervalMillis", 500).Int()) * time.Millisecond
+
 	deadline := time.Now().Add(timeout)
+	ticker := time.NewTicker(pollInterval)
+	defer ticker.Stop()
 
 	for {
-		// ===================== 修复点 1：检查上下文是否取消 =====================
-		select {
-		case <-ctx.Done():
-			// 请求已被取消，直接返回，不继续查库
-			return nil, ctx.Err()
-		default:
-		}
-
-		// 1. 查数据库
 		record, err := dao.ComposeTask.Get(ctx, &entity.ComposeTask{
 			TaskId: taskID,
 		})
 		if err != nil {
-			// ===================== 修复点 2：如果是上下文取消，直接返回 =====================
 			if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
 				return nil, err
 			}
-			return nil, err
+			return nil, fmt.Errorf("查询任务失败: %w", err)
 		}
+
 		if record != nil {
-			switch record.Status {
-			case public.ComposeStatusSuccess:
-				return record, nil
-			case public.ComposeStatusFailed:
-				if strings.TrimSpace(record.ErrorMessage) == "" {
-					return nil, fmt.Errorf("任务失败(taskId=%s)", taskID)
-				}
-				return nil, fmt.Errorf("任务失败(taskId=%s): %s", taskID, record.ErrorMessage)
+			if completed, result := checkTaskCompletion(record); completed {
+				return result, nil
 			}
 		}
 
-		// 2. 查网关状态
-		state, err := gateway.QueryGatewayTaskState(ctx, taskID)
-		if err != nil {
-			// 网关不可达不终止，继续轮询
-			g.Log().Warningf(ctx, "[waitForResult] 查询网关失败 taskId=%s err=%v", taskID, err)
-		} else {
-			switch state {
-			case 2: // 网关成功
-				// 网关已成功，主动更新数据库
-				if record != nil {
-					_, err = dao.ComposeTask.Update(ctx, &entity.ComposeTask{
-						TaskId: taskID,
-						Status: public.ComposeStatusSuccess,
-					})
-					if err != nil {
-						g.Log().Warningf(ctx, "[waitForResult] 更新任务状态失败 taskId=%s err=%v", taskID, err)
-					}
-				}
-			case 3: // 网关失败
-				if record != nil {
-					_, err = dao.ComposeTask.Update(ctx, &entity.ComposeTask{
-						TaskId:       taskID,
-						Status:       public.ComposeStatusFailed,
-						ErrorMessage: "model-gateway 任务执行失败",
-					})
-					if err != nil {
-						g.Log().Warningf(ctx, "[waitForResult] 更新任务状态失败 taskId=%s err=%v", taskID, err)
-					}
-				}
-				return nil, fmt.Errorf("model-gateway 任务执行失败(taskId=%s)", taskID)
-			}
+		if err = syncGatewayTaskState(ctx, taskID, record); err != nil {
+			g.Log().Warningf(ctx, "[waitForResult] 同步网关状态失败 taskId=%s err=%v", taskID, err)
 		}
 
-		// 3. 超时检查
 		if time.Now().After(deadline) {
 			return nil, fmt.Errorf("等待任务回调超时(taskId=%s)", taskID)
 		}
 
-		// ===================== 修复点3：sleep 也要监听 ctx 取消 =====================
 		select {
 		case <-ctx.Done():
 			return nil, ctx.Err()
-		case <-time.After(pollInterval):
+		case <-ticker.C:
 		}
 	}
 }
 
+// checkTaskCompletion 检查任务是否完成
+func checkTaskCompletion(record *entity.ComposeTask) (bool, *entity.ComposeTask) {
+	if record == nil {
+		return false, nil
+	}
+	switch record.Status {
+	case public.ComposeStatusSuccess:
+		return true, record
+	case public.ComposeStatusFailed:
+		errMsg := strings.TrimSpace(record.ErrorMessage)
+		if errMsg == "" {
+			return true, nil
+		}
+		return true, nil
+	default:
+		return false, nil
+	}
+}
+
+// syncGatewayTaskState 同步网关任务状态
+func syncGatewayTaskState(ctx context.Context, taskID string, record *entity.ComposeTask) error {
+	state, err := gateway.QueryGatewayTaskState(ctx, taskID)
+	if err != nil {
+		return fmt.Errorf("查询网关状态失败: %w", err)
+	}
+	switch state {
+	case 2:
+		return updateTaskStatus(ctx, taskID, public.ComposeStatusSuccess, "")
+	case 3:
+		updateTaskStatus(ctx, taskID, public.ComposeStatusFailed, "model-gateway 任务执行失败")
+		return fmt.Errorf("model-gateway 任务执行失败(taskId=%s)", taskID)
+	}
+	return nil
+}
+
+// updateTaskStatus 更新任务状态
+func updateTaskStatus(ctx context.Context, taskID string, status string, errorMsg string) error {
+	task := &entity.ComposeTask{
+		TaskId: taskID,
+		Status: status,
+	}
+	if errorMsg != "" {
+		task.ErrorMessage = errorMsg
+	}
+
+	_, err := dao.ComposeTask.Update(ctx, task)
+	return err
+}
+
 // parsePromptBuild 解析提示词构建结果（BuildType == 1）
-func parsePromptBuild(taskRecord *entity.ComposeTask, model *entity.AsynchModel) map[string]any {
+func parsePromptBuild(taskRecord *entity.ComposeTask, model *entity.AsynchModel) *dto.MultiRoundResult {
 	if taskRecord == nil {
 		return nil
 	}
+	mapped := parseTaskMessages(taskRecord.Messages)
+	if mapped == nil {
+		return createDefaultResult(nil)
+	}
 
-	// 1. 解析 Messages
+	contentField := getContentField(model)
+	contentStr, ok := mapped[contentField].(string)
+	if !ok || contentStr == "" {
+		return createDefaultResult(mapped)
+	}
+
+	if roundsArray := tryParseAsArray(contentStr); roundsArray != nil {
+		return &dto.MultiRoundResult{
+			TotalRounds: len(roundsArray),
+			Rounds:      roundsArray,
+		}
+	}
+
+	if singleRound := tryParseAsObject(contentStr); singleRound != nil {
+		return &dto.MultiRoundResult{
+			TotalRounds: 1,
+			Rounds:      []any{singleRound},
+		}
+	}
+
+	return createDefaultResult(map[string]any{"content": contentStr})
+}
+
+// parseTaskMessages 解析任务消息
+func parseTaskMessages(messages any) map[string]any {
 	var mapped map[string]any
-	switch v := taskRecord.Messages.(type) {
+
+	switch v := messages.(type) {
 	case *gvar.Var:
 		if v != nil {
 			json.Unmarshal([]byte(v.String()), &mapped)
@@ -289,115 +382,137 @@ func parsePromptBuild(taskRecord *entity.ComposeTask, model *entity.AsynchModel)
 		json.Unmarshal(b, &mapped)
 	}
 
-	// 2. 解析模型 ResponseMapping 获取 content 字段名
-	contentField := "content" // 默认值
-	if model != nil {
-		var respMapping map[string]string
-		switch v := model.ResponseMapping.(type) {
-		case *gvar.Var:
-			if v != nil {
-				json.Unmarshal([]byte(v.String()), &respMapping)
-			}
-		case string:
-			json.Unmarshal([]byte(v), &respMapping)
-		case map[string]interface{}:
-			respMapping = make(map[string]string)
-			for k, val := range v {
-				if s, ok := val.(string); ok {
-					respMapping[k] = s
-				}
-			}
-		}
-		// 从映射中找到 content 对应的字段名
-		for k, v := range respMapping {
-			if strings.Contains(v, "content") {
-				contentField = k
-				break
-			}
-		}
-	}
-
-	// 3. 提取 content 的值
-	contentStr, ok := mapped[contentField].(string)
-	if !ok || contentStr == "" {
-		return mapped
-	}
-
-	// 4. 解析 content 内的 JSON
-	var innerData map[string]any
-	json.Unmarshal([]byte(contentStr), &innerData)
-
-	return innerData
+	return mapped
 }
 
-// parseNodeBuild 解析节点构建结果（BuildType == 2）
-func parseNodeBuild(taskRecord *entity.ComposeTask) map[string]any {
-	if taskRecord == nil {
+// tryParseAsArray 尝试将字符串解析为数组
+func tryParseAsArray(contentStr string) []any {
+	var roundsArray []any
+	if err := json.Unmarshal([]byte(contentStr), &roundsArray); err != nil {
 		return nil
 	}
-	var result map[string]any
-	switch v := taskRecord.Messages.(type) {
+	return roundsArray
+}
+
+// tryParseAsObject 尝试将字符串解析为对象
+func tryParseAsObject(contentStr string) any {
+	var singleRound any
+	if err := json.Unmarshal([]byte(contentStr), &singleRound); err != nil {
+		return nil
+	}
+	return singleRound
+}
+
+// createDefaultResult 创建默认结果
+func createDefaultResult(data any) *dto.MultiRoundResult {
+	if data == nil {
+		data = make(map[string]any)
+	}
+	return &dto.MultiRoundResult{
+		TotalRounds: 1,
+		Rounds:      []any{data},
+	}
+}
+
+// getContentField 从模型 ResponseMapping 中获取 content 字段名
+func getContentField(model *entity.AsynchModel) string {
+	if model == nil {
+		return "content"
+	}
+
+	respMapping := parseResponseMapping(model.ResponseMapping)
+	for k, v := range respMapping {
+		if strings.Contains(v, "content") {
+			return k
+		}
+	}
+
+	return "content"
+}
+
+// parseResponseMapping 解析响应映射
+func parseResponseMapping(mapping any) map[string]string {
+	result := make(map[string]string)
+
+	switch v := mapping.(type) {
 	case *gvar.Var:
 		if v != nil {
 			json.Unmarshal([]byte(v.String()), &result)
 		}
 	case string:
 		json.Unmarshal([]byte(v), &result)
-	case map[string]any:
-		result = v
-	default:
-		b, _ := json.Marshal(v)
-		json.Unmarshal(b, &result)
+	case map[string]interface{}:
+		for k, val := range v {
+			if s, ok := val.(string); ok {
+				result[k] = s
+			}
+		}
 	}
+
 	return result
 }
 
+// parseNodeBuild 解析节点构建结果（BuildType == 2）
+func parseNodeBuild(taskRecord *entity.ComposeTask) *dto.MultiRoundResult {
+	if taskRecord == nil {
+		return nil
+	}
+
+	result := parseTaskMessages(taskRecord.Messages)
+	if result == nil {
+		result = make(map[string]any)
+	}
+
+	return &dto.MultiRoundResult{
+		TotalRounds: 1,
+		Rounds:      []any{result},
+	}
+}
+
 // Callback 回调处理
-func Callback(ctx context.Context, req *promptDto.CallbackReq) error {
+func Callback(ctx context.Context, req *dto.CallbackReq) error {
 	g.Log().Infof(ctx, "[Callback][RECV] taskId=%s state=%d ossFile=%s fileType=%s textLen=%d",
 		req.TaskId, req.State, req.OssFile, req.FileType, len(req.Text))
 
-	// ============ 先查任务是否存在 ============
 	task, err := dao.ComposeTask.Get(ctx, &entity.ComposeTask{
 		TaskId: req.TaskId,
 	})
 	if err != nil {
-		return err
+		return fmt.Errorf("查询任务失败: %w", err)
 	}
 	if task == nil {
 		return fmt.Errorf("任务不存在: %s", req.TaskId)
 	}
-	// ============ 根据状态区分处理 ============
+
 	if req.State == 3 {
-		// 失败：直接更新状态
-		_, err = dao.ComposeTask.Update(ctx, &entity.ComposeTask{
-			TaskId:       req.TaskId,
-			Status:       public.ComposeStatusFailed,
-			ErrorMessage: req.ErrorMsg,
-		})
-		return err
-	}
-	// ======================================
-	// 成功：解析模型输出
-	result, err := util.ParseOutput(req.Text)
-	if err != nil {
-		_, updateErr := dao.ComposeTask.Update(ctx, &entity.ComposeTask{
-			TaskId:       req.TaskId,
-			Status:       public.ComposeStatusFailed,
-			ErrorMessage: req.ErrorMsg,
-		})
-		if updateErr != nil {
-			g.Log().Warningf(ctx, "[Callback] 更新失败状态出错 taskId=%s err=%v", req.TaskId, updateErr)
-		}
-		return err
+		return handleCallbackFailure(ctx, req.TaskId, req.ErrorMsg)
+	}
+
+	return handleCallbackSuccess(ctx, req)
+}
+
+// handleCallbackFailure 处理回调失败
+func handleCallbackFailure(ctx context.Context, taskID, errorMsg string) error {
+	_, err := dao.ComposeTask.Update(ctx, &entity.ComposeTask{
+		TaskId:       taskID,
+		Status:       public.ComposeStatusFailed,
+		ErrorMessage: errorMsg,
+	})
+	return err
+}
+
+// handleCallbackSuccess 处理回调成功
+func handleCallbackSuccess(ctx context.Context, req *dto.CallbackReq) error {
+	result, err := util.ParseOutput(req.Text)
+	if err != nil {
+		handleCallbackFailure(ctx, req.TaskId, req.ErrorMsg)
+		return fmt.Errorf("解析模型输出失败: %w", err)
 	}
 
-	// ============ result 可能为 nil ============
 	var messages any
 	if result != nil {
 		messages = result
 	}
-	// =======================================
 
 	_, err = dao.ComposeTask.Update(ctx, &entity.ComposeTask{
 		TaskId:   req.TaskId,
@@ -407,34 +522,43 @@ func Callback(ctx context.Context, req *promptDto.CallbackReq) error {
 	if err != nil {
 		g.Log().Errorf(ctx, "[Callback] 更新任务失败 taskId=%s err=%v", req.TaskId, err)
 	}
+
 	return err
 }
 
 // GetComposeTask 查询任务结果
-func GetComposeTask(ctx context.Context, taskID string) (*promptDto.GetComposeTaskRes, error) {
+func GetComposeTask(ctx context.Context, taskID string) (*dto.GetComposeTaskRes, error) {
 	record, err := dao.ComposeTask.Get(ctx, &entity.ComposeTask{
 		TaskId: taskID,
 	})
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("查询任务失败: %w", err)
 	}
 	if record == nil {
 		return nil, fmt.Errorf("未找到任务(taskId=%s)", taskID)
 	}
 
-	// 如果 Messages 是字符串，反序列化为 JSON 数组
-	messages := record.Messages
-	if str, ok := messages.(string); ok && str != "" {
-		var parsed any
-		if err := json.Unmarshal([]byte(str), &parsed); err == nil {
-			messages = parsed
-		}
-	}
+	messages := parseMessagesForResponse(record.Messages)
 
-	return &promptDto.GetComposeTaskRes{
+	return &dto.GetComposeTaskRes{
 		TaskId:       record.TaskId,
 		Status:       record.Status,
 		ErrorMessage: record.ErrorMessage,
 		Messages:     messages,
 	}, nil
 }
+
+// parseMessagesForResponse 解析用于响应的消息
+func parseMessagesForResponse(messages any) any {
+	str, ok := messages.(string)
+	if !ok || str == "" {
+		return messages
+	}
+
+	var parsed any
+	if err := json.Unmarshal([]byte(str), &parsed); err == nil {
+		return parsed
+	}
+
+	return messages
+}