added a favicon, updated the README.md file
This commit is contained in:
Christian Krakau-Louis
+72
-72
@@ -1,72 +1,72 @@
|
||||
#!/usr/bin/env python3
|
||||
import os
|
||||
import requests
|
||||
import logging
|
||||
import mimetypes
|
||||
from celery import shared_task
|
||||
from app.config import settings
|
||||
from app.tasks.upload_to_s3 import upload_to_s3
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@shared_task
|
||||
def convert_to_pdf(file_path):
|
||||
"""
|
||||
Converts a file to PDF using Gotenberg's API.
|
||||
Determines the appropriate Gotenberg endpoint based on the file's MIME type.
|
||||
On success, saves the PDF locally and enqueues it for S3 upload.
|
||||
"""
|
||||
gotenberg_url = getattr(settings, "gotenberg_url", None)
|
||||
if not gotenberg_url:
|
||||
logger.error("Gotenberg URL is not configured in settings.")
|
||||
return
|
||||
|
||||
# Try to guess the MIME type based on file content (using extension-based fallback)
|
||||
mime_type, encoding = mimetypes.guess_type(file_path)
|
||||
logger.info(f"Guessed MIME type for '{file_path}' is: {mime_type}")
|
||||
|
||||
endpoint = None
|
||||
form_key = "files" # Default form key for most endpoints
|
||||
|
||||
if mime_type:
|
||||
if mime_type == "text/html":
|
||||
endpoint = f"{gotenberg_url}/forms/chromium/convert/html"
|
||||
# The Chromium HTML endpoint expects the HTML file to be provided under the key "index.html"
|
||||
form_key = "index.html"
|
||||
elif mime_type.startswith("image/"):
|
||||
# For images, we use the LibreOffice endpoint (which supports image conversion)
|
||||
endpoint = f"{gotenberg_url}/forms/libreoffice/convert"
|
||||
elif mime_type.startswith("text/plain"):
|
||||
endpoint = f"{gotenberg_url}/forms/libreoffice/convert"
|
||||
elif mime_type in ["text/markdown", "text/x-markdown"]:
|
||||
# Optionally, you could use the Chromium markdown endpoint if you have an HTML wrapper.
|
||||
# For now, we'll fallback to LibreOffice.
|
||||
endpoint = f"{gotenberg_url}/forms/libreoffice/convert"
|
||||
else:
|
||||
# For all other MIME types (e.g. Office documents), use the LibreOffice endpoint.
|
||||
endpoint = f"{gotenberg_url}/forms/libreoffice/convert"
|
||||
else:
|
||||
# If MIME detection fails, fallback to extension-based detection.
|
||||
ext = os.path.splitext(file_path)[1].lower()
|
||||
if ext in [".html", ".htm"]:
|
||||
endpoint = f"{gotenberg_url}/forms/chromium/convert/html"
|
||||
form_key = "index.html"
|
||||
else:
|
||||
endpoint = f"{gotenberg_url}/forms/libreoffice/convert"
|
||||
|
||||
try:
|
||||
with open(file_path, "rb") as f:
|
||||
files = {form_key: f}
|
||||
response = requests.post(endpoint, files=files)
|
||||
|
||||
if response.status_code == 200:
|
||||
converted_file_path = os.path.splitext(file_path)[0] + ".pdf"
|
||||
with open(converted_file_path, "wb") as out_file:
|
||||
out_file.write(response.content)
|
||||
logger.info(f"Converted file saved as PDF: {converted_file_path}")
|
||||
upload_to_s3.delay(converted_file_path)
|
||||
return converted_file_path
|
||||
else:
|
||||
logger.error(f"Conversion failed for {file_path}. Status code: {response.status_code}")
|
||||
except Exception as e:
|
||||
logger.exception(f"Error converting {file_path} to PDF: {e}")
|
||||
#!/usr/bin/env python3
|
||||
import os
|
||||
import requests
|
||||
import logging
|
||||
import mimetypes
|
||||
from celery import shared_task
|
||||
from app.config import settings
|
||||
from app.tasks.upload_to_s3 import upload_to_s3
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@shared_task
|
||||
def convert_to_pdf(file_path):
|
||||
"""
|
||||
Converts a file to PDF using Gotenberg's API.
|
||||
Determines the appropriate Gotenberg endpoint based on the file's MIME type.
|
||||
On success, saves the PDF locally and enqueues it for S3 upload.
|
||||
"""
|
||||
gotenberg_url = getattr(settings, "gotenberg_url", None)
|
||||
if not gotenberg_url:
|
||||
logger.error("Gotenberg URL is not configured in settings.")
|
||||
return
|
||||
|
||||
# Try to guess the MIME type based on file content (using extension-based fallback)
|
||||
mime_type, encoding = mimetypes.guess_type(file_path)
|
||||
logger.info(f"Guessed MIME type for '{file_path}' is: {mime_type}")
|
||||
|
||||
endpoint = None
|
||||
form_key = "files" # Default form key for most endpoints
|
||||
|
||||
if mime_type:
|
||||
if mime_type == "text/html":
|
||||
endpoint = f"{gotenberg_url}/forms/chromium/convert/html"
|
||||
# The Chromium HTML endpoint expects the HTML file to be provided under the key "index.html"
|
||||
form_key = "index.html"
|
||||
elif mime_type.startswith("image/"):
|
||||
# For images, we use the LibreOffice endpoint (which supports image conversion)
|
||||
endpoint = f"{gotenberg_url}/forms/libreoffice/convert"
|
||||
elif mime_type.startswith("text/plain"):
|
||||
endpoint = f"{gotenberg_url}/forms/libreoffice/convert"
|
||||
elif mime_type in ["text/markdown", "text/x-markdown"]:
|
||||
# Optionally, you could use the Chromium markdown endpoint if you have an HTML wrapper.
|
||||
# For now, we'll fallback to LibreOffice.
|
||||
endpoint = f"{gotenberg_url}/forms/libreoffice/convert"
|
||||
else:
|
||||
# For all other MIME types (e.g. Office documents), use the LibreOffice endpoint.
|
||||
endpoint = f"{gotenberg_url}/forms/libreoffice/convert"
|
||||
else:
|
||||
# If MIME detection fails, fallback to extension-based detection.
|
||||
ext = os.path.splitext(file_path)[1].lower()
|
||||
if ext in [".html", ".htm"]:
|
||||
endpoint = f"{gotenberg_url}/forms/chromium/convert/html"
|
||||
form_key = "index.html"
|
||||
else:
|
||||
endpoint = f"{gotenberg_url}/forms/libreoffice/convert"
|
||||
|
||||
try:
|
||||
with open(file_path, "rb") as f:
|
||||
files = {form_key: f}
|
||||
response = requests.post(endpoint, files=files)
|
||||
|
||||
if response.status_code == 200:
|
||||
converted_file_path = os.path.splitext(file_path)[0] + ".pdf"
|
||||
with open(converted_file_path, "wb") as out_file:
|
||||
out_file.write(response.content)
|
||||
logger.info(f"Converted file saved as PDF: {converted_file_path}")
|
||||
upload_to_s3.delay(converted_file_path)
|
||||
return converted_file_path
|
||||
else:
|
||||
logger.error(f"Conversion failed for {file_path}. Status code: {response.status_code}")
|
||||
except Exception as e:
|
||||
logger.exception(f"Error converting {file_path} to PDF: {e}")
|
||||
|
||||
@@ -1,128 +1,128 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import fitz # PyMuPDF for PDF metadata editing
|
||||
import json
|
||||
from app.config import settings
|
||||
from app.tasks.retry_config import BaseTaskWithRetry
|
||||
from app.tasks.finalize_document_storage import finalize_document_storage
|
||||
|
||||
# Import the shared Celery instance
|
||||
from app.celery_app import celery
|
||||
|
||||
def unique_filepath(directory, base_filename, extension=".pdf"):
|
||||
"""
|
||||
Returns a unique filepath in the specified directory.
|
||||
If 'base_filename.pdf' exists, it will append an underscore and counter.
|
||||
"""
|
||||
candidate = os.path.join(directory, base_filename + extension)
|
||||
if not os.path.exists(candidate):
|
||||
return candidate
|
||||
counter = 1
|
||||
while True:
|
||||
candidate = os.path.join(directory, f"{base_filename}_{counter}{extension}")
|
||||
if not os.path.exists(candidate):
|
||||
return candidate
|
||||
counter += 1
|
||||
|
||||
def persist_metadata(metadata, final_pdf_path):
|
||||
"""
|
||||
Saves the metadata dictionary to a JSON file with the same base name as the final PDF.
|
||||
For example, if final_pdf_path is "<workdir>/processed/MyFile.pdf",
|
||||
the metadata will be saved as "<workdir>/processed/MyFile.json".
|
||||
"""
|
||||
base, _ = os.path.splitext(final_pdf_path)
|
||||
json_path = base + ".json"
|
||||
with open(json_path, "w", encoding="utf-8") as f:
|
||||
json.dump(metadata, f, ensure_ascii=False, indent=2)
|
||||
return json_path
|
||||
|
||||
@celery.task(base=BaseTaskWithRetry)
|
||||
def embed_metadata_into_pdf(local_file_path: str, extracted_text: str, metadata: dict):
|
||||
"""
|
||||
Embeds extracted metadata into the PDF's standard metadata fields.
|
||||
The mapping is as follows:
|
||||
- title: uses the extracted metadata "filename"
|
||||
- author: uses "absender" (or "Unknown" if missing)
|
||||
- subject: uses "document_type" (or "Unknown")
|
||||
- keywords: a comma‐separated list from the "tags" field
|
||||
|
||||
After processing, the file is moved to
|
||||
<workdir>/processed/<suggested_filename.pdf>
|
||||
where <suggested_filename.pdf> is derived from metadata["filename"].
|
||||
The output PDF is saved incrementally while preserving its original encryption.
|
||||
Additionally, the metadata is persisted to a JSON file with the same base name.
|
||||
"""
|
||||
# Check for file existence; if not found, try the known shared tmp directory.
|
||||
if not os.path.exists(local_file_path):
|
||||
alt_path = os.path.join(settings.workdir, "tmp", os.path.basename(local_file_path))
|
||||
if os.path.exists(alt_path):
|
||||
local_file_path = alt_path
|
||||
else:
|
||||
print(f"[ERROR] Local file {local_file_path} not found, cannot embed metadata.")
|
||||
return {"error": "File not found"}
|
||||
|
||||
# Work on a safe copy in /tmp
|
||||
tmp_dir = "/tmp"
|
||||
original_file = local_file_path
|
||||
processed_file = os.path.join(tmp_dir, f"processed_{os.path.basename(local_file_path)}")
|
||||
|
||||
# Create a safe copy to work on
|
||||
shutil.copy(original_file, processed_file)
|
||||
|
||||
try:
|
||||
print(f"[DEBUG] Embedding metadata into {processed_file}...")
|
||||
|
||||
# Open the PDF
|
||||
doc = fitz.open(processed_file)
|
||||
# Set PDF metadata using only the standard keys.
|
||||
doc.set_metadata({
|
||||
"title": metadata.get("filename", "Unknown Document"),
|
||||
"author": metadata.get("absender", "Unknown"),
|
||||
"subject": metadata.get("document_type", "Unknown"),
|
||||
"keywords": ", ".join(metadata.get("tags", []))
|
||||
})
|
||||
# Save incrementally and preserve encryption
|
||||
doc.save(processed_file, incremental=True, encryption=fitz.PDF_ENCRYPT_KEEP)
|
||||
doc.close()
|
||||
|
||||
print(f"[INFO] Metadata embedded successfully in {processed_file}")
|
||||
|
||||
# Use the suggested filename from metadata; if not provided, use the original basename.
|
||||
suggested_filename = metadata.get("filename", os.path.splitext(os.path.basename(local_file_path))[0])
|
||||
# Remove any extension and then add .pdf
|
||||
suggested_filename = os.path.splitext(suggested_filename)[0]
|
||||
# Define the final directory based on settings.workdir and ensure it exists.
|
||||
final_dir = os.path.join(settings.workdir, "processed")
|
||||
os.makedirs(final_dir, exist_ok=True)
|
||||
# Get a unique filepath in case of collisions.
|
||||
final_file_path = unique_filepath(final_dir, suggested_filename, extension=".pdf")
|
||||
|
||||
# Move the processed file using shutil.move to handle cross-device moves.
|
||||
shutil.move(processed_file, final_file_path)
|
||||
# Ensure the temporary file is deleted if it still exists.
|
||||
if os.path.exists(processed_file):
|
||||
os.remove(processed_file)
|
||||
|
||||
# Persist the metadata into a JSON file with the same base name.
|
||||
json_path = persist_metadata(metadata, final_file_path)
|
||||
print(f"[INFO] Metadata persisted to {json_path}")
|
||||
|
||||
# Trigger the next step: final storage.
|
||||
finalize_document_storage.delay(original_file, final_file_path, metadata)
|
||||
|
||||
# After triggering final storage, delete the original file if it is in workdir/tmp.
|
||||
workdir_tmp = os.path.join(settings.workdir, "tmp")
|
||||
if original_file.startswith(workdir_tmp) and os.path.exists(original_file):
|
||||
try:
|
||||
os.remove(original_file)
|
||||
print(f"[INFO] Deleted original file from {original_file}")
|
||||
except Exception as e:
|
||||
print(f"[ERROR] Could not delete original file {original_file}: {e}")
|
||||
|
||||
return {"file": final_file_path, "metadata_file": json_path, "status": "Metadata embedded"}
|
||||
|
||||
except Exception as e:
|
||||
print(f"[ERROR] Failed to embed metadata into {processed_file}: {e}")
|
||||
return {"error": str(e)}
|
||||
#!/usr/bin/env python3
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import fitz # PyMuPDF for PDF metadata editing
|
||||
import json
|
||||
from app.config import settings
|
||||
from app.tasks.retry_config import BaseTaskWithRetry
|
||||
from app.tasks.finalize_document_storage import finalize_document_storage
|
||||
|
||||
# Import the shared Celery instance
|
||||
from app.celery_app import celery
|
||||
|
||||
def unique_filepath(directory, base_filename, extension=".pdf"):
|
||||
"""
|
||||
Returns a unique filepath in the specified directory.
|
||||
If 'base_filename.pdf' exists, it will append an underscore and counter.
|
||||
"""
|
||||
candidate = os.path.join(directory, base_filename + extension)
|
||||
if not os.path.exists(candidate):
|
||||
return candidate
|
||||
counter = 1
|
||||
while True:
|
||||
candidate = os.path.join(directory, f"{base_filename}_{counter}{extension}")
|
||||
if not os.path.exists(candidate):
|
||||
return candidate
|
||||
counter += 1
|
||||
|
||||
def persist_metadata(metadata, final_pdf_path):
|
||||
"""
|
||||
Saves the metadata dictionary to a JSON file with the same base name as the final PDF.
|
||||
For example, if final_pdf_path is "<workdir>/processed/MyFile.pdf",
|
||||
the metadata will be saved as "<workdir>/processed/MyFile.json".
|
||||
"""
|
||||
base, _ = os.path.splitext(final_pdf_path)
|
||||
json_path = base + ".json"
|
||||
with open(json_path, "w", encoding="utf-8") as f:
|
||||
json.dump(metadata, f, ensure_ascii=False, indent=2)
|
||||
return json_path
|
||||
|
||||
@celery.task(base=BaseTaskWithRetry)
|
||||
def embed_metadata_into_pdf(local_file_path: str, extracted_text: str, metadata: dict):
|
||||
"""
|
||||
Embeds extracted metadata into the PDF's standard metadata fields.
|
||||
The mapping is as follows:
|
||||
- title: uses the extracted metadata "filename"
|
||||
- author: uses "absender" (or "Unknown" if missing)
|
||||
- subject: uses "document_type" (or "Unknown")
|
||||
- keywords: a comma‐separated list from the "tags" field
|
||||
|
||||
After processing, the file is moved to
|
||||
<workdir>/processed/<suggested_filename.pdf>
|
||||
where <suggested_filename.pdf> is derived from metadata["filename"].
|
||||
The output PDF is saved incrementally while preserving its original encryption.
|
||||
Additionally, the metadata is persisted to a JSON file with the same base name.
|
||||
"""
|
||||
# Check for file existence; if not found, try the known shared tmp directory.
|
||||
if not os.path.exists(local_file_path):
|
||||
alt_path = os.path.join(settings.workdir, "tmp", os.path.basename(local_file_path))
|
||||
if os.path.exists(alt_path):
|
||||
local_file_path = alt_path
|
||||
else:
|
||||
print(f"[ERROR] Local file {local_file_path} not found, cannot embed metadata.")
|
||||
return {"error": "File not found"}
|
||||
|
||||
# Work on a safe copy in /tmp
|
||||
tmp_dir = "/tmp"
|
||||
original_file = local_file_path
|
||||
processed_file = os.path.join(tmp_dir, f"processed_{os.path.basename(local_file_path)}")
|
||||
|
||||
# Create a safe copy to work on
|
||||
shutil.copy(original_file, processed_file)
|
||||
|
||||
try:
|
||||
print(f"[DEBUG] Embedding metadata into {processed_file}...")
|
||||
|
||||
# Open the PDF
|
||||
doc = fitz.open(processed_file)
|
||||
# Set PDF metadata using only the standard keys.
|
||||
doc.set_metadata({
|
||||
"title": metadata.get("filename", "Unknown Document"),
|
||||
"author": metadata.get("absender", "Unknown"),
|
||||
"subject": metadata.get("document_type", "Unknown"),
|
||||
"keywords": ", ".join(metadata.get("tags", []))
|
||||
})
|
||||
# Save incrementally and preserve encryption
|
||||
doc.save(processed_file, incremental=True, encryption=fitz.PDF_ENCRYPT_KEEP)
|
||||
doc.close()
|
||||
|
||||
print(f"[INFO] Metadata embedded successfully in {processed_file}")
|
||||
|
||||
# Use the suggested filename from metadata; if not provided, use the original basename.
|
||||
suggested_filename = metadata.get("filename", os.path.splitext(os.path.basename(local_file_path))[0])
|
||||
# Remove any extension and then add .pdf
|
||||
suggested_filename = os.path.splitext(suggested_filename)[0]
|
||||
# Define the final directory based on settings.workdir and ensure it exists.
|
||||
final_dir = os.path.join(settings.workdir, "processed")
|
||||
os.makedirs(final_dir, exist_ok=True)
|
||||
# Get a unique filepath in case of collisions.
|
||||
final_file_path = unique_filepath(final_dir, suggested_filename, extension=".pdf")
|
||||
|
||||
# Move the processed file using shutil.move to handle cross-device moves.
|
||||
shutil.move(processed_file, final_file_path)
|
||||
# Ensure the temporary file is deleted if it still exists.
|
||||
if os.path.exists(processed_file):
|
||||
os.remove(processed_file)
|
||||
|
||||
# Persist the metadata into a JSON file with the same base name.
|
||||
json_path = persist_metadata(metadata, final_file_path)
|
||||
print(f"[INFO] Metadata persisted to {json_path}")
|
||||
|
||||
# Trigger the next step: final storage.
|
||||
finalize_document_storage.delay(original_file, final_file_path, metadata)
|
||||
|
||||
# After triggering final storage, delete the original file if it is in workdir/tmp.
|
||||
workdir_tmp = os.path.join(settings.workdir, "tmp")
|
||||
if original_file.startswith(workdir_tmp) and os.path.exists(original_file):
|
||||
try:
|
||||
os.remove(original_file)
|
||||
print(f"[INFO] Deleted original file from {original_file}")
|
||||
except Exception as e:
|
||||
print(f"[ERROR] Could not delete original file {original_file}: {e}")
|
||||
|
||||
return {"file": final_file_path, "metadata_file": json_path, "status": "Metadata embedded"}
|
||||
|
||||
except Exception as e:
|
||||
print(f"[ERROR] Failed to embed metadata into {processed_file}: {e}")
|
||||
return {"error": str(e)}
|
||||
|
||||
@@ -1,97 +1,97 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
import json
|
||||
import re
|
||||
from openai import OpenAI
|
||||
from app.config import settings
|
||||
from app.tasks.retry_config import BaseTaskWithRetry
|
||||
from app.tasks.embed_metadata_into_pdf import embed_metadata_into_pdf
|
||||
|
||||
# Import the shared Celery instance
|
||||
from app.celery_app import celery
|
||||
|
||||
client = OpenAI(api_key=settings.openai_api_key)
|
||||
|
||||
def extract_json_from_text(text):
|
||||
"""
|
||||
Try to extract a JSON object from the text.
|
||||
- First, check for a JSON block inside triple backticks.
|
||||
- If not found, try to extract text from the first '{' to the last '}'.
|
||||
"""
|
||||
pattern = r"```(?:json)?\s*(\{.*?\})\s*```"
|
||||
match = re.search(pattern, text, re.DOTALL)
|
||||
if match:
|
||||
return match.group(1)
|
||||
else:
|
||||
start = text.find("{")
|
||||
end = text.rfind("}")
|
||||
if start != -1 and end != -1 and end > start:
|
||||
return text[start:end+1]
|
||||
return None
|
||||
|
||||
@celery.task(base=BaseTaskWithRetry)
|
||||
def extract_metadata_with_gpt(s3_filename: str, cleaned_text: str):
|
||||
"""Uses OpenAI GPT-4o-mini to classify document metadata."""
|
||||
|
||||
prompt = f"""
|
||||
You are a specialized document analyzer trained to extract structured metadata from documents.
|
||||
Your task is to analyze the given text and return a well-structured JSON object.
|
||||
|
||||
Extract and return the following fields:
|
||||
1. **filename**: Machine-readable filename (YYYY-MM-DD_DescriptiveTitle, use only letters, numbers, periods, and underscores).
|
||||
2. **empfaenger**: The recipient, or "Unknown" if not found.
|
||||
3. **absender**: The sender, or "Unknown" if not found.
|
||||
4. **correspondent**: The entity or company that issued the document (shortest possible name, e.g., "Amazon" instead of "Amazon EU SARL, German branch").
|
||||
5. **kommunikationsart**: One of [Behoerdlicher_Brief, Rechnung, Kontoauszug, Vertrag, Quittung, Privater_Brief, Einladung, Gewerbliche_Korrespondenz, Newsletter, Werbung, Sonstiges].
|
||||
6. **kommunikationskategorie**: One of [Amtliche_Postbehoerdliche_Dokumente, Finanz_und_Vertragsdokumente, Geschaeftliche_Kommunikation, Private_Korrespondenz, Sonstige_Informationen].
|
||||
7. **document_type**: Precise classification (e.g., Invoice, Contract, Information, Unknown).
|
||||
8. **tags**: A list of up to 4 relevant thematic keywords.
|
||||
9. **language**: Detected document language (ISO 639-1 code, e.g., "de" or "en").
|
||||
10. **title**: A human-readable title summarizing the document content.
|
||||
11. **confidence_score**: A numeric value (0-100) indicating the confidence level of the extracted metadata.
|
||||
12. **reference_number**: Extracted invoice/order/reference number if available.
|
||||
13. **monetary_amounts**: A list of key monetary values detected in the document.
|
||||
|
||||
### Important Rules:
|
||||
- **OCR Correction**: Assume the text has been corrected for OCR errors.
|
||||
- **Tagging**: Max 4 tags, avoiding generic or overly specific terms.
|
||||
- **Title**: Concise, no addresses, and contains key identifying features.
|
||||
- **Date Selection**: Use the most relevant date if multiple are found.
|
||||
- **Output Language**: Maintain the document's original language.
|
||||
|
||||
Extracted text:
|
||||
{cleaned_text}
|
||||
|
||||
Return only valid JSON with no additional commentary.
|
||||
"""
|
||||
|
||||
try:
|
||||
print(f"[DEBUG] Sending classification request for {s3_filename}...")
|
||||
completion = client.chat.completions.create(
|
||||
model="gpt-4o-mini",
|
||||
messages=[
|
||||
{"role": "system", "content": "You are an intelligent document classifier."},
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
temperature=0
|
||||
)
|
||||
|
||||
content = completion.choices[0].message.content
|
||||
print(f"[DEBUG] Raw classification response for {s3_filename}: {content}")
|
||||
|
||||
json_text = extract_json_from_text(content)
|
||||
if not json_text:
|
||||
print(f"[ERROR] Could not find valid JSON in GPT response for {s3_filename}.")
|
||||
return {}
|
||||
|
||||
metadata = json.loads(json_text)
|
||||
print(f"[DEBUG] Extracted metadata: {metadata}")
|
||||
|
||||
# Trigger the next step: embedding metadata into the PDF
|
||||
embed_metadata_into_pdf.delay(s3_filename, cleaned_text, metadata)
|
||||
|
||||
return {"s3_file": s3_filename, "metadata": metadata}
|
||||
|
||||
except Exception as e:
|
||||
print(f"[ERROR] OpenAI classification failed for {s3_filename}: {e}")
|
||||
return {}
|
||||
#!/usr/bin/env python3
|
||||
|
||||
import json
|
||||
import re
|
||||
from openai import OpenAI
|
||||
from app.config import settings
|
||||
from app.tasks.retry_config import BaseTaskWithRetry
|
||||
from app.tasks.embed_metadata_into_pdf import embed_metadata_into_pdf
|
||||
|
||||
# Import the shared Celery instance
|
||||
from app.celery_app import celery
|
||||
|
||||
client = OpenAI(api_key=settings.openai_api_key)
|
||||
|
||||
def extract_json_from_text(text):
|
||||
"""
|
||||
Try to extract a JSON object from the text.
|
||||
- First, check for a JSON block inside triple backticks.
|
||||
- If not found, try to extract text from the first '{' to the last '}'.
|
||||
"""
|
||||
pattern = r"```(?:json)?\s*(\{.*?\})\s*```"
|
||||
match = re.search(pattern, text, re.DOTALL)
|
||||
if match:
|
||||
return match.group(1)
|
||||
else:
|
||||
start = text.find("{")
|
||||
end = text.rfind("}")
|
||||
if start != -1 and end != -1 and end > start:
|
||||
return text[start:end+1]
|
||||
return None
|
||||
|
||||
@celery.task(base=BaseTaskWithRetry)
|
||||
def extract_metadata_with_gpt(s3_filename: str, cleaned_text: str):
|
||||
"""Uses OpenAI GPT-4o-mini to classify document metadata."""
|
||||
|
||||
prompt = f"""
|
||||
You are a specialized document analyzer trained to extract structured metadata from documents.
|
||||
Your task is to analyze the given text and return a well-structured JSON object.
|
||||
|
||||
Extract and return the following fields:
|
||||
1. **filename**: Machine-readable filename (YYYY-MM-DD_DescriptiveTitle, use only letters, numbers, periods, and underscores).
|
||||
2. **empfaenger**: The recipient, or "Unknown" if not found.
|
||||
3. **absender**: The sender, or "Unknown" if not found.
|
||||
4. **correspondent**: The entity or company that issued the document (shortest possible name, e.g., "Amazon" instead of "Amazon EU SARL, German branch").
|
||||
5. **kommunikationsart**: One of [Behoerdlicher_Brief, Rechnung, Kontoauszug, Vertrag, Quittung, Privater_Brief, Einladung, Gewerbliche_Korrespondenz, Newsletter, Werbung, Sonstiges].
|
||||
6. **kommunikationskategorie**: One of [Amtliche_Postbehoerdliche_Dokumente, Finanz_und_Vertragsdokumente, Geschaeftliche_Kommunikation, Private_Korrespondenz, Sonstige_Informationen].
|
||||
7. **document_type**: Precise classification (e.g., Invoice, Contract, Information, Unknown).
|
||||
8. **tags**: A list of up to 4 relevant thematic keywords.
|
||||
9. **language**: Detected document language (ISO 639-1 code, e.g., "de" or "en").
|
||||
10. **title**: A human-readable title summarizing the document content.
|
||||
11. **confidence_score**: A numeric value (0-100) indicating the confidence level of the extracted metadata.
|
||||
12. **reference_number**: Extracted invoice/order/reference number if available.
|
||||
13. **monetary_amounts**: A list of key monetary values detected in the document.
|
||||
|
||||
### Important Rules:
|
||||
- **OCR Correction**: Assume the text has been corrected for OCR errors.
|
||||
- **Tagging**: Max 4 tags, avoiding generic or overly specific terms.
|
||||
- **Title**: Concise, no addresses, and contains key identifying features.
|
||||
- **Date Selection**: Use the most relevant date if multiple are found.
|
||||
- **Output Language**: Maintain the document's original language.
|
||||
|
||||
Extracted text:
|
||||
{cleaned_text}
|
||||
|
||||
Return only valid JSON with no additional commentary.
|
||||
"""
|
||||
|
||||
try:
|
||||
print(f"[DEBUG] Sending classification request for {s3_filename}...")
|
||||
completion = client.chat.completions.create(
|
||||
model="gpt-4o-mini",
|
||||
messages=[
|
||||
{"role": "system", "content": "You are an intelligent document classifier."},
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
temperature=0
|
||||
)
|
||||
|
||||
content = completion.choices[0].message.content
|
||||
print(f"[DEBUG] Raw classification response for {s3_filename}: {content}")
|
||||
|
||||
json_text = extract_json_from_text(content)
|
||||
if not json_text:
|
||||
print(f"[ERROR] Could not find valid JSON in GPT response for {s3_filename}.")
|
||||
return {}
|
||||
|
||||
metadata = json.loads(json_text)
|
||||
print(f"[DEBUG] Extracted metadata: {metadata}")
|
||||
|
||||
# Trigger the next step: embedding metadata into the PDF
|
||||
embed_metadata_into_pdf.delay(s3_filename, cleaned_text, metadata)
|
||||
|
||||
return {"s3_file": s3_filename, "metadata": metadata}
|
||||
|
||||
except Exception as e:
|
||||
print(f"[ERROR] OpenAI classification failed for {s3_filename}: {e}")
|
||||
return {}
|
||||
|
||||
@@ -1,26 +1,26 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
from app.config import settings
|
||||
from app.tasks.retry_config import BaseTaskWithRetry
|
||||
# Import the shared Celery instance
|
||||
from app.celery_app import celery
|
||||
|
||||
# 1) Import the aggregator task
|
||||
from app.tasks.send_to_all import send_to_all_destinations
|
||||
|
||||
|
||||
@celery.task(base=BaseTaskWithRetry)
|
||||
def finalize_document_storage(original_file: str, processed_file: str, metadata: dict):
|
||||
"""
|
||||
Final storage step after embedding metadata.
|
||||
We will now call 'send_to_all_destinations' to push the final PDF to Dropbox/Nextcloud/Paperless.
|
||||
"""
|
||||
print(f"[INFO] Finalizing document storage for {processed_file}")
|
||||
|
||||
# 2) Enqueue uploads to all destinations (Dropbox, Nextcloud, Paperless)
|
||||
send_to_all_destinations.delay(processed_file)
|
||||
|
||||
return {
|
||||
"status": "Completed",
|
||||
"file": processed_file
|
||||
}
|
||||
#!/usr/bin/env python3
|
||||
|
||||
from app.config import settings
|
||||
from app.tasks.retry_config import BaseTaskWithRetry
|
||||
# Import the shared Celery instance
|
||||
from app.celery_app import celery
|
||||
|
||||
# 1) Import the aggregator task
|
||||
from app.tasks.send_to_all import send_to_all_destinations
|
||||
|
||||
|
||||
@celery.task(base=BaseTaskWithRetry)
|
||||
def finalize_document_storage(original_file: str, processed_file: str, metadata: dict):
|
||||
"""
|
||||
Final storage step after embedding metadata.
|
||||
We will now call 'send_to_all_destinations' to push the final PDF to Dropbox/Nextcloud/Paperless.
|
||||
"""
|
||||
print(f"[INFO] Finalizing document storage for {processed_file}")
|
||||
|
||||
# 2) Enqueue uploads to all destinations (Dropbox, Nextcloud, Paperless)
|
||||
send_to_all_destinations.delay(processed_file)
|
||||
|
||||
return {
|
||||
"status": "Completed",
|
||||
"file": processed_file
|
||||
}
|
||||
|
||||
+412
-412
@@ -1,412 +1,412 @@
|
||||
#!/usr/bin/env python3
|
||||
import os
|
||||
import json
|
||||
import email
|
||||
import imaplib
|
||||
import logging
|
||||
import redis
|
||||
import re
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from celery import shared_task
|
||||
from app.config import settings
|
||||
from app.tasks.upload_to_s3 import upload_to_s3
|
||||
from app.tasks.convert_to_pdf import convert_to_pdf # new conversion task
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Initialize Redis connection using Celery's Redis settings
|
||||
redis_client = redis.StrictRedis.from_url(settings.redis_url, decode_responses=True)
|
||||
|
||||
LOCK_KEY = "imap_lock" # Unique key for locking
|
||||
LOCK_EXPIRE = 300 # Lock expires in 5 minutes
|
||||
|
||||
# Local cache file for tracking processed emails
|
||||
CACHE_FILE = os.path.join(settings.workdir, "processed_mails.json")
|
||||
|
||||
|
||||
def acquire_lock():
|
||||
"""Attempt to acquire a Redis-based lock. If acquired, set an expiration."""
|
||||
lock_acquired = redis_client.setnx(LOCK_KEY, "locked")
|
||||
if lock_acquired:
|
||||
redis_client.expire(LOCK_KEY, LOCK_EXPIRE)
|
||||
logger.info("Lock acquired for IMAP processing.")
|
||||
return True
|
||||
logger.warning("Lock already held. Skipping this cycle.")
|
||||
return False
|
||||
|
||||
|
||||
def release_lock():
|
||||
"""Release the lock by deleting the Redis key."""
|
||||
redis_client.delete(LOCK_KEY)
|
||||
logger.info("Lock released.")
|
||||
|
||||
|
||||
def load_processed_emails():
|
||||
"""Load the list of already processed emails from a local JSON file."""
|
||||
if os.path.exists(CACHE_FILE):
|
||||
try:
|
||||
with open(CACHE_FILE, "r") as f:
|
||||
processed_emails = json.load(f)
|
||||
processed_emails = cleanup_old_entries(processed_emails)
|
||||
return processed_emails
|
||||
except json.JSONDecodeError:
|
||||
logger.warning("Failed to decode JSON, resetting processed emails cache.")
|
||||
return {}
|
||||
return {}
|
||||
|
||||
|
||||
def save_processed_emails(processed_emails):
|
||||
"""Save the processed email IDs to a local JSON file."""
|
||||
with open(CACHE_FILE, "w") as f:
|
||||
json.dump(processed_emails, f, indent=4)
|
||||
|
||||
|
||||
def cleanup_old_entries(processed_emails):
|
||||
"""Remove entries older than 7 days from the cache to avoid infinite growth."""
|
||||
seven_days_ago = datetime.now(timezone.utc) - timedelta(days=7)
|
||||
valid_emails = {}
|
||||
for msg_id, date_str in processed_emails.items():
|
||||
naive_dt = datetime.strptime(date_str, "%Y-%m-%dT%H:%M:%S")
|
||||
aware_dt = naive_dt.replace(tzinfo=timezone.utc)
|
||||
if aware_dt > seven_days_ago:
|
||||
valid_emails[msg_id] = date_str
|
||||
return valid_emails
|
||||
|
||||
|
||||
@shared_task
|
||||
def pull_all_inboxes():
|
||||
"""
|
||||
Periodic Celery task that checks all configured IMAP mailboxes
|
||||
and fetches attachments from new emails.
|
||||
Ensures only one instance runs at a time using Redis-based locking.
|
||||
"""
|
||||
if not acquire_lock():
|
||||
logger.info("Skipping execution: Another instance is running.")
|
||||
return
|
||||
|
||||
try:
|
||||
logger.info("Starting pull_all_inboxes")
|
||||
|
||||
# Mailbox #1 (non-Gmail)
|
||||
check_and_pull_mailbox(
|
||||
mailbox_key="imap1",
|
||||
host=settings.imap1_host,
|
||||
port=settings.imap1_port,
|
||||
username=settings.imap1_username,
|
||||
password=settings.imap1_password,
|
||||
use_ssl=settings.imap1_ssl,
|
||||
delete_after_process=settings.imap1_delete_after_process,
|
||||
)
|
||||
|
||||
# Mailbox #2 (Gmail)
|
||||
check_and_pull_mailbox(
|
||||
mailbox_key="imap2",
|
||||
host=settings.imap2_host,
|
||||
port=settings.imap2_port,
|
||||
username=settings.imap2_username,
|
||||
password=settings.imap2_password,
|
||||
use_ssl=settings.imap2_ssl,
|
||||
delete_after_process=settings.imap2_delete_after_process,
|
||||
)
|
||||
|
||||
logger.info("Finished pull_all_inboxes")
|
||||
|
||||
finally:
|
||||
release_lock()
|
||||
|
||||
|
||||
def check_and_pull_mailbox(
|
||||
mailbox_key: str,
|
||||
host: str | None,
|
||||
port: int | None,
|
||||
username: str | None,
|
||||
password: str | None,
|
||||
use_ssl: bool,
|
||||
delete_after_process: bool,
|
||||
):
|
||||
"""Validates config and invokes pulling from the mailbox if valid."""
|
||||
if not (host and port and username and password):
|
||||
logger.warning(f"Mailbox {mailbox_key} is missing config, skipping.")
|
||||
return
|
||||
|
||||
logger.info(f"Checking mailbox: {mailbox_key}")
|
||||
pull_inbox(
|
||||
mailbox_key=mailbox_key,
|
||||
host=host,
|
||||
port=port,
|
||||
username=username,
|
||||
password=password,
|
||||
use_ssl=use_ssl,
|
||||
delete_after_process=delete_after_process,
|
||||
)
|
||||
|
||||
|
||||
def pull_inbox(mailbox_key, host, port, username, password, use_ssl,
|
||||
delete_after_process):
|
||||
"""
|
||||
Connects to the IMAP inbox, fetches new unread emails from the last 3 days,
|
||||
and processes attachments while preserving the original unread status.
|
||||
|
||||
For Gmail:
|
||||
- Attempts to select the localized All Mail folder.
|
||||
- Runs an X-GM-RAW query: "in:anywhere in:unread newer_than:3d has:attachment".
|
||||
|
||||
For non-Gmail mailboxes, it falls back to selecting the INBOX with a SINCE/UNSEEN filter.
|
||||
"""
|
||||
logger.info("Connecting to %s at %s:%s (SSL=%s)",
|
||||
mailbox_key, host, port, use_ssl)
|
||||
processed_emails = load_processed_emails()
|
||||
|
||||
try:
|
||||
mail = imaplib.IMAP4_SSL(host, port) if use_ssl else imaplib.IMAP4(host, port)
|
||||
mail.login(username, password)
|
||||
|
||||
is_gmail_host = "gmail" in host.lower()
|
||||
if is_gmail_host:
|
||||
# For Gmail, try to select the localized All Mail folder.
|
||||
all_mail_folder = find_all_mail_folder(mail)
|
||||
if all_mail_folder:
|
||||
logger.info("Using Gmail All Mail folder: %s", all_mail_folder)
|
||||
mail.select(f'"{all_mail_folder}"')
|
||||
else:
|
||||
logger.warning("Gmail All Mail folder not found, falling back to INBOX.")
|
||||
mail.select("INBOX")
|
||||
# Use the X-GM-RAW query for Gmail.
|
||||
raw_query = "in:anywhere in:unread newer_than:3d has:attachment"
|
||||
status, search_data = mail.search(None, "X-GM-RAW", f'"{raw_query}"')
|
||||
else:
|
||||
# For non-Gmail, select INBOX and use SINCE/UNSEEN query.
|
||||
mail.select("INBOX")
|
||||
since_date = (datetime.now(timezone.utc) - timedelta(days=3)
|
||||
).strftime("%d-%b-%Y")
|
||||
status, search_data = mail.search(None, f'(SINCE {since_date} UNSEEN)')
|
||||
|
||||
if status != "OK":
|
||||
logger.warning("Search failed on mailbox %s. Status=%s",
|
||||
mailbox_key, status)
|
||||
mail.close()
|
||||
mail.logout()
|
||||
return
|
||||
|
||||
msg_numbers = search_data[0].split()
|
||||
logger.info("Found %d unread emails in %s.", len(msg_numbers), mailbox_key)
|
||||
|
||||
for num in msg_numbers:
|
||||
status, msg_data = mail.fetch(num, "(RFC822)")
|
||||
if status != "OK":
|
||||
logger.warning("Failed to fetch message %s in %s. Status=%s",
|
||||
num, mailbox_key, status)
|
||||
continue
|
||||
|
||||
raw_email = msg_data[0][1]
|
||||
email_message = email.message_from_bytes(raw_email)
|
||||
msg_id = email_message.get("Message-ID")
|
||||
|
||||
if not msg_id:
|
||||
logger.warning("Skipping email without Message-ID in %s", mailbox_key)
|
||||
continue
|
||||
|
||||
if msg_id in processed_emails:
|
||||
logger.info("Skipping already processed email %s in %s", msg_id, mailbox_key)
|
||||
continue
|
||||
|
||||
# For Gmail, check if the email already has the "Ingested" label.
|
||||
if is_gmail_host:
|
||||
if email_already_has_label(mail, num, "Ingested"):
|
||||
logger.info("Skipping email %s in %s, already labeled 'Ingested'.",
|
||||
msg_id, mailbox_key)
|
||||
continue
|
||||
|
||||
# Process attachments (and convert non-PDF files).
|
||||
# We call the function without assigning its return value since it is not used.
|
||||
fetch_attachments_and_enqueue(email_message)
|
||||
|
||||
if is_gmail_host:
|
||||
mark_as_processed_with_star(mail, num)
|
||||
mark_as_processed_with_label(mail, num, label="Ingested")
|
||||
|
||||
processed_emails[msg_id] = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S")
|
||||
save_processed_emails(processed_emails)
|
||||
|
||||
if delete_after_process:
|
||||
logger.info("Deleting message %s from %s", num.decode(), mailbox_key)
|
||||
mail.store(num, "+FLAGS", "\\Deleted")
|
||||
else:
|
||||
mail.store(num, "-FLAGS", "\\Seen")
|
||||
|
||||
if delete_after_process:
|
||||
mail.expunge()
|
||||
|
||||
mail.close()
|
||||
mail.logout()
|
||||
logger.info("Finished processing mailbox %s", mailbox_key)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception("Error pulling mailbox %s: %s", mailbox_key, e)
|
||||
|
||||
|
||||
def fetch_attachments_and_enqueue(email_message):
|
||||
"""
|
||||
Extracts attachments from the email and processes only allowed file types.
|
||||
|
||||
Allowed file types include:
|
||||
- PDF: application/pdf
|
||||
- Microsoft Office files:
|
||||
- Word: application/msword,
|
||||
application/vnd.openxmlformats-officedocument.wordprocessingml.document
|
||||
- Excel: application/vnd.ms-excel,
|
||||
application/vnd.openxmlformats-officedocument.spreadsheetml.sheet
|
||||
- PowerPoint: application/vnd.ms-powerpoint,
|
||||
application/vnd.openxmlformats-officedocument.presentationml.presentation
|
||||
- Other meaningful attachments:
|
||||
- Plain text: text/plain
|
||||
- CSV: text/csv
|
||||
- Rich Text Format: application/rtf, text/rtf
|
||||
|
||||
Attachments not in this list are skipped. Common image MIME types such as
|
||||
image/jpeg, image/png, image/gif, image/bmp, image/tiff, and image/webp are
|
||||
intentionally excluded.
|
||||
|
||||
If the attachment is a PDF, it is enqueued for upload; any other allowed file
|
||||
is enqueued for conversion to PDF.
|
||||
|
||||
Returns True if at least one allowed attachment was processed.
|
||||
"""
|
||||
ALLOWED_MIME_TYPES = {
|
||||
"application/pdf",
|
||||
"application/msword",
|
||||
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||||
"application/vnd.ms-excel",
|
||||
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
||||
"application/vnd.ms-powerpoint",
|
||||
"application/vnd.openxmlformats-officedocument.presentationml.presentation",
|
||||
"text/plain",
|
||||
"text/csv",
|
||||
"application/rtf",
|
||||
"text/rtf",
|
||||
}
|
||||
|
||||
has_attachment = False
|
||||
for part in email_message.walk():
|
||||
if part.get_content_maintype() == "multipart":
|
||||
continue
|
||||
|
||||
filename = part.get_filename()
|
||||
if not filename:
|
||||
continue
|
||||
|
||||
mime_type = part.get_content_type()
|
||||
if mime_type not in ALLOWED_MIME_TYPES:
|
||||
logger.info("Skipping attachment %s with MIME type %s",
|
||||
filename, mime_type)
|
||||
continue
|
||||
|
||||
file_path = os.path.join(settings.workdir, filename)
|
||||
with open(file_path, "wb") as f:
|
||||
f.write(part.get_payload(decode=True))
|
||||
|
||||
if mime_type == "application/pdf":
|
||||
upload_to_s3.delay(file_path)
|
||||
logger.info("Enqueued PDF for upload: %s", filename)
|
||||
elif mime_type in ALLOWED_MIME_TYPES:
|
||||
# Enqueue conversion to PDF using the Gotenberg service.
|
||||
convert_to_pdf.delay(file_path)
|
||||
logger.info("Enqueued file for conversion to PDF: %s", filename)
|
||||
|
||||
has_attachment = True
|
||||
return has_attachment
|
||||
|
||||
|
||||
def email_already_has_label(mail, msg_id, label="Ingested"):
|
||||
"""
|
||||
Checks if the given message (msg_id) has the specified Gmail label.
|
||||
Returns True if the label is found, False otherwise.
|
||||
"""
|
||||
try:
|
||||
label_status, label_data = mail.fetch(msg_id, "(X-GM-LABELS)")
|
||||
if label_status == "OK" and label_data and len(label_data) > 0:
|
||||
raw_labels = label_data[0][1].decode("utf-8", errors="ignore")
|
||||
if label in raw_labels:
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error("Failed to fetch labels for msg_id=%s: %s", msg_id, e)
|
||||
return False
|
||||
|
||||
|
||||
def mark_as_processed_with_star(mail, msg_id):
|
||||
"""Stars the email in Gmail."""
|
||||
try:
|
||||
mail.store(msg_id, "+FLAGS", "\\Flagged")
|
||||
logger.info("Email %s starred in Gmail.", msg_id)
|
||||
except Exception as e:
|
||||
logger.error("Failed to star email %s: %s", msg_id, e)
|
||||
|
||||
|
||||
def mark_as_processed_with_label(mail, msg_id, label="Ingested"):
|
||||
"""Adds a custom label to the email in Gmail."""
|
||||
try:
|
||||
mail.store(msg_id, "+X-GM-LABELS", label)
|
||||
logger.info("Email %s labeled '%s' in Gmail.", msg_id, label)
|
||||
except Exception as e:
|
||||
logger.error("Failed to label email %s with %s: %s", msg_id, label, e)
|
||||
|
||||
|
||||
def find_all_mail_folder(mail):
|
||||
"""
|
||||
Attempts to select the Gmail All Mail folder using known localized names.
|
||||
Falls back to using XLIST if needed.
|
||||
Returns the folder name if found, otherwise None.
|
||||
"""
|
||||
COMMON_ALL_MAIL_NAMES = [
|
||||
"[Gmail]/Alle Nachrichten",
|
||||
"[Gmail]/All Mail",
|
||||
"[Gmail]/Todos",
|
||||
"[Gmail]/Tutte le mail",
|
||||
"[Gmail]/Tous les messages",
|
||||
]
|
||||
for candidate in COMMON_ALL_MAIL_NAMES:
|
||||
status, _ = mail.select(f'"{candidate}"', readonly=True)
|
||||
if status == "OK":
|
||||
return candidate
|
||||
|
||||
capabilities = get_capabilities(mail)
|
||||
if "XLIST" in capabilities:
|
||||
candidate = find_all_mail_xlist(mail)
|
||||
if candidate:
|
||||
return candidate
|
||||
return None
|
||||
|
||||
|
||||
def get_capabilities(mail):
|
||||
"""Returns a list of capabilities supported by the IMAP server."""
|
||||
typ, data = mail.capability()
|
||||
if typ == "OK" and data:
|
||||
caps = data[0].decode("utf-8", errors="ignore").upper().split()
|
||||
return caps
|
||||
return []
|
||||
|
||||
|
||||
def find_all_mail_xlist(mail):
|
||||
"""
|
||||
Uses XLIST to discover the mailbox flagged as All Mail.
|
||||
Returns the folder name if found, otherwise None.
|
||||
"""
|
||||
tag = mail._new_tag().decode("ascii")
|
||||
command_str = f"{tag} XLIST \"\" \"*\""
|
||||
mail.send((command_str + "\r\n").encode("utf-8"))
|
||||
|
||||
all_mail_folder = None
|
||||
while True:
|
||||
line = mail.readline()
|
||||
if not line:
|
||||
break
|
||||
line_str = line.decode("utf-8", errors="ignore").strip()
|
||||
if line_str.upper().startswith("* XLIST ") and "\\ALLMAIL" in line_str.upper():
|
||||
match = re.search(r'"([^"]+)"$', line_str)
|
||||
if match:
|
||||
candidate = match.group(1)
|
||||
logger.info("Found All Mail folder via XLIST: %s", candidate)
|
||||
all_mail_folder = candidate
|
||||
if line_str.startswith(tag):
|
||||
break
|
||||
return all_mail_folder
|
||||
#!/usr/bin/env python3
|
||||
import os
|
||||
import json
|
||||
import email
|
||||
import imaplib
|
||||
import logging
|
||||
import redis
|
||||
import re
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from celery import shared_task
|
||||
from app.config import settings
|
||||
from app.tasks.upload_to_s3 import upload_to_s3
|
||||
from app.tasks.convert_to_pdf import convert_to_pdf # new conversion task
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Initialize Redis connection using Celery's Redis settings
|
||||
redis_client = redis.StrictRedis.from_url(settings.redis_url, decode_responses=True)
|
||||
|
||||
LOCK_KEY = "imap_lock" # Unique key for locking
|
||||
LOCK_EXPIRE = 300 # Lock expires in 5 minutes
|
||||
|
||||
# Local cache file for tracking processed emails
|
||||
CACHE_FILE = os.path.join(settings.workdir, "processed_mails.json")
|
||||
|
||||
|
||||
def acquire_lock():
|
||||
"""Attempt to acquire a Redis-based lock. If acquired, set an expiration."""
|
||||
lock_acquired = redis_client.setnx(LOCK_KEY, "locked")
|
||||
if lock_acquired:
|
||||
redis_client.expire(LOCK_KEY, LOCK_EXPIRE)
|
||||
logger.info("Lock acquired for IMAP processing.")
|
||||
return True
|
||||
logger.warning("Lock already held. Skipping this cycle.")
|
||||
return False
|
||||
|
||||
|
||||
def release_lock():
|
||||
"""Release the lock by deleting the Redis key."""
|
||||
redis_client.delete(LOCK_KEY)
|
||||
logger.info("Lock released.")
|
||||
|
||||
|
||||
def load_processed_emails():
|
||||
"""Load the list of already processed emails from a local JSON file."""
|
||||
if os.path.exists(CACHE_FILE):
|
||||
try:
|
||||
with open(CACHE_FILE, "r") as f:
|
||||
processed_emails = json.load(f)
|
||||
processed_emails = cleanup_old_entries(processed_emails)
|
||||
return processed_emails
|
||||
except json.JSONDecodeError:
|
||||
logger.warning("Failed to decode JSON, resetting processed emails cache.")
|
||||
return {}
|
||||
return {}
|
||||
|
||||
|
||||
def save_processed_emails(processed_emails):
|
||||
"""Save the processed email IDs to a local JSON file."""
|
||||
with open(CACHE_FILE, "w") as f:
|
||||
json.dump(processed_emails, f, indent=4)
|
||||
|
||||
|
||||
def cleanup_old_entries(processed_emails):
|
||||
"""Remove entries older than 7 days from the cache to avoid infinite growth."""
|
||||
seven_days_ago = datetime.now(timezone.utc) - timedelta(days=7)
|
||||
valid_emails = {}
|
||||
for msg_id, date_str in processed_emails.items():
|
||||
naive_dt = datetime.strptime(date_str, "%Y-%m-%dT%H:%M:%S")
|
||||
aware_dt = naive_dt.replace(tzinfo=timezone.utc)
|
||||
if aware_dt > seven_days_ago:
|
||||
valid_emails[msg_id] = date_str
|
||||
return valid_emails
|
||||
|
||||
|
||||
@shared_task
|
||||
def pull_all_inboxes():
|
||||
"""
|
||||
Periodic Celery task that checks all configured IMAP mailboxes
|
||||
and fetches attachments from new emails.
|
||||
Ensures only one instance runs at a time using Redis-based locking.
|
||||
"""
|
||||
if not acquire_lock():
|
||||
logger.info("Skipping execution: Another instance is running.")
|
||||
return
|
||||
|
||||
try:
|
||||
logger.info("Starting pull_all_inboxes")
|
||||
|
||||
# Mailbox #1 (non-Gmail)
|
||||
check_and_pull_mailbox(
|
||||
mailbox_key="imap1",
|
||||
host=settings.imap1_host,
|
||||
port=settings.imap1_port,
|
||||
username=settings.imap1_username,
|
||||
password=settings.imap1_password,
|
||||
use_ssl=settings.imap1_ssl,
|
||||
delete_after_process=settings.imap1_delete_after_process,
|
||||
)
|
||||
|
||||
# Mailbox #2 (Gmail)
|
||||
check_and_pull_mailbox(
|
||||
mailbox_key="imap2",
|
||||
host=settings.imap2_host,
|
||||
port=settings.imap2_port,
|
||||
username=settings.imap2_username,
|
||||
password=settings.imap2_password,
|
||||
use_ssl=settings.imap2_ssl,
|
||||
delete_after_process=settings.imap2_delete_after_process,
|
||||
)
|
||||
|
||||
logger.info("Finished pull_all_inboxes")
|
||||
|
||||
finally:
|
||||
release_lock()
|
||||
|
||||
|
||||
def check_and_pull_mailbox(
|
||||
mailbox_key: str,
|
||||
host: str | None,
|
||||
port: int | None,
|
||||
username: str | None,
|
||||
password: str | None,
|
||||
use_ssl: bool,
|
||||
delete_after_process: bool,
|
||||
):
|
||||
"""Validates config and invokes pulling from the mailbox if valid."""
|
||||
if not (host and port and username and password):
|
||||
logger.warning(f"Mailbox {mailbox_key} is missing config, skipping.")
|
||||
return
|
||||
|
||||
logger.info(f"Checking mailbox: {mailbox_key}")
|
||||
pull_inbox(
|
||||
mailbox_key=mailbox_key,
|
||||
host=host,
|
||||
port=port,
|
||||
username=username,
|
||||
password=password,
|
||||
use_ssl=use_ssl,
|
||||
delete_after_process=delete_after_process,
|
||||
)
|
||||
|
||||
|
||||
def pull_inbox(mailbox_key, host, port, username, password, use_ssl,
|
||||
delete_after_process):
|
||||
"""
|
||||
Connects to the IMAP inbox, fetches new unread emails from the last 3 days,
|
||||
and processes attachments while preserving the original unread status.
|
||||
|
||||
For Gmail:
|
||||
- Attempts to select the localized All Mail folder.
|
||||
- Runs an X-GM-RAW query: "in:anywhere in:unread newer_than:3d has:attachment".
|
||||
|
||||
For non-Gmail mailboxes, it falls back to selecting the INBOX with a SINCE/UNSEEN filter.
|
||||
"""
|
||||
logger.info("Connecting to %s at %s:%s (SSL=%s)",
|
||||
mailbox_key, host, port, use_ssl)
|
||||
processed_emails = load_processed_emails()
|
||||
|
||||
try:
|
||||
mail = imaplib.IMAP4_SSL(host, port) if use_ssl else imaplib.IMAP4(host, port)
|
||||
mail.login(username, password)
|
||||
|
||||
is_gmail_host = "gmail" in host.lower()
|
||||
if is_gmail_host:
|
||||
# For Gmail, try to select the localized All Mail folder.
|
||||
all_mail_folder = find_all_mail_folder(mail)
|
||||
if all_mail_folder:
|
||||
logger.info("Using Gmail All Mail folder: %s", all_mail_folder)
|
||||
mail.select(f'"{all_mail_folder}"')
|
||||
else:
|
||||
logger.warning("Gmail All Mail folder not found, falling back to INBOX.")
|
||||
mail.select("INBOX")
|
||||
# Use the X-GM-RAW query for Gmail.
|
||||
raw_query = "in:anywhere in:unread newer_than:3d has:attachment"
|
||||
status, search_data = mail.search(None, "X-GM-RAW", f'"{raw_query}"')
|
||||
else:
|
||||
# For non-Gmail, select INBOX and use SINCE/UNSEEN query.
|
||||
mail.select("INBOX")
|
||||
since_date = (datetime.now(timezone.utc) - timedelta(days=3)
|
||||
).strftime("%d-%b-%Y")
|
||||
status, search_data = mail.search(None, f'(SINCE {since_date} UNSEEN)')
|
||||
|
||||
if status != "OK":
|
||||
logger.warning("Search failed on mailbox %s. Status=%s",
|
||||
mailbox_key, status)
|
||||
mail.close()
|
||||
mail.logout()
|
||||
return
|
||||
|
||||
msg_numbers = search_data[0].split()
|
||||
logger.info("Found %d unread emails in %s.", len(msg_numbers), mailbox_key)
|
||||
|
||||
for num in msg_numbers:
|
||||
status, msg_data = mail.fetch(num, "(RFC822)")
|
||||
if status != "OK":
|
||||
logger.warning("Failed to fetch message %s in %s. Status=%s",
|
||||
num, mailbox_key, status)
|
||||
continue
|
||||
|
||||
raw_email = msg_data[0][1]
|
||||
email_message = email.message_from_bytes(raw_email)
|
||||
msg_id = email_message.get("Message-ID")
|
||||
|
||||
if not msg_id:
|
||||
logger.warning("Skipping email without Message-ID in %s", mailbox_key)
|
||||
continue
|
||||
|
||||
if msg_id in processed_emails:
|
||||
logger.info("Skipping already processed email %s in %s", msg_id, mailbox_key)
|
||||
continue
|
||||
|
||||
# For Gmail, check if the email already has the "Ingested" label.
|
||||
if is_gmail_host:
|
||||
if email_already_has_label(mail, num, "Ingested"):
|
||||
logger.info("Skipping email %s in %s, already labeled 'Ingested'.",
|
||||
msg_id, mailbox_key)
|
||||
continue
|
||||
|
||||
# Process attachments (and convert non-PDF files).
|
||||
# We call the function without assigning its return value since it is not used.
|
||||
fetch_attachments_and_enqueue(email_message)
|
||||
|
||||
if is_gmail_host:
|
||||
mark_as_processed_with_star(mail, num)
|
||||
mark_as_processed_with_label(mail, num, label="Ingested")
|
||||
|
||||
processed_emails[msg_id] = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S")
|
||||
save_processed_emails(processed_emails)
|
||||
|
||||
if delete_after_process:
|
||||
logger.info("Deleting message %s from %s", num.decode(), mailbox_key)
|
||||
mail.store(num, "+FLAGS", "\\Deleted")
|
||||
else:
|
||||
mail.store(num, "-FLAGS", "\\Seen")
|
||||
|
||||
if delete_after_process:
|
||||
mail.expunge()
|
||||
|
||||
mail.close()
|
||||
mail.logout()
|
||||
logger.info("Finished processing mailbox %s", mailbox_key)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception("Error pulling mailbox %s: %s", mailbox_key, e)
|
||||
|
||||
|
||||
def fetch_attachments_and_enqueue(email_message):
|
||||
"""
|
||||
Extracts attachments from the email and processes only allowed file types.
|
||||
|
||||
Allowed file types include:
|
||||
- PDF: application/pdf
|
||||
- Microsoft Office files:
|
||||
- Word: application/msword,
|
||||
application/vnd.openxmlformats-officedocument.wordprocessingml.document
|
||||
- Excel: application/vnd.ms-excel,
|
||||
application/vnd.openxmlformats-officedocument.spreadsheetml.sheet
|
||||
- PowerPoint: application/vnd.ms-powerpoint,
|
||||
application/vnd.openxmlformats-officedocument.presentationml.presentation
|
||||
- Other meaningful attachments:
|
||||
- Plain text: text/plain
|
||||
- CSV: text/csv
|
||||
- Rich Text Format: application/rtf, text/rtf
|
||||
|
||||
Attachments not in this list are skipped. Common image MIME types such as
|
||||
image/jpeg, image/png, image/gif, image/bmp, image/tiff, and image/webp are
|
||||
intentionally excluded.
|
||||
|
||||
If the attachment is a PDF, it is enqueued for upload; any other allowed file
|
||||
is enqueued for conversion to PDF.
|
||||
|
||||
Returns True if at least one allowed attachment was processed.
|
||||
"""
|
||||
ALLOWED_MIME_TYPES = {
|
||||
"application/pdf",
|
||||
"application/msword",
|
||||
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||||
"application/vnd.ms-excel",
|
||||
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
||||
"application/vnd.ms-powerpoint",
|
||||
"application/vnd.openxmlformats-officedocument.presentationml.presentation",
|
||||
"text/plain",
|
||||
"text/csv",
|
||||
"application/rtf",
|
||||
"text/rtf",
|
||||
}
|
||||
|
||||
has_attachment = False
|
||||
for part in email_message.walk():
|
||||
if part.get_content_maintype() == "multipart":
|
||||
continue
|
||||
|
||||
filename = part.get_filename()
|
||||
if not filename:
|
||||
continue
|
||||
|
||||
mime_type = part.get_content_type()
|
||||
if mime_type not in ALLOWED_MIME_TYPES:
|
||||
logger.info("Skipping attachment %s with MIME type %s",
|
||||
filename, mime_type)
|
||||
continue
|
||||
|
||||
file_path = os.path.join(settings.workdir, filename)
|
||||
with open(file_path, "wb") as f:
|
||||
f.write(part.get_payload(decode=True))
|
||||
|
||||
if mime_type == "application/pdf":
|
||||
upload_to_s3.delay(file_path)
|
||||
logger.info("Enqueued PDF for upload: %s", filename)
|
||||
elif mime_type in ALLOWED_MIME_TYPES:
|
||||
# Enqueue conversion to PDF using the Gotenberg service.
|
||||
convert_to_pdf.delay(file_path)
|
||||
logger.info("Enqueued file for conversion to PDF: %s", filename)
|
||||
|
||||
has_attachment = True
|
||||
return has_attachment
|
||||
|
||||
|
||||
def email_already_has_label(mail, msg_id, label="Ingested"):
|
||||
"""
|
||||
Checks if the given message (msg_id) has the specified Gmail label.
|
||||
Returns True if the label is found, False otherwise.
|
||||
"""
|
||||
try:
|
||||
label_status, label_data = mail.fetch(msg_id, "(X-GM-LABELS)")
|
||||
if label_status == "OK" and label_data and len(label_data) > 0:
|
||||
raw_labels = label_data[0][1].decode("utf-8", errors="ignore")
|
||||
if label in raw_labels:
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error("Failed to fetch labels for msg_id=%s: %s", msg_id, e)
|
||||
return False
|
||||
|
||||
|
||||
def mark_as_processed_with_star(mail, msg_id):
|
||||
"""Stars the email in Gmail."""
|
||||
try:
|
||||
mail.store(msg_id, "+FLAGS", "\\Flagged")
|
||||
logger.info("Email %s starred in Gmail.", msg_id)
|
||||
except Exception as e:
|
||||
logger.error("Failed to star email %s: %s", msg_id, e)
|
||||
|
||||
|
||||
def mark_as_processed_with_label(mail, msg_id, label="Ingested"):
|
||||
"""Adds a custom label to the email in Gmail."""
|
||||
try:
|
||||
mail.store(msg_id, "+X-GM-LABELS", label)
|
||||
logger.info("Email %s labeled '%s' in Gmail.", msg_id, label)
|
||||
except Exception as e:
|
||||
logger.error("Failed to label email %s with %s: %s", msg_id, label, e)
|
||||
|
||||
|
||||
def find_all_mail_folder(mail):
|
||||
"""
|
||||
Attempts to select the Gmail All Mail folder using known localized names.
|
||||
Falls back to using XLIST if needed.
|
||||
Returns the folder name if found, otherwise None.
|
||||
"""
|
||||
COMMON_ALL_MAIL_NAMES = [
|
||||
"[Gmail]/Alle Nachrichten",
|
||||
"[Gmail]/All Mail",
|
||||
"[Gmail]/Todos",
|
||||
"[Gmail]/Tutte le mail",
|
||||
"[Gmail]/Tous les messages",
|
||||
]
|
||||
for candidate in COMMON_ALL_MAIL_NAMES:
|
||||
status, _ = mail.select(f'"{candidate}"', readonly=True)
|
||||
if status == "OK":
|
||||
return candidate
|
||||
|
||||
capabilities = get_capabilities(mail)
|
||||
if "XLIST" in capabilities:
|
||||
candidate = find_all_mail_xlist(mail)
|
||||
if candidate:
|
||||
return candidate
|
||||
return None
|
||||
|
||||
|
||||
def get_capabilities(mail):
|
||||
"""Returns a list of capabilities supported by the IMAP server."""
|
||||
typ, data = mail.capability()
|
||||
if typ == "OK" and data:
|
||||
caps = data[0].decode("utf-8", errors="ignore").upper().split()
|
||||
return caps
|
||||
return []
|
||||
|
||||
|
||||
def find_all_mail_xlist(mail):
|
||||
"""
|
||||
Uses XLIST to discover the mailbox flagged as All Mail.
|
||||
Returns the folder name if found, otherwise None.
|
||||
"""
|
||||
tag = mail._new_tag().decode("ascii")
|
||||
command_str = f"{tag} XLIST \"\" \"*\""
|
||||
mail.send((command_str + "\r\n").encode("utf-8"))
|
||||
|
||||
all_mail_folder = None
|
||||
while True:
|
||||
line = mail.readline()
|
||||
if not line:
|
||||
break
|
||||
line_str = line.decode("utf-8", errors="ignore").strip()
|
||||
if line_str.upper().startswith("* XLIST ") and "\\ALLMAIL" in line_str.upper():
|
||||
match = re.search(r'"([^"]+)"$', line_str)
|
||||
if match:
|
||||
candidate = match.group(1)
|
||||
logger.info("Found All Mail folder via XLIST: %s", candidate)
|
||||
all_mail_folder = candidate
|
||||
if line_str.startswith(tag):
|
||||
break
|
||||
return all_mail_folder
|
||||
|
||||
@@ -1,67 +1,67 @@
|
||||
import os
|
||||
import logging
|
||||
from azure.core.credentials import AzureKeyCredential
|
||||
from azure.ai.documentintelligence import DocumentIntelligenceClient
|
||||
from azure.ai.documentintelligence.models import AnalyzeOutputOption, AnalyzeResult
|
||||
|
||||
from app.config import settings
|
||||
from app.tasks.retry_config import BaseTaskWithRetry
|
||||
from app.tasks.extract_metadata_with_gpt import extract_metadata_with_gpt
|
||||
from app.celery_app import celery
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Initialize Azure Document Intelligence client
|
||||
document_intelligence_client = DocumentIntelligenceClient(
|
||||
endpoint=settings.azure_endpoint,
|
||||
credential=AzureKeyCredential(settings.azure_ai_key)
|
||||
)
|
||||
|
||||
@celery.task(base=BaseTaskWithRetry)
|
||||
def process_with_textract(s3_filename: str):
|
||||
"""
|
||||
Processes a PDF document using Azure Document Intelligence and overlays OCR text onto
|
||||
the local temporary file (stored under <workdir>/tmp).
|
||||
|
||||
Steps:
|
||||
1. Uploads the document for OCR using Azure Document Intelligence.
|
||||
2. Retrieves the processed PDF with embedded text.
|
||||
3. Saves the OCR-processed PDF locally in the same location as before.
|
||||
4. Extracts the text content for metadata processing.
|
||||
5. Triggers downstream metadata extraction by calling extract_metadata_with_gpt.
|
||||
"""
|
||||
try:
|
||||
tmp_file_path = os.path.join(settings.workdir, "tmp", s3_filename)
|
||||
if not os.path.exists(tmp_file_path):
|
||||
raise FileNotFoundError(f"Local file not found: {tmp_file_path}")
|
||||
|
||||
logger.info(f"Processing {s3_filename} with Azure Document Intelligence OCR.")
|
||||
|
||||
# Open and send the document for processing
|
||||
with open(tmp_file_path, "rb") as f:
|
||||
poller = document_intelligence_client.begin_analyze_document(
|
||||
"prebuilt-read", body=f, output=[AnalyzeOutputOption.PDF]
|
||||
)
|
||||
result: AnalyzeResult = poller.result()
|
||||
operation_id = poller.details["operation_id"]
|
||||
|
||||
# Retrieve the processed searchable PDF
|
||||
response = document_intelligence_client.get_analyze_result_pdf(
|
||||
model_id=result.model_id, result_id=operation_id
|
||||
)
|
||||
searchable_pdf_path = tmp_file_path # Overwrite the original PDF location
|
||||
with open(searchable_pdf_path, "wb") as writer:
|
||||
writer.writelines(response)
|
||||
logger.info(f"Searchable PDF saved at: {searchable_pdf_path}")
|
||||
|
||||
# Extract raw text content from the result
|
||||
extracted_text = result.content if result.content else ""
|
||||
logger.info(f"Extracted text for {s3_filename}: {len(extracted_text)} characters")
|
||||
|
||||
# Trigger downstream metadata extraction
|
||||
extract_metadata_with_gpt.delay(s3_filename, extracted_text)
|
||||
|
||||
return {"s3_file": s3_filename, "searchable_pdf": searchable_pdf_path, "cleaned_text": extracted_text}
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing {s3_filename} with Azure Document Intelligence: {e}")
|
||||
raise
|
||||
import os
|
||||
import logging
|
||||
from azure.core.credentials import AzureKeyCredential
|
||||
from azure.ai.documentintelligence import DocumentIntelligenceClient
|
||||
from azure.ai.documentintelligence.models import AnalyzeOutputOption, AnalyzeResult
|
||||
|
||||
from app.config import settings
|
||||
from app.tasks.retry_config import BaseTaskWithRetry
|
||||
from app.tasks.extract_metadata_with_gpt import extract_metadata_with_gpt
|
||||
from app.celery_app import celery
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Initialize Azure Document Intelligence client
|
||||
document_intelligence_client = DocumentIntelligenceClient(
|
||||
endpoint=settings.azure_endpoint,
|
||||
credential=AzureKeyCredential(settings.azure_ai_key)
|
||||
)
|
||||
|
||||
@celery.task(base=BaseTaskWithRetry)
|
||||
def process_with_textract(s3_filename: str):
|
||||
"""
|
||||
Processes a PDF document using Azure Document Intelligence and overlays OCR text onto
|
||||
the local temporary file (stored under <workdir>/tmp).
|
||||
|
||||
Steps:
|
||||
1. Uploads the document for OCR using Azure Document Intelligence.
|
||||
2. Retrieves the processed PDF with embedded text.
|
||||
3. Saves the OCR-processed PDF locally in the same location as before.
|
||||
4. Extracts the text content for metadata processing.
|
||||
5. Triggers downstream metadata extraction by calling extract_metadata_with_gpt.
|
||||
"""
|
||||
try:
|
||||
tmp_file_path = os.path.join(settings.workdir, "tmp", s3_filename)
|
||||
if not os.path.exists(tmp_file_path):
|
||||
raise FileNotFoundError(f"Local file not found: {tmp_file_path}")
|
||||
|
||||
logger.info(f"Processing {s3_filename} with Azure Document Intelligence OCR.")
|
||||
|
||||
# Open and send the document for processing
|
||||
with open(tmp_file_path, "rb") as f:
|
||||
poller = document_intelligence_client.begin_analyze_document(
|
||||
"prebuilt-read", body=f, output=[AnalyzeOutputOption.PDF]
|
||||
)
|
||||
result: AnalyzeResult = poller.result()
|
||||
operation_id = poller.details["operation_id"]
|
||||
|
||||
# Retrieve the processed searchable PDF
|
||||
response = document_intelligence_client.get_analyze_result_pdf(
|
||||
model_id=result.model_id, result_id=operation_id
|
||||
)
|
||||
searchable_pdf_path = tmp_file_path # Overwrite the original PDF location
|
||||
with open(searchable_pdf_path, "wb") as writer:
|
||||
writer.writelines(response)
|
||||
logger.info(f"Searchable PDF saved at: {searchable_pdf_path}")
|
||||
|
||||
# Extract raw text content from the result
|
||||
extracted_text = result.content if result.content else ""
|
||||
logger.info(f"Extracted text for {s3_filename}: {len(extracted_text)} characters")
|
||||
|
||||
# Trigger downstream metadata extraction
|
||||
extract_metadata_with_gpt.delay(s3_filename, extracted_text)
|
||||
|
||||
return {"s3_file": s3_filename, "searchable_pdf": searchable_pdf_path, "cleaned_text": extracted_text}
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing {s3_filename} with Azure Document Intelligence: {e}")
|
||||
raise
|
||||
|
||||
@@ -1,34 +1,34 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
from app.config import settings
|
||||
from openai import OpenAI
|
||||
from app.tasks.retry_config import BaseTaskWithRetry
|
||||
|
||||
# Import the shared Celery instance
|
||||
from app.celery_app import celery
|
||||
|
||||
|
||||
|
||||
client = OpenAI(api_key=settings.openai_api_key)
|
||||
|
||||
@celery.task(base=BaseTaskWithRetry)
|
||||
def refine_text_with_gpt(s3_filename: str, raw_text: str):
|
||||
"""Uses GPT to clean and refine OCR text."""
|
||||
|
||||
# Use the Chat Completions endpoint with 'messages'
|
||||
response = client.chat.completions.create(
|
||||
model="gpt-4",
|
||||
messages=[
|
||||
{"role": "system", "content": "Clean and format the following text. The idea is that the text you see comes from an OCR system and your task is to eliminate OCR errors. Keep the original language when doing so."},
|
||||
{"role": "user", "content": raw_text}
|
||||
]
|
||||
)
|
||||
|
||||
cleaned_text = response.choices[0].message.content
|
||||
|
||||
# Trigger next task (import locally if needed to avoid circular imports)
|
||||
from app.tasks.extract_metadata_with_gpt import extract_metadata_with_gpt
|
||||
extract_metadata_with_gpt.delay(s3_filename, cleaned_text)
|
||||
|
||||
return {"s3_file": s3_filename, "cleaned_text": cleaned_text}
|
||||
|
||||
#!/usr/bin/env python3
|
||||
|
||||
from app.config import settings
|
||||
from openai import OpenAI
|
||||
from app.tasks.retry_config import BaseTaskWithRetry
|
||||
|
||||
# Import the shared Celery instance
|
||||
from app.celery_app import celery
|
||||
|
||||
|
||||
|
||||
client = OpenAI(api_key=settings.openai_api_key)
|
||||
|
||||
@celery.task(base=BaseTaskWithRetry)
|
||||
def refine_text_with_gpt(s3_filename: str, raw_text: str):
|
||||
"""Uses GPT to clean and refine OCR text."""
|
||||
|
||||
# Use the Chat Completions endpoint with 'messages'
|
||||
response = client.chat.completions.create(
|
||||
model="gpt-4",
|
||||
messages=[
|
||||
{"role": "system", "content": "Clean and format the following text. The idea is that the text you see comes from an OCR system and your task is to eliminate OCR errors. Keep the original language when doing so."},
|
||||
{"role": "user", "content": raw_text}
|
||||
]
|
||||
)
|
||||
|
||||
cleaned_text = response.choices[0].message.content
|
||||
|
||||
# Trigger next task (import locally if needed to avoid circular imports)
|
||||
from app.tasks.extract_metadata_with_gpt import extract_metadata_with_gpt
|
||||
extract_metadata_with_gpt.delay(s3_filename, cleaned_text)
|
||||
|
||||
return {"s3_file": s3_filename, "cleaned_text": cleaned_text}
|
||||
|
||||
|
||||
@@ -1,9 +1,9 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
from celery import Task
|
||||
|
||||
class BaseTaskWithRetry(Task):
|
||||
autoretry_for = (Exception,)
|
||||
retry_kwargs = {"max_retries": 3, "countdown": 10} # 3 retries, 10s delay
|
||||
retry_backoff = True # Exponential backoff
|
||||
|
||||
#!/usr/bin/env python3
|
||||
|
||||
from celery import Task
|
||||
|
||||
class BaseTaskWithRetry(Task):
|
||||
autoretry_for = (Exception,)
|
||||
retry_kwargs = {"max_retries": 3, "countdown": 10} # 3 retries, 10s delay
|
||||
retry_backoff = True # Exponential backoff
|
||||
|
||||
|
||||
+21
-21
@@ -1,21 +1,21 @@
|
||||
# app/tasks/send_to_all.py
|
||||
|
||||
from app.celery_app import celery
|
||||
from app.tasks.upload_to_dropbox import upload_to_dropbox
|
||||
from app.tasks.upload_to_nextcloud import upload_to_nextcloud
|
||||
from app.tasks.upload_to_paperless import upload_to_paperless
|
||||
|
||||
@celery.task
|
||||
def send_to_all_destinations(file_path: str):
|
||||
"""
|
||||
Fires off tasks to upload a single file to Dropbox, Nextcloud, and Paperless.
|
||||
These tasks run in parallel (Celery returns immediately from each .delay()).
|
||||
"""
|
||||
upload_to_dropbox.delay(file_path)
|
||||
upload_to_nextcloud.delay(file_path)
|
||||
upload_to_paperless.delay(file_path)
|
||||
|
||||
return {
|
||||
"status": "All upload tasks enqueued",
|
||||
"file_path": file_path
|
||||
}
|
||||
# app/tasks/send_to_all.py
|
||||
|
||||
from app.celery_app import celery
|
||||
from app.tasks.upload_to_dropbox import upload_to_dropbox
|
||||
from app.tasks.upload_to_nextcloud import upload_to_nextcloud
|
||||
from app.tasks.upload_to_paperless import upload_to_paperless
|
||||
|
||||
@celery.task
|
||||
def send_to_all_destinations(file_path: str):
|
||||
"""
|
||||
Fires off tasks to upload a single file to Dropbox, Nextcloud, and Paperless.
|
||||
These tasks run in parallel (Celery returns immediately from each .delay()).
|
||||
"""
|
||||
upload_to_dropbox.delay(file_path)
|
||||
upload_to_nextcloud.delay(file_path)
|
||||
upload_to_paperless.delay(file_path)
|
||||
|
||||
return {
|
||||
"status": "All upload tasks enqueued",
|
||||
"file_path": file_path
|
||||
}
|
||||
|
||||
@@ -1,74 +1,74 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
import os
|
||||
import requests
|
||||
import dropbox
|
||||
from app.config import settings
|
||||
from app.tasks.retry_config import BaseTaskWithRetry
|
||||
from app.celery_app import celery
|
||||
|
||||
def get_dropbox_access_token():
|
||||
"""Refresh the Dropbox access token using the stored refresh token from ENV."""
|
||||
|
||||
token_url = "https://api.dropbox.com/oauth2/token"
|
||||
headers = {"Content-Type": "application/x-www-form-urlencoded"}
|
||||
data = {
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": settings.dropbox_refresh_token, # Now using ENV
|
||||
"client_id": settings.dropbox_app_key,
|
||||
"client_secret": settings.dropbox_app_secret,
|
||||
}
|
||||
|
||||
response = requests.post(token_url, headers=headers, data=data)
|
||||
|
||||
if response.status_code == 200:
|
||||
return response.json()["access_token"]
|
||||
else:
|
||||
error_msg = f"Failed to refresh Dropbox token: {response.status_code} - {response.text}"
|
||||
print(f"[ERROR] {error_msg}")
|
||||
raise Exception(error_msg)
|
||||
|
||||
@celery.task(base=BaseTaskWithRetry)
|
||||
def upload_to_dropbox(file_path: str):
|
||||
"""Uploads a file to Dropbox using the API."""
|
||||
|
||||
if not os.path.exists(file_path):
|
||||
raise FileNotFoundError(f"File not found: {file_path}")
|
||||
|
||||
# Extract filename and set target path
|
||||
filename = os.path.basename(file_path)
|
||||
dropbox_path = f"{settings.dropbox_folder}/{filename}"
|
||||
|
||||
try:
|
||||
# Get fresh access token
|
||||
access_token = get_dropbox_access_token()
|
||||
dbx = dropbox.Dropbox(access_token)
|
||||
|
||||
file_size = os.path.getsize(file_path)
|
||||
chunk_size = 4 * 1024 * 1024 # 4MB chunk size
|
||||
|
||||
with open(file_path, "rb") as file_data:
|
||||
if file_size <= chunk_size:
|
||||
dbx.files_upload(file_data.read(), dropbox_path)
|
||||
else:
|
||||
upload_session_start_result = dbx.files_upload_session_start(file_data.read(chunk_size))
|
||||
cursor = dropbox.files.UploadSessionCursor(
|
||||
session_id=upload_session_start_result.session_id,
|
||||
offset=file_data.tell(),
|
||||
)
|
||||
commit = dropbox.files.CommitInfo(path=dropbox_path)
|
||||
|
||||
while file_data.tell() < file_size:
|
||||
if (file_size - file_data.tell()) <= chunk_size:
|
||||
dbx.files_upload_session_finish(file_data.read(chunk_size), cursor, commit)
|
||||
else:
|
||||
dbx.files_upload_session_append_v2(file_data.read(chunk_size), cursor)
|
||||
cursor.offset = file_data.tell()
|
||||
|
||||
print(f"[INFO] Successfully uploaded {filename} to Dropbox at {dropbox_path}.")
|
||||
return {"status": "Completed", "file": file_path}
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"[ERROR] Failed to upload {filename} to Dropbox: {str(e)}"
|
||||
print(error_msg)
|
||||
raise Exception(error_msg)
|
||||
#!/usr/bin/env python3
|
||||
|
||||
import os
|
||||
import requests
|
||||
import dropbox
|
||||
from app.config import settings
|
||||
from app.tasks.retry_config import BaseTaskWithRetry
|
||||
from app.celery_app import celery
|
||||
|
||||
def get_dropbox_access_token():
|
||||
"""Refresh the Dropbox access token using the stored refresh token from ENV."""
|
||||
|
||||
token_url = "https://api.dropbox.com/oauth2/token"
|
||||
headers = {"Content-Type": "application/x-www-form-urlencoded"}
|
||||
data = {
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": settings.dropbox_refresh_token, # Now using ENV
|
||||
"client_id": settings.dropbox_app_key,
|
||||
"client_secret": settings.dropbox_app_secret,
|
||||
}
|
||||
|
||||
response = requests.post(token_url, headers=headers, data=data)
|
||||
|
||||
if response.status_code == 200:
|
||||
return response.json()["access_token"]
|
||||
else:
|
||||
error_msg = f"Failed to refresh Dropbox token: {response.status_code} - {response.text}"
|
||||
print(f"[ERROR] {error_msg}")
|
||||
raise Exception(error_msg)
|
||||
|
||||
@celery.task(base=BaseTaskWithRetry)
|
||||
def upload_to_dropbox(file_path: str):
|
||||
"""Uploads a file to Dropbox using the API."""
|
||||
|
||||
if not os.path.exists(file_path):
|
||||
raise FileNotFoundError(f"File not found: {file_path}")
|
||||
|
||||
# Extract filename and set target path
|
||||
filename = os.path.basename(file_path)
|
||||
dropbox_path = f"{settings.dropbox_folder}/{filename}"
|
||||
|
||||
try:
|
||||
# Get fresh access token
|
||||
access_token = get_dropbox_access_token()
|
||||
dbx = dropbox.Dropbox(access_token)
|
||||
|
||||
file_size = os.path.getsize(file_path)
|
||||
chunk_size = 4 * 1024 * 1024 # 4MB chunk size
|
||||
|
||||
with open(file_path, "rb") as file_data:
|
||||
if file_size <= chunk_size:
|
||||
dbx.files_upload(file_data.read(), dropbox_path)
|
||||
else:
|
||||
upload_session_start_result = dbx.files_upload_session_start(file_data.read(chunk_size))
|
||||
cursor = dropbox.files.UploadSessionCursor(
|
||||
session_id=upload_session_start_result.session_id,
|
||||
offset=file_data.tell(),
|
||||
)
|
||||
commit = dropbox.files.CommitInfo(path=dropbox_path)
|
||||
|
||||
while file_data.tell() < file_size:
|
||||
if (file_size - file_data.tell()) <= chunk_size:
|
||||
dbx.files_upload_session_finish(file_data.read(chunk_size), cursor, commit)
|
||||
else:
|
||||
dbx.files_upload_session_append_v2(file_data.read(chunk_size), cursor)
|
||||
cursor.offset = file_data.tell()
|
||||
|
||||
print(f"[INFO] Successfully uploaded {filename} to Dropbox at {dropbox_path}.")
|
||||
return {"status": "Completed", "file": file_path}
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"[ERROR] Failed to upload {filename} to Dropbox: {str(e)}"
|
||||
print(error_msg)
|
||||
raise Exception(error_msg)
|
||||
|
||||
@@ -1,37 +1,37 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
import os
|
||||
import requests
|
||||
from app.config import settings
|
||||
from app.tasks.retry_config import BaseTaskWithRetry
|
||||
from app.celery_app import celery
|
||||
|
||||
@celery.task(base=BaseTaskWithRetry)
|
||||
def upload_to_nextcloud(file_path: str):
|
||||
"""Uploads a file to Nextcloud in the configured folder."""
|
||||
|
||||
if not os.path.exists(file_path):
|
||||
raise FileNotFoundError(f"File not found: {file_path}")
|
||||
|
||||
# Extract filename
|
||||
filename = os.path.basename(file_path)
|
||||
|
||||
# Construct the full upload URL
|
||||
nextcloud_url = f"{settings.nextcloud_upload_url}/{settings.nextcloud_folder}/{filename}"
|
||||
|
||||
# Read file content
|
||||
with open(file_path, "rb") as file_data:
|
||||
response = requests.put(
|
||||
nextcloud_url,
|
||||
auth=(settings.nextcloud_username, settings.nextcloud_password),
|
||||
data=file_data
|
||||
)
|
||||
|
||||
# Check if upload was successful
|
||||
if response.status_code in (200, 201):
|
||||
print(f"[INFO] Successfully uploaded {filename} to Nextcloud at {nextcloud_url}.")
|
||||
return {"status": "Completed", "file": file_path}
|
||||
else:
|
||||
error_msg = f"[ERROR] Failed to upload {filename} to Nextcloud: {response.status_code} - {response.text}"
|
||||
print(error_msg)
|
||||
raise Exception(error_msg)
|
||||
#!/usr/bin/env python3
|
||||
|
||||
import os
|
||||
import requests
|
||||
from app.config import settings
|
||||
from app.tasks.retry_config import BaseTaskWithRetry
|
||||
from app.celery_app import celery
|
||||
|
||||
@celery.task(base=BaseTaskWithRetry)
|
||||
def upload_to_nextcloud(file_path: str):
|
||||
"""Uploads a file to Nextcloud in the configured folder."""
|
||||
|
||||
if not os.path.exists(file_path):
|
||||
raise FileNotFoundError(f"File not found: {file_path}")
|
||||
|
||||
# Extract filename
|
||||
filename = os.path.basename(file_path)
|
||||
|
||||
# Construct the full upload URL
|
||||
nextcloud_url = f"{settings.nextcloud_upload_url}/{settings.nextcloud_folder}/{filename}"
|
||||
|
||||
# Read file content
|
||||
with open(file_path, "rb") as file_data:
|
||||
response = requests.put(
|
||||
nextcloud_url,
|
||||
auth=(settings.nextcloud_username, settings.nextcloud_password),
|
||||
data=file_data
|
||||
)
|
||||
|
||||
# Check if upload was successful
|
||||
if response.status_code in (200, 201):
|
||||
print(f"[INFO] Successfully uploaded {filename} to Nextcloud at {nextcloud_url}.")
|
||||
return {"status": "Completed", "file": file_path}
|
||||
else:
|
||||
error_msg = f"[ERROR] Failed to upload {filename} to Nextcloud: {response.status_code} - {response.text}"
|
||||
print(error_msg)
|
||||
raise Exception(error_msg)
|
||||
|
||||
+131
-131
@@ -1,131 +1,131 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
import os
|
||||
import json
|
||||
import time
|
||||
import requests
|
||||
import logging
|
||||
from typing import Dict, Any
|
||||
|
||||
from app.config import settings
|
||||
from app.tasks.retry_config import BaseTaskWithRetry
|
||||
from app.celery_app import celery
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
POLL_MAX_ATTEMPTS = 10
|
||||
POLL_INTERVAL_SEC = 3
|
||||
|
||||
def _get_headers():
|
||||
"""Returns HTTP headers for Paperless-ngx API calls."""
|
||||
return {
|
||||
"Authorization": f"Token {settings.paperless_ngx_api_token}"
|
||||
}
|
||||
|
||||
def _paperless_api_url(path: str) -> str:
|
||||
"""
|
||||
Constructs a full Paperless-ngx API URL using `settings.paperless_host`.
|
||||
Ensures the path is appended with a leading slash if missing.
|
||||
"""
|
||||
host = settings.paperless_host.rstrip("/")
|
||||
if not path.startswith("/"):
|
||||
path = "/" + path
|
||||
return f"{host}{path}"
|
||||
|
||||
def poll_task_for_document_id(task_id: str) -> int:
|
||||
"""
|
||||
Polls /api/tasks/?task_id=<uuid> until we get status=SUCCESS or FAILURE,
|
||||
or until we run out of attempts.
|
||||
|
||||
On SUCCESS: returns the int document_id from 'related_document'.
|
||||
On FAILURE: raises RuntimeError with the task's 'result' message.
|
||||
If times out, raises TimeoutError.
|
||||
"""
|
||||
url = _paperless_api_url("/api/tasks/")
|
||||
attempts = 0
|
||||
|
||||
while attempts < POLL_MAX_ATTEMPTS:
|
||||
try:
|
||||
resp = requests.get(url, headers=_get_headers(), params={"task_id": task_id})
|
||||
resp.raise_for_status()
|
||||
tasks_data = resp.json()
|
||||
except requests.exceptions.RequestException as exc:
|
||||
logger.warning(
|
||||
"Failed to poll for task_id='%s'. Attempt=%d Error=%s",
|
||||
task_id, attempts + 1, exc
|
||||
)
|
||||
time.sleep(POLL_INTERVAL_SEC)
|
||||
attempts += 1
|
||||
continue
|
||||
|
||||
if isinstance(tasks_data, dict) and "results" in tasks_data:
|
||||
tasks_data = tasks_data["results"]
|
||||
|
||||
if tasks_data:
|
||||
task_info = tasks_data[0]
|
||||
status = task_info.get("status")
|
||||
if status == "SUCCESS":
|
||||
doc_str = task_info.get("related_document")
|
||||
if doc_str:
|
||||
return int(doc_str)
|
||||
raise RuntimeError(
|
||||
f"Task {task_id} completed but no doc ID found. Task info: {task_info}"
|
||||
)
|
||||
elif status == "FAILURE":
|
||||
raise RuntimeError(f"Task {task_id} failed: {task_info.get('result')}")
|
||||
|
||||
attempts += 1
|
||||
time.sleep(POLL_INTERVAL_SEC)
|
||||
|
||||
raise TimeoutError(
|
||||
f"Task {task_id} didn't reach SUCCESS within {POLL_MAX_ATTEMPTS} attempts."
|
||||
)
|
||||
|
||||
@celery.task(base=BaseTaskWithRetry)
|
||||
def upload_to_paperless(file_path: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Uploads a PDF to Paperless with minimal metadata (filename and date only).
|
||||
|
||||
1. Extracts the filename and date from the file.
|
||||
2. POSTs the PDF to Paperless => returns a quoted UUID string (task_id).
|
||||
3. Polls /api/tasks/?task_id=<uuid> until SUCCESS or FAILURE => doc_id.
|
||||
|
||||
Returns a dict with status, the paperless_task_id, paperless_document_id, and file_path.
|
||||
"""
|
||||
if not os.path.exists(file_path):
|
||||
raise FileNotFoundError(f"File not found: {file_path}")
|
||||
|
||||
base_name = os.path.basename(file_path)
|
||||
|
||||
# Upload the PDF
|
||||
post_url = _paperless_api_url("/api/documents/post_document/")
|
||||
with open(file_path, "rb") as f:
|
||||
files = {
|
||||
"document": (base_name, f, "application/pdf"),
|
||||
}
|
||||
data = {"title": base_name} # Title = Filename (no additional metadata)
|
||||
|
||||
try:
|
||||
logger.debug("Posting document to Paperless: file=%s", base_name)
|
||||
resp = requests.post(post_url, headers=_get_headers(), files=files, data=data)
|
||||
resp.raise_for_status()
|
||||
except requests.exceptions.RequestException as exc:
|
||||
logger.error(
|
||||
"Failed to upload document '%s' to Paperless. Error: %s. Response=%s",
|
||||
file_path, exc, getattr(exc.response, "text", "<no response>")
|
||||
)
|
||||
raise
|
||||
|
||||
raw_task_id = resp.text.strip().strip('"').strip("'")
|
||||
logger.info(f"Received Paperless task ID: {raw_task_id}")
|
||||
|
||||
# Poll tasks until success/fail => get doc_id
|
||||
doc_id = poll_task_for_document_id(raw_task_id)
|
||||
logger.info(f"Document {file_path} successfully ingested => ID={doc_id}")
|
||||
|
||||
return {
|
||||
"status": "Completed",
|
||||
"paperless_task_id": raw_task_id,
|
||||
"paperless_document_id": doc_id,
|
||||
"file_path": file_path
|
||||
}
|
||||
#!/usr/bin/env python3
|
||||
|
||||
import os
|
||||
import json
|
||||
import time
|
||||
import requests
|
||||
import logging
|
||||
from typing import Dict, Any
|
||||
|
||||
from app.config import settings
|
||||
from app.tasks.retry_config import BaseTaskWithRetry
|
||||
from app.celery_app import celery
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
POLL_MAX_ATTEMPTS = 10
|
||||
POLL_INTERVAL_SEC = 3
|
||||
|
||||
def _get_headers():
|
||||
"""Returns HTTP headers for Paperless-ngx API calls."""
|
||||
return {
|
||||
"Authorization": f"Token {settings.paperless_ngx_api_token}"
|
||||
}
|
||||
|
||||
def _paperless_api_url(path: str) -> str:
|
||||
"""
|
||||
Constructs a full Paperless-ngx API URL using `settings.paperless_host`.
|
||||
Ensures the path is appended with a leading slash if missing.
|
||||
"""
|
||||
host = settings.paperless_host.rstrip("/")
|
||||
if not path.startswith("/"):
|
||||
path = "/" + path
|
||||
return f"{host}{path}"
|
||||
|
||||
def poll_task_for_document_id(task_id: str) -> int:
|
||||
"""
|
||||
Polls /api/tasks/?task_id=<uuid> until we get status=SUCCESS or FAILURE,
|
||||
or until we run out of attempts.
|
||||
|
||||
On SUCCESS: returns the int document_id from 'related_document'.
|
||||
On FAILURE: raises RuntimeError with the task's 'result' message.
|
||||
If times out, raises TimeoutError.
|
||||
"""
|
||||
url = _paperless_api_url("/api/tasks/")
|
||||
attempts = 0
|
||||
|
||||
while attempts < POLL_MAX_ATTEMPTS:
|
||||
try:
|
||||
resp = requests.get(url, headers=_get_headers(), params={"task_id": task_id})
|
||||
resp.raise_for_status()
|
||||
tasks_data = resp.json()
|
||||
except requests.exceptions.RequestException as exc:
|
||||
logger.warning(
|
||||
"Failed to poll for task_id='%s'. Attempt=%d Error=%s",
|
||||
task_id, attempts + 1, exc
|
||||
)
|
||||
time.sleep(POLL_INTERVAL_SEC)
|
||||
attempts += 1
|
||||
continue
|
||||
|
||||
if isinstance(tasks_data, dict) and "results" in tasks_data:
|
||||
tasks_data = tasks_data["results"]
|
||||
|
||||
if tasks_data:
|
||||
task_info = tasks_data[0]
|
||||
status = task_info.get("status")
|
||||
if status == "SUCCESS":
|
||||
doc_str = task_info.get("related_document")
|
||||
if doc_str:
|
||||
return int(doc_str)
|
||||
raise RuntimeError(
|
||||
f"Task {task_id} completed but no doc ID found. Task info: {task_info}"
|
||||
)
|
||||
elif status == "FAILURE":
|
||||
raise RuntimeError(f"Task {task_id} failed: {task_info.get('result')}")
|
||||
|
||||
attempts += 1
|
||||
time.sleep(POLL_INTERVAL_SEC)
|
||||
|
||||
raise TimeoutError(
|
||||
f"Task {task_id} didn't reach SUCCESS within {POLL_MAX_ATTEMPTS} attempts."
|
||||
)
|
||||
|
||||
@celery.task(base=BaseTaskWithRetry)
|
||||
def upload_to_paperless(file_path: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Uploads a PDF to Paperless with minimal metadata (filename and date only).
|
||||
|
||||
1. Extracts the filename and date from the file.
|
||||
2. POSTs the PDF to Paperless => returns a quoted UUID string (task_id).
|
||||
3. Polls /api/tasks/?task_id=<uuid> until SUCCESS or FAILURE => doc_id.
|
||||
|
||||
Returns a dict with status, the paperless_task_id, paperless_document_id, and file_path.
|
||||
"""
|
||||
if not os.path.exists(file_path):
|
||||
raise FileNotFoundError(f"File not found: {file_path}")
|
||||
|
||||
base_name = os.path.basename(file_path)
|
||||
|
||||
# Upload the PDF
|
||||
post_url = _paperless_api_url("/api/documents/post_document/")
|
||||
with open(file_path, "rb") as f:
|
||||
files = {
|
||||
"document": (base_name, f, "application/pdf"),
|
||||
}
|
||||
data = {"title": base_name} # Title = Filename (no additional metadata)
|
||||
|
||||
try:
|
||||
logger.debug("Posting document to Paperless: file=%s", base_name)
|
||||
resp = requests.post(post_url, headers=_get_headers(), files=files, data=data)
|
||||
resp.raise_for_status()
|
||||
except requests.exceptions.RequestException as exc:
|
||||
logger.error(
|
||||
"Failed to upload document '%s' to Paperless. Error: %s. Response=%s",
|
||||
file_path, exc, getattr(exc.response, "text", "<no response>")
|
||||
)
|
||||
raise
|
||||
|
||||
raw_task_id = resp.text.strip().strip('"').strip("'")
|
||||
logger.info(f"Received Paperless task ID: {raw_task_id}")
|
||||
|
||||
# Poll tasks until success/fail => get doc_id
|
||||
doc_id = poll_task_for_document_id(raw_task_id)
|
||||
logger.info(f"Document {file_path} successfully ingested => ID={doc_id}")
|
||||
|
||||
return {
|
||||
"status": "Completed",
|
||||
"paperless_task_id": raw_task_id,
|
||||
"paperless_document_id": doc_id,
|
||||
"file_path": file_path
|
||||
}
|
||||
|
||||
+87
-87
@@ -1,87 +1,87 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
import os
|
||||
import uuid
|
||||
import boto3
|
||||
import shutil
|
||||
import fitz # PyMuPDF for checking embedded text
|
||||
from app.config import settings
|
||||
from app.tasks.retry_config import BaseTaskWithRetry
|
||||
from app.tasks.process_with_textract import process_with_textract
|
||||
from app.tasks.extract_metadata_with_gpt import extract_metadata_with_gpt
|
||||
|
||||
# Import the shared Celery instance
|
||||
from app.celery_app import celery
|
||||
|
||||
# Initialize S3 client
|
||||
s3_client = boto3.client(
|
||||
"s3",
|
||||
aws_access_key_id=settings.aws_access_key_id,
|
||||
aws_secret_access_key=settings.aws_secret_access_key,
|
||||
region_name=settings.aws_region,
|
||||
)
|
||||
|
||||
@celery.task(base=BaseTaskWithRetry)
|
||||
def upload_to_s3(original_local_file: str):
|
||||
"""
|
||||
Uploads a file to S3 with a UUID-based filename and triggers processing.
|
||||
- If the PDF already contains embedded text, skip Textract and extract text locally.
|
||||
- Otherwise, upload to S3 and process with Textract.
|
||||
"""
|
||||
bucket_name = settings.s3_bucket_name
|
||||
if not bucket_name:
|
||||
print("[ERROR] S3 bucket name not set.")
|
||||
return {"error": "Missing S3 bucket name"}
|
||||
|
||||
if not os.path.exists(original_local_file):
|
||||
print(f"[ERROR] File {original_local_file} not found.")
|
||||
return {"error": "File not found"}
|
||||
|
||||
# Generate UUID and create a new filename
|
||||
file_ext = os.path.splitext(original_local_file)[1] # Preserve original file extension
|
||||
file_uuid = str(uuid.uuid4())
|
||||
new_filename = f"{file_uuid}{file_ext}"
|
||||
|
||||
# Construct the new local path using settings.workdir and a 'tmp' subdirectory
|
||||
tmp_dir = os.path.join(settings.workdir, "tmp")
|
||||
new_local_path = os.path.join(tmp_dir, new_filename)
|
||||
|
||||
# Ensure the target tmp directory exists
|
||||
os.makedirs(tmp_dir, exist_ok=True)
|
||||
|
||||
# Copy the file instead of moving it
|
||||
shutil.copy(original_local_file, new_local_path)
|
||||
|
||||
# Check for embedded text
|
||||
pdf_doc = fitz.open(new_local_path)
|
||||
has_text = any(page.get_text() for page in pdf_doc)
|
||||
pdf_doc.close()
|
||||
|
||||
if has_text:
|
||||
print(f"[INFO] PDF {original_local_file} contains embedded text. Skipping Textract.")
|
||||
|
||||
# Extract text locally
|
||||
extracted_text = ""
|
||||
pdf_doc = fitz.open(new_local_path)
|
||||
for page in pdf_doc:
|
||||
extracted_text += page.get_text("text") + "\n"
|
||||
pdf_doc.close()
|
||||
|
||||
# Call metadata extraction directly
|
||||
extract_metadata_with_gpt.delay(new_filename, extracted_text)
|
||||
|
||||
return {"file": new_local_path, "status": "Text extracted locally"}
|
||||
|
||||
try:
|
||||
print(f"[INFO] Uploading {new_local_path} to s3://{bucket_name}/{new_filename}...")
|
||||
s3_client.upload_file(new_local_path, bucket_name, new_filename)
|
||||
print(f"[INFO] File uploaded successfully: {new_filename}")
|
||||
|
||||
# Trigger Textract processing if no embedded text was found
|
||||
process_with_textract.delay(new_filename)
|
||||
|
||||
return {"file": new_local_path, "s3_key": new_filename, "status": "Uploaded to S3 for OCR"}
|
||||
|
||||
except Exception as e:
|
||||
print(f"[ERROR] Failed to upload {new_local_path} to S3: {e}")
|
||||
return {"error": str(e)}
|
||||
#!/usr/bin/env python3
|
||||
|
||||
import os
|
||||
import uuid
|
||||
import boto3
|
||||
import shutil
|
||||
import fitz # PyMuPDF for checking embedded text
|
||||
from app.config import settings
|
||||
from app.tasks.retry_config import BaseTaskWithRetry
|
||||
from app.tasks.process_with_textract import process_with_textract
|
||||
from app.tasks.extract_metadata_with_gpt import extract_metadata_with_gpt
|
||||
|
||||
# Import the shared Celery instance
|
||||
from app.celery_app import celery
|
||||
|
||||
# Initialize S3 client
|
||||
s3_client = boto3.client(
|
||||
"s3",
|
||||
aws_access_key_id=settings.aws_access_key_id,
|
||||
aws_secret_access_key=settings.aws_secret_access_key,
|
||||
region_name=settings.aws_region,
|
||||
)
|
||||
|
||||
@celery.task(base=BaseTaskWithRetry)
|
||||
def upload_to_s3(original_local_file: str):
|
||||
"""
|
||||
Uploads a file to S3 with a UUID-based filename and triggers processing.
|
||||
- If the PDF already contains embedded text, skip Textract and extract text locally.
|
||||
- Otherwise, upload to S3 and process with Textract.
|
||||
"""
|
||||
bucket_name = settings.s3_bucket_name
|
||||
if not bucket_name:
|
||||
print("[ERROR] S3 bucket name not set.")
|
||||
return {"error": "Missing S3 bucket name"}
|
||||
|
||||
if not os.path.exists(original_local_file):
|
||||
print(f"[ERROR] File {original_local_file} not found.")
|
||||
return {"error": "File not found"}
|
||||
|
||||
# Generate UUID and create a new filename
|
||||
file_ext = os.path.splitext(original_local_file)[1] # Preserve original file extension
|
||||
file_uuid = str(uuid.uuid4())
|
||||
new_filename = f"{file_uuid}{file_ext}"
|
||||
|
||||
# Construct the new local path using settings.workdir and a 'tmp' subdirectory
|
||||
tmp_dir = os.path.join(settings.workdir, "tmp")
|
||||
new_local_path = os.path.join(tmp_dir, new_filename)
|
||||
|
||||
# Ensure the target tmp directory exists
|
||||
os.makedirs(tmp_dir, exist_ok=True)
|
||||
|
||||
# Copy the file instead of moving it
|
||||
shutil.copy(original_local_file, new_local_path)
|
||||
|
||||
# Check for embedded text
|
||||
pdf_doc = fitz.open(new_local_path)
|
||||
has_text = any(page.get_text() for page in pdf_doc)
|
||||
pdf_doc.close()
|
||||
|
||||
if has_text:
|
||||
print(f"[INFO] PDF {original_local_file} contains embedded text. Skipping Textract.")
|
||||
|
||||
# Extract text locally
|
||||
extracted_text = ""
|
||||
pdf_doc = fitz.open(new_local_path)
|
||||
for page in pdf_doc:
|
||||
extracted_text += page.get_text("text") + "\n"
|
||||
pdf_doc.close()
|
||||
|
||||
# Call metadata extraction directly
|
||||
extract_metadata_with_gpt.delay(new_filename, extracted_text)
|
||||
|
||||
return {"file": new_local_path, "status": "Text extracted locally"}
|
||||
|
||||
try:
|
||||
print(f"[INFO] Uploading {new_local_path} to s3://{bucket_name}/{new_filename}...")
|
||||
s3_client.upload_file(new_local_path, bucket_name, new_filename)
|
||||
print(f"[INFO] File uploaded successfully: {new_filename}")
|
||||
|
||||
# Trigger Textract processing if no embedded text was found
|
||||
process_with_textract.delay(new_filename)
|
||||
|
||||
return {"file": new_local_path, "s3_key": new_filename, "status": "Uploaded to S3 for OCR"}
|
||||
|
||||
except Exception as e:
|
||||
print(f"[ERROR] Failed to upload {new_local_path} to S3: {e}")
|
||||
return {"error": str(e)}
|
||||
|
||||
Reference in New Issue
Block a user