diff --git a/cognee/api/client.py b/cognee/api/client.py index 4bdc89f10..c504690e4 100644 --- a/cognee/api/client.py +++ b/cognee/api/client.py @@ -12,6 +12,7 @@ from fastapi.encoders import jsonable_encoder from fastapi.responses import JSONResponse, Response from fastapi.middleware.cors import CORSMiddleware from fastapi.exceptions import RequestValidationError +from fastapi.openapi.utils import get_openapi from cognee.exceptions import CogneeApiError from cognee.shared.logging_utils import get_logger, setup_logging @@ -75,6 +76,36 @@ app.add_middleware( ) +def custom_openapi(): + if app.openapi_schema: + return app.openapi_schema + + openapi_schema = get_openapi( + title="Cognee API", + version="1.0.0", + description="Cognee API with Bearer token and Cookie auth", + routes=app.routes, + ) + + openapi_schema["components"]["securitySchemes"] = { + "BearerAuth": {"type": "http", "scheme": "bearer"}, + "CookieAuth": { + "type": "apiKey", + "in": "cookie", + "name": os.getenv("AUTH_TOKEN_COOKIE_NAME", "auth_token"), + }, + } + + openapi_schema["security"] = [{"BearerAuth": []}, {"CookieAuth": []}] + + app.openapi_schema = openapi_schema + + return app.openapi_schema + + +app.openapi = custom_openapi + + @app.exception_handler(RequestValidationError) async def request_validation_exception_handler(request: Request, exc: RequestValidationError): if request.url.path == "/api/v1/auth/login": diff --git a/cognee/modules/users/get_user_manager.py b/cognee/modules/users/get_user_manager.py index 0a9580d50..19734ad33 100644 --- a/cognee/modules/users/get_user_manager.py +++ b/cognee/modules/users/get_user_manager.py @@ -56,6 +56,7 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]): response.body = json.dumps( {"access_token": access_token, "token_type": "bearer"} ).encode(encoding="utf-8") + response.headers.append("Content-Type", "application/json") async def on_after_register(self, user: User, request: Optional[Request] = None): print(f"User {user.id} has registered.")