Merge pull request #186 from VRSEN/dev/refactor-send-message-tool

Allow Modifying Send Message Tool

agency_swarm/threads/thread_async.py

@@ -90,11 +90,10 @@ def check_status(self, run=None):
         return f"""{self.recipient_agent.name}'s Response: '{messages.data[0].content[0].text.value}'"""
 
     def get_last_run(self):
-        if not self.thread:
-            self.init_thread()
+        self.init_thread()
 
         runs = self.client.beta.threads.runs.list(
-            thread_id=self.thread.id,
+            thread_id=self.id,
             order="desc",
         )
 

agency_swarm/tools/BaseTool.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Any, ClassVar
+from typing import Any, ClassVar, Literal, Union
 
 from docstring_parser import parse
 
@@ -11,15 +11,31 @@ class BaseTool(BaseModel, ABC):
     _shared_state: ClassVar[SharedState] = None
     _caller_agent: Any = None
     _event_handler: Any = None
+    _tool_call: Any = None
 
     def __init__(self, **kwargs):
         if not self.__class__._shared_state:
             self.__class__._shared_state = SharedState()
         super().__init__(**kwargs)
+
+        # Ensure all ToolConfig variables are initialized
+        config_defaults = {
+            'strict': False,
+            'one_call_at_a_time': False,
+            'output_as_result': False,
+            'async_mode': None
+        }
+
+        for key, value in config_defaults.items():
+            if not hasattr(self.ToolConfig, key):
+                setattr(self.ToolConfig, key, value)
 
     class ToolConfig:
         strict: bool = False
         one_call_at_a_time: bool = False
+        # return the tool output as assistant message
+        output_as_result: bool = False
+        async_mode: Union[Literal["threading"], None] = None
 
     @classmethod
     @property
@@ -76,5 +92,5 @@ def openai_schema(cls):
         return schema
 
     @abstractmethod
-    def run(self, **kwargs):
+    def run(self):
         pass

agency_swarm/tools/__init__.py

@@ -2,4 +2,4 @@
 from .oai.CodeInterpreter import CodeInterpreter
 from .oai.FileSearch import FileSearch
 from .oai.Retrieval import Retrieval
-from .ToolFactory import ToolFactory
+from .ToolFactory import ToolFactory

agency_swarm/tools/send_message/SendMessage.py

@@ -0,0 +1,44 @@
+from typing import Optional, List
+from pydantic import Field, field_validator, model_validator
+from .SendMessageBase import SendMessageBase
+
+class SendMessage(SendMessageBase):
+    """Use this tool to facilitate direct, synchronous communication between specialized agents within your agency. When you send a message using this tool, you receive a response exclusively from the designated recipient agent. To continue the dialogue, invoke this tool again with the desired recipient agent and your follow-up message. Remember, communication here is synchronous; the recipient agent won't perform any tasks post-response. You are responsible for relaying the recipient agent's responses back to the user, as the user does not have direct access to these replies. Keep engaging with the tool for continuous interaction until the task is fully resolved. Do not send more than 1 message to the same recipient agent at the same time."""
+    my_primary_instructions: str = Field(
+        ..., 
+        description=(
+            "Please repeat your primary instructions step-by-step, including both completed "
+            "and the following next steps that you need to perform. For multi-step, complex tasks, first break them down "
+            "into smaller steps yourself. Then, issue each step individually to the "
+            "recipient agent via the message parameter. Each identified step should be "
+            "sent in a separate message. Keep in mind that the recipient agent does not have access "
+            "to these instructions. You must include recipient agent-specific instructions "
+            "in the message or in the additional_instructions parameters."
+        )
+    )
+    message: str = Field(
+        ..., 
+        description="Specify the task required for the recipient agent to complete. Focus on clarifying what the task entails, rather than providing exact instructions. Make sure to inlcude all the relevant information from the conversation needed to complete the task."
+    )
+    message_files: Optional[List[str]] = Field(
+        default=None,
+        description="A list of file IDs to be sent as attachments to this message. Only use this if you have the file ID that starts with 'file-'.",
+        examples=["file-1234", "file-5678"]
+    )
+    additional_instructions: Optional[str] = Field(
+        default=None,
+        description="Additional context or instructions from the conversation needed by the recipient agent to complete the task."
+    )
+
+    @model_validator(mode='after')
+    def validate_files(self):
+        # prevent hallucinations with agents sending file IDs into incorrect fields
+        if "file-" in self.message or (self.additional_instructions and "file-" in self.additional_instructions):
+            if not self.message_files:
+                raise ValueError("You must include file IDs in message_files parameter.")
+        return self
+
+    def run(self):
+        return self._get_completion(message=self.message,
+                                    message_files=self.message_files,
+                                    additional_instructions=self.additional_instructions)

agency_swarm/tools/send_message/SendMessageAsyncThreading.py

@@ -0,0 +1,8 @@
+from typing import ClassVar, Type
+from agency_swarm.threads.thread_async import ThreadAsync
+from .SendMessage import SendMessage
+
+class SendMessageAsyncThreading(SendMessage):
+    """Use this tool for asynchronous communication with other agents within your agency. Initiate tasks by messaging, and check status and responses later with the 'GetResponse' tool. Relay responses to the user, who instructs on status checks. Continue until task completion."""
+    class ToolConfig:
+        async_mode = "threading"

