app.py

import os
import chromadb
from sentence_transformers import SentenceTransformer
from google import genai
import gradio as gr

# === 환경 설정 ===
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", "YOUR_API_KEY_HERE")

# === Simple RAG 시스템 ===
class SimpleRAGSystem:
    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 []
        emb = self.encoder.encode(query).tolist()
        result = self.collection.query(
            query_embeddings=[emb],
            n_results=top_k,
            include=["documents"]
        )
        return result.get("documents", [[]])[0]

rag = SimpleRAGSystem()

# === Google GenAI 클라이언트 ===
client = genai.Client(api_key=API_KEY)

# === 시스템 메시지 ===
SYSTEM_MSG = """
당신은 경희대학교 미디어학과 전문 상담 AI입니다.
"""

# === 응답 함수 ===
def respond(message, history, system_message, max_tokens, temperature, top_p, model_name):
    docs = rag.search(message) if rag.available else []
    ctx = "\n".join(f"참고문서{i+1}: {d}" for i, d in enumerate(docs))
    sys_msg = system_message + ("\n# 참고문서:\n" + ctx if ctx else "")
    convo = "".join(f"사용자: {u}\nAI: {a}\n" for u, a in history)
    prompt = f"{sys_msg}\n{convo}사용자: {message}\nAI:"
    try:
        response = client.models.generate_content(
            model=model_name,
            contents=prompt,
            config={
                "max_output_tokens": max_tokens,
                "temperature": temperature,
                "top_p": top_p
            }
        )
        return response.text or "응답이 없습니다."
    except Exception as e:
        err = str(e).lower()
        if "quota" in err:
            return "API 할당량을 초과했습니다. 나중에 시도해주세요."
        if "authentication" in err:
            return "인증 오류: API 키를 확인하세요."
        return f"오류 발생: {e}"

# === Gradio 인터페이스 ===
demo = gr.ChatInterface(
    fn=respond,
    title="🎬 경희대학교 미디어학과 AI 상담사",
    description="경희대학교 미디어학과에 대해 물어보세요!",
    additional_inputs=[
        gr.Textbox(value=SYSTEM_MSG, label="시스템 메시지", lines=2),
        gr.Slider(128, 2048, value=1024, step=64, label="최대 토큰"),
        gr.Slider(0.1, 1.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(0.1, 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="모델 선택"
        )
    ],
    theme="soft",
    analytics_enabled=False,
)

if __name__ == "__main__":
    demo.launch(share=False)