RAGStringQueryEngine vs.graph_rag_retriever in query pandas dataframes

[mraguth]

I am using hybrid of RAGStringQueryEngine to query my pandas dataframe like below example:

query_engine = RAGStringQueryEngine(
    retriever=recursive_retriever,
    response_synthesizer=response_synthesizer,
    pandas_prompt=pandas_prompt,
    pandas_output_parser=pandas_output_parser,
    llm=llm,
    qa_prompt=qa_prompt,
    synthesize_response=True,
    response_synthesis_prompt=response_synthesis_prompt,
    callback_manager=Settings.callback_manager,                                  ###### Important for Langfuse.
)

However, someone suggested me to check it out.
Question here is... is it my RAGStringQueryEngine is not better than KnowledgeGraphRAGRetriever? Any suggestions please???

https://docs.llamaindex.ai/en/stable/examples/query_engine/knowledge_graph_rag_query_engine/?h=graph

from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.retrievers import KnowledgeGraphRAGRetriever

graph_rag_retriever = KnowledgeGraphRAGRetriever(
    storage_context=storage_context,
    verbose=True,
)

query_engine = RetrieverQueryEngine.from_args(
    graph_rag_retriever,
)

[dosubot[bot]]

Using the KnowledgeGraphRAGRetriever might be more effective than your current setup with RAGStringQueryEngine for querying pandas dataframes. The KnowledgeGraphRAGRetriever leverages a Knowledge Graph to retrieve information, which can mitigate several issues caused by the "split and embedding" approach of traditional RAG methods. It offers advanced features such as entity extraction, synonym expansion, and different retriever modes (keyword, embedding, hybrid) that can significantly enhance the retrieval process, especially for complex queries and large datasets.

Here is an example of how to set up and use the KnowledgeGraphRAGRetriever:

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.indices.knowledge_graph.retrievers import KnowledgeGraphRAGRetriever
from llama_index.llms.openai import OpenAI
from llama_index.core.prompts import BasePromptTemplate, PromptTemplate

# Load documents
documents = SimpleDirectoryReader("./data/paul_graham/").load_data()
index = VectorStoreIndex.from_documents(documents)

# Define the KnowledgeGraphRAGRetriever
retriever = KnowledgeGraphRAGRetriever(
    storage_context=index.storage_context,
    llm=OpenAI(model="gpt-3.5-turbo"),
    retriever_mode="hybrid",  # You can choose "keyword", "embedding", or "hybrid"
    graph_traversal_depth=2,
    max_knowledge_sequence=30,
    verbose=True
)

# Define the QA prompt
qa_prompt = PromptTemplate(
    "Context information is below.\n"
    "---------------------\n"
    "{context_str}\n"
    "---------------------\n"
    "Given the context information and not prior knowledge, "
    "answer the query.\n"
    "Query: {query_str}\n"
    "Answer: "
)

# Query the retriever
query_str = "What did the author do growing up?"
nodes = retriever.retrieve(query_str)
context_str = "\n\n".join([n.node.get_content() for n in nodes])
response = retriever._llm.complete(
    qa_prompt.format(context_str=context_str, query_str=query_str)
)
print(str(response))

This approach allows you to leverage the advanced capabilities of the KnowledgeGraphRAGRetriever, potentially making it more effective for your specific needs, especially when dealing with complex data structures like pandas dataframes ^[1][2][3].

To continue talking to Dosu, mention @dosu.

[mraguth]

@dosu. Thank you. I did see good and accurate results vai below. Is that possible to include those along with KnowledgeGraphRAGRetriever or do you think it is duplicate efforts? Sorry, I am new to these concepts.

    retriever=recursive_retriever,
    response_synthesizer=response_synthesizer,
    pandas_prompt=pandas_prompt,
    pandas_output_parser=pandas_output_parser,
    synthesize_response=True,

[dosubot[bot]]

