Update app.py

app.py

@@ -31,6 +31,7 @@ except ImportError:
 
 # 환경 변수에서 토큰 가져오기
 FRIENDLI_TOKEN = os.getenv("FRIENDLI_TOKEN", "")
+BRAVE_SEARCH_API_KEY = os.getenv("BRAVE_SEARCH_API_KEY", "")
 API_URL = "https://api.friendli.ai/dedicated/v1/chat/completions"
 MODEL_ID = "dep89a2fld32mcm"
 TEST_MODE = os.getenv("TEST_MODE", "false").lower() == "true"
@@ -40,6 +41,9 @@ 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
 
+if not BRAVE_SEARCH_API_KEY:
+    logger.warning("BRAVE_SEARCH_API_KEY not set. Web search features will be disabled.")
+
 # 전역 변수
 conversation_history = []
 selected_language = "English"  # 기본 언어
@@ -52,8 +56,148 @@ db_lock = threading.Lock()
 WRITER_DRAFT_STAGES = [3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 36, 39, 42, 45, 48]  # 작가 초안
 WRITER_REVISION_STAGES = [5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35, 38, 41, 44, 47, 50]  # 작가 수정본
 
+class WebSearchIntegration:
+    """Brave Search API integration for research"""
+    
+    def __init__(self):
+        self.brave_api_key = BRAVE_SEARCH_API_KEY
+        self.search_url = "https://api.search.brave.com/res/v1/web/search"
+        self.enabled = bool(self.brave_api_key)
+        
+    def search(self, query: str, count: int = 5, language: str = "en") -> List[Dict]:
+        """Perform web search using Brave Search API"""
+        if not self.enabled:
+            return []
+            
+        headers = {
+            "Accept": "application/json",
+            "X-Subscription-Token": self.brave_api_key
+        }
+        
+        # 언어에 따른 검색 설정
+        search_lang = "ko" if language == "Korean" else "en"
+        
+        params = {
+            "q": query,
+            "count": count,
+            "search_lang": search_lang,
+            "text_decorations": False,
+            "safesearch": "moderate"
+        }
+        
+        try:
+            response = requests.get(self.search_url, headers=headers, params=params, timeout=10)
+            if response.status_code == 200:
+                results = response.json()
+                web_results = results.get("web", {}).get("results", [])
+                logger.info(f"Search successful: Found {len(web_results)} results for '{query}'")
+                return web_results
+            else:
+                logger.error(f"Search API error: {response.status_code}")
+                return []
+        except Exception as e:
+            logger.error(f"Search error: {str(e)}")
+            return []
+    
+    def extract_relevant_info(self, results: List[Dict], max_chars: int = 3000) -> str:
+        """Extract relevant information from search results"""
+        if not results:
+            return ""
+            
+        extracted = []
+        total_chars = 0
+        
+        for i, result in enumerate(results[:5], 1):  # 최대 5개 결과
+            title = result.get("title", "")
+            description = result.get("description", "")
+            url = result.get("url", "")
+            
+            # 일부 내용 추출
+            extra_text = ""
+            if "extra_snippets" in result:
+                extra_text = " ".join(result["extra_snippets"][:2])
+            
+            info = f"""[{i}] {title}
+{description}
+{extra_text}
+Source: {url}
+"""
+            
+            if total_chars + len(info) < max_chars:
+                extracted.append(info)
+                total_chars += len(info)
+            else:
+                break
+                
+        return "\n---\n".join(extracted)
+    
+    def create_research_queries(self, topic: str, role: str, stage_info: Dict, language: str = "English") -> List[str]:
+        """Create multiple research queries based on role and context"""
+        queries = []
+        
+        if language == "Korean":
+            if role == "director":
+                queries = [
+                    f"{topic} 소설 배경 설정",
+                    f"{topic} 역사적 사실",
+                    f"{topic} 문화적 특징"
+                ]
+            elif role.startswith("writer"):
+                writer_num = int(role.replace("writer", ""))
+                if writer_num <= 3:  # 초반부 작가
+                    queries = [
+                        f"{topic} 구체적 장면 묘사",
+                        f"{topic} 전문 용어 설명"
+                    ]
+                elif writer_num <= 8:  # 중반부 작가
+                    queries = [
+                        f"{topic} 갈등 상황 사례",
+                        f"{topic} 심리적 측면"
+                    ]
+                else:  # 후반부 작가
+                    queries = [
+                        f"{topic} 해결 방법",
+                        f"{topic} 감동적인 사례"
+                    ]
+            elif role == "critic":
+                queries = [
+                    f"{topic} 문학 작품 분석",
+                    f"{topic} 유사 소설 추천"
+                ]
+        else:
+            if role == "director":
+                queries = [
+                    f"{topic} novel setting ideas",
+                    f"{topic} historical facts",
+                    f"{topic} cultural aspects"
+                ]
+            elif role.startswith("writer"):
+                writer_num = int(role.replace("writer", ""))
+                if writer_num <= 3:  # Early writers
+                    queries = [
+                        f"{topic} vivid scene descriptions",
+                        f"{topic} technical terminology explained"
+                    ]
+                elif writer_num <= 8:  # Middle writers
+                    queries = [
+                        f"{topic} conflict scenarios",
+                        f"{topic} psychological aspects"
+                    ]
+                else:  # Later writers
+                    queries = [
+                        f"{topic} resolution methods",
+                        f"{topic} emotional stories"
+                    ]
+            elif role == "critic":
+                queries = [
+                    f"{topic} literary analysis",
+                    f"{topic} similar novels recommendations"
+                ]
+                
+        return queries
+
 class NovelDatabase:
