💬streamlit과 RAG를 활용한 챗봇

from dotenv import load_dotenv
from langchain_community.chat_models import ChatOpenAI              
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.schema import ChatMessage, HumanMessage, SystemMessage
from langchain.callbacks.base  import BaseCallbackHandler
from pdfminer.high_level import extract_text
import streamlit as st
import os

class MarkdownStreamHandler(BaseCallbackHandler):

    def __init__(self, output_container, initial_content=""):
        self.output_container = output_container 
        self.generated_content = initial_content  

    def on_llm_new_token(self, token: str, **kwargs) -> None:
        self.generated_content += token
        self.output_container.markdown(self.generated_content)

def extract_text_from_pdf(file):
    '''pdfminer를 사용하여 PDF 파일에서 텍스트 추출.'''
    try:
        return extract_text(file)
    except Exception as error:
        st.error(f"PDF 텍스트 추출 중 오류 발생: {error}")
        return ""
    

def handle_uploaded_file(file):
    '''업로드 된 PDF 파일을 처리하고 벡터 스토어 준비'''
    if not file:
        return None, None
    
    # 파일 유형에 따라 텍스트 추출
    document_text = extract_text_from_pdf(file) if file.type == 'application/pdf' else ""
    if not document_text:
        st.error("텍스트 추출 불가")
        return None, None
    
    # 문서를 작은 청크로 나누어 벡터화 준비
    text_splitter = CharacterTextSplitter(
        separator="\n",
        chunk_size=1000,
        chunk_overlap=200
    )

    document_chunks = text_splitter.create_documents([document_text])
    st.info(f"{len(document_chunks)} 개의 문서 단락 생성.")

    # 유사성 검색을 위한 벡터 스토어 생성
    vectorstore = FAISS.from_documents(document_chunks, embedding_model)
    return vectorstore, document_text

def get_rag_response(user_query, vectorstore, callback_handler):
    '''검색된 문서를 기바능로 사용자 질문에 대한 응답 생성'''
    if not vectorstore:
        st.error("No Vectorstore. Upload docs first.")
        return ""
    
    # 가장 유사한 문서 3개 검색
    retrieved_docs = vectorstore.similarity_search(user_query, k=3)
    retrieved_text = "\n".join(f"문서 {i+1} : {doc.page_content}" for i, doc in enumerate(retrieved_docs))

    # LLM 설정
    chat_model = ChatOpenAI(model_name = "gpt-4", temperature=0, streaming=True, callbacks=[callback_handler])

    # RAG 프롬프트 
    rag_prompt = [
        SystemMessage(content = "제공된 문서를 기반으로 사용자의 질문에 답변하세요. 정보가 없으면 '모르겠습니다' 라고 답변하세요"),
        HumanMessage(content=f"질문 : {user_query}\n\n{retrieved_text}")

    ]

    try:
        response = chat_model(rag_prompt)
        return response.content
    
    except Exception as error:
        st.error(f"응답 생성 중 오류 발생: {error}")
        return ""