langroid 0.31.1__py3-none-any.whl → 0.33.3__py3-none-any.whl

langroid/agent/special/lance_doc_chat_agent.py

@@ -1,262 +0,0 @@
-"""
-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 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 NO_ANSWER
-from langroid.utils.pydantic_utils import (
-    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."""
-
-        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 = tbl_pandas.columns.tolist()
-            # drop id, vector, metadata.id, metadata.window_ids, metadata.is_chunk
-            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:
-            {filterable_fields}
-            """
-            )
-            self.config.filter_fields = filterable_fields
-
-        if self.from_dataframe:
-            return self.df_description
-        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")
-
-        # 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():
-            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 = {
-                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) -> AgentDoneTool | 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 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
-        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 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 AgentDoneTool(content=NO_ANSWER)
-            answer = self.vecdb.compute_from_docs(docs, plan.dataframe_calc)
-        else:
-            # pass on the query so LLM can handle it
-            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,
-        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('"', "")
-            .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

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

@@ -1,198 +0,0 @@
-"""
-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.agent.tools.orchestration import AgentDoneTool
-from langroid.utils.constants import 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.
-      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 
-      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.
-      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, 
-    - 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!! 
-    - 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)
-        self.enable_message(AgentDoneTool, use=False, handle=True)
-
-    def init_state(self) -> None:
-        super().init_state()
-        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) -> AgentDoneTool:
-        """Format Valid so 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 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,
-            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/lance_rag_task.py

@@ -1,82 +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_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(),
-            llm=agent.config.llm,
-        )
-        query_plan_agent_config.set_system_message()
-
-        critic_config = QueryPlanCriticConfig(
-            doc_schema=agent._get_clean_vecdb_schema(),
-            llm=agent.config.llm,
-        )
-        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