Numeric citation references [N, p. X] with reference list

app.py

@@ -5,7 +5,13 @@
 import chromadb
 
 from src.config import ConfigError, load_config
-from src.generation.answerer import answer, build_source_elements
+from src.generation.answerer import (
+    answer,
+    build_ref_map,
+    build_reference_list,
+    build_source_elements,
+    strip_llm_references,
+)
 from src.generation.condenser import Condenser
 from src.health_check import check_models, check_ollama
 from src.ingestion.ingest import IngestResult, ingest_folder
@@ -108,6 +114,7 @@ async def on_chat_start():
     condenser = Condenser(model=config.models.llm)
     cl.user_session.set("condenser", condenser)
     cl.user_session.set("chat_history", [])
+    cl.user_session.set("ref_map", {})
 
     # Set up Chainlit settings panel
     settings = await cl.ChatSettings(
@@ -280,18 +287,32 @@ async def on_message(message: cl.Message):
 
         retrieval_step.output = f"Retrieved {len(results)} relevant chunks"
 
+    # Build ref_map for this answer, extending the session-wide map
+    existing_ref_map = cl.user_session.get("ref_map")
+    ref_map = build_ref_map(results, existing_ref_map)
+    cl.user_session.set("ref_map", ref_map)
+
+    # Collect which reference numbers are used in this answer
+    used_nums = {ref_map[r.metadata.get("relative_path", r.metadata.get("filename", "unknown"))] for r in results}
+
     # Stream the answer token by token
     msg = cl.Message(content="")
     async for token in answer(
         query,
         results,
         model=config.models.llm,
+        ref_map=ref_map,
     ):
         await msg.stream_token(token)
+
+    # Strip any reference list the LLM generated, then append ours
+    msg.content = strip_llm_references(msg.content)
+    ref_list = build_reference_list(ref_map, only=used_nums)
+    msg.content += f"\n\n---\n**References:**\n{ref_list}"
     await msg.send()
 
     # Attach expandable source chunks below the answer
-    for el_data in build_source_elements(results):
+    for el_data in build_source_elements(results, ref_map=ref_map):
         element = cl.Text(
             name=el_data["name"],
             content=el_data["content"],

config.yaml

@@ -11,7 +11,7 @@ retrieval:
   semantic_top_k: 20
   rrf_k: 60
   rerank_top_k: 10
-  reranker_model: "bge-reranker-v2-m3"
+  reranker_model: "BAAI/bge-reranker-v2-m3"
 
 paths:
   chroma_db: "~/.multi_doc_query/chroma_db/"

src/config.py

@@ -23,7 +23,7 @@ class RetrievalConfig(BaseModel):
     semantic_top_k: int = 20
     rrf_k: int = 60
     rerank_top_k: int = 10
-    reranker_model: str = "bge-reranker-v2-m3"
+    reranker_model: str = "BAAI/bge-reranker-v2-m3"
 
 
 class PathsConfig(BaseModel):

src/generation/answerer.py

@@ -1,8 +1,9 @@
+import re
 from collections.abc import AsyncIterator
 
 import ollama
 
-from src.retrieval.vector_store import SearchResult
+from src.models import SearchResult
 
 # --- Not via LangChain: prompt construction and LLM calls use direct
 #     ollama.chat() for full control over streaming and prompt format. ---
@@ -11,56 +12,146 @@
     "You are a helpful assistant that answers questions based on the provided "
     "document sources. Follow these rules:\n"
     "1. Base your answer only on the provided sources.\n"
-    "2. Cite sources inline using the exact source label shown in the "
-    "headers, e.g. [filename.pdf, p. 12]. Always use the filename from "
-    "the source header, never invent labels like 'Excerpt 1'.\n"
+    "2. Cite sources inline using numeric references, e.g. [1, p. 12] or "
+    "[2, Section: Methods]. Use the reference numbers shown in the source "
+    "headers and the Reference List. If citing the same document on different "
+    "pages, reuse its number, e.g. [1, p. 5] and [1, p. 22]. Do NOT "
+    "reproduce the Reference List — it will be appended automatically.\n"
     "3. If sources from different documents conflict, highlight the "
     "discrepancy and cite both sources.\n"
     "4. If no sources are relevant to the question, say so clearly."
 )
 
 
-def build_prompt(question: str, results: list[SearchResult]) -> list[dict]:
-    """Build chat messages for Ollama from a question and search results."""
+def build_ref_map(
+    results: list[SearchResult],
+    existing: dict[str, int] | None = None,
+) -> dict[str, int]:
+    """Assign a stable numeric ID to each unique document path.
+
+    Returns a new dict mapping document path to its reference number (1-based).
+    If *existing* is provided, those mappings are preserved and new documents
+    get numbers starting after the current maximum.
+    """
+    ref_map = dict(existing) if existing else {}
+    next_num = max(ref_map.values(), default=0) + 1
+    for r in results:
+        path = _doc_path(r.metadata)
+        if path not in ref_map:
+            ref_map[path] = next_num
+            next_num += 1
+    return ref_map
+
+
+def build_reference_list(
+    ref_map: dict[str, int],
+    only: set[int] | None = None,
+) -> str:
+    """Format a numbered reference list string.
+
+    If *only* is given, include only the reference numbers in that set.
+    """
+    lines = []
+    for path, num in sorted(ref_map.items(), key=lambda x: x[1]):
+        if only is not None and num not in only:
+            continue
+        lines.append(f"[{num}] {path}")
+    return "\n".join(lines)
+
+
+_REFS_TAIL_RE = re.compile(
+    r"\n*(?:---\n)?"          # optional horizontal rule
+    r"\*{0,2}"                # optional bold markers
+    r"[Rr]eferences?:?"      # "Reference:", "References:", bold variants
+    r"\*{0,2}"                # closing bold markers
+    r"\n.*",                  # everything after
+    re.DOTALL,
+)
+
+
+def strip_llm_references(text: str) -> str:
+    """Remove any trailing reference/references section the LLM generated."""
+    return _REFS_TAIL_RE.sub("", text).rstrip()
+
+
+def _doc_path(metadata: dict[str, str | int]) -> str:
+    """Get the document path from metadata, preferring relative_path."""
+    return metadata.get("relative_path", metadata.get("filename", "unknown"))
+
+
+def _source_label(metadata: dict[str, str | int], ref_num: int) -> str:
+    """Build the inline source label for a context header, e.g. [1, p. 5]."""
+    section = metadata.get("section_header")
+    if section:
+        return f"[{ref_num}, Section: {section}]"
+    page = metadata.get("page_number", "?")
+    return f"[{ref_num}, p. {page}]"
+
+
+def _element_name(metadata: dict[str, str | int], ref_num: int) -> str:
+    """Build a display name for a source element, e.g. [1] path, p. 5."""
+    path = _doc_path(metadata)
+    section = metadata.get("section_header")
+    if section:
+        return f"[{ref_num}] {path}, Section: {section}"
+    page = metadata.get("page_number", "?")
+    return f"[{ref_num}] {path}, p. {page}"
+
+
+def build_prompt(
+    question: str,
+    results: list[SearchResult],
+    *,
+    ref_map: dict[str, int] | None = None,
+) -> list[dict]:
+    """Build chat messages for Ollama from a question and search results.
+
+    If *ref_map* is provided, it is used for numbering; otherwise a fresh
+    map is built from *results*.
+    """
+    if ref_map is None:
+        ref_map = build_ref_map(results)
+
     context_parts = []
     for r in results:
-        context_parts.append(
-            f"--- {_source_name(r.metadata)} ---\n{r.text}"
-        )
+        ref_num = ref_map[_doc_path(r.metadata)]
+        label = _source_label(r.metadata, ref_num)
+        context_parts.append(f"--- {label} ---\n{r.text}")
 
     context = "\n\n".join(context_parts)
 
+    # Include reference list in prompt so LLM knows the mapping
+    ref_list = build_reference_list(ref_map)
+
     return [
         {"role": "system", "content": SYSTEM_PROMPT},
         {
             "role": "user",
             "content": (
                 f"Document excerpts:\n\n{context}\n\n"
+                f"Reference List:\n{ref_list}\n\n"
                 f"Question: {question}"
             ),
         },
     ]
 
 
-def _source_name(metadata: dict[str, str | int]) -> str:
-    """Build a display name for a source from its metadata."""
-    path = metadata.get("relative_path", metadata.get("filename", "unknown"))
-    section = metadata.get("section_header")
-    if section:
-        return f"Source: {path} | Section: {section}"
-    page = metadata.get("page_number", "?")
-    return f"Source: {path} | Page {page}"
-
-
-def build_source_elements(results: list[SearchResult]) -> list[dict]:
+def build_source_elements(
+    results: list[SearchResult],
+    *,
+    ref_map: dict[str, int] | None = None,
+) -> list[dict]:
     """Build source element data for Chainlit display.
 
     Returns a list of dicts with keys: name, content, display.
     Order matches the input (relevance-ranked by caller).
     """
+    if ref_map is None:
+        ref_map = build_ref_map(results)
+
     return [
         {
-            "name": _source_name(r.metadata),
+            "name": _element_name(r.metadata, ref_map[_doc_path(r.metadata)]),
             "content": r.text,
             "display": "side",
         }
@@ -73,9 +164,10 @@ async def answer(
     results: list[SearchResult],
     *,
     model: str = "llama3.1:8b",
+    ref_map: dict[str, int] | None = None,
 ) -> AsyncIterator[str]:
     """Stream answer tokens from Ollama."""
-    messages = build_prompt(question, results)
+    messages = build_prompt(question, results, ref_map=ref_map)
     stream = ollama.chat(model=model, messages=messages, stream=True)
     for chunk in stream:
         token = chunk["message"]["content"]