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

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,137 @@
+"""
+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?"),
+            (
+                """
+                Based on the conversation so far, I realize I would get a better
+                response from Bob if rephrase my 2nd message to him to: 
+                '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":

langroid/language_models/mock_lm.py

@@ -0,0 +1,91 @@
+"""Mock Language Model for testing"""
+
+from typing import Callable, Dict, List, Optional, Union
+
+import langroid.language_models as lm
+from langroid.language_models import LLMResponse
+from langroid.language_models.base import LanguageModel, LLMConfig
+
+
+def none_fn(x: str) -> None | str:
+    return None
+
+
+class MockLMConfig(LLMConfig):
+    """
+    Mock Language Model Configuration.
+
+    Attributes:
+        response_dict (Dict[str, str]): A "response rule-book", in the form of a
+            dictionary; if last msg in dialog is x,then respond with response_dict[x]
+    """
+
+    response_dict: Dict[str, str] = {}
+    response_fn: Callable[[str], None | str] = none_fn
+    default_response: str = "Mock response"
+
+    type: str = "mock"
+
+
+class MockLM(LanguageModel):
+
+    def __init__(self, config: MockLMConfig = MockLMConfig()):
+        super().__init__(config)
+        self.config: MockLMConfig = config
+
+    def _response(self, msg: str) -> LLMResponse:
+        # response is based on this fallback order:
+        # - response_dict
+        # - response_fn
+        # - default_response
+        return lm.LLMResponse(
+            message=self.config.response_dict.get(
+                msg,
+                self.config.response_fn(msg) or self.config.default_response,
+            ),
+            cached=False,
+        )
+
+    def chat(
+        self,
+        messages: Union[str, List[lm.LLMMessage]],
+        max_tokens: int = 200,
+        functions: Optional[List[lm.LLMFunctionSpec]] = None,
+        function_call: str | Dict[str, str] = "auto",
+    ) -> lm.LLMResponse:
+        """
+        Mock chat function for testing
+        """
+        last_msg = messages[-1].content if isinstance(messages, list) else messages
+        return self._response(last_msg)
+
+    async def achat(
+        self,
+        messages: Union[str, List[lm.LLMMessage]],
+        max_tokens: int = 200,
+        functions: Optional[List[lm.LLMFunctionSpec]] = None,
+        function_call: str | Dict[str, str] = "auto",
+    ) -> lm.LLMResponse:
+        """
+        Mock chat function for testing
+        """
+        last_msg = messages[-1].content if isinstance(messages, list) else messages
+        return self._response(last_msg)
+
+    def generate(self, prompt: str, max_tokens: int = 200) -> lm.LLMResponse:
+        """
+        Mock generate function for testing
+        """
+        return self._response(prompt)
+
+    async def agenerate(self, prompt: str, max_tokens: int = 200) -> LLMResponse:
+        """
+        Mock generate function for testing
+        """
+        return self._response(prompt)
+
+    def get_stream(self) -> bool:
+        return False
+
+    def set_stream(self, stream: bool) -> bool:
+        return False

langroid/language_models/utils.py

