@@ -10,21 +10,25 @@ import { useAppStore } from '../../../src/stores/appStore';
1010import { imageGenerationService } from '../../../src/services/imageGenerationService' ;
1111import { localDreamGeneratorService } from '../../../src/services/localDreamGenerator' ;
1212import { activeModelService } from '../../../src/services/activeModelService' ;
13+ import { llmService } from '../../../src/services/llm' ;
1314import {
1415 resetStores ,
1516 flushPromises ,
1617 getAppState ,
1718 getChatState ,
1819 setupWithConversation ,
1920} from '../../utils/testHelpers' ;
20- import { createONNXImageModel , createGeneratedImage } from '../../utils/factories' ;
21+ import { createONNXImageModel , createGeneratedImage , createMessage } from '../../utils/factories' ;
22+ import { Message } from '../../../src/types' ;
2123
2224// Mock the services
2325jest . mock ( '../../../src/services/localDreamGenerator' ) ;
2426jest . mock ( '../../../src/services/activeModelService' ) ;
27+ jest . mock ( '../../../src/services/llm' ) ;
2528
2629const mockLocalDreamService = localDreamGeneratorService as jest . Mocked < typeof localDreamGeneratorService > ;
2730const mockActiveModelService = activeModelService as jest . Mocked < typeof activeModelService > ;
31+ const mockLlmService = llmService as jest . Mocked < typeof llmService > ;
2832
2933describe ( 'Image Generation Flow Integration' , ( ) => {
3034 beforeEach ( async ( ) => {
@@ -55,6 +59,11 @@ describe('Image Generation Flow Integration', () => {
5559 } ) ;
5660 mockActiveModelService . loadImageModel . mockResolvedValue ( ) ;
5761
62+ // Default LLM service mocks (for prompt enhancement)
63+ mockLlmService . isModelLoaded . mockReturnValue ( false ) ;
64+ mockLlmService . isCurrentlyGenerating . mockReturnValue ( false ) ;
65+ mockLlmService . stopGeneration . mockResolvedValue ( ) ;
66+
5867 // Reset imageGenerationService state by canceling any in-progress generation
5968 await imageGenerationService . cancelGeneration ( ) . catch ( ( ) => { } ) ;
6069 } ) ;
@@ -512,4 +521,203 @@ describe('Image Generation Flow Integration', () => {
512521 expect ( message ?. generationMeta ?. resolution ) . toBe ( '512x512' ) ;
513522 } ) ;
514523 } ) ;
524+
525+ describe ( 'Prompt Enhancement with Conversation Context' , ( ) => {
526+ const setupEnhancement = ( ) => {
527+ const imageModel = setupImageModelState ( ) ;
528+ mockActiveModelService . getActiveModels . mockReturnValue ( {
529+ text : { model : null , isLoaded : false , isLoading : false } ,
530+ image : { model : imageModel , isLoaded : true , isLoading : false } ,
531+ } ) ;
532+
533+ // Enable enhancement and set up LLM as available
534+ useAppStore . setState ( {
535+ settings : {
536+ ...useAppStore . getState ( ) . settings ,
537+ enhanceImagePrompts : true ,
538+ } ,
539+ } ) ;
540+ mockLlmService . isModelLoaded . mockReturnValue ( true ) ;
541+ mockLlmService . isCurrentlyGenerating . mockReturnValue ( false ) ;
542+ mockLlmService . generateResponse . mockResolvedValue ( 'A beautifully enhanced prompt' ) ;
543+
544+ return imageModel ;
545+ } ;
546+
547+ it ( 'should pass conversation history to enhancement when conversationId provided' , async ( ) => {
548+ setupEnhancement ( ) ;
549+
550+ // Set up a conversation with prior messages
551+ const messages : Message [ ] = [
552+ createMessage ( { role : 'user' , content : 'Draw me a cat' } ) ,
553+ createMessage ( { role : 'assistant' , content : 'Here is a cat image' } ) ,
554+ createMessage ( { role : 'user' , content : 'Make it darker' } ) ,
555+ ] ;
556+ const conversationId = setupWithConversation ( { messages } ) ;
557+
558+ await imageGenerationService . generateImage ( {
559+ prompt : 'Make it darker' ,
560+ conversationId,
561+ } ) ;
562+
563+ // Verify generateResponse was called with conversation context
564+ expect ( mockLlmService . generateResponse ) . toHaveBeenCalled ( ) ;
565+ const callArgs = mockLlmService . generateResponse . mock . calls [ 0 ] ;
566+ const enhancementMessages = callArgs [ 0 ] as Message [ ] ;
567+
568+ // Should have: system + context messages + user enhance prompt
569+ // system (1) + conversation messages (3) + user enhance (1) = 5
570+ expect ( enhancementMessages . length ) . toBe ( 5 ) ;
571+ expect ( enhancementMessages [ 0 ] . role ) . toBe ( 'system' ) ;
572+ expect ( enhancementMessages [ 0 ] . content ) . toContain ( 'conversation history' ) ;
573+ expect ( enhancementMessages [ 1 ] . content ) . toBe ( 'Draw me a cat' ) ;
574+ expect ( enhancementMessages [ 2 ] . content ) . toBe ( 'Here is a cat image' ) ;
575+ expect ( enhancementMessages [ 3 ] . content ) . toBe ( 'Make it darker' ) ;
576+ expect ( enhancementMessages [ 4 ] . role ) . toBe ( 'user' ) ;
577+ expect ( enhancementMessages [ 4 ] . content ) . toBe ( 'Make it darker' ) ;
578+ } ) ;
579+
580+ it ( 'should not include conversation context when no conversationId' , async ( ) => {
581+ setupEnhancement ( ) ;
582+
583+ await imageGenerationService . generateImage ( {
584+ prompt : 'A sunset' ,
585+ } ) ;
586+
587+ expect ( mockLlmService . generateResponse ) . toHaveBeenCalled ( ) ;
588+ const callArgs = mockLlmService . generateResponse . mock . calls [ 0 ] ;
589+ const enhancementMessages = callArgs [ 0 ] as Message [ ] ;
590+
591+ // Should have: system + user enhance prompt only (no context)
592+ expect ( enhancementMessages . length ) . toBe ( 2 ) ;
593+ expect ( enhancementMessages [ 0 ] . role ) . toBe ( 'system' ) ;
594+ expect ( enhancementMessages [ 0 ] . content ) . not . toContain ( 'conversation history' ) ;
595+ expect ( enhancementMessages [ 1 ] . role ) . toBe ( 'user' ) ;
596+ expect ( enhancementMessages [ 1 ] . content ) . toBe ( 'A sunset' ) ;
597+ } ) ;
598+
599+ it ( 'should truncate long messages in conversation context' , async ( ) => {
600+ setupEnhancement ( ) ;
601+
602+ const longContent = 'x' . repeat ( 1000 ) ;
603+ const messages : Message [ ] = [
604+ createMessage ( { role : 'user' , content : longContent } ) ,
605+ ] ;
606+ const conversationId = setupWithConversation ( { messages } ) ;
607+
608+ await imageGenerationService . generateImage ( {
609+ prompt : 'Enhance this' ,
610+ conversationId,
611+ } ) ;
612+
613+ const callArgs = mockLlmService . generateResponse . mock . calls [ 0 ] ;
614+ const enhancementMessages = callArgs [ 0 ] as Message [ ] ;
615+
616+ // The context message should be truncated to 500 chars
617+ const contextMsg = enhancementMessages . find ( m => m . id . startsWith ( 'ctx-' ) ) ;
618+ expect ( contextMsg ) . toBeDefined ( ) ;
619+ expect ( contextMsg ! . content . length ) . toBe ( 500 ) ;
620+ } ) ;
621+
622+ it ( 'should limit conversation context to last 10 messages' , async ( ) => {
623+ setupEnhancement ( ) ;
624+
625+ // Create 15 messages
626+ const messages : Message [ ] = [ ] ;
627+ for ( let i = 0 ; i < 15 ; i ++ ) {
628+ messages . push ( createMessage ( {
629+ role : i % 2 === 0 ? 'user' : 'assistant' ,
630+ content : `Message ${ i + 1 } ,
631+ } ) ) ;
632+ }
633+ const conversationId = setupWithConversation ( { messages } ) ;
634+
635+ await imageGenerationService . generateImage ( {
636+ prompt : 'Generate image' ,
637+ conversationId,
638+ } ) ;
639+
640+ const callArgs = mockLlmService . generateResponse . mock . calls [ 0 ] ;
641+ const enhancementMessages = callArgs [ 0 ] as Message [ ] ;
642+
643+ // system (1) + last 10 context messages + user enhance (1) = 12
644+ expect ( enhancementMessages . length ) . toBe ( 12 ) ;
645+ // First context message should be message 6 (index 5), not message 1
646+ const firstContextMsg = enhancementMessages [ 1 ] ;
647+ expect ( firstContextMsg . content ) . toBe ( 'Message 6' ) ;
648+ } ) ;
649+
650+ it ( 'should skip system messages from conversation context' , async ( ) => {
651+ setupEnhancement ( ) ;
652+
653+ const messages : Message [ ] = [
654+ createMessage ( { role : 'user' , content : 'Hello' } ) ,
655+ createMessage ( { role : 'system' , content : 'Model loaded successfully' } ) ,
656+ createMessage ( { role : 'assistant' , content : 'Hi there' } ) ,
657+ ] ;
658+ const conversationId = setupWithConversation ( { messages } ) ;
659+
660+ await imageGenerationService . generateImage ( {
661+ prompt : 'Draw something' ,
662+ conversationId,
663+ } ) ;
664+
665+ const callArgs = mockLlmService . generateResponse . mock . calls [ 0 ] ;
666+ const enhancementMessages = callArgs [ 0 ] as Message [ ] ;
667+
668+ // system (1) + 2 context (user + assistant, system skipped) + user enhance (1) = 4
669+ expect ( enhancementMessages . length ) . toBe ( 4 ) ;
670+ const contextMessages = enhancementMessages . filter ( m => m . id . startsWith ( 'ctx-' ) ) ;
671+ expect ( contextMessages ) . toHaveLength ( 2 ) ;
672+ expect ( contextMessages . every ( m => m . role !== 'system' ) ) . toBe ( true ) ;
673+ } ) ;
674+
675+ it ( 'should use original prompt when enhancement is disabled' , async ( ) => {
676+ setupImageModelState ( ) ;
677+ mockActiveModelService . getActiveModels . mockReturnValue ( {
678+ text : { model : null , isLoaded : false , isLoading : false } ,
679+ image : { model : setupImageModelState ( ) , isLoaded : true , isLoading : false } ,
680+ } ) ;
681+
682+ // Enhancement disabled (default)
683+ useAppStore . setState ( {
684+ settings : {
685+ ...useAppStore . getState ( ) . settings ,
686+ enhanceImagePrompts : false ,
687+ } ,
688+ } ) ;
689+
690+ const messages : Message [ ] = [
691+ createMessage ( { role : 'user' , content : 'Draw a cat' } ) ,
692+ ] ;
693+ const conversationId = setupWithConversation ( { messages } ) ;
694+
695+ await imageGenerationService . generateImage ( {
696+ prompt : 'Make it blue' ,
697+ conversationId,
698+ } ) ;
699+
700+ // LLM should not be called for enhancement
701+ expect ( mockLlmService . generateResponse ) . not . toHaveBeenCalled ( ) ;
702+ } ) ;
703+
704+ it ( 'should handle empty conversation gracefully' , async ( ) => {
705+ setupEnhancement ( ) ;
706+
707+ const conversationId = setupWithConversation ( { messages : [ ] } ) ;
708+
709+ await imageGenerationService . generateImage ( {
710+ prompt : 'A landscape' ,
711+ conversationId,
712+ } ) ;
713+
714+ const callArgs = mockLlmService . generateResponse . mock . calls [ 0 ] ;
715+ const enhancementMessages = callArgs [ 0 ] as Message [ ] ;
716+
717+ // system + user enhance only (no context from empty conversation)
718+ expect ( enhancementMessages . length ) . toBe ( 2 ) ;
719+ expect ( enhancementMessages [ 0 ] . role ) . toBe ( 'system' ) ;
720+ expect ( enhancementMessages [ 0 ] . content ) . not . toContain ( 'conversation history' ) ;
721+ } ) ;
722+ } ) ;
515723} ) ;
0 commit comments