dream

1
import logging
2
import time
3
from functools import wraps
4

5
from fastapi import FastAPI
6
import uvicorn
7

8
import sentry_sdk
9
from sentry_sdk.integrations.asgi import SentryAsgiMiddleware
10
from sentry_sdk.integrations.logging import ignore_logger
11

12
from comet_commonsense.interface import COMeTFactory
13
from comet_commonsense.config import settings
14
import test_server
15

16
ignore_logger("root")
17

18
sentry_sdk.init(dsn=settings.SENTRY_DSN)
19

20
app = FastAPI()
21

22
logging.basicConfig(format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO)
23
logger = logging.getLogger(__name__)
24

25
comet_engine = COMeTFactory(settings.GRAPH)(settings.PRETRAINED_MODEL, settings.DECODING_ALGO)
26

27

28
def timing(f):
29
    @wraps(f)
30
    def wrap(*args, **kw):
31
        ts = time.time()
32
        result = f(*args, **kw)
33
        total_time = time.time() - ts
34
        logger.info(f"{settings.SERVICE_NAME} exec time = {total_time:.3f}s")
35
        return result
36

37
    return wrap
38

39

40
@timing
41
def handler(data):
42
    try:
43
        return comet_engine.process_request(data)
44
    except Exception as exc:
45
        sentry_sdk.capture_exception(exc)
46
        logger.exception(exc)
47
        raise exc
48

49

50
@timing
51
def annotator_handler(data):
52
    try:
53
        return comet_engine.annotator(data)
54
    except Exception as exc:
55
        sentry_sdk.capture_exception(exc)
56
        logger.exception(exc)
57
        raise exc
58

59

60
@app.post("/comet", response_model=comet_engine.response_model)
61
async def comet_base_handler(input_event: comet_engine.input_event_model):
62
    result = handler(input_event.dict())
63
    return result
64

65

66
@app.post("/comet_annotator", response_model=comet_engine.annotator_response_model)
67
async def comet_annotator_handler(input_event: comet_engine.annotator_input_model):
68
    result = annotator_handler(input_event.dict())
69
    logger.info(f"comet_annotator result: {result}")
70
    return result
71

72

73
app = SentryAsgiMiddleware(app)
74

75
try:
76
    test_server.run_test(handler)
77
    logger.info("test query processed")
78
except Exception as exc:
79
    sentry_sdk.capture_exception(exc)
80
    logger.exception(exc)
81
    raise exc
82

83
logger.info(f"{settings.SERVICE_NAME} is loaded and ready")
84

85
if __name__ == "__main__":
86
    uvicorn.run(app, host="0.0.0.0", port=settings.SERVICE_PORT, log_level="info")
87