Merge pull request #248 from tcdent/framework-swarm

OpenAI Swarm is now supported by AgentStack. `agentstack init <project_name> --framework=openai_swarm`

Author: tcdent

agentstack/__init__.py

@@ -9,8 +9,8 @@
 from pathlib import Path
 from agentstack import conf
 from agentstack.utils import get_framework
-from agentstack.agents import get_agent
-from agentstack.tasks import get_task
+from agentstack.agents import get_agent, get_all_agents, get_all_agent_names
+from agentstack.tasks import get_task, get_all_tasks, get_all_task_names
 from agentstack.inputs import get_inputs
 from agentstack import frameworks
 
@@ -22,7 +22,11 @@
     "get_tags", 
     "get_framework", 
     "get_agent", 
+    "get_all_agents",
+    "get_all_agent_names",
     "get_task", 
+    "get_all_tasks",
+    "get_all_task_names",
     "get_inputs", 
 ]
 

agentstack/frameworks/__init__.py

@@ -16,9 +16,11 @@
 
 CREWAI = 'crewai'
 LANGGRAPH = 'langgraph'
+OPENAI_SWARM = 'openai_swarm'
 SUPPORTED_FRAMEWORKS = [
     CREWAI,
     LANGGRAPH,
+    OPENAI_SWARM,
 ]
 
 
@@ -58,9 +60,9 @@ def remove_tool(self, tool: ToolConfig, agent_name: str) -> None:
         """
         ...
 
-    def get_tool_callables(self, tool_name: str) -> list[Callable]:
+    def wrap_tool(self, tool_func: Callable) -> Callable:
         """
-        Get a tool by name and return it as a list of framework-native callables.
+        Wrap a tool function with framework-specific functionality.
         """
         ...
 
@@ -173,7 +175,43 @@ def get_tool_callables(tool_name: str) -> list[Callable]:
     """
     Get a tool by name and return it as a list of framework-native callables.
     """
-    return get_framework_module(get_framework()).get_tool_callables(tool_name)
+    # TODO: remove after agentops fixes their issue
+    # wrap method with agentops tool event
+    def wrap_method(method: Callable) -> Callable:
+        from inspect import signature
+        
+        original_signature = signature(method)
+        def wrapped_method(*args, **kwargs):
+            import agentops
+            tool_event = agentops.ToolEvent(method.__name__)
+            result = method(*args, **kwargs)
+            agentops.record(tool_event)
+            return result
+
+        # Preserve all original attributes
+        wrapped_method.__name__ = method.__name__
+        wrapped_method.__doc__ = method.__doc__
+        wrapped_method.__module__ = method.__module__
+        wrapped_method.__qualname__ = method.__qualname__
+        wrapped_method.__annotations__ = getattr(method, '__annotations__', {})
+        wrapped_method.__signature__ = original_signature  # type: ignore
+        return wrapped_method
+
+    tool_funcs = []
+    tool_config = ToolConfig.from_tool_name(tool_name)
+    for tool_func_name in tool_config.tools:
+        tool_func = getattr(tool_config.module, tool_func_name)
+
+        assert callable(tool_func), f"Tool function {tool_func_name} is not callable."
+        assert tool_func.__doc__, f"Tool function {tool_func_name} is missing a docstring."
+
+        # First wrap with agentops
+        agentops_wrapped = wrap_method(tool_func)
+        # Then apply framework decorators
+        framework_wrapped = get_framework_module(get_framework()).wrap_tool(agentops_wrapped)
+        tool_funcs.append(framework_wrapped)
+
+    return tool_funcs
 
 
 def get_agent_method_names() -> list[str]:

agentstack/frameworks/crewai.py

@@ -8,6 +8,7 @@
 from agentstack.agents import AgentConfig
 from agentstack.generation import asttools
 from agentstack import graph
+
 if TYPE_CHECKING:
     from agentstack.generation import InsertionPoint
 
@@ -19,6 +20,7 @@ class CrewFile(asttools.File):
     Parses and manipulates the CrewAI entrypoint file.
     All AST interactions should happen within the methods of this class.
     """
+
     def write(self):
         """
         Early versions of the crew entrypoint file used tabs instead of spaces.
@@ -65,7 +67,7 @@ def {task.name}(self) -> Task:
         return Task(
             config=self.tasks_config['{task.name}'],
         )"""
-        
+
         if not self.source[:pos].endswith('\n'):
             code = '\n\n' + code
         if not self.source[pos:].startswith('\n'):
@@ -95,7 +97,7 @@ def {agent.name}(self) -> Agent:
             tools=[], # add tools here or use `agentstack tools add <tool_name>
             verbose=True,
         )"""
-        
+
         if not self.source[:pos].endswith('\n'):
             code = '\n\n' + code
         if not self.source[pos:].startswith('\n'):
@@ -280,7 +282,7 @@ def add_agent(agent: AgentConfig, position: Optional['InsertionPoint'] = None) -
     """
     if position is not None:
         raise NotImplementedError("Agent insertion points are not supported in CrewAI.")
-    
+
     with CrewFile(conf.PATH / ENTRYPOINT) as crew_file:
         crew_file.add_agent_method(agent)
 
@@ -302,52 +304,19 @@ def remove_tool(tool: ToolConfig, agent_name: str):
         crew_file.remove_agent_tools(agent_name, tool)
 
 
