-
Notifications
You must be signed in to change notification settings - Fork 34
/
web.py
144 lines (113 loc) · 3.71 KB
/
web.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
"""
This module initializes the FastAPI application, registers routes, sets up middleware, and configures exception handlers.
"""
import fire
import uvicorn
import logging
import sentry_sdk
from fastapi import FastAPI, Request, status, Depends
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.exceptions import RequestValidationError
from agents_api.common.exceptions import BaseCommonException
from agents_api.exceptions import PromptTooBigError
from pycozo.client import QueryException
from temporalio.service import RPCError
from agents_api.dependencies.auth import get_api_key
from agents_api.env import sentry_dsn
from agents_api.routers import (
agents,
sessions,
users,
jobs,
)
if not sentry_dsn:
print("Sentry DSN not found. Sentry will not be enabled.")
else:
sentry_sdk.init(
dsn=sentry_dsn,
enable_tracing=True,
)
logger = logging.getLogger(__name__)
def make_exception_handler(status: int):
"""
Creates a custom exception handler for the application.
Parameters:
- status (int): The HTTP status code to return for this exception.
Returns:
A callable exception handler that logs the exception and returns a JSON response with the specified status code.
"""
async def _handler(request: Request, exc):
exc_str = f"{exc}".replace("\n", " ").replace(" ", " ")
logger.exception(exc)
content = {"status_code": status, "message": exc_str, "data": None}
return JSONResponse(content=content, status_code=status)
return _handler
def register_exceptions(app: FastAPI):
"""
Registers custom exception handlers for the FastAPI application.
Parameters:
- app (FastAPI): The FastAPI application instance to register the exception handlers for.
"""
app.add_exception_handler(
RequestValidationError,
make_exception_handler(status.HTTP_422_UNPROCESSABLE_ENTITY),
)
app.add_exception_handler(
QueryException,
make_exception_handler(status.HTTP_500_INTERNAL_SERVER_ERROR),
)
app = FastAPI(dependencies=[Depends(get_api_key)])
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
max_age=3600,
)
register_exceptions(app)
app.include_router(agents.router)
app.include_router(sessions.router)
app.include_router(users.router)
app.include_router(jobs.router)
@app.exception_handler(RPCError)
async def validation_error_handler(request: Request, exc: RPCError):
return JSONResponse(
status_code=status.HTTP_400_BAD_REQUEST,
content={
"error": {"message": "job not found or invalid", "code": exc.status.name}
},
)
@app.exception_handler(BaseCommonException)
async def session_not_found_error_handler(request: Request, exc: BaseCommonException):
return JSONResponse(
status_code=exc.http_code,
content={"error": {"message": str(exc)}},
)
@app.exception_handler(PromptTooBigError)
async def prompt_too_big_error(request: Request, exc: PromptTooBigError):
return JSONResponse(
status_code=status.HTTP_400_BAD_REQUEST,
content={"error": {"message": str(exc)}},
)
def main(
host="127.0.0.1",
port=8000,
backlog=4096,
timeout_keep_alive=30,
workers=None,
log_level="info",
):
uvicorn.run(
app,
host=host,
port=port,
log_level=log_level,
timeout_keep_alive=timeout_keep_alive,
backlog=backlog,
workers=workers,
)
# Check if the script is being run directly and, if so, start the Uvicorn server with the specified configuration.
if __name__ == "__main__":
fire.Fire(main)