Update app.py

app.py

@@ -6,6 +6,14 @@ import pandas as pd
 from typing import List, Tuple
 import json
 from datetime import datetime
+from datasets import load_dataset
+
+try:
+    pharmkg_dataset = load_dataset("vinven7/PharmKG")
+    print("PharmKG 데이터셋 로드 완료")
+except Exception as e:
+    print(f"PharmKG 데이터셋 로드 실패: {e}")
+    pharmkg_dataset = None
 
 # 환경 변수 설정
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -154,19 +162,58 @@ def read_uploaded_file(file):
     except Exception as e:
         return f"❌ 파일 읽기 오류: {str(e)}", "error"
 
+
+
+def get_pharmkg_context(query):
+    """PharmKG 데이터셋에서 관련 정보 검색"""
+    if pharmkg_dataset is None:
+        return ""
+    
+    try:
+        # 데이터셋에서 관련 정보 검색 로직 구현
+        relevant_info = []
+        
+        # 약물 정보 검색
+        for item in pharmkg_dataset['train']:
+            if query.lower() in str(item).lower():
+                relevant_info.append(str(item))
+        
+        if relevant_info:
+            return "\n\nPharmKG 참고 정보:\n" + "\n".join(relevant_info[:3])
+        return ""
+    except Exception as e:
+        print(f"PharmKG 검색 오류: {e}")
+        return ""
+
+
 def chat(message, history, uploaded_file, system_message="", max_tokens=4000, temperature=0.7, top_p=0.9):
     if not message:
         return "", history
 
-    system_prefix = """저는 여러분의 친근하고 지적인 AI 어시스턴트 'GiniGEN Pharm'입니다.. 다음과 같은 원칙으로 소통하겠습니다:
-1. 🤝 친근하고 공감적인 태도로 대화
-2. 💡 명확하고 이해하기 쉬운 설명 제공
+    try:
+        # PharmKG 컨텍스트 추가
+        pharmkg_context = get_pharmkg_context(message)
+        system_message = SYSTEM_PREFIX + system_message + pharmkg_context
+        
+
+    SYSTEM_PREFIX = """저는 약리학 전문 AI 어시스턴트 'GiniGEN Pharm'입니다. PharmKG 데이터베이스를 기반으로 
+다음과 같은 전문성을 가지고 소통하겠습니다:
+
+1. 💊 약물 관련 전문 지식 제공
+2. 🔬 약물 상호작용 및 부작용 정보
+3. 🧬 질병-약물 관계 설명
+4. 📊 임상 데이터 기반 분석
+5. ⚕️ 의약품 안전성 정보
+
+다음 원칙으로 소통하겠습니다:
+1. 🤝 전문적이고 신뢰할 수 있는 정보 제공
+2. 💡 이해하기 쉬운 설명과 함께 과학적 근거 제시
 3. 🎯 질문의 의도를 정확히 파악하여 맞춤형 답변
-4. 📚 필요한 경우 업로드된 파일 내용을 참고하여 구체적인 도움 제공
-5. ✨ 추가적인 통찰과 제안을 통한 가치 있는 대화
+4. ⚠️ 의료적 조언이 아닌 정보 제공 목적임을 명시
+5. ✨ 추가 참고자료 및 연구 데이터 제안
+
+의학적 결정은 반드시 전문 의료진과 상담하시기 바랍니다."""
 
-항상 예의 바르고 친절하게 응답하며, 필요한 경우 구체적인 예시나 설명을 추가하여 
-이해를 돕겠습니다."""
 
     try:
         # 파일 업로드 처리