haiku.rag/haiku_rag_slim/haiku/rag/app.py
Yiorgis Gozadinos fdb5710491
Ask and analyze across several databases
Chunk 2 gave search a configured set to fan out over. ask and analyze
covered one database still: the RAG capability had no way to be told which
databases a question spanned, and the analysis sandbox mounted one
document tree.

The selection travels as sources on EvidenceState, beside the filter it
scopes with, so both capabilities read it the same way. clients_covering
is the one rule that turns a selection into clients, used by search, the
sandbox mount and the cite fallback, so a question scoped to some
databases cannot search, mount or cite another. Citations carry the
database they came from, and format_for_agent names it, so the model can
attribute evidence while it answers rather than only afterwards.

The sandbox keeps one flat /documents/{id}/ namespace and resolves each id
to the client holding it, which rests on ids being UUID4. A database
copied from another breaks that, so an id held twice is refused rather
than resolved to whichever arrived last.

On the CLI, search, ask and analyze cover the configured set and label
each result with its database. Every other command works on one, named
with --database NAME (a name reaches a database behind a URI, which --db
cannot) or --db PATH, and refuses a set it cannot choose from instead of
silently reading the default database. Cold databases open together, so a
first query costs the slowest open rather than their sum.
2026-08-24 10:03:46 +03:00

919 lines
35 KiB
Python

