src/streamlit_app.py

import streamlit as st
import google.generativeai as genai
import os
import json
from io import StringIO

# ------------------------------------------------------
# Hugging Face Spaces에서는 환경변수(Secrets)로 API 키를 저장하세요.
# Settings → Repository secrets → GEMINI_API_KEY 등록
# ------------------------------------------------------
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
genai.configure(api_key=GOOGLE_API_KEY)

# 초기 설정
if "messages" not in st.session_state:
    st.session_state.messages = []

if "system_prompt" not in st.session_state:
    st.session_state.system_prompt = "당신은 친절한 AI 어시스턴트입니다."

model = genai.GenerativeModel("gemma-3-27b-it")

# ------------------------------------------------------
# Streamlit UI
# ------------------------------------------------------
st.set_page_config(page_title="Gemini 챗봇", page_icon="🤖", layout="wide")
st.title("🤖 Google Gemini 대화형 챗봇")
st.caption("Hugging Face Spaces + Streamlit + Google Generative AI")

# 사이드바: 시스템 프롬프트 수정 & 로그 다운로드
with st.sidebar:
    st.subheader("⚙️ 설정")
    new_system_prompt = st.text_area("시스템 프롬프트", st.session_state.system_prompt, height=100)
    if st.button("변경 적용"):
        st.session_state.system_prompt = new_system_prompt
        st.success("시스템 프롬프트가 변경되었습니다.")

    st.markdown("---")
    if st.session_state.messages:
        # JSON 변환
        json_data = json.dumps(st.session_state.messages, ensure_ascii=False, indent=2)
        st.download_button(
            label="💾 대화 로그 JSON 다운로드",
            data=json_data,
            file_name="chat_log.json",
            mime="application/json"
        )

# 대화 출력
for msg in st.session_state.messages:
    with st.chat_message("user"):
        st.markdown(msg["user"])
    with st.chat_message("assistant"):
        st.markdown(msg["ai"])

# 사용자 입력
if prompt := st.chat_input("메시지를 입력하세요. 종료하려면 'quit' 입력"):
    if prompt.lower() == "quit":
        st.warning("대화를 종료합니다. 새로고침하면 다시 시작할 수 있습니다.")
    else:
        # 대화 컨텍스트 구성
        context = f"시스템: {st.session_state.system_prompt}\n"
        for msg in st.session_state.messages[-10:]:  # 최근 10개만 유지
            context += f"사용자: {msg['user']}\nAI: {msg['ai']}\n"
        context += f"사용자: {prompt}\nAI:"

        with st.chat_message("user"):
            st.markdown(prompt)

        try:
            response = model.generate_content(context)
            ai_response = response.text
        except Exception as e:
            ai_response = f"⚠️ 오류 발생: {e}"

        with st.chat_message("assistant"):
            st.markdown(ai_response)

        # 대화 기록 저장
        st.session_state.messages.append({"user": prompt, "ai": ai_response})