Adaptive RAG with LlamaIndex

Adaptive RAG with LlamaIndex

복잡한 질의와 단순한 질의를 구분해 처리하는 Adaptive RAG 방식을 LlamaIndex로 구현하는 노트북이에요. RouterQueryEngine과 MistralAI의 FunctionCalling을 사용해 질의의 복잡도에 따라 도구나 인덱스를 다르게 호출합니다.

출처: 문서

본문

사용자 질의는 대체로 복잡한 질의일 수도, 단순한 질의일 수도 있어요. 단순한 질의를 처리하는 데 항상 복잡한 RAG 시스템이 필요한 건 아니죠. Adaptive RAG는 복잡한 질의와 단순한 질의를 처리하는 접근을 제안합니다. 이 노트북에서는 2020·2021·2022년 Lyft의 10k SEC 보고서를 중심으로 복잡·단순 질의를 구분 처리하는 Adaptive RAG와 유사한 방식을 구현합니다.

우리 접근은 RouterQueryEngine과 MistralAI의 FunctionCalling으로 질의 복잡도에 따라 다른 도구나 인덱스를 호출하는 방식이에요.

• 복잡한 질의: 여러 문서의 문맥이 필요한 여러 도구를 활용합니다.
• 단순한 질의: 단일 문서의 문맥이 필요한 단일 도구를 쓰거나, LLM을 직접 사용해 답합니다.

따르는 단계는 다음과 같아요.

1. 데이터 다운로드
2. 데이터 로드
3. 3개 문서의 인덱스 생성
4. 문서와 LLM으로 쿼리 엔진 생성
5. 복잡한 질의용 FunctionCallingAgentWorker 초기화
6. 도구 생성
7. RouterQueryEngine 생성 — 질의 복잡도에 따라 라우팅
8. 질의

설치와 설정

!pip install llama-index
!pip install llama-index-llms-mistralai
!pip install llama-index-embeddings-mistralai

import os
os.environ['MISTRAL_API_KEY'] = '<YOUR MISTRAL API KEY>'

import nest_asyncio
nest_asyncio.apply()

from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
from llama_index.llms.mistralai import MistralAI
from llama_index.embeddings.mistralai import MistralAIEmbedding
from llama_index.core import Settings
from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama_index.core.query_engine.router_query_engine import RouterQueryEngine
from llama_index.core.selectors.llm_selectors import LLMSingleSelector

# Note: Only `mistral-large-latest` supports function calling
llm = MistralAI(model='mistral-large-latest')
embed_model = MistralAIEmbedding()

Settings.llm = llm
Settings.embed_model = embed_model

로깅 (Logging)

참고: 이 설정은 Jupyter 노트북에서만 필요해요. Jupyter는 백그라운드에서 이벤트 루프를 돌리는데, 비동기 질의를 위해 새 이벤트 루프를 시작하면 중첩된 이벤트 루프가 생깁니다. 보통 허용되지 않지만 편의를 위해 nest_asyncio로 허용해요.

# NOTE: This is ONLY necessary in jupyter notebook.
# Details: Jupyter runs an event-loop behind the scenes.
#          This results in nested event-loops when we start an event-loop to make async queries.
#          This is normally not allowed, we use nest_asyncio to allow it for convenience.
import nest_asyncio
nest_asyncio.apply()

import logging
import sys

# Set up the root logger
logger = logging.getLogger()
logger.setLevel(logging.INFO)  # Set logger level to INFO

# Clear out any existing handlers
logger.handlers = []

# Set up the StreamHandler to output to sys.stdout (Colab's output)
handler = logging.StreamHandler(sys.stdout)
handler.setLevel(logging.INFO)  # Set handler level to INFO

# Add the handler to the logger
logger.addHandler(handler)

from IPython.display import display, HTML

데이터 다운로드·로드

2020·2021·2022년 Lyft 10k SEC 보고서를 내려받습니다.

