openai_service.py 2.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879
  1. import os
  2. import base64
  3. import io
  4. import logging
  5. from openai import AsyncOpenAI
  6. from dotenv import load_dotenv
  7. from PIL import Image
  8. from backend.schemas import ExtractionResponse
  9. load_dotenv()
  10. # Setup logging
  11. logging.basicConfig(level=logging.INFO)
  12. logger = logging.getLogger(__name__)
  13. client = AsyncOpenAI(api_key=os.getenv("OPENAI_API_KEY"))
  14. MAX_IMAGE_SIZE = (2000, 2000)
  15. IMAGE_QUALITY = 85
  16. def compress_image(image_content: bytes) -> bytes:
  17. """Resizes and compresses the image to reduce upload size."""
  18. img = Image.open(io.BytesIO(image_content))
  19. # Convert to RGB if necessary (to save as JPEG)
  20. if img.mode in ("RGBA", "P"):
  21. img = img.convert("RGB")
  22. # Resize if larger than max dimensions
  23. img.thumbnail(MAX_IMAGE_SIZE, Image.Resampling.LANCZOS)
  24. # Save to bytes
  25. output_buffer = io.BytesIO()
  26. img.save(output_buffer, format="JPEG", quality=IMAGE_QUALITY, optimize=True)
  27. return output_buffer.getvalue()
  28. async def extract_receipt_data(image_content: bytes, user_name: str, department: str) -> ExtractionResponse:
  29. # 1. Compress Image
  30. compressed_content = compress_image(image_content)
  31. base64_image = base64.b64encode(compressed_content).decode("utf-8")
  32. # 2. Refined Prompt
  33. prompt = (
  34. f"You are an HR data entry assistant helping an employee in Malaysia. "
  35. f"Extract the requested fields from the provided medical receipt image. "
  36. f"The employee submitting this is {user_name} from {department}. "
  37. f"IMPORTANT: The currency is always Ringgit Malaysia (MYR). Extract the total amount and assume it is in MYR. "
  38. f"If the date is missing, look for a 'Payment Date' as a fallback. "
  39. f"Analyze the receipt for authenticity. If the total amount appears altered or if the provider name is missing, "
  40. f"set `needs_manual_review` to `true` and provide a low `confidence_score`."
  41. )
  42. # 3. Async Extraction
  43. completion = await client.beta.chat.completions.parse(
  44. model="gpt-4o",
  45. messages=[
  46. {
  47. "role": "user",
  48. "content": [
  49. {"type": "text", "text": prompt},
  50. {
  51. "type": "image_url",
  52. "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"},
  53. },
  54. ],
  55. }
  56. ],
  57. response_format=ExtractionResponse,
  58. )
  59. result = completion.choices[0].message.parsed
  60. # 4. Logging for Demo
  61. if result:
  62. logger.info(f"Extraction complete for {user_name}. Confidence Score: {result.confidence_score}")
  63. if result.needs_manual_review:
  64. logger.warning(f"Manual review required for receipt submitted by {user_name}")
  65. return result