From 064ac5c7425de91a75c3ef2f2911bb7d942b223f Mon Sep 17 00:00:00 2001
From: Luke K <2609441+pr-0f3t@users.noreply.github.com>
Date: Sun, 30 Apr 2023 06:37:41 +0100
Subject: [PATCH] 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>
---
 autogpt/config/ai_config.py |  9 ++++++--
 tests/test_ai_config.py     | 45 +++++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_ai_config.py

diff --git a/autogpt/config/ai_config.py b/autogpt/config/ai_config.py
index d662429f..88acbfe6 100644
--- a/autogpt/config/ai_config.py
+++ b/autogpt/config/ai_config.py
@@ -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)
diff --git a/tests/test_ai_config.py b/tests/test_ai_config.py
new file mode 100644
index 00000000..a9fcdad6
--- /dev/null
+++ b/tests/test_ai_config.py
@@ -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