Merge pull request #71 from tcdent/file-abstractions

File abstractions for project config and env

Author: bboynton97

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

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()
+

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'))
 

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)

tests/fixtures/.env

@@ -0,0 +1,3 @@
+
+ENV_VAR1=value1
+ENV_VAR2=value2

tests/fixtures/agentstack.json

@@ -0,0 +1,4 @@
+{
+    "framework": "crewai", 
+    "tools": ["tool1", "tool2"]
+}

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")
+