Neo4j 데이터베이스 위에 RAG 구축하기

Neo4j 데이터베이스 위에 RAG 구축하기

Neo4j 그래프 데이터베이스(Movie DBMD 영화 데이터)를 Mistral LLM으로 질의하는 text-2-cypher 에이전트를 만드는 노트북이에요. 자연어 질문을 Cypher 쿼리로 바꾸고, Neo4j에서 실행한 뒤, 그 결과를 다시 LLM이 자연어로 설명하도록 연결합니다.

출처: 문서

본문

이 노트북을 실행하려면 Neo4j Desktop을 다운로드하거나 클라우드를 통해 Neo4j DB에 접근할 수 있어야 해요. Neo4j DB가 localhost에서 실행 중이라면 Google Colab보다는 로컬에서 이 쿡북을 실행하는 걸 권장합니다.

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

! pip install mistralai neo4j
import json
from neo4j import GraphDatabase
from mistralai.client import Mistral
from getpass import getpass

API 키와 Neo4j 접속 정보를 입력합니다.

api_key= getpass("Type your API Key")
neo4j_password =  getpass("Type your neo4j password")
neo4j_user =  getpass("Type your neo4j user name")
neo4j_uri =  getpass("Type your neo4j url")

client = Mistral(api_key=api_key)

1단계: Neo4j Desktop 다운로드와 Python 클라이언트 접근 확인

Neo4j Desktop을 설치하고 기본 프로젝트인 Example Project(영화 DBMD 데이터 포함)를 엽니다. Cypher 쿼리를 실행하는 함수를 정의합니다.

URI = neo4j_uri
AUTH = (neo4j_user, neo4j_password)

def run_cypher_query(cypher_query):
    with GraphDatabase.driver(URI, auth=AUTH) as driver:
        records, _, _ = driver.execute_query(cypher_query, database_="neo4j")
        return records

2단계: Movie DBMD 데이터베이스용 text-2-cypher 에이전트 만들기

자연어 질문을 받아 Cypher 쿼리를 생성하는 함수입니다. Neo4j 스키마(라벨, 관계, 속성)를 프롬프트에 넣고, JSON 형식으로 결과를 돌려주도록 요청해요.

def generate_cypher_query(question):
    prompt = f"""You are a coding agent interacting with a Neo4j database with the following schema :
    - Labels : "Movie", "Person"
    - Relationships : "ACTED_IN", "DIRECTED", "FOLLOWS", "PRODUCED", "REVIEWED", "WROTE"
    "Person" label has the following properties :
    - born
    - name
    "Movie" label has the following properties :
    - title
    - released
    Your will be given as input a query in natural language and your role is to output a cypher query whose output will contain the answer.
    Your output with be in a json format.
    Examples :
    input : When was the movie "The Matrix" released ?
    output : {{"result": "MATCH (n:Movie) WHERE n.title='The Matrix' RETURN n.released"}}
    input : In which movied Tom Hanks played ?
    output : {{"result": "MATCH (p:Person {{name: 'Tom Hanks'}})-[:ACTED_IN]->(m:Movie) RETURN m.title AS movieTitle"}}
    input : What movie Steven Spielber produced ?
    output : {{"result": "MATCH (p:Person {{name: 'Steven Spielberg'}})-[:PRODUCED]->(m:Movie) RETURN m.title AS movieTitle"}}
    Here is the user question :
    {question}
    """
    chat_response = client.chat.complete(
        model= model,
        response_format = {"type": "json_object"},
        messages = [
            {
                "role": "user",
                "content": prompt,
            },
        ]
    )
    return chat_response.choices[0].message.content

Cypher 쿼리 실행 결과를 바탕으로 최종 답변을 생성하는 함수입니다.

def respond_to_query(question, cypher_code, query_output):
    prompt = f"""You are a coding agent interacting with a Neo4j database with the following schema :
    - Labels : "Movie", "Person"
    - Relationships : "ACTED_IN", "DIRECTED", "FOLLOWS", "PRODUCED", "REVIEWED", "WROTE"
    "Person" label has the following properties :
    - born
    - name
    "Movie" label has the following properties :
    - title
    - released
    The user asked the following question :
    {question}
    To answee the question the following cypher query was run on Neo4j :
    {cypher_code}
    The following output was obtained :
    {query_output}
    Based on all these elements answer the initial user question.
    Be straight to the point and concise in your answers.
    Your answer:
    """
    chat_response = client.chat.complete(
        model= model,
        messages = [
            {
                "role": "user",
                "content": prompt,
            },
        ]
    )
    return chat_response.choices[0].message.content

세 함수를 묶은 neo4j_agent입니다. 질문 → Cypher 생성 → 실행 → 응답 생성의 전체 흐름을 한 번에 처리하고, 질문·쿼리·응답을 출력합니다.

def neo4j_agent(question):
    cypher_code = json.loads(generate_cypher_query(question))['result']
    query_result = run_cypher_query(cypher_code)
    response = respond_to_query(question, cypher_code, query_result)

    print(f'Question : \n {question} \n')
    print(f'Query : \n {cypher_code} \n')
    print(f'Response : \n {response} \n')

이제 영화 데이터베이스에 다양한 질문을 던져 볼 수 있어요.

neo4j_agent("When was Keanue Reeves born ?")
neo4j_agent("What actors played in the movie The Matrix ?")
neo4j_agent("Tell me the name of a person that is both an actor and a producer on another movie")
neo4j_agent("List Tom Hanks movies and sort them by release date")
neo4j_agent("We are in 2024, how old is Tom Hanks ?")
neo4j_agent("Give me names of two actors that played together in two differnet films. Give me the names of the associated movies")

더 알아보기 (Learn more)

  • Neo4j 공식 문서 — 그래프 데이터베이스 공식 가이드
  • Neo4j Desktop — 로컬 Neo4j 실행 도구
  • Cypher 쿼리 언어 — MATCH, RETURN 등 그래프 질의 문법
  • mistralai.client.Mistral — Mistral Python SDK의 클라이언트 (여기서는 response_format={"type": "json_object"}로 JSON 출력 유도)