Update app.py

app.py

@@ -1,4 +1,5 @@
 import os
+import time
 import gradio as gr
 from PIL import Image
 import torch
@@ -59,19 +60,20 @@ transform = T.Compose([
 ])
 
 # ======================
-# OpenAI client
+# OpenAI client & Assistant
 # ======================
 
 OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
 if not OPENAI_API_KEY:
-    print("WARNING: OPENAI_API_KEY not set. Set it in HF Space Secrets.")
+    print("WARNING: OPENAI_API_KEY not set. Add it in HF Space Secrets.")
 
 client = OpenAI(api_key=OPENAI_API_KEY)
-OPENAI_MODEL = os.environ.get("OPENAI_MODEL", "gpt-4o-mini")
-APP_FORCE_LANG = os.environ.get("APP_FORCE_LANG", "").strip()
+
+# Assistant PPG ID (quello che mi hai dato)
+ASSISTANT_ID = "asst_20DNMEENkfBsYupFjPCwfijZ"
 
 # ======================
-# Helpers: prediction + PPG advice
+# Helpers
 # ======================
 
 def predict_image(image: Image.Image):
@@ -85,193 +87,92 @@ def predict_image(image: Image.Image):
     confidence = float(probs[idx])
     return label, confidence
 
-def zone_recommendation(zone: str, label: str):
+def call_assistant(label, confidence, zone, note, user_question, thread_id=None):
     """
-    Indicative, non-binding. Placeholder PPG-style suggestions by zone + corrosion label.
-    No warranties, research-only, always verify local specs and standards.
+    Manda input all'assistente OpenAI. Se thread_id è None, crea un nuovo thread.
+    Ritorna (reply, thread_id).
     """
