Unable to Retrieve KnowledgeGraphIndex from Neo4j

[jamesjohnson1025]

I'm attempting to create a KnowledgeGraphIndex, store it in Neo4j, and then retrieve it for querying. However, I'm encountering issues when trying to load the index from Neo4j. The index seems to be created and stored successfully (I guess so), but retrieval fails. The error says ValueError(f"Failed to load index with ID {index_id}")

from llama_index.core import KnowledgeGraphIndex, SimpleDirectoryReader
from llama_index.core import StorageContext
from llama_index.graph_stores.neo4j import Neo4jGraphStore
from llama_index.core import load_index_from_storage
from llama_index.llms.openai import OpenAI
from llama_index.core import Settings
import logging
import sys
import os
import json

os.environ["OPENAI_API_KEY"] = "xxx"
logging.basicConfig(stream=sys.stdout, level=logging.INFO)

# define LLM
llm = OpenAI(temperature=0, model="gpt-4o-mini")
Settings.llm = llm
Settings.chunk_size = 512

username = "neo4j"
password = "xxx"
url = "bolt://localhost:7687"
database = "embeddingneo4j"

def create_and_save_index(index_id, graph_store=None):
    input_data = "/Users/xxx/Workspace/graph_rag/llama_index_graphrag/input"
    
    documents = SimpleDirectoryReader(input_data).load_data()
    
    storage_context = StorageContext.from_defaults(graph_store=graph_store)
    
    graph_index = KnowledgeGraphIndex.from_documents(
        documents[:1],
        storage_context=storage_context,
        max_triplets_per_chunk=2,
        include_embeddings=True
    )
    graph_index.set_index_id(index_id=index_id)
    
    # Save the index to neo4j
    graph_index.storage_context.persist()
    
    return graph_index

def load_and_query(query, neo4j_graph_store, index_id="capillary_index"):
    # Load the storage context
    storage_context = StorageContext.from_defaults(graph_store=neo4j_graph_store)
    
    # Load the index
    loaded_index = load_index_from_storage(storage_context, index_id=index_id)
    
    # Create a query engine
    query_engine = loaded_index.as_query_engine(
        include_text=True, response_mode="tree_summarize",
        embedding_mode="hybrid", similarity_top_k=5,
    )
    
    # Query
    response = query_engine.query(query)
    return response

graph_store = Neo4jGraphStore(
    username=username,
    password=password,
    url=url,
    database=database,
)

index = "capi_v2"
create_and_save_index(index_id=index, graph_store=graph_store)
print(load_and_query("Can we get all the custom fields from audience list filter?", graph_store, index))

Please let me know if I am wrong.

When using Neo4jGraphStore, shouldn't the index be automatically saved in the Neo4j database? If I still need to use a persist directory to save the index locally, I'm curious about the purpose of Neo4jGraphStore. Could you clarify the role of Neo4jGraphStore and explain whether local storage is necessary when using it?

[dosubot[bot]]

Hey there, @jamesjohnson1025! I'm here to help you with any bugs, questions, or contributions you have. Just give me a shout, and I'll do my best to assist you while we wait for a human maintainer.

I found some similar issues and discussions that might help you resolve the problem:

1. ValueError: doc_id not found while using KnowledgeGraphIndex - This issue was resolved by correcting the attribute name from storage_context.document_store to storage_context.docstore ^[1].
2. Big problem on saving and retrieve KnowledgeGraphIndex (Neo4j) - The solution involved using load_index_from_storage instead of load_graph ^[2].

To successfully retrieve a KnowledgeGraphIndex from Neo4j after storing it, you can use the following code:

from llama_index.core import StorageContext, KnowledgeGraphIndex
from llama_index.graph_stores.neo4j import Neo4jGraphStore

# Define your Neo4j connection parameters
username = "neo4j"
password = "xxx"
url = "bolt://localhost:7687"
database = "embeddingneo4j"

# Instantiate the Neo4jGraphStore
graph_store = Neo4jGraphStore(
    username=username,
    password=password,
    url=url,
    database=database,
)

# Create a storage context with the graph store
storage_context = StorageContext.from_defaults(graph_store=graph_store)

# Retrieve the KnowledgeGraphIndex from the storage context
index = KnowledgeGraphIndex.from_storage_context(storage_context=storage_context)

# Now you can use the index to query the knowledge graph
query_engine = index.as_query_engine(include_text=True, response_mode="tree_summarize", embedding_mode="hybrid", similarity_top_k=5)
response = query_engine.query("Can we get all the custom fields from audience list filter?")

from IPython.display import Markdown, display
display(Markdown(f"<b>{response}</b>"))

This code sets up the connection to your Neo4j instance, creates a StorageContext with the Neo4jGraphStore, and retrieves the KnowledgeGraphIndex from the storage context. You can then use this index to perform queries ^[3].

Make sure your Neo4j instance is correctly configured and that the node labels match those expected by the retriever. If you encounter any issues, you might want to increase logging verbosity to get detailed output during the retrieval process ^[4].

To continue talking to Dosu, mention @dosu.