From ce6bd1ee1e32ccf2fb4d199727d8dd3f0881a202 Mon Sep 17 00:00:00 2001
From: Danil Ustinov <vk.vx@yandex.ru>
Date: Tue, 5 Mar 2024 23:01:29 +0800
Subject: [PATCH] CollectorClient returns the response as a dict

---
 src/evidently/collector/client.py | 15 +++++++++------
 src/evidently/ui/utils.py         |  9 +++++++--
 2 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/src/evidently/collector/client.py b/src/evidently/collector/client.py
index ced8269ce1..98b417abc1 100644
--- a/src/evidently/collector/client.py
+++ b/src/evidently/collector/client.py
@@ -1,3 +1,6 @@
+from typing import Any
+from typing import Dict
+
 import pandas as pd
 
 from evidently.collector.config import CollectorConfig
@@ -5,11 +8,11 @@
 
 
 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()
diff --git a/src/evidently/ui/utils.py b/src/evidently/ui/utils.py
index f4e3eccd4b..4902af3f3a 100644
--- a/src/evidently/ui/utils.py
+++ b/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