From bcd90e26b01d1838c686c106d4fb680561f82154 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Moura?= Date: Wed, 19 Feb 2025 12:54:15 -0800 Subject: [PATCH] making flow verbsoe false by default --- src/crewai/flow/persistence/decorators.py | 24 +++++++------- tests/test_flow_persistence.py | 38 ++++++++++++++++++++++- 2 files changed, 50 insertions(+), 12 deletions(-) diff --git a/src/crewai/flow/persistence/decorators.py b/src/crewai/flow/persistence/decorators.py index ebf3778b7..7b3bd447c 100644 --- a/src/crewai/flow/persistence/decorators.py +++ b/src/crewai/flow/persistence/decorators.py @@ -58,7 +58,7 @@ class PersistenceDecorator: _printer = Printer() # Class-level printer instance @classmethod - def persist_state(cls, flow_instance: Any, method_name: str, persistence_instance: FlowPersistence) -> None: + def persist_state(cls, flow_instance: Any, method_name: str, persistence_instance: FlowPersistence, verbose: bool = False) -> None: """Persist flow state with proper error handling and logging. This method handles the persistence of flow state data, including proper @@ -68,6 +68,7 @@ class PersistenceDecorator: flow_instance: The flow instance whose state to persist method_name: Name of the method that triggered persistence persistence_instance: The persistence backend to use + verbose: Whether to log persistence operations Raises: ValueError: If flow has no state or state lacks an ID @@ -88,9 +89,10 @@ class PersistenceDecorator: if not flow_uuid: raise ValueError("Flow state must have an 'id' field for persistence") - # Log state saving with consistent message - cls._printer.print(LOG_MESSAGES["save_state"].format(flow_uuid), color="cyan") - logger.info(LOG_MESSAGES["save_state"].format(flow_uuid)) + # Log state saving only if verbose is True + if verbose: + cls._printer.print(LOG_MESSAGES["save_state"].format(flow_uuid), color="cyan") + logger.info(LOG_MESSAGES["save_state"].format(flow_uuid)) try: persistence_instance.save_state( @@ -115,7 +117,7 @@ class PersistenceDecorator: raise ValueError(error_msg) from e -def persist(persistence: Optional[FlowPersistence] = None): +def persist(persistence: Optional[FlowPersistence] = None, verbose: bool = False): """Decorator to persist flow state. This decorator can be applied at either the class level or method level. @@ -126,6 +128,7 @@ def persist(persistence: Optional[FlowPersistence] = None): Args: persistence: Optional FlowPersistence implementation to use. If not provided, uses SQLiteFlowPersistence. + verbose: Whether to log persistence operations. Defaults to False. Returns: A decorator that can be applied to either a class or method @@ -135,13 +138,12 @@ def persist(persistence: Optional[FlowPersistence] = None): RuntimeError: If state persistence fails Example: - @persist # Class-level persistence with default SQLite + @persist(verbose=True) # Class-level persistence with logging class MyFlow(Flow[MyState]): @start() def begin(self): pass """ - def decorator(target: Union[Type, Callable[..., T]]) -> Union[Type, Callable[..., T]]: """Decorator that handles both class and method decoration.""" actual_persistence = persistence or SQLiteFlowPersistence() @@ -179,7 +181,7 @@ def persist(persistence: Optional[FlowPersistence] = None): @functools.wraps(original_method) async def method_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: result = await original_method(self, *args, **kwargs) - PersistenceDecorator.persist_state(self, method_name, actual_persistence) + PersistenceDecorator.persist_state(self, method_name, actual_persistence, verbose) return result return method_wrapper @@ -199,7 +201,7 @@ def persist(persistence: Optional[FlowPersistence] = None): @functools.wraps(original_method) def method_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: result = original_method(self, *args, **kwargs) - PersistenceDecorator.persist_state(self, method_name, actual_persistence) + PersistenceDecorator.persist_state(self, method_name, actual_persistence, verbose) return result return method_wrapper @@ -228,7 +230,7 @@ def persist(persistence: Optional[FlowPersistence] = None): result = await method_coro else: result = method_coro - PersistenceDecorator.persist_state(flow_instance, method.__name__, actual_persistence) + PersistenceDecorator.persist_state(flow_instance, method.__name__, actual_persistence, verbose) return result for attr in ["__is_start_method__", "__trigger_methods__", "__condition_type__", "__is_router__"]: @@ -240,7 +242,7 @@ def persist(persistence: Optional[FlowPersistence] = None): @functools.wraps(method) def method_sync_wrapper(flow_instance: Any, *args: Any, **kwargs: Any) -> T: result = method(flow_instance, *args, **kwargs) - PersistenceDecorator.persist_state(flow_instance, method.__name__, actual_persistence) + PersistenceDecorator.persist_state(flow_instance, method.__name__, actual_persistence, verbose) return result for attr in ["__is_start_method__", "__trigger_methods__", "__condition_type__", "__is_router__"]: diff --git a/tests/test_flow_persistence.py b/tests/test_flow_persistence.py index e51806b05..b6151de84 100644 --- a/tests/test_flow_persistence.py +++ b/tests/test_flow_persistence.py @@ -17,7 +17,7 @@ class TestState(FlowState): message: str = "" -def test_persist_decorator_saves_state(tmp_path): +def test_persist_decorator_saves_state(tmp_path, caplog): """Test that @persist decorator saves state in SQLite.""" db_path = os.path.join(tmp_path, "test_flows.db") persistence = SQLiteFlowPersistence(db_path) @@ -174,3 +174,39 @@ def test_multiple_method_persistence(tmp_path): final_state = flow2.state assert final_state.counter == 99999 assert final_state.message == "Step 99999" + +def test_persist_decorator_verbose_logging(tmp_path, caplog): + """Test that @persist decorator's verbose parameter controls logging.""" + db_path = os.path.join(tmp_path, "test_flows.db") + persistence = SQLiteFlowPersistence(db_path) + + # Test with verbose=False (default) + class QuietFlow(Flow[Dict[str, str]]): + initial_state = dict() + + @start() + @persist(persistence) # Default verbose=False + def init_step(self): + self.state["message"] = "Hello, World!" + self.state["id"] = "test-uuid-1" + + flow = QuietFlow(persistence=persistence) + flow.kickoff() + assert "Saving flow state to memory for ID: test-uuid-1" not in caplog.text + + # Clear the log + caplog.clear() + + # Test with verbose=True + class VerboseFlow(Flow[Dict[str, str]]): + initial_state = dict() + + @start() + @persist(persistence, verbose=True) + def init_step(self): + self.state["message"] = "Hello, World!" + self.state["id"] = "test-uuid-2" + + flow = VerboseFlow(persistence=persistence) + flow.kickoff() + assert "Saving flow state to memory for ID: test-uuid-2" in caplog.text