FilesExpand file tree

main

/

main.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

66 lines (53 loc) · 2.02 KB

main

/

main.py

Copy path

Top

File metadata and controls

• Code
• Blame

66 lines (53 loc) · 2.02 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

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

# main.py

import uvicorn

from fastapi import FastAPI, Request

from fastapi.exceptions import RequestValidationError

from fastapi.middleware.cors import CORSMiddleware

from contextlib import asynccontextmanager

from starlette.responses import JSONResponse

from app.config import VectorDBType, debug_mode, RAG_HOST, RAG_PORT, CHUNK_SIZE, CHUNK_OVERLAP, PDF_EXTRACT_IMAGES, VECTOR_DB_TYPE, \

LogMiddleware, logger

from app.middleware import security_middleware

from app.routes import document_routes, pgvector_routes

from app.services.database import PSQLDatabase, ensure_custom_id_index_on_embedding

@asynccontextmanager

async def lifespan(app: FastAPI):

# Startup logic goes here

if VECTOR_DB_TYPE == VectorDBType.PGVECTOR:

await PSQLDatabase.get_pool() # Initialize the pool

await ensure_custom_id_index_on_embedding()

yield

app = FastAPI(lifespan=lifespan, debug=debug_mode)

app.add_middleware(

CORSMiddleware,

allow_origins=["*"],

allow_credentials=True,

allow_methods=["*"],

allow_headers=["*"],

)

app.add_middleware(LogMiddleware)

app.middleware("http")(security_middleware)

# Set state variables for use in routes

app.state.CHUNK_SIZE = CHUNK_SIZE

app.state.CHUNK_OVERLAP = CHUNK_OVERLAP

app.state.PDF_EXTRACT_IMAGES = PDF_EXTRACT_IMAGES

# Include routers

app.include_router(document_routes.router)

if debug_mode:

app.include_router(router=pgvector_routes.router)

@app.exception_handler(RequestValidationError)

async def validation_exception_handler(request: Request, exc: RequestValidationError):

body = await request.body()

logger.debug(f"Validation error occurred")

logger.debug(f"Raw request body: {body.decode()}")

logger.debug(f"Validation errors: {exc.errors()}")

return JSONResponse(

status_code=422,

content={

"detail": exc.errors(),

"body": body.decode(),

"message": "Request validation failed",

},

)

if __name__ == "__main__":

uvicorn.run(app, host=RAG_HOST, port=RAG_PORT, log_config=None)