agency_swarm/tools/send_message/SendMessageBase.py

@@ -0,0 +1,41 @@
+from agency_swarm.agents.agent import Agent
+from agency_swarm.threads.thread import Thread
+from typing import ClassVar, Union
+from pydantic import Field, field_validator
+from agency_swarm.threads.thread_async import ThreadAsync
+from agency_swarm.tools import BaseTool
+from abc import ABC
+
+class SendMessageBase(BaseTool, ABC):
+    recipient: str = Field(..., description="Recipient agent that you want to send the message to. This field will be overriden inside the agency class.")
+
+    _agents_and_threads: ClassVar = None
+
+    @field_validator('additional_instructions', mode='before', check_fields=False)
+    @classmethod
+    def validate_additional_instructions(cls, value):
+        # previously the parameter was a list, now it's a string
+        # add compatibility for old code
+        if isinstance(value, list):
+            return "\n".join(value)
+        return value
+
+    def _get_thread(self) -> Thread | ThreadAsync:
+        return self._agents_and_threads[self._caller_agent.name][self.recipient.value]
+
+    def _get_main_thread(self) -> Thread | ThreadAsync:
+        return self._agents_and_threads["main_thread"]
+
+    def _get_recipient_agent(self) -> Agent:
+        return self._agents_and_threads[self._caller_agent.name][self.recipient.value].recipient_agent
+
+    def _get_completion(self, message: Union[str, None] = None, **kwargs):
+        thread = self._get_thread()
+
+        if self.ToolConfig.async_mode == "threading":
+            return thread.get_completion_async(message=message, **kwargs)
+        else:
+            return thread.get_completion(message=message, 
+                                        event_handler=self._event_handler,
+                                        yield_messages=not self._event_handler,
+                                        **kwargs)

agency_swarm/tools/send_message/SendMessageQuick.py

@@ -0,0 +1,14 @@
+from agency_swarm.threads.thread import Thread
+from typing import ClassVar, Optional, List, Type
+from pydantic import Field, field_validator, model_validator
+from .SendMessageBase import SendMessageBase
+
+class SendMessageQuick(SendMessageBase):
+    """Use this tool to facilitate direct, synchronous communication between specialized agents within your agency. When you send a message using this tool, you receive a response exclusively from the designated recipient agent. To continue the dialogue, invoke this tool again with the desired recipient agent and your follow-up message. Remember, communication here is synchronous; the recipient agent won't perform any tasks post-response. You are responsible for relaying the recipient agent's responses back to the user, as the user does not have direct access to these replies. Keep engaging with the tool for continuous interaction until the task is fully resolved. Do not send more than 1 message to the same recipient agent at the same time."""
+    message: str = Field(
+        ..., 
+        description="Specify the task required for the recipient agent to complete. Focus on clarifying what the task entails, rather than providing exact instructions. Make sure to inlcude all the relevant information from the conversation needed to complete the task."
+    )
+
+    def run(self):
+        return self._get_completion(message=self.message)

agency_swarm/tools/send_message/SendMessageSwarm.py

@@ -0,0 +1,48 @@
+from openai import BadRequestError
+from agency_swarm.threads.thread import Thread
+from .SendMessage import SendMessageBase
+
+class SendMessageSwarm(SendMessageBase):
+    """Use this tool to route messages to other agents within your agency. After using this tool, you will be switched to the recipient agent. This tool can only be used once per message. Do not use any other tools together with this tool."""
+
+    class ToolConfig:
+        # set output as result because the communication will be finished after this tool is called
+        output_as_result: bool = True
+        one_call_at_a_time: bool = True
+
+    def run(self):            
+        # get main thread
+        thread = self._get_main_thread()
+
+        # get recipient agent from thread
+        recipient_agent = self._get_recipient_agent()
+
+        # submit tool output
+        try:
+            thread._submit_tool_outputs(
+                tool_outputs=[{"tool_call_id": self._tool_call.id, "output": "The request has been routed. You are now a " + recipient_agent.name + " agent. Please assist the user further with their request."}],
+                poll=False
+            )
+        except BadRequestError as e:
+            raise Exception("You can only call this tool by itself. Do not use any other tools together with this tool.")
+
+        try:
+            # cancel run
+            thread._cancel_run()
+
+            # change recipient agent in thread
+            thread.recipient_agent = recipient_agent
+
+            # change recipient agent in gradio dropdown
+            if self._event_handler:
+                if hasattr(self._event_handler, "change_recipient_agent"):
+                    self._event_handler.change_recipient_agent(self.recipient.value)
+
+            # continue conversation with the new recipient agent
+            message = thread.get_completion(message=None, recipient_agent=recipient_agent, yield_messages=not self._event_handler, event_handler=self._event_handler)
+
+            return message or ""
+        except Exception as e:
+            # we need to catch errors beucase tool outputs are already submitted
+            print("Error in SendMessageSwarm: ", e)
+            return str(e)

agency_swarm/tools/send_message/__init__.py

@@ -0,0 +1,5 @@
+from .SendMessageAsyncThreading import SendMessageAsyncThreading
+from .SendMessageBase import SendMessageBase
+from .SendMessage import SendMessage
+from .SendMessageSwarm import SendMessageSwarm
+from .SendMessageQuick import SendMessageQuick