Update app.py

app.py

@@ -369,8 +369,8 @@ class NovelDatabase:
             return cursor.fetchall()
     
     @staticmethod
-    def get_all_writer_content(session_id: str) -> str:
-        """모든 작가의 수정본 내용을 가져와서 합치기 - 10명 작가"""
+    def get_all_writer_content(session_id: str, test_mode: bool = False) -> str:
+        """모든 작가의 수정본 내용을 가져와서 합치기 - 테스트 모드에서는 1,2만"""
         with NovelDatabase.get_db() as conn:
             cursor = conn.cursor()
             
@@ -378,7 +378,13 @@ class NovelDatabase:
             writer_count = 0
             total_word_count = 0
             
-            for stage_num in WRITER_REVISION_STAGES:
+            # 테스트 모드에서는 writer 1,2만 (stage 5, 8)
+            if test_mode:
+                writer_stages_to_check = [5, 8]  # Writer 1,2의 수정본
+            else:
+                writer_stages_to_check = WRITER_REVISION_STAGES
+            
+            for stage_num in writer_stages_to_check:
                 cursor.execute('''
                     SELECT content, stage_name, word_count FROM stages 
                     WHERE session_id = ? AND stage_number = ?
@@ -399,11 +405,15 @@ class NovelDatabase:
                         logger.info(f"Writer {writer_count} (stage {stage_num}): {word_count} words")
             
             full_content = '\n\n'.join(all_content)
-            logger.info(f"Total: {writer_count} writers, {total_word_count} words")
             
-            # 10명 작가 * 1,450 평균 = 14,500 단어 목표
-            if total_word_count < 12000:  
-                logger.warning(f"Content too short! Only {total_word_count} words instead of ~14,500")
+            if test_mode:
+                logger.info(f"Test mode - Total: {writer_count} writers, {total_word_count} words")
+                if total_word_count < 2800:  # 2 작가 * 1,400 = 2,800 최소
+                    logger.warning(f"Test mode content short! Only {total_word_count} words instead of ~3,000")
+            else:
+                logger.info(f"Total: {writer_count} writers, {total_word_count} words")
+                if total_word_count < 12000:  
+                    logger.warning(f"Content too short! Only {total_word_count} words instead of ~14,500")
             
             return full_content
     
@@ -1026,7 +1036,7 @@ Write a revision reflecting:
 Present the revised final version. Never use page markers.
 You MUST maintain 1,400-1,500 words."""
     
-    def create_test_writer_complete_prompt(self, director_plan: str, writer1_content: str = "", language: str = "English") -> str:
+    def create_test_writer_remaining_prompt(self, director_plan: str, writer1_content: str, language: str) -> str:
         """Test mode - Writer 10 writes remaining novel (chapters 2-10)"""
         if language == "Korean":
             return f"""[테스트 모드] 당신은 나머지 9개 챕터(Chapter 2-10)를 작성하는 특별 작가입니다.
@@ -1142,7 +1152,7 @@ Write Chapters 2-10 now."""
             
             # 작성자들에게는 적절한 토큰 할당
             if role == "writer10" and stage_info and stage_info.get('test_mode'):
-                max_tokens = 30000  # 테스트 모드: 전체 소설 작성 (충분한 토큰)
+                max_tokens = 30000  # 테스트 모드: 9개 챕터 작성 (충분한 토큰)
                 temperature = 0.8   
                 top_p = 0.95
             elif role.startswith("writer"):
@@ -1282,7 +1292,7 @@ Write Chapters 2-10 now."""
             prompts = {
                 "director": "당신은 30페이지 중편 소설을 기획하고 감독하는 문학 감독자입니다. 체계적이고 창의적인 스토리 구조를 만들어냅니다.",
                 "critic": "당신은 날카로운 통찰력을 가진 문학 비평가입니다. 건설적이고 구체적인 피드백을 제공합니다.",
-                "writer10": "[테스트 모드] 당신은 챕터 2-10을 작성하는 특별 작가입니다. 9개 챕터로 구성된 나머지 소설을 작성하세요."
+                "writer10": "[테스트 모드] 당신은 챕터 2-10을 작성하는 특별 작가입니다. 9개 챕터로 구성된 나머지 소설을 작성하세요. 각 챕터는 반드시 1,400-1,500단어로 작성하세요."
             }
             
             # 10명의 작가 프롬프트
@@ -1307,7 +1317,7 @@ Write Chapters 2-10 now."""
             prompts = {
                 "director": "You are a literary director planning and supervising a 30-page novella. You create systematic and creative story structures.",
                 "critic": "You are a literary critic with sharp insights. You provide constructive and specific feedback.",
-                "writer10": "[TEST MODE] You are a special writer creating chapters 2-10. Write the remaining novel organized into 9 chapters."
+                "writer10": "[TEST MODE] You are a special writer creating chapters 2-10. Write the remaining novel organized into 9 chapters. Each chapter MUST be 1,400-1,500 words."
             }
             
             # 10 writer prompts
@@ -1330,7 +1340,7 @@ Write Chapters 2-10 now."""
             return prompts
     
     def get_test_response(self, role: str, language: str) -> str:
-        """Get test response based on role"""
+        """Get test response based on role - updated for writer10 chapters 2-10"""
         if language == "Korean":
             return self.get_korean_test_response(role)
         else:
@@ -1626,111 +1636,16 @@ Seoyeon leaned back in her chair and closed her eyes for a moment. Fatigue penet
                 # Mark stage complete and save to DB
                 stages[stage_idx]["status"] = "complete"
                 
-                # Test mode에서 writer1 추가 저장은 하지 않음 (이미 정상 프로세스로 저장됨)
-                if test_quick_mode and role == "writer10":
-                    logger.info(f"Test mode: Processing remaining chapters 2-10 ({len(stage_content)} chars)")
-                    
-                    # 전체 내용을 9개 챕터로 분할하여 저장 (Chapter 2-10)
-                    # [Chapter 숫자] 패턴으로 분할
-                    chapter_pattern = r'\[(?i)chapter\s+\d+\]'
-                    
-                    # 챕터 마커가 있는지 확인
-                    chapter_markers = re.findall(chapter_pattern, stage_content, re.IGNORECASE)
-                    logger.info(f"Found {len(chapter_markers)} chapter markers")
-                    
-                    if len(chapter_markers) >= 9:
-                        # 챕터 마커로 분할
-                        parts = re.split(chapter_pattern, stage_content, flags=re.IGNORECASE)
-                        chapters = []
-                        
-                        # 첫 번째 부분은 챕터 이전 내용이므로 제외
-                        for i in range(1, len(parts)):
-                            content = parts[i].strip()
-                            if content:
-                                chapters.append(content)
-                    else:
-                        # 챕터 마커가 부족하면 전체를 균등 분할
-                        logger.warning(f"Not enough chapter markers, splitting content evenly")
-                        
-                        # 먼저 단락으로 분할 시도
-                        paragraphs = stage_content.split('\n\n')
-                        paragraphs = [p.strip() for p in paragraphs if p.strip()]
-                        
-                        if len(paragraphs) >= 18:  # 충분한 단락이 있으면
-                            para_per_chapter = len(paragraphs) // 9
-                            chapters = []
-                            for i in range(9):
-                                start_idx = i * para_per_chapter
-                                end_idx = (i + 1) * para_per_chapter if i < 8 else len(paragraphs)
-                                chapter_content = '\n\n'.join(paragraphs[start_idx:end_idx])
-                                chapters.append(chapter_content)
-                        else:
-                            # 단락도 부족하면 단어로 균등 분할
-                            words = stage_content.split()
-                            words_per_chapter = max(len(words) // 9, 100)  # 최소 100단어
-                            chapters = []
-                            for i in range(9):
-                                start_idx = i * words_per_chapter
-                                end_idx = (i + 1) * words_per_chapter if i < 8 else len(words)
-                                chapter_words = words[start_idx:end_idx]
-                                chapters.append(' '.join(chapter_words))
-                    
-                    # 9개 챕터로 제한
-                    chapters = chapters[:9]
-                    
-                    # 부족한 챕터는 마지막 챕터를 분할해서 채우기
-                    while len(chapters) < 9:
-                        if chapters and len(chapters[-1].split()) > 200:
-                            # 마지막 챕터를 반으로 나누기
-                            last_chapter = chapters[-1]
-                            words = last_chapter.split()
-                            mid = len(words) // 2
-                            chapters[-1] = ' '.join(words[:mid])
-                            chapters.append(' '.join(words[mid:]))
-                        else:
-                            # 너무 짧으면 빈 챕터 추가
-                            chapters.append(f"[Chapter {len(chapters)+2} - Content pending]")
-                    
-                    # 각 챕터를 writer 2-10으로 저장
-                    total_saved_words = 0
-                    for i, chapter_content in enumerate(chapters[:9], 2):  # Start from writer 2
-                        writer_stage_num = (i-1) * 3 + 5  # 8, 11, 14, 17, 20, 23, 26, 29, 32
-                        
-                        # 챕터 헤더 추가
-                        full_chapter_content = f"[Chapter {i}]\n\n{chapter_content}"
-                        
-                        NovelDatabase.save_stage(
-                            self.current_session_id,
-                            writer_stage_num,
-                            f"Writer {i}: Revision",
-                            f"writer{i}",
-                            full_chapter_content,
-                            "complete"
-                        )
-                        word_count = len(chapter_content.split())
-                        total_saved_words += word_count
-                        logger.info(f"✅ Test mode: Saved chapter {i} as writer {i} ({word_count} words)")
-                    
-                    logger.info(f"✅ Test mode: Total saved words for chapters 2-10: {total_saved_words}")
-                    
-                    # 원본도 저장 (백업용)
-                    NovelDatabase.save_stage(
-                        self.current_session_id,
-                        stage_idx,
-                        stage_name,
-                        role,
-                        stage_content,
-                        "complete"
-                    )
-                else:
-                    NovelDatabase.save_stage(
-                        self.current_session_id,
-                        stage_idx,
-                        stage_name,
-                        role,
-                        stage_content,
-                        "complete"
-                    )
+                # Test mode에서는 writer1과 writer2만 처리
+                # writer10 관련 특별 처리 제거 - 일반 프로세스와 동일
+                NovelDatabase.save_stage(
+                    self.current_session_id,
+                    stage_idx,
+                    stage_name,
+                    role,
+                    stage_content,
+                    "complete"
+                )
                 
                 # Auto-save notification
                 if role.startswith("writer"):
@@ -1751,7 +1666,7 @@ Seoyeon leaned back in her chair and closed her eyes for a moment. Fatigue penet
                     logger.error(f"Final novel too short! Only {verification['total_words']} words")
             
             # Get complete novel from DB
-            complete_novel = NovelDatabase.get_all_writer_content(self.current_session_id)
+            complete_novel = NovelDatabase.get_all_writer_content(self.current_session_id, test_quick_mode)
             
             # 테스트 모드면 완료 상태 업데이트
             if test_quick_mode and self.current_session_id:
@@ -1762,7 +1677,7 @@ Seoyeon leaned back in her chair and closed her eyes for a moment. Fatigue penet
             
             # Final yield - 화면에는 완료 메시지만 표시
             if test_quick_mode:
-                final_message = f"✅ [TEST MODE] Novel complete! {len(complete_novel.split())} words total. Click Download to save."
+                final_message = f"✅ [TEST MODE] Novel complete! 2 chapters, {len(complete_novel.split())} words total. Click Download to save."
             else:
                 final_message = f"✅ Novel complete! {len(complete_novel.split())} words total. Click Download to save."
             yield final_message, stages
@@ -1850,15 +1765,6 @@ Write a high-quality complete novel."""
             return self.create_director_revision_prompt(
                 stages[0]["content"], stages[1]["content"], language)
         
-        # Test mode - Writer 10 writes remaining chapters
-        elif test_mode and role == "writer10":
-            final_plan = stages[2]["content"]  # Director's final plan
-            # Get Writer 1's content from previous stage
-            writer1_content = ""
-            if len(stages) > 5:  # Writer 1 revision is at index 5
-                writer1_content = stages[5]["content"]
-            return self.create_test_writer_complete_prompt(final_plan, writer1_content, language)
-        
         # Writer stages
         elif role.startswith("writer"):
             writer_num = int(role.replace("writer", ""))
@@ -2039,6 +1945,14 @@ def download_novel(novel_text: str, format: str, language: str, session_id: str
     stages = NovelDatabase.get_stages(session_id)
     logger.info(f"Found {len(stages)} stages in database")
     
+    # 테스트 모드 감지 - Writer 2까지만 있고 전체 stage가 9개면 테스트 모드
+    is_test_mode = False
+    has_writer2 = any(stage['role'] == 'writer2' for stage in stages)
+    has_writer3 = any(stage['role'] == 'writer3' for stage in stages)
+    if has_writer2 and not has_writer3 and len(stages) <= 10:
+        is_test_mode = True
+        logger.info("Test mode detected - will export 2 chapters only")
+    
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
     
     if format == "DOCX" and DOCX_AVAILABLE:
@@ -2047,7 +1961,7 @@ def download_novel(novel_text: str, format: str, language: str, session_id: str
         
         # 제목 페이지
         title_para = doc.add_paragraph()
-        title_run = title_para.add_run('AI Collaborative Novel')
+        title_run = title_para.add_run('AI Collaborative Novel' + (' - Test Mode' if is_test_mode else ''))
         title_run.font.size = Pt(24)
         title_run.font.bold = True
         title_para.alignment = WD_ALIGN_PARAGRAPH.CENTER
@@ -2076,15 +1990,9 @@ def download_novel(novel_text: str, format: str, language: str, session_id: str
         # 각 작가의 수정본만 수집
         writer_contents = []
         
-        # 테스트 모드 체크
-        is_test_mode = False
-        for stage in stages:
-            if stage['role'] == 'writer10' and stage['stage_name'] and 'TEST' in stage['stage_name']:
-                is_test_mode = True
-                logger.info("Test mode detected")
-                break
+        # 테스트 모드에서는 writer 1,2만 검색
+        max_writer = 2 if is_test_mode else 10
         
-        # Writer stages 검색
         for stage in stages:
             logger.debug(f"Checking stage: role={stage['role']}, name={stage['stage_name']}, status={stage['status']}")
             
@@ -2092,6 +2000,8 @@ def download_novel(novel_text: str, format: str, language: str, session_id: str
                 # 작가 번호 추출
                 try:
                     writer_num_from_role = int(stage['role'].replace('writer', ''))
+                    if is_test_mode and writer_num_from_role > max_writer:
+                        continue  # 테스트 모드에서는 writer 3 이상 무시
                     logger.debug(f"Found writer {writer_num_from_role} revision")
                 except:
                     logger.warning(f"Could not extract writer number from role: {stage['role']}")
@@ -2115,42 +2025,9 @@ def download_novel(novel_text: str, format: str, language: str, session_id: str
                     })
                     logger.info(f"Added writer {writer_count}: {word_count} words")
         
-        if not writer_contents:
-            logger.error("No writer content found in database!")
-            
-            # 디버그를 위해 모든 stages 출력
-            for stage in stages:
-                logger.debug(f"Stage {stage['stage_number']}: role={stage['role']}, name={stage['stage_name']}")
-            
-            # 테스트 모드인 경우 writer10 내용 직접 사용
-            if is_test_mode:
-                logger.info("Attempting to use writer10 content directly")
-                for stage in stages:
-                    if stage['role'] == 'writer10':
-                        content = stage['content'] or ""
-                        if content:
-                            doc.add_heading('Complete Novel (Test Mode)', 1)
-                            doc.add_paragraph(f'Total length: {len(content)} characters')
-                            doc.add_page_break()
-                            
-                            paragraphs = content.split('\n\n')
-                            for para_text in paragraphs:
-                                if para_text.strip():
-                                    para = doc.add_paragraph(para_text.strip())
-                                    para.style.font.size = Pt(11)
-                            
-                            # Save
-                            temp_dir = tempfile.gettempdir()
-                            safe_filename = re.sub(r'[^\w\s-]', '', session['user_query'][:30]).strip()
-                            filename = f"Novel_TestMode_{safe_filename}_{timestamp}.docx"
-                            filepath = os.path.join(temp_dir, filename)
-                            doc.save(filepath)
-                            logger.info(f"DOCX saved (test mode): {filepath}")
-                            return filepath
-        
         # 통계 페이지
         doc.add_heading('Novel Statistics', 1)
-        doc.add_paragraph(f'Mode: {"Test Mode" if is_test_mode else "Normal Mode"}')
+        doc.add_paragraph(f'Mode: {"Test Mode (2 chapters)" if is_test_mode else "Full Mode (10 chapters)"}')
         doc.add_paragraph(f'Total Writers: {writer_count}')
         doc.add_paragraph(f'Total Words: {total_words:,}')
         doc.add_paragraph(f'Estimated Pages: ~{total_words/500:.0f}')
@@ -2198,7 +2075,8 @@ def download_novel(novel_text: str, format: str, language: str, session_id: str
         # Save
         temp_dir = tempfile.gettempdir()
         safe_filename = re.sub(r'[^\w\s-]', '', session['user_query'][:30]).strip()
-        filename = f"Novel_Complete_{safe_filename}_{timestamp}.docx"
+        mode_suffix = "_TestMode" if is_test_mode else "_Complete"
+        filename = f"Novel{mode_suffix}_{safe_filename}_{timestamp}.docx"
         filepath = os.path.join(temp_dir, filename)
         doc.save(filepath)
         
@@ -2208,16 +2086,18 @@ def download_novel(novel_text: str, format: str, language: str, session_id: str
         # TXT format
         temp_dir = tempfile.gettempdir()
         safe_filename = re.sub(r'[^\w\s-]', '', session['user_query'][:30]).strip()
-        filename = f"Novel_Complete_{safe_filename}_{timestamp}.txt"
+        mode_suffix = "_TestMode" if is_test_mode else "_Complete"
+        filename = f"Novel{mode_suffix}_{safe_filename}_{timestamp}.txt"
         filepath = os.path.join(temp_dir, filename)
         
         with open(filepath, 'w', encoding='utf-8') as f:
             f.write("="*60 + "\n")
-            f.write("AI COLLABORATIVE NOVEL - COMPLETE VERSION\n")
+            f.write(f"AI COLLABORATIVE NOVEL - {'TEST MODE' if is_test_mode else 'COMPLETE VERSION'}\n")
             f.write("="*60 + "\n")
             f.write(f"Theme: {session['user_query']}\n")
             f.write(f"Language: {session['language']}\n")
             f.write(f"Created: {datetime.now()}\n")
+            f.write(f"Mode: {'Test Mode (2 chapters)' if is_test_mode else 'Full Mode (10 chapters)'}\n")
             f.write("="*60 + "\n\n")
             
             total_words = 0
@@ -2226,6 +2106,14 @@ def download_novel(novel_text: str, format: str, language: str, session_id: str
             # 각 작가의 수정본만 출력
             for stage in stages:
                 if stage['role'] and stage['role'].startswith('writer') and stage['stage_name'] and 'Revision' in stage['stage_name']:
+                    # 테스트 모드에서는 writer 3 이상 무시
+                    try:
+                        writer_num = int(stage['role'].replace('writer', ''))
+                        if is_test_mode and writer_num > 2:
+                            continue
+                    except:
+                        continue
+                        
                     writer_count += 1
                     content = stage['content'] or ""
                     # 페이지 마크 제거
@@ -2244,15 +2132,6 @@ def download_novel(novel_text: str, format: str, language: str, session_id: str
                         f.write(content)
                         f.write("\n\n")
             
-            if writer_count == 0:
-                # 테스트 모드 백업
-                f.write("\n[No writer content found - displaying raw stages]\n\n")
-                for stage in stages:
-                    if stage['content']:
-                        f.write(f"\n--- {stage['stage_name']} ---\n")
-                        f.write(stage['content'])
-                        f.write("\n\n")
-            
             f.write(f"\n{'='*60}\n")
             f.write(f"TOTAL: {writer_count} writers, {total_words:,} words\n")
             f.write(f"{'='*60}\n")
@@ -2372,9 +2251,9 @@ def create_interface():
                     
                     # Test mode checkbox
                     test_mode_check = gr.Checkbox(
-                        label="🧪 Test Mode (Quick: Director + Writer 1 + Writer 10) / 테스트 모드 (빠른 생성)",
+                        label="🧪 Test Mode (Quick: Writer 1 & 2 only) / 테스트 모드 (빠른 생성: 작가 1, 2만)",
                         value=False,
-                        info="Complete Writer 1 process, then Writer 10 creates chapters 2-10 / 작가 1 전체 과정 후 작가 10이 챕터 2-10 작성"
+                        info="Complete Writer 1 & 2 process for 2 chapters only / 작가 1, 2만 작성하여 2개 챕터 생성"
                     )
                     
                     # Web search status indicator