langroid 0.1.85__py3-none-any.whl → 0.1.219__py3-none-any.whl

langroid/agent/special/lance_doc_chat_agent.py

@@ -0,0 +1,258 @@
+"""
+LanceDocChatAgent is a subclass of DocChatAgent that uses LanceDB as a vector store:
+- Uses the DocChatAgentConfig.filter variable
+    (a sql string) in the `where` clause to do filtered vector search.
+- Overrides the get_similar_chunks_bm25() to use LanceDB FTS (Full Text Search).
+
+For usage see:
+ - `tests/main/test_lance_doc_chat_agent.py`.
+ - example script `examples/docqa/lance_rag.py`.
+
+"""
+
+import json
+import logging
+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.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
+
+logger = logging.getLogger(__name__)
+
+
+class LanceDocChatAgent(DocChatAgent):
+    vecdb: LanceDB
+
+    def __init__(self, cfg: DocChatAgentConfig):
+        super().__init__(cfg)
+        self.config: DocChatAgentConfig = cfg
+        self.enable_message(QueryPlanTool, use=False, handle=True)
+
+    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."""
+        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()
+            )
+            # 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)
+            logger.warning(
+                f"""
+            No filter_fields set in config, so using these fields as filterable fields:
+            {filterable_fields}
+            """
+            )
+            self.config.filter_fields = filterable_fields
+
+        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())
+
+        # remove 'content' from filter_fields_set, even if it's not in filter_fields_set
+        filter_fields_set.discard("content")
+
+        # possible values of filterable fields
+        filter_field_values = self.get_field_values(list(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
+        # if self.config.filter_fields is set, restrict to these:
+        if len(self.config.filter_fields) > 0:
+            schema_dict = {
+                k: v for k, v in schema_dict.items() if k in self.config.filter_fields
+            }
+        schema = json.dumps(schema_dict, indent=4)
+
+        schema += f"""
+        NOTE when creating a filter for a query, 
+        ONLY the following fields are allowed:
+        {",".join(self.config.filter_fields)} 
+        """
+        if len(content_fields := self.config.add_fields_to_content) > 0:
+            schema += f"""
+            Additional fields added to `content` as key=value pairs:
+            NOTE that these CAN Help with matching queries!
+            {content_fields}
+            """
+        return schema
+
+    def query_plan(self, msg: QueryPlanTool) -> str:
+        """
+        Handle the LLM's use of the FilterTool.
+        Temporarily set the config filter and either return the final answer
+        in case there's a dataframe_calc, or return the rephrased query
+        so the LLM can handle it.
+        """
+        # create document-subset based on this filter
+        plan = msg.plan
+        try:
+            self.setup_documents(filter=plan.filter or None)
+        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)}
+            """
+
+        # update the filter so it is used in the DocChatAgent
+        self.config.filter = plan.filter or None
+        if plan.dataframe_calc:
+            # we just get relevant docs then do the calculation
+            # TODO if calc causes err, it is captured in result,
+            # and LLM can correct the calc based on the err,
+            # and this will cause retrieval all over again,
+            # which may be wasteful if only the calc part is wrong.
+            # 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.
+                    """
+                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
+        else:
+            # pass on the query so LLM can handle it
+            return plan.query
+
+    def ingest_docs(
+        self,
+        docs: List[Document],
+        split: bool = True,
+        metadata: (
+            List[Dict[str, Any]] | Dict[str, Any] | DocMetaData | List[DocMetaData]
+        ) = [],
+    ) -> int:
+        n = super().ingest_docs(docs, split, metadata)
+        tbl = self.vecdb.client.open_table(self.vecdb.config.collection_name)
+        # We assume "content" is available as top-level field
+        if "content" in tbl.schema.names:
+            tbl.create_fts_index("content", replace=True)
+        return n
+
+    def ingest_dataframe(
+        self,
+        df: pd.DataFrame,
+        content: str = "content",
+        metadata: List[str] = [],
+    ) -> int:
+        """Ingest from a dataframe. Assume we are doing this once, not incrementally"""
+
+        self.from_dataframe = True
+        if df.shape[0] == 0:
+            raise ValueError(
+                """
+                LanceDocChatAgent.ingest_dataframe() received an empty dataframe.
+                """
+            )
+        n = df.shape[0]
+
+        # If any additional fields need to be added to content,
+        # add them as key=value pairs, into the `content` field for all rows.
+        # This helps retrieval for table-like data.
+        # Note we need to do this at stage so that the embeddings
+        # are computed on the full content with these additional fields.
+        fields = [f for f in self.config.add_fields_to_content if f in df.columns]
+        if len(fields) > 0:
+            df[content] = df.apply(
+                lambda row: (",".join(f"{f}={row[f]}" for f in fields))
+                + ", content="
+                + row[content],
+                axis=1,
+            )
+
+        df, metadata = DocChatAgent.document_compatible_dataframe(df, content, metadata)
+        self.df_description = describe_dataframe(
+            df,
+            filter_fields=self.config.filter_fields,
+            n_vals=10,
+        )
+        self.vecdb.add_dataframe(df, content="content", metadata=metadata)
+
+        tbl = self.vecdb.client.open_table(self.vecdb.config.collection_name)
+        # We assume "content" is available as top-level field
+        if "content" in tbl.schema.names:
+            tbl.create_fts_index("content", replace=True)
+        # We still need to do the below so that
+        # other types of searches in DocChatAgent
+        # can work, as they require Document objects
+        docs = dataframe_to_documents(df, content="content", metadata=metadata)
+        self.setup_documents(docs)
+        # mark each doc as already-chunked so we don't try to split them further
+        # TODO later we may want to split large text-columns
+        for d in docs:
+            d.metadata.is_chunk = True
+        return n  # type: ignore
+
+    def get_similar_chunks_bm25(
+        self, query: str, multiple: int
+    ) -> List[Tuple[Document, float]]:
+        """
+        Override the DocChatAgent.get_similar_chunks_bm25()
+        to use LanceDB FTS (Full Text Search).
+        """
+        # Clean up query: replace all newlines with spaces in query,
+        # force special search keywords to lower case, remove quotes,
+        # so it's not interpreted as search syntax
+        query_clean = (
+            query.replace("\n", " ")
+            .replace("AND", "and")
+            .replace("OR", "or")
+            .replace("NOT", "not")
+            .replace("'", "")
+            .replace('"', "")
+        )
+
+        tbl = self.vecdb.client.open_table(self.vecdb.config.collection_name)
+        result = (
+            tbl.search(query_clean)
+            .where(self.config.filter or None)
+            .limit(self.config.parsing.n_similar_docs * multiple)
+        )
+        docs = self.vecdb._lance_result_to_docs(result)
+        scores = [r["score"] for r in result.to_list()]
+        return list(zip(docs, scores))

langroid/agent/special/lance_rag/__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/critic_agent.py

@@ -0,0 +1,136 @@
+"""
+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, PASS
+
+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:
+    
+    {{doc_schema}}
+    
+    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)
+      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.
+      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.).
+         
+    - DATAFRAME CALCULATION, and 
+    - ANSWER recieved from an assistant that used this QUERY PLAN.
+
+    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:
+    - If the ANSWER is in the expected form, then the QUERY PLAN is likely VALID,
+      and your feedback should be EMPTY.
+    - 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: mmaybe 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.
+     
+    
+    ALWAYS use `query_plan_feedback` tool/fn to present your feedback!
+    and DO NOT SAY ANYTHING ELSE OUTSIDE THE TOOL/FN.
+    IF NO REVISION NEEDED, simply give EMPTY FEEBACK, 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}
+    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) -> str:
+        """Format Valid so return to Query Planner"""
+        return DONE + " " + PASS  # return to Query Planner
+
+    def handle_message_fallback(
+        self, msg: str | ChatDocument
+    ) -> str | ChatDocument | None:
+        """Create QueryPlanFeedbackTool since LLM forgot"""
+        if isinstance(msg, ChatDocument) and msg.metadata.sender == Entity.LLM:
+            # our LLM forgot to use the QueryPlanFeedbackTool
+            feedback = QueryPlanFeedbackTool(feedback=msg.content)
+            msg.tool_messages = [feedback]
+            msg.content = DONE
+            return msg
+        return None

