langroid 0.3.1__py3-none-any.whl → 0.5.1__py3-none-any.whl

langroid/agent/base.py

@@ -784,15 +784,51 @@ class Agent(ABC):
         #     ]
         # }
 
+        if not isinstance(json_data, dict):
+            return None
+
         properties = json_data.get("properties")
-        if properties is not None:
+        if isinstance(properties, dict):
             json_data = properties
         request = json_data.get("request")
-        if (
-            request is None
-            or not (isinstance(request, str))
-            or request not in self.llm_tools_handled
-        ):
+
+        if request is None:
+            handled = [self.llm_tools_map[r] for r in self.llm_tools_handled]
+            default_keys = set(ToolMessage.__fields__.keys())
+            request_keys = set(json_data.keys())
+
+            def maybe_parse(tool: type[ToolMessage]) -> Optional[ToolMessage]:
+                all_keys = set(tool.__fields__.keys())
+                non_inherited_keys = all_keys.difference(default_keys)
+                # If the request has any keys not valid for the tool and
+                # does not specify some key specific to the type
+                # (e.g. not just `purpose`), the LLM must explicitly specify `request`
+                if not (
+                    request_keys.issubset(all_keys)
+                    and len(request_keys.intersection(non_inherited_keys)) > 0
+                ):
+                    return None
+
+                try:
+                    return tool.parse_obj(json_data)
+                except ValidationError:
+                    return None
+
+            candidate_tools = list(
+                filter(
+                    lambda t: t is not None,
+                    map(maybe_parse, handled),
+                )
+            )
+
+            # If only one valid candidate exists, we infer
+            # "request" to be the only possible value
+            if len(candidate_tools) == 1:
+                return candidate_tools[0]
+            else:
+                return None
+
+        if not isinstance(request, str) or request not in self.llm_tools_handled:
             return None
 
         message_class = self.llm_tools_map.get(request)

langroid/agent/chat_agent.py

@@ -427,11 +427,11 @@ class ChatAgent(Agent):
                 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 require_recipient and message_class is not None:
+            message_class = message_class.require_recipient()
         super().enable_message_handling(message_class)  # enables handling only
         tools = self._get_tool_list(message_class)
         if message_class is not None:
-            if require_recipient:
-                message_class = message_class.require_recipient()
             request = message_class.default_value("request")
             llm_function = message_class.llm_function_schema(defaults=include_defaults)
             self.llm_functions_map[request] = llm_function

langroid/agent/special/doc_chat_agent.py

@@ -538,12 +538,13 @@ class DocChatAgent(ChatAgent):
         ]
 
     def get_field_values(self, fields: list[str]) -> Dict[str, str]:
-        """Get string-listing of possible values of each filterable field,
+        """Get string-listing of possible values of each field,
         e.g.
         {
             "genre": "crime, drama, mystery, ... (10 more)",
             "certificate": "R, PG-13, PG, R",
         }
