Property Graph의 Extractor와 Retriever

Property Graph의 Extractor와 Retriever

PropertyGraph 인덱스에서 추출기(extractor)와 검색기(retriever)를 정의하는 방법을 다루는 노트북이에요. PropertyGraph를 구축하고 질의하는 다양한 추출기·검색기 조합을 살펴봅니다.

출처: 문서

본문

Property graph는 라벨이 있는 노드(엔티티 카테고리나 텍스트 라벨 같은)와 속성(메타데이터)을 관계로 연결해 구조화된 경로(트리플)를 이룬 모음이에요. LlamaIndex에서 PropertyGraphIndex는 그래프 구축과 그래프 질의라는 두 가지 핵심 역할을 합니다.

from llama_index.core import PropertyGraphIndex

# create
index = PropertyGraphIndex.from_documents(
    documents,
)

# use
retriever = index.as_retriever(
    include_text=True,  # include source chunk with matching paths
    similarity_top_k=2,  # top k for vector kg node retrieval
)
nodes = retriever.retrieve("<QUERY>")

query_engine = index.as_query_engine(
    include_text=True,  # include source chunk with matching paths
    similarity_top_k=2,  # top k for vector kg node retrieval
)
response = query_engine.query("<QUERY>")

# save and load
index.storage_context.persist(persist_dir="./storage")
from llama_index.core import StorageContext, load_index_from_storage
index = load_index_from_storage(
    StorageContext.from_defaults(persist_dir="./storage")
)

# loading from existing graph store (and optional vector store)
# load from existing graph/vector store
index = PropertyGraphIndex.from_existing(
    property_graph_store=graph_store, vector_store=vector_store, llm=llm
)

그래프 구축과 사용

PropertyGraph 구축은 각 청크에 일련의 지식 그래프 추출기를 실행하고, 엔티티와 관계를 각 노드의 메타데이터로 붙이는 과정입니다. 필요한 만큼 여러 추출기를 사용할 수 있고, 모두 적용됩니다.

index = PropertyGraphIndex.from_documents(
    documents,
    kg_extractors=[extractor1, extractor2, ...],
)

# insert additional documents / nodes
index.insert(document)
index.insert_nodes(nodes)

추출기를 지정하지 않으면 기본값은 SimpleLLMPathExtractor와 ImplicitPathExtractor입니다.

SimpleLLMPathExtractor

LLM을 사용해 짧은 문장을 추출하고 (entity1, relation, entity2) 형식의 단일 홉 경로를 파싱합니다. 원하면 프롬프트와 경로 파싱에 쓰는 함수를 모두 커스터마이즈할 수 있어요. 여기 간단한(다소 단순한) 예시가 있습니다.

from llama_index.core.indices.property_graph import SimpleLLMPathExtractor

kg_extractor = SimpleLLMPathExtractor(
    llm=llm,
    max_paths_per_chunk=10,
    num_workers=4,
    show_progress=False,
)

prompt = (
    "Some text is provided below. Given the text, extract up to "
    "{max_paths_per_chunk} "
    "knowledge triples in the form of `subject,predicate,object` on each line. Avoid stopwords.\n"
)

def parse_fn(response_str: str) -> List[Tuple[str, str, str]]:
    lines = response_str.split("\n")
    triples = [line.split(",") for line in lines]
    return triples

kg_extractor = SimpleLLMPathExtractor(
    llm=llm,
    extract_prompt=prompt,
    parse_fn=parse_fn,
)

ImplicitPathExtractor

각 LlamaIndex 노드 객체의 node.relationships 속성을 사용해 경로를 추출합니다. 이 추출 과정은 이미 노드 객체에 존재하는 속성을 파싱할 뿐이므로 LLM이나 임베딩 모델이 필요 없어요.

from llama_index.core.indices.property_graph import ImplicitPathExtractor

kg_extractor = ImplicitPathExtractor()

SchemaLLMPathExtractor

허용되는 엔티티·관계·그들 사이의 연결을 지정하는 엄격한 스키마를 따라 경로를 추출합니다. Pydantic, LLM의 구조화 출력, 그리고 지능적인 검증을 사용해 스키마를 동적으로 정의하고 각 경로(트리플)의 추출을 검증합니다.

