langroid 0.1.234__tar.gz → 0.1.235__tar.gz

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

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

{langroid-0.1.234 → 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.234 → 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.234 → 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.234 → langroid-0.1.235}/langroid/agent/special/table_chat_agent.py

@@ -203,7 +203,6 @@ class TableChatAgent(ChatAgent):
         # Temporarily redirect standard output to our string-based I/O stream
         sys.stdout = code_out
 
-
         # Evaluate the last line and get the result
         try:
             eval_result = pd.eval(exprn, local_dict=local_vars)

{langroid-0.1.234 → 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.234 → 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.234 → 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.234 → langroid-0.1.235}/pyproject.toml

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