86 lines
2.6 KiB
Python
86 lines
2.6 KiB
Python
from fastapi import FastAPI, Request
|
|
from scalar_fastapi import get_scalar_api_reference
|
|
import sentry_sdk
|
|
from sentry_sdk.integrations.loguru import LoguruIntegration, LoggingLevels
|
|
from sentry_sdk.integrations.fastapi import FastApiIntegration
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
|
|
from .utils.KVCache import KVCache
|
|
from .router import ffmpeg, cache, comfyui, google, task
|
|
from .config import WorkerConfig
|
|
|
|
config = WorkerConfig()
|
|
|
|
web_app = FastAPI(title="Modal worker API",
|
|
version="0.1.6",
|
|
summary="Modal Worker的API, 包括缓存视频, 发起生产任务等",
|
|
servers=[
|
|
{
|
|
'url': f'https://bowongai-{config.modal_environment}--{config.modal_app_name}-fastapi-webapp.modal.run',
|
|
'description': '当前Modal API接口endpoint'
|
|
}
|
|
])
|
|
|
|
sentry_sdk.init(dsn="https://dab7b7ae652216282c89f029a76bb10a@sentry.bowongai.com/2",
|
|
send_default_pii=True,
|
|
traces_sample_rate=1.0,
|
|
profiles_sample_rate=1.0,
|
|
add_full_stack=True,
|
|
environment=config.modal_environment,
|
|
integrations=[
|
|
LoguruIntegration(level=LoggingLevels.INFO.value, event_level=LoggingLevels.ERROR.value),
|
|
FastApiIntegration()
|
|
]
|
|
)
|
|
modal_kv_cache = KVCache(kv_name=config.modal_kv_name, environment=config.modal_environment)
|
|
|
|
sentry_header_schema = {
|
|
"x-trace-id": {
|
|
"description": "Sentry Transaction ID",
|
|
"schema": {
|
|
"type": "string",
|
|
}
|
|
},
|
|
"x-baggage": {
|
|
"description": "Sentry Transaction baggage",
|
|
"schema": {
|
|
"type": "string",
|
|
}
|
|
}
|
|
}
|
|
|
|
ALIAS_MAP = {
|
|
"/comfyui": "/comfyui/v2",
|
|
}
|
|
|
|
web_app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
|
|
@web_app.middleware("http")
|
|
@web_app.middleware("https")
|
|
async def alias_middleware(request: Request, call_next):
|
|
if request.url.path in ALIAS_MAP.keys():
|
|
request.scope["path"] = ALIAS_MAP[request.url.path]
|
|
return await call_next(request)
|
|
|
|
|
|
@web_app.get("/scalar", include_in_schema=False)
|
|
async def scalar():
|
|
return get_scalar_api_reference(openapi_url='/openapi.json', title="Modal worker web endpoint")
|
|
|
|
|
|
web_app.include_router(ffmpeg.router)
|
|
web_app.include_router(cache.router)
|
|
|
|
# todo: prod 展示去掉Comfy相关API接口
|
|
# web_app.include_router(comfyui.router)
|
|
|
|
web_app.include_router(google.router)
|
|
web_app.include_router(task.router)
|