Update app.py

app.py

@@ -8,7 +8,6 @@ import email, imaplib, json, time
 import torch, logging
 import uvicorn
 from pydantic import BaseModel 
-import pandas as pd 
 
 app = FastAPI()
 
@@ -17,14 +16,18 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(
 logger = logging.getLogger(__name__)
 
 # Email and database configuration
+IMAP_SERVER = 'imap.gmail.com'
+EMAIL_ADDRESS = '[email protected]'
+PASSWORD = 'gclc wsnx kywt uvqy '  # Store this securely in production
 DB_CONFIG = {
     'host': '0.tcp.in.ngrok.io',
-    'port': 19266,
+    'port': 11329,
     'user': 'root',
-    'password': '',
+    'password': '',  # Add the correct password
     'database': 'shipment_details'
-    }
+}
 
+# JSON format for extracted shipment details
 output_format = {
     "origin": "",
     "destination": "",
@@ -34,25 +37,26 @@ output_format = {
     "Description": "",
     "Quantities": "",
     "Carrier_details": ""
-  }
+}
 
-prompt = f"""
+# Prompt for LLM to process shipment-related emails
+prompt = """
 System prompt: You will be provided with an email containing shipment details. Your task is to extract specific information based on the given instructions.
 
 Instructions:
-1. The input email may contain irrelevant information. Focus only on extracting details about future shipments.
-2. The output should be in JSON format. If a type of information is not found, it should be marked as null.
-3. Extract the following information:
-    - origin: The origin location of the consignment.
-    - destination: The destination location of the consignment.
-    - expected_shipment_datetime: The expected date and time of delivery to the warehouse (format: yyyy-mm-dd hh:mm:ss).
-    - types_of_service: The type of service (AIR, LCL, FCL). AIR can be mentioned as flight, aeroplane, or any mode of air transport. LCL is a Less-Container Load, and FCL is a Full-Container Load.
-    - warehouse: The name of the warehouse.
-    - description: A brief description of the email (ASN).
-    - quantities: The number of items in the shipment.
-    - carrier_details: The details of the carrier.
-4. the output extracted information contains must be in this format:
-{{
+1. Focus only on extracting details about future shipments, ignore irrelevant information.
+2. Output should be in JSON format. Missing information should be marked as null.
+3. Extract the following:
+    - origin
+    - destination
+    - expected_shipment_datetime (format: yyyy-mm-dd hh:mm:ss)
+    - types_of_service (AIR, LCL, FCL)
+    - warehouse
+    - description (max 100 words)
+    - quantities
+    - carrier_details
+4. The output should be formatted as follows:
+{
     "origin": "",
     "destination": "",
     "expected_shipment_datetime": "",
@@ -61,68 +65,9 @@ Instructions:
     "description": "",
     "quantities": "",
     "carrier_details": ""
-}}
-Examples:
-
-1. Email: We are pleased to inform you of an upcoming shipment originating from Hamburg and destined for New York. The shipment is expected to arrive on August 15, 2024. This consignment includes various electronics, with an estimated quantity of 200 units. The service type for this shipment is AIR, provided by our reliable carrier, Sky Logistics.
-   Extracted Information:
-        origin: Hamburg,
-        destination: New York,
-        expected_shipment_datetime: 2024-08-15 00:00:000,
-        types_of_service: AIR,
-        warehouse: Sky Logistics,
-        description: We are pleased to inform you of an upcoming shipment originating from Hamburg and destined for New York. The shipment is expected to arrive on August 15, 2024.,
-        quantities: 200 units,
-        carrier_details: Sky Logistics
-
-2. Email: Please be advised of a shipment from our supplier in Shanghai heading to Los Angeles. The expected date of arrival is July 30, 2024. The shipment consists of mixed goods, mainly textiles, with a total of 500 pieces. This delivery will be handled through LCL service by Ocean Freight Co.
-   Extracted Information:
-        origin: Shanghai,
-        destination: Los Angeles,
-        expected_shipment_datetime: 2024-07-30 00:00:0000,
-        types_of_service: LCL,
-        warehouse: Ocean Freight Co.,
-        description: Please be advised of a shipment from our supplier in Shanghai heading to Los Angeles. The expected date of arrival is July 30, 2024.,
-        quantities: 500 pieces,
-        carrier_details: Ocean Freight Co.
-
-3. Email: A new shipment is on its way from Mumbai to London, scheduled to reach by August 22, 2024. This batch contains furniture items, totaling 150 pieces. It is managed by Global Carriers.
-   Extracted Information:
-        origin: Mumbai,
-        destination: London,
-        expected_shipment_datetime: 2024-08-22 00:00:00000,
-        types_of_service: null,
-        warehouse: Global Carriers,
-        description: A new shipment is on its way from Mumbai to London, scheduled to reach by August 22, 2024.,
-        quantities: 150 pieces,
-        carrier_details: Global Carriers
-
-4. Email: We are notifying you about a shipment dispatched from Tokyo, heading towards Sydney, with an estimated arrival date of September 10, 2024. The cargo includes automotive parts, summing up to 350 units. This shipment will be transported via AIR service, operated by Jet Logistics.
-   Extracted Information:
-        origin: Tokyo,
-        destination: Sydney,
-        expected_shipment_datetime: 2024-09-10 00:00:0000,
-        types_of_service: AIR,
-        warehouse: Jet Logistics,
-        description: We are notifying you about a shipment dispatched from Tokyo, heading towards Sydney, with an estimated arrival date of September 10, 2024.,
-        quantities: 350 units,
-        carrier_details: Jet Logistics
-
-5. Email: Kindly note the details of a forthcoming shipment from Berlin to Toronto. The shipment encompasses various household goods, with a total quantity of 400 items. We have arranged for this to be shipped using LCL service, provided by Sea Wave Transport.
-   Extracted Information:
-        origin: Berlin,
-        destination: Toronto,
-        expected_shipment_datetime: null,
-        types_of_service: LCL,
-        warehouse: Sea Wave Transport,
-        description: Kindly note the details of a forthcoming shipment from Berlin to Toronto. The expected arrival is on August 5, 2024.,
-        quantities: 400 items,
-        carrier_details: Sea Wave Transport
-
-Output: {output_format}
+}
 """
 
-
 # Function to insert extracted shipment details into MySQL database
 def insert_data(extracted_details):
     try:
@@ -172,50 +117,86 @@ def insert_data(extracted_details):
     except Exception as ex:
         logger.error(f"Error inserting data: {ex}")
 
+# Function to extract shipment details using an LLM
+def get_details(mail):
+    try:
+        # Initialize LLM model and tokenizer
+        # Uncomment below if using Hugging Face models, or load your specific model accordingly
+        # pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+        # output = pipe(f"{prompt}\n{mail}", max_new_tokens=200)
+
+        # Using Llama model for completion
+        llm = Llama(model_path="./ggml-model-q8_0.gguf", n_ctx=2048, n_batch=512)
+        response = llm.create_chat_completion(
+            messages=[
+                {"role": "system", "content": prompt},
+                {"role": "user", "content": mail}
+            ],
+            max_tokens=200
+        )
+        return response['choices'][0]['message']['content']
+
+    except Exception as ex:
+        logger.error(f"Error generating details from LLM: {ex}")
+        return None
+
 # Function to read and process unread emails
 def read_email():
     logging.info('ready to read email ! ...')
     try:
-        tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-128k-instruct", trust_remote_code=True, cache_dir='./phi3_mini')
-        
-        # Load the model
-        model = AutoModelForCausalLM.from_pretrained(
-            "microsoft/Phi-3-mini-128k-instruct",
-            load_in_8bit=True,
-            device_map="auto",  # Automatically map layers to available device (GPU/CPU)
-            trust_remote_code=True,
-            cache_dir='./phi3_mini'
-        )
-        
-        df = pd.read_csv('emails.csv')
-        
-        for i in df['Body']:
-          pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
-          prompt_ = f"\n[system]:{prompt}\n[User]: {i}\n[Assistant]:"
-          output = pipe(prompt_, max_new_tokens=200, )
-          print('*'*10,i,'\n')
-          a = output[0]['generated_text']
-          t = a[a.find('[Assistant]:')+len('[Assistant]:') :]
-          jon = t[:t.find('\n}\n')+len('\n}\n')]
-          jon = json.loads(jon)
-          # print(type(jon),'\n',jon)
-          meta_data = {
-                            'sender': 'sender',
-                            'receiver': 'receiver',
-                            'cc': 'cc',
-                            'bcc': 'bcc',
-                            'subject': 'subject'
-                        }
-          jon.update(meta_data)
-          print(jon)
-          insert_data(jon)
+        logging.info('get imap server!')
+        mail = imaplib.IMAP4_SSL(IMAP_SERVER)
+        mail.login(EMAIL_ADDRESS, PASSWORD)
+        mail.select('inbox')
+        logging.info('select mail inbox')
+        status, messages = mail.search(None, 'UNSEEN')
+        message_ids = messages[0].split()
+        logging.info(f"Total unread emails: {len(message_ids)}")
+        print(f"Total unread emails: {len(message_ids)}")
+
+        for message_id in message_ids:
+            try:
+                status, data = mail.fetch(message_id, '(RFC822)')
+                raw_email = data[0][1]
+                email_message = email.message_from_bytes(raw_email)
+
+                # Extract metadata
+                sender = email_message['From']
+                receiver = email_message['To']
+                cc = email_message.get('Cc', '')
+                bcc = email_message.get('Bcc', '')
+                subject = email_message['Subject']
+
+                # Extract email body
+                if email_message.is_multipart():
+                    for part in email_message.walk():
+                        if part.get_content_type() == 'text/plain':
+                            email_body = part.get_payload(decode=True).decode('utf-8')
+                            break
+                else:
+                    email_body = email_message.get_payload(decode=True).decode('utf-8')
+
+                # Extract and store details
+                extracted_details_str = get_details(email_body)
+                extracted_details = json.loads(extracted_details_str)
+                meta_data = {
+                    'sender': sender, 'receiver': receiver, 'cc': cc, 'bcc': bcc, 'subject': subject
+                }
+                extracted_details.update(meta_data)
+                insert_data(extracted_details)
+
+            except Exception as e:
+                logger.error(f"Error processing email {message_id}: {e}")
+
+        mail.close()
+        mail.logout()
 
     except Exception as e:
         logger.error(f"Error reading emails: {e}")
 
-# Global variables
+# Email processing loop
 running = False
-loop_thread = None
+loop_thread = None 
 
 # HTML content for the web interface
 html_content = """
