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

langroid/agent/base.py

@@ -40,9 +40,10 @@ from langroid.mytypes import Entity
 from langroid.parsing.parse_json import extract_top_level_json
 from langroid.parsing.parser import Parser, ParsingConfig
 from langroid.prompts.prompts_config import PromptsConfig
-from langroid.pydantic_v1 import BaseSettings, ValidationError, validator
+from langroid.pydantic_v1 import BaseSettings, Field, ValidationError, validator
 from langroid.utils.configuration import settings
 from langroid.utils.constants import NO_ANSWER
+from langroid.utils.object_registry import ObjectRegistry
 from langroid.utils.output import status
 from langroid.vector_store.base import VectorStore, VectorStoreConfig
 
@@ -64,6 +65,7 @@ class AgentConfig(BaseSettings):
     parsing: Optional[ParsingConfig] = ParsingConfig()
     prompts: Optional[PromptsConfig] = PromptsConfig()
     show_stats: bool = True  # show token usage/cost stats?
+    add_to_registry: bool = True  # register agent in ObjectRegistry?
 
     @validator("name")
     def check_name_alphanum(cls, v: str) -> str:
@@ -90,6 +92,8 @@ class Agent(ABC):
     information about any tool/function-calling messages that have been defined.
     """
 
+    id: str = Field(default_factory=lambda: ObjectRegistry.new_id())
+
     def __init__(self, config: AgentConfig = AgentConfig()):
         self.config = config
         self.lock = asyncio.Lock()  # for async access to update self.llm.usage_cost
@@ -114,6 +118,8 @@ class Agent(ABC):
         self.parser: Optional[Parser] = (
             Parser(config.parsing) if config.parsing else None
         )
+        if config.add_to_registry:
+            ObjectRegistry.register_object(self)
 
         self.callbacks = SimpleNamespace(
             start_llm_stream=lambda: noop_fn,
@@ -128,6 +134,14 @@ class Agent(ABC):
             show_start_response=noop_fn,
         )
 
+    @staticmethod
+    def from_id(id: str) -> "Agent":
+        return cast(Agent, ObjectRegistry.get(id))
+
+    @staticmethod
+    def delete_id(id: str) -> None:
+        ObjectRegistry.remove(id)
+
     def entity_responders(
         self,
     ) -> List[

langroid/agent/chat_agent.py

@@ -21,6 +21,7 @@ from langroid.language_models.base import (
 )
 from langroid.language_models.openai_gpt import OpenAIGPT
 from langroid.utils.configuration import settings
+from langroid.utils.object_registry import ObjectRegistry
 from langroid.utils.output import status
 
 console = Console()
@@ -137,11 +138,22 @@ class ChatAgent(Agent):
         self.llm_functions_usable: Set[str] = set()
         self.llm_function_force: Optional[Dict[str, str]] = None
 
+    @staticmethod
+    def from_id(id: str) -> "ChatAgent":
+        """
+        Get an agent from its ID
+        Args:
+            agent_id (str): ID of the agent
+        Returns:
+            ChatAgent: The agent with the given ID
+        """
+        return cast(ChatAgent, Agent.from_id(id))
+
     def clone(self, i: int = 0) -> "ChatAgent":
         """Create i'th clone of this agent, ensuring tool use/handling is cloned.
         Important: We assume all member variables are in the __init__ method here
         and in the Agent class.
