Update app.py

app.py

@@ -1,4 +1,3 @@
-
 import os
 import base64
 import torch
@@ -15,15 +14,14 @@ from io import BytesIO
 
 from transformers import (
     AutoProcessor,
-    AutoConfig,
-    Qwen3VLForConditionalGeneration
+    Qwen3VLForConditionalGeneration,
 )
 
 # ─────────────────────────────
 # CONFIG
 # ─────────────────────────────
 MODEL_REPO = "Rady10/Plant-Disease-Qwen3VL-2B"
-RAG_REPO = "Rady10/Agriculture-Rag-Data-Index"
+RAG_REPO   = "Rady10/Agriculture-Rag-Data-Index"
 
 DEVICE = "cpu"
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -31,48 +29,39 @@ os.environ["TOKENIZERS_PARALLELISM"] = "false"
 # ─────────────────────────────
 # GLOBALS
 # ─────────────────────────────
-model = None
-processor = None
+model       = None
+processor   = None
 faiss_index = None
-rag_chunks = None
-embedder = None
+rag_chunks  = None
+embedder    = None
 
 # ─────────────────────────────
-# FASTAPI APP
+# LIFESPAN — load everything once
 # ─────────────────────────────
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     global model, processor, faiss_index, rag_chunks, embedder
 
     print("Loading vision model...")
-
     processor = AutoProcessor.from_pretrained(
         MODEL_REPO,
-        trust_remote_code=True
+        trust_remote_code=True,
     )
-
     model = Qwen3VLForConditionalGeneration.from_pretrained(
         MODEL_REPO,
         torch_dtype=torch.float32,
         device_map="cpu",
-        trust_remote_code=True
+        trust_remote_code=True,
     )
-
     model.eval()
 
-    # ───── LOAD RAG ─────
-    print("Loading RAG...")
-
+    print("Loading RAG index...")
     rag_dir = snapshot_download(
         repo_id=RAG_REPO,
         repo_type="dataset",
-        local_dir="./rag"
-    )
-
-    faiss_index = faiss.read_index(
-        os.path.join(rag_dir, "agro.index")
+        local_dir="./rag",
     )
-
+    faiss_index = faiss.read_index(os.path.join(rag_dir, "agro.index"))
     with open(os.path.join(rag_dir, "chunks.json"), "r", encoding="utf-8") as f:
         rag_chunks = json.load(f)
 
@@ -81,121 +70,174 @@ async def lifespan(app: FastAPI):
     )
 
     print("ALL LOADED ✔")
-
     yield
 
 
-app = FastAPI(
-    title="🌿 Plant Disease Vision API",
-    lifespan=lifespan
-)
-
 # ─────────────────────────────
-# REQUEST MODELS
+# APP
 # ─────────────────────────────
-class VisionRequest(BaseModel):
-    image: str
-    text: str = ""
-
-class ChatRequest(BaseModel):
-    messages: list
-    image: str = None
+app = FastAPI(title="🌿 Plant Disease Chat API", lifespan=lifespan)
 
-# ─────────────────────────────
-# IMAGE DECODER
-# ─────────────────────────────
-def decode_image(base64_str):
-    img_data = base64.b64decode(base64_str)
-    return Image.open(BytesIO(img_data)).convert("RGB")
 
 # ─────────────────────────────
-# VISION GENERATION (CHAT FORMAT)
+# REQUEST MODEL
 # ─────────────────────────────
-def generate(image, text):
-
-    if not text.strip():
-        text = "What disease is shown in this plant image?"
+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
+
+
+# ─────────────────────────────
+# 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:
+    """
+    Extract the last user text, embed it, and return the top-k
+    RAG chunks joined as a single string. Returns "" if nothing found.
+    """
+    if not rag_chunks or faiss_index is None:
+        return ""
+
+    # walk backwards to find the latest user text
+    last_user_text = ""
+    for m in reversed(messages):
+        if m.get("role") != "user":
+            continue
+        content = m.get("content", "")
+        if isinstance(content, list):
+            for block in content:
+                if isinstance(block, dict) and block.get("type") == "text":
+                    last_user_text = block["text"]
+                    break
+        elif isinstance(content, str):
+            last_user_text = content
+        if last_user_text:
+            break
+
+    if not last_user_text.strip():
+        return ""
+
+    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)]
+    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(
+            "Use the following retrieved knowledge to inform your answer:\n\n"
+            + rag_context
+        )
+    system_prompt = "\n\n".join(system_parts)
 
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                {"type": "image", "image": image},
-                {"type": "text", "text": text}
-            ]
-        }
+    full_messages = [
+        {"role": "user",      "content": system_prompt},
+        {"role": "assistant", "content": "Understood. I will use this knowledge to help you."},
     ]
 
-    inputs = processor.apply_chat_template(
-        messages,
-        add_generation_prompt=True,
-        return_tensors="pt"
-    )
+    # ── copy conversation; inject image into last user turn ───
+    messages = [dict(m) for m in messages]   # shallow copy so we don't mutate input
 
-    inputs = inputs.to(model.device)
+    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
+                break
 
-    with torch.no_grad():
-        output = model.generate(
-            **inputs,
-            max_new_tokens=256,
-            temperature=0.7,
-            top_p=0.9
-        )
-
-    return processor.decode(output[0], skip_special_tokens=True)
-
-# ─────────────────────────────
-# ROUTES
-# ─────────────────────────────
-@app.get("/")
-def root():
-    return {"status": "vision api running"}
-
-@app.post("/analyze")
-def analyze(req: VisionRequest):
+        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
 
-    image = decode_image(req.image)
-    result = generate(image, req.text)
+    full_messages.extend(messages)
+    return full_messages
 
-    return {"response": result}
 
 # ─────────────────────────────
-# CHAT ENDPOINT (IMAGE + TEXT)
+# SINGLE UNIFIED ENDPOINT
 # ─────────────────────────────
 @app.post("/chat")
 def chat(req: ChatRequest):
-
-    messages = req.messages
-
-    image = None
-    if req.image:
-        image = decode_image(req.image)
-
-    # ───── inject image safely ─────
-    if image:
-        messages[-1]["content"].insert(0, {
-            "type": "image",
-            "image": image
-        })
-
-    # ───── IMPORTANT FIX HERE ─────
+    """
+    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 ─────────────────────────────────────
+    full_messages = build_full_messages(req.messages, image, rag_context)
+
+    # ── tokenise ──────────────────────────────────────────────
     inputs = processor.apply_chat_template(
-        messages,
+        full_messages,
         add_generation_prompt=True,
-        tokenize=True,        # 🔴 THIS FIXES IT
-        return_tensors="pt"
-    )
-
-    # now inputs is a tensor dict (NOT string anymore)
-
-    inputs = {k: v.to(model.device) for k, v in inputs.items()}
+        tokenize=True,
+        return_tensors="pt",
+    ).to(model.device)
 
+    # ── generate ──────────────────────────────────────────────
     with torch.no_grad():
-        output = model.generate(
+        output_ids = model.generate(
             **inputs,
-            max_new_tokens=256
+            max_new_tokens=512,
+            temperature=0.7,
+            top_p=0.9,
         )
 
+    response_text = processor.decode(output_ids[0], skip_special_tokens=True)
+
     return {
-        "response": processor.decode(output[0], skip_special_tokens=True)
-    }
+        "response":   response_text,
+        "rag_used":   bool(rag_context),
+        "image_used": image is not None,
+    }
+
+
+# ─────────────────────────────
+# HEALTH CHECK
+# ─────────────────────────────
+@app.get("/")
+def root():
+    return {"status": "plant disease chat api running"}