Update controllers/ws_controller.py

controllers/ws_controller.py

@@ -2,6 +2,7 @@ import base64
 import json
 import logging
 import re
+from concurrent.futures import ProcessPoolExecutor
 from concurrent.futures import ThreadPoolExecutor
 from typing import Optional, List, Dict
 import requests
@@ -9,7 +10,7 @@ from bs4 import BeautifulSoup
 from models.models import Message, Attachment
 from fastapi import WebSocket
 from services import utils as ut
-from models import supabase_models as sp
+# from models import supabase_models as sp
 import asyncio
 def get_company_type(company_name:str)->str:
     company_types_dict ={'ao yun': 'wines and spirit', 'ardbeg': 'wines and spirit', 'belvedere': 'wines and spirit', 'bodega numanthia': 'wines and spirit', 'chandon': 'wines and spirit', 'château cheval blanc': 'wines and spirit', "château d'yquem": 'wines and spirit', 'château galoupet': 'wines and spirit', 'cheval des andes': 'wines and spirit', 'clos19': 'wines and spirit', 'cloudy bay': 'wines and spirit', 'colgin cellars': 'wines and spirit', 'dom pérignon': 'wines and spirit', 'domaine des lambrays': 'wines and spirit', 'eminente': 'wines and spirit', 'glenmorangie': 'wines and spirit', 'hennessy': 'wines and spirit', 'joseph phelps': 'wines and spirit', 'krug': 'wines and spirit', 'mercier': 'wines and spirit', 'moët & chandon': 'wines and spirit', 'newton vineyard': 'wines and spirit', 'ruinart': 'wines and spirit', 'terrazas de los andes': 'wines and spirit', 'veuve clicquot': 'wines and spirit', 'volcan de mi tierra': 'wines and spirit', 'woodinville': 'wines and spirit' , 'berluti': 'Fashion & Leather Goods', 'celine': 'Fashion & Leather Goods', 'christian dior': 'Fashion & Leather Goods', 'emilio pucci': 'Fashion & Leather Goods', 'fendi': 'Fashion & Leather Goods', 'givenchy': 'Fashion & Leather Goods', 'kenzo': 'Fashion & Leather Goods', 'loewe': 'Fashion & Leather Goods', 'loro piana': 'Fashion & Leather Goods', 'louis vuitton': 'Fashion & Leather Goods', 'marc jacobs': 'Fashion & Leather Goods', 'moynat': 'Fashion & Leather Goods', 'patou': 'Fashion & Leather Goods', 'rimowa': 'Fashion & Leather Goods','acqua di parma': 'Perfumes & Cosmetics', 'benefit cosmetics': 'Perfumes & Cosmetics', 'cha ling': 'Perfumes & Cosmetics', 'fenty beauty by rihanna': 'Perfumes & Cosmetics', 'fresh': 'Perfumes & Cosmetics', 'givenchy parfums': 'Perfumes & Cosmetics', 'guerlain': 'Perfumes & Cosmetics', 'kenzo parfums': 'Perfumes & Cosmetics', 'kvd beauty': 'Perfumes & Cosmetics', 'loewe perfumes': 'Perfumes & Cosmetics', 'maison francis kurkdjian': 'Perfumes & Cosmetics', 'make up for ever': 'Perfumes & Cosmetics', 'officine universelle buly': 'Perfumes & Cosmetics', 'olehenriksen': 'Perfumes & Cosmetics', 'parfums christian dior': 'Perfumes & Cosmetics', 'stella by stella mccartney': 'Perfumes & Cosmetics','bulgari': 'Watches & Jewelry', 'chaumet': 'Watches & Jewelry', 'fred': 'Watches & Jewelry', 'hublot': 'Watches & Jewelry', 'repossi': 'Watches & Jewelry', 'tag heuer': 'Watches & Jewelry', 'tiffany & co.': 'Watches & Jewelry', 'zenith': 'Watches & Jewelry','24s': 'Selective retailing', 'dfs': 'Selective retailing', 'la grande epicerie de paris': 'Selective retailing', 'le bon marché rive gauche': 'Selective retailing', 'sephora': 'Selective retailing','belmond': 'Other activities', 'cheval blanc': 'Other activities', 'connaissance des arts': 'Other activities', 'cova': 'Other activities', 'investir': 'Other activities', "jardin d'acclimatation": 'Other activities', 'le parisien': 'Other activities', 'les echos': 'Other activities', 'radio classique': 'Other activities', 'royal van lent': 'Other activities'}
@@ -28,13 +29,21 @@ async def get_messages(code: str,websocket:WebSocket,brand_name: Optional[str] =
     # gmail_response = requests.get(gmail_url, headers={"Authorization": f"Bearer {access_token}"})
     # gmail_data = gmail_response.json()
     # messages.append(gmail_data['messages'])
-    def fetch_message_wrapper(message_data):
+    def fetch_message_wrapper(message_data,websocket:WebSocket):
         message_id = message_data.get("id")
         if message_id:
             return fetch_message_data(access_token, message_id)
             
         return None
 
+        # if message_id :
+        #     message_json = fetch_message_data(access_token,message_id)
+        #     # await process_message(message_json,websocket,20000)
+
+        # if message_id:
+        #     return await asyncio.to_thread(fetch_message_data, access_token, message_id)
+        # return None
+    # with ProcessPoolExecutor(max_workers=4) as executor:
     while True:
         gmail_url = f"https://www.googleapis.com/gmail/v1/users/me/messages?q={g_query}"
         if page_token:
@@ -42,22 +51,30 @@ async def get_messages(code: str,websocket:WebSocket,brand_name: Optional[str] =
 
         gmail_response = requests.get(gmail_url, headers={"Authorization": f"Bearer {access_token}"})
         gmail_data = gmail_response.json()
-        print(len(gmail_data))
+        print(len(gmail_data['messages']))
         print(gmail_data)
-        
+            
         if "messages" in gmail_data:
             with ThreadPoolExecutor(max_workers=15) as executor:
-                
-                           
-                
-                futures=[executor.submit(fetch_message_wrapper, message_data) for message_data in
-                           gmail_data["messages"]]
+
+                            
+                    
+                futures=[executor.submit(fetch_message_wrapper, message_data,websocket) for message_data in
+                               gmail_data["messages"]]
                 for future in futures:
                     message = future.result()
+                    print(message)
                     if message:
-                        messages.append(message)
-            for message_data in messages:
-                await process_message(message_data,websocket,10000)
+                        # Process and send the message immediately
+                        await process_message(message, websocket, 20000)
+                    # if message:
+                    #     messages.append(message)
+            print("Messages to be sent")
+            print(messages)
+            print(len(messages))
+            # for message_data in messages:
+            #     await process_message(message_data,websocket,10000)
+            
 
         if "nextPageToken" in gmail_data:
             page_token = gmail_data["nextPageToken"]
@@ -69,6 +86,7 @@ async def get_messages(code: str,websocket:WebSocket,brand_name: Optional[str] =
 
 async def process_message(message:Message, websocket:WebSocket, chunk_size:int):
     logging.info("process_message")
+    print(message)
     if message:
         message_json = message.to_json()
         logging.info(f"{message_json}")