Add ToolsRetriever class and convert_retriever_to_tool() function #332
Conversation
|
|
||
| # Define parameters for the static retriever tool | ||
| static_parameters = ObjectParameter( | ||
| description="Parameters for the Neo4j information retriever", |
There was a problem hiding this comment.
Is this description really required? It doesn't seem to have a real added-value.
| ) | ||
|
|
||
| # Convert the retriever to a tool with specific parameters | ||
| static_tool = convert_retriever_to_tool( |
There was a problem hiding this comment.
I'd say this should be a method from the Retriever class, and parameters should be encapsulated in this class as well, these parameters are bound to the search method and won't change from one instance to another.
So something like:
class Retriever:
def get_parameters(self) -> ObjectParameter:
raise NotImplementedError() # need to be implemented in subclasses
def convert_to_tool(self, name: str, description: Optional[str] = None) -> Tool:
# rest of the function goes here
Note: as a future improvement, I think we could infer the parameters from the search method signature without having to redeclare it.
|
|
||
| def __init__( | ||
| self, | ||
| driver: neo4j.Driver, |
There was a problem hiding this comment.
Should we consider updating the Retriever interface if this one does not use the driver?
| # Extract arguments from the tool call | ||
| tool_args = tool_call.arguments or {} | ||
|
|
||
| # Always include the query_text in the arguments for tools that might need it |
There was a problem hiding this comment.
Shouldn't the LLM include the query if needed by the tool? It should be in the tool parameters in this case?
| # Execute the tool with the provided arguments | ||
| tool_result = selected_tool.execute(**tool_args) | ||
| # If the tool result is a RawSearchResult, extract its records | ||
| if hasattr(tool_result, "records"): |
There was a problem hiding this comment.
We can maybe enforce the fact that Tool.execute must return a list of something (or list of strings?), or do we really want to create fake records?
| @@ -211,23 +211,28 @@ def validate_properties(self) -> "ObjectParameter": | |||
| class Tool(ABC): | |||
| """Abstract base class defining the interface for all tools in the neo4j-graphrag library.""" | |||
|
|
|||
| _name: str | |||
There was a problem hiding this comment.
Do you remember why we need to have these declarations here? Mypy stuff?
| """ | ||
| # Use provided name or infer it from the retriever | ||
| if name is None: | ||
| name = getattr(retriever, "name", None) or getattr( |
There was a problem hiding this comment.
If we expected retrievers to have a name attribute, maybe we should add this to the interface?
| # arguments like query_text, top_k, etc., passed as keyword arguments. | ||
| # The Tool's 'parameters' definition (e.g., ObjectParameter) ensures | ||
| # that these arguments are provided in kwargs when Tool.execute is called. | ||
| return retriever.get_search_results(**kwargs) |
There was a problem hiding this comment.
Why not using the search method?
Description
In addition to the major parts of this PR, I've refactored and moved the
tool.pyfile into its own directory to group relevant tool files together. This move touches a lot of files because of updated imports.ToolsRetrieverNow that we have tool support in the LLM interface, let's create a retriever that can use one or more tools to retrieve data.
This retriever follows the Retriever interface so it can be used within the GraphRAG class to get the full e2e experience (see example file).
The ToolsRetriever uses an LLM to decide on what tools to use to find the relevant data.
convert_retriever_to_tool ()This function is a way to convert a
Retrieverto aToolso it can be used within theToolsRetriever. This is useful when you might want to have both aVectorRetrieverand aText2CypherRetreiveras a fallback.See new example files for usage.
Type of Change
How Has This Been Tested?
Checklist
The following requirements should have been met (depending on the changes in the branch):