LangChain을 이용한 RAG - (6) 여러 문서 RAG 태우기

2024.06.03 - [AWS] - [AWS Summit Seoul 2024] LLM의 프롬프트 엔지니어링

2024.07.22 - [데이터 분석/LLM] - LangChain을 이용한 RAG - (5) 생성 편

아이디어

이전에 AWS Summit에서 고급 RAG라고 해서, RAG로 문서와 관계 없는 질문을 하면 성능이 떨어지기 때문에 필요할 때만 RAG를 사용하도록 하라고 했었다.

참고: [발표자료][생성형 AI 및 기계 학습] Amazon Bedrock을 활용한 프롬프트 엔지니어링 모범사례 (awscloud.com)

 

나는 이에 대해 아래와 같은 아이디어를 냈었는데 직접 구현해보기로 했다.

 

1. 주제가 다른 문서마다 벡터스토어를 따로 만들기
2. 사용자의 질문이 들어오면 LLM을 통해 어느 벡터스토어를 거쳐야할지 정하기. 이를 위해 각 벡터스토어마다 꼼꼼한 설명을 적어 프롬프트로 전달해야 함.
3. LLM이 정해준 벡터스토어를 거쳐 답변 내놓기

 

 

구현

Chain 구성

검색기(retriever)가 바뀔 때마다 chain을 선언해줘야 하므로 chain을 만드는 함수를 만들었다.

from langchain_aws import ChatBedrock
from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain.schema.runnable import RunnablePassthrough
from langchain import hub
import boto3

prompt = hub.pull("rlm/rag-prompt")
client = boto3.client('bedrock-runtime', 
    region_name="us-east-1",
)

# Model
llm = ChatBedrock(
    model_id="anthropic.claude-3-5-sonnet-20240620-v1:0",
    model_kwargs={"temperature": 0.1},
    streaming=True,
    client=client,
)


def make_chain(prompt=prompt, llm=llm, retriever=None):
    # Chain
    chain = (
        {'context': retriever | format_docs, 'question': RunnablePassthrough()}
        | prompt
        | llm
        | StrOutputParser()
    )

    return chain

아예 체인을 미리 구성해두면 되지 않을까도 싶었는데 prompt와 llm을 바꿔가며 실험할 수 있는 게 더 좋지 않을까 싶어서 이렇게 정의했다.

 

 

검색기 선택

LLM이 유저의 질문을 보고 검색기를 선택하도록 프롬프트를 작성하였다.

def select_retriever(query, llm=llm):
    prompt = f"""너는 사용자의 질문을 받고 해당 질문이 A와 B 문서 중에 어디에 더 가까운 지 하나를 선택하여 "A" 또는 "B"라는 답을 해야한다.
    만약 둘 다 관계가 없는 경우 "None"이라는 답을 내놓는다.
    <documents>
    A: [한정 수량 판매 제품을 사서 비싼 값에 파는 "리셀" 행위는 소비자의 정당한 권리이며, 암표상과는 다른 문제이다.]라는 주제로 찬성과 반대측이 토론한 내용
    B: AI와 관련한 내용이 담긴 학술지
    </documents>
    

    <question>
    {query}
    </question>"""

    return llm.invoke(prompt).content

이렇게 하면 내가 분류기 모델을 만들 필요 없이 LLM이 알아서 분류를 해준다.

1) 안녕? → None

2) 되팔이는 나쁜 짓이야 → A

3) AI기술 발전이 엔비디아 주가를 어디까지 높일까 → B

 

이후 서로 다른 문서에 대해 벡터스토어를 각각 생성해주고 검색기로 만들어준다.

 

 

답변 생성

retrievers = dict(A=retriever_a, B=retriever_b)

def chat_with_rag(query, prompt=prompt, llm=llm, retrievers=retrievers):
    retriever = select_retriever(query, llm=llm)
    answer = ''
    if retriever == 'None':
        for chunk in llm.stream(query):
            answer += chunk.content
            print(chunk.content, end='')
        return answer
        
    retriever = retrievers[retriever]
    chain = make_chain(prompt, llm, retriever)

    for chunk in chain.stream(query):
        answer += chunk
        print(chunk, end='')

    return answer

1. 앞서 정의한 `select_retriever` 함수를 통해 사용자 질문(`query`)를 `llm`에 보내 검색기를 선택한다. 

2. 만약 선택이 `None`이라면 검색기를 통하지 않고 일반적인 `llm`의 답변을 얻는다.

3. 만약 선택이 `A`나 `B`라면 딕셔너리로부터 검색기를 선택하고 체인을 만든 후 답변을 얻는다.

 

 

1) 안녕 → 안녕하세요! 무엇을 도와드릴까요?

 

2) 되팔이는 나쁜 짓이야 → 되팔이(리셀)에 대해서는 찬성과 반대 의견이 모두 있습니다. 반대 입장에서는 창작자의 권리 침해와 암표상과의 유사성을 지적하며 시장의 공정성을 해친다고 봅니다. 반면 찬성 입장에서는 법적 테두리 내의 활동이며 시장 메커니즘을 통해 문제점을 개선할 수 있다고 주장합니다.

 

3) AI기술 발전이 엔비디아 주가를 어디까지 높일까 → AI 기술의 발전은 엔비디아의 주가에 긍정적인 영향을 미칠 것으로 예상됩니다. IDC에 따르면 AI와 자동화 기술에 대한 기업들의 투자 의지가 확고하며, AI 소프트웨어 시장은 2027년까지 2790억 달러 규모로 성장할 전망입니다. 또한, AI 기술을 보유한 근로자의 임금이 평균 21% 높다는 연구 결과는 AI 기술의 경제적 가치와 수요가 높음을 시사하며, 이는 엔비디아와 같은 AI 관련 기업의 주가 상승으로 이어질 수 있습니다.

 

 

 

생각보다 간단하게 구현이 가능했다.