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
# how many times re-trying query plan in response to feedback:
self.n_retries: int = 0
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
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
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.suggested_fix == "":
self.n_retries = 0
# 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
self.n_retries += 1
if self.n_retries >= self.config.max_retries:
# bail out to avoid infinite loop
self.n_retries = 0
return DONE + " " + NO_ANSWER
return f"""
here is FEEDBACK about your QUERY PLAN, and a SUGGESTED FIX.
Modify the QUERY PLAN if needed:
FEEDBACK: {msg.feedback}
SUGGESTED FIX: {msg.suggested_fix}
"""
|
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( # type: ignore
plan=self.curr_query_plan,
answer=self.result,
)
response_tmpl = self.create_agent_response()
# ... 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
if (
isinstance(msg, ChatDocument)
and not self.has_tool_message_attempt(msg)
and msg.metadata.sender == lr.Entity.LLM
):
# remind LLM to use the QueryPlanFeedbackTool
return """
You forgot to use the `query_plan` tool/function.
Re-try your response using the `query_plan` tool/function.
"""
return None
|