CoRAG (Chain-of-Retrieval Augmented Generation) CoRAG는 복잡한 질문을 여러 개의 하위 쿼리로 분해하고, 각 검색 단계의 결과가 다음 쿼리를 정제하는 연쇄적 검색(Chain-of-Retrieval) 아키텍처입니다. Chain-of-Thought가 추론을 단계적으로 수행하듯, CoRAG는 검색을 단계적으로 수행하여 복잡한 질문에 필요한 컨텍스트를 점진적으로 구축합니다.
핵심 아이디어 기존 RAG는 사용자 질문을 그대로 검색 쿼리로 사용하여 한 번에 모든 관련 문서를 가져옵니다. 그러나 “A와 B의 공통점과 차이점을 비교하라”처럼 여러 정보를 종합해야 하는 질문에서는, 단일 검색으로 필요한 모든 정보를 확보하기 어렵습니다.
CoRAG는 질문을 하위 쿼리 체인으로 분해하고, 이전 검색 결과를 참고하여 다음 쿼리를 생성 합니다. 이를 통해 각 검색 단계가 이전 단계의 맥락 위에서 더 정밀한 정보를 가져옵니다.
CoRAG의 핵심 기여는 검색 자체를 추론 과정으로 취급한다는 점입니다. 단순히 쿼리를 분해하는 것이 아니라, 이전 검색 결과가 다음 쿼리 생성에 영향을 미치는 순차적 의존성 이 있습니다. 이 점이 병렬 쿼리 분해 방식과의 핵심 차이입니다.
동작 방식 Single-shot Retrieval vs Chain-of-Retrieval 항목 Single-shot Retrieval Chain-of-Retrieval (CoRAG) 검색 횟수 1회 N회 (하위 쿼리 수만큼) 쿼리 생성 원본 질문 그대로 사용 이전 검색 결과를 반영하여 순차 생성 컨텍스트 구축 한 번에 확보 단계적으로 누적 멀티홉 질문 취약 (필요 정보 누락 가능) 강점 (단계별 정보 수집) 비교 분석 질문 취약 (한쪽 정보만 검색될 수 있음) 강점 (각 대상을 순차 검색) LLM 호출 비용 낮음 높음 (쿼리 분해 + 단계별 생성)
LangGraph 구현 상태 정의 from typing import TypedDict, List from langchain_core.documents import Document class CoRAGState ( TypedDict ): question: str sub_queries: List[ str ] retrieved_contexts: List[ str ] current_step: int max_steps: int answer: str
노드 함수 from langchain.chat_models import init_chat_model from langchain_openai import OpenAIEmbeddings from langchain_chroma import Chroma from langchain_core.prompts import ChatPromptTemplate from langchain_core.output_parsers import StrOutputParser llm = init_chat_model( "gpt-4o-mini" , temperature = 0 ) embeddings = OpenAIEmbeddings() vectorstore = Chroma( collection_name = "documents" , embedding_function = embeddings) retriever = vectorstore.as_retriever( search_kwargs = { "k" : 3 }) def decompose_query ( state : CoRAGState) -> CoRAGState: """질문을 하위 쿼리 체인으로 분해합니다.""" question = state[ "question" ] prompt = ChatPromptTemplate.from_messages([ ( "system" , ( "복잡한 질문을 답변에 필요한 순서대로 하위 질문으로 분해하세요. \n " "각 하위 질문은 한 줄에 하나씩, 번호 없이 출력하세요. \n " "최대 4개까지만 생성하세요." )), ( "human" , " {question} " ), ]) chain = prompt | llm | StrOutputParser() result = chain.invoke({ "question" : question}) sub_queries = [q.strip() for q in result.strip().split( " \n " ) if q.strip()] return { "sub_queries" : sub_queries[: 4 ], "current_step" : 0 , "max_steps" : min ( len (sub_queries), 4 ), "retrieved_contexts" : [], }
def retrieve_step ( state : CoRAGState) -> CoRAGState: """현재 하위 쿼리에 대해 검색하고, 이전 컨텍스트를 반영한 쿼리를 생성합니다.""" current_step = state[ "current_step" ] sub_queries = state[ "sub_queries" ] previous_contexts = state[ "retrieved_contexts" ] current_query = sub_queries[current_step] # 이전 검색 결과가 있으면 이를 반영하여 쿼리를 정제 if previous_contexts: refine_prompt = ChatPromptTemplate.from_messages([ ( "system" , ( "이전 검색에서 얻은 정보를 참고하여, " "다음 하위 질문을 더 구체적인 검색 쿼리로 변환하세요. \n " "검색 쿼리만 출력하세요." )), ( "human" , ( "이전 검색 결과: \n {previous} \n\n " "하위 질문: {query} " )), ]) refine_chain = refine_prompt | llm | StrOutputParser() refined_query = refine_chain.invoke({ "previous" : " \n --- \n " .join(previous_contexts), "query" : current_query, }) else : refined_query = current_query # 검색 수행 documents = retriever.invoke(refined_query) context = " \n\n " .join(doc.page_content for doc in documents) updated_contexts = list (previous_contexts) + [context] return { "retrieved_contexts" : updated_contexts, "current_step" : current_step + 1 , }
def synthesize ( state : CoRAGState) -> CoRAGState: """모든 검색 결과를 종합하여 최종 답변을 생성합니다.""" question = state[ "question" ] sub_queries = state[ "sub_queries" ] contexts = state[ "retrieved_contexts" ] # 하위 쿼리와 검색 결과를 구조화 structured_context = "" for i, (query, ctx) in enumerate ( zip (sub_queries, contexts)): structured_context += f "[검색 { i + 1 } ] { query } \n { ctx } \n\n " prompt = ChatPromptTemplate.from_messages([ ( "system" , ( "다음은 질문에 대해 단계적으로 검색한 결과입니다. \n " "모든 검색 결과를 종합하여 질문에 답변하세요. \n\n " " {context} " )), ( "human" , " {question} " ), ]) chain = prompt | llm | StrOutputParser() answer = chain.invoke({ "context" : structured_context, "question" : question}) return { "answer" : answer}
조건부 엣지 from typing import Literal def check_complete ( state : CoRAGState) -> Literal[ "continue" , "synthesize" ]: """모든 하위 쿼리에 대한 검색이 완료되었는지 확인합니다.""" if state[ "current_step" ] >= state[ "max_steps" ]: return "synthesize" return "continue"
그래프 구성 from langgraph.graph import StateGraph, START , END workflow = StateGraph(CoRAGState) # 노드 추가 workflow.add_node( "decompose_query" , decompose_query) workflow.add_node( "retrieve_step" , retrieve_step) workflow.add_node( "synthesize" , synthesize) # 엣지 연결 workflow.add_edge( START , "decompose_query" ) workflow.add_edge( "decompose_query" , "retrieve_step" ) workflow.add_conditional_edges( "retrieve_step" , check_complete, { "continue" : "retrieve_step" , "synthesize" : "synthesize" , }) workflow.add_edge( "synthesize" , END ) # 컴파일 app = workflow.compile()
result = app.invoke({ "question" : "Transformer와 Mamba 아키텍처의 핵심 메커니즘을 비교하고, 각각의 장단점을 분석하라" , }) print (result[ "answer" ])
CoRAG는 하위 쿼리 수만큼 검색과 LLM 호출이 반복되므로, 단일 검색 대비 비용이 증가합니다. max_steps를 적절히 설정하여 비용과 품질의 균형을 조절하세요.
적용 시나리오
“A가 B에 영향을 미쳤고, B의 결과로 C가 발생했는데, C의 현재 상태는?” 같은 연쇄적 추론이 필요한 질문에서, 각 단계에서 이전 결과를 참고하여 다음 정보를 검색합니다.
두 개 이상의 대상을 비교하는 질문에서, 각 대상에 대한 정보를 순차적으로 검색하여 편향 없는 비교 컨텍스트를 구축합니다.
여러 측면에서의 분석이 필요한 보고서 생성 시, 각 측면을 하위 쿼리로 분해하여 체계적으로 정보를 수집합니다.
참고 논문 논문 학회 링크 Chain-of-Retrieval Augmented Generation (Shi et al., 2025) - arXiv 2501.14342
