Port to the Graph SDK for authentication scripts (#1510)

* Configure Azure Developer Pipeline

* Configure Azure Developer Pipeline

* Update pricing calculator link

* Port to Graph SDK

* Port to Graph SDK

* Add msgraph to mypy overrides

* Typing fixes

* Use pyproject.toml for everything

* Fix jose jwt exception import

* change import to match stubs

Author: pamelafox

.github/workflows/python-test.yaml

@@ -51,9 +51,9 @@ jobs:
         - name: Check types with mypy
           run: |
             cd scripts/
-            python3 -m mypy .
-            cd ../app/
-            python3 -m mypy .
+            python3 -m mypy . --config-file=../pyproject.toml
+            cd ../app/backend/
+            python3 -m mypy . --config-file=../../pyproject.toml
         - name: Check formatting with black
           run: black . --check --verbose
         - name: Run Python tests

app/backend/app.py

@@ -356,7 +356,7 @@ async def setup_clients():
             vault_url=f"https://{AZURE_KEY_VAULT_NAME}.vault.azure.net", credential=azure_credential
         ) as key_vault_client:
             search_key = (
-                AZURE_SEARCH_SECRET_NAME and (await key_vault_client.get_secret(AZURE_SEARCH_SECRET_NAME)).value
+                AZURE_SEARCH_SECRET_NAME and (await key_vault_client.get_secret(AZURE_SEARCH_SECRET_NAME)).value  # type: ignore[attr-defined]
             )
 
     # Set up clients for AI Search and Storage
@@ -394,6 +394,10 @@ async def setup_clients():
 
     if USE_USER_UPLOAD:
         current_app.logger.info("USE_USER_UPLOAD is true, setting up user upload feature")
+        if not AZURE_USERSTORAGE_ACCOUNT or not AZURE_USERSTORAGE_CONTAINER:
+            raise ValueError(
+                "AZURE_USERSTORAGE_ACCOUNT and AZURE_USERSTORAGE_CONTAINER must be set when USE_USER_UPLOAD is true"
+            )
         user_blob_container_client = FileSystemClient(
             f"https://{AZURE_USERSTORAGE_ACCOUNT}.dfs.core.windows.net",
             AZURE_USERSTORAGE_CONTAINER,
@@ -504,7 +508,8 @@ async def setup_clients():
 
     if USE_GPT4V:
         current_app.logger.info("USE_GPT4V is true, setting up GPT4V approach")
-
+        if not AZURE_OPENAI_GPT4V_MODEL:
+            raise ValueError("AZURE_OPENAI_GPT4V_MODEL must be set when USE_GPT4V is true")
         token_provider = get_bearer_token_provider(azure_credential, "https://cognitiveservices.azure.com/.default")
 
         current_app.config[CONFIG_ASK_VISION_APPROACH] = RetrieveThenReadVisionApproach(
@@ -565,7 +570,7 @@ def create_app():
         # This tracks OpenAI SDK requests:
         OpenAIInstrumentor().instrument()
         # This middleware tracks app route requests:
-        app.asgi_app = OpenTelemetryMiddleware(app.asgi_app)  # type: ignore[method-assign]
+        app.asgi_app = OpenTelemetryMiddleware(app.asgi_app)  # type: ignore[assignment]
 
     # Level should be one of https://docs.python.org/3/library/logging.html#logging-levels
     default_level = "INFO"  # In development, log more verbosely

app/backend/core/authentication.py

@@ -8,6 +8,7 @@
 from azure.search.documents.aio import SearchClient
 from azure.search.documents.indexes.models import SearchIndex
 from jose import jwt
+from jose.exceptions import ExpiredSignatureError, JWTClaimsError
 from msal import ConfidentialClientApplication
 from msal.token_cache import TokenCache
 from tenacity import (
@@ -323,9 +324,9 @@ async def validate_access_token(self, token: str):
 
         try:
             jwt.decode(token, rsa_key, algorithms=["RS256"], audience=audience, issuer=issuer)
-        except jwt.ExpiredSignatureError as jwt_expired_exc:
+        except ExpiredSignatureError as jwt_expired_exc:
             raise AuthError({"code": "token_expired", "description": "token is expired"}, 401) from jwt_expired_exc
-        except jwt.JWTClaimsError as jwt_claims_exc:
+        except JWTClaimsError as jwt_claims_exc:
             raise AuthError(
                 {"code": "invalid_claims", "description": "incorrect claims," "please check the audience and issuer"},
                 401,

app/backend/prepdocs.py

@@ -120,8 +120,8 @@ def setup_embeddings_service(
     azure_credential: AsyncTokenCredential,
     openai_host: str,
     openai_model_name: str,
-    openai_service: str,
-    openai_deployment: str,
+    openai_service: Union[str, None],
+    openai_deployment: Union[str, None],
     openai_dimensions: int,
     openai_key: Union[str, None],
     openai_org: Union[str, None],

app/backend/prepdocslib/embeddings.py

@@ -159,8 +159,8 @@ class AzureOpenAIEmbeddingService(OpenAIEmbeddings):
 
     def __init__(
         self,
-        open_ai_service: str,
-        open_ai_deployment: str,
+        open_ai_service: Union[str, None],
+        open_ai_deployment: Union[str, None],
         open_ai_model_name: str,
         open_ai_dimensions: int,
         credential: Union[AsyncTokenCredential, AzureKeyCredential],

app/backend/requirements.in

@@ -20,9 +20,11 @@ msal
 azure-keyvault-secrets
 cryptography
 python-jose[cryptography]
+types-python-jose
 Pillow
 types-Pillow
 pypdf
 PyMuPDF
 beautifulsoup4
 types-beautifulsoup4
+msgraph-sdk==1.1.0

app/backend/requirements.txt

@@ -7,7 +7,9 @@
 aiofiles==23.2.1
     # via quart
 aiohttp==3.9.3
-    # via -r requirements.in
+    # via
+    #   -r requirements.in
+    #   microsoft-kiota-authentication-azure
 aiosignal==1.3.1
     # via aiohttp
 annotated-types==0.6.0
@@ -35,11 +37,14 @@ azure-core==1.30.1
     #   azure-search-documents
     #   azure-storage-blob
     #   azure-storage-file-datalake
+    #   microsoft-kiota-authentication-azure
     #   msrest
 azure-core-tracing-opentelemetry==1.0.0b11
     # via azure-monitor-opentelemetry
 azure-identity==1.15.0
-    # via -r requirements.in
+    # via
+    #   -r requirements.in
+    #   msgraph-sdk
 azure-keyvault-secrets==4.8.0
     # via -r requirements.in
 azure-monitor-opentelemetry==1.3.0
@@ -104,13 +109,18 @@ h11==0.14.0
     #   uvicorn
     #   wsproto
 h2==4.1.0
-    # via hypercorn
+    # via
+    #   httpx
+    #   hypercorn
 hpack==4.0.0
     # via h2
 httpcore==1.0.4
     # via httpx
-httpx==0.27.0
-    # via openai
+httpx[http2]==0.27.0
+    # via
+    #   microsoft-kiota-http
+    #   msgraph-core
+    #   openai
 hypercorn==0.16.0
     # via quart
 hyperframe==6.0.1
@@ -144,13 +154,37 @@ markupsafe==2.1.5
     #   jinja2
     #   quart
     #   werkzeug
+microsoft-kiota-abstractions==1.3.2
+    # via
+    #   microsoft-kiota-authentication-azure
+    #   microsoft-kiota-http
+    #   microsoft-kiota-serialization-json
+    #   microsoft-kiota-serialization-text
+    #   msgraph-core
+    #   msgraph-sdk
+microsoft-kiota-authentication-azure==1.0.0
+    # via
+    #   msgraph-core
+    #   msgraph-sdk
+microsoft-kiota-http==1.3.1
+    # via
+    #   msgraph-core
+    #   msgraph-sdk
+microsoft-kiota-serialization-json==1.1.0
+    # via msgraph-sdk
+microsoft-kiota-serialization-text==1.0.0
+    # via msgraph-sdk
 msal==1.27.0
     # via
     #   -r requirements.in
     #   azure-identity
     #   msal-extensions
 msal-extensions==1.1.0
     # via azure-identity
+msgraph-core==1.0.0
+    # via msgraph-sdk
+msgraph-sdk==1.1.0
+    # via -r requirements.in
 msrest==0.7.1
     # via azure-monitor-opentelemetry-exporter
 multidict==6.0.5
@@ -170,6 +204,9 @@ opentelemetry-api==1.23.0
     # via
     #   azure-core-tracing-opentelemetry
     #   azure-monitor-opentelemetry-exporter
+    #   microsoft-kiota-abstractions
+    #   microsoft-kiota-authentication-azure
+    #   microsoft-kiota-http
     #   opentelemetry-instrumentation
     #   opentelemetry-instrumentation-aiohttp-client
     #   opentelemetry-instrumentation-asgi
@@ -237,6 +274,9 @@ opentelemetry-resource-detector-azure==0.1.3
 opentelemetry-sdk==1.23.0
     # via
     #   azure-monitor-opentelemetry-exporter
+    #   microsoft-kiota-abstractions
+    #   microsoft-kiota-authentication-azure
+    #   microsoft-kiota-http
     #   opentelemetry-resource-detector-azure
 opentelemetry-semantic-conventions==0.44b0
     # via
@@ -274,6 +314,8 @@ pandas==2.2.1
     # via openai
 pandas-stubs==2.2.0.240218
     # via openai
+pendulum==3.0.0
+    # via microsoft-kiota-serialization-json
 pillow==10.3.0
     # via -r requirements.in
 portalocker==2.8.2
@@ -291,17 +333,19 @@ pydantic==2.6.3
 pydantic-core==2.16.3
     # via pydantic
 pyjwt[crypto]==2.8.0
-    # via
-    #   msal
-    #   pyjwt
+    # via msal
 pymupdf==1.23.26
     # via -r requirements.in
 pymupdfb==1.23.22
     # via pymupdf
 pypdf==4.1.0
     # via -r requirements.in
 python-dateutil==2.9.0.post0
-    # via pandas
+    # via
+    #   microsoft-kiota-serialization-text
+    #   pandas
+    #   pendulum
+    #   time-machine
 python-jose[cryptography]==3.3.0
     # via -r requirements.in
 pytz==2024.1
@@ -338,10 +382,14 @@ sniffio==1.3.1
     #   openai
 soupsieve==2.5
     # via beautifulsoup4
+std-uritemplate==0.0.55
+    # via microsoft-kiota-abstractions
 tenacity==8.2.3
     # via -r requirements.in
 tiktoken==0.6.0
     # via -r requirements.in
+time-machine==2.14.1
+    # via pendulum
 tqdm==4.66.2
     # via openai
 types-beautifulsoup4==4.12.0.20240229
@@ -350,6 +398,10 @@ types-html5lib==1.1.11.20240228
     # via types-beautifulsoup4
 types-pillow==10.2.0.20240213
     # via -r requirements.in
+types-pyasn1==0.6.0.20240402
+    # via types-python-jose
+types-python-jose==3.3.4.20240106
+    # via -r requirements.in
 types-pytz==2024.1.0.20240203
     # via pandas-stubs
 typing-extensions==4.10.0
@@ -364,7 +416,9 @@ typing-extensions==4.10.0
     #   pydantic
     #   pydantic-core
 tzdata==2024.1
-    # via pandas
+    # via
+    #   pandas
+    #   pendulum
 urllib3==2.2.1
     # via requests
 uvicorn==0.27.1

app/mypy.ini

@@ -26,6 +26,9 @@ python_version = 3.9
 
 [[tool.mypy.overrides]]
 module = [
-    "msal.*"
+    "msal.*",
+    "msgraph.*",
+    "kiota_abstractions.*",
+    "kiota.*"
 ]
 ignore_missing_imports = true

pyproject.toml

@@ -1,37 +1,16 @@
 import os
-from typing import Dict, Optional
+from typing import Optional
 
-import aiohttp
-from azure.core.credentials_async import AsyncTokenCredential
+from kiota_abstractions.api_error import APIError
+from msgraph import GraphServiceClient
 
-TIMEOUT = 60
 
-
-async def get_auth_headers(credential: AsyncTokenCredential):
-    token_result = await credential.get_token("https://graph.microsoft.com/.default")
-    return {"Authorization": f"Bearer {token_result.token}"}
-
-
-async def get_application(auth_headers: Dict[str, str], app_id: str) -> Optional[str]:
-    async with aiohttp.ClientSession(headers=auth_headers, timeout=aiohttp.ClientTimeout(total=TIMEOUT)) as session:
-        async with session.get(f"https://graph.microsoft.com/v1.0/applications(appId='{app_id}')") as response:
-            if response.status == 200:
-                response_json = await response.json()
-                return response_json["id"]
-
-    return None
-
-
-async def update_application(auth_headers: Dict[str, str], object_id: str, app_payload: object):
-    async with aiohttp.ClientSession(headers=auth_headers, timeout=aiohttp.ClientTimeout(total=TIMEOUT)) as session:
-        async with session.patch(
-            f"https://graph.microsoft.com/v1.0/applications/{object_id}", json=app_payload
-        ) as response:
-            if not response.ok:
-                response_json = await response.json()
-                raise Exception(response_json)
-
-    return True
+async def get_application(graph_client: GraphServiceClient, client_id: str) -> Optional[str]:
+    try:
+        app = await graph_client.applications_with_app_id(client_id).get()
+        return app.id
+    except APIError:
+        return None
 
 
 def test_authentication_enabled():