langroid 0.8.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

langroid/agent/chat_agent.py

@@ -47,6 +47,8 @@ class ChatAgentConfig(AgentConfig):
             the OpenAI tool-call API is used, rather than the older/deprecated
             function-call API. However the tool-call API has some tricky aspects,
             hence we set this to False by default.
+        enable_orchestration_tool_handling: whether to enable handling of orchestration
+            tools, e.g. ForwardTool, DoneTool, PassTool, etc.
     """
 
     system_message: str = "You are a helpful assistant."
@@ -54,6 +56,7 @@ class ChatAgentConfig(AgentConfig):
     use_tools: bool = False
     use_functions_api: bool = True
     use_tools_api: bool = False
+    enable_orchestration_tool_handling: bool = True
 
     def _set_fn_or_tools(self, fn_available: bool) -> None:
         """
@@ -145,6 +148,30 @@ class ChatAgent(Agent):
         self.llm_functions_usable: Set[str] = set()
         self.llm_function_force: Optional[Dict[str, str]] = None
 
+        if self.config.enable_orchestration_tool_handling:
+            # Only enable HANDLING by `agent_response`, NOT LLM generation of these.
+            # This is useful where tool-handlers or agent_response generate these
+            # tools, and need to be handled.
+            # We don't want enable orch tool GENERATION by default, since that
+            # might clutter-up the LLM system message unnecessarily.
+            from langroid.agent.tools.orchestration import (
+                AgentDoneTool,
+                AgentSendTool,
+                DonePassTool,
+                DoneTool,
+                ForwardTool,
+                PassTool,
+                SendTool,
+            )
+
+            self.enable_message(ForwardTool, use=False, handle=True)
+            self.enable_message(DoneTool, use=False, handle=True)
+            self.enable_message(AgentDoneTool, use=False, handle=True)
+            self.enable_message(PassTool, use=False, handle=True)
+            self.enable_message(DonePassTool, use=False, handle=True)
+            self.enable_message(SendTool, use=False, handle=True)
+            self.enable_message(AgentSendTool, use=False, handle=True)
+
     @staticmethod
     def from_id(id: str) -> "ChatAgent":
         """
@@ -272,19 +299,25 @@ class ChatAgent(Agent):
         Returns:
             str: formatting rules
         """
-        enabled_classes: List[Type[ToolMessage]] = list(self.llm_tools_map.values())
-        if len(enabled_classes) == 0:
+        # ONLY Usable tools (i.e. LLM-generation allowed),
+        usable_tool_classes: List[Type[ToolMessage]] = [
+            t
+            for t in list(self.llm_tools_map.values())
+            if not t._handle_only
+            and t.default_value("request") in self.llm_tools_usable
+        ]
+
+        if len(usable_tool_classes) == 0:
             return "You can ask questions in natural language."
         json_instructions = "\n\n".join(
             [
                 msg_cls.json_instructions(tool=self.config.use_tools)
-                for _, msg_cls in enumerate(enabled_classes)
-                if msg_cls.default_value("request") in self.llm_tools_usable
+                for msg_cls in usable_tool_classes
             ]
         )
         # if any of the enabled classes has json_group_instructions, then use that,
         # else fall back to ToolMessage.json_group_instructions
-        for msg_cls in enabled_classes:
+        for msg_cls in usable_tool_classes:
             if hasattr(msg_cls, "json_group_instructions") and callable(
                 getattr(msg_cls, "json_group_instructions")
             ):
@@ -420,9 +453,16 @@ class ChatAgent(Agent):
         # remove leading and trailing newlines and other whitespace
         return LLMMessage(role=Role.SYSTEM, content=content.strip())
 
+    def unhanded_tools(self) -> set[str]:
+        """The set of tools that are known but not handled.
+        Useful in task flow: an agent can refuse to accept an incoming msg
+        when it only has unhandled tools.
+        """
+        return self.llm_tools_known - self.llm_tools_handled
+
     def enable_message(
         self,
-        message_class: Optional[Type[ToolMessage]],
+        message_class: Optional[Type[ToolMessage] | List[Type[ToolMessage]]],
         use: bool = True,
         handle: bool = True,
         force: bool = False,
@@ -435,8 +475,10 @@ class ChatAgent(Agent):
         - tool HANDLING (i.e. the agent can handle JSON from this tool),
 
         Args:
-            message_class: The ToolMessage class to enable,
+            message_class: The ToolMessage class OR List of such classes to enable,
                 for USE, or HANDLING, or both.
+                If this is a list of ToolMessage classes, then the remain args are
+                applied to all classes.
                 Optional; if None, then apply the enabling to all tools in the
                 agent's toolset that have been enabled so far.
             use: IF True, allow the agent (LLM) to use this tool (or all tools),
@@ -448,12 +490,23 @@ class ChatAgent(Agent):
                  `force` is ignored if `message_class` is None.
             require_recipient: whether to require that recipient be specified
                 when using the tool message (only applies if `use` is True).
-            require_defaults: whether to include fields that have default values,
+            include_defaults: whether to include fields that have default values,
                 in the "properties" section of the JSON format instructions.
                 (Normally the OpenAI completion API ignores these fields,
                 but the Assistant fn-calling seems to pay attn to these,
                 and if we don't want this, we should set this to False.)
         """
+        if message_class is not None and isinstance(message_class, list):
+            for mc in message_class:
+                self.enable_message(
+                    mc,
+                    use=use,
+                    handle=handle,
+                    force=force,
+                    require_recipient=require_recipient,
+                    include_defaults=include_defaults,
+                )
+            return None
         if require_recipient and message_class is not None:
             message_class = message_class.require_recipient()
         super().enable_message_handling(message_class)  # enables handling only
@@ -468,6 +521,8 @@ class ChatAgent(Agent):
                 self.llm_function_force = None
 
         for t in tools:
+            self.llm_tools_known.add(t)
+
             if handle:
                 self.llm_tools_handled.add(t)
                 self.llm_functions_handled.add(t)
@@ -662,7 +717,11 @@ class ChatAgent(Agent):
                 # either the message is a str, or it is a fresh ChatDocument
                 # different from the last message in the history
                 llm_msgs = ChatDocument.to_LLMMessage(message, self.oai_tool_calls)
-
+                # LLM only responds to the content, so only those msgs with
+                # non-empty content should be kept
+                llm_msgs = [m for m in llm_msgs if m.content != ""]
+                if len(llm_msgs) == 0:
+                    return [], 0
                 # process tools if any
                 done_tools = [m.tool_call_id for m in llm_msgs if m.role == Role.TOOL]
                 self.oai_tool_calls = [

langroid/agent/chat_document.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import copy
 import json
+from collections import OrderedDict
 from enum import Enum
 from typing import Any, Dict, List, Optional, Union, cast
 
@@ -96,14 +97,16 @@ class ChatDocument(Document):
     and so does the Task.run() method.
 
     Attributes:
-        oai_tool_calls (List[OpenAIToolCall]): Tool-calls from an OpenAI-compatible API
-        oai_tool_id2results (Dict[str, str]): Results of tool-calls from OpenAI
-            (dict is a map of tool_id -> result)
+        oai_tool_calls (Optional[List[OpenAIToolCall]]):
+            Tool-calls from an OpenAI-compatible API
+        oai_tool_id2results (Optional[OrderedDict[str, str]]):
+            Results of tool-calls from OpenAI (dict is a map of tool_id -> result)
         oai_tool_choice: ToolChoiceTypes | Dict[str, str]: Param controlling how the
             LLM should choose tool-use in its response
             (auto, none, required, or a specific tool)
-        function_call (LLMFunctionCall): Function-call from an OpenAI-compatible API
-                (deprecated; use oai_tool_calls instead)
+        function_call (Optional[LLMFunctionCall]):
+            Function-call from an OpenAI-compatible API
+                (deprecated by OpenAI, in favor of tool-calls)
         tool_messages (List[ToolMessage]): Langroid ToolMessages extracted from
             - `content` field (via JSON parsing),
             - `oai_tool_calls`, or
@@ -113,10 +116,14 @@ class ChatDocument(Document):
     """
 
     oai_tool_calls: Optional[List[OpenAIToolCall]] = None
-    oai_tool_id2result: Optional[Dict[str, str]] = None
+    oai_tool_id2result: Optional[OrderedDict[str, str]] = None
     oai_tool_choice: ToolChoiceTypes | Dict[str, Dict[str, str] | str] = "auto"
     function_call: Optional[LLMFunctionCall] = None
-    tool_messages: List[ToolMessage] = []
+    tool_messages: List[ToolMessage] = []  # only handle-able tools
+    # all known tools in the msg that are in an agent's llm_tools_known list,
+    # even if non-used/handled
+    all_tool_messages: List[ToolMessage] = []
+
     metadata: ChatDocMetaData
     attachment: None | ChatDocAttachment = None
 
@@ -136,6 +143,8 @@ class ChatDocument(Document):
     def deepcopy(doc: ChatDocument) -> ChatDocument:
         new_doc = copy.deepcopy(doc)
         new_doc.metadata.id = ObjectRegistry.new_id()
+        new_doc.metadata.child_id = ""
+        new_doc.metadata.parent_id = ""
         ObjectRegistry.register_object(new_doc)
         return new_doc
 

langroid/agent/openai_assistant.py

@@ -160,7 +160,7 @@ class OpenAIAssistant(ChatAgent):
 
     def enable_message(
         self,
-        message_class: Optional[Type[ToolMessage]],
+        message_class: Optional[Type[ToolMessage] | List[Type[ToolMessage]]],
         use: bool = True,
         handle: bool = True,
         force: bool = False,
@@ -173,6 +173,17 @@ class OpenAIAssistant(ChatAgent):
         fn-calling seems to pay attn to these, and if we don't want this,
         we should set this to False.
         """
+        if message_class is not None and isinstance(message_class, list):
+            for msg_class in message_class:
+                self.enable_message(
+                    msg_class,
+                    use=use,
+                    handle=handle,
+                    force=force,
+                    require_recipient=require_recipient,
+                    include_defaults=include_defaults,
+                )
+            return
         super().enable_message(
             message_class,
             use=use,

langroid/agent/special/lance_doc_chat_agent.py

@@ -17,10 +17,11 @@ from typing import Any, Dict, List, Tuple
 import pandas as pd
 
 from langroid.agent.special.doc_chat_agent import DocChatAgent, DocChatAgentConfig
-from langroid.agent.special.lance_tools import QueryPlanTool
+from langroid.agent.special.lance_tools import AnswerTool, QueryPlanTool
+from langroid.agent.tools.orchestration import AgentDoneTool
 from langroid.mytypes import DocMetaData, Document
 from langroid.parsing.table_loader import describe_dataframe
-from langroid.utils.constants import DONE, NO_ANSWER
+from langroid.utils.constants import NO_ANSWER
 from langroid.utils.pydantic_utils import (
     dataframe_to_documents,
 )
@@ -106,7 +107,7 @@ class LanceDocChatAgent(DocChatAgent):
             """
         return schema
 
-    def query_plan(self, msg: QueryPlanTool) -> str:
+    def query_plan(self, msg: QueryPlanTool) -> AgentDoneTool | str:
         """
         Handle the LLM's use of the FilterTool.
         Temporarily set the config filter and either return the final answer
@@ -120,13 +121,15 @@ class LanceDocChatAgent(DocChatAgent):
         except Exception as e:
             logger.error(f"Error setting up documents: {e}")
             # say DONE with err msg so it goes back to LanceFilterAgent
-            return f"""
-            {DONE} Possible Filter Error:\n {e}
-            
-            Note that only the following fields are allowed in the filter
-            of a query plan: 
-            {", ".join(self.config.filter_fields)}
-            """
+            return AgentDoneTool(
+                content=f"""
+                Possible Filter Error:\n {e}
+                
+                Note that only the following fields are allowed in the filter
+                of a query plan: 
+                {", ".join(self.config.filter_fields)}
+                """
+            )
 
         # update the filter so it is used in the DocChatAgent
         self.config.filter = plan.filter or None
@@ -139,22 +142,25 @@ class LanceDocChatAgent(DocChatAgent):
             # The calc step can later be done with a separate Agent/Tool.
             if plan.query is None or plan.query.strip() == "":
                 if plan.filter is None or plan.filter.strip() == "":
-                    return """DONE
-                    Cannot execute Query Plan since filter as well as 
-                    rephrased query are empty.
-                    """
+                    return AgentDoneTool(
+                        content="""
+                        Cannot execute Query Plan since filter as well as 
+                        rephrased query are empty.                    
+                        """
+                    )
                 else:
                     # no query to match, so just get all docs matching filter
                     docs = self.vecdb.get_all_documents(plan.filter)
             else:
                 _, docs = self.get_relevant_extracts(plan.query)
             if len(docs) == 0:
-                return DONE + " " + NO_ANSWER
-            result = self.vecdb.compute_from_docs(docs, plan.dataframe_calc)
-            return DONE + " " + result
+                return AgentDoneTool(content=NO_ANSWER)
+            answer = self.vecdb.compute_from_docs(docs, plan.dataframe_calc)
         else:
             # pass on the query so LLM can handle it
-            return plan.query
+            response = self.llm_response(plan.query)
+            answer = NO_ANSWER if response is None else response.content
+        return AgentDoneTool(tools=[AnswerTool(answer=answer)])
 
     def ingest_docs(
         self,
@@ -242,6 +248,7 @@ class LanceDocChatAgent(DocChatAgent):
             .replace("NOT", "not")
             .replace("'", "")
             .replace('"', "")
+            .replace(":", "--")
         )
 
         tbl = self.vecdb.client.open_table(self.vecdb.config.collection_name)

langroid/agent/special/lance_rag/critic_agent.py

@@ -26,8 +26,8 @@ from langroid.agent.special.lance_tools import (
     QueryPlanAnswerTool,
     QueryPlanFeedbackTool,
 )
-from langroid.mytypes import Entity
-from langroid.utils.constants import DONE, NO_ANSWER, PASS
+from langroid.agent.tools.orchestration import AgentDoneTool
+from langroid.utils.constants import NO_ANSWER
 
 logger = logging.getLogger(__name__)
 
@@ -50,6 +50,9 @@ class QueryPlanCriticConfig(LanceQueryPlanAgentConfig):
       to create a QUERY PLAN, to be handled by an ASSISTANT.
     - PANDAS-LIKE FILTER, WHICH CAN BE EMPTY (and it's fine if results sound reasonable)
       FILTER SHOULD ONLY BE USED IF EXPLICITLY REQUIRED BY THE QUERY.
+      This filter selects the documents over which the REPHRASED QUERY will be applied,
+      thus naturally, the Re-phrased Query should NOT mention any FILTER fields,
+      since it applies to the documents AFTER FILTERING.
     - REPHRASED QUERY (CANNOT BE EMPTY) that will be used to match against the 
       CONTENT (not filterable) of the documents.
       In general the REPHRASED QUERY should be relied upon to match the CONTENT 
@@ -61,9 +64,31 @@ class QueryPlanCriticConfig(LanceQueryPlanAgentConfig):
         The assistant will answer based on documents whose CONTENTS match the QUERY, 
         possibly REPHRASED. 
         !!!!****THE REPHRASED QUERY SHOULD NEVER BE EMPTY****!!!
+        
+        
     - DATAFRAME CALCULATION, which must be a SINGLE LINE calculation (or empty),
         [NOTE ==> This calculation is applied AFTER the FILTER and REPHRASED QUERY.],
     - ANSWER received from an assistant that used this QUERY PLAN.
+      IT IS TOTALLY FINE FOR THE ANSWER TO NOT MENTION ANY FILTERING CONDITIONS,
+      or if the ANSWER STATEMENT is MISSING SOME CRITERIA in the ORIGINAL QUERY.
+
+        Here is an example of a VALID Plan + Answer:
+        
+        ORIGINAL QUERY: "Which crime novels were written by Russian authors after 1900?"
+        FILTER: "author_nationality == 'Russian' and year_written > 1900"
+        REPHRASED QUERY: "crime novel" [NOTICE NO FILTER FIELDS MENTIONED!!!]
+        DATAFRAME CALC: ""
+        ANSWER: "The Master and Margarita by Mikhail Bulgakov" 
+            [NOTICE the answer does NOT need to say "crime novel" or "russian author"]
+            
+            
+        Other examples of VALID ANSWER for a given ORIGINAL QUERY:
+        
+        ORIGINAL QUERY: "Which mountain is taller than 8000 meters?"
+        ANSWER: "Mount Everest" [NOTICE no mention of "taller than 8000 meters"]
+        
+        ORIGINAL QUERY: "Which country has hosted the most olympics?"
+        ANSWER: "United States" [NOTICE no mention of "most olympics"]
 
     In addition to the above SCHEMA fields there is a `content` field which:
     - CANNOT appear in a FILTER, 
@@ -141,21 +166,28 @@ class QueryPlanCritic(ChatAgent):
         self.config = cfg
         self.enable_message(QueryPlanAnswerTool, use=False, handle=True)
         self.enable_message(QueryPlanFeedbackTool, use=True, handle=True)
+        self.enable_message(AgentDoneTool, use=False, handle=True)
+
+    def init_state(self) -> None:
+        self.expecting_feedback_tool = False
 
     def query_plan_answer(self, msg: QueryPlanAnswerTool) -> str:
         """Present query plan + answer in plain text (not JSON)
         so LLM can give feedback"""
+        self.expecting_feedback_tool = True
         return plain_text_query_plan(msg)
 
-    def query_plan_feedback(self, msg: QueryPlanFeedbackTool) -> str:
+    def query_plan_feedback(self, msg: QueryPlanFeedbackTool) -> AgentDoneTool:
         """Format Valid so return to Query Planner"""
-        return DONE + " " + PASS  # return to Query Planner
+        self.expecting_feedback_tool = False
+        # indicate this task is Done, and return the tool as result
+        return AgentDoneTool(tools=[msg])
 
     def handle_message_fallback(
         self, msg: str | ChatDocument
     ) -> str | ChatDocument | None:
         """Remind the LLM to use QueryPlanFeedbackTool since it forgot"""
-        if isinstance(msg, ChatDocument) and msg.metadata.sender == Entity.LLM:
+        if self.expecting_feedback_tool:
             return """
             You forgot to use the `query_plan_feedback` tool/function.
             Re-try your response using the `query_plan_feedback` tool/function,

langroid/agent/special/lance_rag/query_planner_agent.py

@@ -15,17 +15,19 @@ This agent has access to two tools:
 """
 
 import logging
+from typing import Optional
 
-import langroid as lr
 from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
 from langroid.agent.chat_document import ChatDocument
 from langroid.agent.special.lance_tools import (
+    AnswerTool,
     QueryPlan,
     QueryPlanAnswerTool,
     QueryPlanFeedbackTool,
     QueryPlanTool,
 )
-from langroid.utils.constants import DONE, NO_ANSWER, PASS_TO
+from langroid.agent.tools.orchestration import AgentDoneTool, ForwardTool
+from langroid.utils.constants import NO_ANSWER
 
 logger = logging.getLogger(__name__)
 
@@ -39,14 +41,14 @@ class LanceQueryPlanAgentConfig(ChatAgentConfig):
     max_retries: int = 5  # max number of retries for query plan
     use_functions_api = True
 
-    system_message = f"""
+    system_message = """
     You will receive a QUERY, to be answered based on an EXTREMELY LARGE collection
     of documents you DO NOT have access to, but your ASSISTANT does.
     You only know that these documents have a special `content` field
     and additional FILTERABLE fields in the SCHEMA below, along with the 
     SAMPLE VALUES for each field, and the DTYPE in PANDAS TERMINOLOGY.
     
-    {{doc_schema}}
+    {doc_schema}
     
     Based on the QUERY and the above SCHEMA, your task is to determine a QUERY PLAN,
     consisting of:
@@ -116,10 +118,7 @@ class LanceQueryPlanAgentConfig(ChatAgentConfig):
     You may receive FEEDBACK on your QUERY PLAN and received ANSWER,
     from the 'QueryPlanCritic' who may offer suggestions for
     a better FILTER, REPHRASED QUERY, or DATAFRAME CALCULATION.
-            
-    If you keep getting feedback or keep getting a {NO_ANSWER} from the assistant
-    at least 3 times, then simply say '{DONE} {NO_ANSWER}' and nothing else.
-      
+                  
     At the BEGINNING if there is no query, ASK the user what they want to know.
     """
 
@@ -133,88 +132,128 @@ class LanceQueryPlanAgent(ChatAgent):
     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)
+        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)
+
+    def init_state(self) -> None:
+        self.curr_query_plan: QueryPlan | None = None
+        self.expecting_query_plan: bool = False
+        # how many times re-trying query plan in response to feedback:
+        self.n_retries: int = 0
+        self.n_query_plan_reminders: int = 0
+        self.result: str = ""  # answer received from LanceRAG
+
+    def llm_response(
+        self, message: Optional[str | ChatDocument] = None
+    ) -> Optional[ChatDocument]:
+        self.expecting_query_plan = True
+        return super().llm_response(message)
 
-    def query_plan(self, msg: QueryPlanTool) -> str:
-        """Valid, forward to RAG Agent"""
+    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
-        return PASS_TO + self.config.doc_agent_name
+        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
 
-    def query_plan_feedback(self, msg: QueryPlanFeedbackTool) -> str:
+        # OR
+        # (c) use the ForwardTool:
+        return ForwardTool(agent=self.config.doc_agent_name)
+
+    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 == "":
-            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
+        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 DONE + " " + NO_ANSWER
+            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:
-        FEEDBACK: {msg.feedback}
-        SUGGESTED FIX: {msg.suggested_fix}
+        ANSWER: {self.result}
+        FEEDBACK: {feedback}
+        SUGGESTED FIX: {suggested}
         """
 
+    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
+
     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.
+        Remind to use QueryPlanTool if we are expecting it.
         """
-        # 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
+        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.
-            Re-try your response using the `query_plan` tool/function.
+            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

langroid/agent/special/lance_tools.py

@@ -34,9 +34,17 @@ class QueryPlanTool(ToolMessage):
     plan: QueryPlan
 
 
+class AnswerTool(ToolMessage):
+    """Wrapper for answer from LanceDocChatAgent"""
+
+    purpose: str = "To package the answer from LanceDocChatAgent"
+    request: str = "answer_tool"
+    answer: str
+
+
 class QueryPlanAnswerTool(ToolMessage):
-    request = "query_plan_answer"  # the agent method name that handles this tool
-    purpose = """
+    request: str = "query_plan_answer"  # the agent method name that handles this tool
+    purpose: str = """
     Assemble query <plan> and <answer>
     """
     plan: QueryPlan