Update app.py

app.py

@@ -30,11 +30,16 @@ except ImportError:
     logger.warning("python-docx not installed. DOCX export will be disabled.")
 
 # 환경 변수에서 토큰 가져오기
-FRIENDLI_TOKEN = os.getenv("FRIENDLI_TOKEN", "YOUR_FRIENDLI_TOKEN")
+FRIENDLI_TOKEN = os.getenv("FRIENDLI_TOKEN", "")
 API_URL = "https://api.friendli.ai/dedicated/v1/chat/completions"
 MODEL_ID = "dep89a2fld32mcm"
 TEST_MODE = os.getenv("TEST_MODE", "false").lower() == "true"
 
+# 환경 변수 검증
+if not FRIENDLI_TOKEN and not TEST_MODE:
+    logger.warning("FRIENDLI_TOKEN not set and TEST_MODE is false. Application will run in test mode.")
+    TEST_MODE = True
+
 # 전역 변수
 conversation_history = []
 selected_language = "English"  # 기본 언어
@@ -43,13 +48,20 @@ selected_language = "English"  # 기본 언어
 DB_PATH = "novel_sessions.db"
 db_lock = threading.Lock()
 
+# Stage 번호 상수
+WRITER_DRAFT_STAGES = [3, 6, 9, 12, 15, 18, 21, 24, 27, 30]  # 작가 초안
+WRITER_REVISION_STAGES = [5, 8, 11, 14, 17, 20, 23, 26, 29, 32]  # 작가 수정본
+
 class NovelDatabase:
     """Novel session management database"""
     
     @staticmethod
     def init_db():
-        """Initialize database tables"""
+        """Initialize database tables with WAL mode for better concurrency"""
         with sqlite3.connect(DB_PATH) as conn:
+            # Enable WAL mode for better concurrent access
+            conn.execute("PRAGMA journal_mode=WAL")
+            
             cursor = conn.cursor()
             
             # Sessions table
@@ -58,8 +70,8 @@ class NovelDatabase:
                     session_id TEXT PRIMARY KEY,
                     user_query TEXT NOT NULL,
                     language TEXT NOT NULL,
-                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    created_at TEXT DEFAULT (datetime('now')),
+                    updated_at TEXT DEFAULT (datetime('now')),
                     status TEXT DEFAULT 'active',
                     current_stage INTEGER DEFAULT 0,
                     final_novel TEXT
@@ -76,8 +88,9 @@ class NovelDatabase:
                     role TEXT NOT NULL,
                     content TEXT,
                     status TEXT DEFAULT 'pending',
-                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-                    FOREIGN KEY (session_id) REFERENCES sessions(session_id)
+                    created_at TEXT DEFAULT (datetime('now')),
+                    FOREIGN KEY (session_id) REFERENCES sessions(session_id),
+                    UNIQUE(session_id, stage_number)
                 )
             ''')
             
@@ -90,9 +103,9 @@ class NovelDatabase:
     @staticmethod
     @contextmanager
     def get_db():
-        """Database connection context manager"""
+        """Database connection context manager with timeout"""
         with db_lock:
-            conn = sqlite3.connect(DB_PATH)
+            conn = sqlite3.connect(DB_PATH, timeout=30.0)
             conn.row_factory = sqlite3.Row
             try:
                 yield conn
@@ -121,32 +134,19 @@ class NovelDatabase:
         with NovelDatabase.get_db() as conn:
             cursor = conn.cursor()
             
-            # Check if stage exists
+            # UPSERT operation
             cursor.execute('''
