From 7d0eba9e0c77e9e8857b03fedc076098e090e7a8 Mon Sep 17 00:00:00 2001 From: Travis Dent Date: Thu, 28 Nov 2024 08:01:35 -0800 Subject: [PATCH 1/3] Abstract interactions with agentstack config file and env file into wrappers. --- agentstack/generation/__init__.py | 1 + agentstack/generation/files.py | 121 +++++++++++++++++++++++ agentstack/generation/tool_generation.py | 41 ++++---- agentstack/utils.py | 23 ++--- tests/fixtures/.env | 3 + tests/fixtures/agentstack.json | 4 + tests/test_generation_files.py | 83 ++++++++++++++++ 7 files changed, 237 insertions(+), 39 deletions(-) create mode 100644 agentstack/generation/files.py create mode 100644 tests/fixtures/.env create mode 100644 tests/fixtures/agentstack.json create mode 100644 tests/test_generation_files.py 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..49e88834 --- /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): + # 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): 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 aooend 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): + 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): 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..d34f3958 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') @@ -18,7 +17,8 @@ def get_version(): def verify_agentstack_project(): - if not os.path.isfile('agentstack.json'): + from agentstack.generation import CONFIG_FILENAME + if not os.path.isfile(CONFIG_FILENAME): 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 +33,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 +48,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..e8357414 --- /dev/null +++ b/tests/test_generation_files.py @@ -0,0 +1,83 @@ +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 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_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") + From 0003eb5828304619bcf6249437dd56b9f58c2a81 Mon Sep 17 00:00:00 2001 From: Travis Dent Date: Thu, 28 Nov 2024 08:35:54 -0800 Subject: [PATCH 2/3] Use ConfigFile object to verify existence of agentstack.json fie in project. Add tests for verify_agentstack_project. --- agentstack/utils.py | 8 +++++--- tests/test_generation_files.py | 9 ++++++++- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/agentstack/utils.py b/agentstack/utils.py index d34f3958..a4dccd84 100644 --- a/agentstack/utils.py +++ b/agentstack/utils.py @@ -16,9 +16,11 @@ def get_version(): return "Unknown version" -def verify_agentstack_project(): - from agentstack.generation import CONFIG_FILENAME - if not os.path.isfile(CONFIG_FILENAME): +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") diff --git a/tests/test_generation_files.py b/tests/test_generation_files.py index e8357414..8f8549e3 100644 --- a/tests/test_generation_files.py +++ b/tests/test_generation_files.py @@ -4,7 +4,7 @@ from pathlib import Path import shutil from agentstack.generation.files import ConfigFile, EnvFile -from agentstack.utils import get_framework, get_telemetry_opt_out +from agentstack.utils import verify_agentstack_project, get_framework, get_telemetry_opt_out BASE_PATH = Path(__file__).parent @@ -45,6 +45,13 @@ 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: From d66f78b64264f14fd8f0a810c00e9677dbc154bc Mon Sep 17 00:00:00 2001 From: Travis Dent Date: Thu, 28 Nov 2024 09:29:23 -0800 Subject: [PATCH 3/3] typo, typing --- agentstack/generation/files.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/agentstack/generation/files.py b/agentstack/generation/files.py index 49e88834..0fc1fb14 100644 --- a/agentstack/generation/files.py +++ b/agentstack/generation/files.py @@ -45,7 +45,7 @@ def __init__(self, path: Union[str, Path, None] = None): 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): + 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} @@ -54,7 +54,7 @@ def write(self): with open(self._path / CONFIG_FILENAME, 'w') as f: f.write(json.dumps(self.model_dump(), indent=4)) - def __enter__(self): return self + def __enter__(self) -> 'ConfigFile': return self def __exit__(self, *args): self.write() @@ -62,7 +62,7 @@ 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 aooend new lines to the end of the file. This preseres + 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. @@ -91,7 +91,7 @@ def __setitem__(self, key, value): raise ValueError("EnvFile does not allow overwriting values.") self.append_if_new(key, value) - def __contains__(self, key): + def __contains__(self, key) -> bool: return key in self.variables def append_if_new(self, key, value): @@ -116,6 +116,6 @@ def write(self): for key, value in self._new_variables.items(): f.write(f"\n{key}={value}") - def __enter__(self): return self + def __enter__(self) -> 'EnvFile': return self def __exit__(self, *args): self.write()