From 69668f26f2f9810823119f1e1a615ef1dd529200 Mon Sep 17 00:00:00 2001
From: Willy Douhard <willy.douhard@gmail.com>
Date: Fri, 14 Jun 2024 11:46:59 +0200
Subject: [PATCH] fix: add back get_user_info to AzureADOAuthProvider (#1075)

* fix: add back get_user_info to AzureADOAuthProvider

* prepare patch
---
 CHANGELOG.md                        |  6 +++++
 backend/chainlit/oauth_providers.py | 36 ++++++++++++++++++++++++++---
 backend/pyproject.toml              |  2 +-
 3 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7c178a0ae3..95cb4afaa8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,12 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 Nothing unreleased.
 
+## [1.1.301] - 2024-06-14
+
+### Fixed
+
+- Azure AD oauth get_user_info not implemented error
+
 ## [1.1.300] - 2024-06-13
 
 ### Added
diff --git a/backend/chainlit/oauth_providers.py b/backend/chainlit/oauth_providers.py
index 974a599b95..fe019859b1 100644
--- a/backend/chainlit/oauth_providers.py
+++ b/backend/chainlit/oauth_providers.py
@@ -187,6 +187,36 @@ async def get_token(self, code: str, url: str):
                 )
             return token
 
+    async def get_user_info(self, token: str):
+        async with httpx.AsyncClient() as client:
+            response = await client.get(
+                "https://graph.microsoft.com/v1.0/me",
+                headers={"Authorization": f"Bearer {token}"},
+            )
+            response.raise_for_status()
+
+            azure_user = response.json()
+
+            try:
+                photo_response = await client.get(
+                    "https://graph.microsoft.com/v1.0/me/photos/48x48/$value",
+                    headers={"Authorization": f"Bearer {token}"},
+                )
+                photo_data = await photo_response.aread()
+                base64_image = base64.b64encode(photo_data)
+                azure_user["image"] = (
+                    f"data:{photo_response.headers['Content-Type']};base64,{base64_image.decode('utf-8')}"
+                )
+            except Exception as e:
+                # Ignore errors getting the photo
+                pass
+
+            user = User(
+                identifier=azure_user["userPrincipalName"],
+                metadata={"image": azure_user.get("image"), "provider": "azure-ad"},
+            )
+            return (azure_user, user)
+
 
 class AzureADHybridOAuthProvider(OAuthProvider):
     id = "azure-ad-hybrid"
@@ -258,9 +288,9 @@ async def get_user_info(self, token: str):
                 )
                 photo_data = await photo_response.aread()
                 base64_image = base64.b64encode(photo_data)
-                azure_user[
-                    "image"
-                ] = f"data:{photo_response.headers['Content-Type']};base64,{base64_image.decode('utf-8')}"
+                azure_user["image"] = (
+                    f"data:{photo_response.headers['Content-Type']};base64,{base64_image.decode('utf-8')}"
+                )
             except Exception as e:
                 # Ignore errors getting the photo
                 pass
diff --git a/backend/pyproject.toml b/backend/pyproject.toml
index da89631b47..12247f5625 100644
--- a/backend/pyproject.toml
+++ b/backend/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "chainlit"
-version = "1.1.300"
+version = "1.1.301"
 keywords = [
     'LLM',
     'Agents',