-                SELECT id FROM stages 
-                WHERE session_id = ? AND stage_number = ?
-            ''', (session_id, stage_number))
-            
-            existing = cursor.fetchone()
-            
-            if existing:
-                # Update existing stage
-                cursor.execute('''
-                    UPDATE stages 
-                    SET content = ?, status = ?, stage_name = ?
-                    WHERE session_id = ? AND stage_number = ?
-                ''', (content, status, stage_name, session_id, stage_number))
-            else:
-                # Insert new stage
-                cursor.execute('''
-                    INSERT INTO stages (session_id, stage_number, stage_name, role, content, status)
-                    VALUES (?, ?, ?, ?, ?, ?)
-                ''', (session_id, stage_number, stage_name, role, content, status))
+                INSERT INTO stages (session_id, stage_number, stage_name, role, content, status)
+                VALUES (?, ?, ?, ?, ?, ?)
+                ON CONFLICT(session_id, stage_number) 
+                DO UPDATE SET content=?, status=?, stage_name=?, updated_at=datetime('now')
+            ''', (session_id, stage_number, stage_name, role, content, status,
+                  content, status, stage_name))
             
             # Update session
             cursor.execute('''
                 UPDATE sessions 
-                SET updated_at = CURRENT_TIMESTAMP, current_stage = ?
+                SET updated_at = datetime('now'), current_stage = ?
                 WHERE session_id = ?
             ''', (stage_number, session_id))
             
@@ -179,11 +179,8 @@ class NovelDatabase:
         with NovelDatabase.get_db() as conn:
             cursor = conn.cursor()
             
-            # 작가 수정본만 가져오기 (stage_number 5, 8, 11, 14, 17, 20, 23, 26, 29, 32)
-            writer_revision_stages = [5, 8, 11, 14, 17, 20, 23, 26, 29, 32]
-            
             all_content = []
-            for stage_num in writer_revision_stages:
+            for stage_num in WRITER_REVISION_STAGES:
                 cursor.execute('''
                     SELECT content, stage_name FROM stages 
                     WHERE session_id = ? AND stage_number = ?
@@ -212,7 +209,7 @@ class NovelDatabase:
             cursor = conn.cursor()
             cursor.execute('''
                 UPDATE sessions 
-                SET final_novel = ?, status = 'complete', updated_at = CURRENT_TIMESTAMP
+                SET final_novel = ?, status = 'complete', updated_at = datetime('now')
                 WHERE session_id = ?
             ''', (final_novel, session_id))
             conn.commit()
@@ -231,22 +228,35 @@ class NovelDatabase:
                 LIMIT 10
             ''')
             return cursor.fetchall()
+    
+    @staticmethod
+    def parse_datetime(datetime_str: str) -> datetime:
+        """Parse SQLite datetime string safely"""
+        try:
+            # Try ISO format first
+            return datetime.fromisoformat(datetime_str)
+        except:
+            # Fallback to SQLite default format
+            return datetime.strptime(datetime_str, "%Y-%m-%d %H:%M:%S")
 
 class NovelWritingSystem:
     def __init__(self):
         self.token = FRIENDLI_TOKEN
         self.api_url = API_URL
         self.model_id = MODEL_ID
-        self.test_mode = TEST_MODE or (self.token == "YOUR_FRIENDLI_TOKEN")
+        self.test_mode = TEST_MODE or not self.token
         
         if self.test_mode:
-            logger.warning("Running in test mode.")
+            logger.warning("Running in test mode - no actual API calls will be made.")
+        else:
+            logger.info("Running in production mode with API calls enabled.")
         
         # Initialize database
         NovelDatabase.init_db()
         
         # Session management
         self.current_session_id = None
+        self.total_stages = 0  # Will be set in process_novel_stream
     
     def create_headers(self):
         """API 헤더 생성"""
@@ -931,7 +941,8 @@ Present a complete 50-page novel integrating all writers' contributions."""
                                 content = chunk["choices"][0].get("delta", {}).get("content", "")
                                 if content:
                                     buffer += content
-                                    if len(buffer) > 50 or '\n' in buffer:
+                                    # 더 큰 버퍼로 수정 (긴 단어 대응)
+                                    if len(buffer) > 100 or '\n\n' in buffer:
                                         yield buffer
                                         buffer = ""
                         except json.JSONDecodeError:
@@ -1175,6 +1186,9 @@ Present a complete 50-page novel integrating all writers' contributions."""
                 ("director", f"🎬 {'감독자: 최종 완성본' if language == 'Korean' else 'Director: Final Version'}")
             ])
             
