Simple Fusion Retriever

여러 쿼리와 여러 인덱스에서 가져온 검색 결과를 결합하는 방법을 보여주는 예시예요. 검색된 노드들은 모든 쿼리·인덱스 기준 top-k로 반환되고 중복도 제거돼요.

출처: 문서

본문

이 예시에서는 여러 쿼리와 여러 인덱스에서 가져온 검색 결과를 결합하는 방법을 살펴봅니다.

검색된 노드는 모든 쿼리와 인덱스에 걸쳐 top-k로 반환되며, 어떤 노드의 중복도 제거합니다.

import os
import openai


os.environ["OPENAI_API_KEY"] = "sk-..."
openai.api_key = os.environ["OPENAI_API_KEY"]

Setup

이 노트북에서는 문서의 서로 매우 유사한 두 페이지를 사용하며, 각각 별도의 인덱스에 저장합니다.

from llama_index.core import SimpleDirectoryReader


documents_1 = SimpleDirectoryReader(
    input_files=["../../community/integrations/vector_stores.md"]
).load_data()
documents_2 = SimpleDirectoryReader(
    input_files=["../../module_guides/storing/vector_stores.md"]
).load_data()

from llama_index.core import VectorStoreIndex


index_1 = VectorStoreIndex.from_documents(documents_1)
index_2 = VectorStoreIndex.from_documents(documents_2)

인덱스 퓨전!

이 단계에서는 인덱스들을 단일 리트리버로 융합합니다. 이 리트리버는 또한 원래 질문과 관련된 추가 쿼리를 생성해 우리 쿼리를 증강하고, 그 결과를 집계합니다.

이 설정은 쿼리를 4번 수행하는데, 원래 쿼리로 한 번, 그리고 추가로 3개의 쿼리를 생성합니다.

기본적으로 추가 쿼리를 생성하는 데 다음 프롬프트를 사용합니다:

QUERY_GEN_PROMPT = (
    "You are a helpful assistant that generates multiple search queries based on a "
    "single input query. Generate {num_queries} search queries, one on each line, "
    "related to the following input query:\n"
    "Query: {query}\n"
    "Queries:\n"
)

from llama_index.core.retrievers import QueryFusionRetriever


retriever = QueryFusionRetriever(
    [index_1.as_retriever(), index_2.as_retriever()],
    similarity_top_k=2,
    num_queries=4,  # set this to 1 to disable query generation
    use_async=True,
    verbose=True,
    # query_gen_prompt="...",  # we could override the query generation prompt here
)

# apply nested async to run in a notebook
import nest_asyncio


nest_asyncio.apply()

nodes_with_scores = retriever.retrieve("How do I setup a chroma vector store?")

for node in nodes_with_scores:
    print(f"Score: {node.score:.2f} - {node.text[:100]}...")

Query Engine에서 사용하기!

이제 리트리버를 쿼리 엔진에 연결해 자연어 응답을 합성할 수 있습니다.

from llama_index.core.query_engine import RetrieverQueryEngine


query_engine = RetrieverQueryEngine.from_args(retriever)

response = query_engine.query(
    "How do I setup a chroma vector store? Can you give an example?"
)

from llama_index.core.response.notebook_utils import display_response


display_response(response)

QueryFusionRetriever가 생성한 추가 쿼리들까지 활용해 두 인덱스의 top-k 노드를 결합하므로, 원본 쿼리만으로 검색했을 때보다 더 풍부한 컨텍스트 위에서 응답을 합성할 수 있습니다.

Simple Fusion Retriever