Add lock to swagger UI (#3)

* Add lock to swagger UI

* Fic CI

* Fic CI2

* fix lint

* fix CI

* fix SSL certs

Author: BoBer78

CHANGELOG.md

@@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+- Bearer token insertion in FastAPI UI.
+
 ### Fixed
 - Limit query size.
 - Add OBI copyright.

src/scholarag/app/config.py

@@ -4,33 +4,31 @@
 from typing import Literal
 
 from dotenv import dotenv_values
-from fastapi.openapi.models import OAuthFlowPassword, OAuthFlows
-from pydantic import BaseModel, ConfigDict, SecretStr, model_validator
+from pydantic import BaseModel, ConfigDict, SecretStr
 from pydantic_settings import BaseSettings, SettingsConfigDict
-from typing_extensions import Self
 
 from scholarag.generative_question_answering import MESSAGES
 
 
 class SettingsKeycloak(BaseModel):
     """Class retrieving keycloak info for authorization."""
 
-    issuer: str | None = None
+    issuer: str = "https://openbluebrain.com/auth/realms/SBO"
     validate_token: bool = False
+    # Useful only for service account (dev)
+    client_id: str | None = None
+    username: str | None = None
+    password: SecretStr | None = None
 
     model_config = ConfigDict(frozen=True)
 
-    @model_validator(mode="after")
-    def check_issuer(self) -> Self:
-        """Check if there is an issuer provided for authentication."""
-        if self.validate_token and (self.issuer is None):
-            raise ValueError("No issuer provided.")
-        return self
-
     @property
-    def token_endpoint(self) -> str:
+    def token_endpoint(self) -> str | None:
         """Define the token endpoint."""
-        return f"{self.issuer}/protocol/openid-connect/token"
+        if self.validate_token:
+            return f"{self.issuer}/protocol/openid-connect/token"
+        else:
+            return None
 
     @property
     def user_info_endpoint(self) -> str | None:
@@ -41,13 +39,14 @@ def user_info_endpoint(self) -> str | None:
             return None
 
     @property
-    def flows(self) -> OAuthFlows:
-        """Define the flow to override Fastapi's one."""
-        return OAuthFlows(
-            password=OAuthFlowPassword(
-                tokenUrl=self.token_endpoint,
-            ),
-        )
+    def server_url(self) -> str:
+        """Server url."""
+        return self.issuer.split("/auth")[0] + "/auth/"
+
+    @property
+    def realm(self) -> str:
+        """Realm."""
+        return self.issuer.rpartition("/realms/")[-1]
 
 
 class SettingsDB(BaseModel):

src/scholarag/app/dependencies.py

@@ -7,8 +7,8 @@
 from typing import Annotated, AsyncIterator
 
 from elasticsearch_dsl import Q, Search
-from fastapi import Depends, HTTPException, Query
-from fastapi.security import OAuth2PasswordBearer
+from fastapi import Depends, HTTPException, Query, Request
+from fastapi.security import HTTPBearer
 from httpx import AsyncClient, HTTPStatusError
 from openai import AsyncOpenAI
 from pydantic import constr
@@ -27,10 +27,17 @@
 
 journal_constraints = constr(pattern=r"^\d{4}-\d{3}[0-9X]$")
 
-auth = OAuth2PasswordBearer(
-    tokenUrl="/token",  # Will be overriden
-    auto_error=False,
-)
+
+class HTTPBearerDirect(HTTPBearer):
+    """HTTPBearer class that returns directly the token in the call."""
+
+    async def __call__(self, request: Request) -> str | None:  # type: ignore
+        """Intercept the bearer token in the headers."""
+        auth_credentials = await super().__call__(request)
+        return auth_credentials.credentials if auth_credentials else None
+
+
+auth = HTTPBearerDirect(auto_error=False)
 
 
 @cache
@@ -55,24 +62,30 @@ async def get_httpx_client(
 
 
 async def get_user_id(
-    token: Annotated[str | None, Depends(auth)],
+    request: Request,
+    token: Annotated[str, Depends(auth)],
     settings: Annotated[Settings, Depends(get_settings)],
     httpx_client: Annotated[AsyncClient, Depends(get_httpx_client)],
 ) -> str:
     """Validate JWT token and returns user ID."""
-    if settings.keycloak.validate_token and settings.keycloak.user_info_endpoint:
-        try:
-            response = await httpx_client.get(
-                settings.keycloak.user_info_endpoint,
-                headers={"Authorization": f"Bearer {token}"},
-            )
-            response.raise_for_status()
-            user_info = response.json()
-            return user_info["sub"]
-        except HTTPStatusError:
-            raise HTTPException(
-                status_code=HTTP_401_UNAUTHORIZED, detail="Invalid token."
-            )
+    if hasattr(request.state, "sub"):
+        return request.state.sub
+    if settings.keycloak.validate_token:
+        if settings.keycloak.user_info_endpoint:
+            try:
+                response = await httpx_client.get(
+                    settings.keycloak.user_info_endpoint,
+                    headers={"Authorization": f"Bearer {token}"},
+                )
+                response.raise_for_status()
+                user_info = response.json()
+                return user_info["sub"]
+            except HTTPStatusError:
+                raise HTTPException(
+                    status_code=HTTP_401_UNAUTHORIZED, detail="Invalid token."
+                )
+        else:
+            raise HTTPException(status_code=404, detail="user info url not provided.")
     else:
         return "dev"
 

src/scholarag/app/middleware.py

@@ -250,7 +250,13 @@ async def get_and_set_cache(
     httpx_client = await anext(get_httpx_client(settings))
     # If raises HTTPException return error as json.
     try:
-        await get_user_id(token=token, settings=settings, httpx_client=httpx_client)
+        sub = await get_user_id(
+            request=request,
+            token=token,  # type: ignore
+            settings=settings,
+            httpx_client=httpx_client,
+        )
+        request.state.sub = sub
     except HTTPException as e:
         return JSONResponse(status_code=e.status_code, content=e.detail)
 

src/scholarag/app/routers/qa.py

@@ -22,6 +22,7 @@
     get_reranker,
     get_rts,
     get_settings,
+    get_user_id,
 )
 from scholarag.app.schemas import (
     GenerativeQARequest,
@@ -36,7 +37,9 @@
 from scholarag.retrieve_metadata import MetaDataRetriever
 from scholarag.services import CohereRerankingService, RetrievalService
 
-router = APIRouter(prefix="/qa", tags=["Question answering"])
+router = APIRouter(
+    prefix="/qa", tags=["Question answering"], dependencies=[Depends(get_user_id)]
+)
 
 logger = logging.getLogger(__name__)
 

src/scholarag/app/routers/retrieval.py

@@ -17,6 +17,7 @@
     get_reranker,
     get_rts,
     get_settings,
+    get_user_id,
 )
 from scholarag.app.schemas import (
     ArticleCountResponse,
@@ -28,7 +29,9 @@
 from scholarag.retrieve_metadata import MetaDataRetriever
 from scholarag.services import CohereRerankingService, RetrievalService
 
-router = APIRouter(prefix="/retrieval", tags=["Retrieval"])
+router = APIRouter(
+    prefix="/retrieval", tags=["Retrieval"], dependencies=[Depends(get_user_id)]
+)
 
 logger = logging.getLogger(__name__)
 

src/scholarag/app/routers/suggestions.py

@@ -8,7 +8,12 @@
 from fastapi import APIRouter, Depends, HTTPException
 
 from scholarag.app.config import Settings
-from scholarag.app.dependencies import ErrorCode, get_ds_client, get_settings
+from scholarag.app.dependencies import (
+    ErrorCode,
+    get_ds_client,
+    get_settings,
+    get_user_id,
+)
 from scholarag.app.schemas import (
     ArticleTypeSuggestionResponse,
     AuthorSuggestionRequest,
@@ -19,7 +24,9 @@
 from scholarag.document_stores import AsyncBaseSearch
 from scholarag.utils import format_issn
 
-router = APIRouter(prefix="/suggestions", tags=["Suggestions"])
+router = APIRouter(
+    prefix="/suggestions", tags=["Suggestions"], dependencies=[Depends(get_user_id)]
+)
 
 logger = logging.getLogger(__name__)
 

src/scholarag/scripts/pmc_parse_and_upload.py

@@ -181,7 +181,7 @@ async def run(
             s3_iterator = s3_paginator.paginate(Bucket="pmc-oa-opendata", Prefix=prefix)
             logger.info("Filtering interesting articles.")
             filtered_iterator = s3_iterator.search(
-                f"""Contents[?to_string(LastModified)>='\"{start_date.strftime('%Y-%m-%d %H:%M:%S%')}+00:00\"'
+                f"""Contents[?to_string(LastModified)>='\"{start_date.strftime("%Y-%m-%d %H:%M:%S%")}+00:00\"'
                     && contains(Key, '.xml')]"""
             )
             finished = False

src/scholarag/scripts/pu_consumer.py

@@ -258,7 +258,7 @@ async def run(
                 for file, message in zip(files, batch):
                     url = (
                         parser_url
-                        + f'/{message["MessageAttributes"]["Parser_Endpoint"]["StringValue"]}'
+                        + f"/{message['MessageAttributes']['Parser_Endpoint']['StringValue']}"
                     )
                     result = await parser_service.arun(
                         files=[file],
@@ -309,7 +309,7 @@ async def run(
                     else:
                         logger.info(
                             f"[WORKER {worker_n}] Successfully deleted"
-                            f" {len(to_delete_from_q[i:min(i+10, len(to_delete_from_q))])} entries"
+                            f" {len(to_delete_from_q[i : min(i + 10, len(to_delete_from_q))])} entries"
                             " from the queue."
                         )
     except BaseException as e:

src/scholarag/scripts/pu_producer.py

@@ -142,7 +142,7 @@ async def run(
                 f" date: {start_date}."
             )
             filtered_iterator = s3_iterator.search(
-                f"""Contents[?to_string(LastModified)>='\"{start_date.strftime('%Y-%m-%d %H:%M:%S%')}+00:00\"'
+                f"""Contents[?to_string(LastModified)>='\"{start_date.strftime("%Y-%m-%d %H:%M:%S%")}+00:00\"'
                     && contains(Key, '.{file_extension if file_extension else ""}')]"""
             )
             finished = False

tests/app/test_dependencies.py

@@ -1,3 +1,5 @@
+from unittest.mock import Mock
+
 import pytest
 from fastapi.exceptions import HTTPException
 from httpx import AsyncClient
@@ -66,11 +68,23 @@ async def test_get_user_id(httpx_mock, monkeypatch):
         json=fake_response,
     )
 
+    request = Mock()
     settings = Settings()
     client = AsyncClient()
     token = "eyJgreattoken"
-    user = await get_user_id(token=token, settings=settings, httpx_client=client)
 
+    # Checks for the user sub in the request state.
+    request.state.sub = "test_sub"
+    user = await get_user_id(
+        request=request, token=token, settings=settings, httpx_client=client
+    )
+    assert user == "test_sub"
+
+    # Check for the user sub from keycloack.
+    delattr(request.state, "sub")
+    user = await get_user_id(
+        request=request, token=token, settings=settings, httpx_client=client
+    )
     assert user == "12345"
 
 
@@ -92,8 +106,13 @@ async def test_get_user_id_error(httpx_mock, monkeypatch):
     client = AsyncClient()
     token = "eyJgreattoken"
 
+    request = Mock()
+    delattr(request.state, "sub")
+
     with pytest.raises(HTTPException) as err:
-        await get_user_id(token=token, settings=settings, httpx_client=client)
+        await get_user_id(
+            request=request, token=token, settings=settings, httpx_client=client
+        )
 
     assert err.value.status_code == 401
     assert err.value.detail == "Invalid token."

tests/app/test_qa.py

@@ -355,7 +355,7 @@ async def streamed_response_no_answer(**kwargs):
         "don",
         "'t ",
         "know.",
-        ", " '"paragraphs": ',
+        ', "paragraphs": ',
         "[]}",
     ]
     parsed = {}