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

langroid/agent/special/lance_rag_new/query_planner_agent.py

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

langroid/agent/special/lance_tools.py

@@ -1,16 +1,21 @@
 import logging
 
 from langroid.agent.tool_message import ToolMessage
-from langroid.pydantic_v1 import BaseModel
+from langroid.pydantic_v1 import BaseModel, Field
 
 logger = logging.getLogger(__name__)
 
 
 class QueryPlan(BaseModel):
-    original_query: str
-    query: str
-    filter: str
-    dataframe_calc: str = ""
+    original_query: str = Field(..., description="The original query for reference")
+    query: str = Field(..., description="A possibly NON-EMPTY rephrased query")
+    filter: str = Field(
+        "",
+        description="Filter condition if needed (or empty if no filter is needed)",
+    )
+    dataframe_calc: str = Field(
+        "", description="An optional Pandas-dataframe calculation/aggregation string"
+    )
 
 
 class QueryPlanTool(ToolMessage):
@@ -19,8 +24,9 @@ class QueryPlanTool(ToolMessage):
     Given a user's query, generate a query <plan> consisting of:
     - <original_query> - the original query for reference
     - <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
-        of the documents (can be same as <original_query> if no rephrasing is needed)
+    - <query> - a possibly NON-EMPTY rephrased query that can be used to match the 
+        CONTENT of the documents 
+        (can be same as <original_query> if no rephrasing is needed)
     - <dataframe_calc> - a Pandas-dataframe calculation/aggregation string
         that can be used to calculate the answer 
         (or empty string if no calculation is needed).
@@ -34,7 +40,7 @@ class QueryPlanAnswerTool(ToolMessage):
     Assemble query <plan> and <answer>
     """
     plan: QueryPlan
-    answer: str
+    answer: str = Field(..., description="The answer received from the assistant")
 
 
 class QueryPlanFeedbackTool(ToolMessage):

langroid/agent/special/neo4j/neo4j_chat_agent.py

@@ -363,7 +363,7 @@ class Neo4jChatAgent(ChatAgent):
             content=content,
             metadata=ChatDocMetaData(
                 # source=Entity.AGENT,
-                sender=Entity.LLM,
+                sender=Entity.AGENT,
                 sender_name=sender_name,
                 recipient=recipient,
             ),

langroid/agent/tool_message.py

@@ -35,12 +35,10 @@ class ToolMessage(ABC, BaseModel):
         request (str): name of agent method to map to.
         purpose (str): purpose of agent method, expressed in general terms.
             (This is used when auto-generating the tool instruction to the LLM)
-        result (str): example of result of agent method.
     """
 
     request: str
     purpose: str
-    result: str = ""
 
     class Config:
         arbitrary_types_allowed = False
@@ -48,7 +46,7 @@ class ToolMessage(ABC, BaseModel):
         validate_assignment = True
         # do not include these fields in the generated schema
         # since we don't require the LLM to specify them
-        schema_extra = {"exclude": {"purpose", "result"}}
+        schema_extra = {"exclude": {"purpose"}}
 
     @classmethod
     def instructions(cls) -> str:
@@ -110,13 +108,13 @@ class ToolMessage(ABC, BaseModel):
         return "\n\n".join(examples_jsons)
 
     def to_json(self) -> str:
-        return self.json(indent=4, exclude={"result", "purpose"})
+        return self.json(indent=4, exclude={"purpose"})
 
     def json_example(self) -> str:
-        return self.json(indent=4, exclude={"result", "purpose"})
+        return self.json(indent=4, exclude={"purpose"})
 
     def dict_example(self) -> Dict[str, Any]:
-        return self.dict(exclude={"result", "purpose"})
+        return self.dict(exclude={"purpose"})
 
     @classmethod
     def default_value(cls, f: str) -> Any:
@@ -220,9 +218,7 @@ class ToolMessage(ABC, BaseModel):
                 if "description" not in parameters["properties"][name]:
                     parameters["properties"][name]["description"] = description
 
-        excludes = (
-            ["result", "purpose"] if request else ["request", "result", "purpose"]
-        )
+        excludes = ["purpose"] if request else ["request", "purpose"]
         # exclude 'excludes' from parameters["properties"]:
         parameters["properties"] = {
             field: details
@@ -263,5 +259,5 @@ class ToolMessage(ABC, BaseModel):
         Returns:
             Dict[str, Any]: simplified schema
         """
-        schema = generate_simple_schema(cls, exclude=["result", "purpose"])
+        schema = generate_simple_schema(cls, exclude=["purpose"])
         return schema

langroid/utils/pydantic_utils.py

@@ -9,8 +9,6 @@ from typing import (
     Tuple,
     Type,
     TypeVar,
-    get_args,
-    get_origin,
     no_type_check,
 )
 
@@ -313,54 +311,6 @@ def pydantic_obj_from_flat_dict(
     return model(**nested_data)
 
 
-def clean_schema(model: Type[BaseModel], excludes: List[str] = []) -> Dict[str, Any]:
-    """
-    Generate a simple schema for a given Pydantic model,
-    including inherited fields, with an option to exclude certain fields.
-    Handles cases where fields are Lists or other generic types and includes
-    field descriptions if available.
-
-    Args:
-        model (Type[BaseModel]): The Pydantic model class.
-        excludes (List[str]): A list of field names to exclude.
-
-    Returns:
-        Dict[str, Any]: A dictionary representing the simple schema.
-    """
-    schema = {}
-
-    for field_name, field_info in model.__fields__.items():
-        if field_name in excludes:
-            continue
-
-        field_type = field_info.outer_type_
-        description = field_info.field_info.description or ""
-
-        # Handle generic types like List[...]
-        if get_origin(field_type):
-            inner_types = get_args(field_type)
-            inner_type_names = [
-                t.__name__ if hasattr(t, "__name__") else str(t) for t in inner_types
-            ]
-            field_type_str = (
-                f"{get_origin(field_type).__name__}" f'[{", ".join(inner_type_names)}]'
-            )
-            schema[field_name] = {"type": field_type_str, "description": description}
-        elif issubclass(field_type, BaseModel):
-            # Directly use the nested model's schema,
-            # integrating it into the current level
-            nested_schema = clean_schema(field_type, excludes)
-            schema[field_name] = {**nested_schema, "description": description}
-        else:
-            # For basic types, use 'type'
-            schema[field_name] = {
-                "type": field_type.__name__,
-                "description": description,
-            }
-
-    return schema
-
-
 @contextmanager
 def temp_update(
     pydantic_object: BaseModel, updates: Dict[str, Any]

langroid/vector_store/base.py

@@ -1,14 +1,14 @@
 import copy
 import logging
 from abc import ABC, abstractmethod
-from typing import Dict, List, Optional, Sequence, Tuple
+from typing import Dict, List, Optional, Sequence, Tuple, Type
 
 import numpy as np
 import pandas as pd
 
 from langroid.embedding_models.base import EmbeddingModel, EmbeddingModelsConfig
 from langroid.embedding_models.models import OpenAIEmbeddingsConfig
-from langroid.mytypes import Document
+from langroid.mytypes import DocMetaData, Document
 from langroid.pydantic_v1 import BaseSettings
 from langroid.utils.algorithms.graph import components, topological_sort
 from langroid.utils.configuration import settings
@@ -32,6 +32,9 @@ class VectorStoreConfig(BaseSettings):
     timeout: int = 60
     host: str = "127.0.0.1"
     port: int = 6333
+    # used when parsing search results back as Document objects
+    document_class: Type[Document] = Document
+    metadata_class: Type[DocMetaData] = DocMetaData
     # compose_file: str = "langroid/vector_store/docker-compose-qdrant.yml"
 
 
@@ -113,8 +116,7 @@ class VectorStore(ABC):
         """
 
         self.config.collection_name = collection_name
-        if collection_name not in self.list_collections() or replace:
-            self.create_collection(collection_name, replace=replace)
+        self.config.replace_collection = replace
 
     @abstractmethod
     def create_collection(self, collection_name: str, replace: bool = False) -> None:

langroid/vector_store/chromadb.py

@@ -8,7 +8,7 @@ from langroid.embedding_models.base import (
 )
 from langroid.embedding_models.models import OpenAIEmbeddingsConfig
 from langroid.exceptions import LangroidImportError
-from langroid.mytypes import DocMetaData, Document
+from langroid.mytypes import Document
 from langroid.utils.configuration import settings
 from langroid.utils.output.printing import print_long_text
 from langroid.vector_store.base import VectorStore, VectorStoreConfig
@@ -200,7 +200,9 @@ class ChromaDB(VectorStore):
             else:
                 m["window_ids"] = m["window_ids"].split(",")
         docs = [
-            Document(content=d, metadata=DocMetaData(**m))
+            self.config.document_class(
+                content=d, metadata=self.config.metadata_class(**m)
+            )
             for d, m in zip(contents, metadatas)
         ]
         return docs