Update controllers/ws_controller.py

controllers/ws_controller.py

@@ -20,7 +20,9 @@ def get_company_type(company_name:str)->str:
     return company_types_dict.get(company_name.lower(), 'Others')
 
 async def get_messages(code: str,websocket:WebSocket,start,brand_name: Optional[str] = None):
+    await websocket.send_text("Test text!!")
     access_token = code
+    total_processed = 0
     g_query = f'(subject:"your order" OR subject:receipts OR subject:receipt OR subject:aankoopbon OR subject:reçu OR subject:invoice OR subject:invoices OR category:purchases) has:attachment'
     if brand_name is not None:
         g_query = f'(subject:"your order" OR subject:receipts OR subject:receipt OR  subject: aankoopbon  OR subject:reçu OR subject:invoice OR subject:invoices OR category:purchases OR from:{brand_name}) AND subject:{brand_name} has:attachment'    
@@ -44,7 +46,7 @@ async def get_messages(code: str,websocket:WebSocket,start,brand_name: Optional[
     start1 = time.time()
     while True:
         start2= time.time()
-        gmail_url = f"https://www.googleapis.com/gmail/v1/users/me/messages?q={g_query}&maxResults=15"
+        gmail_url = f"https://www.googleapis.com/gmail/v1/users/me/messages?q={g_query}&maxResults={30}"
         if page_token:
             gmail_url += f"&pageToken={page_token}"
         # print(gmail_url)
@@ -53,57 +55,67 @@ async def get_messages(code: str,websocket:WebSocket,start,brand_name: Optional[
         # print(gmail_response.text)
         end2 = time.time()
 
-        # print("End 2 ")
-        # print(end2-start2)
-        # print("response length")
-        # print(gmail_response.content.__len__())
+        print("End 2 ")
+        print(end2-start2)
+        print("response length")
+        print(gmail_response.content.__len__())
         gmail_data = gmail_response.json()
         # print(gmail_data)
-        # print(len(gmail_data['messages']))
+        print(len(gmail_data['messages']))
        
             
         if "messages" in gmail_data:
-            with ThreadPoolExecutor(max_workers=20) as executor:
+            # for message_data in gmail_data['messages']:
+            #     message = fetch_message_wrapper(message_data,websocket)
+            #     await process_message(message, websocket, 2000000)
+            with ThreadPoolExecutor(max_workers=50) as executor:
 
                             
                     
                 futures=[executor.submit(fetch_message_wrapper, message_data,websocket) for message_data in
                                gmail_data["messages"]]
-                # print(len(futures))
+                print("Futures length")
+                print(len(futures))
                 # print(futures)
                 for future in concurrent.futures.as_completed(futures):
                     message = future.result()
                     # print(message)
                     if message:
+                        total_processed += 1
                         # Process and send the message immediately
-                        await process_message(message, websocket, 50000)
+                        end1 = time.time()
+                        print("time 1")
+                        print("sending the message")
+                        await process_message(message, websocket, 100000)
                         
                       
-                    # if message:
-                    #     messages.append(message)
+            #         # if message:
+            #         #     messages.append(message)
             # print("Messages to be sent")
-            # print(messages)
+            # # print(messages)s
             # print(len(messages))
-            # for message_data in messages:
-            #     await process_message(message_data,websocket,10000)
+            # # for message_data in messages:
+            # #     await process_message(message_data,websocket,10000)
             
 
         if "nextPageToken" in gmail_data:
             page_token = gmail_data["nextPageToken"]
         else:
             break
-    # print("printing messages")
+    print(f"Total messages processed: {total_processed}")
+    logging.info(f"Total Processed Messages : {total_processed}")
+    print("printing messages")
     # print(messages)
     return messages
 
-
 async def process_message(message:Message, websocket:WebSocket, chunk_size:int):
-    # logging.info("process_message")
-    # print(message)
+    logging.info("process_message")
+    print(message)
     if message:
         message_json = message.to_json()
         # logging.info(f"{message_json}")
-        await send_message_in_chunks(websocket, message_json, chunk_size)
+        await send_message_in_chunks(websocket, message_json, 50000)
+        # await websocket.send_text(str(message_json))
         await websocket.send_text("NEXT_MESSAGE")
 
 
@@ -113,6 +125,8 @@ def fetch_message_data(access_token: str, message_id: str) -> Message:
     message_data = message_response.json()
     # print(message_data)
     subject = extract_subject_from_mail(message_data)
+    print("printing message data")
+    print(message_data)
     company_from_mail = extract_domain_name(message_data['payload']['headers'], subject)
 
     body = extract_body_from_mail(message_data)
@@ -132,6 +146,48 @@ def fetch_message_data(access_token: str, message_id: str) -> Message:
     return Message(message_id=message_id, body_len=body_len,body=body, attachments=attachments, company=company_from_mail,high_level_company_type=high_level_company_type,structured_data = structed_attachment_data)
 
 
+# def fetch_message_data(access_token: str, message_id: str) -> Message:
+#     message_url = f"https://www.googleapis.com/gmail/v1/users/me/messages/{message_id}"
+#     message_response = requests.get(message_url, headers={"Authorization": f"Bearer {access_token}"})
+#     message_data = message_response.json()
+
+#     with ThreadPoolExecutor(max_workers=50) as executor:
+#         # Submit tasks to executor
+#         future_subject = executor.submit(extract_subject_from_mail, message_data)
+#         subject = future_subject.result()
+#         future_company_from_mail = executor.submit(extract_domain_name, message_data['payload']['headers'], future_subject.result())
+#         company_from_mail = future_company_from_mail.result()
+#         future_body = executor.submit(extract_body_from_mail, message_data)
+#         body = future_body.result()
+        
+#         # Extract attachments and measure time
+ 
+#         future_attachments = executor.submit(extract_attachments_from_mail, access_token, message_data)
+#         attachments, structured_attachment_data = future_attachments.result()
+  
+
+#         future_high_level_company_type = executor.submit(get_company_type, future_company_from_mail.result())
+
+#         high_level_company_type = future_high_level_company_type.result()
+
+#     body_len = len(body) if body is not None else 0
+
+#     return Message(
+#         message_id=message_id,
+#         body_len=body_len,
+#         body=body,
+#         attachments=attachments,
+#         company=company_from_mail,
+#         high_level_company_type=high_level_company_type,
+#         structured_data=structured_attachment_data
+#     )
+
+
+
+
+
+
+
 
 def extract_subject_from_mail(message_data: dict) -> str:
     if 'payload' in message_data and 'headers' in message_data['payload']:
@@ -215,20 +271,10 @@ def extract_attachments_from_mail(access_token: str, message_data: dict) -> List
                 if filename.endswith(".zip") or filename.endswith(".txt") or filename.endswith(".png") or filename.endswith(".jpg") or filename.endswith(".jpeg") or filename.endswith(".gif"):
                     continue
                 data = attachment_data.get("data", "")
-
-                
-                try:
-                    raw_text = ut.extract_text_from_attachment(filename, data)
-                    struct_data = ut.strcuture_document_data(raw_text)
-                    if struct_data:
-                        structured_data.append(struct_data)
-                except Exception as e:
-                    print(f"Error processing attachment {filename}: {str(e)}")
-                    continue  # Skip this attachment if there's an error
-                # raw_text=ut.extract_text_from_attachment(filename , data)
-                # struct_data = ut.strcuture_document_data(raw_text)
-                # if struct_data:
-                #     structured_data.append(struct_data)
+                raw_text=ut.extract_text_from_attachment(filename , data)
+                struct_data = ut.strcuture_document_data(raw_text,filename)
+                if struct_data:
+                    structured_data.append(struct_data)
                 
                 attachments.append(Attachment(attachment_len = len(attachment_data.get("data", "")),filename=filename, data=attachment_data.get("data", "")))
     return attachments,structured_data
@@ -268,7 +314,23 @@ async def websocket_main(code: str,  websocket: WebSocket,start,brand_name: Opti
     await websocket.close()
 
 
-async def send_message_in_chunks(websocket: WebSocket, message_json: dict, chunk_size: int):         
+async def send_message_in_chunks(websocket: WebSocket, message_json: dict, chunk_size: int):
+
+    # if message_json['attachments'] is not None :
+    #     for attch in message_json['attachments']:
+    #         attachment_len = attch['attachment_len']
+            
+
+    # print(body_len)
+    # print(attachment_len)
+    # if attachment_len == 0:
+    #     attachment_len = None
+    # await websocket.send_json({"body_len":body_len ,"attachment_len":attachment_len})        
+    
     message_str = json.dumps(message_json)
+    # print("Printing message_str")
+    # print(message_str)
+    # logging.info(message_str)
+    # await websocket.send_json({"file_len":len(file)})
     for i in range(0, len(message_str), chunk_size):
         await websocket.send_text(message_str[i:i + chunk_size])