Update app.py

app.py

@@ -1,75 +1,84 @@
-import os
-import uuid  # 고유 세션 ID 생성을 위한 라이브러리 추가
-import gradio as gr
-from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
-from langchain_community.document_loaders import PyPDFLoader
-from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain_chroma import Chroma
-from langchain.chains.combine_documents import create_stuff_documents_chain
-from langchain.chains import create_retrieval_chain
-from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from langchain_core.chat_history import BaseChatMessageHistory, InMemoryChatMessageHistory
-from langchain_core.runnables.history import RunnableWithMessageHistory
-
-# 1. LLM 초기화
-llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
-
-# 2. 문서 로드 및 벡터 DB 구축 (서버 구동 시 1회 실행)
-loader = PyPDFLoader("Maximizing Muscle Hypertrophy.pdf")
-pages = loader.load_and_split()
-text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-splits = text_splitter.split_documents(pages)
-
-embeddings = GoogleGenerativeAIEmbeddings(model="gemini-embedding-001")
-vectorstore = Chroma.from_documents(documents=splits, embedding=embeddings)
-retriever = vectorstore.as_retriever()
-
-# 3. RAG 체인 및 메모리 설정
-store = {}
-def get_session_history(session_id: str) -> BaseChatMessageHistory:
-    if session_id not in store:
-        store[session_id] = InMemoryChatMessageHistory()
-    return store[session_id]
-
-qa_prompt = ChatPromptTemplate.from_messages([
-    ("system", """논문 리뷰 전문가입니다. 제공된 문서를 바탕으로 한국어로 답변하세요.
-문서에 없는 내용은 모른다고 답하세요.
-
-{context}"""),
-    MessagesPlaceholder("chat_history"),
-    ("human", "{input}"),
-])
-
-question_answer_chain = create_stuff_documents_chain(llm, qa_prompt)
-rag_chain = create_retrieval_chain(retriever, question_answer_chain)
-
-conversational_rag_chain = RunnableWithMessageHistory(
-    rag_chain,
-    get_session_history,
-    input_messages_key="input",
-    history_messages_key="chat_history",
-    output_messages_key="answer",
-)
-
-# 4. Gradio 연동 함수 (session_id를 동적으로 받도록 수정)
-def chat_response(message, history, session_id):
-    response = conversational_rag_chain.invoke(
-        {"input": message},
-        config={"configurable": {"session_id": session_id}}
-    )
-    return response["answer"]
-
-# 5. 다중 사용자 환경을 고려한 Gradio UI 실행 설정
-with gr.Blocks() as demo:
-    # 접속하는 사용자(브라우저)마다 고유한 UUID를 생성하여 상태값으로 은닉
-    session_state = gr.State(lambda: str(uuid.uuid4()))
-    
-    gr.ChatInterface(
-        fn=chat_response,
-        additional_inputs=[session_state], # 백엔드 함수에 고유 세션 ID 전달
-        title="💪 근비대 극대화 논문 Q&A 봇",
-        description="'Maximizing Muscle Hypertrophy' 논문에 대해 궁금한 점을 물어보세요!"
-    )
-
-if __name__ == "__main__":
+import os
+import uuid
+import gradio as gr
+from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_chroma import Chroma
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain_core.chat_history import BaseChatMessageHistory, InMemoryChatMessageHistory
+from langchain_core.runnables.history import RunnableWithMessageHistory
+from langchain_core.runnables import RunnablePassthrough
+from langchain_core.output_parsers import StrOutputParser
+
+# 1. LLM 초기화
+llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
+
+# 2. 문서 로드 및 벡터 DB 구축
+loader = PyPDFLoader("Maximizing Muscle Hypertrophy.pdf")
+pages = loader.load_and_split()
+text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+splits = text_splitter.split_documents(pages)
+
+embeddings = GoogleGenerativeAIEmbeddings(model="gemini-embedding-001")
+vectorstore = Chroma.from_documents(documents=splits, embedding=embeddings)
+retriever = vectorstore.as_retriever()
+
+# 검색된 문서를 하나의 문자열로 결합하는 헬퍼 함수
+def format_docs(docs):
+    return "\n\n".join(doc.page_content for doc in docs)
+
+# 3. 프롬프트 정의
+qa_prompt = ChatPromptTemplate.from_messages([
+    ("system", """논문 리뷰 전문가입니다. 제공된 문서를 바탕으로 한국어로 답변하세요.
+문서에 없는 내용은 모른다고 답하세요.
+
+{context}"""),
+    MessagesPlaceholder("chat_history"),
+    ("human", "{input}"),
+])
+
+# 4. 에러가 나던 chains 모듈을 버리고 LCEL(파이프라인) 문법으로 RAG 체인 구축
+rag_chain = (
+    RunnablePassthrough.assign(context=(lambda x: format_docs(retriever.invoke(x["input"]))))
+    | qa_prompt
+    | llm
+    | StrOutputParser()
+)
+
+# 5. 메모리(대화 기록) 연동
+store = {}
+def get_session_history(session_id: str) -> BaseChatMessageHistory:
+    if session_id not in store:
+        store[session_id] = InMemoryChatMessageHistory()
+    return store[session_id]
+
+conversational_rag_chain = RunnableWithMessageHistory(
+    rag_chain,
+    get_session_history,
+    input_messages_key="input",
+    history_messages_key="chat_history",
+)
+
+# 6. Gradio 연동 함수
+def chat_response(message, history, session_id):
+    # LCEL 체인은 딕셔너리가 아닌 문자열을 바로 반환하므로 ["answer"] 추출이 필요 없음
+    response = conversational_rag_chain.invoke(
+        {"input": message},
+        config={"configurable": {"session_id": session_id}}
+    )
+    return response
+
+# 7. 다중 사용자 환경 UI 실행
+with gr.Blocks() as demo:
+    session_state = gr.State(lambda: str(uuid.uuid4()))
+    
+    gr.ChatInterface(
+        fn=chat_response,
+        additional_inputs=[session_state],
+        title="💪 근비대 극대화 논문 Q&A 봇",
+        description="'Maximizing Muscle Hypertrophy' 논문에 대해 궁금한 점을 물어보세요!"
+    )
+
+if __name__ == "__main__":
     demo.launch()