from typing import Literal
from llama_index.core.indices.property_graph import SchemaLLMPathExtractor

# recommended uppercase, underscore separated
entities = Literal["PERSON", "PLACE", "ORGANIZATION"]
relations = Literal["PART_OF", "HAS", "WORKED_AT"]

# schema = {
#     "PERSON": ["PART_OF", "HAS", "WORKED_AT"],
#     "PLACE": ["PART_OF", "HAS"],
#     "ORGANIZATION": ["WORKED_AT"],
# }
schema = [
    ("PLACE", "HAS", "PERSON"),
    ("PERSON", "PART_OF", "PLACE"),
    ("PERSON", "WORKED_AT", "ORGANIZATION")
]

kg_extractor = SchemaLLMPathExtractor(
    llm=llm,
    possible_entities=entities,
    possible_relations=relations,
    kg_validation_schema=schema,
    strict=True,  # if false, will allow triples outside of the schema
    num_workers=4,
    max_paths_per_chunk=10,
    show_progres=False,
)

검색과 질의 (Retrieval and Querying)

라벨이 있는 property graph는 노드와 경로를 검색하는 다양한 질의 방법을 제공합니다. LlamaIndex에서는 여러 노드 검색 기법을 동시에 결합할 수 있어요.

# create a retriever
retriever = index.as_retriever(sub_retrievers=[retriever1, retriever2, ...])

# create a query engine
query_engine = index.as_query_engine(
    sub_retrievers=[retriever1, retriever2, ...]
)

하위 검색기를 지정하지 않으면 기본값은 LLMSynonymRetriever와 (임베딩이 활성화된 경우) VectorContextRetriever입니다. 현재 포함된 검색기는 다음과 같아요.

• LLMSynonymRetriever: LLM이 생성한 키워드와 동의어를 기반으로 노드를 검색합니다.
• VectorContextRetriever: 임베딩된 그래프 노드를 기반으로 노드를 검색합니다.
• TextToCypherRetriever: PropertyGraph의 스키마를 기반으로 LLM이 Cypher 쿼리를 생성하도록 지시합니다.
• CypherTemplateRetriever: LLM이 추론한 파라미터가 있는 Cypher 템플릿을 사용합니다.
• CustomPGRetriever: 커스텀 검색 로직을 구현하도록 쉽게 서브클래싱합니다.

from llama_index.core.indices.property_graph import (
    PGRetriever,
    VectorContextRetriever,
    LLMSynonymRetriever,
)

sub_retrievers = [
    VectorContextRetriever(index.property_graph_store, ...),
    LLMSynonymRetriever(index.property_graph_store, ...),
]

retriever = PGRetriever(sub_retrievers=sub_retrievers)
nodes = retriever.retrieve("<query>")

LLMSynonymRetriever

이 검색기는 입력 질의를 받아 관련 키워드와 동의어를 생성합니다. 이를 사용해 노드를 검색하고 결국 그 노드에 연결된 경로를 가져옵니다. 검색기를 구성에 명시적으로 선언하면 여러 옵션을 커스터마이즈할 수 있어요.

from llama_index.core.indices.property_graph import LLMSynonymRetriever

prompt = (
    "Given some initial query, generate synonyms or related keywords up to {max_keywords} in total, "
    "considering possible cases of capitalization, pluralization, common expressions, etc.\n"
    "Provide all synonyms/keywords separated by '^' symbols: 'keyword1^keyword2^...'\n"
    "Note, result should be in one-line, separated by '^' symbols."
    "----\n"
    "QUERY: {query_str}\n"
    "----\n"
    "KEYWORDS: "
)

def parse_fn(self, output: str) -> list[str]:
    matches = output.strip().split("^")
    # capitalize to normalize with ingestion
    return [x.strip().capitalize() for x in matches if x.strip()]

synonym_retriever = LLMSynonymRetriever(
    index.property_graph_store,
    llm=llm,
    # include source chunk text with retrieved paths
    include_text=False,
    synonym_prompt=prompt,
    output_parsing_fn=parse_fn,
    max_keywords=10,
    # the depth of relations to follow after node retrieval
    path_depth=1,
)

retriever = index.as_retriever(sub_retrievers=[synonym_retriever])

