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

langroid/agent/special/lance_rag_new/lance_rag_task.py

@@ -1,144 +0,0 @@
-"""
-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

langroid/agent/special/lance_rag_new/query_planner_agent.py

@@ -1,222 +0,0 @@
-"""
-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
-
-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 (
-    QueryPlan,
-    QueryPlanAnswerTool,
-    QueryPlanFeedbackTool,
-    QueryPlanTool,
-)
-from langroid.utils.constants import DONE, NO_ANSWER
-
-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
-    max_retries: int = 5  # max number of retries for query plan
-    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, 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 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'"
-        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()",
-        or "df.groupby("director").mean()["rating"]", 
-        or EMPTY string if no calc is needed. 
-        The dataframe calc CAN refer to the `content` field.
-        If a DataFrame calculation is NOT needed, leave this field EMPTY.
-        
-        IMPORTANT: The DataFrame `df` in this calculation is the result of 
-        applying the FILTER AND REPHRASED QUERY to the documents.
-        
-        WATCH OUT!! When deciding the dataframe calc, if any, CAREFULLY
-        note what the query is asking, and ensure that the result of your
-        dataframe calc expression would answer the query.                
-    
-    
-    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()"
-        NOTE!!! The DataFrame `df` in this calculation is the result of
-        applying the FILTER AND REPHRASED QUERY to the documents, 
-        hence this computation will give the total deaths in shoplifting crimes.
-    ------------- END OF EXAMPLE ----------------
-    
-    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.
-    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
-        # 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)
-
-    def query_plan(self, msg: QueryPlanTool) -> ChatDocument:
-        """Valid, forward to RAG Agent"""
-        # save, to be used to assemble QueryPlanResultTool
-        if len(msg.plan.dataframe_calc.split("\n")) > 1:
-            return "DATAFRAME CALCULATION must be a SINGLE LINE; Retry the `query_plan`"
-        self.curr_query_plan = msg.plan
-        # return a ChatDocument with tool_messages set to this tool,
-        # so caller can directly get the tool without parsing
-        return self.create_agent_response(tool_messages=[msg])
-
-    def query_plan_feedback(self, msg: QueryPlanFeedbackTool) -> str:
-        """Process Critic feedback on QueryPlan + Answer from RAG Agent"""
-        # We should have saved answer in self.result by this time,
-        # since this Agent seeks feedback only after receiving RAG answer.
-        if msg.suggested_fix == "":
-            self.n_retries = 0
-            # This means the Query Plan or Result is good, as judged by Critic
-            if self.result == "":
-                # This was feedback for query with no result
-                return "QUERY PLAN LOOKS GOOD!"
-            elif self.result == NO_ANSWER:
-                return NO_ANSWER
-            else:  # non-empty and non-null answer
-                return DONE + " " + self.result
-        self.n_retries += 1
-        if self.n_retries >= self.config.max_retries:
-            # bail out to avoid infinite loop
-            self.n_retries = 0
-            return DONE + " " + NO_ANSWER
-        return f"""
-        here is FEEDBACK about your QUERY PLAN, and a SUGGESTED FIX.
-        Modify the QUERY PLAN if needed:
-        FEEDBACK: {msg.feedback}
-        SUGGESTED FIX: {msg.suggested_fix}
-        """
-
-    def handle_message_fallback(
-        self, msg: str | ChatDocument
-    ) -> str | ChatDocument | None:
-        """
-        Process answer received from RAG Agent:
-         Construct a QueryPlanAnswerTool with the answer,
-         and forward to Critic for feedback.
-        """
-        # TODO we don't need to use this fallback method. instead we can
-        # first call result = super().agent_response(), and if result is None,
-        # then we know there was no tool, so we run below code
-        if (
-            isinstance(msg, ChatDocument)
-            and self.curr_query_plan is not None
-            and msg.metadata.parent is not None
-        ):
-            # save result, to be used in query_plan_feedback()
-            self.result = msg.content
-            # assemble QueryPlanAnswerTool...
-            query_plan_answer_tool = QueryPlanAnswerTool(  # type: ignore
-                plan=self.curr_query_plan,
-                answer=self.result,
-            )
-            response_tmpl = self.create_agent_response()
-            # ... add the QueryPlanAnswerTool to the response
-            # (Notice how the Agent is directly sending a tool, not the LLM)
-            response_tmpl.tool_messages = [query_plan_answer_tool]
-            # set the recipient to the Critic so it can give feedback
-            response_tmpl.metadata.recipient = self.config.critic_name
-            self.curr_query_plan = None  # reset
-            return response_tmpl
-        if (
-            isinstance(msg, ChatDocument)
-            and not self.has_tool_message_attempt(msg)
-            and msg.metadata.sender == lr.Entity.LLM
-        ):
-            # remind LLM to use the QueryPlanFeedbackTool
-            return """
-            You forgot to use the `query_plan` tool/function.
-            Re-try your response using the `query_plan` tool/function.
-            """
-        return None