Update app.py

app.py

@@ -129,21 +129,26 @@ def parse_page_selection(value, num_pages: int) -> int:
         idx = int(m.group(1)) - 1 if m else 0
     return max(0, min(num_pages - 1, idx))
 
-def build_messages(history: List[Tuple[str, str]], user_text: str, images: List[Image.Image]):
+def build_messages(history_msgs: list, user_text: str, images: List[Image.Image]):
     """
-    Construct chat-format messages for tokenizer.apply_chat_template.
+    Compose the full prompt for the model:
+      - system prompt
+      - trimmed history (already in {'role','content'} format)
+      - current user turn with images + text
     """
     messages = [{"role": "system", "content": SYSTEM_PROMPT}]
-    trimmed = history[-4:] if history else []
-    for u, a in trimmed:
-        messages.append({"role": "user", "content": u})
-        messages.append({"role": "assistant", "content": a})
 
+    # Keep last 8 messages to stay lean
+    trimmed = history_msgs[-8:] if history_msgs else []
+    messages.extend(trimmed)
+
+    # Current user turn: images first, then text
     multimodal = []
     for im in images:
         multimodal.append(im)
     if user_text.strip():
         multimodal.append(user_text.strip())
+
     messages.append({"role": "user", "content": multimodal})
     return messages
 
@@ -230,19 +235,37 @@ def page_picker_changed(pages_dropdown, images_state):
     selected = images_state[idx]
     return selected, selected  # preview + selected state
 
-def chat(user_text, history, images_state, selected_img):
+def chat(user_text, history_msgs, images_state, selected_img):
     if not user_text or not user_text.strip():
-        return gr.update(), history
+        return gr.update(), history_msgs
+
     sel_img = selected_img if selected_img is not None else (images_state[0] if images_state else None)
     if sel_img is None:
-        history = history + [(user_text, "Please upload a document first.")]
-        return gr.update(value=history), history
-
-    stream = generate_reply([sel_img], user_text, history)
+        # push a system-style nudge
+        history_msgs = history_msgs + [
+            {"role": "user", "content": user_text},
+            {"role": "assistant", "content": "Please upload a document first."}
+        ]
+        return gr.update(value=history_msgs), history_msgs
+
+    # Stream the assistant reply
+    stream = generate_reply([sel_img], user_text, history_msgs)
     acc = ""
     for chunk in stream:
         acc = chunk
-        yield history + [(user_text, acc)], history + [(user_text, acc)]
+        # do incremental streaming by replacing the last assistant message
+        yield (
+            history_msgs + [
+                {"role": "user", "content": user_text},
+                {"role": "assistant", "content": acc},
+            ],
+            history_msgs + [
+                {"role": "user", "content": user_text},
+                {"role": "assistant", "content": acc},
+            ]
+        )
+
+
 
 # -----------------------------
 # App definition