Update app.py

app.py

@@ -1,222 +1,88 @@
 import os
-import asyncio
-from typing import List, Dict
-
-# ChromaDB 경로 설정
-db_dir = os.path.join(os.getcwd(), "chromadb_KH_media")
-os.environ["CHROMA_DB_DIR"] = db_dir
-
-# === 라이브러리 임포트 ===
 import chromadb
-import gradio as gr
 from sentence_transformers import SentenceTransformer
 from google import genai
+import gradio as gr
 
-# === 현재 교수진 목록 === 
-PROFESSORS = [
-    "이인희", "김태용", "박종민", "홍지아", "이정교",
-    "이기형", "이선영", "조수영", "이종혁", "이두황",
-    "이상원", "이훈", "최수진", "최민아", "김관호"
-]
+# 환경 설정
+DB_DIR = os.getenv("CHROMA_DB_DIR", os.path.join(os.getcwd(), "chromadb_KH_media"))
+os.environ["CHROMA_DB_DIR"] = DB_DIR
+API_KEY = os.getenv("GOOGLE_API_KEY", "AIzaSyBuBsC5k9yw2JwUfVFn1Zu1qM_ifwGx6cM")
 
-# === Simple RAG 시스템 ===
+# RAG 시스템
 class SimpleRAGSystem:
-    def __init__(self, db_path: str = None, collection: str = "KH_media_docs"):
-        db_path = db_path or os.getenv("CHROMA_DB_DIR")
-        try:
-            self.embedding_model = SentenceTransformer("snunlp/KR-SBERT-V40K-klueNLI-augSTS")
-            self.client = chromadb.PersistentClient(path=db_path)
-            self.collection = self.client.get_collection(name=collection)
-            count = self.collection.count()
-            self.available = count > 0
-        except Exception:
-            self.available = False
-
-    def search_similar_docs(self, query: str, top_k: int = 10) -> List[str]:
+    def __init__(self, db_path=None, collection_name="KH_media_docs"):
+        path = db_path or DB_DIR
+        self.encoder = SentenceTransformer("snunlp/KR-SBERT-V40K-klueNLI-augSTS")
+        self.client = chromadb.PersistentClient(path=path)
+        self.collection = self.client.get_collection(name=collection_name)
+        self.available = self.collection.count() > 0
+
+    def search(self, query, top_k=8):
         if not self.available:
             return []
-        try:
-            emb = self.embedding_model.encode(query).tolist()
-            res = self.collection.query(
-                query_embeddings=[emb], n_results=top_k,
-                include=["documents", "metadatas"]
-            )
-            return res["documents"][0] if res["documents"] else []
-        except Exception:
-            return []
-
-# RAG 시스템 초기화
-try:
-    rag_system = SimpleRAGSystem()
-except Exception:
-    rag_system = None
-
-# === Google GenAI Client 설정 ===
-client = genai.Client(api_key="AIzaSyBuBsC5k9yw2JwUfVFn1Zu1qM_ifwGx6cM")
-
-# === 시스템 메시지 ===
-SYSTEM_MESSAGE = """당신은 경희대학교 미디어학과 전문 상담 AI입니다.
-
-# 주요 역할:
-- 제공된 문서 정보를 바탕으로 답변 제공
-- 미디어학과 관련 질문에 친절하고 구체적으로 응답
-- 문서에 없는 내용은 일반 지식으로 보완 (단, 명시)
-
-# 답변 스타일:
-- 자세하고 풍부한 설명을 포함하여 상세하고 길게 답변 제공
-- 친근하고 도움이 되는 상담사 톤
-- 핵심 정보를 명확하게 전달
-- 추가 궁금한 점이 있으면 언제든 물어보라고 안내
-
-# 참고 문서 활용:
-- 문서 내용이 있으면 구체적으로 인용
-- 여러 문서의 정보를 종합하여 답변 작성
-- 정확하지 않은 정보는 추측하지 말고 솔직하게 모른다고 답변
-
-# 현재 경희대학교 미디어학과 교수진:
-이인희, 김태용, 박종민, 홍지아, 이정교, 이기형, 이선영, 조수영, 이종혁, 이두황, 이상원, 이훈, 최수진, 최민아, 김관호
-
-한국어로 자연스럽게 답변해주세요."""
-
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-    model_name,
-):
-    # RAG 컨텍스트 추가
-    context_docs = []
-    if rag_system and rag_system.available:
-        try:
-            context_docs = rag_system.search_similar_docs(message, top_k=8)
-        except Exception:
-            pass
-    
-    # 시스템 메시지에 컨텍스트 추가
-    enhanced_system_message = system_message
-    if context_docs:
-        context_text = "\n".join([f"참고문서 {i+1}: {doc}" for i, doc in enumerate(context_docs)])
-        enhanced_system_message += f"\n\n# 참고 문서:\n{context_text}"
-    
-    # 대화 히스토리를 포함한 전체 컨텍스트 구성
-    full_context = enhanced_system_message + "\n\n"
-    
-    # 이전 대화 히스토리 추가
-    for val in history:
-        if val[0]:
-            full_context += f"사용자: {val[0]}\n"
-        if val[1]:
-            full_context += f"AI 상담사: {val[1]}\n"
-    
-    # 현재 사용자 메시지 추가
-    full_context += f"사용자: {message}\nAI 상담사: "
-    
+        emb = self.encoder.encode(query).tolist()
+        docs = self.collection.query(
+            query_embeddings=[emb], n_results=top_k,
+            include=["documents"]
+        )
+        return docs["documents"][0] if docs.get("documents") else []
+
+rag = SimpleRAGSystem()
+
+# Google GenAI 클라이언트
+client = genai.Client(api_key=API_KEY)
+
+# Gradio 응답 함수
+SYSTEM_MSG = f"""
+당신은 경희대학교 미디어학과 전문 상담 AI입니다.
+"""
+
+def respond(message, history, system_message, max_tokens, temperature, top_p, model_name):
+    # RAG 컨텍스트
+    docs = rag.search(message) if rag.available else []
+    ctx = "\n".join(f"참고문서{i+1}: {d}" for i, d in enumerate(docs))
+    sys = system_message + ("\n# 참고문서:\n" + ctx if ctx else "")
+    # 대화 컨텍스트
+    convo = "".join(f"사용자: {u}\nAI: {a}\n" for u, a in history)
+    prompt = f"{sys}\n{convo}사용자: {message}\nAI:"
+    # API 호출
     try:
