langroid 0.1.265__py3-none-any.whl → 0.2.0__py3-none-any.whl

langroid/agent/task.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import asyncio
 import copy
 import logging
+import threading
 from collections import Counter, deque
 from types import SimpleNamespace
 from typing import (
@@ -13,7 +14,6 @@ from typing import (
     Dict,
     List,
     Optional,
-    Set,
     Tuple,
     Type,
     cast,
@@ -47,6 +47,7 @@ from langroid.utils.constants import (
     USER_QUIT_STRINGS,
 )
 from langroid.utils.logging import RichFileLogger, setup_file_logger
+from langroid.utils.object_registry import scheduled_cleanup
 from langroid.utils.system import hash
 
 logger = logging.getLogger(__name__)
@@ -65,14 +66,17 @@ class TaskConfig(BaseModel):
     we have config classes for `Agent`, `ChatAgent`, `LanguageModel`, etc.
 
     Attributes:
-        inf_loop_cycle_len: max exact-loop cycle length: 0 => no inf loop test
-        inf_loop_dominance_factor: dominance factor for exact-loop detection
-        inf_loop_wait_factor: wait this * cycle_len msgs before loop-check
+        inf_loop_cycle_len (int): max exact-loop cycle length: 0 => no inf loop test
+        inf_loop_dominance_factor (float): dominance factor for exact-loop detection
+        inf_loop_wait_factor (int): wait this * cycle_len msgs before loop-check
+        restart_subtask_run (bool): whether to restart *every* run of this task
+            when run as a subtask.
     """
 
     inf_loop_cycle_len: int = 10
     inf_loop_dominance_factor: float = 1.5
     inf_loop_wait_factor: int = 5
+    restart_as_subtask: bool = False
 
 
 class Task:
@@ -107,6 +111,7 @@ class Task:
 
     # class variable called `cache` that is a RedisCache object
     _cache: RedisCache | None = None
+    _background_tasks_started: bool = False
 
     def __init__(
         self,
@@ -149,7 +154,7 @@ class Task:
                 One run of step() is considered a "turn".
             system_message (str): if not empty, overrides agent's system_message
             user_message (str): if not empty, overrides agent's user_message
-            restart (bool): if true, resets the agent's message history
+            restart (bool): if true, resets the agent's message history *at every run*.
             default_human_response (str): default response from user; useful for
                 testing, to avoid interactive input from user.
                 [Instead of this, setting `interactive` usually suffices]
@@ -187,6 +192,8 @@ class Task:
             set_parent_agent=noop_fn,
         )
         self.config = config
+        # how to behave as a sub-task; can be overriden by `add_sub_task()`
+        self.config_sub_task = copy.deepcopy(config)
         # counts of distinct pending messages in history,
         # to help detect (exact) infinite loops
         self.message_counter: Counter[str] = Counter()
@@ -208,24 +215,25 @@ class Task:
                 the config may affect other agents using the same config.
                 """
             )
-
+        self.restart = restart
+        agent = cast(ChatAgent, agent)
+        self.agent: ChatAgent = agent
         if isinstance(agent, ChatAgent) and len(agent.message_history) == 0 or restart:
-            agent = cast(ChatAgent, agent)
-            agent.clear_history(0)
-            agent.clear_dialog()
+            self.agent.clear_history(0)
+            self.agent.clear_dialog()
             # possibly change the system and user messages
             if system_message:
                 # we always have at least 1 task_message
-                agent.set_system_message(system_message)
+                self.agent.set_system_message(system_message)
             if user_message:
-                agent.set_user_message(user_message)
+                self.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
-        self.agent = agent
+
         self.step_progress = False  # progress in current step?
         self.n_stalled_steps = 0  # how many consecutive steps with no progress?
         self.max_stalled_steps = max_stalled_steps
@@ -305,7 +313,6 @@ class Task:
 
         # other sub_tasks this task can delegate to
         self.sub_tasks: List[Task] = []
-        self.parent_task: Set[Task] = set()
         self.caller: Task | None = None  # which task called this task's `run` method
 
     def clone(self, i: int) -> "Task":
@@ -321,7 +328,7 @@ class Task:
             single_round=self.single_round,
             system_message=self.agent.system_message,
             user_message=self.agent.user_message,
-            restart=False,
+            restart=self.restart,
             default_human_response=self.default_human_response,
             interactive=self.interactive,
             erase_substeps=self.erase_substeps,
@@ -338,6 +345,19 @@ class Task:
             cls._cache = RedisCache(RedisCacheConfig(fake=False))
         return cls._cache
 
+    @classmethod
+    def _start_background_tasks(cls) -> None:
+        """Start background object registry cleanup thread. NOT USED."""
+        if cls._background_tasks_started:
+            return
+        cls._background_tasks_started = True
+        cleanup_thread = threading.Thread(
+            target=scheduled_cleanup,
+            args=(600,),
+            daemon=True,
+        )
+        cleanup_thread.start()
+
     def __repr__(self) -> str:
         return f"{self.name}"
 
@@ -416,24 +436,37 @@ class Task:
     def _leave(self) -> str:
         return self._indent + "<<<"
 
-    def add_sub_task(self, task: Task | List[Task]) -> None:
+    def add_sub_task(
+        self,
+        task: (
+            Task | List[Task] | Tuple[Task, TaskConfig] | List[Tuple[Task, TaskConfig]]
+        ),
+    ) -> None:
         """
         Add a sub-task (or list of subtasks) that this task can delegate
         (or fail-over) to. Note that the sequence of sub-tasks is important,
         since these are tried in order, as the parent task searches for a valid
-        response.
+        response (unless a sub-task is explicitly addressed).
 
         Args:
-            task (Task|List[Task]): sub-task(s) to add
+            task: A task, or list of tasks, or a tuple of task and task config,
+                or a list of tuples of task and task config.
+                These tasks are added as sub-tasks of the current task.
+                The task configs (if any) dictate how the tasks are run when
+                invoked as sub-tasks of other tasks. This allows users to specify
+                behavior applicable only in the context of a particular task-subtask
+                combination.
         """
-
         if isinstance(task, list):
             for t in task:
                 self.add_sub_task(t)
             return
-        assert isinstance(task, Task), f"added task must be a Task, not {type(task)}"
 
-        task.parent_task.add(self)  # add myself to set of parent tasks of `task`
+        if isinstance(task, tuple):
+            task, config = task
+        else:
+            config = TaskConfig()
+        task.config_sub_task = config
         self.sub_tasks.append(task)
         self.name_sub_task_map[task.name] = task
         self.responders.append(cast(Responder, task))
@@ -460,12 +493,28 @@ class Task:
                     sender=Entity.USER,
                 ),
             )
