HugeGraph QA Chain

This notebook shows how to use LLMs to provide a natural language interface to HugeGraph database.

You will need to have a running HugeGraph instance. You can run a local docker container by running the executing the following script:

docker run \
    --name=graph \
    -itd \
    -p 8080:8080 \
    hugegraph/hugegraph

If we want to connect HugeGraph in the application, we need to install python sdk:

pip3 install hugegraph-python

If you are using the docker container, you need to wait a couple of second for the database to start, and then we need create schema and write graph data for the database.

from hugegraph.connection import PyHugeGraph

client = PyHugeGraph("localhost", "8080", user="admin", pwd="admin", graph="hugegraph")

First, we create the schema for a simple movie database:

"""schema"""
schema = client.schema()
schema.propertyKey("name").asText().ifNotExist().create()
schema.propertyKey("birthDate").asText().ifNotExist().create()
schema.vertexLabel("Person").properties(
    "name", "birthDate"
).usePrimaryKeyId().primaryKeys("name").ifNotExist().create()
schema.vertexLabel("Movie").properties("name").usePrimaryKeyId().primaryKeys(
    "name"
).ifNotExist().create()
schema.edgeLabel("ActedIn").sourceLabel("Person").targetLabel(
    "Movie"
).ifNotExist().create()

'create EdgeLabel success, Detail: "b\'{"id":1,"name":"ActedIn","source_label":"Person","target_label":"Movie","frequency":"SINGLE","sort_keys":[],"nullable_keys":[],"index_labels":[],"properties":[],"status":"CREATED","ttl":0,"enable_label_index":true,"user_data":{"~create_time":"2023-07-04 10:48:47.908"}}\'"'

Then we can insert some data.

"""graph"""
g = client.graph()
g.addVertex("Person", {"name": "Al Pacino", "birthDate": "1940-04-25"})
g.addVertex("Person", {"name": "Robert De Niro", "birthDate": "1943-08-17"})
g.addVertex("Movie", {"name": "The Godfather"})
g.addVertex("Movie", {"name": "The Godfather Part II"})
g.addVertex("Movie", {"name": "The Godfather Coda The Death of Michael Corleone"})

g.addEdge("ActedIn", "1:Al Pacino", "2:The Godfather", {})
g.addEdge("ActedIn", "1:Al Pacino", "2:The Godfather Part II", {})
g.addEdge(
    "ActedIn", "1:Al Pacino", "2:The Godfather Coda The Death of Michael Corleone", {}
)
g.addEdge("ActedIn", "1:Robert De Niro", "2:The Godfather Part II", {})

1:Robert De Niro--ActedIn-->2:The Godfather Part II

Creating HugeGraphQAChain

We can now create the HugeGraph and HugeGraphQAChain. To create the HugeGraph we simply need to pass the database object to the HugeGraph constructor.

from langchain.chains import HugeGraphQAChain
from langchain.chat_models import ChatOpenAI
from langchain.graphs import HugeGraph

graph = HugeGraph(
    username="admin",
    password="admin",
    address="localhost",
    port=8080,
    graph="hugegraph",
)

Refresh graph schema information

If the schema of database changes, you can refresh the schema information needed to generate Gremlin statements.

# graph.refresh_schema()

print(graph.get_schema)

Node properties: [name: Person, primary_keys: ['name'], properties: ['name', 'birthDate'], name: Movie, primary_keys: ['name'], properties: ['name']]
Edge properties: [name: ActedIn, properties: []]
Relationships: ['Person--ActedIn-->Movie']

Querying the graph

We can now use the graph Gremlin QA chain to ask question of the graph

chain = HugeGraphQAChain.from_llm(ChatOpenAI(temperature=0), graph=graph, verbose=True)

chain.run("Who played in The Godfather?")

> Entering new  chain...
Generated gremlin:
g.V().has('Movie', 'name', 'The Godfather').in('ActedIn').valueMap(true)
Full Context:
[{'id': '1:Al Pacino', 'label': 'Person', 'name': ['Al Pacino'], 'birthDate': ['1940-04-25']}]

> Finished chain.

'Al Pacino played in The Godfather.'