@@ -62,7 +62,7 @@ def retry_with_exponential_backoff(
                 if num_retries > max_retries:
                     raise Exception(
                         f"Maximum number of retries ({max_retries}) exceeded."
-                        f" Last error: {e}."
+                        f" Last error: {str(e)}."
                     )
 
                 # Increment the delay
@@ -128,6 +128,7 @@ def async_retry_with_exponential_backoff(
                 if num_retries > max_retries:
                     raise Exception(
                         f"Maximum number of retries ({max_retries}) exceeded."
+                        f" Last error: {str(e)}."
                     )
 
                 # Increment the delay

langroid/mytypes.py

@@ -1,10 +1,9 @@
-import hashlib
-import uuid
 from enum import Enum
 from textwrap import dedent
 from typing import Any, Callable, Dict, List, Union
+from uuid import uuid4
 
-from langroid.pydantic_v1 import BaseModel, Extra
+from langroid.pydantic_v1 import BaseModel, Extra, Field
 
 Number = Union[int, float]
 Embedding = List[Number]
@@ -40,7 +39,7 @@ class DocMetaData(BaseModel):
 
     source: str = "context"
     is_chunk: bool = False  # if it is a chunk, don't split
-    id: str = ""  # unique id for the document
+    id: str = Field(default_factory=lambda: str(uuid4()))
     window_ids: List[str] = []  # for RAG: ids of chunks around this one
 
     def dict_bool_int(self, *args: Any, **kwargs: Any) -> Dict[str, Any]:
@@ -67,41 +66,11 @@ class Document(BaseModel):
     content: str
     metadata: DocMetaData
 
-    @staticmethod
-    def hash_id(doc: str) -> str:
-        # Encode the document as UTF-8
-        doc_utf8 = str(doc).encode("utf-8")
-
-        # Create a SHA256 hash object
-        sha256_hash = hashlib.sha256()
-
-        # Update the hash object with the bytes of the document
-        sha256_hash.update(doc_utf8)
-
-        # Get the hexadecimal representation of the hash
-        hash_hex = sha256_hash.hexdigest()
-
-        # Convert the first part of the hash to a UUID
-        hash_uuid = uuid.UUID(hash_hex[:32])
-
-        return str(hash_uuid)
-
-    def _unique_hash_id(self) -> str:
-        return self.hash_id(str(self))
-
     def id(self) -> str:
-        if (
-            hasattr(self.metadata, "id")
-            and self.metadata.id is not None
-            and self.metadata.id != ""
-        ):
-            return self.metadata.id
-        else:
-            return self._unique_hash_id()
+        return self.metadata.id
 
     def __str__(self) -> str:
         # TODO: make metadata a pydantic model to enforce "source"
-        self.metadata.json()
         return dedent(
             f"""
         CONTENT: {self.content}         

langroid/parsing/document_parser.py

@@ -8,6 +8,7 @@ from io import BytesIO
 from typing import TYPE_CHECKING, Any, Generator, List, Tuple
 
 from langroid.exceptions import LangroidImportError
+from langroid.utils.object_registry import ObjectRegistry
 
 try:
     import fitz
@@ -341,6 +342,8 @@ class DocumentParser(Parser):
         split = []  # tokens in curr split
         pages: List[str] = []
         docs: List[Document] = []
+        # metadata.id to be shared by ALL chunks of this document
+        common_id = ObjectRegistry.new_id()
         for i, page in self.iterate_pages():
             page_text = self.extract_text_from_page(page)
             split += self.tokenizer.encode(page_text)
@@ -358,6 +361,7 @@ class DocumentParser(Parser):
                         metadata=DocMetaData(
                             source=f"{self.source} pages {pg}",
                             is_chunk=True,
+                            id=common_id,
                         ),
                     )
                 )
@@ -372,6 +376,7 @@ class DocumentParser(Parser):
                     metadata=DocMetaData(
                         source=f"{self.source} pages {pg}",
                         is_chunk=True,
+                        id=common_id,
                     ),
                 )
             )

langroid/parsing/parser.py

@@ -7,6 +7,7 @@ import tiktoken
 from langroid.mytypes import Document
 from langroid.parsing.para_sentence_split import create_chunks, remove_extra_whitespace
 from langroid.pydantic_v1 import BaseSettings
+from langroid.utils.object_registry import ObjectRegistry
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.WARNING)
@@ -75,11 +76,13 @@ class Parser:
             return
         # The original metadata.id (if any) is ignored since it will be same for all
         # chunks and is useless. We want a distinct id for each chunk.
+        # ASSUMPTION: all chunks c of a doc have same c.metadata.id !
         orig_ids = [c.metadata.id for c in chunks]
-        ids = [Document.hash_id(str(c)) for c in chunks]
+        ids = [ObjectRegistry.new_id() for c in chunks]
         id2chunk = {id: c for id, c in zip(ids, chunks)}
 
         # group the ids by orig_id
+        # (each distinct orig_id refers to a different document)
         orig_id_to_ids: Dict[str, List[str]] = {}
         for orig_id, id in zip(orig_ids, ids):
             if orig_id not in orig_id_to_ids:
@@ -108,6 +111,10 @@ class Parser:
             if d.content.strip() == "":
                 continue
             chunks = remove_extra_whitespace(d.content).split(self.config.separators[0])
+            # note we are ensuring we COPY the document metadata into each chunk,
+            # which ensures all chunks of a given doc have same metadata
+            # (and in particular same metadata.id, which is important later for
+            # add_window_ids)
             chunk_docs = [
                 Document(
                     content=c, metadata=d.metadata.copy(update=dict(is_chunk=True))
@@ -156,6 +163,10 @@ class Parser:
             if d.content.strip() == "":
                 continue
             chunks = create_chunks(d.content, self.config.chunk_size, self.num_tokens)
+            # note we are ensuring we COPY the document metadata into each chunk,
+            # which ensures all chunks of a given doc have same metadata
+            # (and in particular same metadata.id, which is important later for
+            # add_window_ids)
             chunk_docs = [
                 Document(
                     content=c, metadata=d.metadata.copy(update=dict(is_chunk=True))
@@ -171,6 +182,10 @@ class Parser:
         final_docs = []
         for d in docs:
             chunks = self.chunk_tokens(d.content)
+            # note we are ensuring we COPY the document metadata into each chunk,
+            # which ensures all chunks of a given doc have same metadata
+            # (and in particular same metadata.id, which is important later for
+            # add_window_ids)
             chunk_docs = [
                 Document(
                     content=c, metadata=d.metadata.copy(update=dict(is_chunk=True))
@@ -274,7 +289,7 @@ class Parser:
         # we need this to distinguish docs later in add_window_ids
         for d in docs:
             if d.metadata.id in [None, ""]:
-                d.metadata.id = d._unique_hash_id()
+                d.metadata.id = ObjectRegistry.new_id()
         # some docs are already splits, so don't split them further!
         chunked_docs = [d for d in docs if d.metadata.is_chunk]
         big_docs = [d for d in docs if not d.metadata.is_chunk]

langroid/utils/__init__.py

@@ -5,6 +5,7 @@ from . import logging
 from . import pydantic_utils
 from . import system
 from . import output
+from . import object_registry
 
 __all__ = [
     "configuration",
@@ -14,4 +15,5 @@ __all__ = [
     "pydantic_utils",
     "system",
     "output",
+    "object_registry",
 ]

langroid/utils/constants.py

@@ -13,10 +13,11 @@ class Colors(BaseModel):
     RESET: str = "\033[0m"
 
 
-USER_QUIT_STRINGS = ["q", "x", "quit", "exit", "bye"]
 NO_ANSWER = "DO-NOT-KNOW"
 DONE = "DONE"
+USER_QUIT_STRINGS = ["q", "x", "quit", "exit", "bye", DONE]
 PASS = "__PASS__"
 PASS_TO = PASS + ":"
 SEND_TO = "SEND:"
 TOOL = "TOOL"
+AT = "@"

langroid/utils/object_registry.py

@@ -0,0 +1,66 @@
+import time
+from typing import TYPE_CHECKING, Dict, Optional, TypeAlias, TypeVar
+from uuid import uuid4
+
+from langroid.pydantic_v1 import BaseModel
+
+if TYPE_CHECKING:
+    from langroid.agent.base import Agent
+    from langroid.agent.chat_agent import ChatAgent
+    from langroid.agent.chat_document import ChatDocument
+
+    # any derivative of BaseModel that has an id() method or an id attribute
+    ObjWithId: TypeAlias = ChatDocument | ChatAgent | Agent
+else:
+    ObjWithId = BaseModel
+
+# Define a type variable that can be any subclass of BaseModel
+T = TypeVar("T", bound=BaseModel)
+
+
+class ObjectRegistry:
+    """A global registry to hold id -> object mappings."""
+
+    registry: Dict[str, ObjWithId] = {}
+
+    @classmethod
+    def add(cls, obj: ObjWithId) -> str:
+        """Adds an object to the registry, returning the object's ID."""
+        object_id = obj.id() if callable(obj.id) else obj.id
+        cls.registry[object_id] = obj
+        return object_id
+
+    @classmethod
+    def get(cls, obj_id: str) -> Optional[ObjWithId]:
+        """Retrieves an object by ID if it still exists."""
+        return cls.registry.get(obj_id)
+
+    @classmethod
+    def register_object(cls, obj: ObjWithId) -> str:
+        """Registers an object in the registry, returning the object's ID."""
+        return cls.add(obj)
+
+    @classmethod
+    def remove(cls, obj_id: str) -> None:
+        """Removes an object from the registry."""
+        if obj_id in cls.registry:
+            del cls.registry[obj_id]
+
+    @classmethod
+    def cleanup(cls) -> None:
+        """Cleans up the registry by removing entries where the object is None."""
+        to_remove = [key for key, value in cls.registry.items() if value is None]
+        for key in to_remove:
+            del cls.registry[key]
+
+    @staticmethod
+    def new_id() -> str:
+        """Generates a new unique ID."""
+        return str(uuid4())
+
+
+def scheduled_cleanup(interval: int = 600) -> None:
+    """Periodically cleans up the global registry every 'interval' seconds."""
+    while True:
+        ObjectRegistry.cleanup()
+        time.sleep(interval)

langroid/utils/system.py

@@ -181,5 +181,4 @@ def hash(s: str) -> str:
 
 def generate_unique_id() -> str:
     """Generate a unique ID using UUID4."""
-    unique_id = str(uuid.uuid4())
-    return unique_id
+    return str(uuid.uuid4())

langroid/vector_store/base.py

@@ -12,6 +12,7 @@ from langroid.mytypes import Document
 from langroid.pydantic_v1 import BaseSettings
 from langroid.utils.algorithms.graph import components, topological_sort
 from langroid.utils.configuration import settings
+from langroid.utils.object_registry import ObjectRegistry
 from langroid.utils.output.printing import print_long_text
 from langroid.utils.pandas_utils import stringify
 
@@ -163,7 +164,7 @@ class VectorStore(ABC):
         vecdbs don't like having blank ids."""
         for d in documents:
             if d.metadata.id in [None, ""]:
-                d.metadata.id = d._unique_hash_id()
+                d.metadata.id = ObjectRegistry.new_id()
 
     @abstractmethod
     def similar_texts_with_scores(
@@ -254,7 +255,7 @@ class VectorStore(ABC):
                 metadata=metadata,
             )
             # make a fresh id since content is in general different
-            document.metadata.id = document.hash_id(document.content)
+            document.metadata.id = ObjectRegistry.new_id()
             final_docs += [document]
             final_scores += [max(id2max_score[id] for id in w)]
         return list(zip(final_docs, final_scores))