VectorContextRetriever

이 검색기는 벡터 유사도에 기반해 노드를 식별하고, 이후 그 노드에 연결된 경로를 가져옵니다. 그래프 스토어가 벡터 기능을 기본 지원한다면 해당 그래프 스토어만 관리하면 충분해요. 하지만 벡터 지원이 내장되어 있지 않다면 그래프 스토어를 벡터 스토어로 보완해야 합니다. 기본적으로 이 설정은 인메모리 SimpleVectorStore를 사용합니다.

from llama_index.core.indices.property_graph import VectorContextRetriever

vector_retriever = VectorContextRetriever(
    index.property_graph_store,
    # only needed when the graph store doesn't support vector queries
    # vector_store=index.vector_store,
    embed_model=embed_model,
    # include source chunk text with retrieved paths
    include_text=False,
    # the number of nodes to fetch
    similarity_top_k=2,
    # the depth of relations to follow after node retrieval
    path_depth=1,
    # can provide any other kwargs for the VectorStoreQuery class
    ...,
)

retriever = index.as_retriever(sub_retrievers=[vector_retriever])

TextToCypherRetriever

이 검색기는 그래프 스토어 스키마, 질의, 그리고 text-to-cypher 변환용 프롬프트 템플릿을 사용해 Cypher 쿼리를 생성·실행합니다.

참고: SimplePropertyGraphStore는 완전한 그래프 데이터베이스가 아니므로 Cypher 쿼리를 지원하지 않아요. 스키마를 확인하려면 index.property_graph_store.get_schema_str() 메서드를 사용할 수 있습니다.

from llama_index.core.indices.property_graph import TextToCypherRetriever

DEFAULT_RESPONSE_TEMPLATE = (
    "Generated Cypher query:\n{query}\n\n" "Cypher Response:\n{response}"
)
DEFAULT_ALLOWED_FIELDS = ["text", "label", "type"]
DEFAULT_TEXT_TO_CYPHER_TEMPLATE = (
    index.property_graph_store.text_to_cypher_template,
)

cypher_retriever = TextToCypherRetriever(
    index.property_graph_store,
    # customize the LLM, defaults to Settings.llm
    llm=llm,
    # customize the text-to-cypher template.
    # Requires `schema` and `question` template args
    text_to_cypher_template=DEFAULT_TEXT_TO_CYPHER_TEMPLATE,
    # customize how the cypher result is inserted into
    # a text node. Requires `query` and `response` template args
    response_template=DEFAULT_RESPONSE_TEMPLATE,
    # an optional callable that can clean/verify generated cypher
    cypher_validator=None,
    # allowed fields in the resulting
    allowed_output_field=DEFAULT_ALLOWED_FIELDS,
)

CypherTemplateRetriever

TextToCypherRetriever의 제약된 버전이에요. LLM이 어떤 Cypher 문장이든 자유롭게 생성하게 하는 대신, Cypher 템플릿을 제공하고 LLM이 빈칸을 채우게 합니다.

# NOTE: current v1 is needed
from pydantic.v1 import BaseModel, Field
from llama_index.core.indices.property_graph import CypherTemplateRetriever

# write a query with template params
cypher_query = """
MATCH (c:Chunk)-[:MENTIONS]->(o)
WHERE o.name IN $names
RETURN c.text, o.name, o.label;
"""

# create a pydantic class to represent the params for our query
# the class fields are directly used as params for running the cypher query
class TemplateParams(BaseModel):
    """Template params for a cypher query."""
    names: list[str] = Field(
        description="A list of entity names or keywords to use for lookup in a knowledge graph."
    )

template_retriever = CypherTemplateRetriever(
    index.property_graph_store, TemplateParams, cypher_query
)

더 알아보기 (Learn more)

• LlamaIndex PropertyGraph 문서 — 지식 그래프 인덱스 가이드
• SimpleLLMPathExtractor / ImplicitPathExtractor / SchemaLLMPathExtractor — 그래프 추출기
• LLMSynonymRetriever / VectorContextRetriever / TextToCypherRetriever / CypherTemplateRetriever — 그래프 검색기
• CustomPGRetriever — 커스텀 검색 로직 구현 베이스