Introducing ToolsRetriever in the Neo4j GraphRAG Python Package

The Neo4j GraphRAG Python library has introduced a powerful new feature: the ToolsRetriever class and the Retriever.convert_to_tool() method. This addition enhances the library’s flexibility, enabling you to leverage multiple retrievers as tools in a single query, with an LLM intelligently selecting the most appropriate retriever(s) for the task. In this blog post, we’ll explore the motivation behind this feature, its implementation, and how it can be used to build dynamic retrieval-augmented generation (RAG) applications.

Motivation: Flexible and Intelligent Retrieval

RAG systems rely on retrievers to fetch relevant data from a knowledge base to provide context for an LLM’s response. Traditionally, a single retriever is used for a query, but different queries may require different types of data or retrieval strategies. For example, a query about “tomorrow’s schedule” might need a calendar-based retriever, while a query about “tomorrow’s weather” might require a weather-specific retriever.

The ToolsRetriever addresses this challenge by allowing you to combine multiple retrievers (as tools) into a single interface. Instead of manually selecting a retriever, the system uses an LLM to dynamically choose the most suitable tool(s) for a given query. This abstraction simplifies the development process and makes RAG pipelines more adaptive and efficient, enabling a single pass to handle diverse queries.

Additionally, the Retriever.convert_to_tool() method allows any retriever adhering to the Retriever interface to be converted into a tool compatible with ToolsRetriever. This is particularly useful for scenarios where you want to combine different retrieval strategies, such as a VectorRetriever for semantic search and a Text2CypherRetriever for structured graph queries, as a fallback or complementary mechanism.

Key Features

ToolsRetriever

The ToolsRetriever is a new class that integrates multiple tools (including retrievers converted to tools) and uses an LLM to decide which tool(s) to invoke based on the query. It adheres to the Retriever interface, making it seamless to integrate with the existing GraphRAG class for end-to-end RAG workflows.

Here’s an example of how to use ToolsRetriever:

from neo4j_graphrag.llm import OpenAILLM
from neo4j_graphrag.retrievers import ToolsRetriever
from neo4j_graphrag.graphrag import GraphRAG
import neo4j

# Hypothetical imports from a reader's file system
from my_tools import CalendarTool, WeatherTool

# Initialize tools and LLM
calendar_tool = CalendarTool()
weather_tool = WeatherTool()
llm = OpenAILLM()
driver = neo4j.GraphDatabase.driver(…)

# Create ToolsRetriever with multiple tools
tools_retriever = ToolsRetriever(
  driver=driver,
  llm=llm,
  tools=[calendar_tool, weather_tool],
)

# Integrate with GraphRAG
graphrag = GraphRAG(
  llm=llm,
  retriever=tools_retriever,
)
# Perform a search
result = graphrag.search(query_text="Tell me about tomorrow", return_context=False)

In this example, the LLM determines whether the query “Tell me about tomorrow” requires the CalendarTool, the WeatherTool, or both, ensuring the most relevant data is retrieved.

Retriever.convert_to_tool()

The Retriever.convert_to_tool() method allows any retriever (e.g., VectorRetriever, Text2CypherRetriever) to be converted into a tool that can be used by ToolsRetriever. This is particularly powerful for combining different retrieval strategies. For instance, a VectorRetriever might excel at semantic searches, while a Text2CypherRetriever is better suited for precise, structured queries on a Neo4j graph.

Here’s an example of converting a retriever to a tool:

from neo4j_graphrag.retrievers import VectorRetriever, Text2CypherRetriever
from neo4j_graphrag.retrievers import ToolsRetriever
from neo4j_graphrag.llm import OpenAILLM
from neo4j_graphrag.graphrag import GraphRAG
import neo4j

# Initialize retrievers and LLM
driver = neo4j.GraphDatabase.driver(…)
llm = OpenAILLM()
vector_retriever = VectorRetriever(driver, …)
text2cypher_retriever = Text2CypherRetriever(driver, …)

# Convert retrievers to tools
vector_tool = vector_retriever.convert_to_tool()
text2cypher_tool = text2cypher_retriever.convert_to_tool()

# Create ToolsRetriever
tools_retriever = ToolsRetriever(
  driver=driver,
  llm=llm,
  tools=[vector_tool, text2cypher_tool],
)

# Use in GraphRAG
graphrag = GraphRAG(
  llm=llm,
  retriever=tools_retriever,
)

# Perform a search
result = graphrag.search(query_text="Find documents about AI and their authors", return_context=False)

In this case, the LLM decides whether to use the VectorRetriever (for semantic similarity) or the Text2CypherRetriever (for structured queries) based on the query’s intent.

Implementation Details

ToolsRetriever

The ToolsRetriever class was designed to integrate seamlessly with the existing Neo4j GraphRAG framework. Key implementation details include:

• LLM-driven tool selection: The ToolsRetriever uses an LLM to analyze the query and select the most appropriate tool(s). This ensures that the system dynamically adapts to the query’s requirements.
• Tool name uniqueness: The implementation includes validation to ensure that all tools have unique names, preventing conflicts during tool selection.
• Consistent result formatting: The ToolsRetriever handles RetrieverResult objects from tools, ensuring that results are consistently formatted while preserving each retriever’s formatting logic.
• Integration with GraphRAG: By adhering to the Retriever interface, ToolsRetriever can be used as a drop-in replacement for other retrievers in the GraphRAG class.

Retriever.convert_to_tool()

The convert_to_tool() method was added to the Retriever base class, enabling any retriever to be converted into a tool. Key aspects of its implementation include:

• Abstract get_parameters() method: A new abstract method, get_parameters(), was added to the Retriever base class, requiring all concrete retriever classes to define their parameters. This ensures compatibility when converting to tools.
• Search method integration: The convert_to_tool() method uses the retriever’s search() method to ensure consistent result formatting, leveraging the retriever’s built-in result formatter.
• Metadata and attribution: The converted tool maintains metadata and attribution, ensuring traceability of the data source.

Benefits

1. Dynamic tool selection: The LLM intelligently selects the best tool(s) for a query, reducing the need for manual retriever configuration.
2. Flexibility: You can combine multiple retrievers (e.g., VectorRetriever, Text2CypherRetriever) as tools, enabling hybrid retrieval strategies.
3. Seamless integration: The ToolsRetriever integrates with the existing GraphRAG framework, maintaining compatibility with other components.
4. Extensibility: The convert_to_tool() method allows any retriever to be used as a tool, making the system highly extensible.
5. Improved user experience: By abstracting tool selection, you can focus on building applications rather than managing retriever logic.

Summary

The ToolsRetriever and Retriever.convert_to_tool() method represent a step forward for the Neo4j GraphRAG Python library. By enabling dynamic, LLM-driven selection of multiple retrievers as tools, this feature makes RAG pipelines more flexible and powerful. Whether you’re building applications that require diverse data sources or need a fallback mechanism for different retrieval strategies, this feature provides a robust and extensible solution.

To get started, check out the updated examples in the Neo4j GraphRAG repository and experiment with combining your own retrievers as tools.

Install and use it in your application using:

pip install neo4j-graphrag

We look forward to seeing how the community leverages this feature to build future GraphRAG applications!

Resources

• Free book: Essential GraphRAG by Tomaž Bratanic and Oskar Hane
• Blog: What Is Retrieval-Augmented Generation (RAG)?
• Blog: What Is GraphRAG?

---

Introducing ToolsRetriever in the Neo4j GraphRAG Python Package was originally published in Neo4j Developer Blog on Medium, where people are continuing the conversation by highlighting and responding to this story.