-    """Novel session management database"""
+    """Novel session management database with search history"""
     
     @staticmethod
     def init_db():
@@ -95,12 +239,39 @@ class NovelDatabase:
                 )
             ''')
             
+            # Search history table - 검색 이력 저장
+            cursor.execute('''
+                CREATE TABLE IF NOT EXISTS search_history (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    session_id TEXT NOT NULL,
+                    stage_number INTEGER NOT NULL,
+                    role TEXT NOT NULL,
+                    query TEXT NOT NULL,
+                    results TEXT,
+                    created_at TEXT DEFAULT (datetime('now')),
+                    FOREIGN KEY (session_id) REFERENCES sessions(session_id)
+                )
+            ''')
+            
             # Create indices
             cursor.execute('CREATE INDEX IF NOT EXISTS idx_session_id ON stages(session_id)')
             cursor.execute('CREATE INDEX IF NOT EXISTS idx_stage_number ON stages(stage_number)')
+            cursor.execute('CREATE INDEX IF NOT EXISTS idx_search_session ON search_history(session_id)')
             
             conn.commit()
     
+    @staticmethod
+    def save_search_history(session_id: str, stage_number: int, role: str, query: str, results: str):
+        """Save search history"""
+        with sqlite3.connect(DB_PATH) as conn:
+            cursor = conn.cursor()
+            cursor.execute('''
+                INSERT INTO search_history (session_id, stage_number, role, query, results)
+                VALUES (?, ?, ?, ?, ?)
+            ''', (session_id, stage_number, role, query, results))
+            conn.commit()
+    
+    # ... (이전의 모든 NovelDatabase 메서드들은 동일)
     @staticmethod
     @contextmanager
     def get_db():
@@ -299,10 +470,18 @@ class NovelWritingSystem:
         self.model_id = MODEL_ID
         self.test_mode = TEST_MODE or not self.token
         
+        # Web search integration
+        self.web_search = WebSearchIntegration()
+        
         if self.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.")
+            
+        if self.web_search.enabled:
+            logger.info("Web search is enabled with Brave Search API.")
+        else:
+            logger.warning("Web search is disabled. Set BRAVE_SEARCH_API_KEY to enable.")
         
         # Initialize database
         NovelDatabase.init_db()
@@ -318,6 +497,60 @@ class NovelWritingSystem:
             "Content-Type": "application/json"
         }
     
+    def enhance_prompt_with_research(self, original_prompt: str, role: str, 
+                                   topic: str, stage_info: Dict, language: str = "English") -> str:
+        """Enhance prompt with web search results"""
+        if not self.web_search.enabled or self.test_mode:
+            return original_prompt
+        
+        # Create research queries
+        queries = self.web_search.create_research_queries(topic, role, stage_info, language)
+        
+        all_research = []
+        for query in queries[:2]:  # 최대 2개 쿼리
+            logger.info(f"Searching: {query}")
+            results = self.web_search.search(query, count=3, language=language)
+            if results:
+                research_text = self.web_search.extract_relevant_info(results, max_chars=1500)
+                if research_text:
+                    all_research.append(f"### {query}\n{research_text}")
+                    
+                    # Save search history
+                    if self.current_session_id:
+                        NovelDatabase.save_search_history(
+                            self.current_session_id,
+                            stage_info.get('stage_idx', 0),
+                            role,
+                            query,
+                            research_text
+                        )
+        
+        if not all_research:
+            return original_prompt
+        
+        # Add research to prompt
+        if language == "Korean":
+            research_section = f"""
+## 웹 검색 참고 자료:
+{chr(10).join(all_research)}
+
+위의 검색 결과를 참고하여 더욱 사실적이고 구체적인 내용을 작성하세요.
+검색 결과의 정보를 창의적으로 활용하되, 직접 인용은 피하고 소설에 자연스럽게 녹여내세요.
+실제 사실과 창작을 적절히 조화시켜 독자가 몰입할 수 있는 이야기를 만드세요.
+"""
+        else:
+            research_section = f"""
+## Web Search Reference:
+{chr(10).join(all_research)}
+
+Use the above search results to create more realistic and specific content.
+Creatively incorporate the information from search results, but avoid direct quotes and naturally blend them into the novel.
+Balance real facts with creative fiction to create an immersive story for readers.
+"""
+        
+        return original_prompt + "\n\n" + research_section
+    
+    # ... (이전의 모든 프롬프트 생성 메서드들은 동일)
     def create_director_initial_prompt(self, user_query: str, language: str = "English") -> str:
         """Director AI initial prompt - Novel planning for 16 writers"""
         if language == "Korean":
@@ -957,8 +1190,9 @@ Present a complete 48-page novel integrating all writers' contributions."""
             yield chunk + " "
             time.sleep(0.02)
     
-    def call_llm_streaming(self, messages: List[Dict[str, str]], role: str, language: str = "English") -> Generator[str, None, None]:
-        """Streaming LLM API call with improved prompting"""
+    def call_llm_streaming(self, messages: List[Dict[str, str]], role: str, 
+                          language: str = "English", stage_info: Dict = None) -> Generator[str, None, None]:
+        """Streaming LLM API call with web search enhancement"""
         
         if self.test_mode:
             logger.info(f"Test mode streaming - Role: {role}, Language: {language}")
@@ -966,6 +1200,22 @@ Present a complete 48-page novel integrating all writers' contributions."""
             yield from self.simulate_streaming(test_response, role)
             return
         
+        # Extract topic from user query for research
+        topic = ""
+        if stage_info and 'query' in stage_info:
+            topic = stage_info['query']
+        
+        # Enhance prompt with web search if available
+        if messages and messages[-1]["role"] == "user" and self.web_search.enabled:
+            enhanced_prompt = self.enhance_prompt_with_research(
+                messages[-1]["content"],
+                role,
+                topic,
+                stage_info or {},
+                language
+            )
+            messages[-1]["content"] = enhanced_prompt
+        
         # Real API call
         try:
             system_prompts = self.get_system_prompts(language)
@@ -1118,6 +1368,7 @@ Present a complete 48-page novel integrating all writers' contributions."""
             logger.error(f"Error during streaming: {str(e)}")
             yield f"❌ Error occurred: {str(e)}"
     
+    # ... (나머지 메서드들은 동일)
     def get_system_prompts(self, language: str) -> Dict[str, str]:
         """Get system prompts for all 16 writers with enhanced length requirements"""
         if language == "Korean":
@@ -1340,7 +1591,7 @@ Seoyeon leaned back in her chair and closed her eyes for a moment. Fatigue penet
     def process_novel_stream(self, query: str, language: str = "English", 
                            session_id: Optional[str] = None, 
                            resume_from_stage: int = 0) -> Generator[Tuple[str, List[Dict[str, str]]], None, None]:
-        """Process novel writing with streaming updates"""
+        """Process novel writing with streaming updates and web search"""
         try:
             global conversation_history
             
@@ -1403,6 +1654,10 @@ Seoyeon leaned back in her chair and closed her eyes for a moment. Fatigue penet
             for stage_idx in range(resume_from_stage, len(stage_definitions)):
                 role, stage_name = stage_definitions[stage_idx]
                 
+                # Add search indicator if enabled
+                if self.web_search.enabled and not self.test_mode:
+                    stage_name += " 🔍"
+                
                 # Add stage if not already present
                 if stage_idx >= len(stages):
                     stages.append({
@@ -1418,13 +1673,21 @@ Seoyeon leaned back in her chair and closed her eyes for a moment. Fatigue penet
                 # Get appropriate prompt based on stage
                 prompt = self.get_stage_prompt(stage_idx, role, query, language, stages)
                 
+                # Create stage info for web search
+                stage_info = {
+                    'stage_idx': stage_idx,
+                    'query': query,
+                    'stage_name': stage_name
+                }
+                
                 stage_content = ""
                 
-                # Stream content generation
+                # Stream content generation with web search
                 for chunk in self.call_llm_streaming(
                     [{"role": "user", "content": prompt}],
                     role,
-                    language
+                    language,
+                    stage_info
                 ):
                     stage_content += chunk
                     stages[stage_idx]["content"] = stage_content
@@ -1559,7 +1822,7 @@ Seoyeon leaned back in her chair and closed her eyes for a moment. Fatigue penet
         
         return ""
 
-# Gradio Interface Functions
+# Gradio Interface Functions (동일)
 def process_query(query: str, language: str, session_id: str = None) -> Generator[Tuple[str, str, str], None, None]:
     """Process query and yield updates"""
     if not query.strip() and not session_id:
@@ -1765,6 +2028,11 @@ custom_css = """
     border-radius: 8px;
     margin-top: 20px;
 }
+
+.search-indicator {
+    color: #4CAF50;
+    font-weight: bold;
+}
 """
 
 # Create Gradio Interface
@@ -1776,11 +2044,12 @@ def create_interface():
                 📚 SOMA Novel Writing System
             </h1>
             <h3 style="color: #ccc; margin-bottom: 20px;">
-                AI Collaborative Novel Generation - 48 Page Novella Creator
+                AI Collaborative Novel Generation with Web Research - 48 Page Novella Creator
             </h3>
             <p style="font-size: 1.1em; color: #ddd; max-width: 800px; margin: 0 auto;">
                 Enter a theme or prompt, and watch as 19 AI agents collaborate to create a complete 48-page novella.
                 The system includes 1 Director, 1 Critic, and 16 Writers (each writing 3 pages) working in harmony.
+                <span class="search-indicator">🔍 Web search enabled</span> - agents will research relevant information for more realistic content.
                 All progress is automatically saved and can be resumed anytime.
             </p>
         </div>
@@ -1804,6 +2073,11 @@ def create_interface():
                         label="Language / 언어"
                     )
                     
+                    # Web search status indicator
+                    web_search_status = gr.Markdown(
+                        value=f"🔍 **Web Search:** {'Enabled' if WebSearchIntegration().enabled else 'Disabled (Set BRAVE_SEARCH_API_KEY)'}"
+                    )
+                    
                     with gr.Row():
                         submit_btn = gr.Button("🚀 Start Writing / 작성 시작", variant="primary", scale=2)
                         clear_btn = gr.Button("🗑️ Clear / 초기화", scale=1)
@@ -1979,6 +2253,12 @@ if __name__ == "__main__":
     else:
         logger.info(f"Running in PRODUCTION MODE with API endpoint: {API_URL}")
     
+    # Check web search
+    if BRAVE_SEARCH_API_KEY:
+        logger.info("Web search is ENABLED with Brave Search API")
+    else:
+        logger.warning("Web search is DISABLED. Set BRAVE_SEARCH_API_KEY environment variable to enable.")
+    
     # Initialize database on startup
     logger.info("Initializing database...")
     NovelDatabase.init_db()