langroid 0.1.139__py3-none-any.whl → 0.1.219__py3-none-any.whl

langroid/agent/chat_agent.py

@@ -1,5 +1,5 @@
+import copy
 import inspect
-import json
 import logging
 import textwrap
 from contextlib import ExitStack
@@ -7,8 +7,9 @@ from typing import Dict, List, Optional, Set, Tuple, Type, cast
 
 from rich import print
 from rich.console import Console
+from rich.markup import escape
 
-from langroid.agent.base import Agent, AgentConfig
+from langroid.agent.base import Agent, AgentConfig, noop_fn
 from langroid.agent.chat_document import ChatDocument
 from langroid.agent.tool_message import ToolMessage
 from langroid.language_models.base import (
@@ -19,8 +20,9 @@ from langroid.language_models.base import (
 )
 from langroid.language_models.openai_gpt import OpenAIGPT
 from langroid.utils.configuration import settings
+from langroid.utils.output import status
 
-console = Console(quiet=settings.quiet)
+console = Console()
 
 logger = logging.getLogger(__name__)
 
@@ -41,18 +43,36 @@ class ChatAgentConfig(AgentConfig):
 
     system_message: str = "You are a helpful assistant."
     user_message: Optional[str] = None
-    use_tools: bool = True
-    use_functions_api: bool = False
+    use_tools: bool = False
+    use_functions_api: bool = True
 
-    def _switch_fn_to_tools(self) -> None:
+    def _set_fn_or_tools(self, fn_available: bool) -> None:
         """
-        Switch to using our own ToolMessage mechanism,
-        in case the LLM is not an OpenAI model.
+        Enable Langroid Tool or OpenAI-like fn-calling,
+        depending on config settings and availability of fn-calling.
         """
-        if not self.use_functions_api:
+        if self.use_functions_api and not fn_available:
+            logger.debug(
+                """
+                You have enabled `use_functions_api` but the LLM does not support it.
+                So we will enable `use_tools` instead, so we can use 
+                Langroid's ToolMessage mechanism.
+                """
+            )
+            self.use_functions_api = False
+            self.use_tools = True
+
+        if not self.use_functions_api or not self.use_tools:
             return
-        self.use_functions_api = False
-        self.use_tools = True
+        if self.use_functions_api and self.use_tools:
+            logger.debug(
+                """
+                You have enabled both `use_tools` and `use_functions_api`.
+                Turning off `use_tools`, since the LLM supports function-calling.
+                """
+            )
+            self.use_tools = False
+            self.use_functions_api = True
 
 
 class ChatAgent(Agent):
@@ -72,7 +92,9 @@ class ChatAgent(Agent):
     """
 
     def __init__(
-        self, config: ChatAgentConfig, task: Optional[List[LLMMessage]] = None
+        self,
+        config: ChatAgentConfig = ChatAgentConfig(),
+        task: Optional[List[LLMMessage]] = None,
     ):
         """
         Chat-mode agent initialized with task spec as the initial message sequence
@@ -82,23 +104,7 @@ class ChatAgent(Agent):
         """
         super().__init__(config)
         self.config: ChatAgentConfig = config
-        if (
-            self.llm is not None
-            and (
-                not isinstance(self.llm, OpenAIGPT)
-                or not self.llm.is_openai_chat_model()
-            )
-            and self.config.use_functions_api
-        ):
-            # for non-OpenAI models, use Langroid Tool instead of Function-calling
-            logger.warning(
-                f"""
-                Function calling not available for {self.llm.config.chat_model},
-                switching to Langroid Tools instead.
-                """
-            )
-            self.config._switch_fn_to_tools()
-
+        self.config._set_fn_or_tools(self._fn_call_available())
         self.message_history: List[LLMMessage] = []
         self.tool_instructions_added: bool = False
         # An agent's "task" is defined by a system msg and an optional user msg;
@@ -130,8 +136,42 @@ class ChatAgent(Agent):
         self.llm_functions_usable: Set[str] = set()
         self.llm_function_force: Optional[Dict[str, str]] = None
 
+    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
+        changed in possibly many ways. Below is an imperfect solution. Caution advised.
+        Revisit later.
+        """
+        agent_cls = type(self)
+        config_copy = copy.deepcopy(self.config)
+        config_copy.name = f"{config_copy.name}-{i}"
+        new_agent = agent_cls(config_copy)
+        new_agent.system_tool_instructions = self.system_tool_instructions
+        new_agent.system_json_tool_instructions = self.system_json_tool_instructions
+        new_agent.llm_tools_map = self.llm_tools_map
+        new_agent.llm_functions_map = self.llm_functions_map
+        new_agent.llm_functions_handled = self.llm_functions_handled
+        new_agent.llm_functions_usable = self.llm_functions_usable
+        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
+        return new_agent
+
+    def _fn_call_available(self) -> bool:
+        """Does this agent's LLM support function calling?"""
+        return (
+            self.llm is not None
+            and isinstance(self.llm, OpenAIGPT)
+            and self.llm.is_openai_chat_model()
+        )
+
     def set_system_message(self, msg: str) -> None:
         self.system_message = msg
+        if len(self.message_history) > 0:
+            # if there is message history, update the system message in it
+            self.message_history[0].content = msg
 
     def set_user_message(self, msg: str) -> None:
         self.user_message = msg
@@ -188,46 +228,24 @@ class ChatAgent(Agent):
         enabled_classes: List[Type[ToolMessage]] = list(self.llm_tools_map.values())
         if len(enabled_classes) == 0:
             return "You can ask questions in natural language."
-
         json_instructions = "\n\n".join(
             [
-                textwrap.dedent(
-                    f"""
-                TOOL: {msg_cls.default_value("request")}
-                PURPOSE: {msg_cls.default_value("purpose")} 
-                JSON FORMAT: {
-                    json.dumps(
-                        msg_cls.llm_function_schema(request=True).parameters,
-                        indent=4,
-                    )
-                }
-                {"EXAMPLE: " + msg_cls.usage_example() if msg_cls.examples() else ""}
-                """.lstrip()
-                )
-                for i, msg_cls in enumerate(enabled_classes)
+                msg_cls.json_instructions(tool=self.config.use_tools)
+                for _, msg_cls in enumerate(enabled_classes)
                 if msg_cls.default_value("request") in self.llm_tools_usable
             ]
         )
-        return textwrap.dedent(
-            f"""
-            === ALL AVAILABLE TOOLS and THEIR JSON FORMAT INSTRUCTIONS ===
-            You have access to the following TOOLS to accomplish your task:
-
-            {json_instructions}
-            
-            When one of the above TOOLs is applicable, you must express your 
-            request as "TOOL:" followed by the request in the above JSON format.
-            """
-            + """
-            The JSON format will be:
-                \\{
-                    "request": "<tool_name>",
-                    "<arg1>": <value1>,
-                    "<arg2>": <value2>,
-                    ...
-                \\}             
-            ----------------------------
-            """.lstrip()
+        # if any of the enabled classes has json_group_instructions, then use that,
+        # else fall back to ToolMessage.json_group_instructions
+        for msg_cls in enabled_classes:
+            if hasattr(msg_cls, "json_group_instructions") and callable(
+                getattr(msg_cls, "json_group_instructions")
+            ):
+                return msg_cls.json_group_instructions().format(
+                    json_instructions=json_instructions
+                )
+        return ToolMessage.json_group_instructions().format(
+            json_instructions=json_instructions
         )
 
     def tool_instructions(self) -> str:
@@ -301,7 +319,7 @@ class ChatAgent(Agent):
         Useful when we want to replace a long user prompt, that may contain context
         documents plus a question, with just the question.
         Args:
-            message (str): user message
+            message (str): new message to replace with
             role (str): role of message to replace
         """
         if len(self.message_history) == 0:
@@ -337,7 +355,8 @@ class ChatAgent(Agent):
             
             """.lstrip()
         )
-        return LLMMessage(role=Role.SYSTEM, content=content)
+        # remove leading and trailing newlines and other whitespace
+        return LLMMessage(role=Role.SYSTEM, content=content.strip())
 
     def enable_message(
         self,
@@ -443,10 +462,10 @@ class ChatAgent(Agent):
             message_class: The only ToolMessage class to allow
         """
         request = message_class.__fields__["request"].default
-        for r in self.llm_functions_usable:
-            if r != request:
-                self.llm_tools_usable.discard(r)
-                self.llm_functions_usable.discard(r)
+        to_remove = [r for r in self.llm_tools_usable if r != request]
+        for r in to_remove:
+            self.llm_tools_usable.discard(r)
+            self.llm_functions_usable.discard(r)
 
     def llm_response(
         self, message: Optional[str | ChatDocument] = None
@@ -463,6 +482,8 @@ class ChatAgent(Agent):
         if self.llm is None:
             return None
         hist, output_len = self._prep_llm_messages(message)
+        if len(hist) == 0:
+            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
@@ -472,9 +493,7 @@ class ChatAgent(Agent):
         response.metadata.tool_ids = (
             []
             if isinstance(message, str)
-            else message.metadata.tool_ids
-            if message is not None
-            else []
+            else message.metadata.tool_ids if message is not None else []
         )
         return response
 
@@ -497,9 +516,7 @@ class ChatAgent(Agent):
         response.metadata.tool_ids = (
             []
             if isinstance(message, str)
-            else message.metadata.tool_ids
-            if message is not None
-            else []
+            else message.metadata.tool_ids if message is not None else []
         )
         return response
 
@@ -546,8 +563,9 @@ class ChatAgent(Agent):
             if settings.debug:
                 print(
                     f"""
-                [red]LLM Initial Msg History:
-                {self.message_history_str()}
+                [grey37]LLM Initial Msg History:
+                {escape(self.message_history_str())}
+                [/grey37]
                 """
                 )
         else:
@@ -587,7 +605,10 @@ class ChatAgent(Agent):
                         raise ValueError(
                             """
                         The message history is longer than the max chat context 
-                        length allowed, and we have run out of messages to drop."""
+                        length allowed, and we have run out of messages to drop.
+                        HINT: In your `OpenAIGPTConfig` object, try increasing
+                        `chat_context_length` or decreasing `max_output_tokens`.
+                        """
                         )
                     # drop the second message, i.e. first msg after the sys msg
                     # (typically user msg).
@@ -655,10 +676,17 @@ class ChatAgent(Agent):
         """
         assert self.config.llm is not None and self.llm is not None
         output_len = output_len or self.config.llm.max_output_tokens
+        streamer = noop_fn
+        if self.llm.get_stream():
+            streamer = self.callbacks.start_llm_stream()
+        self.llm.config.streamer = streamer
         with ExitStack() as stack:  # for conditionally using rich spinner
-            if not self.llm.get_stream() and not settings.quiet:
+            if not self.llm.get_stream():
                 # show rich spinner only if not streaming!
-                cm = console.status("LLM responding to messages...")
+                cm = status(
+                    "LLM responding to messages...",
+                    log_if_quiet=False,
+                )
                 stack.enter_context(cm)
             if self.llm.get_stream() and not settings.quiet:
                 console.print(f"[green]{self.indent}", end="")
@@ -670,17 +698,31 @@ class ChatAgent(Agent):
                 functions=functions,
                 function_call=fun_call,
             )
+        if self.llm.get_stream():
+            self.callbacks.finish_llm_stream(
+                content=str(response),
+                is_tool=self.has_tool_message_attempt(
+                    ChatDocument.from_LLMResponse(response, displayed=True)
+                ),
+            )
+        self.llm.config.streamer = noop_fn
+        if response.cached:
+            self.callbacks.cancel_llm_stream()
+
         if not self.llm.get_stream() or response.cached:
             # We would have already displayed the msg "live" ONLY if
             # streaming was enabled, AND we did not find a cached response.
             # If we are here, it means the response has not yet been displayed.
             cached = f"[red]{self.indent}(cached)[/red]" if response.cached else ""
-            if response.function_call is not None:
-                response_str = str(response.function_call)
-            else:
-                response_str = response.message
             if not settings.quiet:
-                print(cached + "[green]" + response_str)
+                print(cached + "[green]" + escape(str(response)))
+                self.callbacks.show_llm_response(
+                    content=str(response),
+                    is_tool=self.has_tool_message_attempt(
+                        ChatDocument.from_LLMResponse(response, displayed=True)
+                    ),
+                    cached=response.cached,
+                )
         self.update_token_usage(
             response,
             messages,
@@ -706,24 +748,42 @@ class ChatAgent(Agent):
                 "auto" if self.llm_function_force is None else self.llm_function_force
             )
         assert self.llm is not None
+
+        streamer = noop_fn
+        if self.llm.get_stream():
+            streamer = self.callbacks.start_llm_stream()
+        self.llm.config.streamer = streamer
+
         response = await self.llm.achat(
             messages,
             output_len,
             functions=functions,
             function_call=fun_call,
         )
-
+        if self.llm.get_stream():
+            self.callbacks.finish_llm_stream(
+                content=str(response),
+                is_tool=self.has_tool_message_attempt(
+                    ChatDocument.from_LLMResponse(response, displayed=True)
+                ),
+            )
+        self.llm.config.streamer = noop_fn
+        if response.cached:
+            self.callbacks.cancel_llm_stream()
         if not self.llm.get_stream() or response.cached:
             # We would have already displayed the msg "live" ONLY if
             # streaming was enabled, AND we did not find a cached response.
             # If we are here, it means the response has not yet been displayed.
             cached = f"[red]{self.indent}(cached)[/red]" if response.cached else ""
-            if response.function_call is not None:
-                response_str = str(response.function_call)
-            else:
-                response_str = response.message
             if not settings.quiet:
-                print(cached + "[green]" + response_str)
+                print(cached + "[green]" + escape(str(response)))
+                self.callbacks.show_llm_response(
+                    content=str(response),
+                    is_tool=self.has_tool_message_attempt(
+                        ChatDocument.from_LLMResponse(response, displayed=True)
+                    ),
+                    cached=response.cached,
+                )
 
         self.update_token_usage(
             response,
@@ -785,12 +845,14 @@ class ChatAgent(Agent):
         """
         # explicitly call THIS class's respond method,
         # not a derived class's (or else there would be infinite recursion!)
+        n_msgs = len(self.message_history)
         with StreamingIfAllowed(self.llm, self.llm.get_stream()):  # type: ignore
             response = cast(ChatDocument, ChatAgent.llm_response(self, message))
-        # clear the last two messages, which are the
-        # user message and the assistant response
-        self.message_history.pop()
-        self.message_history.pop()
+        # If there is a response, then we will have two additional
+        # messages in the message history, i.e. the user message and the
+        # assistant response. We want to (carefully) remove these two messages.
+        self.message_history.pop() if len(self.message_history) > n_msgs else None
+        self.message_history.pop() if len(self.message_history) > n_msgs else None
         return response
 
     async def llm_response_forget_async(self, message: str) -> ChatDocument:
@@ -799,14 +861,16 @@ class ChatAgent(Agent):
         """
         # explicitly call THIS class's respond method,
         # not a derived class's (or else there would be infinite recursion!)
+        n_msgs = len(self.message_history)
         with StreamingIfAllowed(self.llm, self.llm.get_stream()):  # type: ignore
             response = cast(
                 ChatDocument, await ChatAgent.llm_response_async(self, message)
             )
-        # clear the last two messages, which are the
-        # user message and the assistant response
-        self.message_history.pop()
-        self.message_history.pop()
+        # If there is a response, then we will have two additional
+        # messages in the message history, i.e. the user message and the
+        # assistant response. We want to (carefully) remove these two messages.
+        self.message_history.pop() if len(self.message_history) > n_msgs else None
+        self.message_history.pop() if len(self.message_history) > n_msgs else None
         return response
 
     def chat_num_tokens(self, messages: Optional[List[LLMMessage]] = None) -> int:

langroid/agent/chat_document.py

@@ -3,6 +3,7 @@ from typing import List, Optional, Union
 
 from pydantic import BaseModel, Extra
 
+from langroid.agent.tool_message import ToolMessage
 from langroid.language_models.base import (
     LLMFunctionCall,
     LLMMessage,
@@ -12,7 +13,7 @@ from langroid.language_models.base import (
 )
 from langroid.mytypes import DocMetaData, Document, Entity
 from langroid.parsing.agent_chats import parse_message
-from langroid.parsing.json import extract_top_level_json, top_level_json_field
+from langroid.parsing.parse_json import extract_top_level_json, top_level_json_field
 from langroid.utils.output.printing import shorten_text
 
 
@@ -53,6 +54,7 @@ class ChatDocLoggerFields(BaseModel):
 
 class ChatDocument(Document):
     function_call: Optional[LLMFunctionCall] = None
+    tool_messages: List[ToolMessage] = []
     metadata: ChatDocMetaData
     attachment: None | ChatDocAttachment = None
 
@@ -82,7 +84,7 @@ class ChatDocument(Document):
             json_data = json.loads(j)
             tool = json_data.get("request")
             if tool is not None:
-                tools.append(tool)
+                tools.append(str(tool))
         return tools
 
     def log_fields(self) -> ChatDocLoggerFields:
@@ -103,6 +105,8 @@ class ChatDocument(Document):
         content = self.content
         sender_entity = self.metadata.sender
         sender_name = self.metadata.sender_name
+        if tool_type == "FUNC":
+            content += str(self.function_call)
         return ChatDocLoggerFields(
             sender_entity=sender_entity,
             sender_name=sender_name,
@@ -140,6 +144,9 @@ class ChatDocument(Document):
             ChatDocument: ChatDocument representation of this LLMResponse.
         """
         recipient, message = response.get_recipient_and_message()
+        message = message.strip()
+        if message in ["''", '""']:
+            message = ""
         return ChatDocument(
             content=message,
             function_call=response.function_call,
@@ -188,6 +195,16 @@ class ChatDocument(Document):
         if isinstance(message, ChatDocument):
             content = message.content
             fun_call = message.function_call
+            if message.metadata.sender == Entity.USER and fun_call is not None:
+                # This may happen when a (parent agent's) LLM generates a
+                # a Function-call, and it ends up being sent to the current task's
+                # LLM (possibly because the function-call is mis-named or has other
+                # issues and couldn't be handled by handler methods).
+                # But a function-call can only be generated by an entity with
+                # Role.ASSISTANT, so we instead put the content of the function-call
+                # in the content of the message.
+                content += " " + str(fun_call)
+                fun_call = None
             sender_name = message.metadata.sender_name
             tool_ids = message.metadata.tool_ids
             tool_id = tool_ids[-1] if len(tool_ids) > 0 else ""

langroid/agent/openai_assistant.py

@@ -8,7 +8,7 @@ from enum import Enum
 from typing import Any, Dict, List, Optional, Tuple, Type, cast, no_type_check
 
 from openai.types.beta import Assistant, Thread
-from openai.types.beta.threads import Run, ThreadMessage
+from openai.types.beta.threads import Message, Run
 from openai.types.beta.threads.runs import RunStep
 from pydantic import BaseModel
 from rich import print
@@ -41,6 +41,8 @@ class AssistantTool(BaseModel):
     def dct(self) -> Dict[str, Any]:
         d = super().dict()
         d["type"] = d["type"].value
+        if self.type != ToolType.FUNCTION:
+            d.pop("function")
         return d
 
 
@@ -78,7 +80,6 @@ class OpenAIAssistantConfig(ChatAgentConfig):
 
 
 class OpenAIAssistant(ChatAgent):
-
     """
     A ChatAgent powered by OpenAI Assistant API:
     mainly, in `llm_response` method, we avoid maintaining conversation state,
@@ -257,22 +258,22 @@ class OpenAIAssistant(ChatAgent):
         self.llm.cache.store(assistant_key, self.assistant.id)
 
     @staticmethod
-    def thread_msg_to_llm_msg(msg: ThreadMessage) -> LLMMessage:
+    def thread_msg_to_llm_msg(msg: Message) -> LLMMessage:
         """
-        Convert a ThreadMessage to an LLMMessage
+        Convert a Message to an LLMMessage
         """
         return LLMMessage(
             content=msg.content[0].text.value,  # type: ignore
-            role=msg.role,
+            role=Role(msg.role),
         )
 
-    def _update_messages_hash(self, msg: ThreadMessage | LLMMessage) -> None:
+    def _update_messages_hash(self, msg: Message | LLMMessage) -> None:
         """
         Update the hash-state in the thread with the given message.
         """
         if self.thread is None:
             raise ValueError("Thread is None")
-        if isinstance(msg, ThreadMessage):
+        if isinstance(msg, Message):
             llm_msg = self.thread_msg_to_llm_msg(msg)
         else:
             llm_msg = msg
@@ -491,7 +492,7 @@ class OpenAIAssistant(ChatAgent):
             LLMMessage(
                 # TODO: could be image, deal with it later
                 content=m.content[0].text.value,  # type: ignore
-                role=m.role,
+                role=Role(m.role),
             )
             for m in thread_msgs
         ]
@@ -646,7 +647,7 @@ class OpenAIAssistant(ChatAgent):
             tool_outputs=tool_outputs,  # type: ignore
         )
 
-    def process_citations(self, thread_msg: ThreadMessage) -> None:
+    def process_citations(self, thread_msg: Message) -> None:
         """
         Process citations in the thread message.
         Modifies the thread message in-place.
@@ -664,7 +665,16 @@ class OpenAIAssistant(ChatAgent):
             )
             # Gather citations based on annotation attributes
             if file_citation := getattr(annotation, "file_citation", None):
-                cited_file = self.client.files.retrieve(file_citation.file_id)
+                try:
+                    cited_file = self.client.files.retrieve(file_citation.file_id)
+                except Exception:
+                    logger.warning(
+                        f"""
+                        Could not retrieve cited file with id {file_citation.file_id}, 
+                        ignoring. 
+                        """
+                    )
+                    continue
                 citations.append(
                     f"[{index}] '{file_citation.quote}',-- from {cited_file.filename}"
                 )

langroid/agent/special/__init__.py

@@ -1,27 +1,50 @@
-from .doc_chat_agent import DocChatAgent, DocChatAgentConfig
-from .recipient_validator_agent import (
-    RecipientValidatorConfig,
-    RecipientValidatorAttachment,
-    RecipientValidator,
+from .relevance_extractor_agent import (
+    RelevanceExtractorAgent,
+    RelevanceExtractorAgentConfig,
 )
+from .doc_chat_agent import DocChatAgent, DocChatAgentConfig
 from .retriever_agent import (
     RecordMetadata,
     RecordDoc,
     RetrieverAgentConfig,
     RetrieverAgent,
 )
+from .lance_doc_chat_agent import LanceDocChatAgent
 from .table_chat_agent import (
     dataframe_summary,
     TableChatAgent,
     TableChatAgentConfig,
     RunCodeTool,
 )
-from .relevance_extractor_agent import (
-    RelevanceExtractorAgent,
-    RelevanceExtractorAgentConfig,
-)
 from . import sql
+from . import relevance_extractor_agent
 from . import doc_chat_agent
-from . import recipient_validator_agent
 from . import retriever_agent
+from . import lance_tools
+from . import lance_doc_chat_agent
+from . import lance_rag
 from . import table_chat_agent
+
+__all__ = [
+    "RelevanceExtractorAgent",
+    "RelevanceExtractorAgentConfig",
+    "DocChatAgent",
+    "DocChatAgentConfig",
+    "RecordMetadata",
+    "RecordDoc",
+    "RetrieverAgentConfig",
+    "RetrieverAgent",
+    "LanceDocChatAgent",
+    "dataframe_summary",
+    "TableChatAgent",
+    "TableChatAgentConfig",
+    "RunCodeTool",
+    "sql",
+    "relevance_extractor_agent",
+    "doc_chat_agent",
+    "retriever_agent",
+    "lance_tools",
+    "lance_doc_chat_agent",
+    "lance_rag",
+    "table_chat_agent",
+]