import logging
from pathlib import Path
from typing import TYPE_CHECKING
from rich.console import Console
from rich.markdown import Markdown
from rich.markup import escape
from rich.progress import (
BarColumn,
DownloadColumn,
Progress,
SpinnerColumn,
TaskID,
TextColumn,
TransferSpeedColumn,
)
from haiku.rag.client import HaikuRAG, RebuildMode
from haiku.rag.config import AppConfig, get_config
from haiku.rag.mcp import create_mcp_server
from haiku.rag.store.models.chunk import SearchType
from haiku.rag.store.models.document import Document
if TYPE_CHECKING:
from haiku.rag.store.engine import Store
from haiku.rag.store.models import SearchResult
from haiku.rag.config import redact_secrets
from haiku.rag.utils import format_bytes, format_citations_rich
logger = logging.getLogger(__name__)
class HaikuRAGApp:
def __init__(
self,
db_path: Path,
config: AppConfig | None = None,
read_only: bool = False,
federated: bool = False,
):
self.db_path = db_path
self.config = config if config is not None else get_config()
self.read_only = read_only
self.federated = federated
self.console = Console()
from haiku.rag.store.engine import ConnectionMode
self._is_local = ConnectionMode.from_config(self.config) == ConnectionMode.LOCAL
self._display_path = self.db_path if self._is_local else self.config.lancedb.uri
@property
def _client_db_path(self) -> "Path | None":
"""None where the command covers `lancedb.databases`, so the client
opens the configured set rather than one path."""
return None if self.federated else self.db_path
async def init(self):
"""Initialize a new database."""
if self._is_local and self.db_path.exists():
self.console.print(
f"[yellow]Database already exists at {self.db_path}[/yellow]"
)
return
async with HaikuRAG(db_path=self.db_path, config=self.config, create=True):
pass
self.console.print(
f"[bold green]Database initialized at {self._display_path}[/bold green]"
)
async def info(self):
"""Display read-only information about the database without modifying it."""
from haiku.rag.store.info import gather_database_info
# Basic: show path/URI
self.console.print("[bold]haiku.rag database info[/bold]")
self.console.print(
f" [repr.attrib_name]path[/repr.attrib_name]: {self._display_path}"
)
if self._is_local and not self.db_path.exists():
self.console.print("[red]Database path does not exist.[/red]")
return
info = await gather_database_info(self.config, self.db_path)
if not info.exists:
self.console.print(
"[red]Database is empty. Use 'haiku-rag init' to initialize.[/red]"
)
return
self.console.print(
f" [repr.attrib_name]haiku.rag version (db)[/repr.attrib_name]: {info.stored_version}"
)
dim_part = (
f"{info.embeddings.vector_dim}"
if info.embeddings.vector_dim is not None
else "unknown"
)
self.console.print(
" [repr.attrib_name]embeddings[/repr.attrib_name]: "
f"{info.embeddings.provider}/{info.embeddings.name} (dim: {dim_part})"
)
tables = {t.name: t for t in info.tables}
# Per-table row counts and sizes. Missing required tables are
# reported as "absent" rather than raising.
for name in ("documents", "document_meta", "chunks", "document_items"):
entry = tables[name]
if entry.exists:
self.console.print(
f" [repr.attrib_name]{name}[/repr.attrib_name]: {entry.num_rows} "
f"({format_bytes(entry.total_bytes)})"
)
else:
self.console.print(
f" [repr.attrib_name]{name}[/repr.attrib_name]: [yellow]absent[/yellow]"
)
# Vector index information
if tables["chunks"].exists:
num_chunks = tables["chunks"].num_rows
if info.vector_index.exists:
self.console.print(
" [repr.attrib_name]vector index[/repr.attrib_name]: ✓ exists"
)
self.console.print(
f" [repr.attrib_name]indexed chunks[/repr.attrib_name]: {info.vector_index.indexed_rows}"
)
if info.vector_index.unindexed_rows > 0:
self.console.print(
f" [repr.attrib_name]unindexed chunks[/repr.attrib_name]: [yellow]{info.vector_index.unindexed_rows}[/yellow] "
"(consider running: haiku-rag create-index)"
)
else:
self.console.print(
f" [repr.attrib_name]unindexed chunks[/repr.attrib_name]: {info.vector_index.unindexed_rows}"
)
else:
if num_chunks >= 256:
self.console.print(
" [repr.attrib_name]vector index[/repr.attrib_name]: [yellow]✗ not created[/yellow] "
"(run: haiku-rag create-index)"
)
else:
self.console.print(
f" [repr.attrib_name]vector index[/repr.attrib_name]: ✗ not created "
f"(need {256 - num_chunks} more chunks)"
)
if tables["documents"].exists:
self.console.print(
f" [repr.attrib_name]versions (documents)[/repr.attrib_name]: "
f"{tables['documents'].num_versions}"
)
if tables["document_meta"].exists:
self.console.print(
f" [repr.attrib_name]versions (document_meta)[/repr.attrib_name]: "
f"{tables['document_meta'].num_versions}"
)
if tables["chunks"].exists:
self.console.print(
f" [repr.attrib_name]versions (chunks)[/repr.attrib_name]: "
f"{tables['chunks'].num_versions}"
)
self.console.rule()
if info.pending_migrations:
self.console.print(
f"[bold yellow]{len(info.pending_migrations)} migration(s) pending.[/bold yellow] "
"Run [cyan]haiku-rag migrate[/cyan] to upgrade."
)
for step in info.pending_migrations:
self.console.print(
f" [yellow]→[/yellow] {step.version}: {step.description}"
)
else:
self.console.print("[green]Database is up to date.[/green]")
self.console.rule()
self.console.print("[bold]Versions[/bold]")
self.console.print(
f" [repr.attrib_name]haiku.rag[/repr.attrib_name]: {info.packages['haiku_rag']}"
)
self.console.print(
f" [repr.attrib_name]lancedb[/repr.attrib_name]: {info.packages['lancedb']}"
)
self.console.print(
f" [repr.attrib_name]docling[/repr.attrib_name]: {info.packages['docling']}"
)
self.console.print(
f" [repr.attrib_name]pydantic-ai[/repr.attrib_name]: {info.packages['pydantic_ai']}"
)
self.console.print(
f" [repr.attrib_name]docling-document schema[/repr.attrib_name]: {info.packages['docling_document_schema']}"
)
async def doctor(self, duplicates_out: Path | None = None) -> bool:
"""Run health checks and print a report. Returns True if any check failed."""
import os
from contextlib import nullcontext
from haiku.rag.doctor import Severity, run_doctor
self.console.print("[bold]haiku.rag doctor[/bold]")
self.console.print(
f" [repr.attrib_name]path[/repr.attrib_name]: {self._display_path}"
)
if self._is_local and not self.db_path.exists():
self.console.print("[red]Database path does not exist.[/red]")
return True
status = (
self.console.status("Running checks") if self.console.is_terminal else None
)
def on_progress(label: str) -> None:
if status is not None:
status.update(f"{label}...")
cm = status if status is not None else nullcontext()
with cm:
report = await run_doctor(
self.config,
self.db_path,
dict(os.environ),
duplicates_out=duplicates_out,
on_progress=on_progress,
)
glyphs = {
Severity.OK: "[green]✓[/green]",
Severity.WARN: "[yellow]![/yellow]",
Severity.FAIL: "[red]✗[/red]",
}
def render(result):
self.console.print(f"{glyphs[result.severity]} {result.message}")
for detail in result.details:
self.console.print(f" [dim]{detail}[/dim]")
if result.remediation:
self.console.print(f" [dim]→ {result.remediation}[/dim]")
database = [r for r in report.results if not r.name.startswith("provider:")]
providers = [r for r in report.results if r.name.startswith("provider:")]
self.console.rule("[bold]Database[/bold]")
for result in database:
render(result)
if providers:
self.console.rule("[bold]Providers[/bold]")
for result in providers:
render(result)
self.console.rule()
self.console.print(
f"[green]{report.count(Severity.OK)} ok[/green], "
f"[yellow]{report.count(Severity.WARN)} warning(s)[/yellow], "
f"[red]{report.count(Severity.FAIL)} failure(s)[/red]"
)
if duplicates_out is not None:
self.console.print(
f"[dim]Duplicate-document groups written to {duplicates_out}[/dim]"
)
return report.failed
async def history(self, table: str | None = None, limit: int | None = None):
"""Display version history for database tables.
Args:
table: Specific table to show history for (documents, chunks, settings).
If None, shows history for all tables.
limit: Maximum number of versions to show per table.
"""
from haiku.rag.store.engine import Store
if self._is_local and not self.db_path.exists():
self.console.print("[red]Database path does not exist.[/red]")
return
async with Store(
self.db_path,
config=self.config,
skip_validation=True,
read_only=True,
skip_migration_check=True,
) as store:
tables = [
"documents",
"document_meta",
"chunks",
"document_items",
"settings",
]
if table:
if table not in tables:
self.console.print(
f"[red]Unknown table: {table}. Must be one of: {', '.join(tables)}[/red]"
)
return
tables = [table]
self.console.print("[bold]Version History[/bold]")
try:
tags = await store.list_tags()
except Exception as exc:
tags = {}
self.console.print(
f"[yellow]Tag annotations unavailable: {escape(str(exc))}[/yellow]"
)
for table_name in tables:
versions = await store.list_table_versions(table_name)
# Sort by version descending (newest first)
versions = sorted(versions, key=lambda v: v["version"], reverse=True)
if limit:
versions = versions[:limit]
version_tags: dict[int, list[str]] = {}
for tag_name, info in tags.items():
tagged_version = info.tables.get(table_name)
if tagged_version is not None:
version_tags.setdefault(tagged_version, []).append(tag_name)
self.console.print(f"\n[bold cyan]{table_name}[/bold cyan]")
if not versions:
self.console.print(" [dim]No versions found[/dim]")
continue
for v in versions:
version_num = v["version"]
timestamp = v["timestamp"]
suffix = ""
if version_num in version_tags:
names = ", ".join(
escape(n) for n in sorted(version_tags[version_num])
)
suffix = f" [magenta]<- {names}[/magenta]"
self.console.print(
f" [repr.attrib_name]v{version_num}[/repr.attrib_name]: {timestamp}{suffix}"
)
def _tag_write_store(self) -> "Store":
"""Writable store for tag create/delete with normal validation and
migration checks.
A coordinated tag is only reliable when the database schema is
current, and a writable open of a legacy database would create
missing tables as a side effect.
"""
from haiku.rag.store.engine import Store
return Store(self.db_path, config=self.config, read_only=self.read_only)
def _tag_read_store(self) -> "Store":
"""Read-only store for tag inspection; works on old or drifted DBs."""
from haiku.rag.store.engine import Store
return Store(
self.db_path,
config=self.config,
skip_validation=True,
skip_migration_check=True,
read_only=True,
)
async def create_tag(self, name: str):
"""Tag the current version of every table."""
if self._is_local and not self.db_path.exists():
raise ValueError(f"Database path does not exist: {self.db_path}")
async with self._tag_write_store() as store:
await store.create_tag(name)
self.console.print(f"[green]Created tag '{escape(name)}'[/green]")
async def list_tags(self):
"""List database tags, flagging partial ones."""
if self._is_local and not self.db_path.exists():
raise ValueError(f"Database path does not exist: {self.db_path}")
async with self._tag_read_store() as store:
tags = await store.list_tags()
if not tags:
self.console.print("No tags")
return
self.console.print("[bold]Tags[/bold]")
for name in sorted(tags):
info = tags[name]
versions = " ".join(f"{t}=v{v}" for t, v in info.tables.items())
line = f" [repr.attrib_name]{escape(name)}[/repr.attrib_name]: {versions}"
if not info.complete:
missing = ", ".join(info.missing_tables)
line += f" [yellow](partial - missing: {missing})[/yellow]"
self.console.print(line)
async def delete_tag(self, name: str):
"""Delete a tag from every table that has it."""
if self._is_local and not self.db_path.exists():
raise ValueError(f"Database path does not exist: {self.db_path}")
async with self._tag_write_store() as store:
await store.delete_tag(name)
self.console.print(f"[green]Deleted tag '{escape(name)}'[/green]")
async def restore_tag(self, name: str):
"""Restore the database to a tagged state and report the outcome.
The Store context exits before anything is printed; no high-level
database access happens after the restore.
Raises:
ValueError: If the database path does not exist.
"""
if self._is_local and not self.db_path.exists():
raise ValueError(f"Database path does not exist: {self.db_path}")
async with self._tag_write_store() as store:
safety_tag = await store.restore_tag(name)
self.console.print(f"[green]Restored database to tag '{escape(name)}'.[/green]")
self.console.print(
f"The previous state is preserved as '{escape(safety_tag)}'."
)
self.console.print(
"The restored state is now live. Later historical versions remain "
"until eligible for vacuum. Run [cyan]haiku-rag migrate[/cyan] if "
"migration is required."
)
async def list_documents(self, filter: str | None = None):
async with HaikuRAG(
db_path=self.db_path,
config=self.config,
read_only=True,
skip_validation=True,
) as self.client:
documents = await self.client.list_documents(filter=filter)
for doc in documents:
self._rich_print_document(doc, truncate=True)
async def add_document_from_text(
self, text: str, title: str | None = None, metadata: dict | None = None
):
async with HaikuRAG(
db_path=self.db_path,
config=self.config,
read_only=self.read_only,
) as self.client:
doc = await self.client.create_document(
text, title=title, metadata=metadata
)
self._rich_print_document(doc, truncate=True)
self.console.print(
f"[bold green]Document {doc.id} added successfully.[/bold green]"
)
async def add_document_from_source(
self, source: str, title: str | None = None, metadata: dict | None = None
):
async with HaikuRAG(
db_path=self.db_path,
config=self.config,
read_only=self.read_only,
) as self.client:
result = await self.client.create_document_from_source(
source, title=title, metadata=metadata
)
if isinstance(result, list):
for doc in result:
self._rich_print_document(doc, truncate=True)
self.console.print(
f"[bold green]{len(result)} documents added successfully.[/bold green]"
)
else:
self._rich_print_document(result, truncate=True)
self.console.print(
f"[bold green]Document {result.id} added successfully.[/bold green]"
)
async def get_document(self, doc_id: str):
async with HaikuRAG(
db_path=self.db_path,
config=self.config,
read_only=True,
skip_validation=True,
) as self.client:
doc = await self.client.get_document_by_id(doc_id)
if doc is None:
self.console.print(f"[red]Document with id {doc_id} not found.[/red]")
return
self._rich_print_document(doc, truncate=False)
async def delete_document(self, doc_id: str):
async with HaikuRAG(
db_path=self.db_path,
config=self.config,
read_only=self.read_only,
skip_validation=True,
) as self.client:
deleted = await self.client.delete_document(doc_id)
if deleted:
self.console.print(
f"[bold green]Document {doc_id} deleted successfully.[/bold green]"
)
else:
self.console.print(
f"[yellow]Document with id {doc_id} not found.[/yellow]"
)
async def search(
self,
query: str | None = None,
limit: int | None = None,
filter: str | None = None,
search_type: SearchType | None = None,
image: Path | None = None,
):
if query is None and image is None:
self.console.print(
"[red]Provide either a query argument or --image PATH.[/red]"
)
return
if query is not None and image is not None:
self.console.print("[red]Pass either a query or --image, not both.[/red]")
return
if query is None and search_type is not None:
self.console.print("[red]Pass --search-type only for text queries[/red]")
return
search_input: str | bytes
if image is not None:
search_input = image.read_bytes()
else:
assert query is not None
search_input = query
async with HaikuRAG(
db_path=self._client_db_path,
config=self.config,
read_only=True,
) as self.client:
results = await self.client.search(
search_input,
limit=limit,
filter=filter,
search_type=search_type,
)
if not results:
self.console.print("[yellow]No results found.[/yellow]")
return
for result in results:
self._rich_print_search_result(result)
async def visualize_chunk(self, chunk_id: str, expand: bool = True):
"""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,
read_only=True,
skip_validation=True,
) as self.client:
chunk = await self.client.get_chunk_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, expand=expand)
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(
self,
question: str,
filter: str | None = None,
images: list[Path] | None = None,
):
"""Ask a question using the RAG system.
Args:
question: The question to ask
filter: SQL WHERE clause to filter documents
images: Paths of images to attach to the question
"""
async with HaikuRAG(
db_path=self._client_db_path,
config=self.config,
read_only=True,
) as self.client:
answer, citations = await self.client.ask(
question,
filter=filter,
images=[path.read_bytes() for path in images] if images else None,
)
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
self.console.print()
self.console.print("[bold green]Answer:[/bold green]")
self.console.print(Markdown(answer))
for renderable in await format_citations_rich(
citations, client=self.client
):
self.console.print(renderable)
async def analyze(
self,
question: str,
filter: str | None = None,
images: list[Path] | None = None,
):
"""Answer a question using the analysis capability.
Args:
question: The question to answer
filter: SQL WHERE clause to filter documents
images: Paths of images to attach to the question
"""
async with HaikuRAG(
db_path=self._client_db_path,
config=self.config,
read_only=True,
) as self.client:
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
self.console.print()
self.console.print(
"[dim]Running analysis capability with code execution...[/dim]"
)
self.console.print()
result = await self.client.analyze(
question,
filter=filter,
images=[path.read_bytes() for path in images] if images else None,
)
self.console.print("[bold green]Answer:[/bold green]")
self.console.print(Markdown(result.answer))
for renderable in await format_citations_rich(
result.citations, client=self.client
):
self.console.print(renderable)
async def rebuild(self, mode: RebuildMode = RebuildMode.FULL):
async with HaikuRAG(
db_path=self.db_path,
config=self.config,
skip_validation=True,
read_only=self.read_only,
) as client:
if mode == RebuildMode.SET_EMBEDDER:
async for _ in client.rebuild_database(mode=mode):
pass
self.console.print(
"[bold green]Stored embedder settings updated.[/bold green]"
)
return
documents = await client.list_documents()
total_docs = len(documents)
if total_docs == 0:
self.console.print("[yellow]No documents found in database.[/yellow]")
return
mode_desc = {
RebuildMode.FULL: "full rebuild",
RebuildMode.RECHUNK: "rechunk",
RebuildMode.EMBED_ONLY: "embed only",
RebuildMode.TITLE_ONLY: "title only",
RebuildMode.DESCRIPTIONS: "picture descriptions",
}[mode]
self.console.print(
f"[bold cyan]Rebuilding database ({mode_desc}) with {total_docs} documents...[/bold cyan]"
)
with Progress() as progress:
task = progress.add_task("Rebuilding...", total=total_docs)
async for _ in client.rebuild_database(mode=mode):
progress.update(task, advance=1)
self.console.print(
"[bold green]Database rebuild completed successfully.[/bold green]"
)
async def vacuum(self):
"""Run database maintenance: optimize and cleanup table history."""
async with HaikuRAG(
db_path=self.db_path,
config=self.config,
skip_validation=True,
read_only=self.read_only,
) as client:
await client.vacuum()
self.console.print("[bold green]Vacuum completed successfully.[/bold green]")
async def migrate(self) -> list[str]:
"""Run pending database migrations.
Returns:
List of descriptions of applied migrations.
"""
from haiku.rag.store.engine import Store
async with Store(
self.db_path,
config=self.config,
skip_validation=True,
skip_migration_check=True,
read_only=self.read_only,
) as store:
return await store.migrate()
async def create_index(self):
"""Create vector index on the chunks table."""
async with HaikuRAG(
db_path=self.db_path,
config=self.config,
skip_validation=True,
read_only=self.read_only,
) as client:
row_count = await client.store.chunks_table.count_rows()
self.console.print(f"Chunks in database: {row_count}")
if row_count < 256:
self.console.print(
f"[yellow]Warning: Need at least 256 chunks to create an index (have {row_count})[/yellow]"
)
return
indices = await client.store.chunks_table.list_indices()
has_vector_index = any("vector" in str(idx).lower() for idx in indices)
if has_vector_index:
self.console.print(
"[yellow]Rebuilding existing vector index...[/yellow]"
)
else:
self.console.print("[bold]Creating vector index...[/bold]")
await client.store._ensure_vector_index()
self.console.print(
"[bold green]Vector index created successfully.[/bold green]"
)
async def download_models(self):
"""Download Docling, HuggingFace tokenizer, and Ollama models per config."""
from haiku.rag.client.downloads import download_models
progress: Progress | None = None
task_id: TaskID | None = None
current_model = ""
current_digest = ""
async for event in download_models(self.config):
if event.status == "start":
self.console.print(
f"[bold blue]Downloading {event.model}...[/bold blue]"
)
elif event.status == "done":
if progress:
progress.stop()
progress = None
task_id = None
self.console.print(f"[green]✓[/green] {event.model}")
current_model = ""
current_digest = ""
elif event.status == "pulling":
self.console.print(f"[bold blue]Pulling {event.model}...[/bold blue]")
current_model = event.model
progress = Progress(
SpinnerColumn(),
TextColumn("[progress.description]{task.description}"),
BarColumn(),
DownloadColumn(),
TransferSpeedColumn(),
console=self.console,
transient=True,
auto_refresh=False,
)
progress.start()
task_id = progress.add_task(event.model, total=None)
elif event.status == "downloading" and progress and task_id is not None:
if event.digest != current_digest:
current_digest = event.digest
short_digest = event.digest[:19] if event.digest else ""
progress.update(
task_id,
description=f"{current_model} ({short_digest})",
total=event.total,
completed=0,
)
progress.update(task_id, completed=event.completed, refresh=True)
elif progress and task_id is not None:
progress.update(
task_id,
description=f"{current_model}: {event.status}",
refresh=True,
)
def show_settings(self):
"""Display current configuration settings."""
self.console.print("[bold]haiku.rag configuration[/bold]")
self.console.print()
# redact_secrets walks the whole dump: masking only top-level names left
# nested api keys, tokens and source passwords printed in full.
for field_name, field_value in redact_secrets(self.config.model_dump()).items():
self.console.print(
f" [repr.attrib_name]{field_name}[/repr.attrib_name]: {field_value}"
)
def _rich_print_document(self, doc: Document, truncate: bool = False):
"""Format a document for display."""
if truncate:
content = doc.content.splitlines()
if len(content) > 3:
content = content[:3] + ["\n"]
content = "\n".join(content)
content = Markdown(content)
else:
content = Markdown(doc.content)
title_part = (
f" [repr.attrib_name]title[/repr.attrib_name]: {doc.title}"
if doc.title
else ""
)
self.console.print(
f"[repr.attrib_name]id[/repr.attrib_name]: {doc.id} "
f"[repr.attrib_name]uri[/repr.attrib_name]: {doc.uri}"
+ title_part
+ f" [repr.attrib_name]meta[/repr.attrib_name]: {doc.metadata}"
)
self.console.print(
f"[repr.attrib_name]created at[/repr.attrib_name]: {doc.created_at} [repr.attrib_name]updated at[/repr.attrib_name]: {doc.updated_at}"
)
self.console.print("[repr.attrib_name]content[/repr.attrib_name]:")
self.console.print(content)
self.console.rule()
def _rich_print_search_result(self, result: "SearchResult"):
"""Format a search result for display."""
content = Markdown(result.content)
self.console.print(
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}"
)
if result.source and len(self.config.lancedb.databases) > 1:
self.console.print(
f"[repr.attrib_name]database[/repr.attrib_name]: {result.source}"
)
if result.document_uri:
self.console.print(
f"[repr.attrib_name]document uri[/repr.attrib_name]: {result.document_uri}"
)
if result.document_title:
self.console.print("[repr.attrib_name]document title[/repr.attrib_name]:")
self.console.print(result.document_title)
if result.page_numbers:
self.console.print("[repr.attrib_name]pages[/repr.attrib_name]:")
self.console.print(", ".join(str(p) for p in result.page_numbers))
if result.headings:
self.console.print("[repr.attrib_name]headings[/repr.attrib_name]:")
self.console.print(" > ".join(result.headings))
self.console.print("[repr.attrib_name]content[/repr.attrib_name]:")
self.console.print(content)
self.console.rule()
async def run_mcp(
self,
transport: str | None = None,
host: str = "127.0.0.1",
port: int = 8001,
):
"""Run the MCP server until interrupted."""
async with HaikuRAG(
self.db_path,
config=self.config,
read_only=self.read_only,
):
server = create_mcp_server(
self.db_path, config=self.config, read_only=self.read_only
)
try:
if transport == "stdio":
await server.run_stdio_async()
else:
logger.info(f"Starting MCP server on {host}:{port}")
await server.run_http_async(
transport="streamable-http", host=host, port=port
)
except KeyboardInterrupt:
pass