[Langchain] Retriever 사용하기

최근의 프롬프트 엔지니어링 영역은 단순히 프롬프트를 잘 작성하는 것에서, 다양한 보조 도구를 사용하는 방법으로 발전해오고 있습니다. 대표적으로 에이전트라는 기능은 사용자가 질문을 할 때 해당 질문을 어떻게 잘 처리해야 할 지에 대해 LLM과 이 LLM이 사용할 도구들을 구현하여, LLM이 도구들을 사용하며 해당 질문을 반복적으로 처리하며 좋은 답변을 생성하는 방법입니다.

 

에이전트에서 자주 사용되는 방법 중 하나는 RAG인데요 검색 증강 생성이라는 방법으로 특정 문서나 DB 등에서 사용자의 질문에 가장 적합한 데이터를 가져와 이를 기반으로 대답하는 방법입니다. 이 RAG에서 가장 중요한 기능은 사용자의 질문에 가장 적합한 데이터를 가져오는 검색기, Retriever입니다.

 

이번에는 랭체인 예제에서 간단하게 리트리버를 살펴보도록 하겠습니다. 예제는 저번 벡터스토어 페이지에서 사용한 예제와 동일합니다.

벡터스토어 페이지의 예제에서 Chroma 벡터스토어를 ./chroma_db에 저장했다고 가정하고 실습을 진행하고 있습니다.

Vector stores and retrievers | 🦜️🔗 LangChain

 

제가 작업한 노트북 파일은 아래 링크에서 확인할 수 있습니다.

https://github.com/saeu5407/langchain/blob/master/retriever 예제.ipynb

langchain/retriever 예제.ipynb at master · saeu5407/langchain

 

먼저 벡터스토어를 만들어뒀다고 가정하고 바로 불러오도록 하겠습니다.

from langchain_chroma import Chroma
from langchain_openai import OpenAIEmbeddings

vectorstore = Chroma(persist_directory="./chroma_db", 
                     embedding_function=OpenAIEmbeddings())

 

Retriever

랭체인에서 VectorStore 객체는 Runnable 객체가 아니라 LangChain Expression Language chains에 연결이 되지 않지만, Retriever는 Runnable 객체이므로 바로 LCEL 체인에 적용될 수 있습니다.

 

RunnableLambda를 사용한 구현

이제 RunnableLambda를 사용하여 간단한 리트리버를 구현해보겠습니다. 벡터스토어 내의 유사도 분석 메서드를 사용했습니다.

from typing import List

from langchain_core.documents import Document
from langchain_core.runnables import RunnableLambda

retriever = RunnableLambda(vectorstore.similarity_search).bind(k=1)  # select top result

retriever.batch(["cat", "shark"])

[[Document(page_content='Cats are independent pets that often enjoy their own space.', metadata={'source': 'mammal-pets-doc'})],
 [Document(page_content='Shrimps are so delicious')]]

 

벡터스토어 as_retriever 메서드로 바로 구현

벡터스토어에는 as_retriever라는 메서드를 사용하여 간단하게 리트리버를 구현할 수 있습니다. 이때는 당연히 러너블람다로 객체를 감쌀 필요가 없습니다. 이미 러너블로 구현되어 있으니까요.

벡터스토어는 다양한 retriever 서치 방법을 구현했는데요. 유사도 기반이나 MMR 등 다양한 방법이 있으며 벡터스토어마다 방식이 다를 수 있으니 사용하는 벡터스토어의 공식 문서에서 확인해보는것이 좋을 것 같습니다.

retriever = vectorstore.as_retriever(
    search_type="similarity",
    search_kwargs={"k": 1},
)

retriever.batch(["cat", "shark"])

[[Document(page_content='Cats are independent pets that often enjoy their own space.', metadata={'source': 'mammal-pets-doc'})],
 [Document(page_content='Shrimps are so delicious')]]

 

LCEL 연결해보기

이제 retriever를 만들어보았으니 체인에 연결해보는 작업을 진행해보겠습니다.

from langchain_openai import ChatOpenAI
import os

os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
llm = ChatOpenAI(
    model_name="gpt-3.5-turbo",
    api_key=os.getenv('OPENAI_API_KEY'),
    max_tokens=1000,
)

프롬프트를 설정하고 Retriever를 연결하여 질문에 대해 리트리버를 통해 Context를 생성한 후 LLM이 응답하는 형식으로 체인을 구현하였습니다.

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough

message = """
Answer this question using the provided context only.

{question}

Context:
{context}
"""

prompt = ChatPromptTemplate.from_messages([("human", message)])

rag_chain = {"context": retriever, "question": RunnablePassthrough()} | prompt | llm

문서에서 데이터를 가져와 적절하게 대답하는 걸 확인할 수 있습니다.

response = rag_chain.invoke("tell me about cats")

print(response.content)

Cats are independent pets that often enjoy their own space.

 

Retriever에 대해서는 다양한 예제가 많으므로 공식 문서를 잘 살펴보는 것이 좋을 것 같습니다.

공식 문서 링크는 다음과 같습니다.

https://python.langchain.com/v0.2/docs/how_to/#retrievers

How-to guides | 🦜️🔗 LangChain