[Agent Study]RAG 성능을 올리기 위한 Context Retrieval 을 적용

요즘 에이전트를 만들다가 다시금 RAG 관련 툴을 만들면서 예전에 적용했었던 기법들을 재활용해보고 있습니다. 코드 레벨에서 구현은 해본 적은 없고, 가짜연구소 9th 잔디심기에서 관련내용을 발표를 한 적은 있는데요. RAG를 하면서 성능이 많이 개선되었던 기억이 있어서 이번에 새로운 Agent 프로젝트에 들어가면서 실제 사례에서 적용한 UseCase를 토대로 글을 작성해 봅니다.

 

Introducing Contextual Retrieval

 

원리는 매우 간단합니다. rag를 구축하기 위해선 데이터를 스플릿 하고 청크를 하게 되는 과정에서 LLM 호출을 한번 해서 데이터 품질을 보강하는 방법입니다. 예로 들어 하나의 도큐먼트 코퍼스가 있을 때 LLM에게 콘텍스트를 주고, 스플릿이 되는 과정에서 청크를 콘텍스트에 추가해서 2가지의 정보를 토대로 summarize 한 형태의 정보를 앞에 입력해 주는 기법이라 할 수 있습니다. 

원본 문서: [전체 문서 내용...]

청크1: "회사의 수익이 전 분기 대비 3% 증가했다"
→ 생성된 컨텍스트1: "이 내용은 ACME사의 2023년 Q2 실적 보고서의 일부입니다. 이전 분기 매출은 3.14억 달러였습니다."
→ 최종 청크1: [컨텍스트1 + 청크1 내용]

청크2: [다른 내용...]
→ 생성된 컨텍스트2: [청크2에 대한 맥락 설명]
→ 최종 청크2: [컨텍스트2 + 청크2 내용]

 

단순하게 청크 1에는 회사 수익이 3% 증가정보만 있고, 어떤 회사인지? 어떤 분 기인 지 다양한 정보가 누락이 되어있습니다. 그렇기 맥락의 손실이 발생하고, 이걸 해결하기 위해 전체 콘텍스트를 정보로 주게 됩니다.

 

여기서 그러면? 만약 500페이지짜리 내용이면 전체콘텍스트를 줄 거냐? 어떤 논문이라고 하면 요즘 롱콘텍스트라 10장 정도는 가볍게 콘텍스트로 줄 수 있겠지만 대량의 페이지는 어려움이 있습니다. 그래서 Document Parser를 활용해서 먼저 데이터 로드를 잘해주는 과정이 필요하겠다 생각이 들었습니다.

GitHub - microsoft/markitdown: Python tool for converting files and office documents to Markdown.

 

결국은 로드가 잘되야 하는 부분인데 저도 생각한 게 마킷다운 형태로 "#" 태그를 구분해서, 서론 본론 결론을 나눌 수 있다면 이문제도 해결이 가능하지 않을까? 생각을 하긴 했습니다. 일단 해봐야 하겠지만 그냥 글을 쓰면서 생각이 나서 적어봤습니다. 본론으로 돌아와서 결국 아래의 그림이 제가 하는 말에 해당합니다.

 

결국 50~100개 사이의 추가 콘텍스트를 만들어주고 Semantic Search + bm25를 섞어서 검색이 잘 걸리게 해라가 이 내용의 key point입니다. 요약 task는 충분히 sLLM이나 Closed 모델에서도 mini , haiku로 처리한다면 비용문제도 어느 정도는 해결될 것입니다. 그리고 애초애 데이터를 잘 만들어주고 힌트를 주는 게 RAG를 구축하는 과정에서 더욱 성능을 올리는 방법이라고 생각합니다.

 

def get_context(self, document_content: str, filename: str) -> Optional[str]:
        try:
            messages = [
                (
                    "system",
                    "이 JSON 문서의 내용을 이해하기 쉽게 설명하는 간단한 문맥을 제공해주세요. "
                    "검색을 개선하기 위한 목적입니다. 간단한 문맥만 답변하고 다른 것은 포함하지 마세요."
                ),
                (
                    "human",
                    f"파일명: {filename}\n\n문서 내용:\n{document_content}"
                )
            ]
            
            response = self.llm.invoke(messages)
            time.sleep(2)
            return response.content
        except Exception as e:
            print(f"Error in get_context: {str(e)}")
            return None

 

