fix(Monitoring): Argilla log middleware can be serialized using pickle

Also, review and format the class name
argilla-io · Nov 15, 2022 · 16a12d2 · 16a12d2
1 parent 91a72c5
commit 16a12d2
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/src/argilla/monitoring/asgi.py b/src/argilla/monitoring/asgi.py
@@ -101,7 +101,7 @@ async def cached_receive() -> Message:
         return self._receive
 
 
-class argillaLogHTTPMiddleware(BaseHTTPMiddleware):
+class ArgillaLogHTTPMiddleware(BaseHTTPMiddleware):
     """An standard starlette middleware that enables argilla logs for http prediction requests"""
 
     def __init__(
@@ -118,7 +118,9 @@ def __init__(
         self._records_mapper = records_mapper or text_classification_mapper
         self._queue = Queue()
         self._worker_task = threading.Thread(
-            target=self.__worker__, name=argillaLogHTTPMiddleware.__name__, daemon=True
+            target=self.__worker__,
+            name=ArgillaLogHTTPMiddleware.__name__,
+            daemon=True,
         )
         self._worker_task.start()
 

diff --git a/tests/client/test_asgi.py b/tests/client/test_asgi.py
@@ -23,7 +23,7 @@
 
 import argilla
 from argilla.monitoring.asgi import (
-    argillaLogHTTPMiddleware,
+    ArgillaLogHTTPMiddleware,
     token_classification_mapper,
 )
 
@@ -35,7 +35,7 @@ def test_argilla_middleware_for_text_classification(monkeypatch):
 
     app = FastAPI()
     app.add_middleware(
-        argillaLogHTTPMiddleware,
+        ArgillaLogHTTPMiddleware,
         api_endpoint=expected_endpoint,
         dataset=expected_dataset_name,
     )
@@ -92,7 +92,7 @@ def test_argilla_middleware_for_token_classification(monkeypatch):
 
     app = Starlette()
     app.add_middleware(
-        argillaLogHTTPMiddleware,
+        ArgillaLogHTTPMiddleware,
         api_endpoint=expected_endpoint,
         dataset=expected_dataset_name,
         records_mapper=token_classification_mapper,