| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879 |
- import os
- import base64
- import io
- import logging
- from openai import AsyncOpenAI
- from dotenv import load_dotenv
- from PIL import Image
- from backend.schemas import ExtractionResponse
- load_dotenv()
- # Setup logging
- logging.basicConfig(level=logging.INFO)
- logger = logging.getLogger(__name__)
- client = AsyncOpenAI(api_key=os.getenv("OPENAI_API_KEY"))
- MAX_IMAGE_SIZE = (2000, 2000)
- IMAGE_QUALITY = 85
- def compress_image(image_content: bytes) -> bytes:
- """Resizes and compresses the image to reduce upload size."""
- img = Image.open(io.BytesIO(image_content))
-
- # Convert to RGB if necessary (to save as JPEG)
- if img.mode in ("RGBA", "P"):
- img = img.convert("RGB")
-
- # Resize if larger than max dimensions
- img.thumbnail(MAX_IMAGE_SIZE, Image.Resampling.LANCZOS)
-
- # Save to bytes
- output_buffer = io.BytesIO()
- img.save(output_buffer, format="JPEG", quality=IMAGE_QUALITY, optimize=True)
- return output_buffer.getvalue()
- async def extract_receipt_data(image_content: bytes, user_name: str, department: str) -> ExtractionResponse:
- # 1. Compress Image
- compressed_content = compress_image(image_content)
- base64_image = base64.b64encode(compressed_content).decode("utf-8")
-
- # 2. Refined Prompt
- prompt = (
- f"You are an HR data entry assistant helping an employee in Malaysia. "
- f"Extract the requested fields from the provided medical receipt image. "
- f"The employee submitting this is {user_name} from {department}. "
- f"IMPORTANT: The currency is always Ringgit Malaysia (MYR). Extract the total amount and assume it is in MYR. "
- f"If the date is missing, look for a 'Payment Date' as a fallback. "
- f"Analyze the receipt for authenticity. If the total amount appears altered or if the provider name is missing, "
- f"set `needs_manual_review` to `true` and provide a low `confidence_score`."
- )
-
- # 3. Async Extraction
- completion = await client.beta.chat.completions.parse(
- model="gpt-4o",
- messages=[
- {
- "role": "user",
- "content": [
- {"type": "text", "text": prompt},
- {
- "type": "image_url",
- "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"},
- },
- ],
- }
- ],
- response_format=ExtractionResponse,
- )
-
- result = completion.choices[0].message.parsed
-
- # 4. Logging for Demo
- if result:
- logger.info(f"Extraction complete for {user_name}. Confidence Score: {result.confidence_score}")
- if result.needs_manual_review:
- logger.warning(f"Manual review required for receipt submitted by {user_name}")
-
- return result
|