Yes, it is possible to include the components from your initial setup (recursive_retriever, response_synthesizer, pandas_prompt, pandas_output_parser, synthesize_response) along with the KnowledgeGraphRAGRetriever. This integration can enhance the retrieval process without redundancy.

Here is an example of how you can configure the KnowledgeGraphRAGRetriever with your existing components:

from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.retrievers import KnowledgeGraphRAGRetriever

# Assuming you have already defined your service_context, storage_context, and other components
kg_retriever = KnowledgeGraphRAGRetriever(
    storage_context=storage_context,
    entity_extract_fn=entity_extract_fn,
    entity_extract_template=pandas_prompt,  # Your pandas_prompt here
    synonym_expand_fn=synonym_expand_fn,
    synonym_expand_template=pandas_output_parser,  # Your pandas_output_parser here
    response_synthesizer=synthesize_response,  # Your synthesize_response here
    retriever_mode="keyword_embedding",  # Example mode
    verbose=True
)

query_engine = RetrieverQueryEngine.from_args(
    kg_retriever,
)

This setup allows you to leverage the capabilities of both your initial components and the KnowledgeGraphRAGRetriever effectively ^[1][2][3].

To continue talking to Dosu, mention @dosu.

[mraguth]

@dosu I am using recursive_retriever as below way. In my recursive_retriever, I have doc_nodes and df (Pandas) nodes as shown below in the code. I did not see you added in your kg_retriever. How I add that?
And also I didn't find these two below from your kg_retriever in your code. Can you give me comlete code along with my recursive_retriever please?

    entity_extract_fn=entity_extract_fn,
    synonym_expand_fn=synonym_expand_fn,

# construct top-level vector index + query engine
vector_index = VectorStoreIndex(doc_nodes + df_nodes)
vector_retriever = vector_index.as_retriever(similarity_top_k=1)


recursive_retriever = RecursiveRetriever(
    "vector",
    retriever_dict={"vector": vector_retriever},
    query_engine_dict=df_id_query_engine_mapping,
    verbose=False,  # Change to see the data
)

[dosubot[bot]]

To integrate your recursive_retriever with the KnowledgeGraphRAGRetriever and include the missing parameters entity_extract_fn and synonym_expand_fn, you can follow the example below:

from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.retrievers import RecursiveRetriever
from llama_index.core.indices.knowledge_graph.retrievers import KnowledgeGraphRAGRetriever
from llama_index.core.storage.storage_context import StorageContext
from llama_index.core.llms.llm import LLM
from llama_index.legacy.schema import QueryBundle, NodeWithScore
from typing import List

# Define the entity extraction function
def entity_extract_fn(text: str) -> List[str]:
    # Example implementation of entity extraction
    return text.split()

# Define the synonym expansion function
def synonym_expand_fn(entity: str) -> List[str]:
    # Example implementation of synonym expansion
    synonyms = {
        "car": ["automobile", "vehicle"],
        "bike": ["bicycle", "cycle"]
    }
    return synonyms.get(entity, [entity])

# Create instances of required contexts
storage_context = StorageContext(graph_store=...)  # Provide your graph store here
llm = LLM(...)  # Provide your LLM instance here

# Initialize the KnowledgeGraphRAGRetriever
kg_retriever = KnowledgeGraphRAGRetriever(
    storage_context=storage_context,
    llm=llm,
    entity_extract_fn=entity_extract_fn,
    synonym_expand_fn=synonym_expand_fn,
    retriever_mode="keyword_embedding",  # Example mode
    verbose=True
)

# Construct top-level vector index + query engine
vector_index = VectorStoreIndex(doc_nodes + df_nodes)
vector_retriever = vector_index.as_retriever(similarity_top_k=1)

# Initialize the RecursiveRetriever
recursive_retriever = RecursiveRetriever(
    root_id="vector",
    retriever_dict={"vector": vector_retriever, "kg": kg_retriever},
    query_engine_dict=df_id_query_engine_mapping,
    verbose=False  # Change to True to see the data
)

