完善训练模块

2026-09-04 17:02:05 +08:00
parent 8dfcb6547c
commit bd7fe63b70
5 changed files with 310 additions and 33 deletions
@@ -213,44 +213,68 @@ func generateOfficialAccountImages(taskID uint, workflow *officialAccountWorkflo
        client := &http.Client{Timeout: 240 * time.Second}
        refs := make([]officialAccountImageRef, 0, len(prompts))
        for idx, item := range prompts {
-                imageSize := officialAccountImageSlotSize(idx)
-                requestSize := officialAccountImageRequestSize(imageSize)
-                started := time.Now()
-                imageURL, reqErr := requestOfficialAccountImage(client, route, item.Prompt, requestSize)
-                latencyMs := int(time.Since(started) / time.Millisecond)
-                ai.LogCall(ai.LogEntry{
-                        UserID:       user.ID,
-                        Capability:   ai.CapabilityImageGen,
-                        Provider:     route.Provider,
-                        RouteID:      route.RouteID,
-                        Model:        route.Model,
-                        Success:      reqErr == nil,
-                        ErrorMessage: errorMessage(reqErr),
-                        LatencyMs:    latencyMs,
-                })
-                if reqErr != nil {
-                        return nil, nil, web.NewLLMError(fmt.Sprintf("图片生成失败：%s", reqErr.Error()))
+                ref, appErr := generateOfficialAccountImageRefWithRoute(client, route, taskID, item, idx, user)
+                if appErr != nil {
+                        return nil, nil, appErr
                }
-                if strings.HasPrefix(strings.TrimSpace(imageURL), "data:image/") {
-                        imageURL, reqErr = persistOfficialAccountDataURL(taskID, item.Key, imageURL)
-                        if reqErr != nil {
-                                return nil, nil, web.NewLLMError(fmt.Sprintf("图片保存失败：%s", reqErr.Error()))
-                        }
-                }
-                refs = append(refs, officialAccountImageRef{
-                        Key:           item.Key,
-                        Section:       item.Section,
-                        Prompt:        item.Prompt,
-                        URL:           imageURL,
-                        ImageSize:     imageSize,
-                        RouteID:       route.RouteID,
-                        RouteModel:    route.Model,
-                        RouteProvider: route.Provider,
-                })
+                refs = append(refs, ref)
        }
        return route, refs, nil
 }

+func generateOfficialAccountImageRef(taskID uint, prompt officialAccountImagePrompt, index int, user *model.User) (*config.RouteConfig, officialAccountImageRef, *web.AppError) {
+        route, err := config.GetRoute("image_gen")
+        if err != nil {
+                return nil, officialAccountImageRef{}, web.NewLLMNotConfigured("未找到可用的图片生成路由，请先检查 ai_config.json 中的 image_gen 配置")
+        }
+        if route.Category != "image" || strings.TrimSpace(route.FullURL) == "" || strings.TrimSpace(route.Model) == "" {
+                return nil, officialAccountImageRef{}, web.NewLLMNotConfigured("图片生成路由配置不完整")
+        }
+        client := &http.Client{Timeout: 240 * time.Second}
+        ref, appErr := generateOfficialAccountImageRefWithRoute(client, route, taskID, prompt, index, user)
+        if appErr != nil {
+                return nil, officialAccountImageRef{}, appErr
+        }
+        return route, ref, nil
+}
+
+func generateOfficialAccountImageRefWithRoute(client *http.Client, route *config.RouteConfig, taskID uint, prompt officialAccountImagePrompt, index int, user *model.User) (officialAccountImageRef, *web.AppError) {
+        imageSize := officialAccountImageSlotSize(index)
+        requestSize := officialAccountImageRequestSize(imageSize)
+        started := time.Now()
+        imageURL, reqErr := requestOfficialAccountImage(client, route, prompt.Prompt, requestSize)
+        latencyMs := int(time.Since(started) / time.Millisecond)
+        ai.LogCall(ai.LogEntry{
+                UserID:       user.ID,
+                Capability:   ai.CapabilityImageGen,
+                Provider:     route.Provider,
+                RouteID:      route.RouteID,
+                Model:        route.Model,
+                Success:      reqErr == nil,
+                ErrorMessage: errorMessage(reqErr),
+                LatencyMs:    latencyMs,
+        })
+        if reqErr != nil {
+                return officialAccountImageRef{}, web.NewLLMError(fmt.Sprintf("图片生成失败：%s", reqErr.Error()))
+        }
+        if strings.HasPrefix(strings.TrimSpace(imageURL), "data:image/") {
+                imageURL, reqErr = persistOfficialAccountDataURL(taskID, prompt.Key, imageURL)
+                if reqErr != nil {
+                        return officialAccountImageRef{}, web.NewLLMError(fmt.Sprintf("图片保存失败：%s", reqErr.Error()))
+                }
+        }
+        return officialAccountImageRef{
+                Key:           prompt.Key,
+                Section:       prompt.Section,
+                Prompt:        prompt.Prompt,
+                URL:           imageURL,
+                ImageSize:     imageSize,
+                RouteID:       route.RouteID,
+                RouteModel:    route.Model,
+                RouteProvider: route.Provider,
+        }, nil
+}
+
 func requestOfficialAccountImage(client *http.Client, route *config.RouteConfig, prompt, size string) (string, error) {
        payload := map[string]any{
                "model":  route.Model,
@@ -516,6 +516,133 @@ func ExecuteOfficialAccountWorkflowStep(c *gin.Context) {
 	})
 }

+func RegenerateOfficialAccountImage(c *gin.Context) {
+        user := middleware.CurrentUser(c)
+        task, ok := loadAccessibleWorkerTask(c, user)
+        if !ok {
+                return
+        }
+        if task.SpecialistKey != officialAccountSpecialistKey {
+                web.Fail(c, web.NewNotFoundError("当前事项不属于公众号助手"))
+                return
+        }
+
+        imageKey := strings.TrimSpace(c.Param("imageKey"))
+        if imageKey == "" {
+                web.Fail(c, web.NewBadRequest("缺少图片标识"))
+                return
+        }
+
+        workflow := parseOfficialAccountWorkflow(task.ContextJSON)
+        article, err := ensureOfficialAccountArticle(task, workflow)
+        if err != nil {
+                web.Fail(c, web.NewBadRequest("读取公众号文章状态失败"))
+                return
+        }
+        syncWorkflowFromOfficialAccountArticle(&workflow, article)
+
+        prompt, index, ok := findOfficialAccountImagePrompt(workflow, imageKey)
+        if !ok {
+                web.Fail(c, web.NewBadRequest("没有找到要重生成的图片"))
+                return
+        }
+
+        now := time.Now()
+        route, ref, appErr := generateOfficialAccountImageRef(article.TaskID, prompt, index, user)
+        if appErr != nil {
+                web.Fail(c, appErr)
+                return
+        }
+
+        workflow.Shared.ImageRefs = upsertOfficialAccountImageRef(workflow.Shared.ImageRefs, workflow.Shared.ImagePrompts, ref)
+        refreshOfficialAccountDerivedAssets(&workflow)
+
+        output := map[string]any{
+                "image_refs":        workflow.Shared.ImageRefs,
+                "count":             len(workflow.Shared.ImageRefs),
+                "regenerated_key":   ref.Key,
+                "regenerated_ref":   ref,
+                "route_id":          route.RouteID,
+                "route_model":       route.Model,
+                "route_provider":    route.Provider,
+        }
+        logs := []string{
+                fmt.Sprintf("%s 已重新生成图片「%s」。", now.Format("15:04:05"), firstNonEmpty(ref.Section, ref.Key)),
+                fmt.Sprintf("%s 预览稿与分页稿中的图片地址已同步刷新。", now.Format("15:04:05")),
+        }
+        setOfficialAccountStepCompleted(&workflow, officialAccountStepKeyImages, output, logs, now)
+
+        task.Status = workerTaskStatusDraft
+        task.Summary = buildOfficialAccountTaskSummaryFromWorkflow(workflow)
+        task.Title = buildOfficialAccountTaskTitleFromWorkflow(workflow)
+        task.ContextJSON = marshalOfficialAccountWorkflow(workflow)
+        task.CurrentResult = fmt.Sprintf("已重新生成图片「%s」。", firstNonEmpty(ref.Section, ref.Key))
+        task.LastTriggeredAt = &now
+        syncOfficialAccountArticleFromWorkflow(article, workflow)
+
+        inputJSON, _ := json.Marshal(gin.H{
+                "image_key": imageKey,
+                "section":   prompt.Section,
+                "prompt":    prompt.Prompt,
+        })
+        outputPayload := cloneOutputMap(output)
+        outputPayload["summary"] = task.CurrentResult
+        outputPayload["keyword"] = workflow.Form.Keyword
+        outputJSON, _ := json.Marshal(outputPayload)
+        logsJSON, _ := json.Marshal(logs)
+        run := model.WorkerRun{
+                TaskID:        task.ID,
+                SpecialistKey: task.SpecialistKey,
+                ActionKey:     "image_regeneration",
+                ActionTitle:   "单张图片重生成",
+                ActionType:    "workflow_step",
+                Status:        "done",
+                InputJSON:     string(inputJSON),
+                OutputJSON:    string(outputJSON),
+                LogsJSON:      string(logsJSON),
+                StartedAt:     now,
+                FinishedAt:    &now,
+        }
+        if err := store.DB.Create(&run).Error; err != nil {
+                web.Fail(c, web.NewBadRequest("保存单张图片重生成记录失败"))
+                return
+        }
+        task.CurrentRunID = &run.ID
+
+        artifact := &model.WorkerArtifact{
+                TaskID:         task.ID,
+                SpecialistKey:  task.SpecialistKey,
+                Title:          fmt.Sprintf("单张重生成图片 - %s", firstNonEmpty(ref.Section, ref.Key)),
+                ArtifactType:   "images",
+                Status:         workerArtifactStatusDraft,
+                ContentText:    buildOfficialAccountImageArtifact([]officialAccountImageRef{ref}),
+                ContentJSON:    string(outputJSON),
+                SourceRefsJSON: "[]",
+                CreatedByRunID: &run.ID,
+        }
+        if err := store.DB.Create(artifact).Error; err != nil {
+                web.Fail(c, web.NewBadRequest("保存单张图片结果失败"))
+                return
+        }
+
+        if err := store.DB.Save(article).Error; err != nil {
+                web.Fail(c, web.NewBadRequest("保存公众号文章状态失败"))
+                return
+        }
+        if err := store.DB.Save(&task).Error; err != nil {
+                web.Fail(c, web.NewBadRequest("更新公众号任务失败"))
+                return
+        }
+
+        web.OK(c, gin.H{
+                "task":      task,
+                "workflow":  workflow,
+                "run":       run,
+                "artifact":  artifact,
+                "image_ref": ref,
+        })
+}
+
 func respondOfficialAccountWorkflow(c *gin.Context, task model.WorkerTask) {
        workflow := parseOfficialAccountWorkflow(task.ContextJSON)
        article, _ := ensureOfficialAccountArticle(task, workflow)
@@ -549,6 +676,87 @@ func respondOfficialAccountWorkflow(c *gin.Context, task model.WorkerTask) {
 	})
 }

