Update app.py

app.py

@@ -9,14 +9,36 @@ from requests.adapters import HTTPAdapter
 from requests.packages.urllib3.util.retry import Retry
 from openai import OpenAI
 from bs4 import BeautifulSoup
-import re  # re 모듈 추가
-import json
-import os
-from datetime import datetime
-import sqlite3
+import re
 import pathlib
+import sqlite3
 
-# DB 초기화 함수
+# 한국 기업 리스트
+KOREAN_COMPANIES = [
+    "SAMSUNG",
+    "HYNIX",
+    "Celltrion",
+    "KIA",
+    "KB",
+    "NAVER",
+    "SHINHAN",
+    "HYUNDAI",
+    "POSCO",
+    "HANA",
+    "KT&G",
+    "WOORI",
+    "LG",
+    "IBK",
+    "POSCO",
+    "KT",
+    "SKT",
+    "DOOSAN",
+    "KAKAO",
+    "HANWHA",
+    "SK"
+]
+
+# DB 관련 함수들
 def init_db():
     db_path = pathlib.Path("search_results.db")
     conn = sqlite3.connect(db_path)
@@ -30,7 +52,6 @@ def init_db():
     conn.commit()
     conn.close()
 
-# 검색 결과 저장 함수
 def save_to_db(keyword, country, results):
     conn = sqlite3.connect("search_results.db")
     c = conn.cursor()
@@ -39,7 +60,6 @@ def save_to_db(keyword, country, results):
     conn.commit()
     conn.close()
 
-# DB에서 검색 결과 불러오기 함수
 def load_from_db(keyword, country):
     conn = sqlite3.connect("search_results.db")
     c = conn.cursor()
@@ -51,22 +71,6 @@ def load_from_db(keyword, country):
         return json.loads(result[0]), result[1]
     return None, None
 
-# 삼성/미국 검색 함수
-def search_samsung_us():
-    error_message, articles = serphouse_search("samsung", "United States")
-    if not error_message and articles:
-        save_to_db("samsung", "United States", articles)
-        return display_results(articles)
-    return "검색 결과가 없습니다."
-
-# DB에서 삼성/미국 결과 불러오기 함수
-def load_samsung_us():
-    results, timestamp = load_from_db("samsung", "United States")
-    if results:
-        return f"저장 시간: {timestamp}\n\n" + display_results(results)
-    return "저장된 결과가 없습니다."
-
-# 결과 표시 함수
 def display_results(articles):
     output = ""
     for idx, article in enumerate(articles, 1):
@@ -77,6 +81,45 @@ def display_results(articles):
         output += f"요약: {article['snippet']}\n\n"
     return output
 
+def search_company(company):
+    error_message, articles = serphouse_search(company, "United States")
+    if not error_message and articles:
+        save_to_db(company, "United States", articles)
+        return display_results(articles)
+    return f"{company}에 대한 검색 결과가 없습니다."
+
+def load_company(company):
+    results, timestamp = load_from_db(company, "United States")
+    if results:
+        return f"### {company} 검색 결과\n저장 시간: {timestamp}\n\n" + display_results(results)
+    return f"{company}에 대한 저장된 결과가 없습니다."
+
+def show_stats():
+    conn = sqlite3.connect("search_results.db")
+    c = conn.cursor()
+    c.execute("""
+        SELECT keyword, COUNT(*) as count, 
+               MAX(timestamp) as last_search
+        FROM searches 
+        WHERE keyword IN ({})
+        GROUP BY keyword 
+        ORDER BY count DESC, last_search DESC
+    """.format(','.join(['?']*len(KOREAN_COMPANIES))), KOREAN_COMPANIES)
+    stats = c.fetchall()
+    conn.close()
+    
+    output = "## 한국 기업 검색 통계\n\n"
+    for keyword, count, last_search in stats:
+        output += f"### {keyword}\n"
+        output += f"- 검색 횟수: {count}회\n"
+        output += f"- 마지막 검색: {last_search}\n\n"
+    return output
+
+
+
+
+
+
     
 ACCESS_TOKEN = os.getenv("HF_TOKEN")
 if not ACCESS_TOKEN:
@@ -1150,30 +1193,67 @@ def continue_writing(history, system_message, max_tokens, temperature, top_p):
     return history
 
 with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css, title="NewsAI 서비스") as iface:
-    init_db()  # DB 초기화
+    init_db()
     
     with gr.Tabs():
         # DB 저장/불러오기 탭
         with gr.Tab("DB 검색"):
-            gr.Markdown("삼성/미국 검색 결과를 DB에 저장하고 불러옵니다.")
+            gr.Markdown("## 한국 주요 기업 미국 뉴스 DB")
+            gr.Markdown("각 기업의 미국 뉴스를 검색하여 DB에 저장하고 불러올 수 있습니다.")
             
-            with gr.Row():
-                search_button = gr.Button("검색: samsung/미국", variant="primary")
-                load_button = gr.Button("출력: samsung/미국", variant="secondary")
-            
-            results_display = gr.Markdown()
+            with gr.Column():
+                for i in range(0, len(KOREAN_COMPANIES), 2):
+                    with gr.Row():
+                        # 첫 번째 열
+                        with gr.Column():
+                            company = KOREAN_COMPANIES[i]
+                            with gr.Group():
+                                gr.Markdown(f"### {company}")
+                                with gr.Row():
+                                    search_btn = gr.Button(f"검색", variant="primary")
+                                    load_btn = gr.Button(f"출력", variant="secondary")
+                                result_display = gr.Markdown()
+                                
+                                search_btn.click(
+                                    fn=lambda c=company: search_company(c),
+                                    outputs=result_display
+                                )
+                                load_btn.click(
+                                    fn=lambda c=company: load_company(c),
+                                    outputs=result_display
+                                )
+                        
+                        # 두 번째 열
+                        if i + 1 < len(KOREAN_COMPANIES):
+                            with gr.Column():
+                                company = KOREAN_COMPANIES[i + 1]
+                                with gr.Group():
+                                    gr.Markdown(f"### {company}")
+                                    with gr.Row():
+                                        search_btn = gr.Button(f"검색", variant="primary")
+                                        load_btn = gr.Button(f"출력", variant="secondary")
+                                    result_display = gr.Markdown()
+                                    
+                                    search_btn.click(
+                                        fn=lambda c=company: search_company(c),
+                                        outputs=result_display
+                                    )
+                                    load_btn.click(
+                                        fn=lambda c=company: load_company(c),
+                                        outputs=result_display
+                                    )
             
-            # 버튼 이벤트 연결
-            search_button.click(
-                fn=search_samsung_us,
-                outputs=results_display
-            )
+            # 전체 검색 통계
+            with gr.Row():
+                stats_btn = gr.Button("전체 검색 통계 보기", variant="secondary")
+                stats_display = gr.Markdown()
             
-            load_button.click(
-                fn=load_samsung_us,
-                outputs=results_display
+            stats_btn.click(
+                fn=show_stats,
+                outputs=stats_display
             )
 
+
         with gr.Tab("국가별"):
             gr.Markdown("검색어를 입력하고 원하는 국가(한국 제외)를를 선택하면, 검색어와 일치하는 24시간 이내 뉴스를 최대 100개 출력합니다.")
             gr.Markdown("국가 선택후 검색어에 '한글'을 입력하면 현지 언어로 번역되어 검색합니다. 예: 'Taiwan' 국가 선택후 '삼성' 입력시 '三星'으로 자동 검색")