Resetting utils subtree.

models/behaviour/mail/oauth_v2.py

@@ -101,6 +101,7 @@ class MailOAuthModel(BaseModel):
             mongo_conn: AsyncMongo,
             user_info: dict,
             client_user_id: dict,
+            auth: dict,
             service_client: Literal["gmail"],
             auth_type: Literal["oauth"],
             sync_freq: Literal[60, 300, 900] = 300,
@@ -114,6 +115,7 @@ class MailOAuthModel(BaseModel):
         :param mongo_conn: The database connection (MongoDB) to use to perform the action.
         :param user_info: The dictionary that has the user's session information.
         :param client_user_id: The way the third-party client recognizes your user.
+        :param auth: The authentication details of the account.
         :param service_client: The name of the company or brand that is providing this service that is being integrated.
         :param auth_type: To identify the type of authentication being done here. This could indicate simple password
             authentication, more advance OAuth2.0 authentication, etc.
@@ -149,6 +151,7 @@ class MailOAuthModel(BaseModel):
                     "authType": auth_type,
                     "user": user_info,
                     "clientUserId": client_user_id,
+                    "auth": auth,
                     "token": None,
                     "firstRefreshTs": None,
                     "lastRefreshTs": None,
@@ -171,8 +174,8 @@ class MailOAuthModel(BaseModel):
                 proc_args = (
                     user_info["entityId"],  # ............................................ 'p_entity_id'
                     service_client,  # ................................................... 'p_provider'
-                    "Auth Requested",  # ................................................. 'p_current_status'
-                    "Auth URL Generated",  # ............................................. 'p_last_action'
+                    "Pending",  # ........................................................ 'p_current_status'
+                    "Auth Requested",  # ................................................. 'p_last_action'
                     None,  # ............................................................. 'p_display_name'
                     None,  # ............................................................. 'p_display_picture'
                     str(mongo_json["_id"]),  # ........................................... 'p_token_id'
@@ -250,8 +253,8 @@ class MailOAuthModel(BaseModel):
         if mongo_json is not None:
             token_notes = {
                 "email": token["email"],
-                "displayName": token["displayName"],
-                "displayPictureUrl": token["displayPictureUrl"],
+                "displayName": token.get("displayName"),
+                "displayPictureUrl": token.get("displayPictureUrl"),
             }
             db_json = await self.call_procedure(
                 db_conn = db_conn,
@@ -259,8 +262,8 @@ class MailOAuthModel(BaseModel):
                 proc_args = (
                     mongo_json["user"]["entityId"],  # ............................... 'p_entity_id'
                     mongo_json["client"],  # ......................................... 'p_provider'
-                    "Auth Granted",  # ............................................... 'p_current_status'
-                    "Set Token",  # .................................................. 'p_last_action'
+                    "Active",  # ..................................................... 'p_current_status'
+                    "Auth Granted",  # ............................................... 'p_last_action'
                     token["displayName"],  # ......................................... 'p_display_name'
                     token["displayPictureUrl"],  # ................................... 'p_display_picture'
                     token_id,  # ..................................................... 'p_token_id'

models/behaviour/mail/retrieve.py

@@ -124,9 +124,7 @@ class MailRetrieveModel(BaseModel):
                 "payload.bcc": True,
                 "payload.parts": True,
                 "payload.attachments": True,
-                "payload.labels": True,
-                "payload.snippet": True,
-                "payload.aiSnippet": True,
+                "payload.labels": True
             }
         )
 
@@ -185,7 +183,10 @@ class MailRetrieveModel(BaseModel):
                 mail_data["payload"]["ts"] = mail_data["payload"]["ts"].isoformat()
                 mail_data["payload"]["readTs"] = mail_data["payload"]["readTs"].isoformat()
                 if ai_snippet := mail_data["payload"].pop("aiSnippet"):
-                    mail_data["payload"]["aiSnippet"] = ai_snippet["snippet"]
+                    mail_data["payload"]["aiSnippet"] = {
+                        "snippet": ai_snippet["snippet"],
+                        "usage": ai_snippet["usage"]
+                    }
 
         # Done here:
         return mails_list

models/behaviour/mail/sync_v2.py

@@ -331,7 +331,12 @@ class MailSyncModel(BaseModel):
             )
             llm_json = {
                 "snippet": llm_response.content,
-                "usage": llm_response.usage_metadata
+                "usage": {
+                    "input": llm_response.usage_metadata["input_tokens"],
+                    "output": llm_response.usage_metadata["output_tokens"],
+                    "total": llm_response.usage_metadata["total_tokens"],
+                },
+                "rawUsage": llm_response.usage_metadata
             }
         client_response.data["aiSnippet"] = llm_json