@@ -237,37 +218,28 @@ html_content = """
             cursor: pointer;
         }
         button.stop { background-color: #f44336; }
-        #status { font-weight: bold; }
     </style>
-    <script>
-        async function startLoop() {
-            const response = await fetch('/start', { method: 'POST' });
-            const result = await response.text();
-            document.getElementById("status").innerHTML = result;
-        }
-
-        async function stopLoop() {
-            const response = await fetch('/stop', { method: 'POST' });
-            const result = await response.text();
-            document.getElementById("status").innerHTML = result;
-        }
-    </script>
 </head>
 <body>
-    <h1>Email Processing Status: <span id="status">{{ status }}</span></h1>
-    <button onclick="startLoop()">Start</button>
-    <button class="stop" onclick="stopLoop()">Stop</button>
+    <h1>Email Processing Status: {{ status }}</h1>
 </body>
 </html>
 """
 
+class ActionModel(BaseModel):
+    action: str  # 'start' or 'stop'
+
+class ModelData(BaseModel):
+    data: str  # Additional email-related information
+
 # Function to process emails in a loop
 def email_processing_loop():
     global running
-    logging.info("Starting email processing loop...")
-    read_email()
-    print("$" * 100)  # Simulate email processing
-    time.sleep(10)  # Check for new emails every 10 seconds
+    logger.info("Starting email processing loop...")
+    while running:
+        # read_email()    # Assuming this is your email processing function
+        print("$"*100)
+        time.sleep(10)  # Check for new emails every 10 seconds
 
 # Endpoint to display the current email processor status
 @app.get("/", response_class=HTMLResponse)
