Visualize chunk command, format citations to show chunk id
This commit is contained in:
parent
910d382748
commit
75a086bae0
4 changed files with 60 additions and 4 deletions
|
|
@ -274,6 +274,38 @@ class HaikuRAGApp:
|
||||||
for result in results:
|
for result in results:
|
||||||
self._rich_print_search_result(result)
|
self._rich_print_search_result(result)
|
||||||
|
|
||||||
|
async def visualize_chunk(self, chunk_id: str):
|
||||||
|
"""Display visual grounding images for a chunk."""
|
||||||
|
from textual_image.renderable import Image as RichImage
|
||||||
|
|
||||||
|
async with HaikuRAG(db_path=self.db_path, config=self.config) as self.client:
|
||||||
|
chunk = await self.client.chunk_repository.get_by_id(chunk_id)
|
||||||
|
if not chunk:
|
||||||
|
self.console.print(f"[red]Chunk with id {chunk_id} not found.[/red]")
|
||||||
|
return
|
||||||
|
|
||||||
|
images = await self.client.visualize_chunk(chunk)
|
||||||
|
if not images:
|
||||||
|
self.console.print(
|
||||||
|
"[yellow]No visual grounding available for this chunk.[/yellow]"
|
||||||
|
)
|
||||||
|
self.console.print(
|
||||||
|
"This may be because the document was converted without page images."
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
self.console.print(f"[bold]Visual grounding for chunk {chunk_id}[/bold]")
|
||||||
|
if chunk.document_uri:
|
||||||
|
self.console.print(
|
||||||
|
f"[repr.attrib_name]document[/repr.attrib_name]: {chunk.document_uri}"
|
||||||
|
)
|
||||||
|
|
||||||
|
for i, img in enumerate(images):
|
||||||
|
self.console.print(
|
||||||
|
f"\n[bold cyan]Page {i + 1}/{len(images)}[/bold cyan]"
|
||||||
|
)
|
||||||
|
self.console.print(RichImage(img))
|
||||||
|
|
||||||
async def ask(
|
async def ask(
|
||||||
self,
|
self,
|
||||||
question: str,
|
question: str,
|
||||||
|
|
@ -624,11 +656,13 @@ class HaikuRAGApp:
|
||||||
content = Markdown(result.content)
|
content = Markdown(result.content)
|
||||||
self.console.print(
|
self.console.print(
|
||||||
f"[repr.attrib_name]document_id[/repr.attrib_name]: {result.document_id} "
|
f"[repr.attrib_name]document_id[/repr.attrib_name]: {result.document_id} "
|
||||||
|
f"[repr.attrib_name]chunk_id[/repr.attrib_name]: {result.chunk_id} "
|
||||||
f"[repr.attrib_name]score[/repr.attrib_name]: {result.score:.4f}"
|
f"[repr.attrib_name]score[/repr.attrib_name]: {result.score:.4f}"
|
||||||
)
|
)
|
||||||
if result.document_uri:
|
if result.document_uri:
|
||||||
self.console.print("[repr.attrib_name]document uri[/repr.attrib_name]:")
|
self.console.print(
|
||||||
self.console.print(result.document_uri)
|
f"[repr.attrib_name]document uri[/repr.attrib_name]: {result.document_uri}"
|
||||||
|
)
|
||||||
if result.document_title:
|
if result.document_title:
|
||||||
self.console.print("[repr.attrib_name]document title[/repr.attrib_name]:")
|
self.console.print("[repr.attrib_name]document title[/repr.attrib_name]:")
|
||||||
self.console.print(result.document_title)
|
self.console.print(result.document_title)
|
||||||
|
|
|
||||||
|
|
@ -261,6 +261,21 @@ def search(
|
||||||
asyncio.run(app.search(query=query, limit=limit, filter=filter))
|
asyncio.run(app.search(query=query, limit=limit, filter=filter))
|
||||||
|
|
||||||
|
|
||||||
|
@cli.command("visualize", help="Show visual grounding for a chunk")
|
||||||
|
def visualize(
|
||||||
|
chunk_id: str = typer.Argument(
|
||||||
|
help="The ID of the chunk to visualize",
|
||||||
|
),
|
||||||
|
db: Path | None = typer.Option(
|
||||||
|
None,
|
||||||
|
"--db",
|
||||||
|
help="Path to the LanceDB database file",
|
||||||
|
),
|
||||||
|
):
|
||||||
|
app = create_app(db)
|
||||||
|
asyncio.run(app.visualize_chunk(chunk_id=chunk_id))
|
||||||
|
|
||||||
|
|
||||||
@cli.command("ask", help="Ask a question using the QA agent")
|
@cli.command("ask", help="Ask a question using the QA agent")
|
||||||
def ask(
|
def ask(
|
||||||
question: str = typer.Argument(
|
question: str = typer.Argument(
|
||||||
|
|
|
||||||
|
|
@ -31,6 +31,8 @@ class ResearchPlan(BaseModel):
|
||||||
class Citation(BaseModel):
|
class Citation(BaseModel):
|
||||||
"""Resolved citation with full metadata for display/visual grounding."""
|
"""Resolved citation with full metadata for display/visual grounding."""
|
||||||
|
|
||||||
|
document_id: str
|
||||||
|
chunk_id: str
|
||||||
document_uri: str
|
document_uri: str
|
||||||
document_title: str | None = None
|
document_title: str | None = None
|
||||||
page_numbers: list[int] = Field(default_factory=list)
|
page_numbers: list[int] = Field(default_factory=list)
|
||||||
|
|
@ -75,6 +77,8 @@ def resolve_citations(
|
||||||
continue
|
continue
|
||||||
citations.append(
|
citations.append(
|
||||||
Citation(
|
Citation(
|
||||||
|
document_id=r.document_id or "",
|
||||||
|
chunk_id=chunk_id,
|
||||||
document_uri=r.document_uri or "",
|
document_uri=r.document_uri or "",
|
||||||
document_title=r.document_title,
|
document_title=r.document_title,
|
||||||
page_numbers=r.page_numbers,
|
page_numbers=r.page_numbers,
|
||||||
|
|
|
||||||
|
|
@ -280,8 +280,11 @@ def format_citations(citations: "list[Citation]") -> str:
|
||||||
return ""
|
return ""
|
||||||
lines = ["## Citations\n"]
|
lines = ["## Citations\n"]
|
||||||
for c in citations:
|
for c in citations:
|
||||||
# Build citation header
|
# Build citation header with document_id and chunk_id
|
||||||
parts = [f"- **{c.document_uri}**"]
|
parts = [
|
||||||
|
f"- document_id: `{c.document_id}` chunk_id: `{c.chunk_id}` "
|
||||||
|
f"uri: **{c.document_uri}**"
|
||||||
|
]
|
||||||
if c.document_title:
|
if c.document_title:
|
||||||
parts.append(f' - "{c.document_title}"')
|
parts.append(f' - "{c.document_title}"')
|
||||||
location_parts = []
|
location_parts = []
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue