Update app.py

app.py

@@ -1,298 +1,313 @@
 # =========================================================
-# KB 금융 RAG 챗봇 (Local Self-Contained Version)
-# =========================================================
-# 이 코드는 서버나 클라우드 DB 없이, 사용자가 직접 PDF를 업로드하여
-# 로컬에서 지식 베이스를 구축하고 질문할 수 있는 구조입니다.
-# Groq(LLM), Google(Voice/Translate) API를 사용하여 무료로 동작합니다.
+# KB AI Challenge - Professional RAG System (Multilingual)
 # =========================================================
 
 import os
 import sys
 import numpy as np
 import traceback
-import fitz  # PyMuPDF (PDF 처리)
+import fitz  # PyMuPDF
 from typing import List
 
 # --- 라이브러리 임포트 ---
 import gradio as gr
 import speech_recognition as sr
+from dotenv import load_dotenv
+
+# .env 로드
+load_dotenv()
+
 from deep_translator import GoogleTranslator
 from sentence_transformers import SentenceTransformer
 from groq import Groq
 from qdrant_client import QdrantClient
 from qdrant_client.models import Distance, VectorParams, PointStruct
+
 try:
     from langchain.text_splitter import RecursiveCharacterTextSplitter
 except ImportError:
-    # langchain 0.2.0 이상에서 구조가 변경된 경우
     from langchain_text_splitters import RecursiveCharacterTextSplitter
 
 # =========================================================
 # 1. 설정 및 초기화
 # =========================================================
 
-# Groq API 키 (필수)
 GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "your_groq_api_key_here")
-if not GROQ_API_KEY or GROQ_API_KEY == "your_groq_api_key_here":
-    print("⚠️ GROQ_API_KEY가 설정되지 않았습니다. RAG 기능 사용 시 오류가 발생할 수 있습니다.")
-
-# 모델 설정
 EMBEDDING_MODEL_NAME = "jhgan/ko-sroberta-multitask"
 GROQ_MODEL_NAME = "llama-3.3-70b-versatile"
 COLLECTION_NAME = "local_kb"
 
-print("🛠️ 모델 및 클라이언트 초기화 중...")
+print("🛠️ 시스템 초기화 중... (System Init)")
 
-# 1. 임베딩 모델 로드 (로컬 실행)
+# 모델 로드
 embedding_model = SentenceTransformer(EMBEDDING_MODEL_NAME)
 embedding_model.max_seq_length = 512
 
-# 2. Qdrant 클라이언트 (로컬 메모리 DB - 프로그램 종료 시 데이터 삭제됨)
-# 영구 저장을 원하면 path="./local_qdrant_db" 로 변경하세요.
-# 여기서는 포트폴리오용 데모를 위해 매번 깨끗한 상태인 ':memory:'를 기본으로 합니다.
+# Qdrant (메모리)
 qdrant_client = QdrantClient(":memory:")
-
-# 컬렉션 생성 (이미 존재하면 삭제 후 재생성)
 try:
     qdrant_client.recreate_collection(
         collection_name=COLLECTION_NAME,
         vectors_config=VectorParams(size=768, distance=Distance.COSINE),
     )
-    print(f"✅ 로컬 Qdrant 컬렉션 '{COLLECTION_NAME}' 생성 완료.")
+    print(f"✅ Qdrant Collection Ready.")
 except Exception as e:
-    print(f"❌ Qdrant 컬렉션 생성 실패: {e}")
+    print(f"❌ Qdrant Error: {e}")
 
-# 3. Groq 클라이언트
-try:
-    groq_client = Groq(api_key=GROQ_API_KEY)
-except Exception as e:
-    groq_client = None
-    print(f"❌ Groq 클라이언트 초기화 실패: {e}")
+# Groq Init
+groq_client = None
+if GROQ_API_KEY and GROQ_API_KEY != "your_groq_api_key_here":
+    try:
+        groq_client = Groq(api_key=GROQ_API_KEY)
+    except Exception as e:
+        print(f"❌ Groq Error: {e}")
+else:
+    print("⚠️ Groq API Key Missing.")
 
-#전역 변수: 문서 ID 카운터
 doc_id_counter = 0
 
-print("✅ 모든 시스템 준비 완료!")
+print("✅ System Ready.")
+
+
+# =========================================================
+# 2. 다국어 지원 로직 (Translation & STT)
+# =========================================================
+
+LANG_MAP = {
+    "한국어 (Korean)": {"code": "ko", "stt": "ko-KR"},
+    "English (영어)": {"code": "en", "stt": "en-US"},
+    "日本語 (Japanese)": {"code": "ja", "stt": "ja-JP"},
+    "中文 (Chinese)": {"code": "zh-CN", "stt": "zh-CN"}
+}
+
+def translate_text(text, target_lang_code):
+    try:
+        if target_lang_code == "ko": return text
+        return GoogleTranslator(source='auto', target=target_lang_code).translate(text)
+    except:
+        return text
 
+def translate_to_korean(text):
+    try:
+        return GoogleTranslator(source='auto', target='ko').translate(text)
+    except:
+        return text
 
 # =========================================================
-# 2. 문서 처리 및 RAG 핵심 로직
+# 3. 핵심 로직 (RAG Pipeline)
 # =========================================================
 
 def process_uploaded_files(files):
-    """PDF 파일을 읽어 텍스트를 추출하고 벡터 DB에 저장"""
+    """PDF 처리 및 임베딩"""
     global doc_id_counter
-    
-    if not files:
-        return "파일이 업로드되지 않았습니다."
+    if not files: return "파일이 선택되지 않았습니다."
     
     total_chunks = 0
     status_msg = ""
-    
-    # 텍스트 분리기 설정
-    text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=500,
-        chunk_overlap=50,
-        length_function=len,
-    )
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50, length_function=len)
 
     for file in files:
         try:
-            # Gradio 버전/설정에 따라 file이 문자열(경로)일 수도 있고 객체일 수도 있음
             file_path = file.name if hasattr(file, 'name') else file
-            
-            # 1. PDF 텍스트 추출
             doc = fitz.open(file_path)
             file_text = ""
-            for page in doc:
-                file_text += page.get_text()
+            for page in doc: file_text += page.get_text()
             
             if not file_text.strip():
-                status_msg += f"⚠️ {os.path.basename(file_path)}: 텍스트 추출 실패 (이미지 PDF일 수 있음)\n"
+                status_msg += f"⚠️ {os.path.basename(file_path)}: 텍스트 없음.\n"
                 continue
                 
-            # 2. 텍스트 분할 (Chunking)
             chunks = text_splitter.split_text(file_text)
-            
-            # 3. 임베딩 및 저장
             points = []
             for i, chunk in enumerate(chunks):
                 vector = embedding_model.encode(chunk).tolist()
-                
-                payload = {
-                    "filename": os.path.basename(file_path),
-                    "text": chunk,
-                    "chunk_id": i
-                }
-                
+                payload = {"filename": os.path.basename(file_path), "text": chunk}
                 points.append(PointStruct(id=doc_id_counter, vector=vector, payload=payload))
                 doc_id_counter += 1
             
-            # Qdrant에 저장
             if points:
-                qdrant_client.upsert(
-                    collection_name=COLLECTION_NAME,
-                    points=points
-                )
+                qdrant_client.upsert(collection_name=COLLECTION_NAME, points=points)
                 total_chunks += len(points)
-                status_msg += f"✅ {os.path.basename(file_path)}: {len(points)}개 지식 저장 완료.\n"
+                status_msg += f"✅ {os.path.basename(file_path)} ({len(points)} 개 저장됨)\n"
             
         except Exception as e:
-            traceback.print_exc()
-            file_name_debug = getattr(file, 'name', str(file))
-            status_msg += f"❌ {os.path.basename(file_name_debug)} 처리 중 오류: {str(e)}\n"
+            status_msg += f"❌ 오류: {os.path.basename(file_path)} - {str(e)}\n"
             
-    print(f"DEBUG: 총 저장된 청크 수: {total_chunks}")
-    if total_chunks == 0:
-        return status_msg + "\n(저장된 데이터가 없습니다. PDF가 비어있거나 이미지일 수 있습니다.)"
-            
-    return f"처리 완료! 총 {total_chunks}개의 지식 조각이 저장되었습니다.\n\n{status_msg}"
+    return f"총 {total_chunks}개 데이터 처리 완료.\n\n{status_msg}"
 
 def search_knowledge_base(query, top_k=5):
-    """로컬 Qdrant에서 관련 문서 검색"""
     try:
         query_vector = embedding_model.encode(query).tolist()
-        # qdrant-client 버전에 따라 .search()가 없거나 다르게 동작할 수 있어 .query_points() 사용
-        search_result = qdrant_client.query_points(
-            collection_name=COLLECTION_NAME,
-            query=query_vector,
-            limit=top_k,
-            with_payload=True
+        res = qdrant_client.query_points(
+            collection_name=COLLECTION_NAME, query=query_vector, limit=top_k, with_payload=True
         )
-        return search_result.points
-    except Exception as e:
-        print(f"검색 오류: {e}")
+        return res.points
+    except:
         return []
 
 def generate_answer_groq(query, context_text):
-    """Groq API를 사용하여 답변 생성"""
-    if not groq_client: 
-        return "Groq API 설정 오류"
-        
+    if not groq_client: return "API 키가 필요합니다."
+    
     system_prompt = """
-    당신은 친절하고 전문적인 금융 AI 어시스턴트입니다.
-    반드시 아래 제공된 [참고자료]만을 바탕으로 질문에 답변하세요.
-    참고자료에 내용이 없다면 솔직하게 모른다고 대답하세요.
-    출처(파일이름)를 답변 끝에 명시해주세요.
+    당신은 KB 금융그룹의 전문 AI 어시스턴트입니다.
+    제공된 [문맥]에 기반하여 질문에 대해 정확하고 전문적인 답변을 작성하세요.
+    모르는 내용은 모른다고 답하고, 추측하지 마세요.
+    답변은 한국어로 작성하세요.
     """
-    
-    user_prompt = f"질문: {query}\n\n[참고자료]\n{context_text}"
-    
+    user_prompt = f"질문: {query}\n\n[문맥]\n{context_text}"
     try:
         response = groq_client.chat.completions.create(
-            messages=[
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": user_prompt},
-            ],
-            model=GROQ_MODEL_NAME,
-            temperature=0.1,
+            messages=[{"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt}],
+            model=GROQ_MODEL_NAME, temperature=0.1
         )
         return response.choices[0].message.content
     except Exception as e:
-        return f"Groq 생성 오류: {e}"
+        return f"응답 생성 오류: {e}"
 
-# RAG 파이프라인 (통합)
-def run_rag_pipeline(text_input, detected_lang='ko'):
-    if not text_input:
-        return "", "", "", ""
-        
-    # 1. 질문 번역 (필요시)
-    korean_query = text_input
-    if detected_lang != 'ko':
-        try:
-            korean_query = GoogleTranslator(source='auto', target='ko').translate(text_input)
-        except: pass
-
-    # 2. 문서 검색
-    hits = search_knowledge_base(korean_query)
+def run_rag_chat(message, history, lang_selection):
+    if not message: return "", history, ""
     
-    if not hits:
-        return korean_query, "저장된 지식이 부족하여 답변할 수 없습니다. PDF를 먼저 업로드해주세요.", "", "참고 문서 없음"
-
-    # 3. 컨텍스트 구성
-    context_text = ""
-    references = []
-    for hit in hits:
-        context_text += f"{hit.payload['text']}\n\n"
-        references.append(f"- {hit.payload['filename']} (유사도: {hit.score:.2f})")
+    target_lang = LANG_MAP[lang_selection]["code"]
     
-    ref_str = "\n".join(references)
+    # 1. 입력 번역 (Target -> Korean)
+    korean_query = message
+    if target_lang != "ko":
+        korean_query = translate_to_korean(message)
     
-    # 4. 답변 생성
-    korean_answer = generate_answer_groq(korean_query, context_text)
+    # 2. 검색 & 답변 생성 (Korean)
+    hits = search_knowledge_base(korean_query)
+    if not hits:
+        bot_response_ko = "죄송합니다. 관련 정보를 찾을 수 없습니다."
+        reference_text = "참고 문서 없음"
+    else:
+        context_text = "\n\n".join([h.payload['text'] for h in hits])
+        # 중복 제거 및 그룹화 (File grouping)
+        ref_data = {}
+        for h in hits:
+            fname = h.payload['filename']
+            if fname not in ref_data:
+                ref_data[fname] = []
+            ref_data[fname].append(h.score)
+            
+        refs = []
+        for fname, scores in ref_data.items():
+            refs.append(f"- {fname} (관련 내용 {len(scores)}건, 최고 유사도: {max(scores):.2f})")
+        reference_text = "\n".join(refs)
+        bot_response_ko = generate_answer_groq(korean_query, context_text)
     
-    # 5. 답변 번역 (필요시)
-    final_answer = korean_answer
-    if detected_lang != 'ko':
-        try:
-            final_answer = GoogleTranslator(source='ko', target=detected_lang).translate(korean_answer)
-        except: pass
-        
-    return korean_query, korean_answer, final_answer, ref_str
-
-
-# =========================================================
-# 3. 음성 및 UI 헬퍼 함수
-# =========================================================
+    # 3. 답변 번역 (Korean -> Target)
+    final_response = bot_response_ko
+    if target_lang != "ko":
+        translated_response = translate_text(bot_response_ko, target_lang)
+        final_response = f"{translated_response}\n\n---\n[한국어 원문]\n{bot_response_ko}"
+    
+    # 히스토리에 추가 (Messages Format for Gradio 6.x)
+    new_history = history + [
+        {"role": "user", "content": message},
+        {"role": "assistant", "content": final_response}
+    ]
+    return "", new_history, reference_text
 
-def voice_to_text(audio_input):
-    """음성 인식 (Google API)"""
-    if audio_input is None: return "음성 입력 없음", None
+def voice_to_text_chat(audio, history, lang_selection):
+    if audio is None: return "", history, "음성 입력 없음"
+    
+    stt_lang = LANG_MAP[lang_selection]["stt"]
     
     try:
-        sample_rate, audio_numpy = audio_input
+        sample_rate, audio_numpy = audio
         if audio_numpy.dtype == np.float32:
             audio_numpy = (audio_numpy * 32767).astype(np.int16)
         if len(audio_numpy.shape) > 1:
             audio_numpy = audio_numpy.mean(axis=1).astype(np.int16)
-            
         audio_data = sr.AudioData(audio_numpy.tobytes(), sample_rate, 2)
         r = sr.Recognizer()
-        text = r.recognize_google(audio_data, language='ko-KR')
-        return text, 'ko'
+        
+        # 선택된 언어로 인식
+        text = r.recognize_google(audio_data, language=stt_lang)
+        
+        # 채팅 함수 호출
+        return run_rag_chat(text, history, lang_selection)
+        
     except sr.UnknownValueError:
-        return "인식 실패 (다시 말해주세요)", None
-    except Exception as e:
-        return f"오류: {e}", None
+        return "", history, "음성을 이해할 수 없습니다."
+    except Exception as e: 
+        return "", history, f"오류: {e}"
 
 # =========================================================
-# 4. Gradio UI 구성
+# 4. UI Layout (Clean Professional Korean)
 # =========================================================
 
-with gr.Blocks(theme=gr.themes.Soft(), title="KB AI Challenge") as demo:
-    gr.Markdown("# KB AI Challenge")
-    gr.Markdown("서버 없이 로컬에서 동작하는 **개인용 RAG 시스템**입니다. PDF를 업로드하면 즉시 학습하여 답변합니다.")
-    
-    with gr.Accordion("📂 1. 지식 베이스 구축 (파일 업로드)", open=True):
-        with gr.Row():
-            file_input = gr.File(label="PDF 업로드 (여러 개 가능)", file_count="multiple", file_types=[".pdf"])
-            upload_btn = gr.Button("저장하기", variant="primary")
-        upload_status = gr.Textbox(label="처리 상태", interactive=False)
-        
-    gr.Markdown("---")
-    gr.Markdown("### 🎤 2. AI와 대화하기")
+theme = gr.themes.Soft(
+    primary_hue="amber",
+    neutral_hue="slate",
+    font=[gr.themes.GoogleFont("Noto Sans KR"), "sans-serif"]
+)
 
+css = """
+footer {visibility: hidden !important;}
+.gradio-container {min-height: 0px !important;}
+"""
+
+with gr.Blocks(theme=theme, title="KB AI Challenge", css=css) as demo:
+    
     with gr.Row():
-        with gr.Column(scale=1):
-            audio_in = gr.Audio(sources=["microphone", "upload"], type="numpy", label="음성 질문")
-            asr_btn = gr.Button("음성 인식 시작", variant="secondary")
-            text_in = gr.Textbox(label="인식된 텍스트 (직접 입력 가능)", lines=3)
-            chat_btn = gr.Button("질문하기", variant="primary")
+        # --- LEFT SIDEBAR ---
+        with gr.Column(scale=1, min_width=300, variant="panel"):
+            gr.Markdown("## KB AI Challenge")
+            gr.Markdown("**다국어 금융 AI 어시스턴트**")
+            
+            with gr.Group():
+                lang_dropdown = gr.Dropdown(
+                    choices=list(LANG_MAP.keys()), 
+                    value="한국어 (Korean)", 
+                    label="언어 설정",
+                    interactive=True
+                )
+                
+                file_input = gr.File(label="지식 베이스 (PDF)", file_count="multiple", file_types=[".pdf"])
+                with gr.Row():
+                    upload_btn = gr.Button("업로드 및 분석", variant="primary", size="sm")
+                    upload_status = gr.Textbox(show_label=False, placeholder="상태 대기 중...", interactive=False, lines=1, max_lines=1)
+            
+            gr.Markdown("### 음성 대화")
+            audio_input = gr.Audio(sources=["microphone"], type="numpy", label="음성 입력", show_label=False)
+            
+            with gr.Accordion("시스템 아키텍처", open=False):
+                gr.Markdown(
+                    """
+                    **최적화 내역**
+                    1. **STT**: Google Speech API
+                    2. **번역**: Google Translate API
+                    3. **LLM**: Groq LPU (Llama 3)
+                    """
+                )
+
+        # --- RIGHT MAIN ---
+        with gr.Column(scale=3):
+            # chatbot (Messages format)
+            chatbot = gr.Chatbot(label="대화", height=500, show_label=False)
             
-        with gr.Column(scale=2):
-            answer_box = gr.Textbox(label="AI 답변 (한국어)", lines=6, interactive=False)
-            ref_box = gr.Textbox(label="참고 문헌", lines=4, interactive=False)
+            # References
+            gr.Markdown("**참고 문서**")
+            ref_output = gr.Textbox(show_label=False, interactive=False, lines=3, max_lines=5, placeholder="관련 문서가 표시됩니다.")
             
-    # 이벤트 연결
+            # Input Area
+            with gr.Row():
+                msg = gr.Textbox(
+                    scale=6, 
+                    show_label=False, 
+                    placeholder="질문을 입력하세요...",
+                    container=False
+                )
+                submit_btn = gr.Button("전송", scale=1, variant="primary")
+
+    # --- Event Handlers ---
     upload_btn.click(process_uploaded_files, inputs=[file_input], outputs=[upload_status])
     
-    asr_btn.click(voice_to_text, inputs=[audio_in], outputs=[text_in, gr.State()])
+    msg.submit(run_rag_chat, [msg, chatbot, lang_dropdown], [msg, chatbot, ref_output])
+    submit_btn.click(run_rag_chat, [msg, chatbot, lang_dropdown], [msg, chatbot, ref_output])
     
-    chat_btn.click(
-        run_rag_pipeline, 
-        inputs=[text_in, gr.State('ko')], # 언어는 기본 한국어로 고정 (단순화)
-        outputs=[gr.State(), answer_box, gr.State(), ref_box]
-    )
+    audio_input.stop_recording(voice_to_text_chat, [audio_input, chatbot, lang_dropdown], [msg, chatbot, ref_output])
 
 if __name__ == "__main__":
     demo.launch(share=True)