Update app.py

app.py

@@ -3,6 +3,8 @@ import os
 import sys
 import json
 import requests
+import re
+from typing import List, Tuple
 
 # 현재 디렉토리와 폰트 파일 경로 설정
 CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
@@ -77,8 +79,88 @@ from wbs_diagram_generator import generate_wbs_diagram
 
 from sample_data import CONCEPT_MAP_JSON, SYNOPTIC_CHART_JSON, RADIAL_DIAGRAM_JSON, PROCESS_FLOW_JSON, WBS_DIAGRAM_JSON
 
-# LLM API 함수
-def call_llm_api(prompt, diagram_type):
+# 브레이브 서치 API 함수
+def brave_search(query: str) -> List[dict]:
+    """브레이브 서치 API를 사용하여 검색 수행"""
+    api_key = os.getenv("BSEARCH_API")
+    if not api_key:
+        print("Warning: BSEARCH_API environment variable not set")
+        return []
+    
+    url = "https://api.search.brave.com/res/v1/web/search"
+    headers = {
+        "Accept": "application/json",
+        "X-Subscription-Token": api_key
+    }
+    params = {
+        "q": query,
+        "count": 5  # 상위 5개 결과만 가져오기
+    }
+    
+    try:
+        response = requests.get(url, headers=headers, params=params, timeout=10)
+        if response.status_code == 200:
+            data = response.json()
+            results = []
+            if "web" in data and "results" in data["web"]:
+                for result in data["web"]["results"][:5]:
+                    results.append({
+                        "title": result.get("title", ""),
+                        "description": result.get("description", ""),
+                        "url": result.get("url", "")
+                    })
+            return results
+        else:
+            print(f"Brave Search API error: {response.status_code}")
+            return []
+    except Exception as e:
+        print(f"Brave Search error: {str(e)}")
+        return []
+
+def extract_keywords(prompt: str, diagram_type: str) -> str:
+    """프롬프트에서 핵심 키워드 추출"""
+    # 다이어그램 타입 제거
+    prompt_clean = prompt.lower()
+    
+    # 일반적인 요청 표현 제거
+    remove_patterns = [
+        r"create\s+a?\s*", r"make\s+a?\s*", r"generate\s+a?\s*", r"build\s+a?\s*",
+        r"design\s+a?\s*", r"develop\s+a?\s*", r"show\s+me\s*", r"i\s+need\s+a?\s*",
+        r"diagram\s*", r"chart\s*", r"map\s*", r"flow\s*", r"를\s*만들어\s*",
+        r"생성해\s*", r"그려\s*", r"보여\s*", r"설계\s*", r"구축\s*"
+    ]
+    
+    for pattern in remove_patterns:
+        prompt_clean = re.sub(pattern, "", prompt_clean)
+    
+    # 주요 명사와 핵심 용어만 추출
+    # 간단한 방법: 길이가 3자 이상인 단어들을 추출
+    words = prompt_clean.split()
+    keywords = []
+    
+    for word in words:
+        # 불필요한 단어 제외
+        if len(word) > 2 and word not in ["the", "and", "for", "with", "that", "this", "from", "about"]:
+            keywords.append(word)
+    
+    # 최대 5개 키워드만 사용
+    return " ".join(keywords[:5])
+
+def enrich_prompt_with_search(prompt: str, diagram_type: str, search_results: List[dict]) -> str:
+    """검색 결과를 바탕으로 프롬프트를 보강"""
+    if not search_results:
+        return prompt
+    
+    # 검색 결과에서 유용한 정보 추출
+    enrichment = "\n\nAdditional context from web search:\n"
+    for i, result in enumerate(search_results[:3]):  # 상위 3개만 사용
+        enrichment += f"- {result['title']}: {result['description']}\n"
+    
+    enriched_prompt = f"{prompt}{enrichment}"
+    return enriched_prompt
+
+# LLM API 함수 수정
+def call_llm_api(prompt, diagram_type, use_search=False, search_results=None):
     """Friendli AI API를 호출하여 JSON 생성"""
     token = os.environ.get("FRIENDLI_TOKEN") or "YOUR_FRIENDLI_TOKEN"
     url = "https://api.friendli.ai/dedicated/v1/chat/completions"
