[langchain] Langchain에서 create_retrieval_chain 사용하기

create_retrieval_chain은 Langchain에서 문서를 검색하고 그 결과를 처리하는 데 매우 유용한 체인이다. 여기에서는 이 함수의 핵심 요소와 사용 방법을 알아보자.

1. `create_retrieval_chain`의 목적

이 함수는 두 가지 주요 작업을 처리한다.

Retriever를 사용해 문서를 검색.
Combine chain을 통해 검색된 문서를 결합하여 최종 결과를 생성.

이 과정을 통해 사용자 질문에 대한 답변을 검색 기반으로 생성할 수 있다.

2. 코드 내용

def create_retrieval_chain(
    retriever: Union[BaseRetriever, Runnable[dict, RetrieverOutput]],
    combine_docs_chain: Runnable[Dict[str, Any], str],
) -> Runnable:
    if not isinstance(retriever, BaseRetriever):
        retrieval_docs: Runnable[dict, RetrieverOutput] = retriever
    else:
        retrieval_docs = (lambda x: x["input"]) | retriever

    retrieval_chain = (
        RunnablePassthrough.assign(
            context=retrieval_docs.with_config(run_name="retrieve_documents"),
        ).assign(answer=combine_docs_chain)
    ).with_config(run_name="retrieval_chain")

    return retrieval_chain

retriever: 문서를 검색하는 역할을 한다. BaseRetriever의 서브클래스이거나 Runnable 객체여야 하며, input 값을 받아서 이를 바탕으로 문서를 반환한다.
combine_docs_chain: 검색된 문서를 처리하여 하나의 문자열로 결합하는 역할을 한다. 이 체인은 원본 입력, 검색된 문서들, 그리고 필요 시 대화 기록(chat history)을 포함한 데이터를 처리한다.

3. 반환 값

이 함수는 LCEL Runnable을 반환하며, 최소한 context와 answer 키를 포함한 딕셔너리 형태의 데이터를 반환한다.

4. 코드 예시

아래는 실제로 이 체인을 구성하고 실행하는 코드 예시이다.

from langchain_community.chat_models import ChatOpenAI
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains import create_retrieval_chain
from langchain_chroma import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langchain import hub

# OpenAI 모델 설정
llm = ChatOpenAI()

# 검색기 설정
vector_store = Chroma(
    persist_directory="./chroma_db",
    embedding_function=embeddings,
)
retriever = vector_store.as_retriever()

# 문서 결합 체인 생성
#retrieval_qa_chat_prompt = hub.pull("langchain-ai/retrieval-qa-chat")
prompt_template = """Answer any use questions based solely on the context below:

{context}

{input}
"""
retrieval_qa_chat_prompt = ChatPromptTemplate(prompt_template)

combine_docs_chain = create_stuff_documents_chain(llm, retrieval_qa_chat_prompt)

# 검색 체인 생성
retrieval_chain = create_retrieval_chain(retriever, combine_docs_chain)

# 체인 실행
result = retrieval_chain.invoke({"input": "사용자 질문"})

5. 핵심 동작

retrieval_docs: retriever가 BaseRetriever의 인스턴스가 아닌 경우, 딕셔너리를 입력으로 받는 Runnable 객체가 된다. 그렇지 않으면, input 키를 추출하여 retriever로 전달하는 람다 함수를 생성한다.
```
 retrieval_docs = (lambda x: x["input"]) | retriever
```

RunnablePassthrough: 검색된 문서를 context로 할당한 후, 이 문서를 기반으로 답변을 생성하는 체인을 구성한다.

 retrieval_chain = (
     RunnablePassthrough.assign(
         context=retrieval_docs.with_config(run_name="retrieve_documents"),
     ).assign(answer=combine_docs_chain)
 ).with_config(run_name="retrieval_chain")