Update app.py

app.py

@@ -52,6 +52,60 @@ Remember not add any description, just return the code only.
 # config.py에서 DEMO_LIST만 import
 from config import DEMO_LIST
 
+import sqlite3
+from datetime import datetime
+
+
+def init_db():
+    try:
+        conn = sqlite3.connect('chat_history.db')
+        c = conn.cursor()
+        c.execute('''CREATE TABLE IF NOT EXISTS sessions
+                     (session_id TEXT PRIMARY KEY, 
+                      created_at TIMESTAMP)''')
+        c.execute('''CREATE TABLE IF NOT EXISTS chat_history
+                     (id INTEGER PRIMARY KEY AUTOINCREMENT,
+                      session_id TEXT,
+                      prompt TEXT,
+                      response TEXT,
+                      timestamp TIMESTAMP,
+                      FOREIGN KEY (session_id) REFERENCES sessions(session_id))''')
+        conn.commit()
+    except sqlite3.Error as e:
+        print(f"Database error: {e}")
+    finally:
+        if conn:
+            conn.close()
+
+# 새로운 세션 생성
+def create_session():
+    session_id = datetime.now().strftime("%Y%m%d_%H%M%S")
+    conn = sqlite3.connect('chat_history.db')
+    c = conn.cursor()
+    c.execute("INSERT INTO sessions VALUES (?, ?)", (session_id, datetime.now()))
+    conn.commit()
+    conn.close()
+    return session_id
+
+# 대화 내용 저장
+def save_chat(session_id, prompt, response):
+    conn = sqlite3.connect('chat_history.db')
+    c = conn.cursor()
+    c.execute("INSERT INTO chat_history (session_id, prompt, response, timestamp) VALUES (?, ?, ?, ?)",
+              (session_id, prompt, response, datetime.now()))
+    conn.commit()
+    conn.close()
+
+# 세션별 히스토리 조회
+def get_session_history(session_id):
+    conn = sqlite3.connect('chat_history.db')
+    c = conn.cursor()
+    c.execute("SELECT prompt, response, timestamp FROM chat_history WHERE session_id = ? ORDER BY timestamp",
+              (session_id,))
+    history = c.fetchall()
+    conn.close()
+    return history
+    
 def get_image_base64(image_path):
     with open(image_path, "rb") as image_file:
         encoded_string = base64.b64encode(image_file.read()).decode()
@@ -136,9 +190,17 @@ async def try_openai_api(openai_messages):
         raise e
 
 
-async def generation_code(query: Optional[str], _setting: Dict[str, str], _history: Optional[History]):
-    if not query or query.strip() == '':
-        query = random.choice(DEMO_LIST)['description']
+
+class Demo:
+    def __init__(self):
+        init_db()
+        self.current_session = create_session()
+
+    async def generation_code(self, query: Optional[str], _setting: Dict[str, str], _history: Optional[History]):
+        if not query or query.strip() == '':
+            query = random.choice(DEMO_LIST)['description']
+
+
     
     if _history is None:
         _history = []