langroid/agent/special/lance_rag/lance_rag_task.py

@@ -0,0 +1,80 @@
+"""
+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_doc_chat_agent import LanceDocChatAgent
+from langroid.agent.special.lance_rag.critic_agent import (
+    QueryPlanCritic,
+    QueryPlanCriticConfig,
+)
+from langroid.agent.special.lance_rag.query_planner_agent import (
+    LanceQueryPlanAgent,
+    LanceQueryPlanAgentConfig,
+)
+from langroid.agent.task import Task
+from langroid.mytypes import Entity
+
+logger = logging.getLogger(__name__)
+
+
+class LanceRAGTaskCreator:
+    @staticmethod
+    def new(
+        agent: LanceDocChatAgent,
+        interactive: bool = True,
+    ) -> Task:
+        """
+        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()
+
+        critic_config = QueryPlanCriticConfig(
+            doc_schema=agent._get_clean_vecdb_schema(),
+        )
+        critic_config.set_system_message()
+
+        query_planner = LanceQueryPlanAgent(query_plan_agent_config)
+        query_plan_task = Task(
+            query_planner,
+            interactive=interactive,
+        )
+        critic_agent = QueryPlanCritic(critic_config)
+        critic_task = Task(
+            critic_agent,
+            interactive=False,
+        )
+        rag_task = Task(
+            agent,
+            name="LanceRAG",
+            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
+        )
+        query_plan_task.add_sub_task([critic_task, rag_task])
+        return query_plan_task

