Update app.py

app.py

@@ -10,7 +10,74 @@ 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 pathlib
+
+# DB 초기화 함수
+def init_db():
+    db_path = pathlib.Path("search_results.db")
+    conn = sqlite3.connect(db_path)
+    c = conn.cursor()
+    c.execute('''CREATE TABLE IF NOT EXISTS searches
+                 (id INTEGER PRIMARY KEY AUTOINCREMENT,
+                  keyword TEXT,
+                  country TEXT,
+                  results TEXT,
+                  timestamp DATETIME DEFAULT CURRENT_TIMESTAMP)''')
+    conn.commit()
+    conn.close()
+
+# 검색 결과 저장 함수
+def save_to_db(keyword, country, results):
+    conn = sqlite3.connect("search_results.db")
+    c = conn.cursor()
+    c.execute("INSERT INTO searches (keyword, country, results) VALUES (?, ?, ?)",
+              (keyword, country, json.dumps(results)))
+    conn.commit()
+    conn.close()
+
+# DB에서 검색 결과 불러오기 함수
+def load_from_db(keyword, country):
+    conn = sqlite3.connect("search_results.db")
+    c = conn.cursor()
+    c.execute("SELECT results, timestamp FROM searches WHERE keyword=? AND country=? ORDER BY timestamp DESC LIMIT 1",
+              (keyword, country))
+    result = c.fetchone()
+    conn.close()
+    if result:
+        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):
+        output += f"### {idx}. {article['title']}\n"
+        output += f"출처: {article['channel']}\n"
+        output += f"시간: {article['time']}\n"
+        output += f"링크: {article['link']}\n"
+        output += f"요약: {article['snippet']}\n\n"
+    return output
 
+    
 ACCESS_TOKEN = os.getenv("HF_TOKEN")
 if not ACCESS_TOKEN:
     raise ValueError("HF_TOKEN environment variable is not set")
@@ -1081,10 +1148,32 @@ def continue_writing(history, system_message, max_tokens, temperature, top_p):
         yield new_history
 
     return history
-    
+
 with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css, title="NewsAI 서비스") as iface:
+    init_db()  # DB 초기화
+    
     with gr.Tabs():
-        # 국가별 탭
+        # DB 저장/불러오기 탭
+        with gr.Tab("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()
+            
+            # 버튼 이벤트 연결
+            search_button.click(
+                fn=search_samsung_us,
+                outputs=results_display
+            )
+            
+            load_button.click(
+                fn=load_samsung_us,
+                outputs=results_display
+            )
+
         with gr.Tab("국가별"):
             gr.Markdown("검색어를 입력하고 원하는 국가(한국 제외)를를 선택하면, 검색어와 일치하는 24시간 이내 뉴스를 최대 100개 출력합니다.")
             gr.Markdown("국가 선택후 검색어에 '한글'을 입력하면 현지 언어로 번역되어 검색합니다. 예: 'Taiwan' 국가 선택후 '삼성' 입력시 '三星'으로 자동 검색")