diff --git a/agentstack/generation/__init__.py b/agentstack/generation/__init__.py index b9aa0941..f1073379 100644 --- a/agentstack/generation/__init__.py +++ b/agentstack/generation/__init__.py @@ -1,3 +1,4 @@ from .agent_generation import generate_agent from .task_generation import generate_task from .tool_generation import add_tool, remove_tool +from .files import ConfigFile, EnvFile, CONFIG_FILENAME \ No newline at end of file diff --git a/agentstack/generation/files.py b/agentstack/generation/files.py new file mode 100644 index 00000000..0fc1fb14 --- /dev/null +++ b/agentstack/generation/files.py @@ -0,0 +1,121 @@ +from typing import Optional, Union +import os +import json +from pathlib import Path +from pydantic import BaseModel + + +DEFAULT_FRAMEWORK = "crewai" +CONFIG_FILENAME = "agentstack.json" +ENV_FILEMANE = ".env" + +class ConfigFile(BaseModel): + """ + Interface for interacting with the agentstack.json file inside a project directory. + Handles both data validation and file I/O. + + `path` is the directory where the agentstack.json file is located. Defaults + to the current working directory. + + Use it as a context manager to make and save edits: + ```python + with ConfigFile() as config: + config.tools.append('tool_name') + ``` + + Config Schema + ------------- + framework: str + The framework used in the project. Defaults to 'crewai'. + tools: list[str] + A list of tools that are currently installed in the project. + telemetry_opt_out: Optional[bool] + Whether the user has opted out of telemetry. + """ + framework: Optional[str] = DEFAULT_FRAMEWORK + tools: list[str] = [] + telemetry_opt_out: Optional[bool] = None + + def __init__(self, path: Union[str, Path, None] = None): + path = Path(path) if path else Path.cwd() + if os.path.exists(path / CONFIG_FILENAME): + with open(path / CONFIG_FILENAME, 'r') as f: + super().__init__(**json.loads(f.read())) + else: + raise FileNotFoundError(f"File {path / CONFIG_FILENAME} does not exist.") + self._path = path # attribute needs to be set after init + + def model_dump(self, *args, **kwargs) -> dict: + # Ignore None values + dump = super().model_dump(*args, **kwargs) + return {key: value for key, value in dump.items() if value is not None} + + def write(self): + with open(self._path / CONFIG_FILENAME, 'w') as f: + f.write(json.dumps(self.model_dump(), indent=4)) + + def __enter__(self) -> 'ConfigFile': return self + def __exit__(self, *args): self.write() + + +class EnvFile: + """ + Interface for interacting with the .env file inside a project directory. + Unlike the ConfigFile, we do not re-write the entire file on every change, + and instead just append new lines to the end of the file. This preseres + comments and other formatting that the user may have added and prevents + opportunities for data loss. + + `path` is the directory where the .env file is located. Defaults to the + current working directory. + `filename` is the name of the .env file, defaults to '.env'. + + Use it as a context manager to make and save edits: + ```python + with EnvFile() as env: + env.append_if_new('ENV_VAR', 'value') + ``` + """ + variables: dict[str, str] + + def __init__(self, path: Union[str, Path, None] = None, filename: str = ENV_FILEMANE): + self._path = Path(path) if path else Path.cwd() + self._filename = filename + self.read() + + def __getitem__(self, key): + return self.variables[key] + + def __setitem__(self, key, value): + if key in self.variables: + raise ValueError("EnvFile does not allow overwriting values.") + self.append_if_new(key, value) + + def __contains__(self, key) -> bool: + return key in self.variables + + def append_if_new(self, key, value): + if not key in self.variables: + self.variables[key] = value + self._new_variables[key] = value + + def read(self): + def parse_line(line): + key, value = line.split('=') + return key.strip(), value.strip() + + if os.path.exists(self._path / self._filename): + with open(self._path / self._filename, 'r') as f: + self.variables = dict([parse_line(line) for line in f.readlines() if '=' in line]) + else: + self.variables = {} + self._new_variables = {} + + def write(self): + with open(self._path / self._filename, 'a') as f: + for key, value in self._new_variables.items(): + f.write(f"\n{key}={value}") + + def __enter__(self) -> 'EnvFile': return self + def __exit__(self, *args): self.write() + diff --git a/agentstack/generation/tool_generation.py b/agentstack/generation/tool_generation.py index 5301ddfb..8eb68692 100644 --- a/agentstack/generation/tool_generation.py +++ b/agentstack/generation/tool_generation.py @@ -8,12 +8,13 @@ from pydantic import BaseModel, ValidationError from agentstack.utils import get_package_path +from agentstack.generation.files import ConfigFile, EnvFile from .gen_utils import insert_code_after_tag, string_in_file from ..utils import open_json_file, get_framework, term_color TOOL_INIT_FILENAME = "src/tools/__init__.py" -AGENTSTACK_JSON_FILENAME = "agentstack.json" + FRAMEWORK_FILENAMES: dict[str, str] = { 'crewai': 'src/crew.py', } @@ -87,9 +88,9 @@ def add_tool(tool_name: str, path: Optional[str] = None): path = './' framework = get_framework(path) - agentstack_json = open_json_file(f'{path}{AGENTSTACK_JSON_FILENAME}') + agentstack_config = ConfigFile(path) - if tool_name in agentstack_json.get('tools', []): + if tool_name in agentstack_config.tools: print(term_color(f'Tool {tool_name} is already installed', 'red')) sys.exit(1) @@ -100,27 +101,20 @@ def add_tool(tool_name: str, path: Optional[str] = None): shutil.copy(tool_file_path, f'{path}src/tools/{tool_name}_tool.py') # Move tool from package to project add_tool_to_tools_init(tool_data, path) # Export tool from tools dir add_tool_to_agent_definition(framework, tool_data, path) # Add tool to agent definition - if tool_data.env: # if the env vars aren't in the .env files, add them - # tool_data.env is a dict, key is the env var name, value is the value - for var, value in tool_data.env.items(): - env_var = f'{var}={value}' - if not string_in_file(f'{path}.env', env_var): - insert_code_after_tag(f'{path}.env', '# Tools', [env_var, ]) - if not string_in_file(f'{path}.env.example', env_var): - insert_code_after_tag(f'{path}.env.example', '# Tools', [env_var, ]) - if tool_data.post_install: - os.system(tool_data.post_install) + if tool_data.env: # add environment variables which don't exist + with EnvFile(path) as env: + for var, value in tool_data.env.items(): + env.append_if_new(var, value) + with EnvFile(path, filename=".env.example") as env: + for var, value in tool_data.env.items(): + env.append_if_new(var, value) if tool_data.post_install: os.system(tool_data.post_install) - - if not agentstack_json.get('tools'): - agentstack_json['tools'] = [] - agentstack_json['tools'].append(tool_name) - with open(f'{path}{AGENTSTACK_JSON_FILENAME}', 'w') as f: - json.dump(agentstack_json, f, indent=4) + with agentstack_config as config: + config.tools.append(tool_name) print(term_color(f'🔨 Tool {tool_name} added to agentstack project successfully', 'green')) if tool_data.cta: @@ -133,9 +127,9 @@ def remove_tool(tool_name: str, path: Optional[str] = None): path = './' framework = get_framework() - agentstack_json = open_json_file(f'{path}{AGENTSTACK_JSON_FILENAME}') + agentstack_config = ConfigFile(path) - if not tool_name in agentstack_json.get('tools', []): + if not tool_name in agentstack_config.tools: print(term_color(f'Tool {tool_name} is not installed', 'red')) sys.exit(1) @@ -152,9 +146,8 @@ def remove_tool(tool_name: str, path: Optional[str] = None): os.system(tool_data.post_remove) # We don't remove the .env variables to preserve user data. - agentstack_json['tools'].remove(tool_name) - with open(f'{path}{AGENTSTACK_JSON_FILENAME}', 'w') as f: - json.dump(agentstack_json, f, indent=4) + with agentstack_config as config: + config.tools.remove(tool_name) print(term_color(f'🔨 Tool {tool_name}', 'green'), term_color('removed', 'red'), term_color('from agentstack project successfully', 'green')) diff --git a/agentstack/utils.py b/agentstack/utils.py index cbe3ea5c..a4dccd84 100644 --- a/agentstack/utils.py +++ b/agentstack/utils.py @@ -8,7 +8,6 @@ from pathlib import Path import importlib.resources - def get_version(): try: return version('agentstack') @@ -17,8 +16,11 @@ def get_version(): return "Unknown version" -def verify_agentstack_project(): - if not os.path.isfile('agentstack.json'): +def verify_agentstack_project(path: Optional[str] = None): + from agentstack.generation import ConfigFile + try: + agentstack_config = ConfigFile(path) + except FileNotFoundError: print("\033[31mAgentStack Error: This does not appear to be an AgentStack project." "\nPlease ensure you're at the root directory of your project and a file named agentstack.json exists. " "If you're starting a new project, run `agentstack init`\033[0m") @@ -33,13 +35,10 @@ def get_package_path() -> Path: def get_framework(path: Optional[str] = None) -> str: + from agentstack.generation import ConfigFile try: - file_path = 'agentstack.json' - if path is not None: - file_path = path + '/' + file_path - - agentstack_data = open_json_file(file_path) - framework = agentstack_data.get('framework') + agentstack_config = ConfigFile(path) + framework = agentstack_config.framework if framework.lower() not in ['crewai', 'autogen', 'litellm']: print(term_color("agentstack.json contains an invalid framework", "red")) @@ -51,14 +50,10 @@ def get_framework(path: Optional[str] = None) -> str: def get_telemetry_opt_out(path: Optional[str] = None) -> str: + from agentstack.generation import ConfigFile try: - file_path = 'agentstack.json' - if path is not None: - file_path = path + '/' + file_path - - agentstack_data = open_json_file(file_path) - opt_out = agentstack_data.get('telemetry_opt_out', False) - return opt_out + agentstack_config = ConfigFile(path) + return bool(agentstack_config.telemetry_opt_out) except FileNotFoundError: print("\033[31mFile agentstack.json does not exist. Are you in the right directory?\033[0m") sys.exit(1) diff --git a/tests/fixtures/.env b/tests/fixtures/.env new file mode 100644 index 00000000..3f1c8b1d --- /dev/null +++ b/tests/fixtures/.env @@ -0,0 +1,3 @@ + +ENV_VAR1=value1 +ENV_VAR2=value2 \ No newline at end of file diff --git a/tests/fixtures/agentstack.json b/tests/fixtures/agentstack.json new file mode 100644 index 00000000..4ca18a10 --- /dev/null +++ b/tests/fixtures/agentstack.json @@ -0,0 +1,4 @@ +{ + "framework": "crewai", + "tools": ["tool1", "tool2"] +} \ No newline at end of file diff --git a/tests/test_generation_files.py b/tests/test_generation_files.py new file mode 100644 index 00000000..8f8549e3 --- /dev/null +++ b/tests/test_generation_files.py @@ -0,0 +1,90 @@ +import os, sys +import unittest +import importlib.resources +from pathlib import Path +import shutil +from agentstack.generation.files import ConfigFile, EnvFile +from agentstack.utils import verify_agentstack_project, get_framework, get_telemetry_opt_out + +BASE_PATH = Path(__file__).parent + +class GenerationFilesTest(unittest.TestCase): + def test_read_config(self): + config = ConfigFile(BASE_PATH / "fixtures") # + agentstack.json + assert config.framework == "crewai" + assert config.tools == ["tool1", "tool2"] + assert config.telemetry_opt_out is None + + def test_write_config(self): + try: + os.makedirs(BASE_PATH/"tmp", exist_ok=True) + shutil.copy(BASE_PATH/"fixtures/agentstack.json", + BASE_PATH/"tmp/agentstack.json") + + with ConfigFile(BASE_PATH/"tmp") as config: + config.framework = "crewai" + config.tools = ["tool1", "tool2"] + config.telemetry_opt_out = True + + tmp_data = open(BASE_PATH/"tmp/agentstack.json").read() + assert tmp_data == """{ + "framework": "crewai", + "tools": [ + "tool1", + "tool2" + ], + "telemetry_opt_out": true +}""" + except Exception as e: + raise e + finally: + os.remove(BASE_PATH / "tmp/agentstack.json") + #os.rmdir(BASE_PATH / "tmp") + + def test_read_missing_config(self): + with self.assertRaises(FileNotFoundError) as context: + config = ConfigFile(BASE_PATH / "missing") + + def test_verify_agentstack_project_valid(self): + verify_agentstack_project(BASE_PATH / "fixtures") + + def test_verify_agentstack_project_invalid(self): + with self.assertRaises(SystemExit) as context: + verify_agentstack_project(BASE_PATH / "missing") + + def test_get_framework(self): + assert get_framework(BASE_PATH / "fixtures") == "crewai" + with self.assertRaises(SystemExit) as context: + get_framework(BASE_PATH / "missing") + + def test_get_telemetry_opt_out(self): + assert get_telemetry_opt_out(BASE_PATH / "fixtures") is False + with self.assertRaises(SystemExit) as context: + get_telemetry_opt_out(BASE_PATH / "missing") + + def test_read_env(self): + env = EnvFile(BASE_PATH / "fixtures") + assert env.variables == {"ENV_VAR1": "value1", "ENV_VAR2": "value2"} + assert env["ENV_VAR1"] == "value1" + assert env["ENV_VAR2"] == "value2" + with self.assertRaises(KeyError) as context: + env["ENV_VAR3"] + + def test_write_env(self): + try: + os.makedirs(BASE_PATH/"tmp", exist_ok=True) + shutil.copy(BASE_PATH/"fixtures/.env", + BASE_PATH/"tmp/.env") + + with EnvFile(BASE_PATH/"tmp") as env: + env.append_if_new("ENV_VAR1", "value100") # Should not be updated + env.append_if_new("ENV_VAR100", "value2") # Should be added + + tmp_data = open(BASE_PATH/"tmp/.env").read() + assert tmp_data == """\nENV_VAR1=value1\nENV_VAR2=value2\nENV_VAR100=value2""" + except Exception as e: + raise e + finally: + os.remove(BASE_PATH / "tmp/.env") + #os.rmdir(BASE_PATH / "tmp") +