-        TODO: We are attempting to close an agent after its state has been
+        TODO: We are attempting to clone an agent after its state has been
         changed in possibly many ways. Below is an imperfect solution. Caution advised.
         Revisit later.
         """
@@ -158,6 +170,9 @@ class ChatAgent(Agent):
         new_agent.llm_function_force = self.llm_function_force
         # Caution - we are copying the vector-db, maybe we don't always want this?
         new_agent.vecdb = self.vecdb
+        new_agent.id = ObjectRegistry.new_id()
+        if self.config.add_to_registry:
+            ObjectRegistry.register_object(new_agent)
         return new_agent
 
     def _fn_call_available(self) -> bool:
@@ -202,6 +217,10 @@ class ChatAgent(Agent):
         if start < 0:
             n = len(self.message_history)
             start = max(0, n + start)
+        dropped = self.message_history[start:]
+        for msg in dropped:
+            # clear out the chat document from the ObjectRegistry
+            ChatDocument.delete_id(msg.chat_document_id)
         self.message_history = self.message_history[:start]
 
     def update_history(self, message: str, response: str) -> None:
@@ -310,10 +329,24 @@ class ChatAgent(Agent):
 
     def last_message_with_role(self, role: Role) -> LLMMessage | None:
         """from `message_history`, return the last message with role `role`"""
-        for i in range(len(self.message_history) - 1, -1, -1):
-            if self.message_history[i].role == role:
-                return self.message_history[i]
-        return None
+        n_role_msgs = len([m for m in self.message_history if m.role == role])
+        if n_role_msgs == 0:
+            return None
+        idx = self.nth_message_idx_with_role(role, n_role_msgs)
+        return self.message_history[idx]
+
+    def nth_message_idx_with_role(self, role: Role, n: int) -> int:
+        """Index of `n`th message in message_history, with specified role.
+        (n is assumed to be 1-based, i.e. 1 is the first message with that role).
+        Return -1 if not found. Index = 0 is the first message in the history.
+        """
+        indices_with_role = [
+            i for i, m in enumerate(self.message_history) if m.role == role
+        ]
+
+        if len(indices_with_role) < n:
+            return -1
+        return indices_with_role[n - 1]
 
     def update_last_message(self, message: str, role: str = Role.USER) -> None:
         """
@@ -488,9 +521,9 @@ class ChatAgent(Agent):
             return None
         with StreamingIfAllowed(self.llm, self.llm.get_stream()):
             response = self.llm_response_messages(hist, output_len)
-        # TODO - when response contains function_call we should include
-        # that (and related fields) in the message_history
         self.message_history.append(ChatDocument.to_LLMMessage(response))