+        The field names may have "metadata." prefix, e.g. "metadata.genre".
         """
         field_values: Dict[str, Set[str]] = {}
         # make empty set for each field
@@ -556,8 +557,11 @@ class DocChatAgent(ChatAgent):
         for d in docs:
             # extract fields from d
             doc_field_vals = extract_fields(d, fields)
-            for field, val in doc_field_vals.items():
-                field_values[field].add(val)
+            # the `field` returned by extract_fields may contain only the last
+            # part of the field name, e.g. "genre" instead of "metadata.genre",
+            # so we use the orig_field name to fill in the values
+            for (field, val), orig_field in zip(doc_field_vals.items(), fields):
+                field_values[orig_field].add(val)
         # For each field make a string showing list of possible values,
         # truncate to 20 values, and if there are more, indicate how many
         # more there are, e.g. Genre: crime, drama, mystery, ... (20 more)
@@ -680,7 +684,13 @@ class DocChatAgent(ChatAgent):
                 )
             return response
         if query_str == "":
-            return None
+            return ChatDocument(
+                content=NO_ANSWER + " since query was empty",
+                metadata=ChatDocMetaData(
+                    source="No query provided",
+                    sender=Entity.LLM,
+                ),
+            )
         elif query_str == "?" and self.response is not None:
             return self.justify_response()
         elif (query_str.startswith(("summar", "?")) and self.response is None) or (

langroid/agent/special/lance_doc_chat_agent.py

@@ -22,7 +22,6 @@ 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.pydantic_utils import (
-    clean_schema,
     dataframe_to_documents,
 )
 from langroid.vector_store.lancedb import LanceDB
@@ -41,24 +40,26 @@ class LanceDocChatAgent(DocChatAgent):
     def _get_clean_vecdb_schema(self) -> str:
         """Get a cleaned schema of the vector-db, to pass to the LLM
         as part of instructions on how to generate a SQL filter."""
+
+        tbl_pandas = (
+            self.vecdb.client.open_table(self.vecdb.config.collection_name)
+            .search()
+            .limit(1)
+            .to_pandas(flatten=True)
+        )
         if len(self.config.filter_fields) == 0:
-            filterable_fields = (
-                self.vecdb.client.open_table(self.vecdb.config.collection_name)
-                .search()
-                .limit(1)
-                .to_pandas(flatten=True)
-                .columns.tolist()
-            )
+            filterable_fields = tbl_pandas.columns.tolist()
             # drop id, vector, metadata.id, metadata.window_ids, metadata.is_chunk
-            for fields in [
-                "id",
-                "vector",
-                "metadata.id",
-                "metadata.window_ids",
-                "metadata.is_chunk",
-            ]:
-                if fields in filterable_fields:
-                    filterable_fields.remove(fields)
+            filterable_fields = list(
+                set(filterable_fields)
+                - {
+                    "id",
+                    "vector",
+                    "metadata.id",
+                    "metadata.window_ids",
+                    "metadata.is_chunk",
+                }
+            )
             logger.warning(
                 f"""
             No filter_fields set in config, so using these fields as filterable fields:
@@ -69,15 +70,7 @@ class LanceDocChatAgent(DocChatAgent):
 
         if self.from_dataframe:
             return self.df_description
-        schema_dict = clean_schema(
-            self.vecdb.schema,
-            excludes=["id", "vector"],
-        )
-        # intersect config.filter_fields with schema_dict.keys() in case
-        # there are extraneous fields in config.filter_fields
-        filter_fields_set = set(
-            self.config.filter_fields or schema_dict.keys()
-        ).intersection(schema_dict.keys())
+        filter_fields_set = set(self.config.filter_fields)
 
         # remove 'content' from filter_fields_set, even if it's not in filter_fields_set
         filter_fields_set.discard("content")
@@ -85,10 +78,14 @@ class LanceDocChatAgent(DocChatAgent):
         # possible values of filterable fields
         filter_field_values = self.get_field_values(list(filter_fields_set))
 
+        schema_dict: Dict[str, Dict[str, Any]] = dict(
+            (field, {}) for field in filter_fields_set
+        )
         # add field values to schema_dict as another field `values` for each field
         for field, values in filter_field_values.items():
