Build [v1.9.19]

backend/ai_vision.py

@@ -1,6 +1,5 @@
-import os
-from dotenv import load_dotenv
-from .ai import gemini, claude
+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__))
@@ -12,51 +11,74 @@ def extract_label_info(image_bytes: bytes, mode: str = "item"):
     Orchestrates extraction across multiple AI providers.
     Modes: 'item' (full technical extraction), 'box' (container discovery)
     """
-    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.
+    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)
         
-        Return ONLY a valid JSON object:
-        {
-          "box_label": "The identified container name",
-          "name": "Same as box_label",
-          "category": "Storage",
-          "specs": "Brief description if useful",
-          "quantity": 1
-        }
-        """
-    else:
-        prompt = """
-        Extract technical inventory information from this label image.
+        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()
         
-        CRITICAL INSTRUCTIONS:
-        1. Look at the most prominent text (usually top 1-2 rows). This is the product NAME and MODEL.
-        2. Extract the PART NUMBER (P/N, Model No, Type). If no explicit Part Number is found, synthesize one from the most unique identifier in the header (e.g. 'OM4-MMF-DX').
-        3. Separate the COLOR (e.g. Turquoise, Yellow, Black).
-        4. Extract CATEGORY based on the item type (e.g. Patchcord, SFP, Connector).
-        5. Extract technical SPECS (e.g. '2.0mm', '10G', '850nm').
-        
-        Return ONLY a valid JSON object:
-        {
-          "name": "Full descriptive name from header",
-          "part_number": "Unique identifier for fast scanning",
-          "category": "Broad category",
-          "color": "Color if present",
-          "specs": "Brief tech specs list",
-          "barcode": "Barcode value if visible",
-          "quantity": 1
-        }
-        """
-    
-    # 1. Try Gemini
-    result = gemini.extract(image_bytes, prompt)
-    if result:
-        # Maintenance: Ensure fields are mapped if mode was box
-        if mode == "box" and "box_label" in result and "name" not in result:
-             result["name"] = result["box_label"]
-        return result
+    # 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)