update code

tests/metagpt/actions/test_sd_design.py

@@ -0,0 +1,54 @@
+# -*- coding: utf-8 -*-
+# @Date    : 2023/7/22 02:40
+# @Author  : stellahong (stellahong@fuzhi.ai)
+#
+
+import pytest
+from typing import List
+from metagpt.actions.design import SDPromptOptimize, SDPromptImprove
+from metagpt.actions.ui_design import ModelSelection
+
+
+@pytest.mark.asyncio
+async def test_ui_model_selection():
+    ms = ModelSelection()
+    model_name, domain = await ms.run("Pokémon Go")
+    assert model_name == "pixelmix_v10"
+
+
+@pytest.mark.asyncio
+async def test_ui_sd_generation():
+    pass
+
+
+@pytest.mark.asyncio
+async def test_ui_sd_prompt_optimize():
+    sd_po = SDPromptOptimize()
+    resp = await sd_po.run(query="Pokémon Go", domain="Anime", answer_count=1)
+    assert type(resp) == List
+    assert len(resp) == 1
+
+
+@pytest.mark.asyncio
+async def test_ui_sd_optimize_answer_count():
+    sd_po = SDPromptOptimize()
+    answer_count = 2
+    resp = await sd_po.run(query="Pokémon Go", domain="Anime", answer_count=2)
+    assert type(resp) == List
+    assert len(resp) == answer_count
+
+@pytest.mark.asyncio
+async def test_ui_sd_improve_answer_count():
+    sd_pi = SDPromptImprove()
+    answer_count = 2
+    resp = await sd_pi.run(query="Pokémon Go", domain="Anime", answer_count=2)
+    assert type(resp) == List
+    assert len(resp) == answer_count
+    
+
+@pytest.mark.asyncio
+async def test_ui_sd_prompt_improve():
+    sd_pi = SDPromptImprove()
+    resp = await sd_pi.run(query="Pokémon Go", domain="Anime", answer_count=1)
+    assert type(resp) == List
+    assert len(resp) == 1

tests/metagpt/utils/test_flatten_json_object.py

@@ -0,0 +1,39 @@
+import unittest
+import json5
+
+def flatten_json_object(obj, parent_key='', sep=', '):
+    if isinstance(obj, str):
+        return dict([("value", obj)])
+    
+    if isinstance(obj, list):
+        return dict([("value", sep.join(str(v) for v in obj))])
+    
+    items = []
+    for key, value in obj.items():
+        new_key = f"{parent_key}{sep}{key}" if parent_key else key
+        if isinstance(value, dict):
+            items.extend(flatten_json_object(value, new_key, sep=sep).items())
+        elif isinstance(value, list):
+            items.append((new_key, sep.join(str(v) for v in value)))
+        else:
+            items.append((new_key, value))
+    return dict(items)
+
+class TestFlattenJsonObject(unittest.TestCase):
+    def test_flatten_json_object(self):
+        json_obj = json5.loads('{"a": 1, "b": {"c": 2, "d": {"e": 3, "f": 4}}, "g": [5, 6, 7]}')
+        expected_result = {'a': 1, 'b, c': 2, 'b, d, e': 3, 'b, d, f': 4, 'g': '5, 6, 7'}
+        self.assertEqual(flatten_json_object(json_obj), expected_result)
+
+    def test_flatten_json_object_with_string(self):
+        json_obj = json5.loads('{"a": "hello"}')
+        expected_result = {'a': 'hello'}
+        self.assertEqual(flatten_json_object(json_obj), expected_result)
+
+    def test_flatten_json_object_with_list(self):
+        json_obj = json5.loads('{"a": [1, 2, 3]}')
+        expected_result = {'a': '1, 2, 3'}
+        self.assertEqual(flatten_json_object(json_obj), expected_result)
+
+if __name__ == '__main__':
+    unittest.main()

tests/metagpt/utils/test_flatten_json_structure_json.py

