feat: add observability plugin system (#227)

Co-authored-by: Michael Neale <michael.neale@gmail.com>
Co-authored-by: Lifei Zhou <lifei@squareup.com>
Co-authored-by: Alice Hau <ahau@squareup.com>

packages/exchange/tests/test_langfuse_wrapper.py

@@ -1,48 +0,0 @@
-import pytest
-from unittest.mock import patch, MagicMock
-from exchange.langfuse_wrapper import observe_wrapper
-
-
-@pytest.fixture
-def mock_langfuse_context():
-    with patch("exchange.langfuse_wrapper.langfuse_context") as mock:
-        yield mock
-
-
-@patch("exchange.langfuse_wrapper.auth_check")
-def test_function_is_wrapped(mock_auth_check, mock_langfuse_context):
-    mock_observe = MagicMock(side_effect=lambda *args, **kwargs: lambda fn: fn)
-    mock_auth_check.return_value = True
-    mock_langfuse_context.observe = mock_observe
-
-    def original_function(x: int, y: int) -> int:
-        return x + y
-
-    # test function before we decorate it with
-    # @observe_wrapper("arg1", kwarg1="kwarg1")
-    assert not hasattr(original_function, "__wrapped__")
-
-    # ensure we args get passed along (e.g. @observe(capture_input=False, capture_output=False))
-    decorated_function = observe_wrapper("arg1", kwarg1="kwarg1")(original_function)
-    assert hasattr(decorated_function, "__wrapped__")
-    assert decorated_function.__wrapped__ is original_function, "Function is not properly wrapped"
-
-    assert decorated_function(2, 3) == 5
-    mock_observe.assert_called_once()
-    mock_observe.assert_called_with("arg1", kwarg1="kwarg1")
-
-
-@patch("exchange.langfuse_wrapper.auth_check")
-def test_function_is_not_wrapped(mock_auth_check, mock_langfuse_context):
-    mock_observe = MagicMock(return_value=lambda f: f)
-    mock_auth_check.return_value = False
-    mock_langfuse_context.observe = mock_observe
-
-    @observe_wrapper("arg1", kwarg1="kwarg1")
-    def hello() -> str:
-        return "Hello"
-
-    assert not hasattr(hello, "__wrapped__")
-    assert hello() == "Hello"
-
-    mock_observe.assert_not_called()

packages/exchange/tests/test_observer.py

@@ -0,0 +1,61 @@
+from exchange.observers import ObserverManager, observe_wrapper
+from exchange.observers.base import Observer
+
+
+class MockObserver(Observer):
+    def __init__(self):
+        self.initialized = False
+        self.args = None
+        self.kwargs = None
+        self.finalized = False
+
+    def initialize(self):
+        pass
+
+    def observe_wrapper(self, *args, **kwargs):
+        def wrapper(func):
+            self.args = args
+            self.kwargs = kwargs
+            return func
+
+        return wrapper
+
+    def finalize(self):
+        pass
+
+
+def test_wrapper_is_invoked():
+    manager = ObserverManager.get_instance()
+    mock_observer = MockObserver()
+    manager.initialize(True, [mock_observer])
+
+    @observe_wrapper("arg0", arg1="arg2")
+    def wrapped(x: int, y: int) -> int:
+        return x + y
+
+    # code in decorator hasn't run yet
+    assert mock_observer.args is None
+    assert mock_observer.kwargs is None
+
+    ret_val = wrapped(2, 3)
+    assert ret_val == 5
+
+    # decorator has been run since `wrapped` was called
+    assert mock_observer.args == ("arg0",)
+    assert mock_observer.kwargs == {"arg1": "arg2"}
+
+
+def test_multiple_wrappers():
+    manager = ObserverManager.get_instance()
+    mock_observer_1 = MockObserver()
+    mock_observer_2 = MockObserver()
+    manager.initialize(True, [mock_observer_1, mock_observer_2])
+
+    @observe_wrapper("arg0")
+    def wrapped(x: int, y: int) -> int:
+        return x + y
+
+    wrapped(2, 3)
+
+    assert mock_observer_1.args == ("arg0",)
+    assert mock_observer_2.args == ("arg0",)