Update app.py

app.py

@@ -309,6 +309,8 @@ class NovelDatabase:
         """Save stage content with word count"""
         word_count = len(content.split()) if content else 0
         
+        logger.info(f"Saving stage: session={session_id}, stage_num={stage_number}, role={role}, stage_name={stage_name}, words={word_count}")
+        
         with NovelDatabase.get_db() as conn:
             cursor = conn.cursor()
             
@@ -331,7 +333,7 @@ class NovelDatabase:
             ''', (stage_number, stage_number, session_id))
             
             conn.commit()
-            logger.info(f"Saved stage {stage_number} for session {session_id}, content length: {len(content)}, words: {word_count}")
+            logger.info(f"Stage saved successfully")
     
     @staticmethod
     def get_session(session_id: str) -> Optional[Dict]:
@@ -1035,18 +1037,24 @@ You MUST maintain 1,400-1,500 words."""
 **중요 지침:**
 1. 전체 30페이지 분량의 완성된 소설을 작성하세요
 2. 총 14,000-15,000 단어로 작성하세요
-3. 10개의 자연스러운 챕터로 구성하세요 (각 챕터 약 1,400-1,500 단어)
+3. 정확히 10개의 챕터로 구성하세요 (각 챕터 약 1,400-1,500 단어)
 4. 시작부터 결말까지 완전한 이야기를 만드세요
 5. 마스터플랜의 모든 요소를 포함하세요
 
-챕터 구분은 다음과 같이 표시하세요:
+**필수 형식:**
+반드시 아래와 같이 챕터를 명확히 구분하세요:
+
 [Chapter 1]
-내용...
+(1,400-1,500 단어의 내용)
 
 [Chapter 2]
-내용...
+(1,400-1,500 단어의 내용)
 
-완성도 높은 전체 소설을 작성하세요."""
+...이런 식으로 [Chapter 10]까지...
+
+각 챕터는 반드시 [Chapter 숫자] 형식으로 시작해야 합니다.
+챕터 사이에는 빈 줄을 넣어 구분하세요.
+전체 소설을 작성하세요."""
         else:
             return f"""[TEST MODE] You are a special writer creating the complete 30-page novel at once.
 
@@ -1056,18 +1064,24 @@ Director's Masterplan:
 **CRITICAL INSTRUCTIONS:**
 1. Write a complete 30-page novel
 2. Total 14,000-15,000 words
-3. Organize into 10 natural chapters (each chapter ~1,400-1,500 words)
+3. Organize into EXACTLY 10 chapters (each chapter ~1,400-1,500 words)
 4. Create a complete story from beginning to end
 5. Include all elements from the masterplan
 
-Mark chapters as follows:
+**MANDATORY FORMAT:**
+You MUST clearly separate chapters as follows:
+
 [Chapter 1]
-content...
+(1,400-1,500 words of content)
 
 [Chapter 2]
-content...
+(1,400-1,500 words of content)
 
-Write a high-quality complete novel."""
+...continue this way until [Chapter 10]...
+
+Each chapter MUST start with [Chapter number] format.
+Leave blank lines between chapters for separation.
+Write the complete novel now."""
     
     def simulate_streaming(self, text: str, role: str) -> Generator[str, None, None]:
         """Simulate streaming in test mode"""
@@ -1393,10 +1407,13 @@ Write a high-quality complete novel."""
         if role == "writer10":
             full_novel = ""
             for i in range(1, 11):
-                full_novel += f"\n\n[Chapter {i}]\n\n"
-                full_novel += sample_story + f"\n\n이것은 챕터 {i}의 내용입니다. "
-                full_novel += "이야기는 계속 전개되었고, 인물들의 감정은 점점 더 복잡해졌다. " * 100
-                full_novel += "\n\n"
+                full_novel += f"[Chapter {i}]\n\n"
+                # 각 챕터마다 약 300단어씩 5단락
+                for j in range(5):
+                    full_novel += sample_story + f"\n\n그것은 챕터 {i}의 {j+1}번째 단락이었다. "
+                    full_novel += "이야기는 계속 전개되었고, 인물들의 감정은 점점 더 복잡해졌다. " * 20
+                    full_novel += "\n\n"
+                full_novel += "\n\n"  # 챕터 간 구분
             test_responses["writer10"] = full_novel
             
         return test_responses.get(role, "테스트 응답입니다.")
@@ -1471,10 +1488,13 @@ Seoyeon leaned back in her chair and closed her eyes for a moment. Fatigue penet
         if role == "writer10":
             full_novel = ""
             for i in range(1, 11):
-                full_novel += f"\n\n[Chapter {i}]\n\n"
-                full_novel += sample_story + f"\n\nThis was chapter {i} content. "
-                full_novel += "The story continued to unfold, and the characters' emotions grew increasingly complex. " * 100
-                full_novel += "\n\n"
+                full_novel += f"[Chapter {i}]\n\n"
+                # About 300 words per paragraph, 5 paragraphs per chapter
+                for j in range(5):
+                    full_novel += sample_story + f"\n\nThis was paragraph {j+1} of chapter {i}. "
+                    full_novel += "The story continued to unfold, and the characters' emotions grew increasingly complex. " * 20
+                    full_novel += "\n\n"
+                full_novel += "\n\n"  # Chapter separation
             test_responses["writer10"] = full_novel
             
         return test_responses.get(role, "Test response.")
@@ -1495,11 +1515,13 @@ Seoyeon leaned back in her chair and closed her eyes for a moment. Fatigue penet
                     query = session['user_query']
                     language = session['language']
                     resume_from_stage = session['current_stage'] + 1
+                    logger.info(f"Resuming session {session_id} from stage {resume_from_stage}")
             else:
                 self.current_session_id = NovelDatabase.create_session(query, language)
                 resume_from_stage = 0
+                logger.info(f"Created new session: {self.current_session_id}")
             
-            logger.info(f"Processing novel for session {self.current_session_id}, starting from stage {resume_from_stage}")
+            logger.info(f"Processing novel for session {self.current_session_id}, starting from stage {resume_from_stage}, test_mode={test_quick_mode}")
             
             # Initialize conversation
             conversation_history = [{
@@ -1597,46 +1619,100 @@ Seoyeon leaned back in her chair and closed her eyes for a moment. Fatigue penet
                 
                 # 테스트 모드에서 writer10 저장 시 특별 처리
                 if test_quick_mode and role == "writer10":
+                    logger.info(f"Test mode: Processing complete novel ({len(stage_content)} chars)")
+                    
                     # 전체 내용을 10개 챕터로 분할하여 저장
-                    # [Chapter 숫자] 패턴으로 분할
-                    chapter_pattern = r'\[Chapter\s+\d+\]'
-                    chapters_with_headers = re.split(f'({chapter_pattern})', stage_content)
+                    # [Chapter 숫자] 또는 [chapter 숫자] 패턴으로 분할
+                    chapter_pattern = r'\[(?i)chapter\s+\d+\]'
                     
-                    # 챕터 내용만 추출 (헤더와 내용을 쌍으로)
-                    chapters = []
-                    for i in range(1, len(chapters_with_headers), 2):
-                        if i+1 < len(chapters_with_headers):
-                            chapter_content = chapters_with_headers[i+1].strip()
-                            if chapter_content:
-                                chapters.append(chapter_content)
+                    # 챕터 마커가 있는지 확인
+                    chapter_markers = re.findall(chapter_pattern, stage_content, re.IGNORECASE)
+                    logger.info(f"Found {len(chapter_markers)} chapter markers")
                     
-                    # 챕터가 10개가 안 되면 균등 분할
-                    if len(chapters) < 10:
-                        logger.warning(f"Only {len(chapters)} chapters found, splitting evenly")
-                        # 전체 텍스트를 10개로 균등 분할
-                        total_text = stage_content.replace('[Chapter', '\n[Chapter')
-                        words = total_text.split()
-                        words_per_chapter = len(words) // 10
+                    if len(chapter_markers) >= 10:
+                        # 챕터 마커로 분할
+                        parts = re.split(chapter_pattern, stage_content, flags=re.IGNORECASE)
                         chapters = []
-                        for i in range(10):
-                            start_idx = i * words_per_chapter
-                            end_idx = (i + 1) * words_per_chapter if i < 9 else len(words)
-                            chapter_words = words[start_idx:end_idx]
-                            chapters.append(' '.join(chapter_words))
+                        
+                        # 첫 번째 부분은 챕터 이전 내용이므로 제외
+                        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) >= 20:  # 충분한 단락이 있으면
+                            para_per_chapter = len(paragraphs) // 10
+                            chapters = []
+                            for i in range(10):
+                                start_idx = i * para_per_chapter
+                                end_idx = (i + 1) * para_per_chapter if i < 9 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) // 10, 100)  # 최소 100단어
+                            chapters = []
+                            for i in range(10):
+                                start_idx = i * words_per_chapter
+                                end_idx = (i + 1) * words_per_chapter if i < 9 else len(words)
+                                chapter_words = words[start_idx:end_idx]
+                                chapters.append(' '.join(chapter_words))
+                    
+                    # 10개 챕터로 제한
+                    chapters = chapters[:10]
+                    
+                    # 부족한 챕터는 마지막 챕터를 분할해서 채우기
+                    while len(chapters) < 10:
+                        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)+1} - Content pending]")
                     
                     # 각 챕터를 개별 writer로 저장
+                    total_saved_words = 0
                     for i, chapter_content in enumerate(chapters[:10], 1):
                         writer_stage_num = (i-1) * 3 + 5  # 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}",
-                            chapter_content,
+                            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: {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,
@@ -1812,13 +1888,13 @@ Write a high-quality complete novel."""
         return ""
 
 # Gradio Interface Functions
-def process_query(query: str, language: str, session_id: str = None, test_mode: bool = False) -> Generator[Tuple[str, str, str], None, None]:
-    """Process query and yield updates"""
+def process_query(query: str, language: str, session_id: str = None, test_mode: bool = False) -> Generator[Tuple[str, str, str, str], None, None]:
+    """Process query and yield updates with session ID"""
     if not query.strip() and not session_id:
         if language == "Korean":
-            yield "", "", "❌ 소설 주제를 입력해주세요."
+            yield "", "", "❌ 소설 주제를 입력해주세요.", None
         else:
-            yield "", "", "❌ Please enter a novel theme."
+            yield "", "", "❌ Please enter a novel theme.", None
         return
     
     system = NovelWritingSystem()
@@ -1854,14 +1930,15 @@ def process_query(query: str, language: str, session_id: str = None, test_mode:
                 else:
                     status = f"🔄 Processing... ({completed}/{total} - {progress_percent:.1f}%){mode_status}"
             
-            yield stages_display, final_novel, status
+            # Return current session ID
+            yield stages_display, final_novel, status, system.current_session_id
                     
     except Exception as e:
         logger.error(f"Error in process_query: {str(e)}", exc_info=True)
         if language == "Korean":
-            yield "", "", f"❌ 오류 발생: {str(e)}"
+            yield "", "", f"❌ 오류 발생: {str(e)}", None
         else:
-            yield "", "", f"❌ Error occurred: {str(e)}"
+            yield "", "", f"❌ Error occurred: {str(e)}", None
 
 def format_stages_display(stages: List[Dict[str, str]], language: str) -> str:
     """Format stages into simple display with writer save status"""
@@ -1902,13 +1979,13 @@ def get_active_sessions(language: str) -> List[Tuple[str, str]]:
         logger.error(f"Error getting active sessions: {str(e)}", exc_info=True)
         return []
 
-def resume_session(session_id: str, language: str, test_mode: bool = False) -> Generator[Tuple[str, str, str], None, None]:
+def resume_session(session_id: str, language: str, test_mode: bool = False) -> Generator[Tuple[str, str, str, str], None, None]:
     """Resume an existing session"""
     if not session_id:
         if language == "Korean":
-            yield "", "", "❌ 세션을 선택해주세요."
+            yield "", "", "❌ 세션을 선택해주세요.", None
         else:
-            yield "", "", "❌ Please select a session."
+            yield "", "", "❌ Please select a session.", None
         return
     
     # Process with existing session ID
@@ -1936,6 +2013,8 @@ def download_novel(novel_text: str, format: str, language: str, session_id: str
         logger.error("No session_id provided for download")
         return None
     
+    logger.info(f"Starting download for session: {session_id}, format: {format}")
+    
     # DB에서 세션 정보 가져오기
     session = NovelDatabase.get_session(session_id)
     if not session:
@@ -1944,6 +2023,7 @@ def download_novel(novel_text: str, format: str, language: str, session_id: str
     
     # DB에서 모든 스테이지 가져오기
     stages = NovelDatabase.get_stages(session_id)
+    logger.info(f"Found {len(stages)} stages in database")
     
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
     
@@ -1981,10 +2061,31 @@ 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'].startswith('writer') and 'Revision' in stage['stage_name']:
+            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 stages 검색
+        for stage in stages:
+            logger.debug(f"Checking stage: role={stage['role']}, name={stage['stage_name']}, status={stage['status']}")
+            
+            if stage['role'] and stage['role'].startswith('writer') and stage['stage_name'] and 'Revision' in stage['stage_name']:
+                # 작가 번호 추출
+                try:
+                    writer_num_from_role = int(stage['role'].replace('writer', ''))
+                    logger.debug(f"Found writer {writer_num_from_role} revision")
+                except:
+                    logger.warning(f"Could not extract writer number from role: {stage['role']}")
+                    continue
+                    
                 writer_count += 1
                 content = stage['content'] or ""
+                
                 # 페이지 마크 제거
                 content = re.sub(r'\[(?:페이지|Page|page)\s*\d+\]', '', content)
                 content = re.sub(r'(?:페이지|Page)\s*\d+:', '', content)
@@ -1998,9 +2099,44 @@ def download_novel(novel_text: str, format: str, language: str, session_id: str
                         'content': content,
                         'word_count': word_count
                     })
+                    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'Total Writers: {writer_count}')
         doc.add_paragraph(f'Total Words: {total_words:,}')
         doc.add_paragraph(f'Estimated Pages: ~{total_words/500:.0f}')
@@ -2008,32 +2144,33 @@ def download_novel(novel_text: str, format: str, language: str, session_id: str
         doc.add_page_break()
         
         # 목차
-        doc.add_heading('Table of Contents', 1)
-        for i in range(1, writer_count + 1):
-            doc.add_paragraph(f'Chapter {i}: Pages {(i-1)*3+1}-{i*3}')
-        doc.add_page_break()
-        
-        # 각 작가의 내용 추가
-        for writer_data in writer_contents:
-            writer_num = writer_data['writer_num']
-            content = writer_data['content']
-            word_count = writer_data['word_count']
-            
-            # 챕터 헤더
-            doc.add_heading(f'Chapter {writer_num}', 1)
-            doc.add_paragraph(f'Pages {(writer_num-1)*3+1}-{writer_num*3}')
-            doc.add_paragraph(f'Word Count: {word_count:,}')
-            doc.add_paragraph()
-            
-            # 작가 내용 추가
-            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)
-            
-            if writer_num < writer_count:  # 마지막 작가 후에는 페이지 구분 없음
-                doc.add_page_break()
+        if writer_contents:
+            doc.add_heading('Table of Contents', 1)
+            for i in range(1, writer_count + 1):
+                doc.add_paragraph(f'Chapter {i}: Pages {(i-1)*3+1}-{i*3}')
+            doc.add_page_break()
+            
+            # 각 작가의 내용 추가
+            for writer_data in writer_contents:
+                writer_num = writer_data['writer_num']
+                content = writer_data['content']
+                word_count = writer_data['word_count']
+                
+                # 챕터 헤더
+                doc.add_heading(f'Chapter {writer_num}', 1)
+                doc.add_paragraph(f'Pages {(writer_num-1)*3+1}-{writer_num*3}')
+                doc.add_paragraph(f'Word Count: {word_count:,}')
+                doc.add_paragraph()
+                
+                # 작가 내용 추가
+                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)
+                
+                if writer_num < writer_count:  # 마지막 작가 후에는 페이지 구분 없음
+                    doc.add_page_break()
         
         # 페이지 설정
         for section in doc.sections:
@@ -2051,7 +2188,7 @@ def download_novel(novel_text: str, format: str, language: str, session_id: str
         filepath = os.path.join(temp_dir, filename)
         doc.save(filepath)
         
-        logger.info(f"DOCX saved successfully: {filepath} ({total_words} words)")
+        logger.info(f"DOCX saved successfully: {filepath} ({total_words} words, {writer_count} writers)")
         return filepath
     else:
         # TXT format
@@ -2074,7 +2211,7 @@ def download_novel(novel_text: str, format: str, language: str, session_id: str
             
             # 각 작가의 수정본만 출력
             for stage in stages:
-                if stage['role'].startswith('writer') and 'Revision' in stage['stage_name']:
+                if stage['role'] and stage['role'].startswith('writer') and stage['stage_name'] and 'Revision' in stage['stage_name']:
                     writer_count += 1
                     content = stage['content'] or ""
                     # 페이지 마크 제거
@@ -2093,6 +2230,15 @@ 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")
@@ -2308,7 +2454,7 @@ def create_interface():
         submit_btn.click(
             fn=process_query,
             inputs=[query_input, language_select, current_session_id, test_mode_check],
-            outputs=[stages_display, novel_output, status_text]
+            outputs=[stages_display, novel_output, status_text, current_session_id]
         )
         
         # Update novel text state and session ID when novel output changes
@@ -2325,7 +2471,7 @@ def create_interface():
         ).then(
             fn=resume_session,
             inputs=[current_session_id, language_select, test_mode_check],
-            outputs=[stages_display, novel_output, status_text]
+            outputs=[stages_display, novel_output, status_text, current_session_id]
         )
         
         auto_recover_btn.click(
@@ -2335,7 +2481,7 @@ def create_interface():
         ).then(
             fn=resume_session,
             inputs=[current_session_id, language_select, test_mode_check],
-            outputs=[stages_display, novel_output, status_text]
+            outputs=[stages_display, novel_output, status_text, current_session_id]
         )
         
         refresh_btn.click(
@@ -2349,13 +2495,19 @@ def create_interface():
         )
         
         def handle_download(format_type, language, session_id, novel_text):
+            logger.info(f"Download requested: format={format_type}, session_id={session_id}, language={language}")
+            logger.info(f"Session ID type: {type(session_id)}, value: {repr(session_id)}")
+            
             if not session_id:
+                logger.error("No session_id available for download")
                 return gr.update(visible=False)
             
             file_path = download_novel(novel_text, format_type, language, session_id)
             if file_path:
+                logger.info(f"Download successful: {file_path}")
                 return gr.update(value=file_path, visible=True)
             else:
+                logger.error("Download failed - no file generated")
                 return gr.update(visible=False)
         
         download_btn.click(