@@ -275,33 +247,40 @@ async def home():
     global running
     status = "Running" if running else "Stopped"
     return HTMLResponse(content=html_content.replace("{{ status }}", status), status_code=200)
-
-# Endpoint to start the email processing loop
-@app.post("/start")
-async def start_email_loop():
+async def control_email_loop(action: ActionModel, data: ModelData):
     global running, loop_thread
-    if not running:
-        running = True
-        loop_thread = threading.Thread(target=email_processing_loop, daemon=True)
-        loop_thread.start()
-        logging.info("Email processing loop started.")
-        return "Running"
+    logger.info(action.action)
+    
+    if action.action == "start":
+        if not running:
+            running = True
+            email_data = data.data  # This is already a string, no need to call .dict()
+            logger.info(f"Email Data: {email_data}")
+            loop_thread = threading.Thread(target=email_processing_loop, daemon=True)
+            loop_thread.start()
+            logger.info("Email processing loop started.")
+        else:
+            logger.info("Email processing loop is already running.")
+    
+    elif action.action == "stop":
+        if running:
+            running = False
+            logger.info("Email processing loop stopped.")
+        else:
+            logger.info("Email processing loop is not running.")
+    
     else:
-        return "Already running"
+        raise HTTPException(status_code=400, detail="Invalid action. Use 'start' or 'stop'.")
+    
+    status = "Running" if running else "Stopped"
+    return HTMLResponse(content=html_content.replace("{{ status }}", status), status_code=200)
 
-# Endpoint to stop the email processing loop
-@app.post("/stop")
-async def stop_email_loop():
-    global running
-    if running:
-        running = False
-        logging.info("Email processing loop stopped.")
-        return "Stopped"
-    else:
-        return "Already stopped"
 
 if __name__ == "__main__":
-    logging.basicConfig(level=logging.INFO)
-    logging.info("Starting FastAPI server...")
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    
+    print('starting')
+    logging.info('starting project!...')
+    # running = True
+    # threading.Thread(target=email_processing_loop, daemon=True).start()
+    logging.info('...')
+    uvicorn.run()