!wget "https://www.dropbox.com/scl/fi/ywc29qvt66s8i97h1taci/lyft-10k-2020.pdf?rlkey=d7bru2jno7398imeirn09fey5&dl=0" -q -O ./lyft_10k_2020.pdf
!wget "https://www.dropbox.com/scl/fi/lpmmki7a9a14s1l5ef7ep/lyft-10k-2021.pdf?rlkey=ud5cwlfotrii6r5jjag1o3hvm&dl=0" -q -O ./lyft_10k_2021.pdf
!wget "https://www.dropbox.com/scl/fi/iffbbnbw9h7shqnnot5es/lyft-10k-2022.pdf?rlkey=grkdgxcrib60oegtp4jn8hpl8&dl=0" -q -O ./lyft_10k_2022.pdf

# Lyft 2020 docs
lyft_2020_docs = SimpleDirectoryReader(input_files=["./lyft_10k_2020.pdf"]).load_data()

# Lyft 2021 docs
lyft_2021_docs = SimpleDirectoryReader(input_files=["./lyft_10k_2021.pdf"]).load_data()

# Lyft 2022 docs
lyft_2022_docs = SimpleDirectoryReader(input_files=["./lyft_10k_2022.pdf"]).load_data()

인덱스와 쿼리 엔진 만들기

각 연도 문서에 대한 인덱스를 만듭니다.

# Index on Lyft 2020 Document
lyft_2020_index = VectorStoreIndex.from_documents(lyft_2020_docs)

# Index on Lyft 2021 Document
lyft_2021_index = VectorStoreIndex.from_documents(lyft_2021_docs)

# Index on Lyft 2022 Document
lyft_2022_index = VectorStoreIndex.from_documents(lyft_2022_docs)

# Query Engine on Lyft 2020 Docs Index
lyft_2020_query_engine = lyft_2020_index.as_query_engine(similarity_top_k=5)

# Query Engine on Lyft 2021 Docs Index
lyft_2021_query_engine = lyft_2021_index.as_query_engine(similarity_top_k=5)

# Query Engine on Lyft 2022 Docs Index
lyft_2022_query_engine = lyft_2022_index.as_query_engine(similarity_top_k=5)

일반 질의에 답하기 위해 LLM을 그대로 사용하는 쿼리 엔진(LLMQueryEngine)도 만듭니다.

from llama_index.core.query_engine import CustomQueryEngine

class LLMQueryEngine(CustomQueryEngine):
    """RAG String Query Engine."""
    llm: llm

    def custom_query(self, query_str: str):
        response = self.llm.complete(query_str)
        return str(response)

llm_query_engine = LLMQueryEngine(llm=llm)

FunctionCallingAgentWorker 초기화

복잡한 질의는 여러 문서를 다루므로 에이전트가 필요합니다. 이 도구들은 여러 문서를 포함하는 복잡한 질의를 답하는 데 사용됩니다.

# These tools are used to answer complex queries involving multiple documents.
query_engine_tools = [
    QueryEngineTool(
        query_engine=lyft_2020_query_engine,
        metadata=ToolMetadata(
            name="lyft_2020_10k_form",
            description="Annual report of Lyft's financial activities in 2020",
        ),
    ),
    QueryEngineTool(
        query_engine=lyft_2021_query_engine,
        metadata=ToolMetadata(
            name="lyft_2021_10k_form",
            description="Annual report of Lyft's financial activities in 2021",
        ),
    ),
    QueryEngineTool(
        query_engine=lyft_2022_query_engine,
        metadata=ToolMetadata(
            name="lyft_2022_10k_form",
            description="Annual report of Lyft's financial activities in 2022",
        ),
    ),
]

from llama_index.core.agent import FunctionCallingAgentWorker
from llama_index.core.agent import AgentRunner

agent_worker = FunctionCallingAgentWorker.from_tools(
    query_engine_tools,
    llm=llm,
    verbose=True,
    allow_parallel_tool_calls=True,
)
agent = AgentRunner(agent_worker)

도구 만들기

