Update app.py

app.py

@@ -1,9 +1,11 @@
 import os
+import time
+import google.generativeai as genai
 import streamlit as st
+import re
 from streamlit_chat import message
-import google.generativeai as genai
 
-# Gemini API 키 설정
+# Google Generative AI API 설정
 genai.configure(api_key=os.environ["GEMINI_API_KEY"])
 
 # 모델 설정
@@ -14,43 +16,90 @@ generation_config = {
     "max_output_tokens": 8192,
     "response_mime_type": "text/plain",
 }
+safety_settings = [
+    {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+    {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+    {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+    {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+]
 
-model = genai.GenerativeModel(
-    model_name="learnlm-1.5-pro-experimental",
-    generation_config=generation_config,
-    system_instruction="""교사가 학습 목표를 입력하면, 학생들이 학습 목표에 도달할 수 있게 질문을 계속 해줘, 
+# 전역 변수로 chat_session 초기화
+chat_session = None
+
+def start_new_chat_session():
+    global chat_session
+    chat_session = genai.GenerativeModel(
+        model_name="learnlm-1.5-pro-experimental",
+        safety_settings=safety_settings,
+        generation_config=generation_config,
+        system_instruction="""
+        교사가 학습 목표를 입력하면, 학생들이 학습 목표에 도달할 수 있게 질문을 계속 해줘, 
         학생들은 네 질문에 답을 할 것이야. 그러면 바로 답을 제시하지 말고 질문형태로 피드백을 해줘. 
-        질문을 통해서 학생들이 스스로 학습 목표를 도달할 수 있도록 도와줘.""",
+        질문을 통해서 학생들이 스스로 학습 목표를 도달할 수 있도록 도와줘.
+        """,
+    ).start_chat(history=[])
+
+# 초기 세션 시작
+start_new_chat_session()
+
+# Streamlit 앱 설정
+st.set_page_config(page_title="질문하는 AI 선생님", layout="centered")
+
+st.markdown(
+    """
+    <div style='font-size: 30px; font-weight: bold; text-align: center;'>질문하는 AI 선생님</div>
+    <div style='text-align: center; margin-bottom: 20px;'>학습주제를 입력하면, AI가 끊임 없는 질문을 통해 스스로 학습하게 도와줍니다.</div>
+    """,
+    unsafe_allow_html=True
 )
 
-# Streamlit 세션 상태 초기화
+# 채팅 히스토리 관리
+if "history" not in st.session_state:
+    st.session_state.history = []
+
+# 채팅 세션 초기화
 if "chat_session" not in st.session_state:
-    st.session_state.chat_session = model.start_chat(history=[])
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-
-# 웹앱 제목 설정
-st.title("AI 학습 튜터")
-
-# 학습 목표 입력
-if "learning_goal" not in st.session_state:
-    st.session_state.learning_goal = st.text_input("학습 목표를 입력해주세요:")
-    if st.session_state.learning_goal:
-        st.session_state.chat_session.send_message(st.session_state.learning_goal)
-        response = st.session_state.chat_session.last
-        st.session_state.messages.append({"role": "ai", "content": response.text})
-
-# 채팅 메시지 표시
-if st.session_state.learning_goal:
-    for i, msg in enumerate(st.session_state.messages):
-        message(msg["content"], is_user=msg["role"] == "user", key=str(i))
-
-    # 사용자 입력
-    if prompt := st.chat_input("질문을 입력해주세요:"):
-        st.session_state.messages.append({"role": "user", "content": prompt})
-        message(prompt, is_user=True, key=str(len(st.session_state.messages)))
-
-        # Gemini API를 통해 답변 생성
-        response = st.session_state.chat_session.send_message(prompt)
-        st.session_state.messages.append({"role": "ai", "content": response.text})
-        message(response.text, is_user=False, key=str(len(st.session_state.messages)))
+    st.session_state.chat_session = chat_session
+
+# 사용자 입력 받기
+user_input = st.text_input(
+    "입력", key="user_input", placeholder="여기에 학습 목표나 질문을 입력하세요...",
+    help="학습 목표를 입력하거나 학습하고자 하는 주제에 대한 질문을 입력하세요."
+)
+
+def respond(user_input):
+    st.session_state.history.append({"user": user_input, "ai": ""})
+    try:
+        response = st.session_state.chat_session.send_message(user_input, stream=True)
+        full_text = ""
+        
+        # 스트리밍 응답 처리
+        for chunk in response:
+            chunk_text = chunk.text
+            
+            # 단어 단위로 분리하여 실시간 출력
+            for word in chunk_text.split():
+                full_text += word + " "
+                
+                # 히스토리 최신 항목 업데이트
+                st.session_state.history[-1]["ai"] = full_text
+                time.sleep(0.05)  # 출력 속도 조절
+
+    except Exception as e:
+        st.session_state.history[-1]["ai"] = f"에러 발생: {str(e)}"
+
+# 응답 생성 및 출력
+if user_input:
+    respond(user_input)
+
+# 채팅 메시지 출력
+for chat in st.session_state.history:
+    if "user" in chat:
+        message(chat["user"], is_user=True, key=f"user_{chat['user']}")
+    if "ai" in chat:
+        message(chat["ai"], is_user=False, key=f"ai_{chat['ai']}")
+
+# 초기화 버튼 클릭 시
+if st.button("초기화"):
+    start_new_chat_session()
+    st.session_state.history = []