# Example usage of the recursive retriever
query_bundle = QueryBundle(query_str="What is the capital of France?")
nodes_with_scores = recursive_retriever._retrieve(query_bundle)
for node_with_score in nodes_with_scores:
    print(node_with_score)

This code integrates your recursive_retriever with the KnowledgeGraphRAGRetriever and includes the entity_extract_fn and synonym_expand_fn functions. The recursive_retriever is configured to use both the vector retriever and the knowledge graph retriever, allowing you to leverage the capabilities of both retrievers effectively ^[1][2].

To continue talking to Dosu, mention @dosu.

[mraguth]

@dosu Thank you. However, I am not clear about tow things here....

1. How to create graph_store

# Create instances of required contexts
storage_context = StorageContext(graph_store=...)  # Provide your graph store here

2. At the end of query_bundle, how I get the response?

# Example usage of the recursive retriever
query_bundle = QueryBundle(query_str="What is the capital of France?")
nodes_with_scores = recursive_retriever._retrieve(query_bundle)
for node_with_score in nodes_with_scores:
    print(node_with_score)

[dosubot[bot]]

To address your questions:

1. Creating the graph_store:
   You can create a graph_store using the Neo4jGraphStore class. Here is an example of how to initialize it:

   from llama_index.core.storage.storage_context import StorageContext
   from llama_index.core.graph_stores.neo4j import Neo4jGraphStore
   
   # Initialize the Neo4jGraphStore with required parameters
   graph_store = Neo4jGraphStore(
       username="your_username",
       password="your_password",
       url="bolt://localhost:7687",
       database="neo4j"  # Optional, default is "neo4j"
   )
   
   # Create the storage context with the graph store
   storage_context = StorageContext(graph_store=graph_store)

2. Retrieving the final response from the query_bundle:
   After retrieving the nodes with scores using the recursive_retriever, you can process these nodes to get the final response. Here is an example of how to do this:

   from llama_index.core.schema import QueryBundle
   from llama_index.core.query_engine.knowledge_graph_query_engine import KnowledgeGraphQueryEngine
   
   # Example usage of the recursive retriever
   query_bundle = QueryBundle(query_str="What is the capital of France?")
   nodes_with_scores = recursive_retriever._retrieve(query_bundle)
   
   # Initialize the KnowledgeGraphQueryEngine with the storage context
   query_engine = KnowledgeGraphQueryEngine(storage_context=storage_context)
   
   # Retrieve the final response
   final_response = query_engine._query(query_bundle)
   
   # Print the final response
   print(final_response)

Here is the complete code integrating both parts:

from llama_index.core.storage.storage_context import StorageContext
from llama_index.core.graph_stores.neo4j import Neo4jGraphStore
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.retrievers import RecursiveRetriever
from llama_index.core.indices.knowledge_graph.retrievers import KnowledgeGraphRAGRetriever
from llama_index.core.llms.llm import LLM
from llama_index.core.schema import QueryBundle
from llama_index.core.query_engine.knowledge_graph_query_engine import KnowledgeGraphQueryEngine
from llama_index.legacy.schema import NodeWithScore
from typing import List

# Define the entity extraction function
def entity_extract_fn(text: str) -> List[str]:
    return text.split()

# Define the synonym expansion function
def synonym_expand_fn(entity: str) -> List[str]:
    synonyms = {
        "car": ["automobile", "vehicle"],
        "bike": ["bicycle", "cycle"]
    }
    return synonyms.get(entity, [entity])

# Initialize the Neo4jGraphStore
graph_store = Neo4jGraphStore(
    username="your_username",
    password="your_password",
    url="bolt://localhost:7687",
    database="neo4j"
)

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

# Initialize the LLM
llm = LLM(...)  # Provide your LLM instance here