@@ -202,8 +264,12 @@ async def generation_code(query: Optional[str], _setting: Dict[str, str], _histo
                 await asyncio.sleep(0)  # UI 업데이트를 위한 양보
             collected_content = content
 
+
         if collected_content:
-            _history = messages_to_history([
+        # 채팅 내용 저장
+            save_chat(self.current_session, query, collected_content)    
+            _history = messages_to_history([...])
+        
                 {'role': Role.SYSTEM, 'content': system_message}
             ] + claude_messages + [{
                 'role': Role.ASSISTANT,
@@ -220,11 +286,15 @@ async def generation_code(query: Optional[str], _setting: Dict[str, str], _histo
         else:
             raise ValueError("No content was generated from either API")
     
-    except Exception as e:
-        print(f"Error details: {str(e)}")
-        raise ValueError(f'Error calling APIs: {str(e)}')
-        
 
+            
+        except Exception as e:
+            print(f"Error details: {str(e)}")
+            raise ValueError(f'Error calling APIs: {str(e)}')
+
+    def clear_history(self):
+        self.current_session = create_session()
+        return []
 
 
 def remove_code_block(text):
@@ -238,8 +308,6 @@ def remove_code_block(text):
 def history_render(history: History):
     return gr.update(open=True), history
 
-def clear_history():
-    return []
 
 def send_to_sandbox(code):
     encoded_html = base64.b64encode(code.encode('utf-8')).decode('utf-8')
@@ -249,7 +317,19 @@ def send_to_sandbox(code):
 
 theme = gr.themes.Soft()
 
-
+def clear_expired_sessions():
+    """30일 이상 된 세션 삭제"""
+    conn = sqlite3.connect('chat_history.db')
+    try:
+        c = conn.cursor()
+        c.execute("""DELETE FROM chat_history WHERE session_id IN 
+                    (SELECT session_id FROM sessions 
+                     WHERE created_at < datetime('now', '-30 days'))""")
+        c.execute("DELETE FROM sessions WHERE created_at < datetime('now', '-30 days')")
+        conn.commit()
+    finally:
+        conn.close()
+        
 with gr.Blocks(css_paths="app.css",theme=theme) as demo:
     history = gr.State([])
     setting = gr.State({
@@ -401,7 +481,15 @@ with gr.Blocks(css_paths="app.css",theme=theme) as demo:
                         with antd.Flex(gap="small", elem_classes="setting-buttons"):
                             codeBtn = antd.Button("🧑‍💻 view code", type="default")
                             historyBtn = antd.Button("📜 history", type="default")
-                            
+                            sessionBtn = antd.Button("📚 sessions", type="default")  # 세션 버튼 추가
+
+# 세션 버튼 이벤트 핸들러
+                        sessionBtn.click(
+                            lambda: (gr.update(open=True), update_session_list()),
+                            inputs=[],
+                            outputs=[session_drawer, session_list]
+                        )
+                        
                         gr.HTML('<div class="render_header"><span class="header_btn"></span><span class="header_btn"></span><span class="header_btn"></span></div>')
                         with antd.Tabs(active_key="empty", render_tab_bar="() => null") as state_tab:
                             with antd.Tabs.Item(key="empty"):
@@ -446,16 +534,44 @@ with gr.Blocks(css_paths="app.css",theme=theme) as demo:
             history_drawer.close(lambda: gr.update(
                 open=False), inputs=[], outputs=[history_drawer])
 
-# 수정된 코드
+            
             btn.click(
-                generation_code,  # 비동기 함수를 직접 전달
+                demo_instance.generation_code,  # Demo 인스턴스의 메서드 사용
                 inputs=[input, setting, history],
                 outputs=[code_output, history, sandbox, state_tab, code_drawer]
             )
 
+            clear_btn.click(
+                demo_instance.clear_history,  # Demo 인스턴스의 메서드 사용
+                inputs=[],
+                outputs=[history]
+            )
+
 
-            
-            clear_btn.click(clear_history, inputs=[], outputs=[history])
 
+
+    with antd.Drawer(open=False, title="Session History", placement="left", width="900px") as session_drawer:
+        session_list = gr.Dropdown(label="Select Session", choices=[])
+        session_history = legacy.Chatbot(show_label=False, height=960)
+        
+    def update_session_list():
+        conn = sqlite3.connect('chat_history.db')
+        c = conn.cursor()
+        c.execute("SELECT session_id FROM sessions ORDER BY created_at DESC")
+        sessions = [row[0] for row in c.fetchall()]
+        conn.close()
+        return gr.update(choices=sessions)
+        
+    def load_session_history(session_id):
+        history = get_session_history(session_id)
+        return [[p, r] for p, r, _ in history]
+        
+    session_list.change(load_session_history, 
+                       inputs=[session_list], 
+                       outputs=[session_history])
+    
 if __name__ == "__main__":
+    init_db()
+    clear_expired_sessions()  # 만료된 세션 정리
+    demo_instance = Demo()
     demo.queue(default_concurrency_limit=20).launch(ssr_mode=False)