228 lines
8.1 KiB
Python
228 lines
8.1 KiB
Python
"""
|
|
OCR Service — extracts transaction data from receipt images using Groq vision.
|
|
|
|
Model: meta-llama/llama-4-scout-17b-16e-instruct (multimodal)
|
|
Input: image file path or bytes
|
|
Output: dict with amount, date, merchant, category_suggestion, notes
|
|
"""
|
|
|
|
import base64
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
from datetime import date, datetime
|
|
from flask import current_app
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
GROQ_VISION_MODEL = 'meta-llama/llama-4-scout-17b-16e-instruct'
|
|
SUPPORTED_MIME_TYPES = {
|
|
'.jpg': 'image/jpeg',
|
|
'.jpeg': 'image/jpeg',
|
|
'.png': 'image/png',
|
|
'.gif': 'image/gif',
|
|
'.webp': 'image/webp',
|
|
}
|
|
|
|
RECEIPT_PROMPT = """You are a receipt data extractor. Analyze this receipt image and extract the following information.
|
|
|
|
Respond ONLY with a valid JSON object, no explanation, no markdown, no code fences. Just raw JSON.
|
|
|
|
Required format:
|
|
{
|
|
"amount": <total amount as a number, no currency symbols, e.g. 85.50>,
|
|
"date": "<date in YYYY-MM-DD format, or null if not found>",
|
|
"merchant": "<store/restaurant/vendor name, or null if not found>",
|
|
"category_suggestion": "<one of: Food & Dining, Transport, Shopping, Health, Entertainment, Utilities, Housing, Education, Personal Care, Travel, Other>",
|
|
"notes": "<brief description, e.g. 'Lunch at McDonald\\'s' or 'Grocery shopping'>"
|
|
}
|
|
|
|
Rules:
|
|
- amount: use the TOTAL/GRAND TOTAL amount, not subtotal. Numbers only, no commas.
|
|
- date: format as YYYY-MM-DD. If only month/day visible, use current year.
|
|
- merchant: the business name only, no address.
|
|
- category_suggestion: pick the single most appropriate category from the list.
|
|
- notes: keep under 60 characters.
|
|
- If you cannot read the receipt clearly, still return your best guess with available data.
|
|
- Never return null for amount if any number is visible."""
|
|
|
|
|
|
def extract_from_file(file_path):
|
|
"""
|
|
Extract receipt data from an image file on disk.
|
|
Returns dict or None on failure.
|
|
"""
|
|
ext = os.path.splitext(file_path)[1].lower()
|
|
mime_type = SUPPORTED_MIME_TYPES.get(ext)
|
|
|
|
if not mime_type:
|
|
log.warning(f'[ocr] Unsupported file type: {ext}')
|
|
return _error_result(f'Unsupported file type: {ext}')
|
|
|
|
try:
|
|
with open(file_path, 'rb') as f:
|
|
image_bytes = f.read()
|
|
return extract_from_bytes(image_bytes, mime_type)
|
|
except FileNotFoundError:
|
|
log.error(f'[ocr] File not found: {file_path}')
|
|
return _error_result('Receipt file not found')
|
|
except Exception as e:
|
|
log.error(f'[ocr] File read error: {e}')
|
|
return _error_result(str(e))
|
|
|
|
|
|
def extract_from_bytes(image_bytes, mime_type='image/jpeg'):
|
|
"""
|
|
Extract receipt data from raw image bytes.
|
|
Returns dict with keys: amount, date, merchant, category_suggestion, notes, error
|
|
"""
|
|
api_key = current_app.config.get('GROQ_API_KEY', '')
|
|
if not api_key:
|
|
return _error_result('GROQ_API_KEY not configured')
|
|
|
|
# Encode to base64
|
|
b64 = base64.b64encode(image_bytes).decode('utf-8')
|
|
data_url = f'data:{mime_type};base64,{b64}'
|
|
|
|
try:
|
|
import requests
|
|
resp = requests.post(
|
|
'https://api.groq.com/openai/v1/chat/completions',
|
|
headers={
|
|
'Authorization': f'Bearer {api_key}',
|
|
'Content-Type': 'application/json',
|
|
},
|
|
json={
|
|
'model': GROQ_VISION_MODEL,
|
|
'messages': [
|
|
{
|
|
'role': 'user',
|
|
'content': [
|
|
{
|
|
'type': 'image_url',
|
|
'image_url': {'url': data_url},
|
|
},
|
|
{
|
|
'type': 'text',
|
|
'text': RECEIPT_PROMPT,
|
|
},
|
|
],
|
|
}
|
|
],
|
|
'max_tokens': 512,
|
|
'temperature': 0.1, # low temp for consistent structured output
|
|
},
|
|
timeout=30,
|
|
)
|
|
resp.raise_for_status()
|
|
content = resp.json()['choices'][0]['message']['content'].strip()
|
|
log.info(f'[ocr] Raw response: {content[:200]}')
|
|
return _parse_response(content)
|
|
|
|
except requests.exceptions.HTTPError as e:
|
|
if e.response.status_code == 400:
|
|
# Model may not support this image — try to parse error
|
|
try:
|
|
err_body = e.response.json()
|
|
msg = err_body.get('error', {}).get('message', str(e))
|
|
except Exception:
|
|
msg = str(e)
|
|
log.error(f'[ocr] Groq 400 error: {msg}')
|
|
return _error_result(f'Image processing failed: {msg[:100]}')
|
|
elif e.response.status_code == 429:
|
|
return _error_result('Rate limit reached. Try again shortly.')
|
|
elif e.response.status_code == 401:
|
|
return _error_result('Invalid Groq API key.')
|
|
else:
|
|
log.error(f'[ocr] HTTP {e.response.status_code}: {e}')
|
|
return _error_result(f'API error ({e.response.status_code})')
|
|
except Exception as e:
|
|
log.error(f'[ocr] Unexpected error: {e}')
|
|
return _error_result('OCR service temporarily unavailable')
|
|
|
|
|
|
def _parse_response(content):
|
|
"""Parse the JSON response from Groq into a clean dict."""
|
|
# Strip any accidental markdown fences
|
|
content = re.sub(r'^```(?:json)?\s*', '', content, flags=re.MULTILINE)
|
|
content = re.sub(r'\s*```$', '', content, flags=re.MULTILINE)
|
|
content = content.strip()
|
|
|
|
try:
|
|
data = json.loads(content)
|
|
except json.JSONDecodeError:
|
|
# Try to extract JSON object from surrounding text
|
|
match = re.search(r'\{.*\}', content, re.DOTALL)
|
|
if match:
|
|
try:
|
|
data = json.loads(match.group())
|
|
except json.JSONDecodeError:
|
|
log.error(f'[ocr] Could not parse JSON: {content[:200]}')
|
|
return _error_result('Could not parse OCR response')
|
|
else:
|
|
return _error_result('No JSON found in OCR response')
|
|
|
|
# Normalise amount
|
|
amount = data.get('amount')
|
|
if amount is not None:
|
|
try:
|
|
amount = float(str(amount).replace(',', '').strip())
|
|
if amount <= 0:
|
|
amount = None
|
|
except (ValueError, TypeError):
|
|
amount = None
|
|
|
|
# Normalise date
|
|
raw_date = data.get('date')
|
|
parsed_date = None
|
|
if raw_date:
|
|
for fmt in ('%Y-%m-%d', '%m/%d/%Y', '%d/%m/%Y', '%Y/%m/%d'):
|
|
try:
|
|
parsed_date = datetime.strptime(str(raw_date).strip(), fmt).date().isoformat()
|
|
break
|
|
except ValueError:
|
|
continue
|
|
|
|
# Map category suggestion to our system category names
|
|
cat_map = {
|
|
'food & dining': 'Food & Dining',
|
|
'food': 'Food & Dining',
|
|
'dining': 'Food & Dining',
|
|
'transport': 'Transport',
|
|
'transportation':'Transport',
|
|
'shopping': 'Shopping',
|
|
'health': 'Health',
|
|
'entertainment': 'Entertainment',
|
|
'utilities': 'Utilities',
|
|
'housing': 'Housing',
|
|
'education': 'Education',
|
|
'personal care': 'Personal Care',
|
|
'travel': 'Travel',
|
|
'other': 'Other',
|
|
}
|
|
raw_cat = str(data.get('category_suggestion', '')).lower().strip()
|
|
category = cat_map.get(raw_cat, data.get('category_suggestion', 'Other'))
|
|
|
|
return {
|
|
'amount': amount,
|
|
'date': parsed_date or date.today().isoformat(),
|
|
'merchant': data.get('merchant') or '',
|
|
'category_suggestion': category,
|
|
'notes': data.get('notes') or '',
|
|
'error': None,
|
|
'raw': data,
|
|
}
|
|
|
|
|
|
def _error_result(msg):
|
|
return {
|
|
'amount': None,
|
|
'date': date.today().isoformat(),
|
|
'merchant': '',
|
|
'category_suggestion': 'Other',
|
|
'notes': '',
|
|
'error': msg,
|
|
'raw': {},
|
|
}
|