+        elif msg is None and len(self.agent.message_history) > 1:
+            # if agent has a history beyond system msg, set the
+            # pending message to the ChatDocument linked from
+            # last message in the history
+            last_agent_msg = self.agent.message_history[-1]
+            self.pending_message = ChatDocument.from_id(last_agent_msg.chat_document_id)
+            if self.pending_message is not None:
+                self.pending_sender = self.pending_message.metadata.sender
         else:
-            self.pending_message = copy.deepcopy(msg)
+            if isinstance(msg, ChatDocument):
+                # carefully deep-copy: fresh metadata.id, register
+                # as new obj in registry
+                self.pending_message = ChatDocument.deepcopy(msg)
             if self.pending_message is not None and self.caller is not None:
                 # msg may have come from `caller`, so we pretend this is from
                 # the CURRENT task's USER entity
                 self.pending_message.metadata.sender = Entity.USER
+                # update parent, child, agent pointers
+                if msg is not None:
+                    msg.metadata.child_id = self.pending_message.metadata.id
+                    self.pending_message.metadata.parent_id = msg.metadata.id
+                self.pending_message.metadata.agent_id = self.agent.id
 
         self._show_pending_message_if_debug()
 
@@ -484,6 +533,13 @@ class Task:
         self.log_message(Entity.USER, self.pending_message)
         return self.pending_message
 
+    def reset_all_sub_tasks(self) -> None:
+        """Recursively reset message history of own agent and all sub-tasks"""
+        self.agent.clear_history(0)
+        self.agent.clear_dialog()
+        for t in self.sub_tasks:
+            t.reset_all_sub_tasks()
+
     def run(
         self,
         msg: Optional[str | ChatDocument] = None,
@@ -495,6 +551,14 @@ class Task:
     ) -> Optional[ChatDocument]:
         """Synchronous version of `run_async()`.
         See `run_async()` for details."""
+        if (self.restart and caller is None) or (
+            self.config_sub_task.restart_as_subtask and caller is not None
+        ):
+            # We are either at top level, with restart = True, OR
+            # we are a sub-task with restart_as_subtask = True,
+            # so reset own agent and recursively for all sub-tasks
+            self.reset_all_sub_tasks()
+
         self.task_progress = False
         self.n_stalled_steps = 0
         self.max_cost = max_cost
@@ -597,6 +661,18 @@ class Task:
         # have come from another LLM), as far as this agent is concerned, the initial
         # message can be considered to be from the USER
         # (from the POV of this agent's LLM).
