Merge pull request #118 from tcdent/template-export

Export agentstack project as a template

agentstack/agents.py

@@ -29,22 +29,22 @@ class AgentConfig(pydantic.BaseModel):
     -------------
     name: str
         The name of the agent; used for lookup.
-    role: Optional[str]
+    role: str
         The role of the agent.
-    goal: Optional[str]
+    goal: str
         The goal of the agent.
-    backstory: Optional[str]
+    backstory: str
         The backstory of the agent.
-    llm: Optional[str]
+    llm: str
         The model this agent should use.
         Adheres to the format set by the framework.
     """
 
     name: str
-    role: Optional[str] = ""
-    goal: Optional[str] = ""
-    backstory: Optional[str] = ""
-    llm: Optional[str] = ""
+    role: str = ""
+    goal: str = ""
+    backstory: str = ""
+    llm: str = ""
 
     def __init__(self, name: str, path: Optional[Path] = None):
         if not path:
@@ -96,3 +96,18 @@ def __enter__(self) -> 'AgentConfig':
 
     def __exit__(self, *args):
         self.write()
+
+
+def get_all_agent_names(path: Optional[Path] = None) -> list[str]:
+    if not path:
+        path = Path()
+    filename = path / AGENTS_FILENAME
+    if not os.path.exists(filename):
+        return []
+    with open(filename, 'r') as f:
+        data = yaml.load(f) or {}
+    return list(data.keys())
+
+
+def get_all_agents(path: Optional[Path] = None) -> list[AgentConfig]:
+    return [AgentConfig(name, path) for name in get_all_agent_names(path)]

agentstack/cli/__init__.py

@@ -1 +1 @@
-from .cli import init_project_builder, list_tools, configure_default_model, run_project
+from .cli import init_project_builder, list_tools, configure_default_model, run_project, export_template

agentstack/cli/cli.py

@@ -15,6 +15,7 @@
 from cookiecutter.main import cookiecutter
 from dotenv import load_dotenv
 import subprocess
+from packaging.metadata import Metadata
 
 from .agentstack_data import (
     FrameworkData,
@@ -25,11 +26,13 @@
 from agentstack.logger import log
 from agentstack.utils import get_package_path
 from agentstack.tools import get_all_tools
-from agentstack.generation.files import ConfigFile
+from agentstack.generation.files import ConfigFile, ProjectFile
 from agentstack import frameworks
 from agentstack import packaging
 from agentstack import generation
-from agentstack.utils import open_json_file, term_color, is_snake_case
+from agentstack.agents import get_all_agents
+from agentstack.tasks import get_all_tasks
+from agentstack.utils import open_json_file, term_color, is_snake_case, get_framework
 from agentstack.update import AGENTSTACK_PACKAGE
 from agentstack.proj_templates import TemplateConfig
 
@@ -62,13 +65,13 @@ def init_project_builder(
             try:
                 template_data = TemplateConfig.from_url(template)
             except Exception as e:
-                print(term_color(f"Failed to fetch template data from {template}", 'red'))
+                print(term_color(f"Failed to fetch template data from {template}.\n{e}", 'red'))
                 sys.exit(1)
         else:
             try:
                 template_data = TemplateConfig.from_template_name(template)
             except Exception as e:
-                print(term_color(f"Failed to load template {template}", 'red'))
+                print(term_color(f"Failed to load template {template}.\n{e}", 'red'))
                 sys.exit(1)
 
     if template_data:
@@ -81,11 +84,11 @@ def init_project_builder(
         }
         framework = template_data.framework
         design = {
-            'agents': template_data.agents,
-            'tasks': template_data.tasks,
+            'agents': [agent.model_dump() for agent in template_data.agents],
+            'tasks': [task.model_dump() for task in template_data.tasks],
             'inputs': template_data.inputs,
         }
-        tools = template_data.tools
+        tools = [tools.model_dump() for tools in template_data.tools]
 
     elif use_wizard:
         welcome_message()
@@ -465,3 +468,85 @@ def list_tools():
 
     print("\n\n✨ Add a tool with: agentstack tools add <tool_name>")
     print("   https://docs.agentstack.sh/tools/core")
+
+
+def export_template(output_filename: str, path: str = ''):
+    """
+    Export the current project as a template.
+    """
+    _path = Path(path)
+    framework = get_framework(_path)
+
+    try:
+        metadata = ProjectFile(_path)
+    except Exception as e:
+        print(term_color(f"Failed to load project metadata: {e}", 'red'))
+        sys.exit(1)
+
+    # Read all the agents from the project's agents.yaml file
+    agents: list[TemplateConfig.Agent] = []
+    for agent in get_all_agents(_path):
+        agents.append(
+            TemplateConfig.Agent(
+                name=agent.name,
+                role=agent.role,
+                goal=agent.goal,
+                backstory=agent.backstory,
+                model=agent.llm,  # TODO consistent naming (llm -> model)
+            )
+        )
+
+    # Read all the tasks from the project's tasks.yaml file
+    tasks: list[TemplateConfig.Task] = []
+    for task in get_all_tasks(_path):
+        tasks.append(
+            TemplateConfig.Task(
+                name=task.name,
+                description=task.description,
+                expected_output=task.expected_output,
+                agent=task.agent,
+            )
+        )
+
+    # Export all of the configured tools from the project
+    tools_agents: dict[str, list[str]] = {}
+    for agent_name in frameworks.get_agent_names(framework, _path):
+        for tool_name in frameworks.get_agent_tool_names(framework, agent_name, _path):
+            if not tool_name:
+                continue
+            if tool_name not in tools_agents:
+                tools_agents[tool_name] = []
+            tools_agents[tool_name].append(agent_name)
+
+    tools: list[TemplateConfig.Tool] = []
+    for tool_name, agent_names in tools_agents.items():
+        tools.append(
+            TemplateConfig.Tool(
+                name=tool_name,
+                agents=agent_names,
+            )
+        )
+
+    inputs: list[str] = []
+    # TODO extract inputs from project
+    # for input in frameworks.get_input_names():
+    #     inputs.append(input)
+
+    template = TemplateConfig(
+        template_version=1,
+        name=metadata.project_name,
+        description=metadata.project_description,
+        framework=framework,
+        method="sequential",  # TODO this needs to be stored in the project somewhere
+        agents=agents,
+        tasks=tasks,
+        tools=tools,
+        inputs=inputs,
+    )
+
+    try:
+        template.write_to_file(_path / output_filename)
+        print(term_color(f"Template saved to: {_path / output_filename}", 'green'))
+    except Exception as e:
+        print(term_color(f"Failed to write template to file: {e}", 'red'))
+        sys.exit(1)

agentstack/frameworks/__init__.py

@@ -28,6 +28,12 @@ def validate_project(self, path: Optional[Path] = None) -> None:
         """
         ...
 
