langroid 0.1.233__tar.gz → 0.1.235__tar.gz

{langroid-0.1.233 → langroid-0.1.235}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langroid
-Version: 0.1.233
+Version: 0.1.235
 Summary: Harness LLMs with Multi-Agent Programming
 License: MIT
 Author: Prasad Chalasani

{langroid-0.1.233 → langroid-0.1.235}/langroid/__init__.py

@@ -27,6 +27,7 @@ from .agent.batch import (
 )
 
 from .agent.chat_document import (
+    StatusCode,
     ChatDocument,
     ChatDocMetaData,
 )
@@ -77,6 +78,7 @@ __all__ = [
     "AgentConfig",
     "ChatAgent",
     "ChatAgentConfig",
+    "StatusCode",
     "ChatDocument",
     "ChatDocMetaData",
     "Task",

{langroid-0.1.233 → langroid-0.1.235}/langroid/agent/batch.py

@@ -53,6 +53,7 @@ def run_batch_task_gen(
         message (Optional[str]): optionally overrides the console status messages
         handle_exceptions: bool: Whether to replace exceptions with outputs of None
         max_cost: float: maximum cost to run the task (default 0.0 for unlimited)
+        max_tokens: int: maximum token usage (in and out) (default 0 for unlimited)
 
 
     Returns:

{langroid-0.1.233 → langroid-0.1.235}/langroid/agent/chat_document.py

@@ -1,4 +1,5 @@
 import json
+from enum import Enum
 from typing import List, Optional, Union
 
 from pydantic import BaseModel, Extra
@@ -23,6 +24,23 @@ class ChatDocAttachment(BaseModel):
         extra = Extra.allow
 
 
+class StatusCode(str, Enum):
+    """Codes meant to be returned by task.run(). Some are not used yet."""
+
+    OK = "OK"
+    ERROR = "ERROR"
+    DONE = "DONE"
+    STALLED = "STALLED"
+    INF_LOOP = "INF_LOOP"
+    KILL = "KILL"
+    MAX_TURNS = "MAX_TURNS"
+    MAX_COST = "MAX_COST"
+    MAX_TOKENS = "MAX_TOKENS"
+    TIMEOUT = "TIMEOUT"
+    NO_ANSWER = "NO_ANSWER"
+    USER_QUIT = "USER_QUIT"
+
+
 class ChatDocMetaData(DocMetaData):
     parent: Optional["ChatDocument"] = None
     sender: Entity
@@ -35,6 +53,7 @@ class ChatDocMetaData(DocMetaData):
     usage: Optional[LLMTokenUsage]
     cached: bool = False
     displayed: bool = False
+    status: Optional[StatusCode] = None
 
 
 class ChatDocLoggerFields(BaseModel):

{langroid-0.1.233 → langroid-0.1.235}/langroid/agent/special/__init__.py

@@ -14,7 +14,7 @@ from .table_chat_agent import (
     dataframe_summary,
     TableChatAgent,
     TableChatAgentConfig,
-    RunCodeTool,
+    PandasEvalTool,
 )
 from . import sql
 from . import relevance_extractor_agent
@@ -38,7 +38,7 @@ __all__ = [
     "dataframe_summary",
     "TableChatAgent",
     "TableChatAgentConfig",
-    "RunCodeTool",
+    "PandasEvalTool",
     "sql",
     "relevance_extractor_agent",
     "doc_chat_agent",

{langroid-0.1.233 → langroid-0.1.235}/langroid/agent/special/table_chat_agent.py

@@ -2,10 +2,11 @@
 Agent that supports asking queries about a tabular dataset, internally
 represented as a Pandas dataframe. The `TableChatAgent` is configured with a
 dataset, which can be a Pandas df, file or URL. The delimiter/separator
-is auto-detected. In response to a user query, the Agent's LLM generates Pandas
-code to answer the query. The code is passed via the `run_code` tool/function-call,
-which is handled by the Agent's `run_code` method. This method executes/evaluates
-the code and returns the result as a string.
+is auto-detected. In response to a user query, the Agent's LLM generates a Pandas
+expression (involving a dataframe `df`) to answer the query.
+The expression is passed via the `pandas_eval` tool/function-call,
+which is handled by the Agent's `pandas_eval` method. This method evaluates
+the expression and returns the result as a string.
 """
 
 import io
@@ -35,25 +36,26 @@ DEFAULT_TABLE_CHAT_SYSTEM_MESSAGE = f"""
 You are a savvy data scientist, with expertise in analyzing tabular datasets,
 using Python and the Pandas library for dataframe manipulation.
 Since you do not have access to the dataframe 'df', you
-will need to use the `run_code` tool/function-call to answer my questions.
+will need to use the `pandas_eval` tool/function-call to answer my questions.
 Here is a summary of the dataframe:
 {{summary}}
 Do not assume any columns other than those shown.
-In the code you submit to the `run_code` tool/function, 
-do not forget to include any necessary imports, such as `import pandas as pd`.
-Sometimes you may not be able to answer the question in a single call to `run_code`,
-so you can use a series of calls to `run_code` to build up the answer. 
+In the expression you submit to the `pandas_eval` tool/function, 
+you are allowed to use the variable 'df' to refer to the dataframe.
+
+Sometimes you may not be able to answer the question in a single call to `pandas_eval`,
+so you can use a series of calls to `pandas_eval` to build up the answer. 
 For example you may first want to know something about the possible values in a column.
 
 If you receive a null or other unexpected result, see if you have made an assumption
-in your code, and try another way, or use `run_code` to explore the dataframe 
+in your code, and try another way, or use `pandas_eval` to explore the dataframe 
 before submitting your final code. 
 
 Once you have the answer to the question, possibly after a few steps,
 say {DONE} and show me the answer. If you receive an error message, 
-try using the `run_code` tool/function again with the corrected code. 
+try using the `pandas_eval` tool/function again with the corrected code. 
 
-VERY IMPORTANT: When using the `run_code` tool/function, DO NOT EXPLAIN ANYTHING,
+VERY IMPORTANT: When using the `pandas_eval` tool/function, DO NOT EXPLAIN ANYTHING,
    SIMPLY USE THE TOOL, with the CODE.
 Start by asking me what I want to know about the data.
 """
@@ -123,22 +125,22 @@ class TableChatAgentConfig(ChatAgentConfig):
     )
 
 
-class RunCodeTool(ToolMessage):
-    """Tool/function to run code on a dataframe named `df`"""
+class PandasEvalTool(ToolMessage):
+    """Tool/function to evaluate a pandas expression involving a dataframe `df`"""
 
-    request: str = "run_code"
+    request: str = "pandas_eval"
     purpose: str = """
-            To run <code> on the dataframe 'df' and 
+            To eval a pandas <expression> on the dataframe 'df' and 
             return the results to answer a question.
-            IMPORTANT: ALL the code should be in the <code> field.
+            IMPORTANT: the <expression> field should be a valid pandas expression.
             """
-    code: str
+    expression: str
 
     @classmethod
     def examples(cls) -> List["ToolMessage"]:
         return [
-            cls(code="df.head()"),
-            cls(code="df[(df['gender'] == 'Male')]['income'].mean()"),
+            cls(expression="df.head()"),
+            cls(expression="df[(df['gender'] == 'Male')]['income'].mean()"),
         ]
 
 
@@ -147,7 +149,7 @@ class TableChatAgent(ChatAgent):
     Agent for chatting with a collection of documents.
     """
 
-    sent_code: bool = False
+    sent_expression: bool = False
 
     def __init__(self, config: TableChatAgentConfig):
         if isinstance(config.data, pd.DataFrame):
@@ -170,8 +172,8 @@ class TableChatAgent(ChatAgent):
             {self.df.columns}
             """
         )
-        # enable the agent to use and handle the RunCodeTool
-        self.enable_message(RunCodeTool)
+        # enable the agent to use and handle the PandasEvalTool
+        self.enable_message(PandasEvalTool)
 
     def user_response(
         self,
@@ -179,44 +181,31 @@ class TableChatAgent(ChatAgent):
     ) -> Optional[ChatDocument]:
         response = super().user_response(msg)
         if response is not None and response.content != "":
-            self.sent_code = False
+            self.sent_expression = False
         return response
 
-    def run_code(self, msg: RunCodeTool) -> str:
+    def pandas_eval(self, msg: PandasEvalTool) -> str:
         """
-        Handle a RunCodeTool message by running the code and returning the result.
+        Handle a PandasEvalTool message by evaluating the `expression` field
+            and returning the result.
         Args:
-            msg (RunCodeTool): The tool-message to handle.
+            msg (PandasEvalTool): The tool-message to handle.
 
         Returns:
             str: The result of running the code along with any print output.
         """
-        self.sent_code = True
-        code = msg.code
-        # Create a dictionary that maps 'df' to the actual DataFrame
+        self.sent_expression = True
+        exprn = msg.expression
         local_vars = {"df": self.df}
-
         # Create a string-based I/O stream
         code_out = io.StringIO()
 
         # Temporarily redirect standard output to our string-based I/O stream
         sys.stdout = code_out
 
-        # Split the code into lines
-        lines = code.strip().split("\n")
-
-        lines = [
-            "import pandas as pd",
-            "import numpy as np",
-        ] + lines
-
-        # Run all lines as statements except for the last one
-        for line in lines[:-1]:
-            exec(line, {}, local_vars)
-
         # Evaluate the last line and get the result
         try:
-            eval_result = pd.eval(lines[-1], local_dict=local_vars)
+            eval_result = pd.eval(exprn, local_dict=local_vars)
         except Exception as e:
             eval_result = f"ERROR: {type(e)}: {e}"
 
@@ -242,15 +231,15 @@ class TableChatAgent(ChatAgent):
     def handle_message_fallback(
         self, msg: str | ChatDocument
     ) -> str | ChatDocument | None:
-        """Handle scenario where LLM forgets to say DONE or forgets to use run_code"""
+        """Handle scenario where LLM forgets to say DONE or
+        forgets to use pandas_eval"""
         if isinstance(msg, ChatDocument) and msg.metadata.sender == lr.Entity.LLM:
-            if self.sent_code:
+            if self.sent_expression:
                 return DONE
             else:
                 return """
-                    You forgot to use the `run_code` tool/function to find the answer.
-                    Try again using the `run_code` tool/function.
-                    Remember that ALL your code, including imports, 
-                    should be in the `code` field.
+                    You forgot to use the `pandas_eval` tool/function 
+                    to find the answer.
+                    Try again using the `pandas_eval` tool/function.
                     """
         return None

{langroid-0.1.233 → langroid-0.1.235}/langroid/agent/task.py

@@ -27,11 +27,20 @@ from langroid.agent.chat_document import (
     ChatDocLoggerFields,
     ChatDocMetaData,
     ChatDocument,
+    StatusCode,
 )
+from langroid.cachedb.redis_cachedb import RedisCache, RedisCacheConfig
 from langroid.mytypes import Entity
 from langroid.parsing.parse_json import extract_top_level_json
 from langroid.utils.configuration import settings
-from langroid.utils.constants import DONE, NO_ANSWER, PASS, PASS_TO, SEND_TO, USER_QUIT
+from langroid.utils.constants import (
+    DONE,
+    NO_ANSWER,
+    PASS,
+    PASS_TO,
+    SEND_TO,
+    USER_QUIT_STRINGS,
+)
 from langroid.utils.logging import RichFileLogger, setup_file_logger
 
 logger = logging.getLogger(__name__)
@@ -73,6 +82,8 @@ class Task:
     the value of `result()`, which is the final result of the task.
     """
 
+    # class variable called `cache` that is a RedisCache object
+    cache: RedisCache = RedisCache(RedisCacheConfig(fake=False))
     def __init__(
         self,
         agent: Optional[Agent] = None,
@@ -141,7 +152,6 @@ class Task:
         """
         if agent is None:
             agent = ChatAgent()
-
         self.callbacks = SimpleNamespace(
             show_subtask_response=noop_fn,
             set_parent_agent=noop_fn,
@@ -172,6 +182,7 @@ class Task:
                 agent.set_user_message(user_message)
         self.max_cost: float = 0
         self.max_tokens: int = 0
+        self.session_id: str = ""
         self.logger: None | RichFileLogger = None
         self.tsv_logger: None | logging.Logger = None
         self.color_log: bool = False if settings.notebook else True
@@ -285,6 +296,54 @@ class Task:
     def __str__(self) -> str:
         return f"{self.name}"
 
+    def _cache_session_store(self, key: str, value: str) -> None:
+        """
+        Cache a key-value pair for the current session.
+        E.g. key = "kill", value = "1"
+        """
+        try:
+            self.cache.store(f"{self.session_id}:{key}", value)
+        except Exception as e:
+            logging.error(f"Error in Task._cache_session_store: {e}")
+
+    def _cache_session_lookup(self, key: str) -> Dict[str, Any] | str | None:
+        """
+        Retrieve a value from the cache for the current session.
+        """
+        session_id_key = f"{self.session_id}:{key}"
+        try:
+            cached_val = self.cache.retrieve(session_id_key)
+        except Exception as e:
+            logging.error(f"Error in Task._cache_session_lookup: {e}")
+            return None
+        return cached_val
+
+    def _is_kill(self) -> bool:
+        """
+        Check if the current session is killed.
+        """
+        return self._cache_session_lookup("kill") == "1"
+
+    def _set_alive(self) -> None:
+        """
+        Initialize the kill status of the current session.
+        """
+        self._cache_session_store("kill", "0")
+
+    @classmethod
+    def kill_session(cls, session_id=""):
+        """
+        Kill the current session.
+        """
+        session_id_kill_key = f"{session_id}:kill"
+        cls.cache.store(session_id_kill_key, "1")
+
+    def kill(self) -> None:
+        """
+        Kill the task run associated with the current session.
+        """
+        self._cache_session_store("kill", "1")
+
     @property
     def _level(self) -> int:
         if self.caller is None:
@@ -378,6 +437,7 @@ class Task:
         caller: None | Task = None,
         max_cost: float = 0,
         max_tokens: int = 0,
+        session_id: str = "",
     ) -> Optional[ChatDocument]:
         """Synchronous version of `run_async()`.
         See `run_async()` for details."""
@@ -385,6 +445,9 @@ class Task:
         self.n_stalled_steps = 0
         self.max_cost = max_cost
         self.max_tokens = max_tokens
+        self.session_id = session_id
+        self._set_alive()
+
         assert (
             msg is None or isinstance(msg, str) or isinstance(msg, ChatDocument)
         ), f"msg arg in Task.run() must be None, str, or ChatDocument, not {type(msg)}"
@@ -406,15 +469,19 @@ class Task:
         i = 0
         while True:
             self.step()
-            if self.done():
+            done, status = self.done()
+            if done:
                 if self._level == 0 and not settings.quiet:
                     print("[magenta]Bye, hope this was useful!")
                 break
             i += 1
             if turns > 0 and i >= turns:
+                status = StatusCode.MAX_TURNS
                 break
 
         final_result = self.result()
+        if final_result is not None:
+            final_result.metadata.status = status
         self._post_run_loop()
         return final_result
 
@@ -425,6 +492,7 @@ class Task:
         caller: None | Task = None,
         max_cost: float = 0,
         max_tokens: int = 0,
+        session_id: str = "",
     ) -> Optional[ChatDocument]:
         """
         Loop over `step()` until task is considered done or `turns` is reached.
@@ -443,6 +511,7 @@ class Task:
             caller (Task|None): the calling task, if any
             max_cost (float): max cost allowed for the task (default 0 -> no limit)
             max_tokens (int): max tokens allowed for the task (default 0 -> no limit)
+            session_id (str): session id for the task
 
         Returns:
             Optional[ChatDocument]: valid result of the task.
@@ -456,6 +525,9 @@ class Task:
         self.n_stalled_steps = 0
         self.max_cost = max_cost
         self.max_tokens = max_tokens
+        self.session_id = session_id
+        self._set_alive()
+
         if (
             isinstance(msg, ChatDocument)
             and msg.metadata.recipient != ""
@@ -473,15 +545,19 @@ class Task:
         i = 0
         while True:
             await self.step_async()
-            if self.done():
+            done, status = self.done()
+            if done:
                 if self._level == 0 and not settings.quiet:
                     print("[magenta]Bye, hope this was useful!")
                 break
             i += 1
             if turns > 0 and i >= turns:
+                status = StatusCode.MAX_TURNS
                 break
 
         final_result = self.result()
+        if final_result is not None:
+            final_result.metadata.status = status
         self._post_run_loop()
         return final_result
 
@@ -942,6 +1018,7 @@ class Task:
         recipient = result_msg.metadata.recipient if result_msg else None
         responder = result_msg.metadata.parent_responder if result_msg else None
         tool_ids = result_msg.metadata.tool_ids if result_msg else []
+        status = result_msg.metadata.status if result_msg else None
 
         # regardless of which entity actually produced the result,
         # when we return the result, we set entity to USER
@@ -954,6 +1031,7 @@ class Task:
                 source=Entity.USER,
                 sender=Entity.USER,
                 block=block,
+                status=status,
                 parent_responder=responder,
                 sender_name=self.name,
                 recipient=recipient,
@@ -1036,7 +1114,7 @@ class Task:
 
     def done(
         self, result: ChatDocument | None = None, r: Responder | None = None
-    ) -> bool:
+    ) -> Tuple[bool, StatusCode]:
         """
         Check if task is done. This is the default behavior.
         Derived classes can override this.
@@ -1046,26 +1124,29 @@ class Task:
                 Not used here, but could be used by derived classes.
         Returns:
             bool: True if task is done, False otherwise
+            StatusCode: status code indicating why task is done
         """
+        if self._is_kill():
+            return (True, StatusCode.KILL)
         result = result or self.pending_message
         user_quit = (
             result is not None
-            and result.content in USER_QUIT
+            and result.content in USER_QUIT_STRINGS
             and result.metadata.sender == Entity.USER
         )
         if self._level == 0 and self.only_user_quits_root:
             # for top-level task, only user can quit out
-            return user_quit
+            return (user_quit, StatusCode.USER_QUIT if user_quit else StatusCode.OK)
 
         if self.is_done:
-            return True
+            return (True, StatusCode.DONE)
 
         if self.n_stalled_steps >= self.max_stalled_steps:
             # we are stuck, so bail to avoid infinite loop
             logger.warning(
                 f"Task {self.name} stuck for {self.max_stalled_steps} steps; exiting."
             )
-            return True
+            return (True, StatusCode.STALLED)
 
         if self.max_cost > 0 and self.agent.llm is not None:
             try:
@@ -1073,7 +1154,7 @@ class Task:
                     logger.warning(
                         f"Task {self.name} cost exceeded {self.max_cost}; exiting."
                     )
-                    return True
+                    return (True, StatusCode.MAX_COST)
             except Exception:
                 pass
 
@@ -1083,10 +1164,10 @@ class Task:
                     logger.warning(
                         f"Task {self.name} uses > {self.max_tokens} tokens; exiting."
                     )
-                    return True
+                    return (True, StatusCode.MAX_TOKENS)
             except Exception:
                 pass
-        return (
+        final = (
             # no valid response from any entity/agent in current turn
             result is None
             # An entity decided task is done
@@ -1103,6 +1184,7 @@ class Task:
             # )
             or user_quit
         )
+        return (final, StatusCode.OK)
 
     def valid(
         self,
@@ -1120,7 +1202,7 @@ class Task:
 
         # if task would be considered done given responder r's `result`,
         # then consider the result valid.
-        if result is not None and self.done(result, r):
+        if result is not None and self.done(result, r)[0]:
             return True
         return (
             result is not None

{langroid-0.1.233 → langroid-0.1.235}/langroid/cachedb/redis_cachedb.py

@@ -109,7 +109,7 @@ class RedisCache(CacheDB):
             key (str): The key to retrieve the value for.
 
         Returns:
-            dict: The value associated with the key.
+            dict|str|None: The value associated with the key.
         """
         with self.redis_client() as client:  # type: ignore
             try:

{langroid-0.1.233 → langroid-0.1.235}/langroid/utils/constants.py

@@ -13,7 +13,7 @@ class Colors(BaseModel):
     RESET: str = "\033[0m"
 
 
-USER_QUIT = ["q", "x", "quit", "exit", "bye"]
+USER_QUIT_STRINGS = ["q", "x", "quit", "exit", "bye"]
 NO_ANSWER = "DO-NOT-KNOW"
 DONE = "DONE"
 PASS = "__PASS__"

{langroid-0.1.233 → langroid-0.1.235}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langroid"
-version = "0.1.233"
+version = "0.1.235"
 description = "Harness LLMs with Multi-Agent Programming"
 authors = ["Prasad Chalasani <pchalasani@gmail.com>"]
 readme = "README.md"