Update app.py

app.py

@@ -12,12 +12,10 @@ MODEL = "gpt-5.1"
 client = OpenAI(api_key=API_KEY)
 
 
-# ---------------- PDF Upload ----------------
 def upload_pdf(path):
     return client.files.create(file=open(path, "rb"), purpose="assistants").id
 
 
-# ---------------- Prompt ----------------
 def prompt():
     return (
         "Extract structured JSON from the attached logistics document. Return ONLY valid JSON.\n"
@@ -51,83 +49,80 @@ def prompt():
         "}\n\n"
 
         "SHIP FROM RULES:\n"
-        "- If explicit fields like 'Origin', 'Ship From' exist, extract that value.\n"
-        "- If email header block exists:\n"
-        "      From: Name <email>\n"
-        "  then ship_from_name = Name, ship_from_email = email.\n"
-        "- If only email exists, set both fields to email.\n"
-        "- If both Origin and email sender exist, name = Origin and email under custom_fields.\n"
+        "- If header shows 'From: Name <email>' → ship_from_name, ship_from_email.\n"
+        "- If 'Origin:' exists, use that as ship_from_name.\n"
         "- Priority: Origin → Email Name → Mill → Sender block → null.\n\n"
 
-        "CARRIER / EQUIPMENT RULE:\n"
-        "- If table shows 'Equipment id = X' and 'Mark = Y', then X = rail_car_number.\n"
-        "- 'Mark' must never be used as the railcar number.\n\n"
+        "EQUIPMENT RULE:\n"
+        "- 'Equipment id' = rail_car_number. Never use 'Mark'.\n\n"
 
         "INVENTORY RULES:\n"
-        "- Each dimension group must remain separate.\n"
-        "- pieces_per_pkg, packages, pieces, fbm must be exact.\n"
-        "- total_pcs = sum of pieces.\n"
-        "- total_fbm = sum of fbm.\n\n"
-
-        "TOTAL QUANTITY RULE:\n"
-        "- Use explicit totals. If no total, leave null.\n\n"
-
-        "CUSTOM FIELDS RULE:\n"
-        "- Capture leftover meaningful text.\n\n"
+        "- Each dimension group is separate.\n"
+        "- Do NOT merge table rows.\n\n"
 
         "Return ONLY the JSON."
     )
 
 
-# ---------------- Extraction ----------------
 def extract(file):
     path = Path(file.name)
-    suffix = path.suffix.lower()
+    ext = path.suffix.lower()
 
-    if suffix == ".pdf":
+    if ext == ".pdf":
         fid = upload_pdf(path)
-        content = [
+        payload = [
             {"type": "text", "text": prompt()},
             {"type": "file", "file": {"file_id": fid}}
         ]
     else:
         b64 = base64.b64encode(path.read_bytes()).decode()
-        content = [
+        payload = [
             {"type": "text", "text": prompt()},
             {
                 "type": "image_url",
-                "image_url": {"url": f"data:image/{suffix[1:]};base64,{b64}"}
+                "image_url": {"url": f"data:image/{ext[1:]};base64,{b64}"}
             }
         ]
 
     r = client.chat.completions.create(
         model=MODEL,
-        messages=[{"role": "user", "content": content}]
+        messages=[{"role": "user", "content": payload}]
     )
 
-    text = r.choices[0].message.content
-    s = text.find("{")
-    e = text.rfind("}")
-    return text[s:e+1]
+    output = r.choices[0].message.content
+    s = output.find("{")
+    e = output.rfind("}")
+    return output[s:e+1]
 
 
 def ui(file):
     return extract(file)
 
 
-# ---------------- Sample Images (Preview Enabled) ----------------
-examples = [
+### -----------------------------------------
+### REAL IMAGE PREVIEW (your two sample imgs)
+### -----------------------------------------
+sample_images = [
     ["IMG_0001.jpg"],
     ["IMG_0002.jpg"]
 ]
 
 
-# ---------------- Gradio App ----------------
-gr.Interface(
-    fn=ui,
-    inputs=gr.File(label="Upload PDF or Image"),
-    outputs=gr.JSON(label="Extracted JSON"),
-    title="Logistics OCR Data Extractor (GPT-5.1)",
-    examples=examples,
-    examples_per_page=2
-).launch(share=True)
+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")
+
+    submit_btn = gr.Button("Submit")
+
+    submit_btn.click(fn=ui, inputs=file_input, outputs=output_json)
+
+    gr.Examples(
+        examples=sample_images,
+        inputs=file_input,
+        label="Sample Inputs",
+        examples_per_page=2
+    )
+
+demo.launch(share=True)