-    # Surface prep suggestions (very generic)
-    prep = {
-        "light": "Mechanical St 3 or power tool clean to sound substrate; remove salts/oil; profile ~25–50 μm.",
-        "medium": "Abrasive blast to ISO 8501-1 Sa 2 or Sa 2½; profile 50–75 μm; soluble salts to low levels.",
-        "heavy": "High standard abrasive blast Sa 2½; profile 75–100 μm; thorough decontamination and edge rounding."
-    }
-
-    # Very simplified decisioning
-    zone = zone or "Other / Not sure"
-    base = []
-    prep_note = prep["medium"]
-
-    if zone == "Below waterline (hull/AF area)":
-        base.append("Primer: epoxy anticorrosive (e.g., Sigmaprime 200).")
-        base.append("Tie-coat: epoxy tie-coat where required.")
-        base.append("AF: Ecofleet 530 or silicone-based (e.g., SigmaGlide 2390) depending on performance needs.")
-        prep_note = prep["medium"]
-
-    elif zone == "Waterline / Splash zone":
-        base.append("High-build epoxy barrier (e.g., Sigmashield 1200) + polyurethane finish (Sigmadur 550) above WL.")
-        base.append("Consider elastomeric or glass-flake epoxy in severe splash.")
-        prep_note = prep["heavy"]
-
-    elif zone == "Topsides / Boot-top":
-        base.append("Epoxy build (Sigmacover 456) + polyurethane finish (Sigmadur 550).")
-        prep_note = prep["medium"]
-
-    elif zone == "Deck / Weather deck":
-        base.append("Abrasion-resistant epoxy (Sigmashield 1200) + non-skid where needed + PU topcoat.")
-        prep_note = prep["heavy"]
-
-    elif zone == "Superstructure / Accommodation":
-        base.append("Epoxy primer (Sigmacover 280/456) + PU finish (Sigmadur 550 / One).")
-        prep_note = prep["light"]
-
-    elif zone == "Ballast tanks (immersed)":
-        base.append("Certified tank lining epoxy system (consult tank-lining spec).")
-        prep_note = prep["heavy"]
-
-    elif zone == "Cargo holds / Dry bulk":
-        base.append("Impact/abrasion resistant epoxy lining; verify cargo compatibility.")
-        prep_note = prep["heavy"]
-
-    elif zone == "Engine room / Hot surfaces":
-        base.append("Heat-resistant coating per temperature class; check T-ratings.")
-        prep_note = prep["medium"]
-
-    elif zone == "Pipes / Under insulation (UIC/CUI)":
-        base.append("CUI-rated epoxy/novolac system; consider high temp and cyclic wet/dry.")
-        prep_note = prep["heavy"]
-
-    else:
-        base.append("Generic: epoxy primer + epoxy build + PU finish. Validate per environment.")
-        prep_note = prep["medium"]
-
-    # Tiny tweak by corrosion label for nuance
-    extra = []
-    if label in ["pitting_corrosion", "crevice_corrosion"]:
-        extra.append("Pay attention to edges, welds, crevices; stripe coat before full coats.")
-    if label in ["galvanic_corrosion"]:
-        extra.append("Isolate dissimilar metals; ensure dielectric barrier.")
-    if label in ["mic_corrosion"]:
-        extra.append("Address microbiological contamination; thorough cleaning/sanitization.")
-    if label in ["under_insulation_corrosion"]:
-        extra.append("Remove/inspect insulation; ensure dry, sealed design before recoating.")
-
-    rec = "- " + "\n- ".join(base + extra) if base else "Consult dedicated spec."
-    return rec, prep_note
-
-# ======================
-# OpenAI prompts
-# ======================
-
-def build_dual_lang_prompt(label: str, confidence: float, user_note: str):
-    # Second language selection logic
-    if APP_FORCE_LANG:
-        second_lang_clause = f"Then provide the same content in {APP_FORCE_LANG}."
-    else:
-        if user_note and user_note.strip():
-            second_lang_clause = "Then repeat the same content in the same language used by the user note."
-        else:
-            second_lang_clause = "Then repeat the same content in Italian."
-
-    return f"""
-You are a marine coatings technical assistant (PPG style). Image classification:
-- Corrosion type: {label}
-- Model confidence: {round(confidence*100,2)}%
-
-Provide a technical advisory in two parts:
-
-1) English section (concise, bullets):
-- Diagnosis
-- Likely causes
-- Generic paint system suggestion (no unrealistic promises)
-- Warnings about substrate condition / surface preparation
+    # se non esiste un thread, crealo
+    if not thread_id:
+        thread = client.beta.threads.create()
+        thread_id = thread.id
+
+    # messaggio utente con contesto
+    context = f"""
+Classification result: {label} ({round(confidence*100,2)}%).
+Zone: {zone or "Not specified"}.
+Note: {note or "(none)"}.
+
+User question:
+{user_question}
+
+Always act as a PPG marine coatings technical specialist.
+Keep answers concise, technical, and always add:
+"Research use only; verify with official PPG specs."
+"""
+    client.beta.threads.messages.create(
+        thread_id=thread_id,
+        role="user",
+        content=context
+    )
 
-2) {second_lang_clause}
+    # run con l'assistente
+    run = client.beta.threads.runs.create(
+        thread_id=thread_id,
+        assistant_id=ASSISTANT_ID
+    )
 
-User note (optional): {user_note or "(none)"}
+    # polling fino a completamento
+    while True:
+        r = client.beta.threads.runs.retrieve(thread_id=thread_id, run_id=run.id)
+        if r.status in ["completed", "failed", "cancelled", "expired"]:
+            break
+        time.sleep(0.8)
 
-Be pragmatic and brief.
-"""
+    # recupera ultimo messaggio
+    msgs = client.beta.threads.messages.list(thread_id=thread_id)
+    reply = None
+    for m in msgs.data:
+        if m.role == "assistant":
+            reply = m.content[0].text.value
+            break
 
-def openai_answer(messages):
-    """Call OpenAI with a list of messages [{'role': 'user'|'assistant'|'system', 'content': str}, ...]."""
-    try:
-        resp = client.chat.completions.create(
-            model=OPENAI_MODEL,
-            messages=messages,
-            temperature=0.3
-        )
-        return resp.choices[0].message.content.strip()
-    except Exception as e:
-        return f"OpenAI error: {e}"
+    return reply or "No reply from Assistant.", thread_id
 
 # ======================
-# Pipelines: analysis + chat
+# Pipelines
 # ======================
 
-def run_analysis(image, note, zone, history):
+def run_analysis(image, note, zone, chat_history, thread_state):
     if image is None:
-        return "No image received.", history, {"label": None, "confidence": 0.0, "zone": zone or ""}
+        return "No image received.", chat_history, thread_state
 
     label, conf = predict_image(image)
+    user_question = "Provide initial advisory based on classification and note."
+    reply, thread_id = call_assistant(label, conf, zone, note, user_question)
 
-    # Dual-language advisory
-    initial_prompt = build_dual_lang_prompt(label, conf, note or "")
-    advisory = openai_answer([{"role": "user", "content": initial_prompt}])
-
-    # Zone-based recommendation block
-    rec, prep_note = zone_recommendation(zone or "", label)
     header = f"**Model result:** `{label}` — confidence **{round(conf*100,2)}%**\n\n"
-    zone_blk = f"### Zone-based indicative advice\n- **Zone:** {zone or 'Not specified'}\n- **Prep:** {prep_note}\n{rec}\n\n"
-    disclaimer = "> **Disclaimer:** Research & experimental use only. No warranty. Validate all advice with certified specs and local regulations.\n"
-
-    # UI markdown
-    out_text = header + advisory + "\n\n" + zone_blk + disclaimer
-
-    # Initialize/seed chat with the first assistant message (advisory)
-    new_history = history or []
-    new_history.append(("", f"{advisory}\n\n{zone_blk}\n{disclaimer}"))
-
-    # Save last result in a dict state
-    result_state = {"label": label, "confidence": conf, "zone": zone or ""}
-
-    return out_text, new_history, result_state
-
-def continue_chat(user_msg, history, result_state, user_note, zone):
-    if not user_msg or not user_msg.strip():
-        return history, ""
+    out_text = header + (reply or "")
 
-    label = (result_state or {}).get("label") or "unknown"
-    conf = (result_state or {}).get("confidence") or 0.0
-    current_zone = zone or (result_state or {}).get("zone") or "Not specified"
+    new_history = chat_history or []
+    new_history.append(("", reply))
 
-    # Provide a focused system prompt to keep bot on-topic
-    system_prompt = f"""
-You are a corrosion & marine coatings assistant. Stay on topic (corrosion, surface prep, PPG-style paint systems).
-Context:
-- Last classification: {label} (confidence {round(conf*100,2)}%)
-- Zone: {current_zone}
-- Note: {user_note or "(none)"}
+    return out_text, new_history, {"thread_id": thread_id, "label": label, "confidence": conf, "zone": zone or ""}
 
-When recommending, be indicative and short, and include surface preparation standards where relevant (e.g., ISO 8501-1).
-Always include a short safety/legal disclaimer at the end: "Research use only; verify with official specs."
-"""
-
-    # Convert history (list[tuple(user, assistant)]) to messages
-    msgs = [{"role": "system", "content": system_prompt}]
-    for u, a in history:
-        if u:
-            msgs.append({"role": "user", "content": u})
-        if a:
-            msgs.append({"role": "assistant", "content": a})
+def continue_chat(user_msg, chat_history, thread_state, note, zone):
+    if not user_msg.strip():
+        return chat_history, ""
 
-    msgs.append({"role": "user", "content": user_msg})
+    label = (thread_state or {}).get("label") or "unknown"
+    conf = (thread_state or {}).get("confidence") or 0.0
+    current_zone = zone or (thread_state or {}).get("zone") or "Not specified"
+    thread_id = (thread_state or {}).get("thread_id")
 
-    reply = openai_answer(msgs)
+    reply, thread_id = call_assistant(label, conf, current_zone, note, user_msg, thread_id)
 
-    # Update history
-    history.append((user_msg, reply))
-    return history, ""  # clear textbox
+    chat_history.append((user_msg, reply))
+    thread_state["thread_id"] = thread_id
+    return chat_history, ""
 
 # ======================
 # UI
