dream

1
#!/usr/bin/env python
2
import time
3
import asyncio
4
import logging
5
import requests
6
import sentry_sdk
7

8
from typing import Callable, Dict
9
from os import getenv
10

11

12
logging.basicConfig(format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO)
13
logger = logging.getLogger(__name__)
14

15
sentry_sdk.init(getenv("SENTRY_DSN"))
16
headers = {"Content-Type": "application/json;charset=utf-8"}
17

18

19
class BatchConnector:
20
    def __init__(self, url: str):
21
        self._url = url
22

23
    async def send(self, payload: Dict, callback: Callable):
24
        try:
25
            st_time = time.time()
26
            toxic_result = requests.request(
27
                url=self._url, headers=headers, json=payload["payload"], method="POST", timeout=1.0
28
            ).json()
29
            result = [res[0] for res in toxic_result]
30
            total_time = time.time() - st_time
31
            logger.info(f"DeepPavlovToxicClassification batch connector exec time: {total_time:.3f}s")
32
            # In connector [result] leads to bug, so it is not inside array like on
33
            # conv eval and badlist annotator
34
            asyncio.create_task(callback(task_id=payload["task_id"], response={"batch": result}))
35
        except Exception as e:
36
            logger.exception(e)
37
            sentry_sdk.capture_exception(e)
38
            asyncio.create_task(callback(task_id=payload["task_id"], response=e))
39