89 lines
3.5 KiB
Python
89 lines
3.5 KiB
Python
from . import models
|
|
from .database import SessionLocal
|
|
|
|
# Load environment variables from the directory where this file resides
|
|
base_dir = os.path.dirname(os.path.abspath(__file__))
|
|
dotenv_path = os.path.join(base_dir, ".env")
|
|
load_dotenv(dotenv_path)
|
|
|
|
def extract_label_info(image_bytes: bytes, mode: str = "item"):
|
|
"""
|
|
Orchestrates extraction across multiple AI providers.
|
|
Modes: 'item' (full technical extraction), 'box' (container discovery)
|
|
"""
|
|
db = SessionLocal()
|
|
try:
|
|
if mode == "box":
|
|
prompt = """
|
|
Identify the CONTAINER or BOX name from this image.
|
|
Look for large, prominent, bold, or hand-written text that identifies a storage unit.
|
|
Ignore small technical details, quantities, or fine print.
|
|
|
|
Return ONLY a valid JSON object:
|
|
{
|
|
"box_label": "The identified container name",
|
|
"name": "Same as box_label",
|
|
"category": "Storage",
|
|
"description": "Brief description if useful",
|
|
"quantity": 1
|
|
}
|
|
"""
|
|
else:
|
|
# Fetch custom prompt from DB
|
|
setting = db.query(models.SystemSetting).filter(models.SystemSetting.key == "ai_extraction_prompt").first()
|
|
if setting:
|
|
prompt = setting.value
|
|
else:
|
|
# Fallback to a sensible default if DB is not ready
|
|
prompt = "Extract technical specs. Return JSON with name, category, description, connector, size, color, part_number, ocr_text, quantity."
|
|
|
|
# 1. Try Gemini
|
|
result = gemini.extract(image_bytes, prompt)
|
|
|
|
if result:
|
|
# Map user-defined prompt keys to model fields if needed
|
|
# User keys: Item, Type, Description, Category, Connector, Size, Color, PartNr, OCR
|
|
mapping = {
|
|
"Item": "name",
|
|
"Type": "type",
|
|
"Description": "description",
|
|
"Category": "category",
|
|
"Connector": "connector",
|
|
"Size": "size",
|
|
"Color": "color",
|
|
"PartNr": "part_number",
|
|
"OCR": "ocr_text"
|
|
}
|
|
|
|
final_result = {}
|
|
for ai_key, model_key in mapping.items():
|
|
if ai_key in result:
|
|
final_result[model_key] = result[ai_key]
|
|
elif model_key in result: # Already mapped or using model keys
|
|
final_result[model_key] = result[model_key]
|
|
|
|
# Ensure quantity and barcode are handled if returned or default
|
|
final_result["quantity"] = result.get("quantity", 1)
|
|
final_result["barcode"] = result.get("barcode", result.get("PartNr", result.get("part_number", "")))
|
|
|
|
# Handle Box mode specifically
|
|
if mode == "box":
|
|
final_result["box_label"] = result.get("box_label", result.get("name", "Unknown Box"))
|
|
final_result["name"] = final_result["box_label"]
|
|
|
|
return final_result
|
|
|
|
finally:
|
|
db.close()
|
|
|
|
# 2. Try Claude (Fallback) - Note: Mapping logic would need to be replicated here if enabled
|
|
# For now, keeping it simple
|
|
return {"error": "AI extraction failed or no data returned. Check your API key and Prompt."}
|
|
|
|
# 2. Try Claude (Fallback)
|
|
result = claude.extract(image_bytes, prompt)
|
|
if result:
|
|
return result
|
|
|
|
return {"error": "All AI providers failed or no API keys configured. Check your .env file."}
|