@@ -193,9 +275,16 @@ def call_llm_api(prompt, diagram_type):
 }"""
     }
     
+    # 검색 결과가 있으면 시스템 프롬프트에 추가
+    search_context = ""
+    if use_search and search_results:
+        search_context = "\n\nUse the following search results to enrich the diagram content:\n"
+        for result in search_results[:3]:
+            search_context += f"- {result['title']}: {result['description']}\n"
+    
     system_prompt = f"""You are a helpful assistant that generates JSON structures for diagrams.
 {json_guides.get(diagram_type, '')}
-
+{search_context}
 Important rules:
 1. Generate ONLY valid JSON without any explanation or markdown formatting
 2. The JSON must follow the EXACT structure shown above - do not change field names
@@ -205,7 +294,8 @@ Important rules:
 6. Ensure all connections reference existing node IDs
 7. For Process Flow: 'type' can be: "process", "decision", "start", "end", "io"
 8. For nested structures (Concept Map, Synoptic Chart, Radial, WBS), include empty 'subnodes' or 'subtasks' arrays when there are no children
-9. Do not add any additional fields not shown in the example structure"""
+9. Do not add any additional fields not shown in the example structure
+10. If search results are provided, incorporate relevant information into the diagram content"""
 
     payload = {
         "model": "dep89a2fld32mcm",
@@ -270,13 +360,29 @@ Important rules:
         print(f"LLM API Error: {str(e)}")
         return json.dumps({"error": str(e)})
 
-def generate_with_llm(prompt, diagram_type, output_format):
+def generate_with_llm(prompt, diagram_type, output_format, use_search):
     """LLM으로 JSON을 생성하고 다이어그램을 생성"""
     if not prompt:
-        return None, "Please enter a prompt"
+        return None, "Please enter a prompt", ""
+    
+    search_results = []
+    search_info = ""
+    
+    # 브레이브 서치 사용 시
+    if use_search:
+        keywords = extract_keywords(prompt, diagram_type)
+        if keywords:
+            search_info = f"🔍 Searching for: {keywords}\n"
+            search_results = brave_search(keywords)
+            if search_results:
+                search_info += f"✅ Found {len(search_results)} relevant results\n\n"
+                for i, result in enumerate(search_results[:3]):
+                    search_info += f"{i+1}. {result['title']}\n   {result['description'][:100]}...\n\n"
+            else:
+                search_info += "❌ No search results found\n\n"
     
-    # LLM으로 JSON 생성
-    generated_json = call_llm_api(prompt, diagram_type)
+    # LLM으로 JSON 생성 (검색 결과 포함)
+    generated_json = call_llm_api(prompt, diagram_type, use_search, search_results)
     
     try:
         # JSON 유효성 검사
@@ -286,13 +392,13 @@ def generate_with_llm(prompt, diagram_type, output_format):
         # JSON 구조 검증
         if diagram_type in ["Concept Map", "Synoptic Chart", "Radial Diagram"]:
             if "central_node" not in json_data or "nodes" not in json_data:
-                return None, f"Invalid JSON structure for {diagram_type}. Missing 'central_node' or 'nodes'. Generated JSON:\n{json_str}"
+                return None, f"Invalid JSON structure for {diagram_type}. Missing 'central_node' or 'nodes'. Generated JSON:\n{json_str}", search_info
         elif diagram_type == "Process Flow":
             if "start_node" not in json_data or "nodes" not in json_data or "connections" not in json_data:
-                return None, f"Invalid JSON structure for Process Flow. Missing 'start_node', 'nodes', or 'connections'. Generated JSON:\n{json_str}"
+                return None, f"Invalid JSON structure for Process Flow. Missing 'start_node', 'nodes', or 'connections'. Generated JSON:\n{json_str}", search_info
         elif diagram_type == "WBS Diagram":
             if "project_title" not in json_data or "phases" not in json_data:
-                return None, f"Invalid JSON structure for WBS. Missing 'project_title' or 'phases'. Generated JSON:\n{json_str}"
+                return None, f"Invalid JSON structure for WBS. Missing 'project_title' or 'phases'. Generated JSON:\n{json_str}", search_info
         
         # 다이어그램 생성
         try:
@@ -307,21 +413,21 @@ def generate_with_llm(prompt, diagram_type, output_format):
             elif diagram_type == "WBS Diagram":
                 diagram = generate_wbs_diagram(json_str, output_format)
             else:
-                return None, "Invalid diagram type"
+                return None, "Invalid diagram type", search_info
             
             # 에러 메시지가 반환된 경우
             if isinstance(diagram, str) and diagram.startswith("Error:"):
-                return None, f"Diagram generation error: {diagram}\n\nGenerated JSON:\n{json_str}"
+                return None, f"Diagram generation error: {diagram}\n\nGenerated JSON:\n{json_str}", search_info
             
-            return diagram, json_str
+            return diagram, json_str, search_info
             
         except Exception as e:
-            return None, f"Error generating diagram: {str(e)}\n\nGenerated JSON:\n{json_str}"
+            return None, f"Error generating diagram: {str(e)}\n\nGenerated JSON:\n{json_str}", search_info
         
     except json.JSONDecodeError as e:
-        return None, f"Invalid JSON generated: {str(e)}\n\nGenerated content:\n{generated_json}"
+        return None, f"Invalid JSON generated: {str(e)}\n\nGenerated content:\n{generated_json}", search_info
     except Exception as e:
-        return None, f"Unexpected error: {str(e)}\n\nGenerated content:\n{generated_json}"
+        return None, f"Unexpected error: {str(e)}\n\nGenerated content:\n{generated_json}", search_info
 
 if __name__ == "__main__":
     DEFAULT_BASE_COLOR = '#19191a' 
@@ -482,6 +588,23 @@ if __name__ == "__main__":
             margin-bottom: 20px;
         }
         