+        response.metadata.msg_idx = len(self.message_history) - 1
+        response.metadata.agent_id = self.id
         # Preserve trail of tool_ids for OpenAI Assistant fn-calls
         response.metadata.tool_ids = (
             []
@@ -511,9 +544,9 @@ class ChatAgent(Agent):
         hist, output_len = self._prep_llm_messages(message)
         with StreamingIfAllowed(self.llm, self.llm.get_stream()):
             response = await self.llm_response_messages_async(hist, output_len)
-        # TODO - when response contains function_call we should include
-        # that (and related fields) in the message_history
         self.message_history.append(ChatDocument.to_LLMMessage(response))
+        response.metadata.msg_idx = len(self.message_history) - 1
+        response.metadata.agent_id = self.id
         # Preserve trail of tool_ids for OpenAI Assistant fn-calls
         response.metadata.tool_ids = (
             []
@@ -522,6 +555,16 @@ class ChatAgent(Agent):
         )
         return response
 
+    def init_message_history(self) -> None:
+        """
+        Initialize the message history with the system message and user message
+        """
+        self.message_history = [self._create_system_and_tools_message()]
+        if self.user_message:
+            self.message_history.append(
+                LLMMessage(role=Role.USER, content=self.user_message)
+            )
+
     def _prep_llm_messages(
         self,
         message: Optional[str | ChatDocument] = None,
@@ -555,11 +598,7 @@ class ChatAgent(Agent):
 
         if len(self.message_history) == 0:
             # initial messages have not yet been loaded, so load them
-            self.message_history = [self._create_system_and_tools_message()]
-            if self.user_message:
-                self.message_history.append(
-                    LLMMessage(role=Role.USER, content=self.user_message)
-                )
+            self.init_message_history()
 
             # for debugging, show the initial message history
             if settings.debug:
@@ -576,8 +615,14 @@ class ChatAgent(Agent):
             self.message_history[0] = self._create_system_and_tools_message()
 
         if message is not None:
-            llm_msg = ChatDocument.to_LLMMessage(message)
-            self.message_history.append(llm_msg)
+            if (
+                isinstance(message, str)
+                or message.id() != self.message_history[-1].chat_document_id
+            ):
+                # either the message is a str, or it is a fresh ChatDocument
+                # different from the last message in the history
+                llm_msg = ChatDocument.to_LLMMessage(message)
+                self.message_history.append(llm_msg)
 
         hist = self.message_history
         output_len = self.config.llm.max_output_tokens
@@ -614,6 +659,7 @@ class ChatAgent(Agent):
                         )
                     # drop the second message, i.e. first msg after the sys msg
                     # (typically user msg).
+                    ChatDocument.delete_id(hist[1].chat_document_id)
                     hist = hist[:1] + hist[2:]
 
                 if len(hist) < len(self.message_history):
@@ -650,6 +696,12 @@ class ChatAgent(Agent):
                 and the response may be truncated.
                 """
             )
+        if isinstance(message, ChatDocument):
+            # record the position of the corresponding LLMMessage in
+            # the message_history
+            message.metadata.msg_idx = len(hist) - 1
+            message.metadata.agent_id = self.id
+
         return hist, output_len
 
     def _function_args(

langroid/agent/chat_document.py

@@ -1,6 +1,9 @@
+from __future__ import annotations
+
+import copy
 import json
 from enum import Enum
-from typing import List, Optional, Union
+from typing import Any, List, Optional, Union, cast
 
 from langroid.agent.tool_message import ToolMessage
 from langroid.language_models.base import (
@@ -14,6 +17,7 @@ from langroid.mytypes import DocMetaData, Document, Entity
 from langroid.parsing.agent_chats import parse_message
 from langroid.parsing.parse_json import extract_top_level_json, top_level_json_field
 from langroid.pydantic_v1 import BaseModel, Extra
+from langroid.utils.object_registry import ObjectRegistry
 from langroid.utils.output.printing import shorten_text
 
 
@@ -41,8 +45,11 @@ class StatusCode(str, Enum):
 
 
 class ChatDocMetaData(DocMetaData):
-    parent: Optional["ChatDocument"] = None
-    sender: Entity
+    parent_id: str = ""  # msg (ChatDocument) to which this is a response
+    child_id: str = ""  # ChatDocument that has response to this message
+    agent_id: str = ""  # ChatAgent that generated this message
+    msg_idx: int = -1  # index of this message in the agent `message_history`
+    sender: Entity  # sender of the message
     tool_ids: List[str] = []  # stack of tool_ids; used by OpenAIAssistant
     block: None | Entity = None
     sender_name: str = ""
@@ -53,6 +60,14 @@ class ChatDocMetaData(DocMetaData):
     has_citation: bool = False
     status: Optional[StatusCode] = None
 
+    @property
+    def parent(self) -> Optional["ChatDocument"]:
+        return ChatDocument.from_id(self.parent_id)
+
+    @property
+    def child(self) -> Optional["ChatDocument"]:
+        return ChatDocument.from_id(self.child_id)
+
 
 class ChatDocLoggerFields(BaseModel):
     sender_entity: Entity = Entity.USER
@@ -75,6 +90,41 @@ class ChatDocument(Document):
     metadata: ChatDocMetaData
     attachment: None | ChatDocAttachment = None
 
+    def __init__(self, **data: Any):
+        super().__init__(**data)
+        ObjectRegistry.register_object(self)
+
+    @property
+    def parent(self) -> Optional["ChatDocument"]:
+        return ChatDocument.from_id(self.metadata.parent_id)
+
+    @property
+    def child(self) -> Optional["ChatDocument"]:
+        return ChatDocument.from_id(self.metadata.child_id)
+
+    @staticmethod
+    def deepcopy(doc: ChatDocument) -> ChatDocument:
+        new_doc = copy.deepcopy(doc)
+        new_doc.metadata.id = ObjectRegistry.new_id()
+        ObjectRegistry.register_object(new_doc)
+        return new_doc
+
+    @staticmethod
+    def from_id(id: str) -> Optional["ChatDocument"]:
+        return cast(ChatDocument, ObjectRegistry.get(id))
+
+    @staticmethod
+    def delete_id(id: str) -> None:
+        """Remove ChatDocument with given id from ObjectRegistry,
+        and all its descendants.
+        """
+        chat_doc = ChatDocument.from_id(id)
+        # first delete all descendants
+        while chat_doc is not None:
+            next_chat_doc = chat_doc.child
+            ObjectRegistry.remove(chat_doc.id())
+            chat_doc = next_chat_doc
+
     def __str__(self) -> str:
         fields = self.log_fields()
         tool_str = ""
@@ -224,6 +274,7 @@ class ChatDocument(Document):
         sender_role = Role.USER
         fun_call = None
         tool_id = ""
+        chat_document_id: str = ""
         if isinstance(message, ChatDocument):
             content = message.content
             fun_call = message.function_call
@@ -240,6 +291,7 @@ class ChatDocument(Document):
             sender_name = message.metadata.sender_name
             tool_ids = message.metadata.tool_ids
             tool_id = tool_ids[-1] if len(tool_ids) > 0 else ""
+            chat_document_id = message.id()
             if message.metadata.sender == Entity.SYSTEM:
                 sender_role = Role.SYSTEM
             if (
@@ -260,7 +312,9 @@ class ChatDocument(Document):
             content=content,
             function_call=fun_call,
             name=sender_name,
+            chat_document_id=chat_document_id,
         )
 
 
+LLMMessage.update_forward_refs()
 ChatDocMetaData.update_forward_refs()

langroid/agent/special/doc_chat_agent.py

@@ -35,7 +35,6 @@ from langroid.embedding_models.models import (
     OpenAIEmbeddingsConfig,
     SentenceTransformerEmbeddingsConfig,
 )
-from langroid.exceptions import LangroidImportError
 from langroid.language_models.base import StreamingIfAllowed
 from langroid.language_models.openai_gpt import OpenAIChatModel, OpenAIGPTConfig
 from langroid.mytypes import DocMetaData, Document, Entity
@@ -54,6 +53,7 @@ from langroid.parsing.utils import batched
 from langroid.prompts.prompts_config import PromptsConfig
 from langroid.prompts.templates import SUMMARY_ANSWER_PROMPT_GPT4
 from langroid.utils.constants import NO_ANSWER
+from langroid.utils.object_registry import ObjectRegistry
 from langroid.utils.output import show_if_debug, status
 from langroid.utils.output.citations import (
     extract_markdown_references,
@@ -101,29 +101,6 @@ oai_embed_config = OpenAIEmbeddingsConfig(
     dims=1536,
 )
 
-vecdb_config: VectorStoreConfig = QdrantDBConfig(
-    collection_name="doc-chat-qdrantdb",
-    replace_collection=True,
-    storage_path=".qdrantdb/data/",
-    embedding=hf_embed_config if has_sentence_transformers else oai_embed_config,
-)
-
-try:
-    import lancedb
-
-    lancedb  # appease mypy
-    from langroid.vector_store.lancedb import LanceDBConfig
-
-    vecdb_config = LanceDBConfig(
-        collection_name="doc-chat-lancedb",
-        replace_collection=True,
-        storage_path=".lancedb/data/",
-        embedding=(hf_embed_config if has_sentence_transformers else oai_embed_config),
-    )
-
-except (ImportError, LangroidImportError):
-    pass
-
 
 class DocChatAgentConfig(ChatAgentConfig):
     system_message: str = DEFAULT_DOC_CHAT_SYSTEM_MESSAGE
@@ -201,7 +178,12 @@ class DocChatAgentConfig(ChatAgentConfig):
     )
 
     # Allow vecdb to be None in case we want to explicitly set it later
-    vecdb: Optional[VectorStoreConfig] = vecdb_config
+    vecdb: Optional[VectorStoreConfig] = QdrantDBConfig(
+        collection_name="doc-chat-qdrantdb",
+        replace_collection=True,
+        storage_path=".qdrantdb/data/",
+        embedding=hf_embed_config if has_sentence_transformers else oai_embed_config,
+    )
 
     llm: OpenAIGPTConfig = OpenAIGPTConfig(
         type="openai",
@@ -414,7 +396,7 @@ class DocChatAgent(ChatAgent):
             raise ValueError("Parser not set")
         for d in docs:
             if d.metadata.id in [None, ""]:
-                d.metadata.id = d._unique_hash_id()
+                d.metadata.id = ObjectRegistry.new_id()
         if split:
             docs = self.parser.split(docs)
         else: