diff --git a/flo_ai/state/flo_json_output_collector.py b/flo_ai/state/flo_json_output_collector.py index 71eeaf18..896c5ed0 100644 --- a/flo_ai/state/flo_json_output_collector.py +++ b/flo_ai/state/flo_json_output_collector.py @@ -15,13 +15,59 @@ def __init__(self, strict: bool = False): def append(self, agent_output): self.data.append(self.__extract_jsons(agent_output)) + def __strip_comments(self, json_str: str) -> str: + cleaned = [] + length = len(json_str) + i = 0 + + while i < length: + char = json_str[i] + + if char not in '"/*': + cleaned.append(char) + i += 1 + continue + + if char == '"': + cleaned.append(char) + i += 1 + + while i < length: + char = json_str[i] + cleaned.append(char) + i += 1 + if char == '"' and json_str[i - 2] != '\\': + break + continue + + if char == '/' and i + 1 < length: + next_char = json_str[i + 1] + + if next_char == '/': + i += 2 + while i < length and json_str[i] != '\n': + i += 1 + continue + elif next_char == '*': + i += 2 + while i + 1 < length: + if json_str[i] == '*' and json_str[i + 1] == '/': + i += 2 + break + i += 1 + continue + + cleaned.append(char) + i += 1 + return ''.join(cleaned) + def __extract_jsons(self, llm_response): json_pattern = r'\{(?:[^{}]|(?R))*\}' json_matches = regex.findall(json_pattern, llm_response) json_object = {} for json_str in json_matches: try: - json_obj = json.loads(json_str) + json_obj = json.loads(self.__strip_comments(json_str)) json_object.update(json_obj) except json.JSONDecodeError as e: get_logger().error(f'Invalid JSON in response: {json_str}') diff --git a/pyproject.toml b/pyproject.toml index 61b38997..dd8c4d5d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "flo-ai" -version = "0.0.5-rc2" +version = "0.0.5-rc3" description = "A easy way to create structured AI agents" authors = ["vizsatiz "] license = "MIT" diff --git a/setup.py b/setup.py index 65e80a59..77c24d71 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setuptools.setup( name='flo-ai', - version='0.0.5-rc2', + version='0.0.5-rc3', author='Rootflo', description='Create composable AI agents', long_description=long_description, diff --git a/tests/test_json_output_collection.py b/tests/test_json_output_collection.py new file mode 100644 index 00000000..28f4fe78 --- /dev/null +++ b/tests/test_json_output_collection.py @@ -0,0 +1,144 @@ +import pytest +import json +from flo_ai.error.flo_exception import FloException +from flo_ai.state.flo_output_collector import FloOutputCollector +from flo_ai.state.flo_json_output_collector import FloJsonOutputCollector + + +class TestFloJsonOutputCollector: + @pytest.fixture + def collector(self): + return FloJsonOutputCollector(strict=False) + + @pytest.fixture + def strict_collector(self): + return FloJsonOutputCollector(strict=True) + + def test_initialization(self, collector): + assert isinstance(collector, FloOutputCollector) + assert collector.strict is False + assert collector.data == [] + + def test_append_single_json(self, collector): + test_input = '{"key": "value"}' + collector.append(test_input) + assert collector.data == [{'key': 'value'}] + + def test_append_multiple_jsons(self, collector): + test_input = '{"key1": "value1"} Some text {"key2": "value2"}' + collector.append(test_input) + assert collector.data == [{'key1': 'value1', 'key2': 'value2'}] + + def test_append_nested_json(self, collector): + test_input = '{"outer": {"inner": "value"}}' + collector.append(test_input) + assert collector.data == [{'outer': {'inner': 'value'}}] + + def test_strip_comments(self, collector): + test_input = """ + { + // Single line comment + "key1": "value1", + /* Multi-line + comment */ + "key2": "value2" + } + """ + collector.append(test_input) + assert collector.data == [{'key1': 'value1', 'key2': 'value2'}] + + def test_string_with_comment_chars(self, collector): + test_input = '{"key": "This // is not a comment", "url": "http://example.com"}' + collector.append(test_input) + assert collector.data == [ + {'key': 'This // is not a comment', 'url': 'http://example.com'} + ] + + def test_strict_mode_no_json(self, strict_collector): + with pytest.raises(FloException) as exc_info: + strict_collector.append('No JSON here') + assert exc_info.value.error_code == 1099 + + def test_strict_mode_with_json(self, strict_collector): + test_input = '{"key": "value"}' + strict_collector.append(test_input) + assert strict_collector.data == [{'key': 'value'}] + + def test_pop_operation(self, collector: FloJsonOutputCollector): + test_input1 = '{"key1": "value1"}' + test_input2 = '{"key2": "value2"}' + collector.append(test_input1) + collector.append(test_input2) + + popped = collector.pop() + assert popped == {'key2': 'value2'} + assert len(collector.data) == 1 + + def test_peek_operation(self, collector: FloJsonOutputCollector): + test_input = '{"key": "value"}' + collector.append(test_input) + + peeked = collector.peek() + assert peeked == {'key': 'value'} + assert len(collector.data) == 1 + + def test_peek_empty_collector(self, collector): + assert collector.peek() is None + + def test_fetch_operation(self, collector: FloJsonOutputCollector): + test_input1 = '{"key1": "value1"}' + test_input2 = '{"key2": "value2"}' + collector.append(test_input1) + collector.append(test_input2) + + result = collector.fetch() + assert result == {'key1': 'value1', 'key2': 'value2'} + + def test_fetch_with_overlapping_keys(self, collector: FloJsonOutputCollector): + test_input1 = '{"key": "value1"}' + test_input2 = '{"key": "value2"}' + collector.append(test_input1) + collector.append(test_input2) + + result = collector.fetch() + assert result == {'key': 'value2'} # Later values should override earlier ones + + def test_invalid_json(self, collector: FloJsonOutputCollector): + test_input = '{"key": "value",}' # Invalid JSON with trailing comma + with pytest.raises(json.JSONDecodeError): + collector.append(test_input) + + def test_complex_nested_structure(self, collector: FloJsonOutputCollector): + test_input = """ + { + "array": [1, 2, 3], + "nested": { + "deep": { + "deeper": "value" + } + }, + "mixed": [{"key": "value"}, 42, "string"] + } + """ + collector.append(test_input) + expected = { + 'array': [1, 2, 3], + 'nested': {'deep': {'deeper': 'value'}}, + 'mixed': [{'key': 'value'}, 42, 'string'], + } + assert collector.data == [expected] + + @pytest.mark.parametrize( + 'test_input,expected', + [ + ('{"a": 1}', [{'a': 1}]), + ('{"a": 1, "b": 2}', [{'a': 1, 'b': 2}]), + ('{"a": 1} {"b": 2}', [{'a': 1, 'b': 2}]), + ('No JSON', [{}]), + ], + ) + def test_various_inputs( + self, collector: FloJsonOutputCollector, test_input, expected + ): + collector.append(test_input) + assert collector.data == expected