Update app.py

app.py

@@ -18,12 +18,31 @@ models = {
 def get_client(model):
     return InferenceClient(model=model, token=hf_token)
 
+# 히스토리 관리 클래스
+class PromptHistory:
+    def __init__(self):
+        self.history = []
+
+    def add_entry(self, prompt, response, model, system_message):
+        self.history.append({
+            "prompt": prompt,
+            "response": response,
+            "model": model,
+            "system_message": system_message
+        })
+
+    def get_history(self):
+        return self.history
+
+# 히스토리 인스턴스 생성
+prompt_history = PromptHistory()
+
 # 응답 생성 함수
 def respond(prompt, system_message, max_tokens, temperature, top_p, selected_model):
     stop_event.clear()
     client = get_client(selected_model)
     
-    # 프롬프트 설정 - 시스템 메시지를 자유롭게 설정 가능
+    # 프롬프트 설정
     messages = [
         {"role": "system", "content": system_message},
         {"role": "user", "content": prompt}
@@ -46,6 +65,9 @@ def respond(prompt, system_message, max_tokens, temperature, top_p, selected_mod
             if chunk:
                 response += chunk
                 total_tokens_used += len(chunk.split())  # 청크당 사용된 토큰 수 추산
+
+        # 히스토리에 저장
+        prompt_history.add_entry(prompt, response, selected_model, system_message)
         
         return response
     
@@ -61,6 +83,12 @@ def continue_writing(chatbot, system_message, max_tokens, temperature, top_p, se
 def stop_generation():
     stop_event.set()
 
+# 히스토리 확인 함수
+def view_history():
+    history = prompt_history.get_history()
+    history_str = "\n\n".join([f"모델: {entry['model']}\n프롬프트: {entry['prompt']}\n응답: {entry['response']}\n시스템 메시지: {entry['system_message']}" for entry in history])
+    return history_str if history_str else "히스토리가 없습니다."
+
 # Gradio UI 구성
 with gr.Blocks() as demo:
     gr.Markdown("# 프롬프트 최적화 Playground")
@@ -99,6 +127,7 @@ with gr.Blocks() as demo:
                 continue_btn = gr.Button("계속 작성")
                 stop = gr.Button("🛑 생성 중단")
                 clear = gr.Button("🗑️ 대화 내역 지우기")
+                view_history_btn = gr.Button("🗂️ 히스토리 보기")
 
     # Event handlers
     send.click(respond, inputs=[prompt, system_message, max_tokens, temperature, top_p, model], outputs=[chatbot])
@@ -108,6 +137,7 @@ with gr.Blocks() as demo:
                        outputs=[chatbot])
     stop.click(stop_generation)
     clear.click(lambda: None, outputs=[chatbot])
+    view_history_btn.click(view_history, outputs=[gr.Textbox(label="히스토리")])
 
 # UI 실행
 demo.launch()