From 5e6bbef0b300134d67b132007462ef4d95a38bda Mon Sep 17 00:00:00 2001 From: Boris Arzentar Date: Mon, 30 Jun 2025 12:27:53 +0200 Subject: [PATCH 1/2] fix: add auth methods to swagger authtorize --- cognee/api/client.py | 35 +++++++++++++++++++ .../users/methods/get_authenticated_user.py | 10 +++++- 2 files changed, 44 insertions(+), 1 deletion(-) diff --git a/cognee/api/client.py b/cognee/api/client.py index 4bdc89f10..1ffa85ba2 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,40 @@ 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/methods/get_authenticated_user.py b/cognee/modules/users/methods/get_authenticated_user.py index b60ddfe28..486f92c1b 100644 --- a/cognee/modules/users/methods/get_authenticated_user.py +++ b/cognee/modules/users/methods/get_authenticated_user.py @@ -1,9 +1,17 @@ +from cognee.modules.users.authentication.get_api_auth_backend import get_api_auth_backend +from cognee.modules.users.authentication.get_client_auth_backend import get_client_auth_backend from ..get_fastapi_users import get_fastapi_users fastapi_users = get_fastapi_users() -get_authenticated_user = fastapi_users.current_user(active=True) +def get_enabled_backends(): + api_auth_backend = get_api_auth_backend() + client_auth_backend = get_client_auth_backend() + + return [api_auth_backend, client_auth_backend] + +get_authenticated_user = fastapi_users.current_user(active=True, get_enabled_backends=get_enabled_backends) # from types import SimpleNamespace From 64edb38c4368516acfd7669e1268977e6758b419 Mon Sep 17 00:00:00 2001 From: Boris Arzentar Date: Mon, 30 Jun 2025 15:09:13 +0200 Subject: [PATCH 2/2] fix: add custom openauth schema --- cognee/api/client.py | 16 ++++++---------- cognee/modules/users/get_user_manager.py | 1 + .../users/methods/get_authenticated_user.py | 10 +--------- 3 files changed, 8 insertions(+), 19 deletions(-) diff --git a/cognee/api/client.py b/cognee/api/client.py index 1ffa85ba2..c504690e4 100644 --- a/cognee/api/client.py +++ b/cognee/api/client.py @@ -88,28 +88,24 @@ def custom_openapi(): ) openapi_schema["components"]["securitySchemes"] = { - "BearerAuth": { - "type": "http", - "scheme": "bearer" - }, + "BearerAuth": {"type": "http", "scheme": "bearer"}, "CookieAuth": { "type": "apiKey", "in": "cookie", - "name": os.getenv("AUTH_TOKEN_COOKIE_NAME", "auth_token") - } + "name": os.getenv("AUTH_TOKEN_COOKIE_NAME", "auth_token"), + }, } - openapi_schema["security"] = [ - {"BearerAuth": []}, - {"CookieAuth": []} - ] + 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.") diff --git a/cognee/modules/users/methods/get_authenticated_user.py b/cognee/modules/users/methods/get_authenticated_user.py index 486f92c1b..b60ddfe28 100644 --- a/cognee/modules/users/methods/get_authenticated_user.py +++ b/cognee/modules/users/methods/get_authenticated_user.py @@ -1,17 +1,9 @@ -from cognee.modules.users.authentication.get_api_auth_backend import get_api_auth_backend -from cognee.modules.users.authentication.get_client_auth_backend import get_client_auth_backend from ..get_fastapi_users import get_fastapi_users fastapi_users = get_fastapi_users() -def get_enabled_backends(): - api_auth_backend = get_api_auth_backend() - client_auth_backend = get_client_auth_backend() - - return [api_auth_backend, client_auth_backend] - -get_authenticated_user = fastapi_users.current_user(active=True, get_enabled_backends=get_enabled_backends) +get_authenticated_user = fastapi_users.current_user(active=True) # from types import SimpleNamespace