From 972523465e704450aa9f8ff7e4a85ed3b31b982f Mon Sep 17 00:00:00 2001 From: vizsatiz Date: Tue, 28 Jan 2025 18:48:03 +0530 Subject: [PATCH 1/2] Fix for output collector --- flo_ai/state/__init__.py | 4 ++-- flo_ai/state/flo_json_output_collector.py | 9 ++++++--- flo_ai/state/flo_output_collector.py | 7 +++++++ 3 files changed, 15 insertions(+), 5 deletions(-) diff --git a/flo_ai/state/__init__.py b/flo_ai/state/__init__.py index fdd6c459..9830a9d0 100644 --- a/flo_ai/state/__init__.py +++ b/flo_ai/state/__init__.py @@ -1,4 +1,4 @@ from flo_ai.state.flo_json_output_collector import FloJsonOutputCollector -from flo_ai.state.flo_output_collector import FloOutputCollector +from flo_ai.state.flo_output_collector import FloOutputCollector, CollectionStatus -__all__ = ['FloJsonOutputCollector', 'FloOutputCollector'] +__all__ = ['FloJsonOutputCollector', 'FloOutputCollector', 'CollectionStatus'] diff --git a/flo_ai/state/flo_json_output_collector.py b/flo_ai/state/flo_json_output_collector.py index 896c5ed0..40b88fc9 100644 --- a/flo_ai/state/flo_json_output_collector.py +++ b/flo_ai/state/flo_json_output_collector.py @@ -3,13 +3,14 @@ from flo_ai.error.flo_exception import FloException from typing import Dict, List, Any from flo_ai.common.flo_logger import get_logger -from flo_ai.state.flo_output_collector import FloOutputCollector +from flo_ai.state.flo_output_collector import FloOutputCollector, CollectionStatus class FloJsonOutputCollector(FloOutputCollector): def __init__(self, strict: bool = False): super().__init__() self.strict = strict + self.status = CollectionStatus.success self.data: List[Dict[str, Any]] = [] def append(self, agent_output): @@ -70,9 +71,11 @@ def __extract_jsons(self, llm_response): 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}') - raise e + self.status = CollectionStatus.partial + get_logger().error(f'Invalid JSON in response: {json_str}, {e}') if self.strict and len(json_matches) == 0: + self.status = CollectionStatus.error + get_logger().error(f'Error while finding json in -- {llm_response}') raise FloException( 'JSON response expected in collector model: strict', error_code=1099 ) diff --git a/flo_ai/state/flo_output_collector.py b/flo_ai/state/flo_output_collector.py index 459ead20..5416acf9 100644 --- a/flo_ai/state/flo_output_collector.py +++ b/flo_ai/state/flo_output_collector.py @@ -1,6 +1,13 @@ +from enum import Enum from abc import ABC, abstractmethod +class CollectionStatus(Enum): + success = 'success' + partial = 'partial' + error = 'error' + + class FloOutputCollector(ABC): @abstractmethod def append(): From e8e67263680ee023de1fcbaf6fed4527a79d50ef Mon Sep 17 00:00:00 2001 From: vizsatiz Date: Tue, 28 Jan 2025 18:53:39 +0530 Subject: [PATCH 2/2] Fix for test --- tests/test_json_output_collection.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/tests/test_json_output_collection.py b/tests/test_json_output_collection.py index 28f4fe78..11325671 100644 --- a/tests/test_json_output_collection.py +++ b/tests/test_json_output_collection.py @@ -1,5 +1,4 @@ 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 @@ -103,11 +102,6 @@ def test_fetch_with_overlapping_keys(self, collector: FloJsonOutputCollector): 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 = """ {