CollectorClient returns the response as a dict

src/evidently/collector/client.py

@@ -1,15 +1,18 @@
+from typing import Any
+from typing import Dict
+
 import pandas as pd
 
 from evidently.collector.config import CollectorConfig
 from evidently.ui.utils import RemoteClientBase
 
 
 class CollectorClient(RemoteClientBase):
-    def create_collector(self, id: str, collector: CollectorConfig):
-        self._request(f"/{id}", "POST", body=collector.dict())
+    def create_collector(self, id: str, collector: CollectorConfig) -> Dict[str, Any]:
+        return self._request(f"/{id}", "POST", body=collector.dict()).json()
 
-    def send_data(self, id: str, data: pd.DataFrame):
-        self._request(f"/{id}/data", "POST", body=data.to_dict())
+    def send_data(self, id: str, data: pd.DataFrame) -> Dict[str, Any]:
+        return self._request(f"/{id}/data", "POST", body=data.to_dict()).json()
 
-    def set_reference(self, id: str, reference: pd.DataFrame):
-        self._request(f"/{id}/reference", "POST", body=reference.to_dict())
+    def set_reference(self, id: str, reference: pd.DataFrame) -> Dict[str, Any]:
+        return self._request(f"/{id}/reference", "POST", body=reference.to_dict()).json()

src/evidently/ui/utils.py

@@ -1,6 +1,9 @@
 import json
 import urllib.parse
 from typing import Optional
+from typing import Type
+from typing import TypeVar
+from typing import Union
 
 import requests
 
@@ -9,6 +12,8 @@
 from evidently.ui.storage.common import SECRET_HEADER_NAME
 from evidently.utils import NumpyEncoder
 
+T = TypeVar("T", bound=BaseModel)
+
 
 class RemoteClientBase:
     def __init__(self, base_url: str, secret: str = None):
@@ -21,8 +26,8 @@ def _request(
         method: str,
         query_params: Optional[dict] = None,
         body: Optional[dict] = None,
-        response_model=None,
-    ):
+        response_model: Optional[Type[T]] = None,
+    ) -> Union[T, requests.Response]:
         # todo: better encoding
         headers = {SECRET_HEADER_NAME: self.secret}
         data = None