main.py 2.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384
  1. import time
  2. from fastapi import FastAPI, Request
  3. from fastapi.staticfiles import StaticFiles
  4. from fastapi import HTTPException
  5. from fastapi.middleware.wsgi import WSGIMiddleware
  6. from fastapi.middleware.cors import CORSMiddleware
  7. from starlette.exceptions import HTTPException as StarletteHTTPException
  8. from litellm.proxy.proxy_server import app as litellm_app
  9. from apps.ollama.main import app as ollama_app
  10. from apps.openai.main import app as openai_app
  11. from apps.audio.main import app as audio_app
  12. from apps.images.main import app as images_app
  13. from apps.rag.main import app as rag_app
  14. from apps.web.main import app as webui_app
  15. from config import ENV, FRONTEND_BUILD_DIR
  16. class SPAStaticFiles(StaticFiles):
  17. async def get_response(self, path: str, scope):
  18. try:
  19. return await super().get_response(path, scope)
  20. except (HTTPException, StarletteHTTPException) as ex:
  21. if ex.status_code == 404:
  22. return await super().get_response("index.html", scope)
  23. else:
  24. raise ex
  25. app = FastAPI(docs_url="/docs" if ENV == "dev" else None, redoc_url=None)
  26. origins = ["*"]
  27. app.add_middleware(
  28. CORSMiddleware,
  29. allow_origins=origins,
  30. allow_credentials=True,
  31. allow_methods=["*"],
  32. allow_headers=["*"],
  33. )
  34. @app.middleware("http")
  35. async def check_url(request: Request, call_next):
  36. start_time = int(time.time())
  37. response = await call_next(request)
  38. process_time = int(time.time()) - start_time
  39. response.headers["X-Process-Time"] = str(process_time)
  40. return response
  41. app.mount("/api/v1", webui_app)
  42. app.mount("/litellm/api", litellm_app)
  43. app.mount("/ollama/api", ollama_app)
  44. app.mount("/openai/api", openai_app)
  45. app.mount("/images/api/v1", images_app)
  46. app.mount("/audio/api/v1", audio_app)
  47. app.mount("/rag/api/v1", rag_app)
  48. @app.get("/api/config")
  49. async def get_app_config():
  50. return {
  51. "status": True,
  52. "images": images_app.state.ENABLED,
  53. "default_models": webui_app.state.DEFAULT_MODELS,
  54. "default_prompt_suggestions": webui_app.state.DEFAULT_PROMPT_SUGGESTIONS,
  55. }
  56. app.mount(
  57. "/",
  58. SPAStaticFiles(directory=FRONTEND_BUILD_DIR, html=True),
  59. name="spa-static-files",
  60. )