위처럼 llm을 활용해서 답변을 만드는 함수를 만듭니다. 그 후 스플릿을 지정하고, 그에 따라 맞는 로직을 구성해 줍니다. 저는 json을 하나의 Document로 받아서 document에 대한 desc를 작성할 수 있도록 만들었습니다.

def process_json_files(self, directory_path: str) -> List[Document]:
        all_documents = []
        
        if not os.path.exists(directory_path):
            raise FileNotFoundError(f"Directory not found: {directory_path}")
        
        for filename in os.listdir(directory_path):
            if not filename.endswith('.json'):
                continue
                
            file_path = os.path.join(directory_path, filename)
            
            try:
                loader = JSONLoader(
                    file_path=file_path,
                    jq_schema='.',
                    text_content=False
                )
                
                documents = loader.load()
                
                for doc in documents:
                    context = self.get_context(doc.page_content, filename)
                    if context is not None:
                        enhanced_doc = Document(
                            page_content=f"{context}\n\n{doc.page_content}",
                            metadata={
                                "source": filename,
                                "type": "document"
                            }
                        )
                        all_documents.append(enhanced_doc)
                        
                print(f"Successfully processed {filename}")
                print(f"응답생성후 대기중 {filename}")
                time.sleep(10)
                print(f"응답생성후 대기중 완료 {filename}")
            except Exception as e:
                print(f"Error processing file {filename}: {str(e)}")
                continue

 

위로직처럼 만들어진 context는 document 앞부분에 저장이 되어 나옵니다. 

 

소스: internet.json
내용: 이 JSON 문서는 인터넷 및 TV 요금제와 관련된 다양한 정보를 제공합니다. 인터넷 단독 요금제와 TV 번들 요금제의 속도별 가격, 할인 유형, 가족 할인 플랜, 셋톱박스와 와이파이 라우터 장비 가격, 설치 비용 등의 정보가 포함되어 있습니다.

