ToolsRetriever

정의: 각 단계를 함수(도구)로 정의하고, LLM이 자동으로 이 도구들을 선택·실행하도록 하는 패턴. GraphRAG 파이프라인의 자동화 엔진.

핵심 개념

기본 구조

도구 정의
  ├─ extract_entities(text, types)
  ├─ extract_relationships(entities, types)
  ├─ create_cypher_query(entities, relations)
  ├─ execute_cypher(query)
  └─ generate_answer(query, results)
       ↓
LLM 호출
  → "이 도구들 중 어떤 것을 실행할까?"
       ↓
자동 실행
  → 도구 결과를 받아 다음 단계로
       ↓
순환
  → 결과에 따라 다음 도구 결정

장점

수동 개입 최소화 — 파이프라인이 자동으로 진행 ✅ 적응형 처리 — 쿼리/데이터에 따라 다른 도구 사용 ✅ 재사용성 — 한 번 정의된 도구를 여러 파이프라인에서 활용 ✅ 일관성 — 프롬프트보다 코드 기반이라 결정론적

실제 예시: GraphRAG 파이프라인

도구 정의

def extract_entities(text: str, entity_types: List[str]) -> List[Dict]:
    """텍스트에서 엔티티 추출"""
    prompt = f"""
    텍스트: {text}
    찾을 타입: {entity_types}
    
    JSON 형식으로 반환해줘.
    """
    result = llm.invoke(prompt)
    return json.loads(result)
 
def extract_relationships(entities: List[Dict], rel_types: List[str]) -> List[Dict]:
    """엔티티 간 관계 추출"""
    prompt = f"""
    엔티티: {entities}
    관계 타입: {rel_types}
    
    [{{"source": "...", "target": "...", "relation": "..."}}]
    """
    result = llm.invoke(prompt)
    return json.loads(result)
 
def create_cypher_query(entities: List[Dict], relationships: List[Dict]) -> str:
    """Cypher 쿼리 자동 생성"""
    cypher = "CREATE "
    for e in entities:
        cypher += f"(:{e['type']} {{{json.dumps(e['attrs'])}}}), "
    for r in relationships:
        cypher += f"(source)-[:{r['type']}]->(target), "
    return cypher
 
def execute_cypher(query: str) -> List[Dict]:
    """Neo4j에서 쿼리 실행"""
    return neo4j.run(query)
 
def generate_answer(user_query: str, cypher_results: List[Dict]) -> str:
    """최종 답변 생성"""
    prompt = f"""
    사용자 질문: {user_query}
    조회 결과: {cypher_results}
    
    자연어로 답변해줘.
    """
    return llm.invoke(prompt)

LLM의 자동 호출

tools = [
    extract_entities,
    extract_relationships,
    create_cypher_query,
    execute_cypher,
    generate_answer
]
 
user_input = "Alice의 매니저는 누구인가?"
 
# LLM이 도구를 순차적으로 선택·실행
step1 = extract_entities(user_input, ["Person", "Company"])
# → [{"entity": "Alice", "type": "Person"}, ...]
 
step2 = extract_relationships(step1, ["MANAGES", "WORKS_AT"])
# → [{"source": "Bob", "target": "Alice", "relation": "MANAGES"}]
 
step3 = create_cypher_query(step1, step2)
# → "CREATE (alice:Person {...})-[:MANAGES]->(bob:Person {...})"
 
step4 = execute_cypher(step3)
# → [{"name": "Bob"}]
 
step5 = generate_answer(user_input, step4)
# → "Alice의 매니저는 Bob입니다."

ToolsRetriever vs 단순 Prompting

단순 Prompting의 문제

prompt = """
텍스트: {text}
 
1. 엔티티를 찾아줘
2. 관계를 찾아줘
3. Cypher 쿼리를 만들어줘
4. 쿼리를 실행해줘
5. 답변해줘
"""
 
result = llm.invoke(prompt)  # 한 번에 모든 것 시도 → 환각 위험

문제: ❌ 각 단계 검증 불가 ❌ 중간에 잘못되면 후속 단계 오류 전파 ❌ 코드 실행이 아니라 텍스트 생성 (Cypher 문법 오류)

ToolsRetriever의 장점

for tool_name, inputs in llm_decisions:
    tool_func = tools[tool_name]
    result = tool_func(*inputs)  # 실제 함수 실행
    assert result is not None  # 검증 가능

장점: ✅ 각 단계의 결과를 검증할 수 있음 ✅ 중간 오류 감지 및 대체 경로 실행 가능 ✅ 코드 기반이라 문법 정확성 보장 ✅ 로깅·모니터링 가능

실무 구현 패턴

LangChain의 Agent

from langchain.agents import initialize_agent, Tool
from langchain.llms import OpenAI
 
tools = [
    Tool(name="Extract Entities", func=extract_entities),
    Tool(name="Extract Relationships", func=extract_relationships),
    Tool(name="Create Cypher", func=create_cypher_query),
    Tool(name="Execute Cypher", func=execute_cypher),
    Tool(name="Generate Answer", func=generate_answer)
]
 
agent = initialize_agent(
    tools,
    llm=OpenAI(),
    agent="zero-shot-react-description",
    verbose=True
)
 
result = agent.run("Alice의 매니저는 누구인가?")

LangGraph의 StateGraph

from langgraph.graph import StateGraph
 
def create_pipeline():
    graph = StateGraph(GraphRAGState)
    
    graph.add_node("extract_entities", extract_entities)
    graph.add_node("extract_relationships", extract_relationships)
    graph.add_node("create_cypher", create_cypher_query)
    graph.add_node("execute", execute_cypher)
    graph.add_node("generate", generate_answer)
    
    graph.add_edge("extract_entities", "extract_relationships")
    graph.add_edge("extract_relationships", "create_cypher")
    graph.add_edge("create_cypher", "execute")
    graph.add_edge("execute", "generate")
    
    return graph.compile()
 
pipeline = create_pipeline()
result = pipeline.invoke({"query": "Alice의 매니저는?"})

에러 처리 및 재시도

실패 복구 패턴

def execute_with_retry(tool, inputs, max_retries=3):
    for attempt in range(max_retries):
        try:
            return tool(*inputs)
        except ValidationError as e:
            if attempt < max_retries - 1:
                # 수정된 입력으로 재시도
                inputs = refine_inputs(inputs, error=e)
            else:
                raise

대체 경로 (Fallback)

if not entities:
    # 엔티티 추출 실패 시
    entities = fallback_entity_extraction(text)
 
if not relationships:
    # 관계 추출 실패 시
    relationships = fallback_relationship_extraction(entities, text)

성능 최적화

최적화 전략설명
캐싱동일 쿼리는 도구 재실행 안 함
병렬 실행독립적 도구는 동시 실행
배치 처리여러 엔티티를 한 번에 추출
조기 종료충분한 신뢰도 도달 시 멈춤
샘플링결과의 일부만 처리

관련 개념


출처: data-to-kg — GraphRAG 파이프라인에서의 ToolsRetriever 활용