src/retriever.py

from pathlib import Path
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document

# 임베딩 모델 로드 (BGE-M3 추천)
model_name = "BAAI/bge-m3"
model_kwargs = {'device': 'cpu'} # cpu 활용
encode_kwargs = {'normalize_embeddings': True}
embeddings = HuggingFaceEmbeddings(
    model_name=model_name,
    model_kwargs=model_kwargs,
    encode_kwargs=encode_kwargs
)

# 문서 전처리 및 청킹 (Chunking)
def prepare_documents(text_path):
    all_docs = []
    # OCR로 추출한 공지사항
    for file in Path(text_path).glob('*.txt'):
        with open(file, 'r', encoding='utf-8') as f:
            content = f.read()
            print(content)
            all_docs.append(Document(page_content=content, metadata={"source": file.name}))

    # 문서를 적절한 크기로 Cutting
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=600, chunk_overlap=100)
    splits = text_splitter.split_documents(all_docs)
    print(splits)
    return splits

# 1. 경로 설정
DB_PATH = "chroma_db" # 아주대 프로젝트용 DB 저장 경로

# 2. 벡터 DB 생성 및 저장 함수 (데이터가 새로 추가됐을 때만 실행)
def build_vector_store(text_path):
    splits = prepare_documents(text_path)
    
    # Chroma DB를 생성하면서 동시에 디스크에 저장합니다.
    vectorstore = Chroma.from_documents(
        documents=splits, 
        embedding=embeddings,
        persist_directory=DB_PATH
    )
    print(f"벡터 DB가 {DB_PATH}에 저장되었습니다.")
    return vectorstore

def retriever_context(user_input):
    # 저장된 DB 불러오기
    vectorstore = Chroma(
        persist_directory=DB_PATH, 
        embedding_function=embeddings
    )
    
    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
    relevant_docs = retriever.invoke(user_input)
    
    context = "\n\n".join([doc.page_content for doc in relevant_docs])
    source_info = [doc.metadata['source'] for doc in relevant_docs]
    
    return context, source_info