LlamaIndex와 Mistral AI로 Router Query Engine 만들기

LlamaIndex와 Mistral AI로 Router Query Engine 만들기

사용자 질의를 질문 유형에 따라 적절한 인덱스로 라우팅하는 RouterQueryEngine을 LlamaIndex와 MistralAI로 구현하는 노트북이에요. 특정 문맥 질문과 요약 질문을 자동으로 구분해 다른 인덱스에 보내므로 더 정확한 답을 얻을 수 있습니다.

출처: 문서

본문

VectorStoreIndex는 특정 문맥과 관련된 질의를 처리하도록 설계됐고, SummaryIndex는 요약 질의에 최적화되어 있어요. 그런데 실제 상황에서는 사용자의 질의가 문맥 특화 응답을 요구할 수도, 요약을 요구할 수도 있죠. 그래서 시스템이 질의 유형에 따라 적절한 인덱스로 라우팅해 주는 것이 필요합니다. 이 노트북에서는 RouterQueryEngine을 이용해 질의 유형에 맞는 인덱스로 사용자 질의를 보내 봅니다.

설치와 API 키 설정부터 시작합니다.

!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()

필요한 모듈을 임포트하고 LLM·임베딩 모델을 설정합니다.

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

llm = MistralAI(model='mistral-large')
embed_model = MistralAIEmbedding()

Settings.llm = llm
Settings.embed_model = embed_model

이번 예제는 Uber 10K SEC 보고서를 사용합니다.

!mkdir -p 'data/10k/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10k/uber_2021.pdf' -O 'data/10k/uber_2021.pdf'

uber_docs = SimpleDirectoryReader(input_files=["./data/10k/uber_2021.pdf"]).load_data()
  • VectorStoreIndex → 특정 문맥 질의
  • SummaryIndex → 요약 질의

인덱스와 쿼리 엔진을 만듭니다.

uber_vector_index = VectorStoreIndex.from_documents(uber_docs)
uber_summary_index = VectorStoreIndex.from_documents(uber_docs)

uber_vector_query_engine = uber_vector_index.as_query_engine(similarity_top_k = 5)
uber_summary_query_engine = uber_summary_index.as_query_engine()

두 쿼리 엔진을 각각 도구(tool)로 만들고 이름과 설명을 붙입니다.

query_engine_tools = [
    QueryEngineTool(
        query_engine=uber_vector_query_engine,
        metadata=ToolMetadata(
            name="vector_engine",
            description=(
                "Provides information about Uber financials for year 2021."
            ),
        ),
    ),
    QueryEngineTool(
        query_engine=uber_summary_query_engine,
        metadata=ToolMetadata(
            name="summary_engine",
            description=(
                "Provides Summary about Uber financials for year 2021."
            ),
        ),
    ),
]

RouterQueryEngine을 만들고 LLMSingleSelector로 질의 유형을 판별하게 합니다.

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

요약 질의는 SummaryIndex를 사용하는 걸 확인할 수 있어요.

response = query_engine.query("What is the summary of the Uber Financials in 2021?")
print(response)

특정 문맥 질의는 VectorStoreIndex로 답하는 걸 볼 수 있습니다.

response = query_engine.query("What is the the revenue of Uber in 2021?")
print(response)

더 알아보기 (Learn more)

  • LlamaIndex RouterQueryEngine 문서 — 질의 유형별 라우팅 쿼리 엔진
  • LLMSingleSelector — LLM 기반 단일 셀렉터
  • QueryEngineTool / ToolMetadata — 쿼리 엔진을 도구로 감싸는 구성 요소