+    def get_tool_names(self, path: Optional[Path] = None) -> list[str]:
+        """
+        Get a list of tool names in the user's project.
+        """
+        ...
+
     def add_tool(self, tool: ToolConfig, agent_name: str, path: Optional[Path] = None) -> None:
         """
         Add a tool to an agent in the user's project.
@@ -46,6 +52,12 @@ def get_agent_names(self, path: Optional[Path] = None) -> list[str]:
         """
         ...
 
+    def get_agent_tool_names(self, agent_name: str, path: Optional[Path] = None) -> list[str]:
+        """
+        Get a list of tool names in an agent in the user's project.
+        """
+        ...
+
     def add_agent(self, agent: AgentConfig, path: Optional[Path] = None) -> None:
         """
         Add an agent to the user's project.
@@ -102,6 +114,12 @@ def get_agent_names(framework: str, path: Optional[Path] = None) -> list[str]:
     """
     return get_framework_module(framework).get_agent_names(path)
 
+def get_agent_tool_names(framework: str, agent_name: str, path: Optional[Path] = None) -> list[str]:
+    """
+    Get a list of tool names in the user's project.
+    """
+    return get_framework_module(framework).get_agent_tool_names(agent_name, path)
+
 def add_agent(framework: str, agent: AgentConfig, path: Optional[Path] = None):
     """
     Add an agent to the user's project.

agentstack/frameworks/crewai.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, Any
 from pathlib import Path
 import ast
 from agentstack import ValidationError
@@ -259,6 +259,18 @@ def get_agent_names(path: Optional[Path] = None) -> list[str]:
     return [method.name for method in crew_file.get_agent_methods()]
 
 
+def get_agent_tool_names(agent_name: str, path: Optional[Path] = None) -> list[Any]:
+    """
+    Get a list of tools used by an agent.
+    """
+    if path is None:
+        path = Path()
+    with CrewFile(path / ENTRYPOINT) as crew_file:
+        tools = crew_file.get_agent_tools(agent_name)
+    print([node for node in tools.elts])
+    return [asttools.get_node_value(node) for node in tools.elts]
+
+
 def add_agent(agent: AgentConfig, path: Optional[Path] = None) -> None:
     """
     Add an agent method to the CrewAI entrypoint.

agentstack/generation/agent_generation.py

