feat: KG RAG query engine, enable Graph RAG on existing KGs

[wey-gu]

Description

• Now we could do Graph RAG towards existing KGs
• Enabled nl2cypher for older NebulaGraph, with improved schema fetching impl., reported by @pachgadehardik via here

The shape now is RFC @logan-markewich @jerryjliu , in case good to go I'll impl. the async functions, too.

Type of Change

Please delete options that are not relevant.

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

How Has This Been Tested?

Please describe the tests that you ran to verify your changes. Provide instructions so we can reproduce. Please also list any relevant details for your test configuration

• Added new unit/integration tests
• Added new notebook (that tests end-to-end)
• I stared at the code and made sure it makes sense

Suggested Checklist:

• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[logan-markewich]

I think this looks pretty good so far! Thanks for contributing this! 🙏🏻

Just to confirm how I think this is working

• Take a query, extract out keywords and synonyms of those keywords
• Use those to extract triplets that mention these keywords/synonyms
• Optionally combine those retrieved triplets with a text2cypher query result?

[logan-markewich]

can _get_entities and _expand_synonyms be the same function? It seems like the only difference is the mode and prompt, which could be inputs to a unified function?

Would help clean up the code a bit I think!

[logan-markewich]

nodes_keyword or nodes_embedding might not be defined here depending on the mode? We might want to just extend directly after calling each retriever

i.e
nodes.extend(self._retrieve_keyword(query_bundle))

[logan-markewich]

can we move this into a from_defaults() function? Since a majority of the arguments are used only in the retriever, we can make the retriever an argument to the class instead, and handle setting up the retriever in from_defaults()?

[logan-markewich]

Actually, is this entire class really needed? Seems like to me only the retriever is needed, and the retriever could be used with the normal RetrieverQueryEngine()?

[wey-gu]

Yes, the whole thing of RAG is actually a retriever, will do so with RetrieverQueryEngine!(👍🏻)

[wey-gu]

I think this looks pretty good so far! Thanks for contributing this! 🙏🏻

Just to confirm how I think this is working

• Take a query, extract out keywords and synonyms of those keywords
• Use those to extract triplets that mention these keywords/synonyms
• Optionally combine those retrieved triplets with a text2cypher query result?

Yes, exactly!

And for the entity/keyword search of step 1, in the future, we could introduce the embedding-based capability to enable semantic searching, too.

For step 2: extract triplets from KG, or we could say it's getting n-depth subgraphs of those entities.

And for step 3, the text2cypher could be helpful but incomplete for the retrieval, and vice visa, thus options to combine the two could be one strategy to go with on KG.

Also now all comments are addressed! Thanks @logan-markewich !!!

For rendered notebook, see here

Also, I put a demo video in the notebook, too(not rendered in colab).

demo_text2cypher_vs_graph_rag.mp4

[logan-markewich]

nit: can we give these a more descriptive name?

[wey-gu]

Will do so :)

[logan-markewich]

Thanks for the updates @wey-gu !! I just made a bit of a change to move remove the RAG query engine, and just move the retriever to be used in a RetrieverQueryEngine (Since the RAG query engine was not doing anything that the RetrieverQueryEngine was not doing.

I updated the notebooks to confirm they still work as well, but you may need to update the screenshot of the streamline app :)

After that, I think it's basically good to merge!

[wey-gu]

Thanks for the updates @wey-gu !! I just made a bit of a change to move remove the RAG query engine, and just move the retriever to be used in a RetrieverQueryEngine (Since the RAG query engine was not doing anything that the RetrieverQueryEngine was not doing.

I updated the notebooks to confirm they still work as well, but you may need to update the screenshot of the streamline app :)

After that, I think it's basically good to merge!

Now it's elegant :) It should have been settled in pure retrievers!

The demo video could be as-is as it was created from the kg_index --> as query engine, but generationally enough to compare the two(text2cyper vs graph rag).

[logan-markewich]

Thanks for the amazing work on this @wey-gu 💪💪

[wey-gu]

Thanks for the amazing work on this @wey-gu 💪💪

Thanks a ton @logan-markewich for your amazing work and support 🤩

[pachgadehardik]

Thanks a ton guys. Really appreciate the updates done in this thread.

[BleakStone]

@wey-gu
awesome!!!
This implementation helped me a lot，but the rendered notebook here seems to have some version adaptation problems
Here is some of the code in rendered notebook：

from llama_index.query_engine import (
    KnowledgeGraphRAGRetriever,
    KnowledgeGraphRAGQueryEngine,
)

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

query_engine = KnowledgeGraphRAGQueryEngine(retriever=graph_rag_retriever)

1. I dont find the cls KnowledgeGraphRAGQueryEngine in llama-index==0.8.3，instead RetrieverQueryEngine

query_engine = RetrieverQueryEngine.from_args(
    graph_rag_retriever, service_context=service_context
)

2. Did parameters llm=llm have effect ? And the service_context also has an llm attribute

[wey-gu]

Wow, good to know it helps @BleakStone !

Sorry, we made refactor to not introduce a query engine but only a retriever(as it is only a retriever), please instead follow the final version of the notebook here instead.