Update controllers/ws_controller_1.py

controllers/ws_controller_1.py

@@ -9,10 +9,9 @@ from bs4 import BeautifulSoup
 from fastapi import  WebSocket
 from services import utils as ut
 from models.models import Message, Attachment
-from concurrent.futures import ThreadPoolExecutor,ProcessPoolExecutor
+from concurrent.futures import ThreadPoolExecutor
 import functools
-thread_pool = ThreadPoolExecutor(max_workers=20)  # Increased from 10
-process_pool = ProcessPoolExecutor(max_workers=4)
+thread_pool = ThreadPoolExecutor(max_workers=10) 
 
 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'}
@@ -145,11 +144,10 @@ async def get_messages(code: str, websocket: WebSocket,start, brand_name: Option
     
     page_token = None
     processed_count = 0
-    batch_size = 100  # Increased from 50
 
     async with aiohttp.ClientSession() as session:
         while True:
-            gmail_url = f"https://www.googleapis.com/gmail/v1/users/me/messages?q={g_query}&maxResults={batch_size}"
+            gmail_url = f"https://www.googleapis.com/gmail/v1/users/me/messages?q={g_query}&maxResults=50"
             if page_token:
                 gmail_url += f"&pageToken={page_token}"
 
@@ -160,23 +158,25 @@ async def get_messages(code: str, websocket: WebSocket,start, brand_name: Option
                 await websocket.send_text(json.dumps({"status": "No messages found"}))
                 break
 
-            message_ids = [message["id"] for message in gmail_data["messages"]]
-            full_messages = await fetch_messages_batch(session, access_token, message_ids)
+            message_tasks = [fetch_message(session, access_token, message["id"]) for message in gmail_data["messages"]]
+            full_messages = await asyncio.gather(*message_tasks)
 
             processing_tasks = [process_message(session, access_token, message) for message in full_messages]
             
             for future in asyncio.as_completed(processing_tasks):
                 processed_message = await future
                 await websocket.send_text(json.dumps({
-                    "message": processed_message.to_json()
+                    "status": "processing",
+                    "message": processed_message.to_json(),
+                    "processed_count": processed_count + 1
                 }))
                 processed_count += 1
 
             if "nextPageToken" in gmail_data:
                 page_token = gmail_data["nextPageToken"]
-                # await websocket.send_text(json.dumps({"status": "next_page"}))
+                await websocket.send_text(json.dumps({"status": "next_page"}))
             else:
-                # await websocket.send_text(json.dumps({"status": "completed", "total_processed": processed_count}))
+                await websocket.send_text(json.dumps({"status": "completed", "total_processed": processed_count}))
                 break
 
 async def send_message_in_chunks(websocket: WebSocket, message_json: dict, chunk_size: int):
@@ -196,22 +196,11 @@ async def websocket_main(code: str,  websocket: WebSocket,start,brand_name: Opti
 async def run_in_thread(func, *args, **kwargs):
     return await asyncio.get_event_loop().run_in_executor(thread_pool, functools.partial(func, *args, **kwargs))
 
-async def run_in_process(func, *args, **kwargs):
-    return await asyncio.get_event_loop().run_in_executor(process_pool, functools.partial(func, *args, **kwargs))
-
-async def fetch_messages_batch(session: aiohttp.ClientSession, access_token: str, message_ids: List[str]):
-    tasks = [fetch_message(session, access_token, message_id) for message_id in message_ids]
-    return await asyncio.gather(*tasks)
-
 async def fetch_message(session: aiohttp.ClientSession, access_token: str, message_id: str):
     message_url = f"https://www.googleapis.com/gmail/v1/users/me/messages/{message_id}"
     async with session.get(message_url, headers={"Authorization": f"Bearer {access_token}"}) as response:
         return await response.json()
 
-async def fetch_attachments_batch(session: aiohttp.ClientSession, access_token: str, attachments_info: List[tuple]):
-    tasks = [fetch_attachment(session, access_token, msg_id, att_id) for msg_id, att_id in attachments_info]
-    return await asyncio.gather(*tasks)
-
 async def fetch_attachment(session: aiohttp.ClientSession, access_token: str, message_id: str, attachment_id: str):
     attachment_url = f"https://www.googleapis.com/gmail/v1/users/me/messages/{message_id}/attachments/{attachment_id}"
     async with session.get(attachment_url, headers={"Authorization": f"Bearer {access_token}"}) as response:
@@ -230,21 +219,23 @@ async def process_message(session: aiohttp.ClientSession, access_token: str, mes
     structured_data = []
 
     if "payload" in message_data and "parts" in message_data["payload"]:
-        attachments_info = []
+        attachment_tasks = []
         for part in message_data["payload"]["parts"]:
             if "body" in part and "attachmentId" in part["body"]:
                 attachment_id = part["body"]["attachmentId"]
                 filename = part.get("filename", "untitled.txt")
-                if not filename.endswith((".zip", ".txt", ".png", ".jpg", ".jpeg", ".gif")):
-                    attachments_info.append((message_data["id"], attachment_id))
+                if filename.endswith((".zip", ".txt", ".png", ".jpg", ".jpeg", ".gif")):
+                    continue
+                
+                task = fetch_attachment(session, access_token, message_data["id"], attachment_id)
+                attachment_tasks.append((filename, task))
         
-        attachment_data_list = await fetch_attachments_batch(session, access_token, attachments_info)
+        attachment_results = await asyncio.gather(*(task for _, task in attachment_tasks))
         
         processing_tasks = []
-        for (_, _), attachment_data in zip(attachments_info, attachment_data_list):
+        for (filename, _), attachment_data in zip(attachment_tasks, attachment_results):
             data = attachment_data.get("data", "")
-            filename = next((part["filename"] for part in message_data["payload"]["parts"] if part["body"].get("attachmentId") == attachment_data["id"]), "untitled.txt")
-            processing_tasks.append(run_in_process(process_attachment, filename, data))
+            processing_tasks.append(run_in_thread(process_attachment, filename, data))
             attachments.append(Attachment(attachment_len=len(data), filename=filename, data=data))
         
         structured_data = await asyncio.gather(*processing_tasks)