Update app.py

app.py

@@ -7,6 +7,7 @@ import matplotlib.pyplot as plt
 import csv
 import io
 import matplotlib.font_manager as fm
+from neo4j import GraphDatabase
 
 # 한국어 처리를 위한 KoSentence-BERT 모델 로드
 model = SentenceTransformer('jhgan/ko-sbert-sts')
@@ -16,6 +17,30 @@ font_path = "/root/fonts/NanumBarunGothic.ttf"  # Hugging Face 루트에 저장
 fontprop = fm.FontProperties(fname=font_path)
 plt.rc('font', family=fontprop.get_name())
 
+# Neo4j 데이터베이스 연결 클래스
+class Neo4jConnection:
+    def __init__(self, uri, user, pwd):
+        self.driver = GraphDatabase.driver(uri, auth=(user, pwd))
+
+    def close(self):
+        self.driver.close()
+
+    def query(self, query, parameters=None, db=None):
+        session = None
+        response = None
+        try:
+            session = self.driver.session(database=db) if db else self.driver.session() 
+            response = list(session.run(query, parameters))
+        except Exception as e:
+            print("Query failed:", e)
+        finally:
+            if session:
+                session.close()
+        return response
+
+# Neo4j 연결 설정
+conn = Neo4jConnection(uri="bolt://localhost:7687", user="neo4j", pwd="your_password")
+
 # 추천 결과를 실제 파일로 저장하는 함수
 def save_recommendations_to_file(recommendations):
     file_path = "recommendations.csv"
@@ -60,7 +85,8 @@ def validate_and_get_columns(employee_df, program_df):
     return None, employee_cols, program_cols
 
 # 직원 데이터를 분석하여 교육 프로그램을 추천하고, 테이블과 그래프를 생성하는 함수
-def analyze_data(employee_file, program_file):
+def hybrid_rag(employee_file, program_file):
+    # 1. VectorRAG: KoSentence-BERT를 이용한 유사도 계산
     employee_df = pd.read_csv(employee_file.name)
     program_df = pd.read_csv(program_file.name)
 
@@ -92,6 +118,19 @@ def analyze_data(employee_file, program_file):
 
         recommendations.append(recommendation)
 
+    # 2. GraphRAG: Neo4j에서 프로그램 추천을 가져옴
+    query = """
+    MATCH (e:Employee)-[:HAS_SKILL]->(p:Program)
+    RETURN e.name AS employee_name, p.name AS program_name, p.duration AS duration
+    """
+    graph_rag_results = conn.query(query)
+
+    # GraphRAG 결과 추가
+    for record in graph_rag_results:
+        for row in recommendation_rows:
+            if record['employee_name'] == row[1]:
+                row[2] += f", {record['program_name']} (GraphRAG)"
+
     G = nx.Graph()
     for employee in employee_df[employee_cols['employee_name']]:
         G.add_node(employee, type='employee')
@@ -136,7 +175,7 @@ with gr.Blocks(css=".gradio-button {background-color: #007bff; color: white;} .g
             chart_output = gr.Plot(label="시각화 차트")
 
     # 분석 버튼 클릭 시 테이블, 차트, 파일 다운로드를 업데이트
-    analyze_button.click(analyze_data, inputs=[employee_file, program_file], outputs=[output_table, chart_output, csv_download])
+    analyze_button.click(hybrid_rag, inputs=[employee_file, program_file], outputs=[output_table, chart_output, csv_download])
 
 # Gradio 인터페이스 실행
 demo.launch()