+func findOfficialAccountImagePrompt(workflow officialAccountWorkflowState, imageKey string) (officialAccountImagePrompt, int, bool) {
+        prompts := workflow.Shared.ImagePrompts
+        if len(prompts) == 0 {
+                prompts = buildOfficialAccountImagePrompts(&workflow)
+                workflow.Shared.ImagePrompts = prompts
+        }
+        for idx, item := range prompts {
+                if strings.EqualFold(strings.TrimSpace(item.Key), imageKey) {
+                        return item, idx, true
+                }
+        }
+        for idx, item := range workflow.Shared.ImageRefs {
+                if strings.EqualFold(strings.TrimSpace(item.Key), imageKey) {
+                        return officialAccountImagePrompt{
+                                Key:     item.Key,
+                                Section: item.Section,
+                                Prompt:  item.Prompt,
+                        }, idx, true
+                }
+        }
+        return officialAccountImagePrompt{}, -1, false
+}
+
+func upsertOfficialAccountImageRef(existing []officialAccountImageRef, prompts []officialAccountImagePrompt, next officialAccountImageRef) []officialAccountImageRef {
+        byKey := make(map[string]officialAccountImageRef, len(existing)+1)
+        for _, item := range existing {
+                key := strings.TrimSpace(item.Key)
+                if key == "" {
+                        continue
+                }
+                byKey[key] = item
+        }
+        byKey[strings.TrimSpace(next.Key)] = next
+
+        result := make([]officialAccountImageRef, 0, len(byKey))
+        used := map[string]struct{}{}
+        for _, prompt := range prompts {
+                key := strings.TrimSpace(prompt.Key)
+                if key == "" {
+                        continue
+                }
+                item, ok := byKey[key]
+                if !ok {
+                        continue
+                }
+                result = append(result, item)
+                used[key] = struct{}{}
+        }
+        for _, item := range existing {
+                key := strings.TrimSpace(item.Key)
+                if key == "" {
+                        continue
+                }
+                if _, ok := used[key]; ok {
+                        continue
+                }
+                if updated, exists := byKey[key]; exists {
+                        result = append(result, updated)
+                        used[key] = struct{}{}
+                }
+        }
+        if _, ok := used[strings.TrimSpace(next.Key)]; !ok {
+                result = append(result, next)
+        }
+        return result
+}
+
+func refreshOfficialAccountDerivedAssets(workflow *officialAccountWorkflowState) {
+        if workflow == nil {
+                return
+        }
+        if strings.TrimSpace(workflow.Shared.PreviewMarkdown) != "" || strings.TrimSpace(workflow.Shared.PreviewHTML) != "" || strings.TrimSpace(workflow.Shared.PageMarkdown) != "" || len(workflow.Shared.Pages) > 0 {
+                workflow.Shared.PreviewMarkdown = buildOfficialAccountPreviewMarkdown(workflow)
+                workflow.Shared.PreviewHTML = renderOfficialAccountHTML(workflow.Shared.PreviewMarkdown)
+        }
+        if strings.TrimSpace(workflow.Shared.PageMarkdown) != "" || len(workflow.Shared.Pages) > 0 {
+                workflow.Shared.Pages = buildOfficialAccountPages(workflow)
+                workflow.Shared.PageMarkdown = buildOfficialAccountPageMarkdown(workflow.Shared.Pages)
+        }
+}
+
 func newOfficialAccountWorkflow(req officialAccountTaskCreateReq) officialAccountWorkflowState {
 	workflow := officialAccountWorkflowState{
 		WorkflowType: "wechat_official_account",
@@ -51,6 +51,7 @@ func RegisterRoutes(r *gin.Engine, cfg *config.Config) {
 	r.GET("/api/official-account/tasks/:id/workflow", middleware.Auth(cfg), GetOfficialAccountWorkflow)
 	r.PUT("/api/official-account/tasks/:id", middleware.Auth(cfg), UpdateOfficialAccountTask)
 	r.POST("/api/official-account/tasks/:id/steps/:stepKey", middleware.Auth(cfg), ExecuteOfficialAccountWorkflowStep)
+   r.POST("/api/official-account/tasks/:id/images/:imageKey/regenerate", middleware.Auth(cfg), RegenerateOfficialAccountImage)
        r.GET("/api/official-account/tasks/:id/export", middleware.Auth(cfg), ExportOfficialAccountDocument)
        r.GET("/api/official-account/generated-images/:filename", middleware.Auth(cfg), ServeOfficialAccountGeneratedImage)
 	r.GET("/api/connectors", middleware.Auth(cfg), ListConnectors)