+
+        if (
+            self.restart
+            and caller is None
+            or self.config_sub_task.restart_as_subtask
+            and caller is not None
+        ):
+            # We are either at top level, with restart = True, OR
+            # we are a sub-task with restart_as_subtask = True,
+            # so reset own agent and recursively for all sub-tasks
+            self.reset_all_sub_tasks()
+
         self.task_progress = False
         self.n_stalled_steps = 0
         self.max_cost = max_cost
@@ -701,6 +777,23 @@ class Task:
             if self.erase_substeps:
                 # TODO I don't like directly accessing agent message_history. Revisit.
                 # (Pchalasani)
+                # Note: msg history will consist of:
+                # - H: the original msg history, ending at idx= self.message_history_idx
+                # - R: this agent's response, which presumably leads to:
+                # - X: a series of back-and-forth msgs (including with agent's own
+                #     responders and with sub-tasks)
+                # - F: the final result message, from this agent.
+                # Here we are deleting all of [X] from the agent's message history,
+                # so that it simply looks as if the sub-tasks never happened.
+
+                dropped = self.agent.message_history[
+                    self.message_history_idx + 2 : n_messages - 1
+                ]
+                # first delete the linked ChatDocuments (and descendants) from
+                # ObjectRegistry
+                for msg in dropped:
+                    ChatDocument.delete_id(msg.chat_document_id)
+                # then delete the messages from the agent's message_history
                 del self.agent.message_history[
                     self.message_history_idx + 2 : n_messages - 1
                 ]
@@ -750,9 +843,11 @@ class Task:
 
         if (
             Entity.USER in self.responders
+            and self.interactive
             and not self.human_tried
             and not self.agent.has_tool_message_attempt(self.pending_message)
         ):
+            # When in interactive mode,
             # Give human first chance if they haven't been tried in last step,
             # and the msg is not a tool-call attempt;
             # This ensures human gets a chance to respond,
@@ -778,6 +873,8 @@ class Task:
                         recipient=recipient,
                     ),
                 )
+                # no need to register this dummy msg in ObjectRegistry
+                ChatDocument.delete_id(log_doc.id())
                 self.log_message(r, log_doc)
                 continue
             self.human_tried = r == Entity.USER
@@ -844,6 +941,7 @@ class Task:
 
         if (
             Entity.USER in self.responders
+            and self.interactive
             and not self.human_tried
             and not self.agent.has_tool_message_attempt(self.pending_message)
         ):
@@ -870,6 +968,8 @@ class Task:
                         recipient=recipient,
                     ),
                 )
+                # no need to register this dummy msg in ObjectRegistry
+                ChatDocument.delete_id(log_doc.id())
                 self.log_message(r, log_doc)
                 continue
             self.human_tried = r == Entity.USER
@@ -905,10 +1005,26 @@ class Task:
         # Contrast this with self.pending_message.metadata.sender, which is an ENTITY
         # of this agent, or a sub-task's agent.
         if not self.is_pass_thru:
-            self.pending_sender = r
-        result.metadata.parent = parent
-        if not self.is_pass_thru:
+            if (
+                self.pending_message is not None
+                and self.pending_message.metadata.agent_id == self.agent.id
+            ):
+                # when pending msg is from our own agent, respect the sender set there,
+                # since sometimes a response may "mock" as if the response is from
+                # another entity (e.g when using RewindTool, the agent handler
+                # returns a result as if it were from the LLM).
+                self.pending_sender = result.metadata.sender
+            else:
+                # when pending msg is from a sub-task, the sender is the sub-task
+                self.pending_sender = r
             self.pending_message = result
+        # set the parent/child links ONLY if not already set by agent internally,
+        # which may happen when using the RewindTool
+        if parent is not None and not result.metadata.parent_id:
+            result.metadata.parent_id = parent.id()
+        if parent is not None and not parent.metadata.child_id:
+            parent.metadata.child_id = result.id()
+
         self.log_message(self.pending_sender, result, mark=True)
         self.step_progress = True
         self.task_progress = True
@@ -941,9 +1057,10 @@ class Task:
             responder = (
                 Entity.LLM if self.pending_sender == Entity.USER else Entity.USER
             )
+            parent_id = "" if parent is None else parent.id()
             self.pending_message = ChatDocument(
                 content=NO_ANSWER,
-                metadata=ChatDocMetaData(sender=responder, parent=parent),
+                metadata=ChatDocMetaData(sender=responder, parent_id=parent_id),
             )
             self.pending_sender = responder
         self.log_message(self.pending_sender, self.pending_message, mark=True)