@@ -0,0 +1,84 @@
+import unittest
+import json5
+import re
+
+
+def flatten_json_object(obj, parent_key='', sep=', '):
+    if isinstance(obj, str):
+        return dict([("value", obj)])
+    
+    if isinstance(obj, list):
+        return dict([("value", sep.join(str(v) for v in obj))])
+    
+    items = []
+    for key, value in obj.items():
+        new_key = f"{parent_key}{sep}{key}" if parent_key else key
+        if isinstance(value, dict):
+            items.extend(flatten_json_object(value, new_key, sep=sep).items())
+        elif isinstance(value, list):
+            items.append((new_key, sep.join(str(v) for v in value)))
+        else:
+            items.append((new_key, value))
+    return dict(items)
+
+def flatten_json_structure(json_array):
+    if (isinstance(json_array, list) and len(json_array) == 1 and not isinstance(json_array[0], str)):
+        return flatten_json_structure(json_array[0])
+    
+    if (isinstance(json_array, dict) and len(json_array.values()) == 1 and not isinstance(list(json_array.values())[0],
+                                                                                          str)):
+        return flatten_json_structure(list(json_array.values())[0])
+    
+    flattened_json_array = []
+    
+    if (isinstance(json_array, dict)):
+        json_array = json_array.values()
+    
+    for json_object in json_array:
+        flattened_dict = flatten_json_object(json_object)
+        flattened_values = ", ".join(str(v) for v in flattened_dict.values())
+        flattened_json_array.append(flattened_values)
+    
+    return flattened_json_array
+
+class TestFlattenJson(unittest.TestCase):
+    def test_flatten_json_structure(self):
+        input_json = [
+            {
+                "name": "John",
+                "age": 30,
+                "city": "New York"
+            },
+            {
+                "name": "Jane",
+                "age": 25,
+                "city": "Chicago"
+            }
+        ]
+        expected_output = ["John, 30, New York", "Jane, 25, Chicago"]
+        self.assertEqual(flatten_json_structure(input_json), expected_output)
+
+    def test_flatten_json_structure_with_nested_json(self):
+        input_json = [
+            {
+                "name": "John",
+                "age": 30,
+                "address": {
+                    "city": "New York",
+                    "state": "NY"
+                }
+            },
+            {
+                "name": "Jane",
+                "age": 25,
+                "address": {
+                    "city": "Chicago",
+                    "state": "IL"
+                }
+            }
+        ]
+        expected_output = ["John, 30, New York, NY", "Jane, 25, Chicago, IL"]
+        self.assertEqual(flatten_json_structure(input_json), expected_output)
+
+if __name__ == '__main__':
+    unittest.main()

tests/metagpt/utils/test_try_parse_json.py

@@ -0,0 +1,58 @@
+import unittest
+import json5
+import re
+
+def try_parse_json(input_text):
+    input_text.index
+    start_index_brackets = input_text.find('[')
+    end_index_brackets = input_text.rfind(']')
+    start_index_curly = input_text.find('{')
+    end_index_curly = input_text.rfind('}')
+    
+    start_index = start_index_brackets
+    end_index = end_index_brackets
+    
+    if (start_index_curly != -1 and (start_index_curly < start_index_brackets or start_index_brackets < 0)):
+        start_index = start_index_curly
+        end_index = end_index_curly
+    
+    if start_index >= 0 and end_index > 0:
+        json_string = input_text[start_index:end_index + 1]
+        json_string = re.sub(r'\}[\s]*\{', '}, {', json_string)
+        json_string = re.sub(r'\][\s]*\[', '], [', json_string)
+        json_string = re.sub(r'"[\s]*"', '", "', json_string)
+        
+        try:
+            json_object = json5.loads(json_string)
+        except ValueError:
+            json_object = json5.loads(f"[{json_string}]")
+        
+        return json_object
+    
+    raise Exception("No JSON object found in input text.")
+
+
+class TestTryParseJson(unittest.TestCase):
+    def test_valid_json(self):
+        input_text = '{"name": "John", "age": 30, "city": "New York"}'
+        expected_output = {"name": "John", "age": 30, "city": "New York"}
+        self.assertEqual(try_parse_json(input_text), expected_output)
+
+    def test_invalid_json(self):
+        input_text = 'This is not a JSON string'
+        with self.assertRaises(Exception) as context:
+            try_parse_json(input_text)
+        self.assertTrue('No JSON object found in input text.' in str(context.exception))
+
+    def test_empty_json(self):
+        input_text = '{}'
+        expected_output = {}
+        self.assertEqual(try_parse_json(input_text), expected_output)
+
+    def test_nested_json(self):
+        input_text = '{"name": "John", "age": 30, "city": "New York", "friends": ["Mike", "Anna"]}'
+        expected_output = {"name": "John", "age": 30, "city": "New York", "friends": ["Mike", "Anna"]}
+        self.assertEqual(try_parse_json(input_text), expected_output)
+
+if __name__ == '__main__':
+    unittest.main()