Update rag/engine.py

rag/engine.py

@@ -117,6 +117,12 @@ async def agent_stream(message: str, machine_id: str, history: list[dict]):
     except Exception as e:
         print(f"Error in context expansion: {e}")
         expanded_chunks = top_reranked
+
+    # Lock context to the top document to prevent cross-document answer mixing
+    top_filename = top_reranked[0]["filename"] if top_reranked else None
+    if top_filename:
+        expanded_chunks = [c for c in expanded_chunks if c["filename"] == top_filename]
+        print(f"RAG v2: Context locked to '{top_filename}' ({len(expanded_chunks)} chunks)", flush=True)
         
     # Yield retrieved chunks metadata for frontend citation panel (use expanded_chunks so all cited pages are present)
     retrieved_event_data = []
@@ -280,6 +286,12 @@ async def agent_chat(message: str, machine_id: str, history: list[dict]) -> dict
     except Exception as e:
         print(f"[RAG Diagnostics] ERROR in context expansion: {e}", flush=True)
         expanded_chunks = top_reranked
+
+    # Lock context to the top document to prevent cross-document answer mixing
+    top_filename = top_reranked[0]["filename"] if top_reranked else None
+    if top_filename:
+        expanded_chunks = [c for c in expanded_chunks if c["filename"] == top_filename]
+        print(f"[RAG Diagnostics] Context locked to '{top_filename}' ({len(expanded_chunks)} chunks).", flush=True)
         
     # Format sources and searches
     sources = []