# Initialize the KnowledgeGraphRAGRetriever
kg_retriever = KnowledgeGraphRAGRetriever(
    storage_context=storage_context,
    llm=llm,
    entity_extract_fn=entity_extract_fn,
    synonym_expand_fn=synonym_expand_fn,
    retriever_mode="keyword_embedding",
    verbose=True
)

# Construct top-level vector index + query engine
vector_index = VectorStoreIndex(doc_nodes + df_nodes)
vector_retriever = vector_index.as_retriever(similarity_top_k=1)

# Initialize the RecursiveRetriever
recursive_retriever = RecursiveRetriever(
    root_id="vector",
    retriever_dict={"vector": vector_retriever, "kg": kg_retriever},
    query_engine_dict=df_id_query_engine_mapping,
    verbose=False
)

# Example usage of the recursive retriever
query_bundle = QueryBundle(query_str="What is the capital of France?")
nodes_with_scores = recursive_retriever._retrieve(query_bundle)

# Initialize the KnowledgeGraphQueryEngine with the storage context
query_engine = KnowledgeGraphQueryEngine(storage_context=storage_context)

# Retrieve the final response
final_response = query_engine._query(query_bundle)

# Print the final response
print(final_response)

This code sets up the graph_store, integrates it into the storage_context, and retrieves the final response from the query_bundle using the KnowledgeGraphQueryEngine ^[1][2].

To continue talking to Dosu, mention @dosu.

[mraguth]

@dosu Do I need to install Neo4jGraphStore locally? Please advise.

[dosubot[bot]]

Yes, you need to have Neo4j installed and running locally or on a remote server to use Neo4jGraphStore. Here are the steps to set it up:

1. Install Neo4j:

   • You can download and install Neo4j from the official website.
   • Follow the installation instructions for your operating system.

2. Start Neo4j:

   • Once installed, start the Neo4j server. You can usually do this via the Neo4j Desktop application or by running the appropriate command in your terminal.

3. Install the Neo4j Python Driver:

   • You need the Neo4j Python driver to connect to the Neo4j database from your Python code. Install it using pip:

     pip install neo4j

4. Configure Neo4jGraphStore:

   • Ensure you have the correct connection details (username, password, URL) for your Neo4j instance.

Here is an example of how to set up and use Neo4jGraphStore:

from llama_index.core.storage.storage_context import StorageContext
from llama_index.core.graph_stores.neo4j import Neo4jGraphStore

# Initialize the Neo4jGraphStore with required parameters
graph_store = Neo4jGraphStore(
    username="your_username",
    password="your_password",
    url="bolt://localhost:7687",  # Adjust the URL if your Neo4j instance is remote
    database="neo4j"  # Optional, default is "neo4j"
)

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

# Now you can use storage_context in your retriever setup

Make sure Neo4j is running and accessible at the specified URL. This setup will allow you to use Neo4jGraphStore as part of your KnowledgeGraphRAGRetriever configuration.

To continue talking to Dosu, mention @dosu.

[mraguth]

@dosu... Now I am using Nebula Graph store. The core of new Nebula is like below... it needs documents.
However, my data is in pandas dataframe and I am using this way to get vector index. Question is.. how I pass this to GraphDatabase?

df_nodes = [
    IndexNode(text=summary, index_id=f"pandas{idx}")
    for idx, summary in enumerate(summaries)
]

df_id_query_engine_mapping = {
    f"pandas{idx}": df_query_engine
    for idx, df_query_engine in enumerate(df_query_engines)
}
# construct top-level vector index + query engine
vector_index = VectorStoreIndex(df_nodes)
vector_retriever = vector_index.as_retriever(similarity_top_k=1)

As per this link, Nebula needed, documnets... the request is how I pass my df_nodes to 'knowledgegraphindex' and use in the bundle as you suggested above.

https://docs.llamaindex.ai/en/stable/examples/index_structs/knowledge_graph/NebulaGraphKGIndexDemo/?h=nebu

from llama_index.core import StorageContext
from llama_index.graph_stores.nebula import NebulaGraphStore
graph_store = NebulaGraphStore(
    space_name=space_name,
    edge_types=edge_types,
    rel_prop_names=rel_prop_names,
    tags=tags,
)

storage_context = StorageContext.from_defaults(graph_store=graph_store)

index = KnowledgeGraphIndex.from_documents(
    documents,
    storage_context=storage_context,
    max_triplets_per_chunk=2,
    include_embeddings=True,
)

Your suggestion from the top:

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

# Initialize the LLM
llm = LLM(...)  # Provide your LLM instance here

# Initialize the KnowledgeGraphRAGRetriever
kg_retriever = KnowledgeGraphRAGRetriever(
    storage_context=storage_context,
    llm=llm,
    entity_extract_fn=entity_extract_fn,
    synonym_expand_fn=synonym_expand_fn,
    retriever_mode="keyword_embedding",
    verbose=True
)

# Construct top-level vector index + query engine
vector_index = VectorStoreIndex(doc_nodes + df_nodes)
vector_retriever = vector_index.as_retriever(similarity_top_k=1)

# Initialize the RecursiveRetriever
recursive_retriever = RecursiveRetriever(
    root_id="vector",
    retriever_dict={"vector": vector_retriever, "kg": kg_retriever},
    query_engine_dict=df_id_query_engine_mapping,
    verbose=False
)

# Example usage of the recursive retriever
query_bundle = QueryBundle(query_str="What is the capital of France?")
nodes_with_scores = recursive_retriever._retrieve(query_bundle)

# Initialize the KnowledgeGraphQueryEngine with the storage context
query_engine = KnowledgeGraphQueryEngine(storage_context=storage_context)

# Retrieve the final response
final_response = query_engine._query(query_bundle)

# Print the final response
print(final_response)

[dosubot[bot]]

To pass your df_nodes to KnowledgeGraphIndex and use it effectively in the query process with Nebula Graph store, you can follow these steps:

1. Install necessary packages:

   %pip install llama-index-llms-openai
   %pip install llama-index-embeddings-openai
   %pip install llama-index-graph-stores-nebula
   %pip install llama-index-llms-azure-openai
   %pip install nebula3-python

2. Prepare the NebulaGraphStore and StorageContext:

   from llama_index.core import StorageContext
   from llama_index.graph_stores.nebula import NebulaGraphStore
   
   space_name = "llamaindex"
   edge_types = ["relationship"]
   rel_prop_names = ["relationship"]
   tags = ["entity"]
   
   graph_store = NebulaGraphStore(
       space_name=space_name,
       edge_types=edge_types,
       rel_prop_names=rel_prop_names,
       tags=tags,
   )
   storage_context = StorageContext.from_defaults(graph_store=graph_store)

3. Generate a KnowledgeGraphIndex with your df_nodes:
   Assuming df_nodes is a DataFrame containing your nodes data, you can convert it to a format suitable for KnowledgeGraphIndex.

   from llama_index.core import KnowledgeGraphIndex
   from llama_index.core.schema import BaseNode
   
   # Assuming df_nodes is a DataFrame with your nodes data
   documents = df_nodes.to_dict(orient='records')
   nodes = [BaseNode(**doc) for doc in documents]
   
   kg_index = KnowledgeGraphIndex(
       nodes=nodes,
       storage_context=storage_context,
       max_triplets_per_chunk=10,
       include_embeddings=True,
   )

4. Set up the LLM (OpenAI or Azure OpenAI):
   For OpenAI:

   import os
   from llama_index.llms.openai import OpenAI
   from llama_index.core import Settings
   
   os.environ["OPENAI_API_KEY"] = "INSERT OPENAI KEY"
   llm = OpenAI(temperature=0, model="gpt-3.5-turbo")
   
   Settings.llm = llm
   Settings.chunk_size = 512

   For Azure OpenAI:

   import os
   import openai
   from llama_index.llms.azure_openai import AzureOpenAI
   from llama_index.embeddings.openai import OpenAIEmbedding
   from llama_index.core import Settings
   
   openai.api_type = "azure"
   openai.api_base = "https://<foo-bar>.openai.azure.com"
   openai.api_version = "2022-12-01"
   os.environ["OPENAI_API_KEY"] = "<your-openai-key>"
   openai.api_key = os.getenv("OPENAI_API_KEY")
   
   llm = AzureOpenAI(
       model="<foo-bar-model>",
       engine="<foo-bar-deployment>",
       temperature=0,
       api_key=openai.api_key,
       api_type=openai.api_type,
       api_base=openai.api_base,
       api_version=openai.api_version,
   )
   embedding_model = OpenAIEmbedding(
       model="text-embedding-ada-002",
       deployment_name="<foo-bar-deployment>",
       api_key=openai.api_key,
       api_base=openai.api_base,
       api_type=openai.api_type,
       api_version=openai.api.version,
   )
   
   Settings.llm = llm
   Settings.chunk_size = 512
   Settings.embed_model = embedding_model

5. Query the Knowledge Graph using Natural Language:

   from llama_index.core.query_engine import KnowledgeGraphQueryEngine
   
   query_engine = KnowledgeGraphQueryEngine(
       storage_context=storage_context,
       llm=llm,  # Assuming llm is already defined
       verbose=True,
   )
   
   response = query_engine.query("Tell me about Peter Quill?")
   print(response)

Alternatively, you can use the NebulaGraphQueryEnginePack for more advanced query capabilities:

from llama_index.packs.nebulagraph_query_engine.base import NebulaGraphQueryEnginePack, NebulaGraphQueryEngineType

# Initialize the query engine pack
query_engine_pack = NebulaGraphQueryEnginePack(
    username="your_nebula_username",
    password="your_nebula_password",
    ip_and_port="127.0.0.1:9669",
    space_name=space_name,
    edge_types=edge_types,
    rel_prop_names=rel_prop_names,
    tags=tags,
    max_triplets_per_chunk=10,
    docs=nodes,  # Pass your nodes here
    query_engine_type=NebulaGraphQueryEngineType.KG_QE
)

# Run a query
response = query_engine_pack.run("Tell me about Peter Quill?")
print(response)

This setup will allow you to pass your df_nodes to KnowledgeGraphIndex and use it effectively in the query process with Nebula Graph store ^[1][2][3][4].

To continue talking to Dosu, mention @dosu.

[mraguth]

@dosu In this below code, I need help. I have just dataframe called 'df'. How I convert that to 'DataFrame with nodes data'?

# Assuming df_nodes is a DataFrame with your nodes data
documents = df_nodes.to_dict(orient='records')
nodes = [BaseNode(**doc) for doc in documents]

[mraguth]

@dosu.. my pandas dataframe has 3 columns: ID, Pathway, Process ---> That table need to be converted to nodes.

[dosubot[bot]]

To convert your pandas dataframe df with columns ID, Pathway, and Process into the required format for nodes to be used with KnowledgeGraphIndex, you can use the following code:

import pandas as pd

# Assuming df is your dataframe
# df = pd.DataFrame({
#     'ID': [...],
#     'Pathway': [...],
#     'Process': [...]
# })

# Create the nodes dictionary
nodes = {}
for index, row in df.iterrows():
    nodes[row['ID']] = {
        "doc_id": row['ID'],
        "embedding": None,
        "extra_info": None,
        "index": index,
        "child_indices": [],
        "ref_doc_id": None,
        "node_info": {
            "Pathway": row['Pathway'],
            "Process": row['Process']
        }
    }

# Now nodes is in the required format
print(nodes)

This code will create a dictionary nodes where each key is the ID from your dataframe, and the value is a dictionary containing the required fields for KnowledgeGraphIndex ^[1].

To continue talking to Dosu, mention @dosu.