@@ -27,7 +27,7 @@ def add_agent(
         config.role = role or "Add your role here"
         config.goal = goal or "Add your goal here"
         config.backstory = backstory or "Add your backstory here"
-        config.llm = llm or agentstack_config.default_model
+        config.llm = llm or agentstack_config.default_model or ""
 
     try:
         frameworks.add_agent(framework, agent, path)

agentstack/generation/asttools.py

@@ -9,7 +9,7 @@
 functions that are useful for the specific tasks we need to accomplish.
 """
 
-from typing import TypeVar, Optional, Union, Iterable
+from typing import TypeVar, Optional, Union, Iterable, Any
 from pathlib import Path
 import ast
 import astor
@@ -175,3 +175,16 @@ def find_decorated_method_in_class(classdef: ast.ClassDef, decorator_name: str)
 def create_attribute(base_name: str, attr_name: str) -> ast.Attribute:
     """Create an AST node for an attribute"""
     return ast.Attribute(value=ast.Name(id=base_name, ctx=ast.Load()), attr=attr_name, ctx=ast.Load())
+
+
+def get_node_value(node: Union[ast.expr, ast.Attribute, ast.Constant, ast.Str, ast.Num]) -> Any:
+    if isinstance(node, ast.Constant):
+        return node.value
+    elif isinstance(node, ast.Attribute):
+        prefix = get_node_value(node.value)
+        if prefix:
+            return '.'.join([prefix, node.attr])
+        else:
+            return node.attr
+    else:
+        return None

agentstack/generation/files.py

@@ -1,14 +1,20 @@
 from typing import Optional, Union
-import os
+import os, sys
 import json
 from pathlib import Path
 from pydantic import BaseModel
+
+if sys.version_info >= (3, 11):
+    import tomllib
+else:
+    import tomli as tomllib
 from agentstack.utils import get_version
 
 
 DEFAULT_FRAMEWORK = "crewai"
 CONFIG_FILENAME = "agentstack.json"
 ENV_FILEMANE = ".env"
+PYPROJECT_FILENAME = "pyproject.toml"
 
 
 class ConfigFile(BaseModel):
@@ -140,3 +146,45 @@ def __enter__(self) -> 'EnvFile':
 
     def __exit__(self, *args):
         self.write()
+
+
+class ProjectFile:
+    """
+    Interface for interacting with pyproject.toml files inside of a project directory.
+    This class is read-only and does not support writing changes back to the file.
+    We expose project metadata as properties to support migration to other formats
+    in the future.
+    """
+
+    _data: dict
+
+    def __init__(self, path: Union[str, Path, None] = None, filename: str = PYPROJECT_FILENAME):
+        self._path = Path(path) if path else Path.cwd()
+        self._filename = filename
+        self.read()
+
+    @property
+    def project_metadata(self) -> dict:
+        try:
+            return self._data['tool']['poetry']
+        except KeyError:
+            raise KeyError("No poetry metadata found in pyproject.toml.")
+
+    @property
+    def project_name(self) -> str:
+        return self.project_metadata.get('name', '')
+
+    @property
+    def project_version(self) -> str:
+        return self.project_metadata.get('version', '')
+
+    @property
+    def project_description(self) -> str:
+        return self.project_metadata.get('description', '')
+
+    def read(self):
+        if os.path.exists(self._path / self._filename):
+            with open(self._path / self._filename, 'rb') as f:
+                self._data = tomllib.load(f)
+        else:
+            raise FileNotFoundError(f"File {self._path / self._filename} does not exist.")

agentstack/main.py

@@ -2,7 +2,13 @@
 import os
 import sys
 
-from agentstack.cli import init_project_builder, list_tools, configure_default_model, run_project
+from agentstack.cli import (
+    init_project_builder,
+    list_tools,
+    configure_default_model,
+    run_project,
+    export_template,
+)
 from agentstack.telemetry import track_cli_command
 from agentstack.utils import get_version, get_framework
 from agentstack import generation
@@ -83,6 +89,9 @@ def main():
     tools_remove_parser = tools_subparsers.add_parser("remove", aliases=["r"], help="Remove a tool")
     tools_remove_parser.add_argument("name", help="Name of the tool to remove")
 
+    export_parser = subparsers.add_parser('export', aliases=['e'], help='Export your agent as a template')
+    export_parser.add_argument('filename', help='The name of the file to export to')
+
     update = subparsers.add_parser('update', aliases=['u'], help='Check for updates')
 
     # Parse arguments
@@ -128,6 +137,8 @@ def main():
             generation.remove_tool(args.name)
         else:
             tools_parser.print_help()
+    elif args.command in ['export', 'e']:
+        export_template(args.filename)
     elif args.command in ['update', 'u']:
         pass  # Update check already done
     else: