Update app.py

app.py

@@ -47,82 +47,116 @@ def prompt():
         "  ],\n"
         "  \"custom_fields\": {}\n"
         "}\n\n"
-
         "SHIP FROM RULES:\n"
-        "- If header shows 'From: Name <email>' → ship_from_name, ship_from_email.\n"
-        "- If 'Origin:' exists, use that as ship_from_name.\n"
+        "- If explicit fields like 'Origin', 'Ship From' exist, extract that value.\n"
+        "- If the document is an email-style inbound notice (header block) and shows:\n"
+        "    From: Name <email>\n"
+        "  then ship_from_name = Name, ship_from_email = email.\n"
+        "- If only an email exists and no human name, set both fields to that email.\n"
+        "- If both Origin and an email sender exist, use Origin for ship_from_name and still capture the email under ship_from_email.\n"
         "- Priority: Origin → Email Name → Mill → Sender block → null.\n\n"
+        "CARRIER / EQUIPMENT RULE:\n"
+        "- If the table contains:\n"
+        "      Equipment id = <value>\n"
+        "      Mark = <value>\n"
+        "  then ALWAYS treat 'Equipment id' as the railcar number.\n"
+        "- NEVER use 'Mark' as railcar number.\n"
+        "- Carrier type must match the carrier text exactly (e.g., CHICAGO RAIL LINK).\n\n"
+        "INVENTORY RULES:\n"
+        "- Do not merge length groups. Each unique length or dimension is its own variant.\n"
+        "- Extract pcs_per_pkg, packages, pieces, fbm exactly as written.\n"
+        "- total_pcs = sum of pieces.\n"
+        "- total_fbm = sum of fbm.\n\n"
+        "TOTAL QUANTITY RULE:\n"
+        "- Use explicit totals if they appear.\n"
+        "- If no explicit total quantity appears, leave null.\n\n"
+        "CUSTOM FIELDS RULE:\n"
+        "- Capture all meaningful leftover fields not part of main schema.\n\n"
+        "Return ONLY the JSON."
+    )
 
-        "EQUIPMENT RULE:\n"
-        "- 'Equipment id' = rail_car_number. Never use 'Mark'.\n\n"
 
-        "INVENTORY RULES:\n"
-        "- Each dimension group is separate.\n"
-        "- Do NOT merge table rows.\n\n"
 
-        "Return ONLY the JSON."
+def extract_image(img):
+    """Process image input"""
+    ext = "png"  # always PNG internally from Gradio
+
+    b64 = base64.b64encode(img).decode()
+
+    content = [
+        {"type": "text", "text": prompt()},
+        {
+            "type": "image_url",
+            "image_url": {"url": f"data:image/{ext};base64,{b64}"}
+        }
+    ]
+
+    r = client.chat.completions.create(
+        model=MODEL,
+        messages=[{"role": "user", "content": content}]
     )
 
+    text = r.choices[0].message.content
+    s = text.find("{")
+    e = text.rfind("}")
+    return text[s:e+1]
+
 
-def extract(file):
+def extract_pdf(file):
+    """Process PDF input"""
     path = Path(file.name)
-    ext = path.suffix.lower()
-
-    if ext == ".pdf":
-        fid = upload_pdf(path)
-        payload = [
-            {"type": "text", "text": prompt()},
-            {"type": "file", "file": {"file_id": fid}}
-        ]
-    else:
-        b64 = base64.b64encode(path.read_bytes()).decode()
-        payload = [
-            {"type": "text", "text": prompt()},
-            {
-                "type": "image_url",
-                "image_url": {"url": f"data:image/{ext[1:]};base64,{b64}"}
-            }
-        ]
+    fid = upload_pdf(path)
+
+    content = [
+        {"type": "text", "text": prompt()},
+        {"type": "file", "file": {"file_id": fid}}
+    ]
 
     r = client.chat.completions.create(
         model=MODEL,
-        messages=[{"role": "user", "content": payload}]
+        messages=[{"role": "user", "content": content}]
     )
 
-    output = r.choices[0].message.content
-    s = output.find("{")
-    e = output.rfind("}")
-    return output[s:e+1]
+    text = r.choices[0].message.content
+    s = text.find("{")
+    e = text.rfind("}")
+    return text[s:e+1]
 
 
-def ui(file):
-    return extract(file)
+def process(image_input, pdf_input):
+    if image_input is not None:
+        return extract_image(image_input)
 
+    if pdf_input is not None:
+        return extract_pdf(pdf_input)
 
-### -----------------------------------------
-### REAL IMAGE PREVIEW (your two sample imgs)
-### -----------------------------------------
-sample_images = [
-    ["IMG_0001.jpg"],
-    ["IMG_0002.jpg"]
-]
+    return "{}"
 
 
 with gr.Blocks() as demo:
     gr.Markdown("# **Logistics OCR Data Extractor (GPT-5.1)**")
 
-    file_input = gr.File(label="Upload PDF or Image")
-    output_json = gr.JSON(label="Extracted JSON")
+    with gr.Row():
+        image_input = gr.Image(type="bytes", label="Upload Image")
+        pdf_input = gr.File(type="file", label="Upload PDF")
+
+    output = gr.JSON(label="Extracted JSON")
 
-    submit_btn = gr.Button("Submit")
+    submit = gr.Button("Submit")
 
-    submit_btn.click(fn=ui, inputs=file_input, outputs=output_json)
+    submit.click(
+        fn=process,
+        inputs=[image_input, pdf_input],
+        outputs=output
+    )
 
     gr.Examples(
-        examples=sample_images,
-        inputs=file_input,
-        label="Sample Inputs",
-        examples_per_page=2
+        examples=[
+            ["IMG_0001.jpg", None],
+            ["IMG_0002.jpg", None]
+        ],
+        inputs=[image_input, pdf_input],
+        label="Sample Images",
     )
 
 demo.launch(share=True)