CLI citations: compact panel, inline figures, doc/chunk IDs in footer
This commit is contained in:
Yiorgis Gozadinos
parent
3858ab905a
commit
a317a951d9
3 changed files with 184 additions and 52 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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: <ref>]`` 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: <ref>]`` 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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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 ---
|
||||
|
|
|
|||
Loading…
Reference in a new issue