@@ -1089,7 +1206,7 @@ class Task:
         # regardless of which entity actually produced the result,
         # when we return the result, we set entity to USER
         # since to the "parent" task, this result is equivalent to a response from USER
-        return ChatDocument(
+        result_doc = ChatDocument(
             content=content,
             function_call=fun_call,
             tool_messages=tool_messages,
@@ -1101,8 +1218,14 @@ class Task:
                 sender_name=self.name,
                 recipient=recipient,
                 tool_ids=tool_ids,
+                parent_id=result_msg.id() if result_msg else "",
+                agent_id=str(self.agent.id),
             ),
         )
+        if self.pending_message is not None:
+            self.pending_message.metadata.child_id = result_doc.id()
+
+        return result_doc
 
     def _is_empty_message(self, msg: str | ChatDocument | None) -> bool:
         """

langroid/agent/tools/__init__.py

@@ -1,7 +1,9 @@
 from . import google_search_tool
 from . import recipient_tool
+from . import rewind_tool
 from .google_search_tool import GoogleSearchTool
 from .recipient_tool import AddRecipientTool, RecipientTool
+from .rewind_tool import RewindTool
 
 __all__ = [
     "GoogleSearchTool",
@@ -9,4 +11,6 @@ __all__ = [
     "RecipientTool",
     "google_search_tool",
     "recipient_tool",
+    "rewind_tool",
+    "RewindTool",
 ]

langroid/agent/tools/rewind_tool.py

@@ -0,0 +1,136 @@
+"""
+The `rewind_tool` is used to rewind to the `n`th previous Assistant message
+and replace it with a new `content`. This is useful in several scenarios and
+- saves token-cost + inference time,
+- reduces distracting clutter in chat history, which helps improve response quality.
+
+This is intended to mimic how a human user might use a chat interface, where they
+go down a conversation path, and want to go back in history to "edit and re-submit"
+a previous message, to get a better response.
+
+See usage examples in `tests/main/test_rewind_tool.py`.
+"""
+
+from typing import List, Tuple
+
+import langroid.language_models as lm
+from langroid.agent.chat_agent import ChatAgent
+from langroid.agent.chat_document import ChatDocument
+from langroid.agent.tool_message import ToolMessage
+
+
+def prune_messages(agent: ChatAgent, idx: int) -> ChatDocument | None:
+    """
+    Clear the message history of agent, starting at index `idx`,
+    taking care to first clear all dependent messages (possibly from other agents'
+    message histories) that are linked to the message at `idx`, via the `child_id` field
+    of the `metadata` field of the ChatDocument linked from the message at `idx`.
+
+    Args:
+        agent (ChatAgent): The agent whose message history is to be pruned.
+        idx (int): The index from which to start clearing the message history.
+
+    Returns:
+        The parent ChatDocument of the ChatDocument linked from the message at `idx`,
+        if it exists, else None.
+
+    """
+    assert idx >= 0, "Invalid index for message history!"
+    chat_doc_id = agent.message_history[idx].chat_document_id
+    chat_doc = ChatDocument.from_id(chat_doc_id)
+    assert chat_doc is not None, "ChatDocument not found in registry!"
+
+    parent = ChatDocument.from_id(chat_doc.metadata.parent_id)  # may be None
+    # We're invaliding the msg at idx,
+    # so starting with chat_doc, go down the child links
+    # and clear history of each agent, to the msg_idx
+    curr_doc = chat_doc
+    while child_doc := curr_doc.metadata.child:
+        if child_doc.metadata.msg_idx >= 0:
+            child_agent = ChatAgent.from_id(child_doc.metadata.agent_id)
+            if child_agent is not None:
+                child_agent.clear_history(child_doc.metadata.msg_idx)
+        curr_doc = child_doc
+
+    # Clear out ObjectRegistry entries for this ChatDocuments
+    # and all descendants (in case they weren't already cleared above)
+    ChatDocument.delete_id(chat_doc.id())
+
+    # Finally, clear this agent's history back to idx,
+    # and replace the msg at idx with the new content
+    agent.clear_history(idx)
+    return parent
+
+
+class RewindTool(ToolMessage):
+    """
+    Used by LLM to rewind (i.e. backtrack) to the `n`th Assistant message
+    and replace with a new msg.
+    """
+
+    request: str = "rewind_tool"
+    purpose: str = """
+        To rewind the conversation and replace the 
+        <n>'th Assistant message with <content>
+        """
+    n: int
+    content: str
+
+    @classmethod
+    def examples(cls) -> List["ToolMessage" | Tuple[str, "ToolMessage"]]:
+        return [
+            cls(n=1, content="What are the 3 major causes of heart disease?"),
+            (
+                """
+                I want to change my 2nd message to Bob, to say 
+                'who wrote the book Grime and Banishment?'
+                """,
+                cls(n=2, content="who wrote the book 'Grime and Banishment'?"),
+            ),
+        ]
+
+    def response(self, agent: ChatAgent) -> str | ChatDocument:
+        """
+        Define the tool-handler method for this tool here itself,
+        since it is a generic tool whose functionality should be the
+        same for any agent.
+
+        When LLM has correctly used this tool, rewind this agent's
+        `message_history` to the `n`th assistant msg, and replace it with `content`.
+        We need to mock it as if the LLM is sending this message.
+
+        Within a multi-agent scenario, this also means that any other messages dependent
+        on this message will need to be invalidated --
+        so go down the chain of child messages and clear each agent's history
+        back to the `msg_idx` corresponding to the child message.
+
+        Returns:
+            (ChatDocument): with content set to self.content.
+        """
+        idx = agent.nth_message_idx_with_role(lm.Role.ASSISTANT, self.n)
+        if idx < 0:
+            # set up a corrective message from AGENT
+            msg = f"""
+                Could not rewind to {self.n}th Assistant message!
+                Please check the value of `n` and try again.
+                Or it may be too early to use the `rewind_tool`.
+                """
+            return agent.create_agent_response(msg)
+
+        parent = prune_messages(agent, idx)
+
+        # create ChatDocument with new content, to be returned as result of this tool
+        result_doc = agent.create_llm_response(self.content)
+        result_doc.metadata.parent_id = "" if parent is None else parent.id()
+        result_doc.metadata.agent_id = agent.id
+        result_doc.metadata.msg_idx = idx
+
+        # replace the message at idx with this new message
+        agent.message_history.append(ChatDocument.to_LLMMessage(result_doc))
+
+        # set the replaced doc's parent's child to this result_doc
+        if parent is not None:
+            # first remove the this parent's child from registry
+            ChatDocument.delete_id(parent.metadata.child_id)
+            parent.metadata.child_id = result_doc.id()
+        return result_doc

langroid/language_models/__init__.py

@@ -20,6 +20,7 @@ from .openai_gpt import (
     OpenAIGPTConfig,
     OpenAIGPT,
 )
+from .mock_lm import MockLM, MockLMConfig
 from .azure_openai import AzureConfig, AzureGPT
 
 
@@ -43,4 +44,6 @@ __all__ = [
     "OpenAIGPT",
     "AzureConfig",
     "AzureGPT",
+    "MockLM",
+    "MockLMConfig",
 ]

langroid/language_models/base.py

@@ -4,7 +4,17 @@ import logging
 from abc import ABC, abstractmethod
 from datetime import datetime
 from enum import Enum
-from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    Union,
+    cast,
+)
 
 from langroid.cachedb.base import CacheDBConfig
 from langroid.parsing.agent_chats import parse_message
@@ -134,12 +144,15 @@ class LLMMessage(BaseModel):
     content: str
     function_call: Optional[LLMFunctionCall] = None
     timestamp: datetime = Field(default_factory=datetime.utcnow)
+    # link to corresponding chat document, for provenance/rewind purposes
+    chat_document_id: str = ""
 
     def api_dict(self) -> Dict[str, Any]:
         """