@@ -280,12 +181,12 @@ Always include a short safety/legal disclaimer at the end: "Research use only; v
 WELCOME = """
 # Corrosion Assistant — Beta
 
-**Welcome!** This demo runs a custom **ResNet50 corrosion classifier** and a chat assistant on top.
-- **Model**: ResNet50 classifier, **trained locally** on **~9,000 images**.
-- **Data collection**: a public **contribution link** for new images will open **soon**.
-- **Scope**: research & experimental only. No professional advice, no warranty.
+**Welcome!** This demo runs a custom **ResNet50 corrosion classifier** and connects to a dedicated **PPG Assistant** on OpenAI.
+- **Model**: ResNet50 classifier, **trained locally** on ~9,000 images  
+- **Data collection**: a public link for contributing images will open soon  
+- **Disclaimer**: research & experimental only. No professional advice, no warranty.  
 
-> This interface first returns an English advisory, then repeats in your input language (or Italian if none).
+After image analysis you can continue chatting with the assistant.
 """
 
 with gr.Blocks(title="Corrosion Assistant", theme=gr.themes.Soft()) as demo:
@@ -294,46 +195,43 @@ with gr.Blocks(title="Corrosion Assistant", theme=gr.themes.Soft()) as demo:
     with gr.Row():
         with gr.Column(scale=2):
             img = gr.Image(type="pil", sources=["upload","webcam"], label="Upload or webcam")
-            note = gr.Textbox(label="Notes / Context (optional)", placeholder="Write in your language. EN comes first, then your language.")
+            note = gr.Textbox(label="Notes / Context (optional)", placeholder="Write in your language.")
             zone = gr.Dropdown(choices=ZONES, label="Zone (indicative)", value="Other / Not sure")
             analyze_btn = gr.Button("Analyze image", variant="primary")
         with gr.Column(scale=3):
             out_md = gr.Markdown(label="Analysis")
 
-    gr.Markdown(
-        "### Continue the conversation\nAsk follow-up questions. The assistant will use the last classification and the selected zone."
-    )
+    gr.Markdown("### Continue the conversation with the PPG Assistant")
 
     with gr.Row():
         with gr.Column(scale=3):
             chat = gr.Chatbot(height=320, label="Advisor chat")
-            chat_in = gr.Textbox(label="Type your message", placeholder="Ask about prep, products, recoats, etc.")
+            chat_in = gr.Textbox(label="Your message", placeholder="Ask about prep, products, recoats, etc.")
             send_btn = gr.Button("Send")
             clear_btn = gr.Button("Clear chat")
         with gr.Column(scale=2):
             gr.Markdown(
-                "> **Disclaimer:** Research & experimental use only. Validate with certified specifications, class rules and local regulations. "
+                "> **Disclaimer:** Research & experimental use only. Validate with certified PPG specs and local regulations. "
                 "No professional advice is provided and no responsibility is assumed."
             )
 
     # States
-    chat_state = gr.State([])         # list of (user, assistant)
-    result_state = gr.State({"label": None, "confidence": 0.0, "zone": ""})
+    chat_state = gr.State([])
+    thread_state = gr.State({"thread_id": None, "label": None, "confidence": 0.0, "zone": ""})
 
-    # Wire up
     analyze_btn.click(
         fn=run_analysis,
-        inputs=[img, note, zone, chat_state],
-        outputs=[out_md, chat_state, result_state]
+        inputs=[img, note, zone, chat_state, thread_state],
+        outputs=[out_md, chat_state, thread_state]
     ).then(
-        lambda h: h,  # reflect chat state to UI
+        lambda h: h,
         inputs=[chat_state],
         outputs=[chat]
     )
 
     send_btn.click(
         fn=continue_chat,
-        inputs=[chat_in, chat_state, result_state, note, zone],
+        inputs=[chat_in, chat_state, thread_state, note, zone],
         outputs=[chat, chat_in]
     )
 
@@ -343,13 +241,8 @@ with gr.Blocks(title="Corrosion Assistant", theme=gr.themes.Soft()) as demo:
         outputs=[chat, chat_in]
     )
 
-    gr.Markdown(
-        "<small>Tip: keep photos under ~2MB. Secrets are stored server-side. "
-        "The HF outer shell cannot be customized; this inner Gradio UI can.</small>"
-    )
-
-    # Programmatic API for /api/predict etc.
     demo.api_mode = "enabled"
 
 if __name__ == "__main__":
     demo.launch()
+