amer003100 commited on
Commit
1fa8476
·
1 Parent(s): 29a9e61

fix(rag): prioritize document tree root title and exclude generic headings in research queries

Browse files
Files changed (1) hide show
  1. app/core/rag_pipeline.py +16 -17
app/core/rag_pipeline.py CHANGED
@@ -637,30 +637,29 @@ class VectorlessRAG:
637
  yield f'<!--STAGE:{{"stage":"retrieved","duration":{t_retrieval:.2f},"nodes":{len(all_selected_nodes)},"docs":{len(docs)}}}-->\n'
638
 
639
  # Determine contextual document topic for web and academic queries
 
640
  doc_topic = None
641
- if all_selected_nodes:
 
 
 
 
 
 
 
 
 
 
 
642
  for n in all_selected_nodes:
643
  t = n.get("title", "")
644
  clean_t = re.sub(r'^(?:\d+[\.\)]*\s*|Abstract\b|Introduction\b)', '', t, flags=re.IGNORECASE).strip()
645
  clean_t = re.sub(r'[\(\[\{].*?[\)\]\}]', '', clean_t).strip()
646
- if len(clean_t) > 4:
647
  doc_topic = clean_t
648
  break
649
- if not doc_topic and docs:
650
- for dobj in docs:
651
- dt = (dobj.get("tree") or [{}])[0].get("title", "")
652
- if dt and len(dt) > 4 and not dt.lower().startswith("section"):
653
- doc_topic = dt
654
- break
655
- sections = dobj.get("flat_sections") or []
656
- for s in sections[:3]:
657
- st = s.get("title", "")
658
- clean_st = re.sub(r'^(?:\d+[\.\)]*\s*|Abstract\b|Introduction\b)', '', st, flags=re.IGNORECASE).strip()
659
- if len(clean_st) > 4:
660
- doc_topic = clean_st
661
- break
662
- if doc_topic:
663
- break
664
  if not doc_topic and doc_names_str and not doc_names_str.startswith("doc_"):
665
  doc_topic = re.sub(r'\.pdf$', '', doc_names_str)
666
 
 
637
  yield f'<!--STAGE:{{"stage":"retrieved","duration":{t_retrieval:.2f},"nodes":{len(all_selected_nodes)},"docs":{len(docs)}}}-->\n'
638
 
639
  # Determine contextual document topic for web and academic queries
640
+ GENERIC_HEADINGS = {"abstract", "introduction", "conclusion", "references", "overview", "background", "related work", "preliminaries", "methods"}
641
  doc_topic = None
642
+
643
+ # 1. Prefer document tree root title (which is the document title)
644
+ if docs:
645
+ for dobj in docs:
646
+ dt = (dobj.get("tree") or [{}])[0].get("title", "")
647
+ clean_dt = re.sub(r'[\(\[\{].*?[\)\]\}]', '', dt).strip()
648
+ if clean_dt and len(clean_dt) > 4 and clean_dt.lower() not in GENERIC_HEADINGS and not clean_dt.lower().startswith("section"):
649
+ doc_topic = clean_dt
650
+ break
651
+
652
+ # 2. Fallback to non-generic section titles
653
+ if not doc_topic and all_selected_nodes:
654
  for n in all_selected_nodes:
655
  t = n.get("title", "")
656
  clean_t = re.sub(r'^(?:\d+[\.\)]*\s*|Abstract\b|Introduction\b)', '', t, flags=re.IGNORECASE).strip()
657
  clean_t = re.sub(r'[\(\[\{].*?[\)\]\}]', '', clean_t).strip()
658
+ if len(clean_t) > 4 and clean_t.lower() not in GENERIC_HEADINGS:
659
  doc_topic = clean_t
660
  break
661
+
662
+ # 3. Fallback to document filename
 
 
 
 
 
 
 
 
 
 
 
 
 
663
  if not doc_topic and doc_names_str and not doc_names_str.startswith("doc_"):
664
  doc_topic = re.sub(r'\.pdf$', '', doc_names_str)
665