Jangrae/Streamlit

1 2import streamlit as st3from langchain_openai import ChatOpenAI, OpenAIEmbeddings4from langchain_community.vectorstores import FAISS5from langchain_core.prompts import ChatPromptTemplate6from langchain_core.output_parsers import StrOutputParser7from langchain_core.messages import HumanMessage, AIMessage8from operator import itemgetter9 10# 데이터 및 리트리버 불러오기11FAISS_DIR = "faiss_index"12 13# 체인 캐싱14@st.cache_resource15def get_chain():16    s_msg = "당신은 주어진 규정을 기반으로 답변하는 규정 도우미입니다.\n\n[규정]\n\n{context}"17    h_msg = "{input}"18    messages = [("system", s_msg), ("human", h_msg)]19    prompt = ChatPromptTemplate.from_messages(messages)20    llm = ChatOpenAI(model="gpt-4o-mini", streaming=True, temperature=0)21    rag_chain = (22        {23            "context": itemgetter("input") | retriever | format_docs,24            "input": itemgetter("input")25        }26        | prompt27        | llm28        | StrOutputParser()29    )30    return rag_chain31 32# Embedding 모델, 검색기33@st.cache_resource34def get_retriever():35    embedding = OpenAIEmbeddings(model="text-embedding-3-small")36    vectorstore = FAISS.load_local(FAISS_DIR,37                                   embedding,38                                   allow_dangerous_deserialization=True)39    return vectorstore.as_retriever(search_kwargs={"k": 3})40 41# 문서 포맷팅 함수42def format_docs(docs):43    return "\n\n".join(doc.page_content for doc in docs)44 45# 검색기 및 RAG 체인 준비46retriever = get_retriever()47chain = get_chain()48 49# 페이지 설정50st.set_page_config(page_title="FAQ 챗봇", page_icon="🏢")51st.title("🏢 EduVerse Academy FAQ 챗봇")52 53# 세션 상태 초기화54if "messages" not in st.session_state:55    st.session_state.messages = []56 57# 기존 대화 내용 출력58for message in st.session_state.messages:59    role = "user" if isinstance(message, HumanMessage) else "assistant"60    with st.chat_message(role):61        st.markdown(message.content)62 63# 사용자 입력 처리64user_input = st.chat_input("메시지를 입력하세요.")65if user_input:66    # 사용자 메시지 표시67    st.chat_message("user").markdown(user_input)68 69    # 응답 출력70    with st.chat_message("assistant"):71        response_stream = chain.stream({"input": user_input})72        full_response = st.write_stream(response_stream)73 74    # 대화 기록 저장75    st.session_state.messages.append(HumanMessage(content=user_input))76    st.session_state.messages.append(AIMessage(content=full_response))77