Refactor AIConfig to Sanitize Input for Goal Parameters (#3492)

* Update remove_color_codes to handle non-string input

The `remove_color_codes` function now accepts any type of input that can be cast to a string. Previously, it was only accepting string input and not casting non-string types to string which was causing errors in some cases.

The changes were made to both logs.py and its corresponding test file.

* Refactor AIConfig to Sanitize Input for Goal Parameters

Details:
- Modified `ai_config.py` to correctly handle and sanitize user input for AI goals and convert them to formatted strings, to fix an issue where some specially formatted ai_settings.yaml files were causing goals to load as list[dict]
- `test_ai_config.py` includes a test for the `sanitize_input` function in `AIConfig` class.
- Removed unnecessary tests from `test_logs.py`

* Update for readabiity

* Update for readabiity

* Updates for conciceness

* Updated tests to confirm AIConfig saves goals as strings

* FIxed trailing space at end of line

---------

Co-authored-by: Luke Kyohere <lkyohere@mfsafrica.com>
Co-authored-by: James Collins <collijk@uw.edu>
This commit is contained in:
Luke K
2023-04-30 06:37:41 +01:00
committed by GitHub
parent 06ae4684c8
commit 064ac5c742
2 changed files with 52 additions and 2 deletions

View File

@@ -7,7 +7,7 @@ from __future__ import annotations
import os
import platform
from pathlib import Path
from typing import Optional, Type
from typing import Any, Optional, Type
import distro
import yaml
@@ -79,7 +79,12 @@ class AIConfig:
ai_name = config_params.get("ai_name", "")
ai_role = config_params.get("ai_role", "")
ai_goals = config_params.get("ai_goals", [])
ai_goals = [
str(goal).strip("{}").replace("'", "").replace('"', "")
if isinstance(goal, dict)
else str(goal)
for goal in config_params.get("ai_goals", [])
]
api_budget = config_params.get("api_budget", 0.0)
# type: Type[AIConfig]
return AIConfig(ai_name, ai_role, ai_goals, api_budget)

45
tests/test_ai_config.py Normal file
View File

@@ -0,0 +1,45 @@
from autogpt.config.ai_config import AIConfig
"""
Test cases for the AIConfig class, which handles loads the AI configuration
settings from a YAML file.
"""
def test_goals_are_always_lists_of_strings(tmp_path):
"""Test if the goals attribute is always a list of strings."""
yaml_content = """
ai_goals:
- Goal 1: Make a sandwich
- Goal 2, Eat the sandwich
- Goal 3 - Go to sleep
- "Goal 4: Wake up"
ai_name: McFamished
ai_role: A hungry AI
api_budget: 0.0
"""
config_file = tmp_path / "ai_settings.yaml"
config_file.write_text(yaml_content)
ai_config = AIConfig.load(config_file)
assert len(ai_config.ai_goals) == 4
assert ai_config.ai_goals[0] == "Goal 1: Make a sandwich"
assert ai_config.ai_goals[1] == "Goal 2, Eat the sandwich"
assert ai_config.ai_goals[2] == "Goal 3 - Go to sleep"
assert ai_config.ai_goals[3] == "Goal 4: Wake up"
config_file.write_text("")
ai_config.save(config_file)
yaml_content2 = """ai_goals:
- 'Goal 1: Make a sandwich'
- Goal 2, Eat the sandwich
- Goal 3 - Go to sleep
- 'Goal 4: Wake up'
ai_name: McFamished
ai_role: A hungry AI
api_budget: 0.0
"""
assert config_file.read_text() == yaml_content2