-        Convert to dictionary for API request.
-        DROP the tool_id, since it is only for use in the Assistant API,
-        not the completion API.
+        Convert to dictionary for API request, keeping ONLY
+        the fields that are expected in an API call!
+        E.g., DROP the tool_id, since it is only for use in the Assistant API,
+            not the completion API.
         Returns:
             dict: dictionary representation of LLM message
         """
@@ -155,8 +168,10 @@ class LLMMessage(BaseModel):
                 dict_no_none["function_call"]["arguments"] = json.dumps(
                     dict_no_none["function_call"]["arguments"]
                 )
+        # IMPORTANT! drop fields that are not expected in API call
         dict_no_none.pop("tool_id", None)
         dict_no_none.pop("timestamp", None)
+        dict_no_none.pop("chat_document_id", None)
         return dict_no_none
 
     def __str__(self) -> str:
@@ -268,11 +283,15 @@ class LanguageModel(ABC):
                 """
             )
         from langroid.language_models.azure_openai import AzureGPT
+        from langroid.language_models.mock_lm import MockLM, MockLMConfig
         from langroid.language_models.openai_gpt import OpenAIGPT
 
         if config is None or config.type is None:
             return None
 
+        if config.type == "mock":
+            return MockLM(cast(MockLMConfig, config))
+
         openai: Union[Type[AzureGPT], Type[OpenAIGPT]]
 
         if config.type == "azure":