Skip to content

Y
yaoai-video
Commit b29011ef authored by yaoke.yk's avatar yaoke.yk
Browse files
Options

fix(storyboard): harden prompt persistence and invalidation

parent a19ab7b9
Hide whitespace changes
Inline Side-by-side
Showing with 123 additions and 22 deletions
yaoai-comic-studio/yaoai-bootstrap/src/test/java/com/yaoai/pipeline/service/impl/StoryboardPipelineServiceImplTest.java
View file @ b29011ef
package com.yaoai.pipeline.service.impl;
import com.baomidou.mybatisplus.core.conditions.AbstractWrapper;
import com.baomidou.mybatisplus.core.conditions.update.UpdateWrapper;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.yaoai.ai.core.model.ChatRequest;
import com.yaoai.ai.core.service.LlmService;
import com.yaoai.billing.service.BillingService;
import com.yaoai.common.context.UserContext;
import com.yaoai.common.exception.BizException;
import com.yaoai.domain.entity.Episode;
import com.yaoai.domain.entity.Scene;
import com.yaoai.domain.entity.Storyboard;
......@@ -18,16 +21,22 @@ import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.ArgumentCaptor;
import org.mockito.Mock;
import org.mockito.InOrder;
import org.mockito.junit.jupiter.MockitoExtension;
import org.springframework.transaction.annotation.Transactional;
import java.util.List;
import java.util.Map;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertSame;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.ArgumentMatchers.argThat;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.inOrder;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
......@@ -78,14 +87,15 @@ class StoryboardPipelineServiceImplTest {
episode.setSummary("summary");
when(episodeMapper.selectById(11L)).thenReturn(episode);
when(characterMapper.findByProject(22L, 33L)).thenReturn(List.of());
when(sceneMapper.findByProject(22L, 33L)).thenReturn(List.of(scene("cafe"), scene("street")));
when(sceneMapper.findByProject(22L, 33L)).thenReturn(List.of(scene("cafe", null), scene("street", "outdoor")));
when(storyboardMapper.insert(any(Storyboard.class))).thenAnswer(invocation -> 1);
when(llmService.chat(any(ChatRequest.class))).thenReturn("""
[{
"sequence_num": 1,
"scene_number": "001",
"short_description": "shot",
"detailed_description": "detail",
"characters": "@hero, @cafe, @friend",
"characters": "@hero, @street, @cafe, @friend",
"dialogues": "",
"camera_direction": "push",
"composition_guide": "medium",
......@@ -98,8 +108,14 @@ class StoryboardPipelineServiceImplTest {
List<Storyboard> storyboards = service.generateStoryboards(11L, 22L, 33L);
assertEquals(1, storyboards.size());
assertEquals("@cafe", storyboards.get(0).getSceneRef());
assertEquals("@street", storyboards.get(0).getSceneRef());
assertEquals("@hero,@friend", storyboards.get(0).getCharacters());
ArgumentCaptor<ChatRequest> requestCaptor = ArgumentCaptor.forClass(ChatRequest.class);
verify(llmService).chat(requestCaptor.capture());
String systemMessage = requestCaptor.getValue().getMessages().get(0).getContent();
assertTrue(systemMessage.contains("- @cafe ():"));
assertTrue(!systemMessage.contains("- @cafe (outdoor):"));
}
@Test
......@@ -119,6 +135,7 @@ class StoryboardPipelineServiceImplTest {
storyboard.setCompositionGuide("medium");
storyboard.setDurationSeconds(5);
when(storyboardMapper.selectById(7L)).thenReturn(storyboard);
when(storyboardMapper.updateById(any(Storyboard.class))).thenReturn(1);
when(episodeMapper.selectById(11L)).thenReturn(new Episode());
when(characterMapper.findByProject(22L, 33L)).thenReturn(List.of());
when(sceneMapper.findByProject(22L, 33L)).thenReturn(List.of());
......@@ -134,6 +151,10 @@ class StoryboardPipelineServiceImplTest {
assertEquals(7L, persisted.getId());
assertEquals("persisted prompt", persisted.getVideoPrompt());
InOrder inOrder = inOrder(billingService, storyboardMapper);
inOrder.verify(billingService).charge(any());
inOrder.verify(storyboardMapper).updateById(any(Storyboard.class));
ArgumentCaptor<ChatRequest> requestCaptor = ArgumentCaptor.forClass(ChatRequest.class);
verify(llmService).chat(requestCaptor.capture());
String userMessage = requestCaptor.getValue().getMessages().get(1).getContent();
......@@ -156,9 +177,10 @@ class StoryboardPipelineServiceImplTest {
storyboard.setCompositionGuide("wide");
storyboard.setDurationSeconds(6);
when(storyboardMapper.selectById(8L)).thenReturn(storyboard);
when(storyboardMapper.updateById(any(Storyboard.class))).thenReturn(1);
when(episodeMapper.selectById(11L)).thenReturn(new Episode());
when(characterMapper.findByProject(22L, 33L)).thenReturn(List.of());
when(sceneMapper.findByProject(22L, 33L)).thenReturn(List.of(scene("cafe")));
when(sceneMapper.findByProject(22L, 33L)).thenReturn(List.of(scene("cafe", null)));
when(llmService.chat(any(ChatRequest.class))).thenReturn("legacy prompt");
service.generatePrompt(8L, 33L);
......@@ -168,6 +190,8 @@ class StoryboardPipelineServiceImplTest {
String userMessage = requestCaptor.getValue().getMessages().get(1).getContent();
assertTrue(userMessage.contains("Character refs: @hero"));
assertTrue(userMessage.contains("Scene refs: @cafe"));
assertTrue(userMessage.contains("- @cafe ():"));
assertTrue(!userMessage.contains("- @cafe (outdoor):"));
}
@Test
......@@ -180,6 +204,8 @@ class StoryboardPipelineServiceImplTest {
Storyboard refreshed = new Storyboard();
refreshed.setId(9L);
refreshed.setTenantId(33L);
when(storyboardMapper.updateById(any(Storyboard.class))).thenReturn(1);
when(storyboardMapper.update(org.mockito.ArgumentMatchers.isNull(), any())).thenReturn(1);
when(storyboardMapper.selectById(9L)).thenReturn(existing, refreshed);
Storyboard patch = new Storyboard();
......@@ -195,10 +221,16 @@ class StoryboardPipelineServiceImplTest {
assertEquals(9L, updated.getId());
assertEquals(33L, updated.getTenantId());
ArgumentCaptor<com.baomidou.mybatisplus.core.conditions.Wrapper<Storyboard>> wrapperCaptor =
ArgumentCaptor.forClass(com.baomidou.mybatisplus.core.conditions.Wrapper.class);
ArgumentCaptor<UpdateWrapper<Storyboard>> wrapperCaptor = ArgumentCaptor.forClass(UpdateWrapper.class);
verify(storyboardMapper).update(org.mockito.ArgumentMatchers.isNull(), wrapperCaptor.capture());
assertTrue(wrapperCaptor.getValue().getSqlSet().contains("video_prompt"));
UpdateWrapper<Storyboard> wrapper = wrapperCaptor.getValue();
assertTrue(wrapper.getSqlSet().contains("video_prompt"));
assertTrue(wrapper.getSqlSegment().contains("id"));
assertTrue(wrapper.getSqlSegment().contains("tenant_id"));
Map<String, Object> paramPairs = ((AbstractWrapper<Storyboard, ?, ?>) wrapper).getParamNameValuePairs();
assertTrue(paramPairs.containsValue(9L));
assertTrue(paramPairs.containsValue(33L));
}
@Test
......@@ -207,6 +239,7 @@ class StoryboardPipelineServiceImplTest {
existing.setId(9L);
existing.setTenantId(33L);
existing.setShortDescription("same");
when(storyboardMapper.updateById(any(Storyboard.class))).thenReturn(1);
when(storyboardMapper.selectById(9L)).thenReturn(existing, existing);
Storyboard patch = new Storyboard();
......@@ -220,6 +253,45 @@ class StoryboardPipelineServiceImplTest {
}
@Test
void update_shouldBeTransactionalAndFailWhenPromptInvalidationFails() throws Exception {
Storyboard existing = new Storyboard();
existing.setId(9L);
existing.setTenantId(33L);
existing.setShortDescription("old");
when(storyboardMapper.selectById(9L)).thenReturn(existing);
when(storyboardMapper.updateById(any(Storyboard.class))).thenReturn(1);
when(storyboardMapper.update(org.mockito.ArgumentMatchers.isNull(), any())).thenReturn(0);
Storyboard patch = new Storyboard();
patch.setShortDescription("new");
assertThrows(BizException.class, () -> service.update(9L, 33L, patch));
assertTrue(StoryboardPipelineServiceImpl.class
.getMethod("update", Long.class, Long.class, Storyboard.class)
.isAnnotationPresent(Transactional.class));
}
@Test
void generatePrompt_shouldNotPersistVideoPromptWhenChargeFails() throws Exception {
Storyboard storyboard = promptStoryboard(7L, 11L, 22L, 33L);
storyboard.setSceneRef("@cafe");
when(storyboardMapper.selectById(7L)).thenReturn(storyboard);
when(episodeMapper.selectById(11L)).thenReturn(new Episode());
when(characterMapper.findByProject(22L, 33L)).thenReturn(List.of());
when(sceneMapper.findByProject(22L, 33L)).thenReturn(List.of());
when(llmService.chat(any(ChatRequest.class))).thenReturn("persisted prompt");
doThrow(new BizException("TEST", "charge failed")).when(billingService).charge(any());
assertThrows(BizException.class, () -> service.generatePrompt(7L, 33L));
verify(storyboardMapper, never()).updateById(org.mockito.ArgumentMatchers.<Storyboard>argThat(storyboardArg ->
"persisted prompt".equals(storyboardArg.getVideoPrompt())
));
assertTrue(StoryboardPipelineServiceImpl.class
.getMethod("generatePrompt", Long.class, Long.class)
.isAnnotationPresent(Transactional.class));
}
@Test
void populateMissingVideoPromptsByEpisode_shouldFillOnlyMissingRows() {
Storyboard missingA = new Storyboard();
missingA.setId(1L);
......@@ -227,6 +299,7 @@ class StoryboardPipelineServiceImplTest {
missingB.setId(2L);
Storyboard full = new Storyboard();
full.setId(3L);
when(storyboardMapper.updateById(any(Storyboard.class))).thenReturn(1);
when(storyboardMapper.findMissingVideoPromptByEpisode(44L, 33L)).thenReturn(List.of(missingA, missingB));
when(storyboardMapper.findByEpisode(44L, 33L)).thenReturn(List.of(missingA, missingB, full));
when(storyboardMapper.selectById(1L)).thenReturn(promptStoryboard(1L, 44L, 22L, 33L));
......@@ -239,12 +312,16 @@ class StoryboardPipelineServiceImplTest {
List<Storyboard> result = service.populateMissingVideoPromptsByEpisode(44L, 33L);
assertEquals(3, result.size());
verify(storyboardMapper, times(2)).updateById(any(Storyboard.class));
verify(storyboardMapper, never()).selectById(3L);
ArgumentCaptor<Storyboard> updateCaptor = ArgumentCaptor.forClass(Storyboard.class);
verify(storyboardMapper, times(2)).updateById(updateCaptor.capture());
assertEquals(List.of(1L, 2L), updateCaptor.getAllValues().stream().map(Storyboard::getId).toList());
}
private static Scene scene(String name) {
private static Scene scene(String name, String sceneType) {
Scene scene = new Scene();
scene.setName(name);
scene.setSceneType(sceneType);
return scene;
}
......
yaoai-comic-studio/yaoai-pipeline/src/main/java/com/yaoai/pipeline/service/impl/StoryboardPipelineServiceImpl.java
View file @ b29011ef
......@@ -23,6 +23,7 @@ import com.yaoai.pipeline.service.StoryboardPipelineService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import java.util.ArrayList;
import java.util.Arrays;
......@@ -193,6 +194,7 @@ public class StoryboardPipelineServiceImpl implements StoryboardPipelineService
}
@Override
@Transactional
public Storyboard update(Long id, Long tenantId, Storyboard patch) {
Storyboard existing = storyboardMapper.selectById(id);
if (existing == null || !tenantId.equals(existing.getTenantId())) {
......@@ -202,7 +204,10 @@ public class StoryboardPipelineServiceImpl implements StoryboardPipelineService
boolean bodyFieldsChanged = hasStoryboardBodyChanges(existing, patch);
patch.setId(id);
patch.setTenantId(tenantId);
storyboardMapper.updateById(patch);
int updatedRows = storyboardMapper.updateById(patch);
if (updatedRows != 1) {
throw new BizException(ErrorCode.INTERNAL_ERROR, "Failed to update storyboard");
}
if (bodyFieldsChanged) {
clearVideoPrompt(id, tenantId);
}
......@@ -226,6 +231,7 @@ public class StoryboardPipelineServiceImpl implements StoryboardPipelineService
}
@Override
@Transactional
public String generatePrompt(Long storyboardId, Long tenantId) {
Storyboard storyboard = storyboardMapper.selectById(storyboardId);
if (storyboard == null || !tenantId.equals(storyboard.getTenantId())) {
......@@ -271,7 +277,7 @@ public class StoryboardPipelineServiceImpl implements StoryboardPipelineService
userMsg.append(String.format(
"- @%s (%s): %s%n",
nullSafe(scene.getName()),
"indoor".equals(scene.getSceneType()) ? "indoor" : "outdoor",
sceneTypeLabel(scene.getSceneType()),
nullSafe(scene.getDescription())
));
}
......@@ -331,11 +337,6 @@ public class StoryboardPipelineServiceImpl implements StoryboardPipelineService
String prompt = llmService.chat(request);
Storyboard promptUpdate = new Storyboard();
promptUpdate.setId(storyboard.getId());
promptUpdate.setVideoPrompt(prompt);
storyboardMapper.updateById(promptUpdate);
billingService.charge(BillingChargeRequest.builder()
.tenantId(tenantId)
.userId(UserContext.require())
......@@ -348,6 +349,14 @@ public class StoryboardPipelineServiceImpl implements StoryboardPipelineService
.unitCount(1)
.refId(String.valueOf(storyboardId))
.build());
Storyboard promptUpdate = new Storyboard();
promptUpdate.setId(storyboard.getId());
promptUpdate.setVideoPrompt(prompt);
int updatedRows = storyboardMapper.updateById(promptUpdate);
if (updatedRows != 1) {
throw new BizException(ErrorCode.INTERNAL_ERROR, "Failed to persist storyboard prompt");
}
return prompt;
}
......@@ -388,7 +397,7 @@ public class StoryboardPipelineServiceImpl implements StoryboardPipelineService
context.append(String.format(
"- @%s (%s): %s%n",
nullSafe(scene.getName()),
"indoor".equals(scene.getSceneType()) ? "indoor" : "outdoor",
sceneTypeLabel(scene.getSceneType()),
nullSafe(scene.getDescription())
));
}
......@@ -412,17 +421,19 @@ public class StoryboardPipelineServiceImpl implements StoryboardPipelineService
private StoryboardRefs splitStoryboardRefs(String rawRefs, Set<String> sceneRefs) {
List<String> refs = parseRefs(rawRefs);
List<String> characterRefs = new ArrayList<>();
List<String> matchedSceneRefs = new ArrayList<>();
String sceneRef = null;
for (String ref : refs) {
if (sceneRefs.contains(ref)) {
matchedSceneRefs.add(ref);
if (sceneRef == null) {
sceneRef = ref;
}
} else {
characterRefs.add(ref);
}
}
return new StoryboardRefs(joinRefs(characterRefs), joinRefs(matchedSceneRefs));
return new StoryboardRefs(joinRefs(characterRefs), sceneRef);
}
private StoryboardRefs promptRefsForCurrentShot(Storyboard storyboard, List<Scene> projectScenes) {
......@@ -459,13 +470,16 @@ public class StoryboardPipelineServiceImpl implements StoryboardPipelineService
}
private void clearVideoPrompt(Long storyboardId, Long tenantId) {
storyboardMapper.update(
int updatedRows = storyboardMapper.update(
null,
new UpdateWrapper<Storyboard>()
.eq("id", storyboardId)
.eq("tenant_id", tenantId)
.set("video_prompt", null)
);
if (updatedRows != 1) {
throw new BizException(ErrorCode.INTERNAL_ERROR, "Failed to clear storyboard prompt");
}
}
private boolean hasStoryboardBodyChanges(Storyboard existing, Storyboard patch) {
......@@ -508,6 +522,16 @@ public class StoryboardPipelineServiceImpl implements StoryboardPipelineService
return value != null ? value : "";
}
private String sceneTypeLabel(String sceneType) {
if ("indoor".equals(sceneType)) {
return "indoor";
}
if ("outdoor".equals(sceneType)) {
return "outdoor";
}
return "";
}
private record StoryboardRefs(String characters, String sceneRef) {
}
}
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment