Implement redirect url to captn

google_sheets/app.py

@@ -59,23 +59,22 @@ async def get_login_url(
     user_id: Annotated[
         int, Query(description="The user ID for which the data is requested")
     ],
+    conv_uuid: Annotated[
+        Optional[str], Query(description="The conversation UUID")
+    ] = None,
     force_new_login: Annotated[bool, Query(description="Force new login")] = False,
 ) -> Dict[str, str]:
+    _check_parameters_are_not_none({"conv_uuid": conv_uuid})
     if not force_new_login:
         is_authenticated = await is_authenticated_for_ads(user_id=user_id)
         if is_authenticated:
             return {"login_url": "User is already authenticated"}
 
-    google_oauth_url = get_google_oauth_url(user_id)
+    google_oauth_url = get_google_oauth_url(user_id, conv_uuid)  # type: ignore
     markdown_url = f"To navigate Google Ads waters, I require access to your account. Please [click here]({google_oauth_url}) to grant permission."
     return {"login_url": markdown_url}
 
 
-@app.get("/login/success", description="Get the success message after login")
-async def get_login_success() -> Dict[str, str]:
-    return {"login_success": "You have successfully logged in"}
-
-
 def _check_parameters_are_not_none(kwargs: Dict[str, Any]) -> None:
     error_message = "The following parameters are required: "
     missing_parameters = [key for key, value in kwargs.items() if value is None]
@@ -84,6 +83,9 @@ def _check_parameters_are_not_none(kwargs: Dict[str, Any]) -> None:
         raise HTTPException(status_code=400, detail=error_message)
 
 
+REDIRECT_DOMAIN = environ.get("REDIRECT_DOMAIN", "http://localhost:3000")
+
+
 # Route 2: Save user credentials/token to a JSON file
 @app.get("/login/callback")
 async def login_callback(
@@ -94,9 +96,11 @@ async def login_callback(
     state: Annotated[Optional[str], Query(description="State")] = None,
 ) -> RedirectResponse:
     _check_parameters_are_not_none({"state": state})
-    if not state.isdigit():  # type: ignore
+    user_id_and_chat_uuid = state.split(":")  # type: ignore
+    if not user_id_and_chat_uuid[0].isdigit():  # type: ignore
         raise HTTPException(status_code=400, detail="User ID must be an integer")
-    user_id = int(state)  # type: ignore
+    user_id = int(user_id_and_chat_uuid[0])
+    chat_uuid = user_id_and_chat_uuid[1]
 
     token_request_data = get_token_request_data(code)
 
@@ -132,12 +136,9 @@ async def login_callback(
             },
         )
 
-    # redirect_domain = environ.get("REDIRECT_DOMAIN", "https://captn.ai")
-    # logged_in_message = "I have successfully logged in"
-    # redirect_uri = f"{redirect_domain}/chat/{chat_uuid}?msg={logged_in_message}"
-    # return RedirectResponse(redirect_uri)
-    # redirect to success page
-    return RedirectResponse(url=f"{base_url}/login/success")
+    logged_in_message = "I have successfully logged in"
+    redirect_uri = f"{REDIRECT_DOMAIN}/chat/{chat_uuid}?msg={logged_in_message}"
+    return RedirectResponse(redirect_uri)
 
 
 @app.get("/get-sheet", description="Get data from a Google Sheet")

google_sheets/google_api/oauth_settings.py

@@ -20,12 +20,13 @@
 }
 
 
-def get_google_oauth_url(user_id: int) -> str:
+def get_google_oauth_url(user_id: int, conv_uuid: str) -> str:
+    state = f"{user_id}:{conv_uuid}"
     google_oauth_url = (
         f"{oauth2_settings['auth_uri']}?client_id={oauth2_settings['clientId']}"
         f"&redirect_uri={oauth2_settings['redirectUri']}&response_type=code"
         f"&scope={urllib.parse.quote_plus('email https://www.googleapis.com/auth/spreadsheets https://www.googleapis.com/auth/drive.metadata.readonly')}"
-        f"&access_type=offline&prompt=consent&state={user_id}"
+        f"&access_type=offline&prompt=consent&state={state}"
     )
     return google_oauth_url
 

tests/app/test_app.py

@@ -278,7 +278,6 @@ def test_openapi(self) -> None:
         paths = response.json()["paths"]
         expected_path_keys = [
             "/login",
-            "/login/success",
             "/login/callback",
             "/get-sheet",
             "/update-sheet",