sofhiaazzhr Claude Opus 5.5 commited on
Commit
6f6628d
·
1 Parent(s): 155a3cf

[NOTICKET] fix(retrieval): no page citation for TXT documents

Browse files

- New helper `citable_page(data)` in retrieval/base.py: the chunk's page to show, or None when there is no page or the file type has no real pages
- TXT ingest stores every chunk as page_label 1 (ingest_v2.chunks_from_text), so without this the chip and the answer would always say "p. 1"
- Used for the traceability chip label, the chunk labels sent to the answer model, and the planner's document context; PDF and DOCX unchanged
- Raw `page_label` stays in the traceability payload; file type falls back to the filename extension
- Local tests 12/12 for this change, ruff clean; the 6 failing unit tests fail the same way before this change (stale retriever fake, one planner validator test)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

src/agents/chat_handler.py CHANGED
@@ -37,7 +37,7 @@ from typing import TYPE_CHECKING, Any
37
  from langchain_core.messages import BaseMessage
38
 
39
  from src.middlewares.logging import get_logger
40
- from src.retrieval.base import RetrievalResult
41
  from src.traceability import TraceabilityScratchpad, TraceabilityToolInvoker
42
 
43
  from ..charts import chart_envelopes
@@ -1080,8 +1080,9 @@ def _render_document_context(results: list[RetrievalResult]) -> str:
1080
  for r in results:
1081
  data = r.metadata.get("data", {}) if isinstance(r.metadata, dict) else {}
1082
  where = data.get("filename") or "document"
1083
- if data.get("page_label") is not None:
1084
- where += f", p. {data['page_label']}"
 
1085
  text = r.content.strip()
1086
  if len(text) > _DOC_CONTEXT_MAX_CHARS:
1087
  text = text[:_DOC_CONTEXT_MAX_CHARS].rstrip() + " …"
@@ -1110,12 +1111,12 @@ def _normalize_chunks(raw: Any) -> list[DocumentChunk]:
1110
  DocumentChunk(
1111
  content=str(item.get("content", "")),
1112
  filename=item.get("filename"),
1113
- page_label=item.get("page_label"),
1114
  )
1115
  )
1116
  elif isinstance(item, RetrievalResult):
1117
  data = item.metadata.get("data", {})
1118
- page = data.get("page_label")
1119
  chunks.append(DocumentChunk(
1120
  content=item.content,
1121
  filename=data.get("filename"),
 
37
  from langchain_core.messages import BaseMessage
38
 
39
  from src.middlewares.logging import get_logger
40
+ from src.retrieval.base import RetrievalResult, citable_page
41
  from src.traceability import TraceabilityScratchpad, TraceabilityToolInvoker
42
 
43
  from ..charts import chart_envelopes
 
1080
  for r in results:
1081
  data = r.metadata.get("data", {}) if isinstance(r.metadata, dict) else {}
1082
  where = data.get("filename") or "document"
1083
+ page = citable_page(data) # None for TXT (every chunk is "page 1")
1084
+ if page is not None:
1085
+ where += f", p. {page}"
1086
  text = r.content.strip()
1087
  if len(text) > _DOC_CONTEXT_MAX_CHARS:
1088
  text = text[:_DOC_CONTEXT_MAX_CHARS].rstrip() + " …"
 
1111
  DocumentChunk(
1112
  content=str(item.get("content", "")),
1113
  filename=item.get("filename"),
1114
+ page_label=citable_page(item),
1115
  )
1116
  )
1117
  elif isinstance(item, RetrievalResult):
1118
  data = item.metadata.get("data", {})
1119
+ page = citable_page(data) # None for TXT (every chunk is "page 1")
1120
  chunks.append(DocumentChunk(
1121
  content=item.content,
1122
  filename=data.get("filename"),
src/retrieval/base.py CHANGED
@@ -13,6 +13,26 @@ class RetrievalResult:
13
  source_type: str # "document" | "database"
14
 
15
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
16
  class BaseRetriever(ABC):
17
  @abstractmethod
18
  async def retrieve(
 
13
  source_type: str # "document" | "database"
14
 
15
 
16
+ # File types whose `page_label` is not a real page: TXT ingest stores every chunk as
17
+ # page 1 (`document/ingest_v2.chunks_from_text`), so citing it would always say "p. 1".
18
+ _NO_PAGE_TYPES = frozenset({"txt"})
19
+
20
+
21
+ def citable_page(data: dict[str, Any]) -> Any:
22
+ """The chunk's page to show the user (chip label, answer citation), or None.
23
+
24
+ `data` is a chunk's `metadata["data"]`. None when there is no page, or when the
25
+ file type has no real pages (TXT).
26
+ """
27
+ file_type = str(data.get("file_type") or "").lower()
28
+ if not file_type:
29
+ name = str(data.get("filename") or "").lower()
30
+ file_type = name.rsplit(".", 1)[-1] if "." in name else ""
31
+ if file_type in _NO_PAGE_TYPES:
32
+ return None
33
+ return data.get("page_label")
34
+
35
+
36
  class BaseRetriever(ABC):
37
  @abstractmethod
38
  async def retrieve(
src/traceability/scratchpad.py CHANGED
@@ -17,6 +17,7 @@ from typing import Any
17
  from pydantic import BaseModel
18
 
19
  from src.middlewares.logging import get_logger
 
20
 
21
  from .schemas import (
22
  KnowledgeUsed,
@@ -356,7 +357,7 @@ class TraceabilityScratchpad:
356
  continue
357
  self._doc_seen.add(key)
358
  filename = data.get("filename", "Unknown")
359
- page = data.get("page_label")
360
  source: dict[str, Any] = {
361
  "type": "document",
362
  # Display label. The FE citation chip reads `name` (then title/source/
@@ -365,7 +366,7 @@ class TraceabilityScratchpad:
365
  "name": f"{filename}, p. {page}" if page is not None else filename,
366
  "document_id": data.get("document_id"),
367
  "filename": filename,
368
- "page_label": page,
369
  "query": _truncate(query),
370
  }
371
  snippet = data.get("snippet") or data.get("content") or data.get("text")
 
17
  from pydantic import BaseModel
18
 
19
  from src.middlewares.logging import get_logger
20
+ from src.retrieval.base import citable_page
21
 
22
  from .schemas import (
23
  KnowledgeUsed,
 
357
  continue
358
  self._doc_seen.add(key)
359
  filename = data.get("filename", "Unknown")
360
+ page = citable_page(data) # None for TXT (every chunk is "page 1")
361
  source: dict[str, Any] = {
362
  "type": "document",
363
  # Display label. The FE citation chip reads `name` (then title/source/
 
366
  "name": f"{filename}, p. {page}" if page is not None else filename,
367
  "document_id": data.get("document_id"),
368
  "filename": filename,
369
+ "page_label": data.get("page_label"),
370
  "query": _truncate(query),
371
  }
372
  snippet = data.get("snippet") or data.get("content") or data.get("text")