147 lines
6.2 KiB
Python
147 lines
6.2 KiB
Python
import os
|
|
import time
|
|
from dotenv import load_dotenv
|
|
from . import models
|
|
from .database import SessionLocal
|
|
from .ai import gemini, claude
|
|
|
|
# Note: Environment variables are managed centrally by config_loader.py
|
|
base_dir = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
class PromptManager:
|
|
"""Manages AI prompt with auto-reloading from file."""
|
|
def __init__(self, file_path: str):
|
|
self.file_path = file_path
|
|
self.last_mtime = 0
|
|
self.cached_prompt = None
|
|
|
|
def get_prompt(self) -> str:
|
|
if not os.path.exists(self.file_path):
|
|
return None
|
|
|
|
try:
|
|
current_mtime = os.path.getmtime(self.file_path)
|
|
if current_mtime > self.last_mtime or self.cached_prompt is None:
|
|
with open(self.file_path, 'r', encoding='utf-8') as f:
|
|
self.cached_prompt = f.read().strip()
|
|
self.last_mtime = current_mtime
|
|
# Use print or log for visibility in dev
|
|
print(f"🔄 AI Vision prompt reloaded from {self.file_path} (mtime: {current_mtime})")
|
|
return self.cached_prompt
|
|
except Exception as e:
|
|
print(f"⚠️ Failed to reload AI prompt: {e}")
|
|
return self.cached_prompt
|
|
|
|
# The prompt file is located in the global /config directory
|
|
# We go up one level from backend/ to reach project root, then into config/
|
|
PROJECT_ROOT = os.path.dirname(os.path.abspath(base_dir))
|
|
PROMPT_FILE_PATH = os.path.join(PROJECT_ROOT, "config", "ai_prompt.md")
|
|
prompt_mgr = PromptManager(PROMPT_FILE_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:
|
|
# 1. Try fetching from the configuration file first (SSOT)
|
|
prompt = prompt_mgr.get_prompt()
|
|
|
|
if not prompt:
|
|
# 2. Fallback to Database if file is missing
|
|
setting = db.query(models.SystemSetting).filter(models.SystemSetting.key == "ai_extraction_prompt").first()
|
|
if setting:
|
|
prompt = setting.value
|
|
else:
|
|
# 3. Final fallback to hardcoded default
|
|
prompt = "Extract technical specs. Return JSON with name, category, description, connector, size, color, part_number, ocr_text, quantity."
|
|
|
|
# 0. Get Active Provider from DB
|
|
provider_setting = db.query(models.SystemSetting).filter(models.SystemSetting.key == "ai_provider").first()
|
|
active_provider = provider_setting.value if provider_setting else "gemini"
|
|
|
|
# 1. Execute extraction based on selection
|
|
result = None
|
|
if active_provider == "claude":
|
|
print(f"📡 Using Anthropic Claude for extraction...")
|
|
result = claude.extract(image_bytes, prompt)
|
|
else:
|
|
print(f"📡 Using Google Gemini for extraction...")
|
|
result = gemini.extract(image_bytes, prompt)
|
|
|
|
if result:
|
|
# Check if AI returned a list of items or a singular object
|
|
raw_items = result.get("items") or result.get("Items")
|
|
was_list = isinstance(raw_items, list)
|
|
items_to_map = raw_items if was_list else [result]
|
|
|
|
mapping = {
|
|
"Item": "name",
|
|
"Type": "type",
|
|
"Description": "description",
|
|
"Category": "category",
|
|
"Connector": "connector",
|
|
"Size": "size",
|
|
"Color": "color",
|
|
"PartNr": "part_number",
|
|
"OCR": "ocr_text"
|
|
}
|
|
|
|
mapped_items = []
|
|
for item_data in items_to_map:
|
|
final_item = {}
|
|
for ai_key, model_key in mapping.items():
|
|
val = item_data.get(ai_key) or item_data.get(model_key)
|
|
if val and isinstance(val, str):
|
|
final_item[model_key] = val.strip()
|
|
else:
|
|
final_item[model_key] = val
|
|
|
|
# Default fields
|
|
final_item["quantity"] = item_data.get("quantity", 1)
|
|
raw_barcode = item_data.get("barcode") or item_data.get("PartNr") or item_data.get("part_number") or item_data.get("Part Number")
|
|
final_item["barcode"] = str(raw_barcode).strip() if raw_barcode else f"AI-{int(time.time()*100)}"
|
|
|
|
# Handle Box mode specifically inside mapping
|
|
if mode == "box":
|
|
final_item["box_label"] = final_item.get("box_label") or item_data.get("Box") or final_item.get("name") or "Unknown Box"
|
|
final_item["name"] = final_item["box_label"]
|
|
|
|
mapped_items.append(final_item)
|
|
|
|
# Return either the whole list wrapper or the first item (legacy compatibility)
|
|
if was_list:
|
|
return {"items": mapped_items}
|
|
return mapped_items[0] if mapped_items else {"error": "No items after mapping"}
|
|
|
|
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."}
|