Router Retriever

주어진 쿼리를 실행하기 위해 하나 이상의 후보 리트리버를 선택하는 커스텀 라우터 리트리버를 정의하는 가이드예요. LLM이 어떤 검색 도구를 쓸지 동적으로 결정하는 방식이 핵심이에요.

출처: 문서

본문

이 가이드에서는 주어진 쿼리를 실행하기 위해 하나 이상의 후보 리트리버를 선택하는 커스텀 라우터 리트리버를 정의합니다.

라우터(BaseSelector) 모듈은 LLM을 사용해 어떤 기본 검색 도구를 사용할지 동적으로 결정합니다. 이는 다양한 데이터 소스 중에서 하나를 선택하는 데 유용할 수 있습니다. 또한 (multi-selector 모듈을 사용한다면) 다양한 데이터 소스 전반의 검색 결과를 집계하는 데도 유용할 수 있습니다.

이 노트북은 RouterQueryEngine 노트북과 매우 유사합니다.

설정 (Setup)

이 노트북을 colab에서 여는 경우 LlamaIndex 🦙를 설치해야 할 수 있습니다.

%pip install llama-index-llms-openai

!pip install llama-index

# 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


logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().handlers = []
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))


from llama_index.core import (
    VectorStoreIndex,
    SimpleDirectoryReader,
    StorageContext,
    SimpleKeywordTableIndex,
)
from llama_index.core import SummaryIndex
from llama_index.core.node_parser import SentenceSplitter
from llama_index.llms.openai import OpenAI

데이터 다운로드

!mkdir -p 'data/paul_graham/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'

데이터 로드

먼저 Document를 Node 집합으로 변환하고 DocumentStore에 삽입하는 방법을 보여줍니다.

# load documents
documents = SimpleDirectoryReader("./data/paul_graham/").load_data()

# initialize LLM + splitter
llm = OpenAI(model="gpt-4")
splitter = SentenceSplitter(chunk_size=1024)
nodes = splitter.get_nodes_from_documents(documents)

# initialize storage context (by default it's in-memory)
storage_context = StorageContext.from_defaults()
storage_context.docstore.add_documents(nodes)

# define
summary_index = SummaryIndex(nodes, storage_context=storage_context)
vector_index = VectorStoreIndex(nodes, storage_context=storage_context)
keyword_index = SimpleKeywordTableIndex(nodes, storage_context=storage_context)

list_retriever = summary_index.as_retriever()
vector_retriever = vector_index.as_retriever()
keyword_retriever = keyword_index.as_retriever()

from llama_index.core.tools import RetrieverTool


list_tool = RetrieverTool.from_defaults(
    retriever=list_retriever,
    description=(
        "Will retrieve all context from Paul Graham's essay on What I Worked"
        " On. Don't use if the question only requires more specific context."
    ),
)
vector_tool = RetrieverTool.from_defaults(
    retriever=vector_retriever,
    description=(
        "Useful for retrieving specific context from Paul Graham essay on What"
        " I Worked On."
    ),
)
keyword_tool = RetrieverTool.from_defaults(
    retriever=keyword_retriever,
    description=(
        "Useful for retrieving specific context from Paul Graham essay on What"
        " I Worked On (using entities mentioned in query)"
    ),
)

라우팅용 Selector 모듈 정의

사용할 수 있는 selectors가 여러 가지 있으며, 각각 고유한 특성을 가집니다.

LLM selectors는 LLM을 사용해 파싱할 JSON을 출력하고, 해당하는 인덱스들이 질의됩니다.

Pydantic selectors(현재 gpt-4-0613과 gpt-3.5-turbo-0613만 지원, 기본값)는 raw JSON을 파싱하는 대신 OpenAI Function Call API를 사용해 pydantic selection 객체를 생성합니다.

여기서는 PydanticSingleSelector/PydanticMultiSelector를 사용하지만 LLM 버전도 사용할 수 있습니다.

from llama_index.core.selectors import LLMSingleSelector, LLMMultiSelector
from llama_index.core.selectors import (
    PydanticMultiSelector,
    PydanticSingleSelector,
)
from llama_index.core.retrievers import RouterRetriever
from llama_index.core.response.notebook_utils import display_source_node

PydanticSingleSelector

retriever = RouterRetriever(
    selector=PydanticSingleSelector.from_defaults(llm=llm),
    retriever_tools=[
        list_tool,
        vector_tool,
    ],
)

# will retrieve all context from the author's life
nodes = retriever.retrieve(
    "Can you give me all the context regarding the author's life?"
)
for node in nodes:
    display_source_node(node)

PydanticSingleSelector는 단 하나의 리트리버만 선택합니다. 위 쿼리처럼 "저자의 삶 전체 맥락"을 요청하면 전체 컨텍스트를 가져오는 list_tool(인덱스 0)을 선택합니다. 반면 구체적 사실을 묻는 "RISD 이후에 Paul Graham이 무엇을 했나요?" 같은 질문은 구체적 컨텍스트에 맞는 vector_tool(인덱스 1)을 선택해 상위 유사도 노드를 반환합니다.

nodes = retriever.retrieve("What did Paul Graham do after RISD?")
for node in nodes:
    display_source_node(node)

PydanticMultiSelector

retriever = RouterRetriever(
    selector=PydanticMultiSelector.from_defaults(llm=llm),
    retriever_tools=[list_tool, vector_tool, keyword_tool],
)

nodes = retriever.retrieve(
    "What were noteable events from the authors time at Interleaf and YC?"
)
for node in nodes:
    display_source_node(node)

PydanticMultiSelector는 여러 리트리버를 동시에 선택할 수 있습니다. 위 쿼리처럼 "Interleaf와 YC에서의 주목할 만한 사건"을 묻는 질의에서는 vector_tool(인덱스 1)과 keyword_tool(인덱스 2)을 함께 선택해 두 도구의 결과를 결합합니다. 쿼리에서 'interleaf'와 'yc' 같은 엔티티가 추출되어 키워드 검색에 활용되는 것을 볼 수 있습니다.

nodes = retriever.retrieve(
    "What were noteable events from the authors time at Interleaf and YC?"
)
for node in nodes:
    display_source_node(node)

비동기 검색도 동일하게 aretrieve를 사용해 수행할 수 있습니다.

nodes = await retriever.aretrieve(
    "What were noteable events from the authors time at Interleaf and YC?"
)
for node in nodes:
    display_source_node(node)

Router Retriever