From a317a951d984f385c4a567e7795b494564b742ef Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 20 May 2026 12:00:24 +0300 Subject: [PATCH] CLI citations: compact panel, inline figures, doc/chunk IDs in footer --- haiku_rag_slim/haiku/rag/app.py | 8 +- haiku_rag_slim/haiku/rag/utils.py | 143 ++++++++++++++++++++++-------- tests/test_utils.py | 85 +++++++++++++++--- 3 files changed, 184 insertions(+), 52 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 1f68deba..15ff412d 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -462,7 +462,9 @@ class HaikuRAGApp: # pragma: no cover self.console.print() self.console.print("[bold green]Answer:[/bold green]") self.console.print(Markdown(answer)) - for renderable in format_citations_rich(citations): + for renderable in await format_citations_rich( + citations, client=self.client + ): self.console.print(renderable) async def analyze( @@ -493,7 +495,9 @@ class HaikuRAGApp: # pragma: no cover self.console.print("[bold green]Answer:[/bold green]") self.console.print(Markdown(result.answer)) - for renderable in format_citations_rich(result.citations): + for renderable in await format_citations_rich( + result.citations, client=self.client + ): self.console.print(renderable) async def research( diff --git a/haiku_rag_slim/haiku/rag/utils.py b/haiku_rag_slim/haiku/rag/utils.py index d6d104d3..861c14df 100644 --- a/haiku_rag_slim/haiku/rag/utils.py +++ b/haiku_rag_slim/haiku/rag/utils.py @@ -12,6 +12,7 @@ if TYPE_CHECKING: from rich.console import RenderableType from haiku.rag.agents.research.models import Citation + from haiku.rag.client import HaikuRAG from haiku.rag.config.models import AppConfig, ModelConfig @@ -338,10 +339,34 @@ def format_bytes(num_bytes: int) -> str: return f"{size:.1f} PB" +CITATION_PREVIEW_CHARS = 300 + + +def _citation_pages(c: "Citation") -> str | None: + if not c.page_numbers: + return None + if len(c.page_numbers) == 1: + return f"p. {c.page_numbers[0]}" + return f"pp. {c.page_numbers[0]}-{c.page_numbers[-1]}" + + +def _citation_section(c: "Citation") -> str | None: + if c.headings: + return c.headings[-1] + return None + + +def _citation_label(c: "Citation") -> str: + if c.document_title and c.document_uri: + return f"{c.document_title} ({c.document_uri})" + return c.document_title or c.document_uri + + def format_citations(citations: "list[Citation]") -> str: """Format citations as plain text with preserved formatting. Used by things like the MCP server where Rich renderables are not available. + Pictures referenced by the chunk are surfaced as ``[Figure: ]`` markers. """ if not citations: return "" @@ -353,34 +378,42 @@ def format_citations(citations: "list[Citation]") -> str: title = c.document_title or c.document_uri header = f"[{idx}] {title}" - # Location info location_parts = [] - if c.page_numbers: - if len(c.page_numbers) == 1: - location_parts.append(f"p. {c.page_numbers[0]}") - else: - location_parts.append(f"pp. {c.page_numbers[0]}-{c.page_numbers[-1]}") - if c.headings: - location_parts.append(f"Section: {c.headings[-1]}") + pages = _citation_pages(c) + if pages: + location_parts.append(pages) + section = _citation_section(c) + if section: + location_parts.append(f"Section: {section}") source = c.document_uri if location_parts: source += f" - {', '.join(location_parts)}" lines.append(f"{header} {source}") + for ref in c.picture_refs: + lines.append(f"[Figure: {ref}]") lines.append(c.content) lines.append("") return "\n".join(lines) -def format_citations_rich(citations: "list[Citation]") -> "list[RenderableType]": - """Format citations as Rich renderables. +async def format_citations_rich( + citations: "list[Citation]", + client: "HaikuRAG | None" = None, +) -> "list[RenderableType]": + """Format citations as Rich renderables for terminal display. - Returns a list of Rich Panel objects for direct console printing, - with content rendered as markdown for syntax highlighting. + Each citation becomes a Panel with a compact header (``[N] Title (URI) — locator``), + a body holding any referenced figures followed by a truncated text preview, and + a dimmed footer that exposes the document and chunk IDs. + + When ``client`` is supplied, picture bytes for ``picture_refs`` are fetched and + rendered inline via ``textual_image``. Without a client, picture refs appear as + ``[Figure: ]`` text markers. """ - from rich.markdown import Markdown + from rich.console import Group from rich.panel import Panel from rich.text import Text @@ -388,35 +421,49 @@ def format_citations_rich(citations: "list[Citation]") -> "list[RenderableType]" return [] renderables: list[RenderableType] = [] - renderables.append(Text("Citations", style="bold")) + renderables.append(Text("")) + renderables.append(Text("Citations", style="bold green")) + renderables.append(Text("")) - for c in citations: - # Build header with IDs - header = Text() - header.append("doc: ", style="dim") - header.append(c.document_id, style="cyan") - header.append(" chunk: ", style="dim") - header.append(c.chunk_id, style="cyan") + for i, c in enumerate(citations): + if i > 0: + renderables.append(Text("")) + idx = c.index if c.index is not None else (i + 1) - # Location info for subtitle - location_parts = [] - if c.page_numbers: - if len(c.page_numbers) == 1: - location_parts.append(f"p. {c.page_numbers[0]}") - else: - location_parts.append(f"pp. {c.page_numbers[0]}-{c.page_numbers[-1]}") - if c.headings: - location_parts.append(f"Section: {c.headings[-1]}") + header_parts: list[str] = [f"[{idx}] {_citation_label(c)}"] + pages = _citation_pages(c) + if pages: + header_parts.append(pages) + section = _citation_section(c) + if section: + header_parts.append(f"§{section}") + header = Text(" — ".join(header_parts), style="bold") + + body: list[RenderableType] = [] + for ref in c.picture_refs: + image_renderable = await _render_picture(client, c.document_id, ref) + body.append( + image_renderable + if image_renderable + else Text(f"[Figure: {ref}]", style="italic dim") + ) + + preview = c.content + if len(preview) > CITATION_PREVIEW_CHARS: + preview = preview[:CITATION_PREVIEW_CHARS].rstrip() + "…" + body.append(Text(preview)) + + footer = Text() + footer.append("doc: ", style="dim") + footer.append(c.document_id, style="dim cyan") + footer.append(" chunk: ", style="dim") + footer.append(c.chunk_id, style="dim cyan") - subtitle = c.document_uri - if c.document_title: - subtitle = f"{c.document_title} ({c.document_uri})" - if location_parts: - subtitle += f" - {', '.join(location_parts)}" panel = Panel( - Markdown(c.content), + Group(*body), title=header, - subtitle=subtitle, + title_align="left", + subtitle=footer, subtitle_align="left", border_style="dim", ) @@ -425,6 +472,28 @@ def format_citations_rich(citations: "list[Citation]") -> "list[RenderableType]" return renderables +async def _render_picture( + client: "HaikuRAG | None", document_id: str, ref: str +) -> "RenderableType | None": + """Fetch a picture and return a Rich renderable, or None on failure/no client.""" + if client is None: + return None + from io import BytesIO + + from PIL import Image as PILImage + from textual_image.renderable import Image as RichImage + + data = await client.document_item_repository.get_picture_bytes(document_id, ref) + if not data: + return None + try: + pil = PILImage.open(BytesIO(data)) + pil.load() + except Exception: + return None + return RichImage(pil) + + def get_default_data_dir() -> Path: """Get the user data directory for the current system platform. diff --git a/tests/test_utils.py b/tests/test_utils.py index 59324264..0af52224 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -618,19 +618,82 @@ def test_format_citations_sequential_indices(): assert "[2] Second" in result +# --- format_citations tests (pictures) --- + + +def test_format_citations_picture_refs_render_as_markers(): + from haiku.rag.agents.research.models import Citation + from haiku.rag.utils import format_citations + + citation = Citation( + document_id="doc1", + chunk_id="chunk1", + document_uri="test://doc", + document_title="Test Doc", + content="text body", + picture_refs=["#/pictures/0", "#/pictures/3"], + ) + result = format_citations([citation]) + assert "[Figure: #/pictures/0]" in result + assert "[Figure: #/pictures/3]" in result + + # --- format_citations_rich tests --- -def test_format_citations_rich_empty(): +def _render_rich(renderables: list) -> str: + from rich.console import Console + + console = Console(record=True, width=200) + for r in renderables: + console.print(r) + return console.export_text() + + +async def test_format_citations_rich_empty(): from haiku.rag.utils import format_citations_rich - assert format_citations_rich([]) == [] + assert await format_citations_rich([]) == [] -def test_format_citations_rich_with_citation(): - from rich.panel import Panel - from rich.text import Text +async def test_format_citations_rich_header_and_footer(): + from haiku.rag.agents.research.models import Citation + from haiku.rag.utils import format_citations_rich + citation = Citation( + document_id="doc-uuid-1", + chunk_id="chunk-uuid-1", + document_uri="test://doc", + document_title="Test Doc", + content="Body", + page_numbers=[1, 2, 3], + headings=["Intro", "Background"], + ) + output = _render_rich(await format_citations_rich([citation])) + assert "Citations" in output + assert "[1] Test Doc (test://doc)" in output + assert "pp. 1-3" in output + assert "§Background" in output + assert "doc: doc-uuid-1" in output + assert "chunk: chunk-uuid-1" in output + + +async def test_format_citations_rich_truncates_long_content(): + from haiku.rag.agents.research.models import Citation + from haiku.rag.utils import CITATION_PREVIEW_CHARS, format_citations_rich + + citation = Citation( + document_id="doc1", + chunk_id="chunk1", + document_uri="test://doc", + content="A" * (CITATION_PREVIEW_CHARS + 200), + ) + output = _render_rich(await format_citations_rich([citation])) + assert "…" in output + assert "A" * (CITATION_PREVIEW_CHARS + 1) not in output + + +async def test_format_citations_rich_picture_marker_without_client(): from haiku.rag.agents.research.models import Citation from haiku.rag.utils import format_citations_rich @@ -638,15 +701,11 @@ def test_format_citations_rich_with_citation(): document_id="doc1", chunk_id="chunk1", document_uri="test://doc", - document_title="Test Doc", - content="Some content", - page_numbers=[1, 2], - headings=["Intro"], + content="body", + picture_refs=["#/pictures/0"], ) - result = format_citations_rich([citation]) - assert len(result) == 2 - assert isinstance(result[0], Text) - assert isinstance(result[1], Panel) + output = _render_rich(await format_citations_rich([citation])) + assert "[Figure: #/pictures/0]" in output # --- get_default_data_dir tests ---