app.py

import os
import streamlit as st
import google.generativeai as genai

# Gemini API 설정
genai.configure(api_key=os.environ["GEMINI_API_KEY"])

# 모델 설정
generation_config = {
    "temperature": 0.5,
    "top_p": 0.95,
    "top_k": 40,
    "max_output_tokens": 8192,
    "response_mime_type": "text/plain",
}

model = genai.GenerativeModel(
    model_name="gemini-2.0-flash",
    generation_config=generation_config,
    system_instruction="""당신은 뛰어난 전문성을 갖춘 초등학교 교사, 사회, 과학 등 다양한 교과 내용 전문가입니다. 
    학생이 교과 내용에 대한 질문을 하면, 다양하고 구체적인 예시를 통해 친절하게 설명해주세요. 
    학생들이 모르는 내용에 대해 잘 이해할 수 있도록 쉽게 설명해주세요. 정보의 출처도 표시가 가능하면 꼭 해주세요.""",
)

# 채팅 세션 초기화
if "chat_session" not in st.session_state:
    st.session_state.chat_session = model.start_chat(history=[])

# 스트림릿 UI 설정
st.title("AI 선생님")

# 채팅 메시지 표시
if "messages" not in st.session_state:
    st.session_state.messages = []

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# 사용자 입력 받기
if prompt := st.chat_input("학습 내용을 질문해 보세요. AI 선생님이 구체적인 예시와 설명으로 도와드립니다."):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    # Gemini API 호출
    response = st.session_state.chat_session.send_message(prompt)
    with st.chat_message("assistant"):
        st.markdown(response.text)
    st.session_state.messages.append({"role": "assistant", "content": response.text})

if st.button("초기화"):
    st.session_state.chat_session = model.start_chat(history=[])
    st.session_state.messages = []
    st.rerun()