+        /* 검색 정보 스타일 */
+        .search-info {
+            background: #f8f9fa;
+            padding: 15px;
+            border-radius: 10px;
+            margin-top: 10px;
+            border-left: 4px solid #667eea;
+        }
+        
+        /* 체크박스 스타일 */
+        .gr-checkbox {
+            margin-top: 15px;
+            padding: 10px;
+            background: #f8f9fa;
+            border-radius: 8px;
+        }
+        
         /* 반응형 디자인 */
         @media (max-width: 768px) {
             .main-container {
@@ -541,8 +664,22 @@ if __name__ == "__main__":
                                 interactive=True
                             )
                             
+                            use_search_checkbox = gr.Checkbox(
+                                label="🔍 Use Brave Search to enhance content",
+                                value=False,
+                                info="Search the web for relevant information to enrich your diagram"
+                            )
+                            
                             generate_btn = gr.Button("✨ Generate with AI", variant="primary", size="lg")
                             
+                            search_info_output = gr.Textbox(
+                                label="🔍 Search Information",
+                                lines=8,
+                                interactive=False,
+                                visible=False,
+                                elem_classes=["search-info"]
+                            )
+                            
                             generated_json_output = gr.Textbox(
                                 label="📄 Generated JSON",
                                 lines=15,
@@ -558,10 +695,20 @@ if __name__ == "__main__":
                                 height=600
                             )
                     
+                    # 검색 체크박스 변경 시 검색 정보 출력 표시/숨김
+                    def toggle_search_info(use_search):
+                        return gr.update(visible=use_search)
+                    
+                    use_search_checkbox.change(
+                        fn=toggle_search_info,
+                        inputs=[use_search_checkbox],
+                        outputs=[search_info_output]
+                    )
+                    
                     generate_btn.click(
                         fn=generate_with_llm,
-                        inputs=[prompt_input, diagram_type_select, output_format_radio],
-                        outputs=[ai_output_image, generated_json_output]
+                        inputs=[prompt_input, diagram_type_select, output_format_radio, use_search_checkbox],
+                        outputs=[ai_output_image, generated_json_output, search_info_output]
                     )
                     
                     with gr.Row(elem_classes=["panel-section"]):