{"internet_tv_plans": {"internet_only": {"speeds": {"100M": {"base_price": 22000, "mobile_bundle_price": 16500}, "500M": {"base_price": 33000, "mobile_bundle_price": 22000}, "1G": {"base_price": 38500, "mobile_bundle_price": 27500}}}, "internet_tv_bundles": {"basic": {"channels": 236, "speeds": {"100M": {"base_price": 38500, "mobile_bundle_price": 33000}, "500M": {"base_price": 44000, "mobile_bundle_price": 38500}, "1G": {"base_price": 49500, "mobile_bundle_price": 44000}}}, "light": {"channels": 240, "speeds": {"100M": {"base_price": 39600, "mobile_bundle_price": 34100} ...

 

위처럼 Document는 구성이 되고 이제 벡터서치를 해주면 됩니다.

 

하이브리드서치

from typing import List, Dict, Any
from langchain.embeddings import OpenAIEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.retrievers import EnsembleRetriever
from langchain_core.runnables import ConfigurableField
from langchain_teddynote.retrievers import KiwiBM25Retriever

class HybridSearchSystem:
    def __init__(
        self,
        vectorstore_path: str,
        embedding_model_name: str = "text-embedding-3-small",
        weights: List[float] = [0.7, 0.3],
        semantic_k: int = 4,
        bm25_k: int = 4,
    ):
        """
        초기화 함수
        
        Args:
            vectorstore_path: 저장된 벡터스토어 경로
            embedding_model_name: 임베딩 모델 이름 
            weights: [KiwiBM25 가중치, Semantic 가중치]
            k: 검색 결과 개수
        """
        self.semantic_k = semantic_k
        self.bm25_k = bm25_k
        self.weights = weights
        
   
        print("벡터스토어 로드 중...")
        self.embeddings = OpenAIEmbeddings(model=embedding_model_name)
        self.vectorstore = FAISS.load_local(vectorstore_path, self.embeddings, allow_dangerous_deserialization=True)
        
        docs = self.vectorstore.docstore._dict.values()
        texts = [doc.page_content for doc in docs]
        
        self.bm25_retriever = KiwiBM25Retriever.from_texts(texts)
        self.bm25_retriever.k = self.bm25_k
        
        self.semantic_retriever = self.vectorstore.as_retriever(
            search_kwargs={"k": self.semantic_k}
        )
        self.ensemble_retriever = EnsembleRetriever(
            retrievers=[self.bm25_retriever, self.semantic_retriever],
            weights=weights
        ).configurable_fields(
            weights=ConfigurableField(
                id="ensemble_weights",
                name="Ensemble Weights",
                description="검색기 가중치 설정 [KiwiBM25, Semantic]"
            )
        )

    def search(self, query: str, weights: List[float] = None) -> List[Dict[str, Any]]:
        """
        검색 수행
        
        Args:
            query: 검색어
            weights: [KiwiBM25 가중치, Semantic 가중치] (선택적)
            
        Returns:
            검색 결과 리스트
        """
        config = None
        if weights is not None:
            config = {"configurable": {"ensemble_weights": weights}}
            
        docs = self.ensemble_retriever.invoke(
            query,
            config=config
        )
        
        results = []
        for doc in docs:
            results.append({
                "content": doc.page_content,
                "metadata": doc.metadata
            })
        
        return results

def pretty_print(results: List[Dict[str, Any]]):
    """결과 출력 함수"""
    for i, result in enumerate(results, 1):
        print(f"\n[결과 {i}]")
        print(f"내용: {result['content']}")
        if 'source' in result['metadata']:
            print(f"소스: {result['metadata']['source']}")
        if 'score' in result['metadata']:
            print(f"점수: {result['metadata']['score']:.4f}")
        print("-" * 80)

if __name__ == "__main__":
    try:
        # 기존 벡터스토어 경로
        vectorstore_path = "/Users/kdb/Desktop/agent/faiss_index"
        
        # 검색 시스템 초기화
        print("시스템 초기화 중...")
        search_system = HybridSearchSystem(
            vectorstore_path=vectorstore_path,
            semantic_k=4,
            bm25_k=2,
            weights=[0.7, 0.3]
        )
        
        # 검색할 쿼리 설정
        query = "인터넷 요금제"
        print(f"\n검색 쿼리: {query}")
        # KiwiBM25 위주 검색
        print("\nKiwiBM25 위주 검색 결과 (KiwiBM25: 0.9, Semantic: 0.1):")
        results = search_system.search(
            query=query, 
            weights=[0.9, 0.1]
        )
        pretty_print(results)
            
    except Exception as e:
        print(f"오류 발생: {str(e)}")

 

앙상블 리트리버를 활용해서  kiwi bm25와 Semantic Search를 적용한 로직입니다.

 

문서에서는 20개의 데이터를 가지고 Rank Fusion을 통해 겹치는 것은 제외한 후 리트리버를 하는 방식이지만 RAG를 무작정 쪼갠다기보단 데이터에 맞게 쪼개는 게 저는 맞다고 생각을 해서, 어느 정도 LLM의 콘텍스트에 다 들어갈 데이터라면 하나의 Docs 형태로 구성해서 RAG 처리를 하는 편입니다.

 

사용하는 Task에 맞게 서치 개수도 조절해서, 콘텍스트에 주어야 noise가 덜 발생할 거라고 생각이 들기 때문에 k는 조절해서 사용해 보면 좋고, 위 글에서는 Re-ranker를 통해서 개선을 했을 때 더욱 개선이 되었다는 결과가 있습니다. 리랭커는 150개의 청크를 뽑아서 리랭크하는 방식으로 적용을 했다고 하는데 일단 예제에서는 그만한 데이터가 없어서 이 정도로 해보았습니다. 지금 개발하고자 하는 프로덕트가 대용량이 될 것 같기에 나중에 리랭커를 꼭 넣어야 하지 않을까 생각은 듭니다. 이제 RAG 할 때 콘텍스트 리트리버는 필수 아닐까 조심스럽게 생각은 드는데 데이터가 너무 많은 경우 그만큼 돈이 드는 것도 있기 때문에 잘 고려해서 쓰시면 좋을 것 같습니다.