Add new tasks as delegates of the last agent in the project.

Author: tcdent

agentstack/frameworks/openai_swarm.py

@@ -1,16 +1,15 @@
-from typing import TYPE_CHECKING, Optional, Any, Callable
+from typing import Optional, Any, Callable
 from pathlib import Path
 import ast
 from agentstack import conf, log
 from agentstack.exceptions import ValidationError
+from agentstack.generation import InsertionPoint
 from agentstack._tools import ToolConfig
 from agentstack.tasks import TaskConfig
 from agentstack.agents import AgentConfig
 from agentstack.generation import asttools
 from agentstack import graph
 
-if TYPE_CHECKING:
-    from agentstack.generation import InsertionPoint
 
 NAME: str = "OpenAI Swarm"
 ENTRYPOINT: Path = Path('src/stack.py')
@@ -62,19 +61,54 @@ def add_task_method(self, task: TaskConfig):
             pos, _ = self.get_node_range(main_method)
 
         code = f"""    @agentstack.task
-    def {task.name}(self):
+    def {task.name}(self, messages: list[str] = []) -> Agent:
         task_config = agentstack.get_task('{task.name}')
+        agent = getattr(self, task_config.agent)
         messages = [
+            *messages, 
             task_config.prompt, 
         ]
-        agent = getattr(self, task_config.agent)
         return agent(messages)"""
 
         if not self.source[:pos].endswith('\n'):
             code = '\n\n' + code
         if not self.source[pos:].startswith('\n'):
             code += '\n\n'
         self.edit_node_range(pos, pos, code)
+        
+        # add a new task to the last agent in the stack
+        existing_agent_methods = self.get_agent_methods()
+        if not len(existing_agent_methods):
+            return  # no agents to update
+        
+        # add a call to `self._handoff(task_name)` to the front of the update_method's
+        # `function` argument which is a list of functions
+        update_method = existing_agent_methods[-1]
+        try:
+            agent_instance = asttools.find_method_calls(update_method, 'Agent')[0]
+        except IndexError:
+            raise ValidationError(f"Agent method `{update_method.name}` does not instantiate `Agent` in {ENTRYPOINT}")
+
+        existing_agent_tools = asttools.find_kwarg_in_method_call(agent_instance, 'functions')
+        if not existing_agent_tools:
+            raise ValidationError(
+                f"`@agent` method `{update_method.name}` does not have a keyword argument `functions` in {ENTRYPOINT}"
+            )
+        
+        assert isinstance(existing_agent_tools.value, ast.List)
+        existing_elts = existing_agent_tools.value.elts
+        existing_elts.insert(0, ast.Call(
+            func=ast.Attribute(
+                value=ast.Name(id='self', ctx=ast.Load()),
+                attr='_handoff',
+                ctx=ast.Load(),
+            ),
+            args=[ast.Constant(value=task.name)],
+            keywords=[],
+        ))
+        new_node = ast.List(elts=existing_elts, ctx=ast.Load())
+        start, end = self.get_node_range(existing_agent_tools.value)
+        self.edit_node_range(start, end, new_node)
 
     def get_agent_methods(self) -> list[ast.FunctionDef]:
         """An `agent` method is a method decorated with `@agent`."""
@@ -92,12 +126,16 @@ def add_agent_method(self, agent: AgentConfig) -> None:
             pos, _ = self.get_node_range(main_method)
 
         code = f"""    @agentstack.agent
-    def {agent.name}(self, messages: list[str] = []):
+    def {agent.name}(self, messages: list[str] = []) -> Agent:
         agent_config = agentstack.get_agent('{agent.name}')
+        messages = [
+            agent_config.prompt,
+            *messages,
+        ]
         return Agent(
             name=agent_config.name, 
             model=agent_config.model, 
-            instructions='\\n'.join([agent_config.prompt, *messages, ]),
+            instructions='\\n'.join(messages),
             functions=[],
         )"""
 
@@ -259,7 +297,7 @@ def get_agent_tool_names(agent_name: str) -> list[Any]:
         return entrypoint.get_agent_tool_names(agent_name)
 
 
-def add_agent(agent: AgentConfig, position: Optional['InsertionPoint'] = None) -> None:
+def add_agent(agent: AgentConfig, position: Optional[InsertionPoint] = None) -> None:
     """
     Add an agent method to the entrypoint.
     """

agentstack/templates/openai_swarm/{{cookiecutter.project_metadata.project_slug}}/src/stack.py

@@ -4,14 +4,15 @@
 
 class {{ cookiecutter.project_metadata.project_name|replace('-', '')|replace('_', '')|capitalize }}Stack:
 
-    def _handoff(self, agent_name: str):
-        """Return a tool for handing off to another agent."""
-        agent = getattr(self, agent_name)
-        def func(context_variables: list[str]):
-            return agent(context_variables)
+    def _handoff(self, task_name: str):
+        """Return a task formatted as a tool for handing off to another agent."""
+        task = getattr(self, task_name)
+        def func(messages: list[str] = []) -> Agent:
+            return task(messages=messages)
+        func.__name__ = task_name
         return func
 
-    def _get_first_task(self):
+    def _get_first_task(self) -> Agent:
         """Get the first task."""
         task_name = agentstack.get_all_task_names()[0]
         return getattr(self, task_name)()
@@ -26,5 +27,11 @@ def run(self, inputs: list[str]):
         )
 
         for message in response.messages:
-            agentstack.log.info(message['content'])
+            if message.get('tool_calls'):
+                for tool_call in message['tool_calls']:
+                    agentstack.log.notify(f"Calling tool `{tool_call['function']['name']}`")
+                    agentstack.log.debug(tool_call['function']['arguments'])
+            elif message.get('role') != 'tool':
+                agentstack.log.notify(f"{message['role']}:")
+                agentstack.log.info(message['content'])