-        # Google GenAI API 호출
-        api_response = client.models.generate_content(
+        res = client.models.generate_content(
             model=model_name,
-            contents=full_context,
-            config={
-                "temperature": temperature,
-                "top_p": top_p,
-                "max_output_tokens": max_tokens,
-            }
+            contents=prompt,
+            config={"max_output_tokens": max_tokens, "temperature": temperature, "top_p": top_p}
         )
-        
-        if hasattr(api_response, 'text') and api_response.text:
-            yield api_response.text
-        else:
-            yield "응답이 비어있습니다. 다른 모델을 시도해보세요."
-        
+        return res.text or "응답이 없습니다."
     except Exception as e:
-        if "no api" in str(e).lower():
-            error_msg = "API 접근 오류가 발생했습니다. Google AI Studio에서 API 키를 다시 확인해주세요."
-        elif "quota" in str(e).lower():
-            error_msg = "API 할당량이 초과되었습니다. 잠시 후 다시 시도해주세요."
-        elif "authentication" in str(e).lower():
-            error_msg = "인증 오류입니다. API 키를 확인해주세요."
-        else:
-            error_msg = f"응답 생성 중 오류가 발생했습니다: {str(e)}"
-        
-        yield error_msg
-
-# === Gradio 인터페이스 ===
+        err = str(e).lower()
+        if "quota" in err:
+            return "API 할당량을 초과했습니다. 나중에 시도해주세요."
+        if "authentication" in err:
+            return "인증 오류: API 키를 확인하세요."
+        return f"오류 발생: {e}"
+
+# Gradio 인터페이스
 demo = gr.ChatInterface(
-    respond,
+    fn=respond,
     title="🎬 경희대학교 미디어학과 AI 상담사",
-    description="""
-    경희대학교 미디어학과에 대한 모든 궁금한 점을 물어보세요!
-    """,
+    description="경희대학교 미디어학과에 대해 물어보세요!",
     additional_inputs=[
-        gr.Slider(
-            minimum=128,
-            maximum=2048,
-            value=1024,
-            step=64,
-            label="최대 토큰 수"
-        ),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.7,
-            step=0.1,
-            label="Temperature (창의성)"
-        ),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.9,
-            step=0.05,
-            label="Top-p (다양성)"
-        ),
-        gr.Dropdown(
-            choices=[
-                "gemini-2.0-flash",
-                "gemini-2.0-flash-lite", 
-                "gemini-1.5-flash",
-                "gemini-1.5-pro",
-                "gemma-3-27b-it",
-                "gemma-3-12b-it", 
-                "gemma-3-4b-it"
-            ],
-            value="gemini-2.0-flash",
-            label="모델 선택"
-        ),
-    ],
-    additional_inputs_accordion="🔧 고급 설정",
-    examples=[
-        ["미디어학과에서 배우는 주요 과목들은 무엇인가요?"],
-        ["미디어학과 교수진을 소개해주세요."],
-        ["미디어학과 졸업 후 진로는 어떻게 되나요?"],
-        ["미디어학과 입학 전형에 대해 알려주세요."],
-        ["미디어학과 동아리나 학생 활동은 어떤 것들이 있나요?"]
+        gr.Slider(128, 2048, 1024, step=64, label="최대 토큰"),
+        gr.Slider(0.1, 1.0, 0.7, step=0.1, label="Temperature"),
+        gr.Slider(0.1, 1.0, 0.9, step=0.05, label="Top-p"),
+        gr.Dropdown([
+            "gemini-2.0-flash", "gemini-2.0-flash-lite",
+            "gemini-1.5-flash", "gemini-1.5-pro",
+            "gemma-3-27b-it", "gemma-3-12b-it", "gemma-3-4b-it"
+        ], value="gemini-2.0-flash", label="모델 선택")
     ],
     theme="soft",
     analytics_enabled=False,
 )
 
-# === 시스템 정보 표시 ===
-with demo:
-    with gr.Row():
-        with gr.Column():
-            gr.Markdown(f"""
-            ---
-            ### ⚙️ 시스템 정보
-            **언어 모델**: Google Gemini 2.0 Flash + Gemma 3 (4B/12B/27B) 선택 가능
-            **임베딩 모델**: snunlp/KR-SBERT-V40K-klueNLI-augSTS (한국어 특화)
-            **RAG 상���**: {"✅ 활성화" if rag_system and rag_system.available else "❌ 비활성화"}
-            **문서 수**: {rag_system.collection.count() if rag_system and rag_system.available else "0"}개
-            
-            💡 **사용 팁**: 구체적인 질문일수록 더 정확한 답변을 받을 수 있습니다!
-            """)
+def main():
+    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", 7860)), share=False)
 
 if __name__ == "__main__":
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=int(os.environ.get("PORT", 7860)),
-        share=False,
-        show_api=False
-    )
+    main()