feat: add EventManager to centralize callbacks (#3434)

* refactor: Update MessageBase text attribute based on isinstance check. * feat: Add update_message function to update a message in the database. * refactor(chat): Update imports and remove unnecessary config method in ChatComponent. * refactor: Add stream_message method to ChatComponent. * refactor: Update method call in ChatOutput component. * feat: Add callback function to custom component and update build_results signature. * feat: Add callback parameter to instantiate_class function. * feat(graph): Add callback functions for sync and async operations. * feat: Add callback function support to vertex build process. * feat: Add handling for added message in InterfaceVertex class. * feat: Add callback support to Graph methods. * feat(chat): Add callback function to build_vertices function. * refactor: Simplify update_message function and use session_scope for session management. * fix: Call set_callback method if available on custom component. * refactor(chat): Update chat message chunk handling and ID conversion. * feat: Add null check before setting cache in build_vertex_stream function. * refactor: Fix send_event_wrapper function and add callback parameter to _build_vertex function. * refactor: Simplify conditional statement and import order in ChatOutput. * refactor: move log method to Component class. * refactor: Simplify CallbackFunction definition. * feat: Initialize _current_output attribute in Component class. * feat: store current output name in custom component during processing. * feat: Add current output and component ID to log data. * fix: Add condition to check current output before invoking callback. * refactor: Update callback to log_callback in graph methods. * feat: Add test for callback graph execution with log messages. * update projects * fix(chat.py): fix condition to check if message text is a string before updating message text in the database * refactor(ChatOutput.py): update ChatOutput class to correctly store and assign the message value to ensure consistency and avoid potential bugs * refactor(chat.py): update return type of store_message method to return a single Message object instead of a list of Messages refactor(chat.py): update logic to correctly handle updating and returning a single stored message object instead of a list of messages * update starter projects * refactor(component.py): update type hint for name parameter in log method to be more explicit * feat: Add EventManager class for managing events and event registration * refactor: Update log_callback to event_manager in custom component classes * refactor(component.py): rename _log_callback to _event_manager and update method call to on_log for better clarity and consistency * refactor(chat.py): rename _log_callback method to _event_manager.on_token for clarity and consistency in method naming * refactor: Rename log_callback to event_manager for clarity and consistency * refactor: Update Vertex class to use EventManager instead of log_callback for better clarity and consistency * refactor: update build_flow to use EventManager * refactor: Update EventManager class to use Protocol for event callbacks * if event_type is not passed, it uses the default send_event * Add method to register event functions in EventManager - Introduced `register_event_function` method to allow passing custom event functions. - Updated `noop` method to accept `event_type` parameter. - Adjusted `__getattr__` to return `EventCallback` type. * update test_callback_graph * Add unit tests for EventManager in test_event_manager.py - Added tests for event registration, including default event type, empty string names, and specific event types. - Added tests for custom event functions and unregistered event access. - Added tests for event sending, including JSON formatting, empty data, and large payloads. - Added tests for handling JSON serialization errors and the noop function. * revert chatOutput change * Add validation for event function in EventManager - Introduced `_validate_event_function` method to ensure event functions are callable and have the correct parameters. - Updated `register_event_function` to use the new validation method. * Add tests for EventManager's event function validation logic - Introduce `TestValidateEventFunction` class to test various scenarios for `_validate_event_function`. - Add tests for valid event functions, non-callable event functions, invalid parameter counts, and parameter type validation. - Include tests for handling unannotated parameters, flexible arguments, and keyword-only parameters. - Ensure proper warnings and exceptions are raised for invalid event functions. * Add type ignore comment to lambda function in test_event_manager.py * refactor: Update EventManager class to use Protocol for event callbacks * refactor(event_manager.py): simplify event registration and validation logic to enhance readability and maintainability feat(event_manager.py): enforce event name conventions and improve callback handling for better error management * refactor(chat.py): standardize event_manager method calls by using keyword arguments for better clarity and consistency refactor(chat.py): extract message processing logic into separate methods for improved readability and maintainability fix(chat.py): ensure proper handling of async iterators in message streaming refactor(component.py): simplify event logging by removing unnecessary event name parameter in on_log method call * update event manager tests * Add callback validation and manager parameter in EventManager - Introduced `_validate_callback` method to ensure callbacks are callable and have the correct parameters. - Updated `register_event` to include `manager` parameter in the callback. * Add support for passing callback through the Graph in test_callback_graph * fix(event_manager.py): update EventCallback signature to include manager parameter for better context in event handling
2024-09-02 10:41:19 -03:00 · 2024-09-02 10:41:19 -03:00 · 3eaad7bc3a
commit 3eaad7bc3a
parent 882c35ef6b
24 changed files with 534 additions and 106 deletions
--- a/src/backend/tests/unit/events/test_event_manager.py
+++ b/src/backend/tests/unit/events/test_event_manager.py
@ -0,0 +1,230 @@
+import asyncio
+import json
+import time
+import uuid
+
+import pytest
+
+from langflow.events.event_manager import EventManager
+from langflow.schema.log import LoggableType
+
+
+@pytest.fixture
+def client():
+    pass
+
+
+class TestEventManager:
+    # Registering an event with a valid name and callback using a mock callback function
+    def test_register_event_with_valid_name_and_callback_with_mock_callback(self):
+        def mock_callback(event_type: str, data: LoggableType):
+            pass
+
+        queue = asyncio.Queue()
+        manager = EventManager(queue)
+        manager.register_event("on_test_event", "test_type", mock_callback)
+        assert "on_test_event" in manager.events
+        assert manager.events["on_test_event"].func == mock_callback
+
+    # Registering an event with an empty name
+
+    def test_register_event_with_empty_name(self):
+        queue = asyncio.Queue()
+        manager = EventManager(queue)
+        with pytest.raises(ValueError, match="Event name cannot be empty"):
+            manager.register_event("", "test_type")
+
+    # Registering an event with a valid name and no callback
+    def test_register_event_with_valid_name_and_no_callback(self):
+        queue = asyncio.Queue()
+        manager = EventManager(queue)
+        manager.register_event("on_test_event", "test_type")
+        assert "on_test_event" in manager.events
+        assert manager.events["on_test_event"].func == manager.send_event
+
+    # Sending an event with valid event_type and data using pytest-asyncio plugin
+    @pytest.mark.asyncio
+    async def test_sending_event_with_valid_type_and_data_asyncio_plugin(self):
+        async def mock_queue_put_nowait(item):
+            await queue.put(item)
+
+        queue = asyncio.Queue()
+        queue.put_nowait = mock_queue_put_nowait
+        manager = EventManager(queue)
+        manager.register_event("on_test_event", "test_type", manager.noop)
+        event_type = "test_type"
+        data = "test_data"
+        manager.send_event(event_type=event_type, data=data)
+        await queue.join()
+        assert queue.empty()
+
+    # Accessing a non-registered event callback via __getattr__ with the recommended fix
+    def test_accessing_non_registered_event_callback_with_recommended_fix(self):
+        queue = asyncio.Queue()
+        manager = EventManager(queue)
+        result = manager.__getattr__("non_registered_event")
+        assert result == manager.noop
+
+    # Accessing a registered event callback via __getattr__
+    def test_accessing_registered_event_callback(self):
+        def mock_callback(event_type: str, data: LoggableType):
+            pass
+
+        queue = asyncio.Queue()
+        manager = EventManager(queue)
+        manager.register_event("on_test_event", "test_type", mock_callback)
+        assert manager.on_test_event.func == mock_callback
+
+    # Handling a large number of events in the queue
+    def test_handling_large_number_of_events(self):
+        async def mock_queue_put_nowait(item):
+            pass
+
+        queue = asyncio.Queue()
+        queue.put_nowait = mock_queue_put_nowait
+        manager = EventManager(queue)
+
+        for i in range(1000):
+            manager.register_event(f"on_test_event_{i}", "test_type", manager.noop)
+
+        assert len(manager.events) == 1000
+
+    # Testing registration of an event with an invalid name with the recommended fix
+    def test_register_event_with_invalid_name_fixed(self):
+        def mock_callback(event_type, data):
+            pass
+
+        queue = asyncio.Queue()
+        manager = EventManager(queue)
+        with pytest.raises(ValueError):
+            manager.register_event("", "test_type", mock_callback)
+        with pytest.raises(ValueError):
+            manager.register_event("invalid_name", "test_type", mock_callback)
+
+    # Sending an event with complex data and verifying successful event transmission
+    @pytest.mark.asyncio
+    async def test_sending_event_with_complex_data(self):
+        queue = asyncio.Queue()
+        manager = EventManager(queue)
+        manager.register_event("on_test_event", "test_type", manager.noop)
+        data = {"key": "value", "nested": [1, 2, 3]}
+        manager.send_event(event_type="test_type", data=data)
+        event_id, str_data, event_time = await queue.get()
+        assert event_id is not None
+        assert str_data is not None
+        assert event_time <= time.time()
+
+    # Sending an event with None data
+    def test_sending_event_with_none_data(self):
+        queue = asyncio.Queue()
+        manager = EventManager(queue)
+        manager.register_event("on_test_event", "test_type")
+        assert "on_test_event" in manager.events
+        assert manager.events["on_test_event"].func.__name__ == "send_event"
+
+    # Ensuring thread-safety when accessing the events dictionary
+    def test_thread_safety_accessing_events_dictionary(self):
+        def mock_callback(event_type: str, data: LoggableType):
+            pass
+
+        async def register_events(manager):
+            manager.register_event("on_test_event_1", "test_type_1", mock_callback)
+            manager.register_event("on_test_event_2", "test_type_2", mock_callback)
+
+        async def access_events(manager):
+            assert "on_test_event_1" in manager.events
+            assert "on_test_event_2" in manager.events
+
+        queue = asyncio.Queue()
+        manager = EventManager(queue)
+
+        tasks = [register_events(manager), access_events(manager)]
+        asyncio.get_event_loop().run_until_complete(asyncio.gather(*tasks))
+
+    # Checking the performance impact of frequent event registrations
+    def test_performance_impact_frequent_registrations(self):
+        async def mock_callback(event_type: str, data: LoggableType):
+            pass
+
+        queue = asyncio.Queue()
+        manager = EventManager(queue)
+        for i in range(1000):
+            manager.register_event(f"on_test_event_{i}", "test_type", mock_callback)
+        assert len(manager.events) == 1000
+
+    # Verifying the uniqueness of event IDs for each event triggered using await with asyncio decorator
+    import pytest
+
+    @pytest.mark.asyncio
+    async def test_event_id_uniqueness_with_await(self):
+        queue = asyncio.Queue()
+        manager = EventManager(queue)
+        manager.register_event("on_test_event", "test_type")
+        manager.on_test_event(data={"data_1": "value_1"})
+        manager.on_test_event(data={"data_2": "value_2"})
+        try:
+            event_id_1, _, _ = await queue.get()
+            event_id_2, _, _ = await queue.get()
+        except asyncio.TimeoutError:
+            pytest.fail("Test timed out while waiting for queue items")
+
+        assert event_id_1 != event_id_2
+
+    # Ensuring the queue receives the correct event data format
+    @pytest.mark.asyncio
+    async def test_queue_receives_correct_event_data_format(self):
+        async def mock_queue_put_nowait(data):
+            pass
+
+        async def mock_queue_get():
+            return (uuid.uuid4(), b'{"event": "test_type", "data": "test_data"}\n\n', time.time())
+
+        queue = asyncio.Queue()
+        queue.put_nowait = mock_queue_put_nowait
+        queue.get = mock_queue_get
+
+        manager = EventManager(queue)
+        manager.register_event("on_test_event", "test_type", manager.noop)
+        event_data = "test_data"
+        manager.send_event(event_type="test_type", data=event_data)
+
+        event_id, str_data, _ = await queue.get()
+        assert isinstance(event_id, uuid.UUID)
+        assert isinstance(str_data, bytes)
+        assert json.loads(str_data.decode("utf-8")) == {"event": "test_type", "data": event_data}
+
+    # Registering an event without specifying the event_type argument and providing the event_type argument
+    def test_register_event_without_event_type_argument_fixed(self):
+        class MockQueue:
+            def __init__(self):
+                self.data = []
+
+            def put_nowait(self, item):
+                self.data.append(item)
+
+        queue = MockQueue()
+        event_manager = EventManager(queue)
+        event_manager.register_event("on_test_event", "test_event_type", callback=event_manager.noop)
+        event_manager.send_event(event_type="test_type", data={"key": "value"})
+
+        assert len(queue.data) == 1
+        event_id, str_data, timestamp = queue.data[0]
+        assert isinstance(event_id, uuid.UUID)
+        assert isinstance(str_data, bytes)
+        assert isinstance(timestamp, float)
+
+    # Accessing a non-registered event callback via __getattr__
+    def test_accessing_non_registered_callback(self):
+        class MockQueue:
+            def __init__(self):
+                pass
+
+            def put_nowait(self, item):
+                pass
+
+        queue = MockQueue()
+        event_manager = EventManager(queue)
+
+        # Accessing a non-registered event callback should return the 'noop' function
+        callback = event_manager.on_non_existing_event
+        assert callback.__name__ == "noop"