Graph Data Store

What's a Graph Data Store?

Graph data stores are specialized databases designed to store and query data in the form of nodes (entities) and edges (relationships). They are essential for applications that need to model and analyze complex relationships between data points. Graph data stores are particularly useful for:

• Knowledge graphs and semantic networks

• Social network analysis and recommendation systems

• Fraud detection and pattern recognition

• Complex relationship modeling and traversal

• Hierarchical data representation

Available Implementations:

• Neo4j Graph Data Store

• Nebula Graph Data Store

• LlamaIndex Neo4j Graph RAG Data Store

• LightRAG Postgres Graph RAG Data Store

Installation

Linux, MacOS, or Windows WSL

# you can use a Conda environment
pip install --extra-index-url https://oauth2accesstoken:$(gcloud auth print-access-token)@glsdk.gdplabs.id/gen-ai-internal/simple/ gllm-datastore --extras kg

Windows Powershell

Windows Command Prompt

Managing Graph Data with Graph Data Store

The Graph Data Store provides comprehensive operations for managing graph data. Here's how to use the core operations effectively:

from gllm_datastore.graph_data_store import Neo4jGraphDataStore

# Initialize graph store
graph_store = Neo4jGraphDataStore(
    uri="bolt://localhost:7687",
    user="neo4j",
    password="password"
)

# Create nodes and relationships
await graph_store.upsert_node(
    label="Person",
    identifier_key="name",
    identifier_value="John Doe",
    properties={"age": 30, "occupation": "Engineer"}
)

await graph_store.upsert_node(
    label="Company",
    identifier_key="name",
    identifier_value="TechCorp",
    properties={"industry": "Technology", "founded": 2010}
)

# Create a relationship between nodes
await graph_store.upsert_relationship(
    node_source_key="name",
    node_source_value="John Doe",
    relation="WORKS_AT",
    node_target_key="name",
    node_target_value="TechCorp",
    properties={"since": "2018-05-20", "position": "Senior Engineer"}
)

You can then query the graph using the query method:

query = """
    MATCH (p:Person)-[r:WORKS_AT]->(c:Company)
    RETURN p.name, r.position, c.name
"""
results = await graph_store.query(query)

Node Operations

Creating and Updating Nodes

The upsert_node method creates a new node if it doesn't exist, or updates an existing node if it does. This is useful for maintaining data consistency.

# Create or update a node
await graph_store.upsert_node(
    label="Person",              # Node label/type
    identifier_key="email",     # Property used as unique identifier
    identifier_value="[e-mail]",  # Value of the identifier
    properties={                # Additional properties
        "name": "John Smith",
        "age": 35,
        "skills": ["Python", "Data Science"]
    }
)

Deleting Nodes

The delete_node method removes a node from the graph. Note that in most graph databases, this will also delete all relationships connected to this node.

# Delete a node
await graph_store.delete_node(
    label="Person",
    identifier_key="email",
    identifier_value="[e-mail]"
)

Relationship Operations

Creating and Updating Relationships

The upsert_relationship method creates or updates a relationship between two existing nodes.

# Create or update a relationship
await graph_store.upsert_relationship(
    node_source_key="email",
    node_source_value="[e-mail]",
    relation="MANAGES",
    node_target_key="email",
    node_target_value="[e-mail]",
    properties={"since": "2022-01-15", "department": "Engineering"}
)

Deleting Relationships

The delete_relationship method removes a specific relationship between two nodes.

# Delete a relationship
await graph_store.delete_relationship(
    node_source_key="email",
    node_source_value="[e-mail]",
    relation="MANAGES",
    node_target_key="email",
    node_target_value="[e-mail]"
)

Querying the Graph

The query method allows you to execute native graph query language statements (e.g., Cypher for Neo4j) to retrieve data from the graph.

# Execute a Cypher query (Neo4j)
results = await graph_store.query(
    query="MATCH (p:Person)-[r:WORKS_AT]->(c:Company) WHERE c.industry = $industry RETURN p.name, r.position, c.name",
    parameters={"industry": "Technology"}
)

# Process results
for record in results:
    print(f"{record['p.name']} works as {record['r.position']} at {record['c.name']}")

Graph Traversal

The traverse_graph method provides a powerful way to explore the graph structure starting from a specific node and following relationships to a specified depth.

# Traverse the graph starting from a specific node
node_properties = {"name": "John Doe"}
node_props_to_extract = ["name", "age", "occupation"]
rel_props_to_extract = ["since", "position"]

nodes, relationships = await graph_store.traverse_graph(
    node_properties=node_properties,
    extracted_node_properties=node_props_to_extract,
    extracted_relationship_properties=rel_props_to_extract,
    depth=2  # How many hops to traverse
)

# Process the traversal results
print(f"Found {len(nodes)} nodes and {len(relationships)} relationships")

Graph RAG Data Store

The Graph RAG (Retrieval-Augmented Generation) Data Store extends the basic graph data store with capabilities specifically designed for AI applications. It allows you to query the graph using natural language and integrate with LLM workflows. It also helps to link graph components with documents, allowing for a synchronized data between Vector DB and Graph DB.

LightRAG Postgres Graph RAG Data Store

The LightRAG Data Store provides a PostgreSQL-backed implementation for graph-based RAG with automatic knowledge graph construction and natural language querying capabilities.

from gllm_inference.em_invoker import OpenAIEMInvoker
from gllm_inference.lm_invoker import OpenAILMInvoker
from gllm_core.schema import Chunk
from gllm_datastore.graph_data_store.light_rag_postgres_data_store import LightRAGPostgresDataStore

# Initialize the invokers
em_invoker = OpenAIEMInvoker(model_name="text-embedding-3-small")
lm_invoker = OpenAILMInvoker(model_name="gpt-4o-mini")

async def main():
    # Initialize the LightRAG data store
    data_store = LightRAGPostgresDataStore(
        lm_invoker=lm_invoker,
        em_invoker=em_invoker,
        postgres_db_host="localhost",
        postgres_db_port=5432,
        postgres_db_name="rag",
    )
    
    # Ensure the data store is initialized
    await data_store.ensure_initialized()
    
    # Insert chunks - LightRAG automatically builds the knowledge graph
    await data_store.insert(
        chunks=[
            Chunk(content="Joko Widodo is the 7th President of Indonesia.", id="joko_widodo"),
            Chunk(content="Prabowo Subianto is the 8th President of Indonesia.", id="prabowo_subianto"),
        ],
    )
    
    # Query using natural language
    print("Performing query 'Who is Prabowo Subianto?':")
    print(await data_store.query("Who is Prabowo Subianto?"))
    
    # Delete specific chunks
    await data_store.delete("joko_widodo")
    await data_store.delete("prabowo_subianto")

# Run the async function
asyncio.run(main())

LlamaIndex Neo4j Graph RAG Data Store

You can instantiate a LlamaIndex Neo4j Graph RAG Data Store which can then be used for retrieval (see LlamaIndex Graph RAG Retriever).

from gllm_datastore.graph_data_store import LlamaIndexNeo4jGraphRAGDataStore

# Initialize the Graph RAG store
graph_rag_store = LlamaIndexNeo4jGraphRAGDataStore(
    url="bolt://localhost:7687",
    username="neo4j",
    password="password"
)