query_planner_agent

langroid/agent/special/lance_rag/query_planner_agent.py

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.model_dump(), 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}
    """

answer_tool(msg)

Handle AnswerTool received from LanceRagAgent: Construct a QueryPlanAnswerTool with the answer, addressed to the Critic.

The recipient is explicit on purpose. Returning the tool bare leaves it with no addressee, and Task.step then offers the turn to THIS agent's own LLM before the Critic sub-task. That reply is accepted when it is non-null, and since llm_response above sets expecting_query_plan unconditionally, the LLM is told to produce a query plan again: it emits a duplicate query_plan, the Critic never receives the answer, query_plan_feedback never arrives, and the task spins until the inf-loop guard or max_turns (Task.result() then returns None). Naming the Critic makes _recipient_mismatch skip our own LLM, so the handoff is deterministic on any model.

Source code in langroid/agent/special/lance_rag/query_planner_agent.py

def answer_tool(self, msg: AnswerTool) -> ChatDocument:
    """Handle AnswerTool received from LanceRagAgent:
    Construct a QueryPlanAnswerTool with the answer, addressed to the Critic.

    The recipient is explicit on purpose. Returning the tool bare leaves it
    with no addressee, and `Task.step` then offers the turn to THIS agent's
    own LLM before the Critic sub-task. That reply is accepted when it is
    non-null, and since `llm_response` above sets `expecting_query_plan`
    unconditionally, the LLM is told to produce a query plan again: it
    emits a duplicate `query_plan`, the Critic never receives the answer,
    `query_plan_feedback` never arrives, and the task spins until the
    inf-loop guard or max_turns (`Task.result()` then returns None).
    Naming the Critic makes `_recipient_mismatch` skip our own LLM, so the
    handoff is deterministic on any model.
    """
    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 self.create_agent_response(
        recipient=self.config.critic_name,
        tool_messages=[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