From 1a8dee372ae986cb7e0d3721b3847e909f4be31a Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Sat, 27 Jan 2024 13:57:04 -0800 Subject: [PATCH 1/3] (feat) add google login to litellm proxy --- litellm/proxy/proxy_server.py | 65 +++++++++++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d9b8fcce2d9..ac1dbcc2415 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -104,6 +104,7 @@ from fastapi.responses import ( ORJSONResponse, JSONResponse, ) +from fastapi.responses import RedirectResponse from fastapi.middleware.cors import CORSMiddleware from fastapi.security.api_key import APIKeyHeader import json @@ -2856,6 +2857,70 @@ async def user_auth(request: Request): return "Email sent!" +@app.get("/google-login", tags=["experimental"]) +async def google_login(): + GOOGLE_REDIRECT_URI = "http://localhost:4000/google-callback" + GOOGLE_CLIENT_ID = ( + "246483686424-clje5sggkjma26ilktj6qssakqhoon0m.apps.googleusercontent.com" + ) + google_auth_url = f"https://accounts.google.com/o/oauth2/auth?client_id={GOOGLE_CLIENT_ID}&redirect_uri={GOOGLE_REDIRECT_URI}&response_type=code&scope=openid%20profile%20email" + return RedirectResponse(url=google_auth_url) + + +@app.get("/google-callback", tags=["experimental"]) +async def google_callback(code: str): + import httpx + + GOOGLE_REDIRECT_URI = "http://localhost:4000/google-callback" + GOOGLE_CLIENT_ID = ( + "246483686424-clje5sggkjma26ilktj6qssakqhoon0m.apps.googleusercontent.com" + ) + # Exchange code for access token + async with httpx.AsyncClient() as client: + token_url = f"https://oauth2.googleapis.com/token" + data = { + "code": code, + "client_id": GOOGLE_CLIENT_ID, + "client_secret": "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", + "redirect_uri": GOOGLE_REDIRECT_URI, + "grant_type": "authorization_code", + } + response = await client.post(token_url, data=data) + + # Process the response, extract user info, etc. + if response.status_code == 200: + access_token = response.json()["access_token"] + + # Fetch user info using the access token + async with httpx.AsyncClient() as client: + user_info_url = "https://www.googleapis.com/oauth2/v1/userinfo" + headers = {"Authorization": f"Bearer {access_token}"} + user_info_response = await client.get(user_info_url, headers=headers) + + # Process user info response + if user_info_response.status_code == 200: + user_info = user_info_response.json() + user_email = user_info.get("email") + user_name = user_info.get("name") + + # we can use user_email on litellm proxy now + + # TODO: Handle user info as needed, for example, store it in a database, authenticate the user, etc. + return JSONResponse( + content={"user_email": user_email, "user_name": user_name}, + status_code=200, + ) + else: + # Handle user info retrieval error + raise HTTPException( + status_code=user_info_response.status_code, + detail=user_info_response.text, + ) + else: + # Handle the error from the token exchange + raise HTTPException(status_code=response.status_code, detail=response.text) + + @router.get( "/user/info", tags=["user management"], dependencies=[Depends(user_api_key_auth)] ) From 180a1a7f9c730a66fad297d9caaa780c6b1e4ec0 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Sat, 27 Jan 2024 14:17:03 -0800 Subject: [PATCH 2/3] (fix) return dict response --- litellm/proxy/proxy_server.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ac1dbcc2415..749705157c0 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2867,7 +2867,7 @@ async def google_login(): return RedirectResponse(url=google_auth_url) -@app.get("/google-callback", tags=["experimental"]) +@app.get("/google-callback", tags=["experimental"], response_model=GenerateKeyResponse) async def google_callback(code: str): import httpx @@ -2906,10 +2906,16 @@ async def google_callback(code: str): # we can use user_email on litellm proxy now # TODO: Handle user info as needed, for example, store it in a database, authenticate the user, etc. - return JSONResponse( - content={"user_email": user_email, "user_name": user_name}, - status_code=200, + response = await generate_key_helper_fn( + **{"duration": "24hr", "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": user_email, "team_id": "litellm-dashboard"} # type: ignore ) + + key = response["token"] # type: ignore + user_id = response["user_id"] # type: ignore + { + "key": key, + "user_id": user_id, + } else: # Handle user info retrieval error raise HTTPException( From a7bbc0de0ff5bf4d3801ba73957600629146323d Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Sat, 27 Jan 2024 14:19:11 -0800 Subject: [PATCH 3/3] (fix) google-login auth flow proxy --- litellm/proxy/proxy_server.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 749705157c0..61103e079cb 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2912,10 +2912,10 @@ async def google_callback(code: str): key = response["token"] # type: ignore user_id = response["user_id"] # type: ignore - { - "key": key, - "user_id": user_id, - } + return JSONResponse( + content={"key": key, "user_id": user_id}, status_code=200 + ) + else: # Handle user info retrieval error raise HTTPException(