+            # Store total stages for get_stage_prompt
+            self.total_stages = len(stage_definitions)
+            
             # Process stages starting from resume point
             for stage_idx in range(resume_from_stage, len(stage_definitions)):
                 role, stage_name = stage_definitions[stage_idx]
@@ -1229,7 +1243,7 @@ Present a complete 50-page novel integrating all writers' contributions."""
             yield final_novel, stages
             
         except Exception as e:
-            logger.error(f"Error in process_novel_stream: {str(e)}")
+            logger.error(f"Error in process_novel_stream: {str(e)}", exc_info=True)
             
             # Save error state to DB
             if self.current_session_id:
@@ -1315,7 +1329,7 @@ Present a complete 50-page novel integrating all writers' contributions."""
                         )
         
         # Director final - DB에서 모든 작가 내용 가져오기
-        elif stage_idx == len(stage_definitions) - 1:
+        elif stage_idx == self.total_stages - 1:  # Fixed: using self.total_stages
             critic_final_idx = stage_idx - 1
             all_writer_content = NovelDatabase.get_all_writer_content(self.current_session_id)
             logger.info(f"Final director compilation with {len(all_writer_content)} characters of content")
@@ -1349,7 +1363,7 @@ def process_query(query: str, language: str, session_id: str = None) -> Generato
             yield stages_display, final_novel, status
                     
     except Exception as e:
-        logger.error(f"Error in process_query: {str(e)}")
+        logger.error(f"Error in process_query: {str(e)}", exc_info=True)
         if language == "Korean":
             yield "", "", f"❌ 오류 발생: {str(e)}"
         else:
@@ -1378,7 +1392,7 @@ def get_active_sessions(language: str) -> List[Tuple[str, str]]:
         
         choices = []
         for session in sessions:
-            created = datetime.fromisoformat(session['created_at'])
+            created = NovelDatabase.parse_datetime(session['created_at'])
             date_str = created.strftime("%Y-%m-%d %H:%M")
             query_preview = session['user_query'][:50] + "..." if len(session['user_query']) > 50 else session['user_query']
             label = f"[{date_str}] {query_preview} (Stage {session['current_stage']})"
@@ -1386,7 +1400,7 @@ def get_active_sessions(language: str) -> List[Tuple[str, str]]:
         
         return choices
     except Exception as e:
-        logger.error(f"Error getting active sessions: {str(e)}")
+        logger.error(f"Error getting active sessions: {str(e)}", exc_info=True)
         return []
 
 def resume_session(session_id: str, language: str) -> Generator[Tuple[str, str, str], None, None]:
@@ -1416,9 +1430,11 @@ def download_novel(novel_text: str, format: str, language: str) -> str:
         lines = novel_text.split('\n')
         for line in lines:
             if line.startswith('#'):
-                level = len(line.split()[0])
+                # Safe heading level extraction
+                level = min(len(line.split()[0].strip('#')), 9)
                 text = line.lstrip('#').strip()
-                doc.add_heading(text, level)
+                if text:
+                    doc.add_heading(text, level)
             elif line.strip():
                 doc.add_paragraph(line)
         
@@ -1613,7 +1629,7 @@ def create_interface():
                 sessions = get_active_sessions("English")
                 return gr.update(choices=sessions)
             except Exception as e:
-                logger.error(f"Error refreshing sessions: {str(e)}")
+                logger.error(f"Error refreshing sessions: {str(e)}", exc_info=True)
                 return gr.update(choices=[])
         
         submit_btn.click(
@@ -1674,6 +1690,12 @@ def create_interface():
 if __name__ == "__main__":
     logger.info("Starting SOMA Novel Writing System...")
     
+    # Check environment
+    if TEST_MODE:
+        logger.warning("Running in TEST MODE - no actual API calls will be made")
+    else:
+        logger.info(f"Running in PRODUCTION MODE with API endpoint: {API_URL}")
+    
     # Initialize database on startup
     logger.info("Initializing database...")
     NovelDatabase.init_db()