앞서 만든 쿼리 엔진들과 FunctionCallingAgentWorker로 도구를 만듭니다. 복잡한 질의용 도구에는 여러 연도를 아우르는 에이전트 도구(lyft_2020_2021_2022_10k_form)가 포함되고, 일반 질의용 도구(general_queries)도 포함됩니다.

query_engine_tools = [
    QueryEngineTool(
        query_engine=lyft_2020_query_engine,
        metadata=ToolMetadata(
            name="lyft_2020_10k_form",
            description="Queries related to only 2020 Lyft's financial activities.",
        ),
    ),
    QueryEngineTool(
        query_engine=lyft_2021_query_engine,
        metadata=ToolMetadata(
            name="lyft_2021_10k_form",
            description="Queries related to only 2021 Lyft's financial activities.",
        ),
    ),
    QueryEngineTool(
        query_engine=lyft_2022_query_engine,
        metadata=ToolMetadata(
            name="lyft_2022_10k_form",
            description="Queries related to only 2022 Lyft's financial activities.",
        ),
    ),
    QueryEngineTool(
        query_engine=agent,
        metadata=ToolMetadata(
            name="lyft_2020_2021_2022_10k_form",
            description=(
                "Useful for queries that span multiple years from 2020 to 2022 for Lyft's financial activities."
            ),
        ),
    ),
    QueryEngineTool(
        query_engine=llm_query_engine,
        metadata=ToolMetadata(
            name="general_queries",
            description=(
                "Provides information about general queries other than lyft."
            ),
        ),
    ),
]

RouterQueryEngine 만들기

RouterQueryEngine은 질의의 복잡도에 따라 사용자가 고른 도구 중 하나로 질의를 라우팅합니다.

from llama_index.core.query_engine import RouterQueryEngine
from llama_index.core.selectors import LLMSingleSelector

query_engine = RouterQueryEngine(
    selector=LLMSingleSelector.from_defaults(),
    query_engine_tools=query_engine_tools,
    verbose = True
)

질의하기 (Querying)

일반 질의라서 LLM 도구를 사용한 걸 볼 수 있어요.

response = query_engine.query("What is the capital of France?")
display(HTML(f'<p style="font-size:20px">{response.response}</p>'))

lyft_2022 도구로 답합니다.

response = query_engine.query("What did Lyft do in R&D in 2022?")
display(HTML(f'<p style="font-size:20px">{response.response}</p>'))

lyft_2021 도구로 답합니다.

response = query_engine.query("What did Lyft do in R&D in 2021?")
display(HTML(f'<p style="font-size:20px">{response.response}</p>'))

lyft_2020 도구로 답합니다.

response = query_engine.query("What did Lyft do in R&D in 2020?")
display(HTML(f'<p style="font-size:20px">{response.response}</p>'))

이제 여러 도구가 필요한 질의를 테스트합니다. FunctionCallingAgent로 lyft_2020과 lyft_2022 도구를 사용한 걸 볼 수 있어요.

response = query_engine.query("What did Lyft do in R&D in 2022 vs 2020?")
display(HTML(f'<p style="font-size:20px">{response.response}</p>'))

FunctionCallingAgent로 lyft_2020과 lyft_2021 도구를 사용합니다.

response = query_engine.query("What did Lyft do in R&D in 2020 vs 2021?")
display(HTML(f'<p style="font-size:20px">{response.response}</p>'))

FunctionCallingAgent로 lyft_2020, lyft_2021, lyft_2022 도구를 모두 사용합니다.

response = query_engine.query("What did Lyft do in R&D in 2022 vs 2021 vs 2020?")
display(HTML(f'<p style="font-size:20px">{response.response}</p>'))

더 알아보기 (Learn more)

• LlamaIndex RouterQueryEngine 문서 — 질의 라우팅 쿼리 엔진
• FunctionCallingAgentWorker / AgentRunner — 함수 호출 기반 에이전트
• CustomQueryEngine — 커스텀 쿼리 엔진 정의
• mistral-large-latest — function calling을 지원하는 사용 모델