LanceQueryPlanAgent is a ChatAgent created with a specific document schema.
Given a QUERY, the LLM constructs a Query Plan consisting of:
- filter condition if needed (or empty string if no filter is needed)
- query - a possibly rephrased query that can be used to match the content field
- dataframe_calc - a Pandas-dataframe calculation/aggregation string, possibly empty
- original_query - the original query for reference
This agent has access to two tools:
- QueryPlanTool, which is used to generate the Query Plan, and the handler of
this tool simply passes it on to the RAG agent named in config.doc_agent_name.
- QueryPlanFeedbackTool, which is used to handle feedback on the Query Plan and
Result from the RAG agent. The QueryPlanFeedbackTool is used by
the QueryPlanCritic, who inserts feedback into the feedback field
LanceQueryPlanAgent(config)
Bases: ChatAgent
Source code in langroid/agent/special/lance_rag/query_planner_agent.py
| def __init__(self, config: LanceQueryPlanAgentConfig):
super().__init__(config)
self.config: LanceQueryPlanAgentConfig = config
self.curr_query_plan: QueryPlan | None = None
self.result: str = "" # answer received from LanceRAG
# This agent should generate the QueryPlanTool
# as well as handle it for validation
self.enable_message(QueryPlanTool, use=True, handle=True)
self.enable_message(QueryPlanFeedbackTool, use=False, handle=True)
|
query_plan(msg)
Valid, forward to RAG Agent
Source code in langroid/agent/special/lance_rag/query_planner_agent.py
| def query_plan(self, msg: QueryPlanTool) -> str:
"""Valid, forward to RAG Agent"""
# save, to be used to assemble QueryPlanResultTool
self.curr_query_plan = msg.plan
return PASS_TO + self.config.doc_agent_name
|
query_plan_feedback(msg)
Process Critic feedback on QueryPlan + Answer from RAG Agent
Source code in langroid/agent/special/lance_rag/query_planner_agent.py
| def query_plan_feedback(self, msg: QueryPlanFeedbackTool) -> str:
"""Process Critic feedback on QueryPlan + Answer from RAG Agent"""
# We should have saved answer in self.result by this time,
# since this Agent seeks feedback only after receiving RAG answer.
if msg.feedback == "":
# This means the Query Plan or Result is good, as judged by Critic
if self.result == "":
# This was feedback for query with no result
return "QUERY PLAN LOOKS GOOD!"
elif self.result == NO_ANSWER:
return NO_ANSWER
else: # non-empty and non-null answer
return DONE + " " + self.result
return f"""
here is FEEDBACK about your QUERY PLAN. Modify it if needed:
{msg.feedback}
"""
|
handle_message_fallback(msg)
Process answer received from RAG Agent
Construct a QueryPlanAnswerTool with the answer,
and forward to Critic for feedback.
Source code in langroid/agent/special/lance_rag/query_planner_agent.py
| def handle_message_fallback(
self, msg: str | ChatDocument
) -> str | ChatDocument | None:
"""
Process answer received from RAG Agent:
Construct a QueryPlanAnswerTool with the answer,
and forward to Critic for feedback.
"""
# TODO we don't need to use this fallback method. instead we can
# first call result = super().agent_response(), and if result is None,
# then we know there was no tool, so we run below code
if (
isinstance(msg, ChatDocument)
and self.curr_query_plan is not None
and msg.metadata.parent is not None
):
# save result, to be used in query_plan_feedback()
self.result = msg.content
# assemble QueryPlanAnswerTool...
query_plan_answer_tool = QueryPlanAnswerTool(
plan=self.curr_query_plan,
answer=self.result,
)
response_tmpl = self.agent_response_template()
# ... add the QueryPlanAnswerTool to the response
# (Notice how the Agent is directly sending a tool, not the LLM)
response_tmpl.tool_messages = [query_plan_answer_tool]
# set the recipient to the Critic so it can give feedback
response_tmpl.metadata.recipient = self.config.critic_name
self.curr_query_plan = None # reset
return response_tmpl
return None
|