langroid/agent/special/lance_rag/query_planner_agent.py

@@ -0,0 +1,180 @@
+"""
+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
+"""
+
+import logging
+
+from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig
+from langroid.agent.chat_document import ChatDocument
+from langroid.agent.special.lance_tools import (
+    QueryPlan,
+    QueryPlanAnswerTool,
+    QueryPlanFeedbackTool,
+    QueryPlanTool,
+)
+from langroid.utils.constants import DONE, NO_ANSWER, PASS_TO
+
+logger = logging.getLogger(__name__)
+
+
+class LanceQueryPlanAgentConfig(ChatAgentConfig):
+    name: str = "LancePlanner"
+    critic_name: str = "QueryPlanCritic"
+    doc_agent_name: str = "LanceRAG"
+    doc_schema: str = ""
+    use_tools = False
+    use_functions_api = True
+
+    system_message = f"""
+    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:  
+    
+    {{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.
+        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.
+        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 answer will answer based on documents whose CONTENTS match the QUERY, 
+        possibly REPHRASED. 
+    - a Pandas-dataframe calculation/aggregation string that can be used to calculate
+        the answer to the original query, e.g. "df["rating"].mean()",
+        or "df.groupby("director").mean()["rating"]", etc, or empty string if no calc
+        is needed. The dataframe calc CAN refer to the `content` field. 
+    
+    
+    EXAMPLE:
+    ------- 
+    Suppose there is a document-set about crime reports, where:
+     CONTENT = crime report,
+     Filterable SCHEMA consists of City, Year, num_deaths.
+    
+    Then given this ORIGINAL QUERY: 
+    
+        Total deaths in shoplifting crimes in Los Angeles in 2023?
+    
+    A POSSIBLE QUERY PLAN could be:
+    
+    FILTER: "City LIKE '%Los Angeles%' AND Year = 2023"
+    REPHRASED QUERY: "shoplifting crime" --> this will be used to MATCH content of docs
+         [NOTE: we dropped the FILTER fields City and Year since the 
+         ASSISTANT does not know about them and only uses the query to 
+         match the CONTENT of the docs.]
+    DATAFRAME CALCULATION: "df["num_deaths"].sum()"
+
+    ------------- END OF EXAMPLE ----------------
+    
+    The FILTER must be a SQL-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.
+    E.g. "CEO LIKE '%Jobs%'" instead of "CEO = 'Steve Jobs'"
+    Use DOT NOTATION to refer to nested fields, e.g. `metadata.year`, etc. 
+        
+    You must FIRST present the QUERY PLAN using the `query_plan` tool/function.
+    This will be handled by your document assistant, who will produce an ANSWER.
+            
+    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.
+    """
+
+    def set_system_message(self) -> None:
+        self.system_message = self.system_message.format(
+            doc_schema=self.doc_schema,
+        )
+
+
+class LanceQueryPlanAgent(ChatAgent):
+    def __init__(self, config: LanceQueryPlanAgentConfig):
+        super().__init__(config)
+        self.config: LanceQueryPlanAgentConfig = config
+        self.curr_query_plan: QueryPlan | None = None
+        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)
+
+    def query_plan(self, msg: QueryPlanTool) -> str:
+        """Valid, forward to RAG Agent"""
+        # save, to be used to assemble QueryPlanResultTool
+        self.curr_query_plan = msg.plan
+        return PASS_TO + self.config.doc_agent_name
+
+    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.feedback == "":
+            # 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
+        return f"""
+        here is FEEDBACK about your QUERY PLAN. Modify it if needed:
+        {msg.feedback}
+        """
+
+    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(
+                plan=self.curr_query_plan,
+                answer=self.result,
+            )
+            response_tmpl = self.agent_response_template()
+            # ... 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
+        return None