Update email_rag/rag_retrieval.py

email_rag/rag_retrieval.py

@@ -17,7 +17,7 @@ EMAIL_PAT = re.compile(r"\b[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}\b")
 AMOUNT_PAT = re.compile(r"\b(?:\$|USD\s*)?\d{1,3}(?:,\d{3})*(?:\.\d+)?\b")
 DATE_PAT = re.compile(r"\b\d{1,2}/\d{1,2}/\d{2,4}\b")  # very simple date pattern
 
-
+# making query more richer, this help both BM25 and the embedding model to see thread_ids and important names.
 def rewrite_query(user_text: str, session: dict) -> str:
     """
     Rewrite user query by injecting thread ID and a light summary
@@ -28,6 +28,7 @@ def rewrite_query(user_text: str, session: dict) -> str:
 
     key_bits = []
 
+    #For each entity type (people, files, amounts, dates), Take up to first 3 values.
     people = mem.get("people") or []
     if people:
         key_bits.append(f"people: {', '.join(people[:3])}")
@@ -50,7 +51,7 @@ def rewrite_query(user_text: str, session: dict) -> str:
 
     return f"In thread {tid}, {context_str}answer this question: {user_text}"
 
-
+# Tokenize rewrite querry, give one score per chunks
 def retrieve_chunks(rewrite: str, session: dict, search_outside_thread: bool):
     """
     Hybrid retrieval: BM25 + semantic similarity over precomputed embeddings.
@@ -60,7 +61,7 @@ def retrieve_chunks(rewrite: str, session: dict, search_outside_thread: bool):
 
     # Semantic query vector
     q_vec = sem_model.encode([rewrite], normalize_embeddings=True)[0]  # (D,)
-    sem_scores = embeddings @ q_vec  # cosine similarity
+    sem_scores = embeddings @ q_vec  # cosine similarity, dot product with every chunk vector
 
     # Normalize to [0,1]
     bm25_norm = bm25_scores / bm25_scores.max() if bm25_scores.max() > 0 else bm25_scores
@@ -77,12 +78,13 @@ def retrieve_chunks(rewrite: str, session: dict, search_outside_thread: bool):
         bm25_norm = bm25_norm[mask]
         sem_norm = sem_norm[mask]
 
-    combined = 0.6 * bm25_norm + 0.4 * sem_norm
-    order = np.argsort(-combined)
+    combined = 0.6 * bm25_norm + 0.4 * sem_norm # 60% BM25, 40% semantic
+    order = np.argsort(-combined) # indices sorted descending by combined score
 
     top_k = 8
     top_indices = indices[order[:top_k]]
 
+    #For each top chunk: Copy all key metadata, Add Score
     retrieved = []
     for local_rank, idx in enumerate(top_indices):
         c = chunks[idx]
@@ -116,8 +118,9 @@ def build_answer(user_text: str, rewrite: str, retrieved):
             "I couldn’t find any emails or content in this thread that clearly answer your question.",
             []
         )
-
-    # ---- Heuristic: check scores + keyword overlap ----
+    
+    # If the retrieved chunks are low-scoring or don’t share keywords with the question, the system refuses to guess and returns a polite “no clear answer” instead of hallucinating
+    # Heuristic: check scores + keyword overlap 
     question_tokens = {t.lower() for t in user_text.split() if len(t) > 3}
 
     def snippet_has_overlap(snippet: str) -> bool:
@@ -177,8 +180,10 @@ def build_answer(user_text: str, rewrite: str, retrieved):
         page_no = r.get("page_no")
         snippet = r["text"].replace("\n", " ")
         snippet = (snippet[:300] + "…") if len(snippet) > 300 else snippet
-
+        # Shorten each chunk to a 300-character snippet
+        
         key = (msg_id, snippet)
+        # avoid printing the same snippet twice
         if key in seen:
             continue
         seen.add(key)
@@ -215,7 +220,7 @@ def extract_entities_for_turn(user_text: str, retrieved):
     amounts = set()
     dates = set()
 
-    # from/to emails are good 'people' proxies
+    # Extracts email addresses and adds them to people (from_addr <-> to_addr)
     for r in retrieved:
         for field in ("from_addr", "to_addr"):
             val = r.get(field)
@@ -245,7 +250,7 @@ def extract_entities_for_turn(user_text: str, retrieved):
     entities = {k: v for k, v in entities.items() if v}
     return entities
 
-
+#Logs every interaction (query, rewrite, retrieved chunks, answer, citations) into runs/trace.jsonl for evaluation and debugging.
 def log_trace(session_id: str, user_text: str, rewrite: str, retrieved, answer, citations):
     trace_path = RUNS_DIR / "trace.jsonl"