-            if field in schema_dict:
-                schema_dict[field]["values"] = values
+            schema_dict[field]["values"] = values
+            dtype = tbl_pandas[field].dtype.name
+            schema_dict[field]["dtype"] = dtype
         # if self.config.filter_fields is set, restrict to these:
         if len(self.config.filter_fields) > 0:
             schema_dict = {

langroid/agent/special/lance_rag/critic_agent.py

@@ -37,20 +37,30 @@ class QueryPlanCriticConfig(LanceQueryPlanAgentConfig):
     system_message = f"""
     You are an expert at carefully planning a query that needs to be answered
     based on a large collection of documents. These docs have a special `content` field
-    and additional FILTERABLE fields in the SCHEMA below:
+    and additional FILTERABLE fields in the SCHEMA below, along with the 
+    SAMPLE VALUES for each field, and the DTYPE in PANDAS TERMINOLOGY.
     
     {{doc_schema}}
     
+    The ORIGINAL QUERY is handled by a QUERY PLANNER who sends the PLAN to an ASSISTANT,
+    who returns an ANSWER.
+    
     You will receive a QUERY PLAN consisting of:
-    - ORIGINAL QUERY, 
-    - SQL-Like FILTER, WHICH CAN BE EMPTY (and it's fine if results sound reasonable)
+    - ORIGINAL QUERY from the user, which a QUERY PLANNER processes,
+      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.
-    - REPHRASED QUERY that will be used to match against the CONTENT (not filterable)
-         of the documents.
+    - 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 
       of the docs. Thus the REPHRASED QUERY itself acts like a 
       SEMANTIC/LEXICAL/FUZZY FILTER since the Assistant is able to use it to match 
-      the CONTENT of the docs in various ways (semantic, lexical, fuzzy, etc.).         
+      the CONTENT of the docs in various ways (semantic, lexical, fuzzy, etc.). 
+        Keep in mind that the ASSISTANT does NOT know anything about the FILTER fields,
+        so the REPHRASED QUERY should NOT mention ANY FILTER fields.
+        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.

langroid/agent/special/lance_rag/query_planner_agent.py

@@ -43,23 +43,27 @@ class LanceQueryPlanAgentConfig(ChatAgentConfig):
     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:  
+    and additional FILTERABLE fields in the SCHEMA below, along with the 
+    SAMPLE VALUES for each field, and the DTYPE in PANDAS TERMINOLOGY.
     
     {{doc_schema}}
     
     Based on the QUERY and the above SCHEMA, your task is to determine a QUERY PLAN,
     consisting of:
-    -  a FILTER (can be empty string) that would help the ASSISTANT to answer the query.
+    -  a PANDAS-TYPE FILTER (can be empty string) that would help the ASSISTANT to 
+        answer the query.
         Remember the FILTER can refer to ANY fields in the above SCHEMA
         EXCEPT the `content` field of the documents. 
         ONLY USE A FILTER IF EXPLICITLY MENTIONED IN THE QUERY.
         TO get good results, for STRING MATCHES, consider using LIKE instead of =, e.g.
         "CEO LIKE '%Jobs%'" instead of "CEO = 'Steve Jobs'"
-    - a possibly REPHRASED QUERY to be answerable given the FILTER.
+        YOUR FILTER MUST BE A PANDAS-TYPE FILTER, respecting the shown DTYPES.        
+    - a possibly REPHRASED QUERY (CANNOT BE EMPTY) to be answerable given the FILTER.
         Keep in mind that the ASSISTANT does NOT know anything about the FILTER fields,
         so the REPHRASED QUERY should NOT mention ANY FILTER fields.
         The assistant will answer based on documents whose CONTENTS match the QUERY, 
         possibly REPHRASED. 
+        !!!!****THE REPHRASED QUERY SHOULD NEVER BE EMPTY****!!!
     - an OPTIONAL SINGLE-LINE Pandas-dataframe calculation/aggregation string 
         that can be used to calculate the answer to the original query, 
         e.g. "df["rating"].mean()",
@@ -99,7 +103,7 @@ class LanceQueryPlanAgentConfig(ChatAgentConfig):
         hence this computation will give the total deaths in shoplifting crimes.
     ------------- END OF EXAMPLE ----------------
     
-    The FILTER must be a SQL-like condition, e.g. 
+    The FILTER must be a PANDAS-like condition, e.g. 
     "year > 2000 AND genre = 'ScienceFiction'".
     To ensure you get useful results, you should make your FILTER 
     NOT TOO STRICT, e.g. look for approximate match using LIKE, etc.

langroid/agent/special/lance_rag_new/__init__.py

@@ -0,0 +1,9 @@
+from . import query_planner_agent
+from . import critic_agent
+from . import lance_rag_task
+
+__all__ = [
+    "query_planner_agent",
+    "critic_agent",
+    "lance_rag_task",
+]

langroid/agent/special/lance_rag_new/critic_agent.py

@@ -0,0 +1,171 @@
+"""
+QueryPlanCritic is a ChatAgent that is created with a specific document schema.
+
+Its role is to provide feedback on a Query Plan, which consists 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
+- result - the answer received from an assistant that used this QUERY PLAN.
+
+This agent has access to two tools:
+- QueryPlanTool: The handler method for this tool re-writes the query plan
+  in plain text (non-JSON) so the LLM can provide its feedback using the
+  QueryPlanFeedbackTool.
+- QueryPlanFeedbackTool: LLM uses this tool to provide feedback on the Query Plan
+"""
+
+import logging
+
+from langroid.agent.chat_agent import ChatAgent
+from langroid.agent.chat_document import ChatDocument
+from langroid.agent.special.lance_rag.query_planner_agent import (
+    LanceQueryPlanAgentConfig,
+)
+from langroid.agent.special.lance_tools import (
+    QueryPlanAnswerTool,
+    QueryPlanFeedbackTool,
+)
+from langroid.mytypes import Entity
+from langroid.utils.constants import DONE, NO_ANSWER
+
+logger = logging.getLogger(__name__)
+
+
+class QueryPlanCriticConfig(LanceQueryPlanAgentConfig):
+    name = "QueryPlanCritic"
+    system_message = f"""
+    You are an expert at carefully planning a query that needs to be answered
+    based on a large collection of documents. These docs 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}}
+    
+    The ORIGINAL QUERY is handled by a QUERY PLANNER who sends the PLAN to an ASSISTANT,
+    who returns an ANSWER.
+    
+    You will receive a QUERY PLAN consisting of:
+    - ORIGINAL QUERY from the user, which a QUERY PLANNER processes,
+      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.
+    - 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 
+      of the docs. Thus the REPHRASED QUERY itself acts like a 
+      SEMANTIC/LEXICAL/FUZZY FILTER since the Assistant is able to use it to match 
+      the CONTENT of the docs in various ways (semantic, lexical, fuzzy, etc.). 
+        Keep in mind that the ASSISTANT does NOT know anything about the FILTER fields,
+        so the REPHRASED QUERY should NOT mention ANY FILTER fields.
+        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.
+      NOTE -- the ANSWER will usually NOT contain any references to FILTERING conditions,
+      and this is ALLOWED, since the ANSWER is based on documents AFTER FILTERING.
+
+    In addition to the above SCHEMA fields there is a `content` field which:
+    - CANNOT appear in a FILTER, 
+    - CAN appear in the DATAFRAME CALCULATION.
+    THERE ARE NO OTHER FIELDS IN THE DOCUMENTS or in the RESULTING DATAFRAME.
+        
+    Your job is to act as a CRITIC and provide feedback, 
+    ONLY using the `query_plan_feedback` tool, and DO NOT SAY ANYTHING ELSE.
+    
+    Here is how you must examine the QUERY PLAN + ANSWER:
+    - ALL filtering conditions in the original query must be EXPLICITLY 
+      mentioned in the FILTER, and the QUERY field should not be used for filtering.
+    - If the ANSWER contains an ERROR message, then this means that the query
+      plan execution FAILED, and your feedback should say INVALID along 
+      with the ERROR message, `suggested_fix` that aims to help the assistant 
+      fix the problem (or simply equals "address the the error shown in feedback")
+    - Ask yourself, is the ANSWER in the expected form, e.g. 
+        if the question is asking for the name of an ENTITY with max SIZE,
+        then the answer should be the ENTITY name, NOT the SIZE!! 
+    - It is perfectly FINE if the ANSWER does NOT contain any references to FILTERING 
+      conditions, since the ANSWER is obtained from documents AFTER FILTERING!
+    - If the ANSWER is in the expected form, then the QUERY PLAN is likely VALID,
+      and your feedback should say VALID, with empty `suggested_fix`.
+      ===> HOWEVER!!! Watch out for a spurious correct-looking answer, for EXAMPLE:
+      the query was to find the ENTITY with a maximum SIZE, 
+      but the dataframe calculation is find the SIZE, NOT the ENTITY!!      
+    - If the ANSWER is {NO_ANSWER} or of the wrong form, 
+      then try to DIAGNOSE the problem IN THE FOLLOWING ORDER:
+      - DATAFRAME CALCULATION -- is it doing the right thing?
+        Is it finding the Index of a row instead of the value in a column?
+        Or another example: maybe it is finding the maximum population
+           rather than the CITY with the maximum population?
+        If you notice a problem with the DATAFRAME CALCULATION, then
+        ONLY SUBMIT FEEDBACK ON THE DATAFRAME CALCULATION, and DO NOT
+        SUGGEST ANYTHING ELSE.
+      - If the DATAFRAME CALCULATION looks correct, then check if 
+        the REPHRASED QUERY makes sense given the ORIGINAL QUERY and FILTER.
+        If this is the problem, then ONLY SUBMIT FEEDBACK ON THE REPHRASED QUERY,
+        and DO NOT SUGGEST ANYTHING ELSE.
+      - If the REPHRASED QUERY looks correct, then check if the FILTER makes sense.
+        REMEMBER: A filter should ONLY be used if EXPLICITLY REQUIRED BY THE QUERY.
+     
+     
+     IMPORTANT!! The DATAFRAME CALCULATION is done AFTER applying the 
+         FILTER and REPHRASED QUERY! Keep this in mind when evaluating 
+         the correctness of the DATAFRAME CALCULATION.
+    
+    ALWAYS use `query_plan_feedback` tool/fn to present your feedback
+    in the `feedback` field, and if any fix is suggested,
+    present it in the `suggested_fix` field.
+    DO NOT SAY ANYTHING ELSE OUTSIDE THE TOOL/FN.
+    IF NO REVISION NEEDED, simply leave the `suggested_fix` field EMPTY,
+    and SAY NOTHING ELSE
+    and DO NOT EXPLAIN YOURSELF.        
+    """
+
+
+def plain_text_query_plan(msg: QueryPlanAnswerTool) -> str:
+    plan = f"""
+    OriginalQuery: {msg.plan.original_query}
+    Filter: {msg.plan.filter}
+    Rephrased Query: {msg.plan.query}
+    DataframeCalc: {msg.plan.dataframe_calc}
+    Answer: {msg.answer}
+    """
+    return plan
+
+
+class QueryPlanCritic(ChatAgent):
+    """
+    Critic for LanceQueryPlanAgent, provides feedback on
+    query plan + answer.
+    """
+
+    def __init__(self, cfg: LanceQueryPlanAgentConfig):
+        super().__init__(cfg)
+        self.config = cfg
+        self.enable_message(QueryPlanAnswerTool, use=False, handle=True)
+        self.enable_message(QueryPlanFeedbackTool, use=True, handle=True)
+
+    def query_plan_answer(self, msg: QueryPlanAnswerTool) -> str:
+        """Present query plan + answer in plain text (not JSON)
+        so LLM can give feedback"""
+        return plain_text_query_plan(msg)
+
+    def query_plan_feedback(self, msg: QueryPlanFeedbackTool) -> ChatDocument:
+        """Format Valid so return to Query Planner"""
+        doc = self.create_agent_response(DONE)
+        doc.tool_messages = [msg]
+        return doc
+
+    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:
+            return """
+            You forgot to use the `query_plan_feedback` tool/function.
+            Re-try your response using the `query_plan_feedback` tool/function,
+            remember to provide feedback in the `feedback` field,
+            and if any fix is suggested, provide it in the `suggested_fix` field.
+            """
+        return None

langroid/agent/special/lance_rag_new/lance_rag_task.py

@@ -0,0 +1,144 @@
+"""
+The LanceRAGTaskCreator.new() method creates a 3-Agent system that uses this agent.
+It takes a LanceDocChatAgent instance as argument, and adds two more agents:
+- LanceQueryPlanAgent, which is given the LanceDB schema in LanceDocChatAgent,
+and based on this schema, for a given user query, creates a Query Plan
+using the QueryPlanTool, which contains a filter, a rephrased query,
+and a dataframe_calc.
+- QueryPlanCritic, which is given the LanceDB schema in LanceDocChatAgent,
+ and gives feedback on the Query Plan and Result using the QueryPlanFeedbackTool.
+
+The LanceRAGTaskCreator.new() method sets up the given LanceDocChatAgent and
+QueryPlanCritic as sub-tasks of the LanceQueryPlanAgent's task.
+
+Langroid's built-in task orchestration ensures that:
+- the LanceQueryPlanAgent reformulates the plan based
+    on the QueryPlanCritics's feedback,
+- LLM deviations are corrected via tools and overrides of ChatAgent methods.
+"""
+
+import logging
+
+from langroid.agent.special.lance_tools import (
+    QueryPlanAnswerTool,
+)
+from langroid.agent.task import Task
+from langroid.mytypes import Entity
+from langroid.utils.constants import NO_ANSWER
+
+from ..lance_doc_chat_agent import LanceDocChatAgent
+from .critic_agent import (
+    QueryPlanCritic,
+    QueryPlanCriticConfig,
+)
+from .query_planner_agent import (
+    LanceQueryPlanAgent,
+    LanceQueryPlanAgentConfig,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def run_lance_rag_task(
+    query: str,
+    agent: LanceDocChatAgent,
+    interactive: bool = True,
+) -> str:
+    """
+    Add a LanceFilterAgent to the LanceDocChatAgent,
+    set up the corresponding Tasks, connect them,
+    and return the top-level query_plan_task.
+    """
+    doc_agent_name = "LanceRAG"
+    critic_name = "QueryPlanCritic"
+    query_plan_agent_config = LanceQueryPlanAgentConfig(
+        critic_name=critic_name,
+        doc_agent_name=doc_agent_name,
+        doc_schema=agent._get_clean_vecdb_schema(),
+    )
+    query_plan_agent_config.set_system_message()
+
+    query_planner = LanceQueryPlanAgent(query_plan_agent_config)
+    query_plan_task = Task(
+        query_planner,
+        interactive=interactive,
+        restart=False,
+        done_if_response=[Entity.AGENT],
+    )
+    # TODO - figure out how to define the fns so we avoid re-creating
+    # agents in each invocation. Right now we are defining the fn
+    # inside this context, which may not be great.
+
+    rag_task = Task(
+        agent,
+        name="LanceRAG",
+        restart=True,  # default; no need to accumulate dialog
+        interactive=False,
+        done_if_response=[Entity.LLM],  # done when non-null response from LLM
+        done_if_no_response=[Entity.LLM],  # done when null response from LLM
+    )
+
+    critic_config = QueryPlanCriticConfig(
+        doc_schema=agent._get_clean_vecdb_schema(),
+    )
+    critic_config.set_system_message()
+
+    critic_agent = QueryPlanCritic(critic_config)
+    critic_task = Task(
+        critic_agent,
+        interactive=False,
+        restart=True,  # default; no need to accumulate dialog
+    )
+
+    no_answer = False
+    feedback = None
+    i = 0
+    while i := i + 1 < 5:
+        # query, feedback (QueryPlanFeedbackTool) => ChatDocument[QueryPlanTool]
+        if feedback is not None and feedback.suggested_fix != "":
+            prompt = f"""
+                A Critic has seen your Query Plan and the Answer, and has given the 
+                following feedback. Take it into account and re-generate your Query Plan
+                for the QUERY:
+                
+                QUERY: {query}
+                FEEDBACK: {feedback.feedback}
+                SUGGESTED FIX: {feedback.suggested_fix}
+                """
+        elif no_answer:
+            prompt = f"There was a {NO_ANSWER} response; try a different query plan"
+        else:
+            prompt = query
+
+        while True:
+            plan_doc = query_plan_task.run(prompt)
+            if len(plan_doc.tool_messages) > 0:
+                break
+            # forgot to use QueryPlanTool
+            prompt = """You forgot to use the `query_plan` tool/function. Try again."""
+
+        # TODO if plan_doc does NOT have a QueryPlan, remind the agent
+
+        # ChatDocument with QueryPlanTool => ChatDocument with answer
+        rag_answer_doc = rag_task.run(plan_doc)
+
+        if rag_answer_doc is None:
+            rag_answer_doc = rag_task.agent.create_llm_response(NO_ANSWER)
+        # QueryPlan, answer => QueryPlanAnswerTool
+        plan_answer_tool = QueryPlanAnswerTool(
+            plan=plan_doc.tool_messages[0].plan,
+            answer=rag_answer_doc.content,
+        )
+        # QueryPlanAnswerTool => ChatDocument[QueryPlanAnswerTool]
+        plan_answer_doc = agent.create_agent_response(tool_messages=[plan_answer_tool])
+
+        # ChatDocument[QueryPlanAnswerTool] => ChatDocument[QueryPlanFeedbackTool]
+        feedback_doc = critic_task.run(plan_answer_doc)
+        # ChatDocument[QueryPlanFeedbackTool] => QueryPlanFeedbackTool
+        feedback = feedback_doc.tool_messages[0]  # QueryPlanFeedbackTool
+        no_answer = NO_ANSWER in rag_answer_doc.content
+        if feedback.suggested_fix == "" and not no_answer:
+            break
+
+    # query_plan_task.add_sub_task([critic_task, rag_task])
+    return rag_answer_doc.content