Update app.py

app.py

@@ -36,7 +36,7 @@ rag_chunks  = None
 embedder    = None
 
 # ─────────────────────────────
-# LIFESPAN — load everything once
+# LIFESPAN
 # ─────────────────────────────
 @asynccontextmanager
 async def lifespan(app: FastAPI):
@@ -83,29 +83,41 @@ app = FastAPI(title="🌿 Plant Disease Chat API", lifespan=lifespan)
 # REQUEST MODEL
 # ─────────────────────────────
 class ChatRequest(BaseModel):
-    messages: list          # full conversation history in OpenAI-style format
-    image: str = None       # base64-encoded image (optional)
-    use_rag: bool = True    # set False to skip RAG retrieval
+    messages: list
+    image: str = None
+    # image present → RAG skipped automatically
 
 
 # ─────────────────────────────
 # HELPERS
 # ─────────────────────────────
 def decode_image(base64_str: str) -> Image.Image:
-    """Decode a base64 string into a PIL RGB image."""
     img_bytes = base64.b64decode(base64_str)
     return Image.open(BytesIO(img_bytes)).convert("RGB")
 
 
-def retrieve_rag_context(messages: list, k: int = 3) -> str:
+def chunk_to_text(chunk) -> str:
     """
-    Extract the last user text, embed it, and return the top-k
-    RAG chunks joined as a single string. Returns "" if nothing found.
+    Safely convert a chunk to plain string regardless of its type.
+    chunks.json may contain strings, dicts, or other structures.
     """
+    if isinstance(chunk, str):
+        return chunk
+    if isinstance(chunk, dict):
+        # common keys used in RAG datasets — try in order
+        for key in ("text", "content", "passage", "chunk", "body"):
+            if key in chunk and isinstance(chunk[key], str):
+                return chunk[key]
+        # fallback: join all string values
+        return " ".join(str(v) for v in chunk.values())
+    return str(chunk)
+
+
+def retrieve_rag_context(messages: list, k: int = 3) -> str:
     if not rag_chunks or faiss_index is None:
         return ""
 
-    # walk backwards to find the latest user text
+    # find last user text
     last_user_text = ""
     for m in reversed(messages):
         if m.get("role") != "user":
@@ -126,18 +138,16 @@ def retrieve_rag_context(messages: list, k: int = 3) -> str:
 
     query_vec = embedder.encode([last_user_text])
     _, indices = faiss_index.search(query_vec, k=k)
-    chunks = [rag_chunks[i] for i in indices[0] if i < len(rag_chunks)]
+
+    chunks = [
+        chunk_to_text(rag_chunks[i])
+        for i in indices[0]
+        if i < len(rag_chunks)
+    ]
     return "\n\n".join(chunks)
 
 
 def build_full_messages(messages: list, image: Image.Image, rag_context: str) -> list:
-    """
-    Combine system prompt (RAG context), conversation history, and optional
-    image into a single message list ready for apply_chat_template.
-    """
-    # ── system as a fake user/assistant pair ──────────────────
-    # Qwen3VL's apply_chat_template does not support a 'system' role,
-    # so we simulate it with a leading exchange.
     system_parts = ["You are a plant disease expert assistant."]
     if rag_context:
         system_parts.append(
@@ -151,65 +161,34 @@ def build_full_messages(messages: list, image: Image.Image, rag_context: str) ->
         {"role": "assistant", "content": "Understood. I will use this knowledge to help you."},
     ]
 
-    # ── copy conversation; inject image into last user turn ───
-    messages = [dict(m) for m in messages]   # shallow copy so we don't mutate input
+    messages = [dict(m) for m in messages]
 
     if image is not None:
-        last_user_idx = None
         for i in range(len(messages) - 1, -1, -1):
             if messages[i].get("role") == "user":
-                last_user_idx = i
+                content = messages[i].get("content", "")
+                if isinstance(content, str):
+                    content = [{"type": "text", "text": content}]
+                content = [{"type": "image", "image": image}] + content
+                messages[i]["content"] = content
                 break
 
-        if last_user_idx is not None:
-            content = messages[last_user_idx].get("content", "")
-            if isinstance(content, str):
-                content = [{"type": "text", "text": content}]
-            # prepend image block
-            content = [{"type": "image", "image": image}] + content
-            messages[last_user_idx]["content"] = content
-
     full_messages.extend(messages)
     return full_messages
 
 
 # ─────────────────────────────
-# SINGLE UNIFIED ENDPOINT
+# UNIFIED ENDPOINT
 # ─────────────────────────────
 @app.post("/chat")
 def chat(req: ChatRequest):
-    """
-    Unified chat endpoint. Handles three modes transparently:
-
-      1. RAG only     — pass messages, use_rag=true, no image
-      2. Image only   — pass messages + image, use_rag=false
-      3. Image + RAG  — pass messages + image, use_rag=true (default)
-
-    Request body
-    ────────────
-    messages  : list of {"role": "user"|"assistant", "content": str | list}
-    image     : base64-encoded image string (optional)
-    use_rag   : bool, default true
-
-    Response
-    ────────
-    {
-      "response"  : str,
-      "rag_used"  : bool,
-      "image_used": bool
-    }
-    """
-
-    # ── decode image ──────────────────────────────────────────
     image = decode_image(req.image) if req.image else None
 
-    # ── RAG retrieval ─────────────────────────────────────────
-    rag_context = retrieve_rag_context(req.messages) if req.use_rag else ""
-
-    # ── assemble messages ─────────────────────────────────────
+    # image present → use model's own vision training only (no RAG)
+    # no image      → use RAG to ground the text answer
+    rag_context = "" if image else retrieve_rag_context(req.messages)
     full_messages = build_full_messages(req.messages, image, rag_context)
 
-    # ── tokenise ──────────────────────────────────────────────
     inputs = processor.apply_chat_template(
         full_messages,
         add_generation_prompt=True,
@@ -217,7 +196,6 @@ def chat(req: ChatRequest):
         return_tensors="pt",
     ).to(model.device)
 
-    # ── generate ──────────────────────────────────────────────
     with torch.no_grad():
         output_ids = model.generate(
             **inputs,