PropertyGraph에서 커스텀 Retriever 정의하기

PropertyGraph를 위한 커스텀 retriever를 정의하는 노트북이에요. 표준 그래프 검색기보다는 복잡하지만 검색 과정을 세밀하게 제어할 수 있고, 애플리케이션의 요구에 딱 맞게 커스터마이징할 수 있습니다. 벡터 검색과 text-to-Cypher 검색을 함께 수행하고, 그 결과를 reranking 모듈로 통합하는 고급 검색 워크플로도 다룹니다.

출처: 문서

본문

Cohere reranker를 사용하므로 cohere API key가 필요해요.

필요한 패키지를 설치합니다.

%pip install llama-index-core
%pip install llama-index-llms-mistralai
%pip install llama-index-embeddings-mistralai
%pip install llama-index-graph-stores-neo4j
%pip install llama-index-postprocessor-cohere-rerank

import nest_asyncio
nest_asyncio.apply()
from IPython.display import Markdown, display

설정 (Setup)

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

from llama_index.embeddings.mistralai import MistralAIEmbedding
from llama_index.llms.mistralai import MistralAI

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

데이터 다운로드·로드

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

from llama_index.core import SimpleDirectoryReader

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

Docker 설정

APOC 플러그인을 활성화한 Neo4j 컨테이너를 띄웁니다. 첫 실행 시 로그인하고 비밀번호를 설정해야 해요. (neo4j / neo4j)

!docker run \
    -p 7474:7474 -p 7687:7687 \
    -v $PWD/data:/data -v $PWD/plugins:/plugins \
    --name neo4j-apoc \
    -e NEO4J_apoc_export_file_enabled=true \
    -e NEO4J_apoc_import_file_enabled=true \
    -e NEO4J_apoc_import_file_use__neo4j__config=true \
    -e NEO4JLABS_PLUGINS=\[\"apoc\"\] \
    neo4j:latest

Neo4j GraphStore 설정

from llama_index.graph_stores.neo4j import Neo4jPropertyGraphStore

graph_store = Neo4jPropertyGraphStore(
    username="neo4j",
    password="llamaindex",
    url="bolt://localhost:7687",
)

PropertyGraphIndex 구성

from llama_index.core import PropertyGraphIndex

index = PropertyGraphIndex.from_documents(
    documents,
    llm=llm,
    embed_model=embed_model,
    property_graph_store=graph_store,
    show_progress=True,
)

CustomRetriever (커스텀 검색기)

VectorContextRetriever·TextToCypherRetriever·Reranker를 결합한 커스텀 검색기를 정의합니다.

from llama_index.core.retrievers import (
    CustomPGRetriever,
    VectorContextRetriever,
    TextToCypherRetriever,
)
from llama_index.core.graph_stores import PropertyGraphStore
from llama_index.core.vector_stores.types import VectorStore
from llama_index.core.embeddings import BaseEmbedding
from llama_index.core.prompts import PromptTemplate
from llama_index.core.llms import LLM
from llama_index.postprocessor.cohere_rerank import CohereRerank
from typing import Optional, Any, Union

class CustomRetriever(CustomPGRetriever):
    """Custom retriever with cohere reranking."""

    def __init__(
        self,
        ## vector context retriever params
        embed_model: Optional[BaseEmbedding] = None,
        vector_store: Optional[VectorStore] = None,
        similarity_top_k: int = 4,
        path_depth: int = 1,
        ## text-to-cypher params
        llm: Optional[LLM] = None,
        text_to_cypher_template: Optional[Union[PromptTemplate, str]] = None,
        ## cohere reranker params
        cohere_api_key: Optional[str] = None,
        cohere_top_n: int = 2,
        **kwargs: Any,
    ) -> None:
        """Uses any kwargs passed in from class constructor."""
        self.vector_retriever = VectorContextRetriever(
            self.graph_store,
            include_text=self.include_text,
            embed_model=embed_model,
            vector_store=vector_store,
            similarity_top_k=similarity_top_k,
            path_depth=path_depth,
        )
        self.cypher_retriever = TextToCypherRetriever(
            self.graph_store,
            llm=llm,
            text_to_cypher_template=text_to_cypher_template
            ## NOTE: you can attach other parameters here if you'd like
        )
        self.reranker = CohereRerank(
            api_key=cohere_api_key, top_n=cohere_top_n
        )

    def custom_retrieve(self, query_str: str) -> str:
        """Define custom retriever with reranking.
        Could return `str`, `TextNode`, `NodeWithScore`, or a list of those.
        """
        nodes_1 = self.vector_retriever.retrieve(query_str)
        nodes_2 = self.cypher_retriever.retrieve(query_str)

        reranked_nodes = self.reranker.postprocess_nodes(
            nodes_1 + nodes_2, query_str=query_str
        )

        ## TMP: please change
        final_text = "\n\n".join(
            [n.get_content(metadata_mode="llm") for n in reranked_nodes]
        )
        return final_text
        # optional async method
        # async def acustom_retrieve(self, query_str: str) -> str:
        #     ...

from llama_index.core import Settings
Settings.llm = llm
Settings.embed_model=embed_model

정의한 커스텀 검색기를 인스턴스화합니다. cohere API key가 필요해요.

custom_sub_retriever = CustomRetriever(
    index.property_graph_store,
    include_text=True,
    vector_store=index.vector_store,
    cohere_api_key="YOUR COHERE API KEY",
)

QueryEngine

커스텀 검색기를 붙인 RetrieverQueryEngine을 만들어 질의합니다.

from llama_index.core.query_engine import RetrieverQueryEngine

query_engine = RetrieverQueryEngine.from_args(
    index.as_retriever(sub_retrievers=[custom_sub_retriever]), llm=llm
)

response = query_engine.query("What did author do at Interleaf?")
display(Markdown(f"{response.response}"))

더 알아보기 (Learn more)

LlamaIndex PropertyGraph 문서 — 지식 그래프 인덱스 가이드
CustomPGRetriever — 커스텀 PropertyGraph 검색기 베이스 클래스
TextToCypherRetriever — 질의를 Cypher로 변환해 검색하는 검색기
CohereRerank — Cohere 기반 재순위화 포스트프로세서