-def get_tool_callables(tool_name: str) -> list[Callable]:
+def wrap_tool(tool_func: Callable) -> Callable:
     """
-    Get a tool implementations for use directly by a CrewAI agent.
+    Wrap a tool function with framework-specific functionality.
     """
     try:
         from crewai.tools import tool as _crewai_tool_decorator
     except ImportError:
         raise ValidationError("Could not import `crewai`. Is this an AgentStack CrewAI project?")
 
-    # TODO: remove after agentops fixes their issue
-    # wrap method with agentops tool event
-    def wrap_method(method: Callable) -> Callable:
-        def wrapped_method(*args, **kwargs):
-            import agentops
-            tool_event = agentops.ToolEvent(method.__name__)
-            result = method(*args, **kwargs)
-            agentops.record(tool_event)
-            return result
-
-        # Preserve all original attributes
-        wrapped_method.__name__ = method.__name__
-        wrapped_method.__doc__ = method.__doc__
-        wrapped_method.__module__ = method.__module__
-        wrapped_method.__qualname__ = method.__qualname__
-        wrapped_method.__annotations__ = getattr(method, '__annotations__', {})
-        return wrapped_method
-
-    tool_funcs = []
-    tool_config = ToolConfig.from_tool_name(tool_name)
-    for tool_func_name in tool_config.tools:
-        tool_func = getattr(tool_config.module, tool_func_name)
-
-        assert callable(tool_func), f"Tool function {tool_func_name} is not callable."
-        assert tool_func.__doc__, f"Tool function {tool_func_name} is missing a docstring."
-
-        # First wrap with agentops
-        agentops_wrapped = wrap_method(tool_func)
-        # Then apply CrewAI decorator last so it properly inherits from BaseTool
-        crewai_wrapped = _crewai_tool_decorator(agentops_wrapped)
-        tool_funcs.append(crewai_wrapped)
-
-    return tool_funcs
+    return _crewai_tool_decorator(tool_func)
 
 
 def get_graph() -> list[graph.Edge]:
     """Get the graph of the user's project."""
     log.debug("CrewAI does not support graph generation.")
     return []
-

agentstack/frameworks/langgraph.py

@@ -18,7 +18,11 @@
 GRAPH_NODE_END = 'END'
 GRAPH_NODE_TOOLS = 'tools'  # references the `ToolNode` instance
 GRAPH_NODE_TOOLS_CONDITION = 'tools_condition'
-GRAPH_NODES_SPECIAL = (GRAPH_NODE_START, GRAPH_NODE_END, GRAPH_NODE_TOOLS_CONDITION, )
+GRAPH_NODES_SPECIAL = (
+    GRAPH_NODE_START,
+    GRAPH_NODE_END,
+    GRAPH_NODE_TOOLS_CONDITION,
+)
 
 
 @dataclass
@@ -347,7 +351,7 @@ def _get_node_name(node: ast.expr) -> str:
         for node in nodes:
             source, target = node.args
             source_name = _get_node_name(source)
-            #target_name = _get_node_name(target)
+            # target_name = _get_node_name(target)
             if source_name == GRAPH_NODE_TOOLS:  # TODO this is a bit brittle
                 nodes.remove(node)
             # if target_name == GRAPH_NODE_TOOLS:
@@ -430,7 +434,7 @@ def add_conditional_edge(self, edge: graph.Edge):
         else:
             graph_instance = asttools.find_method_calls(self.get_run_method(), 'StateGraph')[0]
             _, end = self.get_node_range(graph_instance)
-        
+
         source, target = edge.source.name, edge.target.name
         # wrap the node names in quotes if they are not special nodes
         if edge.source.type != graph.NodeType.SPECIAL:
@@ -745,36 +749,12 @@ def remove_tool(tool: ToolConfig, agent_name: str):
         entrypoint.remove_agent_tools(agent_name, tool)
 
 
-def get_tool_callables(tool_name: str) -> list[Callable]:
+def wrap_tool(tool_func: Callable) -> Callable:
     """
-    Get a tool by name and return it as a list of framework-native callables.
+    Wrap a tool function with framework-specific functionality.
     """
-    # LangGraph accepts functions as tools, so we can return them directly
-    tool_funcs = []
-    tool_config = ToolConfig.from_tool_name(tool_name)
-
-    # TODO: remove after agentops supports langgraph
-    # wrap method with agentops tool event
-    def wrap_method(method: Callable) -> Callable:
-        @wraps(method)  # This preserves the original function's metadata
-        def wrapped_method(*args, **kwargs):
-            import agentops
-            tool_event = agentops.ToolEvent(method.__name__)
-            result = method(*args, **kwargs)
-            agentops.record(tool_event)
-            return result
-
-        return wrapped_method
-
-    for tool_func_name in tool_config.tools:
-        tool_func = getattr(tool_config.module, tool_func_name)
-
-        assert callable(tool_func), f"Tool function {tool_func_name} is not callable."
-        assert tool_func.__doc__, f"Tool function {tool_func_name} is missing a docstring."
-
-        tool_funcs.append(wrap_method(tool_func))
-
-    return tool_funcs
+    # LangGraph accepts bare functions as tools, so we don't need to do anything here.
+    return tool_func
 
 
 def get_graph() -> list[graph.Edge]: