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
# 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)
self.enable_message(AnswerTool, use=False, handle=True)
# neither use nor handle! Added to "known" tools so that the Planner agent
# can avoid processing it
self.enable_message(QueryPlanAnswerTool, use=False, handle=False)
# LLM will not use this, so set use=False (Agent generates it)
self.enable_message(AgentDoneTool, use=False, handle=True)
|
query_plan(msg)
Valid, tool msg, forward chat_doc to RAG Agent.
Note this chat_doc will already have the
QueryPlanTool in its tool_messages list.
We just update the recipient to the doc_agent_name.
Source code in langroid/agent/special/lance_rag/query_planner_agent.py
| def query_plan(self, msg: QueryPlanTool) -> ForwardTool | str:
"""Valid, tool msg, forward chat_doc to RAG Agent.
Note this chat_doc will already have the
QueryPlanTool in its tool_messages list.
We just update the recipient to the doc_agent_name.
"""
# save, to be used to assemble QueryPlanResultTool
if len(msg.plan.dataframe_calc.split("\n")) > 1:
return "DATAFRAME CALCULATION must be a SINGLE LINE; Retry the `query_plan`"
self.curr_query_plan = msg.plan
self.expecting_query_plan = False
# To forward the QueryPlanTool to doc_agent, we could either:
# (a) insert `recipient` in the QueryPlanTool:
# QPWithRecipient = QueryPlanTool.require_recipient()
# qp = QPWithRecipient(**msg.dict(), recipient=self.config.doc_agent_name)
# return qp
#
# OR
#
# (b) create an agent response with recipient and tool_messages.
# response = self.create_agent_response(
# recipient=self.config.doc_agent_name, tool_messages=[msg]
# )
# return response
# OR
# (c) use the ForwardTool:
return ForwardTool(agent=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 | AgentDoneTool:
"""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.suggested_fix == ""
and NO_ANSWER not in self.result
and self.result != ""
):
# This means the result is good AND Query Plan is fine,
# as judged by Critic
# (Note sometimes critic may have empty suggested_fix even when
# the result is NO_ANSWER)
self.n_retries = 0 # good answer, so reset this
return AgentDoneTool(content=self.result)
self.n_retries += 1
if self.n_retries >= self.config.max_retries:
# bail out to avoid infinite loop
self.n_retries = 0
return AgentDoneTool(content=NO_ANSWER)
# there is a suggested_fix, OR the result is empty or NO_ANSWER
if self.result == "" or NO_ANSWER in self.result:
# if result is empty or NO_ANSWER, we should retry the query plan
feedback = """
There was no answer, which might mean there is a problem in your query.
"""
suggested = "Retry the `query_plan` to try to get a non-null answer"
else:
feedback = msg.feedback
suggested = msg.suggested_fix
self.expecting_query_plan = True
return f"""
here is FEEDBACK about your QUERY PLAN, and a SUGGESTED FIX.
Modify the QUERY PLAN if needed:
ANSWER: {self.result}
FEEDBACK: {feedback}
SUGGESTED FIX: {suggested}
"""
|
Handle AnswerTool received from LanceRagAgent:
Construct a QueryPlanAnswerTool with the answer
Source code in langroid/agent/special/lance_rag/query_planner_agent.py
| def answer_tool(self, msg: AnswerTool) -> QueryPlanAnswerTool:
"""Handle AnswerTool received from LanceRagAgent:
Construct a QueryPlanAnswerTool with the answer"""
self.result = msg.answer # save answer to interpret feedback later
assert self.curr_query_plan is not None
query_plan_answer_tool = QueryPlanAnswerTool(
plan=self.curr_query_plan,
answer=msg.answer,
)
self.curr_query_plan = None # reset
return query_plan_answer_tool
|
handle_message_fallback(msg)
Remind to use QueryPlanTool if we are expecting it.
Source code in langroid/agent/special/lance_rag/query_planner_agent.py
| def handle_message_fallback(
self, msg: str | ChatDocument
) -> str | ChatDocument | None:
"""
Remind to use QueryPlanTool if we are expecting it.
"""
if self.expecting_query_plan and self.n_query_plan_reminders < 5:
self.n_query_plan_reminders += 1
return """
You FORGOT to use the `query_plan` tool/function,
OR you had a WRONG JSON SYNTAX when trying to use it.
Re-try your response using the `query_plan` tool/function CORRECTLY.
"""
self.n_query_plan_reminders = 0 # reset
return None
|