Feat: Desacoplamento

code/services/document_extractor.py

@@ -0,0 +1,75 @@
+import asyncio
+import io
+from pathlib import Path
+from urllib.parse import urlparse
+
+import boto3
+from PyPDF2 import PdfReader, PdfWriter
+
+from utils.config import AWS_REGION
+from utils.secrets_manager import SECRETS
+
+_s3_input = boto3.client(
+    "s3",
+    aws_access_key_id=SECRETS["AWS_ACCESS_KEY"],
+    aws_secret_access_key=SECRETS["AWS_SECRET_KEY"],
+    region_name=AWS_REGION,
+)
+_textract = boto3.client("textract", region_name=AWS_REGION)
+
+
+def parse_s3_uri(s3_uri: str) -> tuple[str, str]:
+    parsed = urlparse(s3_uri)
+    if parsed.scheme != "s3":
+        raise ValueError(f"Not an S3 URI: {s3_uri}")
+    bucket = parsed.netloc
+    key = parsed.path.lstrip("/")
+    if not bucket or not key:
+        raise ValueError(f"Invalid S3 URI: {s3_uri}")
+    return bucket, key
+
+
+def _extract_text_from_textract_response(response: dict) -> str:
+    if not response:
+        return ""
+    return "\n".join(
+        block["Text"] for block in response.get("Blocks", [])
+        if block["BlockType"] == "LINE"
+    )
+
+
+def _split_pdf_pages(pdf_bytes: bytes) -> list[bytes]:
+    reader = PdfReader(io.BytesIO(pdf_bytes))
+    pages = []
+    for page in reader.pages:
+        writer = PdfWriter()
+        writer.add_page(page)
+        buf = io.BytesIO()
+        writer.write(buf)
+        pages.append(buf.getvalue())
+    return pages
+
+
+def _textract_detect_bytes(file_bytes: bytes) -> str:
+    response = _textract.detect_document_text(Document={"Bytes": file_bytes})
+    return _extract_text_from_textract_response(response)
+
+
+async def extract_text_from_s3_document(bucket: str, key: str) -> tuple[str, int]:
+    file_bytes = await asyncio.to_thread(
+        lambda: _s3_input.get_object(Bucket=bucket, Key=key)["Body"].read()
+    )
+    file_ext = Path(key).suffix.lower()
+
+    if file_ext in [".png", ".jpg", ".jpeg"]:
+        text = await asyncio.to_thread(_textract_detect_bytes, file_bytes)
+        return text, 1
+
+    if file_ext == ".pdf":
+        page_bytes_list = await asyncio.to_thread(_split_pdf_pages, file_bytes)
+        texts = await asyncio.gather(*[
+            asyncio.to_thread(_textract_detect_bytes, p) for p in page_bytes_list
+        ])
+        return "\n".join(texts), len(page_bytes_list)
+
+    return "", 0