Update app.py

app.py

@@ -16,15 +16,31 @@ plt.rc('font', family='NanumBarunGothic')
 
 # 전역 변수로 분석 결과 저장
 global_recommendations = None
+global_csv_string = None
 
 # CSV 문자열 생성 함수
-def recommendations_to_csv(recommendations):
+def create_csv_string(recommendations):
     output = io.StringIO()
     writer = csv.writer(output)
     writer.writerow(["Employee ID", "Employee Name", "Recommended Programs"])
-    writer.writerows(recommendations)
+    for rec in recommendations:
+        writer.writerow(rec)
     return output.getvalue()
 
+# 차트 생성 함수
+def create_chart(G):
+    plt.figure(figsize=(10, 8))
+    pos = nx.spring_layout(G)
+    nx.draw(G, pos, with_labels=True, node_color='lightblue', node_size=3000, font_size=10, font_weight='bold', edge_color='gray')
+    plt.title("직원과 프로그램 간의 관계", fontsize=14, fontweight='bold')
+    plt.tight_layout()
+    
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png')
+    buf.seek(0)
+    plt.close()
+    return buf
+
 # 열 매칭 함수
 def auto_match_columns(df, required_cols):
     matched_cols = {}
@@ -58,6 +74,7 @@ def validate_and_get_columns(employee_df, program_df):
 # 직원 데이터를 분석하여 교육 프로그램을 추천하고, 테이블과 그래프를 생성하는 함수
 def hybrid_rag(employee_file, program_file):
     global global_recommendations
+    global global_csv_string
     
     # 1. VectorRAG: KoSentence-BERT를 이용한 유사도 계산
     employee_df = pd.read_csv(employee_file.name)
@@ -105,19 +122,16 @@ def hybrid_rag(employee_file, program_file):
             if similarities[i][j] > 0.5:
                 G.add_edge(employee[employee_cols['employee_name']], program[program_cols['program_name']])
 
-    plt.figure(figsize=(10, 8))
-    pos = nx.spring_layout(G)
-    nx.draw(G, pos, with_labels=True, node_color='lightblue', node_size=3000, font_size=10, font_weight='bold', edge_color='gray')
-    plt.title("직원과 프로그램 간의 관계", fontsize=14, fontweight='bold')
-    plt.tight_layout()
-
-    # CSV 문자열로 추천 결과 반환
-    csv_output = recommendations_to_csv(recommendation_rows)
+    # 차트 생성
+    chart_buffer = create_chart(G)
+    
+    # CSV 문자열 생성
+    global_csv_string = create_csv_string(recommendation_rows)
     
     # 결과 테이블 데이터프레임 생성
     result_df = pd.DataFrame(recommendation_rows, columns=["Employee ID", "Employee Name", "Recommended Programs"])
     
-    return result_df, plt.gcf(), csv_output
+    return result_df, chart_buffer, gr.File.update(visible=True)
 
 # 채팅 응답 함수
 def chat_response(message, history):
@@ -131,6 +145,13 @@ def chat_response(message, history):
     
     return "죄송합니다. 해당 직원의 정보를 찾을 수 없습니다. 다른 직원 이름을 입력해주세요."
 
+# CSV 다운로드 함수
+def download_csv():
+    global global_csv_string
+    if global_csv_string is None:
+        return gr.File.update(visible=False)
+    return gr.File.update(value=global_csv_string, visible=True, filename="recommendations.csv")
+
 # Gradio 블록
 with gr.Blocks(css=".gradio-button {background-color: #007bff; color: white;} .gradio-textbox {border-color: #6c757d;}") as demo:
     gr.Markdown("<h1 style='text-align: center; color: #2c3e50;'>💼 HybridRAG 시스템</h1>")
@@ -142,11 +163,11 @@ with gr.Blocks(css=".gradio-button {background-color: #007bff; color: white;} .g
             program_file = gr.File(label="교육 프로그램 데이터 업로드", interactive=True)
             analyze_button = gr.Button("분석 시작", elem_classes="gradio-button")
             output_table = gr.DataFrame(label="분석 결과 (테이블)")
-            csv_download = gr.File(label="추천 결과 다운로드")
+            csv_download = gr.File(label="추천 결과 다운로드", visible=False)
 
         with gr.Column(scale=2, min_width=500):
             gr.Markdown("<h3 style='color: #34495e;'>2. 분석 결과 및 시각화</h3>")
-            chart_output = gr.Plot(label="시각화 차트")
+            chart_output = gr.Image(label="시각화 차트")
 
     gr.Markdown("<h3 style='color: #34495e;'>3. 직원별 추천 프로그램 확인</h3>")
     chatbot = gr.Chatbot()
@@ -156,6 +177,9 @@ with gr.Blocks(css=".gradio-button {background-color: #007bff; color: white;} .g
     # 분석 버튼 클릭 시 테이블, 차트, 파일 다운로드를 업데이트
     analyze_button.click(hybrid_rag, inputs=[employee_file, program_file], outputs=[output_table, chart_output, csv_download])
 
+    # CSV 다운로드 버튼
+    csv_download.click(download_csv, inputs=[], outputs=[csv_download])
+
     # 채팅 기능
     msg.submit(chat_response, [msg, chatbot], [chatbot])
     clear.click(lambda: None, None, chatbot, queue=False)