Merge pull request #236 from ggozad/chore/cleanup
Improve test coverage and remove unnecessary defensive code
This commit is contained in:
commit
df48f47b6e
10 changed files with 387 additions and 235 deletions
|
|
@ -9,6 +9,10 @@
|
||||||
- Default `None` preserves backwards compatibility (bare state emission)
|
- Default `None` preserves backwards compatibility (bare state emission)
|
||||||
- **Page Image Generation Control**: New `generate_page_images` option in `ConversionOptions` to control PDF page image extraction
|
- **Page Image Generation Control**: New `generate_page_images` option in `ConversionOptions` to control PDF page image extraction
|
||||||
|
|
||||||
|
### Changed
|
||||||
|
|
||||||
|
- **CLI Error Handling**: Commands (`rebuild`, `vacuum`, `create-index`, `ask`, `research`) now propagate errors with proper exit codes instead of swallowing exceptions
|
||||||
|
|
||||||
### Fixed
|
### Fixed
|
||||||
|
|
||||||
- **Embed-only rebuild with changed vector dimensions**: Fixed `haiku-rag rebuild --embed-only` failing when the configured embedding model has different dimensions than the database
|
- **Embed-only rebuild with changed vector dimensions**: Fixed `haiku-rag rebuild --embed-only` failing when the configured embedding model has different dimensions than the database
|
||||||
|
|
|
||||||
|
|
@ -78,12 +78,7 @@ class HaikuRAGApp:
|
||||||
return
|
return
|
||||||
|
|
||||||
# Connect without going through Store to avoid upgrades/validation writes
|
# Connect without going through Store to avoid upgrades/validation writes
|
||||||
try:
|
db = lancedb.connect(self.db_path)
|
||||||
db = lancedb.connect(self.db_path)
|
|
||||||
table_names = set(db.table_names())
|
|
||||||
except Exception as e:
|
|
||||||
self.console.print(f"[red]Failed to open database: {e}[/red]")
|
|
||||||
return
|
|
||||||
|
|
||||||
versions = get_package_versions()
|
versions = get_package_versions()
|
||||||
|
|
||||||
|
|
@ -94,24 +89,17 @@ class HaikuRAGApp:
|
||||||
table_stats = store.get_stats()
|
table_stats = store.get_stats()
|
||||||
|
|
||||||
# Read settings after Store init (migrations have run)
|
# Read settings after Store init (migrations have run)
|
||||||
stored_version = "unknown"
|
settings_tbl = db.open_table("settings")
|
||||||
embed_provider: str | None = None
|
arrow = settings_tbl.search().where("id = 'settings'").limit(1).to_arrow()
|
||||||
embed_model: str | None = None
|
rows = arrow.to_pylist()
|
||||||
vector_dim: int | None = None
|
raw = rows[0].get("settings") or "{}"
|
||||||
|
data = json.loads(raw) if isinstance(raw, str) else (raw or {})
|
||||||
if "settings" in table_names:
|
stored_version = str(data.get("version", "unknown"))
|
||||||
settings_tbl = db.open_table("settings")
|
embeddings = data.get("embeddings", {})
|
||||||
arrow = settings_tbl.search().where("id = 'settings'").limit(1).to_arrow()
|
embed_model_obj = embeddings.get("model", {})
|
||||||
rows = arrow.to_pylist() if arrow is not None else []
|
embed_provider = embed_model_obj.get("provider", "unknown")
|
||||||
if rows:
|
embed_model = embed_model_obj.get("name", "unknown")
|
||||||
raw = rows[0].get("settings") or "{}"
|
vector_dim = embed_model_obj.get("vector_dim")
|
||||||
data = json.loads(raw) if isinstance(raw, str) else (raw or {})
|
|
||||||
stored_version = str(data.get("version", stored_version))
|
|
||||||
embeddings = data.get("embeddings", {})
|
|
||||||
embed_model_obj = embeddings.get("model", {})
|
|
||||||
embed_provider = embed_model_obj.get("provider")
|
|
||||||
embed_model = embed_model_obj.get("name")
|
|
||||||
vector_dim = embed_model_obj.get("vector_dim")
|
|
||||||
|
|
||||||
store.close()
|
store.close()
|
||||||
|
|
||||||
|
|
@ -126,32 +114,17 @@ class HaikuRAGApp:
|
||||||
num_unindexed_rows = table_stats["chunks"].get("num_unindexed_rows", 0)
|
num_unindexed_rows = table_stats["chunks"].get("num_unindexed_rows", 0)
|
||||||
|
|
||||||
# Table versions per table (direct API)
|
# Table versions per table (direct API)
|
||||||
doc_versions = (
|
doc_versions = len(list(db.open_table("documents").list_versions()))
|
||||||
len(list(db.open_table("documents").list_versions()))
|
chunk_versions = len(list(db.open_table("chunks").list_versions()))
|
||||||
if "documents" in table_names
|
|
||||||
else 0
|
|
||||||
)
|
|
||||||
chunk_versions = (
|
|
||||||
len(list(db.open_table("chunks").list_versions()))
|
|
||||||
if "chunks" in table_names
|
|
||||||
else 0
|
|
||||||
)
|
|
||||||
|
|
||||||
self.console.print(
|
self.console.print(
|
||||||
f" [repr.attrib_name]haiku.rag version (db)[/repr.attrib_name]: {stored_version}"
|
f" [repr.attrib_name]haiku.rag version (db)[/repr.attrib_name]: {stored_version}"
|
||||||
)
|
)
|
||||||
if embed_provider or embed_model or vector_dim:
|
dim_part = f"{vector_dim}" if vector_dim is not None else "unknown"
|
||||||
provider_part = embed_provider or "unknown"
|
self.console.print(
|
||||||
model_part = embed_model or "unknown"
|
" [repr.attrib_name]embeddings[/repr.attrib_name]: "
|
||||||
dim_part = f"{vector_dim}" if vector_dim is not None else "unknown"
|
f"{embed_provider}/{embed_model} (dim: {dim_part})"
|
||||||
self.console.print(
|
)
|
||||||
" [repr.attrib_name]embeddings[/repr.attrib_name]: "
|
|
||||||
f"{provider_part}/{model_part} (dim: {dim_part})"
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
self.console.print(
|
|
||||||
" [repr.attrib_name]embeddings[/repr.attrib_name]: unknown"
|
|
||||||
)
|
|
||||||
self.console.print(
|
self.console.print(
|
||||||
f" [repr.attrib_name]documents[/repr.attrib_name]: {num_docs} "
|
f" [repr.attrib_name]documents[/repr.attrib_name]: {num_docs} "
|
||||||
f"({format_bytes(doc_bytes)})"
|
f"({format_bytes(doc_bytes)})"
|
||||||
|
|
@ -418,50 +391,47 @@ class HaikuRAGApp:
|
||||||
read_only=self.read_only,
|
read_only=self.read_only,
|
||||||
before=self.before,
|
before=self.before,
|
||||||
) as self.client:
|
) as self.client:
|
||||||
try:
|
citations = []
|
||||||
citations = []
|
if deep:
|
||||||
if deep:
|
graph = build_research_graph(config=self.config)
|
||||||
graph = build_research_graph(config=self.config)
|
context = ResearchContext(original_question=question)
|
||||||
context = ResearchContext(original_question=question)
|
state = ResearchState.from_config(
|
||||||
state = ResearchState.from_config(
|
context=context,
|
||||||
context=context,
|
config=self.config,
|
||||||
config=self.config,
|
max_iterations=2,
|
||||||
max_iterations=2,
|
confidence_threshold=0.0,
|
||||||
confidence_threshold=0.0,
|
)
|
||||||
)
|
state.search_filter = filter
|
||||||
state.search_filter = filter
|
deps = ResearchDeps(client=self.client)
|
||||||
deps = ResearchDeps(client=self.client)
|
|
||||||
|
|
||||||
report = await graph.run(state=state, deps=deps)
|
report = await graph.run(state=state, deps=deps)
|
||||||
|
|
||||||
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
|
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
|
||||||
self.console.print()
|
self.console.print()
|
||||||
if report:
|
if report:
|
||||||
self.console.print("[bold green]Answer:[/bold green]")
|
|
||||||
self.console.print(Markdown(report.executive_summary))
|
|
||||||
if report.main_findings:
|
|
||||||
self.console.print()
|
|
||||||
self.console.print("[bold cyan]Key Findings:[/bold cyan]")
|
|
||||||
for finding in report.main_findings:
|
|
||||||
self.console.print(f"• {finding}")
|
|
||||||
if report.sources_summary:
|
|
||||||
self.console.print()
|
|
||||||
self.console.print("[bold cyan]Sources:[/bold cyan]")
|
|
||||||
self.console.print(report.sources_summary)
|
|
||||||
else:
|
|
||||||
self.console.print("[yellow]No answer generated.[/yellow]")
|
|
||||||
else:
|
|
||||||
answer, citations = await self.client.ask(question, filter=filter)
|
|
||||||
|
|
||||||
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
|
|
||||||
self.console.print()
|
|
||||||
self.console.print("[bold green]Answer:[/bold green]")
|
self.console.print("[bold green]Answer:[/bold green]")
|
||||||
self.console.print(Markdown(answer))
|
self.console.print(Markdown(report.executive_summary))
|
||||||
if cite and citations:
|
if report.main_findings:
|
||||||
for renderable in format_citations_rich(citations):
|
self.console.print()
|
||||||
self.console.print(renderable)
|
self.console.print("[bold cyan]Key Findings:[/bold cyan]")
|
||||||
except Exception as e:
|
for finding in report.main_findings:
|
||||||
self.console.print(f"[red]Error: {e}[/red]")
|
self.console.print(f"• {finding}")
|
||||||
|
if report.sources_summary:
|
||||||
|
self.console.print()
|
||||||
|
self.console.print("[bold cyan]Sources:[/bold cyan]")
|
||||||
|
self.console.print(report.sources_summary)
|
||||||
|
else:
|
||||||
|
self.console.print("[yellow]No answer generated.[/yellow]")
|
||||||
|
else:
|
||||||
|
answer, citations = await self.client.ask(question, filter=filter)
|
||||||
|
|
||||||
|
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))
|
||||||
|
if cite and citations:
|
||||||
|
for renderable in format_citations_rich(citations):
|
||||||
|
self.console.print(renderable)
|
||||||
|
|
||||||
async def research(self, question: str, filter: str | None = None):
|
async def research(self, question: str, filter: str | None = None):
|
||||||
"""Run research via the pydantic-graph pipeline.
|
"""Run research via the pydantic-graph pipeline.
|
||||||
|
|
@ -476,77 +446,73 @@ class HaikuRAGApp:
|
||||||
read_only=self.read_only,
|
read_only=self.read_only,
|
||||||
before=self.before,
|
before=self.before,
|
||||||
) as client:
|
) as client:
|
||||||
try:
|
self.console.print("[bold cyan]Starting research[/bold cyan]")
|
||||||
self.console.print("[bold cyan]Starting research[/bold cyan]")
|
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
|
||||||
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
|
self.console.print()
|
||||||
|
|
||||||
|
graph = build_research_graph(config=self.config)
|
||||||
|
context = ResearchContext(original_question=question)
|
||||||
|
state = ResearchState.from_config(context=context, config=self.config)
|
||||||
|
state.search_filter = filter
|
||||||
|
deps = ResearchDeps(client=client)
|
||||||
|
|
||||||
|
report = await graph.run(state=state, deps=deps)
|
||||||
|
|
||||||
|
if report is None:
|
||||||
|
self.console.print("[red]Research did not produce a report.[/red]")
|
||||||
|
return
|
||||||
|
|
||||||
|
# Display the report
|
||||||
|
self.console.print("[bold green]Research Report[/bold green]")
|
||||||
|
self.console.rule()
|
||||||
|
|
||||||
|
# Title and Executive Summary
|
||||||
|
self.console.print(f"[bold]{report.title}[/bold]")
|
||||||
|
self.console.print()
|
||||||
|
self.console.print("[bold cyan]Executive Summary:[/bold cyan]")
|
||||||
|
self.console.print(report.executive_summary)
|
||||||
|
self.console.print()
|
||||||
|
|
||||||
|
# Confidence (from last evaluation)
|
||||||
|
if state.last_eval:
|
||||||
|
conf = state.last_eval.confidence_score # type: ignore[attr-defined]
|
||||||
|
self.console.print(f"[bold cyan]Confidence:[/bold cyan] {conf:.1%}")
|
||||||
self.console.print()
|
self.console.print()
|
||||||
|
|
||||||
graph = build_research_graph(config=self.config)
|
# Main Findings
|
||||||
context = ResearchContext(original_question=question)
|
if report.main_findings:
|
||||||
state = ResearchState.from_config(context=context, config=self.config)
|
self.console.print("[bold cyan]Main Findings:[/bold cyan]")
|
||||||
state.search_filter = filter
|
for finding in report.main_findings:
|
||||||
deps = ResearchDeps(client=client)
|
self.console.print(f"• {finding}")
|
||||||
|
|
||||||
report = await graph.run(state=state, deps=deps)
|
|
||||||
|
|
||||||
if report is None:
|
|
||||||
self.console.print("[red]Research did not produce a report.[/red]")
|
|
||||||
return
|
|
||||||
|
|
||||||
# Display the report
|
|
||||||
self.console.print("[bold green]Research Report[/bold green]")
|
|
||||||
self.console.rule()
|
|
||||||
|
|
||||||
# Title and Executive Summary
|
|
||||||
self.console.print(f"[bold]{report.title}[/bold]")
|
|
||||||
self.console.print()
|
|
||||||
self.console.print("[bold cyan]Executive Summary:[/bold cyan]")
|
|
||||||
self.console.print(report.executive_summary)
|
|
||||||
self.console.print()
|
self.console.print()
|
||||||
|
|
||||||
# Confidence (from last evaluation)
|
# (Themes section removed)
|
||||||
if state.last_eval:
|
|
||||||
conf = state.last_eval.confidence_score # type: ignore[attr-defined]
|
|
||||||
self.console.print(f"[bold cyan]Confidence:[/bold cyan] {conf:.1%}")
|
|
||||||
self.console.print()
|
|
||||||
|
|
||||||
# Main Findings
|
# Conclusions
|
||||||
if report.main_findings:
|
if report.conclusions:
|
||||||
self.console.print("[bold cyan]Main Findings:[/bold cyan]")
|
self.console.print("[bold cyan]Conclusions:[/bold cyan]")
|
||||||
for finding in report.main_findings:
|
for conclusion in report.conclusions:
|
||||||
self.console.print(f"• {finding}")
|
self.console.print(f"• {conclusion}")
|
||||||
self.console.print()
|
self.console.print()
|
||||||
|
|
||||||
# (Themes section removed)
|
# Recommendations
|
||||||
|
if report.recommendations:
|
||||||
|
self.console.print("[bold cyan]Recommendations:[/bold cyan]")
|
||||||
|
for rec in report.recommendations:
|
||||||
|
self.console.print(f"• {rec}")
|
||||||
|
self.console.print()
|
||||||
|
|
||||||
# Conclusions
|
# Limitations
|
||||||
if report.conclusions:
|
if report.limitations:
|
||||||
self.console.print("[bold cyan]Conclusions:[/bold cyan]")
|
self.console.print("[bold yellow]Limitations:[/bold yellow]")
|
||||||
for conclusion in report.conclusions:
|
for limitation in report.limitations:
|
||||||
self.console.print(f"• {conclusion}")
|
self.console.print(f"• {limitation}")
|
||||||
self.console.print()
|
self.console.print()
|
||||||
|
|
||||||
# Recommendations
|
# Sources Summary
|
||||||
if report.recommendations:
|
if report.sources_summary:
|
||||||
self.console.print("[bold cyan]Recommendations:[/bold cyan]")
|
self.console.print("[bold cyan]Sources:[/bold cyan]")
|
||||||
for rec in report.recommendations:
|
self.console.print(report.sources_summary)
|
||||||
self.console.print(f"• {rec}")
|
|
||||||
self.console.print()
|
|
||||||
|
|
||||||
# Limitations
|
|
||||||
if report.limitations:
|
|
||||||
self.console.print("[bold yellow]Limitations:[/bold yellow]")
|
|
||||||
for limitation in report.limitations:
|
|
||||||
self.console.print(f"• {limitation}")
|
|
||||||
self.console.print()
|
|
||||||
|
|
||||||
# Sources Summary
|
|
||||||
if report.sources_summary:
|
|
||||||
self.console.print("[bold cyan]Sources:[/bold cyan]")
|
|
||||||
self.console.print(report.sources_summary)
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
self.console.print(f"[red]Error during research: {e}[/red]")
|
|
||||||
|
|
||||||
async def rebuild(self, mode: RebuildMode = RebuildMode.FULL):
|
async def rebuild(self, mode: RebuildMode = RebuildMode.FULL):
|
||||||
async with HaikuRAG(
|
async with HaikuRAG(
|
||||||
|
|
@ -556,89 +522,76 @@ class HaikuRAGApp:
|
||||||
read_only=self.read_only,
|
read_only=self.read_only,
|
||||||
before=self.before,
|
before=self.before,
|
||||||
) as client:
|
) as client:
|
||||||
try:
|
documents = await client.list_documents()
|
||||||
documents = await client.list_documents()
|
total_docs = len(documents)
|
||||||
total_docs = len(documents)
|
|
||||||
|
|
||||||
if total_docs == 0:
|
if total_docs == 0:
|
||||||
self.console.print(
|
self.console.print("[yellow]No documents found in database.[/yellow]")
|
||||||
"[yellow]No documents found in database.[/yellow]"
|
return
|
||||||
)
|
|
||||||
return
|
|
||||||
|
|
||||||
mode_desc = {
|
mode_desc = {
|
||||||
RebuildMode.FULL: "full rebuild",
|
RebuildMode.FULL: "full rebuild",
|
||||||
RebuildMode.RECHUNK: "rechunk",
|
RebuildMode.RECHUNK: "rechunk",
|
||||||
RebuildMode.EMBED_ONLY: "embed only",
|
RebuildMode.EMBED_ONLY: "embed only",
|
||||||
}[mode]
|
}[mode]
|
||||||
|
|
||||||
self.console.print(
|
self.console.print(
|
||||||
f"[bold cyan]Rebuilding database ({mode_desc}) with {total_docs} documents...[/bold cyan]"
|
f"[bold cyan]Rebuilding database ({mode_desc}) with {total_docs} documents...[/bold cyan]"
|
||||||
)
|
)
|
||||||
with Progress() as progress:
|
with Progress() as progress:
|
||||||
task = progress.add_task("Rebuilding...", total=total_docs)
|
task = progress.add_task("Rebuilding...", total=total_docs)
|
||||||
async for _ in client.rebuild_database(mode=mode):
|
async for _ in client.rebuild_database(mode=mode):
|
||||||
progress.update(task, advance=1)
|
progress.update(task, advance=1)
|
||||||
|
|
||||||
self.console.print(
|
self.console.print(
|
||||||
"[bold green]Database rebuild completed successfully.[/bold green]"
|
"[bold green]Database rebuild completed successfully.[/bold green]"
|
||||||
)
|
)
|
||||||
except Exception as e:
|
|
||||||
self.console.print(f"[red]Error rebuilding database: {e}[/red]")
|
|
||||||
|
|
||||||
async def vacuum(self):
|
async def vacuum(self):
|
||||||
"""Run database maintenance: optimize and cleanup table history."""
|
"""Run database maintenance: optimize and cleanup table history."""
|
||||||
try:
|
async with HaikuRAG(
|
||||||
async with HaikuRAG(
|
db_path=self.db_path,
|
||||||
db_path=self.db_path,
|
config=self.config,
|
||||||
config=self.config,
|
skip_validation=True,
|
||||||
skip_validation=True,
|
read_only=self.read_only,
|
||||||
read_only=self.read_only,
|
before=self.before,
|
||||||
before=self.before,
|
) as client:
|
||||||
) as client:
|
await client.vacuum()
|
||||||
await client.vacuum()
|
self.console.print("[bold green]Vacuum completed successfully.[/bold green]")
|
||||||
self.console.print(
|
|
||||||
"[bold green]Vacuum completed successfully.[/bold green]"
|
|
||||||
)
|
|
||||||
except Exception as e:
|
|
||||||
self.console.print(f"[red]Error during vacuum: {e}[/red]")
|
|
||||||
|
|
||||||
async def create_index(self):
|
async def create_index(self):
|
||||||
"""Create vector index on the chunks table."""
|
"""Create vector index on the chunks table."""
|
||||||
try:
|
async with HaikuRAG(
|
||||||
async with HaikuRAG(
|
db_path=self.db_path,
|
||||||
db_path=self.db_path,
|
config=self.config,
|
||||||
config=self.config,
|
skip_validation=True,
|
||||||
skip_validation=True,
|
read_only=self.read_only,
|
||||||
read_only=self.read_only,
|
before=self.before,
|
||||||
before=self.before,
|
) as client:
|
||||||
) as client:
|
row_count = client.store.chunks_table.count_rows()
|
||||||
row_count = client.store.chunks_table.count_rows()
|
self.console.print(f"Chunks in database: {row_count}")
|
||||||
self.console.print(f"Chunks in database: {row_count}")
|
|
||||||
|
|
||||||
if row_count < 256:
|
if row_count < 256:
|
||||||
self.console.print(
|
|
||||||
f"[yellow]Warning: Need at least 256 chunks to create an index (have {row_count})[/yellow]"
|
|
||||||
)
|
|
||||||
return
|
|
||||||
|
|
||||||
# Check if index already exists
|
|
||||||
indices = 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]")
|
|
||||||
|
|
||||||
client.store._ensure_vector_index()
|
|
||||||
self.console.print(
|
self.console.print(
|
||||||
"[bold green]Vector index created successfully.[/bold green]"
|
f"[yellow]Warning: Need at least 256 chunks to create an index (have {row_count})[/yellow]"
|
||||||
)
|
)
|
||||||
except Exception as e:
|
return
|
||||||
self.console.print(f"[red]Error creating index: {e}[/red]")
|
|
||||||
|
# Check if index already exists
|
||||||
|
indices = 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]")
|
||||||
|
|
||||||
|
client.store._ensure_vector_index()
|
||||||
|
self.console.print(
|
||||||
|
"[bold green]Vector index created successfully.[/bold green]"
|
||||||
|
)
|
||||||
|
|
||||||
async def download_models(self):
|
async def download_models(self):
|
||||||
"""Download Docling, HuggingFace tokenizer, and Ollama models per config."""
|
"""Download Docling, HuggingFace tokenizer, and Ollama models per config."""
|
||||||
|
|
|
||||||
|
|
@ -16,7 +16,7 @@ def run_chat(
|
||||||
"""
|
"""
|
||||||
try:
|
try:
|
||||||
from haiku.rag.chat.app import ChatApp
|
from haiku.rag.chat.app import ChatApp
|
||||||
except ImportError as e:
|
except ImportError as e: # pragma: no cover
|
||||||
raise ImportError(
|
raise ImportError(
|
||||||
"textual is not installed. Please install it with `pip install 'haiku.rag-slim[tui]'` or use the full haiku.rag package."
|
"textual is not installed. Please install it with `pip install 'haiku.rag-slim[tui]'` or use the full haiku.rag package."
|
||||||
) from e
|
) from e
|
||||||
|
|
|
||||||
|
|
@ -128,7 +128,7 @@ def get_embedder(config: AppConfig = Config) -> EmbedderWrapper:
|
||||||
if provider == "voyageai":
|
if provider == "voyageai":
|
||||||
try:
|
try:
|
||||||
from haiku.rag.embeddings.voyageai import VoyageAIEmbeddingModel
|
from haiku.rag.embeddings.voyageai import VoyageAIEmbeddingModel
|
||||||
except ImportError:
|
except ImportError: # pragma: no cover
|
||||||
raise ImportError(
|
raise ImportError(
|
||||||
"VoyageAI embedder requires the 'voyageai' package. "
|
"VoyageAI embedder requires the 'voyageai' package. "
|
||||||
"Please install haiku.rag with the 'voyageai' extra: "
|
"Please install haiku.rag with the 'voyageai' extra: "
|
||||||
|
|
|
||||||
|
|
@ -1,6 +1,6 @@
|
||||||
try:
|
try:
|
||||||
from haiku.rag.inspector.app import run_inspector
|
from haiku.rag.inspector.app import run_inspector
|
||||||
except ImportError as e:
|
except ImportError as e: # pragma: no cover
|
||||||
raise ImportError(
|
raise ImportError(
|
||||||
"textual is not installed. Please install it with `pip install 'haiku.rag-slim[tui]'` or use the full haiku.rag package."
|
"textual is not installed. Please install it with `pip install 'haiku.rag-slim[tui]'` or use the full haiku.rag package."
|
||||||
) from e
|
) from e
|
||||||
|
|
|
||||||
|
|
@ -30,7 +30,7 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
|
||||||
|
|
||||||
os.environ["TOKENIZERS_PARALLELISM"] = "true"
|
os.environ["TOKENIZERS_PARALLELISM"] = "true"
|
||||||
reranker = MxBAIReranker()
|
reranker = MxBAIReranker()
|
||||||
except ImportError:
|
except ImportError: # pragma: no cover
|
||||||
reranker = None
|
reranker = None
|
||||||
|
|
||||||
elif config.reranking.model and config.reranking.model.provider == "cohere":
|
elif config.reranking.model and config.reranking.model.provider == "cohere":
|
||||||
|
|
@ -38,7 +38,7 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
|
||||||
from haiku.rag.reranking.cohere import CohereReranker
|
from haiku.rag.reranking.cohere import CohereReranker
|
||||||
|
|
||||||
reranker = CohereReranker()
|
reranker = CohereReranker()
|
||||||
except ImportError:
|
except ImportError: # pragma: no cover
|
||||||
reranker = None
|
reranker = None
|
||||||
|
|
||||||
elif config.reranking.model and config.reranking.model.provider == "vllm":
|
elif config.reranking.model and config.reranking.model.provider == "vllm":
|
||||||
|
|
@ -49,7 +49,7 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
|
||||||
if not base_url:
|
if not base_url:
|
||||||
raise ValueError("vLLM reranker requires base_url in reranking.model")
|
raise ValueError("vLLM reranker requires base_url in reranking.model")
|
||||||
reranker = VLLMReranker(config.reranking.model.name, base_url)
|
reranker = VLLMReranker(config.reranking.model.name, base_url)
|
||||||
except ImportError:
|
except ImportError: # pragma: no cover
|
||||||
reranker = None
|
reranker = None
|
||||||
|
|
||||||
elif config.reranking.model and config.reranking.model.provider == "zeroentropy":
|
elif config.reranking.model and config.reranking.model.provider == "zeroentropy":
|
||||||
|
|
@ -59,7 +59,7 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
|
||||||
# Use configured model or default to zerank-1
|
# Use configured model or default to zerank-1
|
||||||
model = config.reranking.model.name or "zerank-1"
|
model = config.reranking.model.name or "zerank-1"
|
||||||
reranker = ZeroEntropyReranker(model)
|
reranker = ZeroEntropyReranker(model)
|
||||||
except ImportError:
|
except ImportError: # pragma: no cover
|
||||||
reranker = None
|
reranker = None
|
||||||
|
|
||||||
_reranker_cache[config_id] = reranker
|
_reranker_cache[config_id] = reranker
|
||||||
|
|
|
||||||
|
|
@ -3,7 +3,7 @@ from haiku.rag.store.models.chunk import Chunk
|
||||||
|
|
||||||
try:
|
try:
|
||||||
import cohere
|
import cohere
|
||||||
except ImportError as e:
|
except ImportError as e: # pragma: no cover
|
||||||
raise ImportError(
|
raise ImportError(
|
||||||
"cohere is not installed. Please install it with `pip install cohere` or use the cohere optional dependency."
|
"cohere is not installed. Please install it with `pip install cohere` or use the cohere optional dependency."
|
||||||
) from e
|
) from e
|
||||||
|
|
|
||||||
|
|
@ -85,12 +85,7 @@ class SettingsRepository:
|
||||||
|
|
||||||
if existing:
|
if existing:
|
||||||
# Preserve existing version if present to avoid interfering with upgrade flow
|
# Preserve existing version if present to avoid interfering with upgrade flow
|
||||||
try:
|
existing_settings = json.loads(existing[0].settings)
|
||||||
existing_settings = (
|
|
||||||
json.loads(existing[0].settings) if existing[0].settings else {}
|
|
||||||
)
|
|
||||||
except Exception:
|
|
||||||
existing_settings = {}
|
|
||||||
if "version" in existing_settings:
|
if "version" in existing_settings:
|
||||||
current_config["version"] = existing_settings["version"]
|
current_config["version"] = existing_settings["version"]
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -506,3 +506,95 @@ async def test_history_nonexistent_db(tmp_path, monkeypatch):
|
||||||
|
|
||||||
calls = [str(c) for c in mock_print.call_args_list]
|
calls = [str(c) for c in mock_print.call_args_list]
|
||||||
assert any("does not exist" in c for c in calls)
|
assert any("does not exist" in c for c in calls)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_init_creates_database(tmp_path, monkeypatch):
|
||||||
|
"""Test init creates a new database."""
|
||||||
|
db_path = tmp_path / "new.lancedb"
|
||||||
|
app = HaikuRAGApp(db_path=db_path)
|
||||||
|
mock_print = MagicMock()
|
||||||
|
monkeypatch.setattr(app.console, "print", mock_print)
|
||||||
|
|
||||||
|
assert not db_path.exists()
|
||||||
|
await app.init()
|
||||||
|
|
||||||
|
assert db_path.exists()
|
||||||
|
calls = [str(c) for c in mock_print.call_args_list]
|
||||||
|
assert any("initialized" in c for c in calls)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_init_existing_database(tmp_path, monkeypatch):
|
||||||
|
"""Test init with existing database shows warning."""
|
||||||
|
from haiku.rag.store.engine import Store
|
||||||
|
|
||||||
|
db_path = tmp_path / "existing.lancedb"
|
||||||
|
store = Store(db_path, create=True)
|
||||||
|
store.close()
|
||||||
|
|
||||||
|
app = HaikuRAGApp(db_path=db_path)
|
||||||
|
mock_print = MagicMock()
|
||||||
|
monkeypatch.setattr(app.console, "print", mock_print)
|
||||||
|
|
||||||
|
await app.init()
|
||||||
|
|
||||||
|
calls = [str(c) for c in mock_print.call_args_list]
|
||||||
|
assert any("already exists" in c for c in calls)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_vacuum(tmp_path, monkeypatch):
|
||||||
|
"""Test vacuum operation."""
|
||||||
|
from haiku.rag.store.engine import Store
|
||||||
|
|
||||||
|
db_path = tmp_path / "test.lancedb"
|
||||||
|
store = Store(db_path, create=True)
|
||||||
|
store.close()
|
||||||
|
|
||||||
|
app = HaikuRAGApp(db_path=db_path)
|
||||||
|
mock_print = MagicMock()
|
||||||
|
monkeypatch.setattr(app.console, "print", mock_print)
|
||||||
|
|
||||||
|
await app.vacuum()
|
||||||
|
|
||||||
|
calls = [str(c) for c in mock_print.call_args_list]
|
||||||
|
assert any("Vacuum completed" in c for c in calls)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_create_index_insufficient_chunks(tmp_path, monkeypatch):
|
||||||
|
"""Test create_index with insufficient chunks shows warning."""
|
||||||
|
from haiku.rag.store.engine import Store
|
||||||
|
|
||||||
|
db_path = tmp_path / "test.lancedb"
|
||||||
|
store = Store(db_path, create=True)
|
||||||
|
store.close()
|
||||||
|
|
||||||
|
app = HaikuRAGApp(db_path=db_path)
|
||||||
|
mock_print = MagicMock()
|
||||||
|
monkeypatch.setattr(app.console, "print", mock_print)
|
||||||
|
|
||||||
|
await app.create_index()
|
||||||
|
|
||||||
|
calls = [str(c) for c in mock_print.call_args_list]
|
||||||
|
assert any("Need at least 256 chunks" in c for c in calls)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_rebuild_empty_database(tmp_path, monkeypatch):
|
||||||
|
"""Test rebuild with empty database shows warning."""
|
||||||
|
from haiku.rag.store.engine import Store
|
||||||
|
|
||||||
|
db_path = tmp_path / "test.lancedb"
|
||||||
|
store = Store(db_path, create=True)
|
||||||
|
store.close()
|
||||||
|
|
||||||
|
app = HaikuRAGApp(db_path=db_path)
|
||||||
|
mock_print = MagicMock()
|
||||||
|
monkeypatch.setattr(app.console, "print", mock_print)
|
||||||
|
|
||||||
|
await app.rebuild()
|
||||||
|
|
||||||
|
calls = [str(c) for c in mock_print.call_args_list]
|
||||||
|
assert any("No documents found" in c for c in calls)
|
||||||
|
|
|
||||||
|
|
@ -1,4 +1,7 @@
|
||||||
from haiku.rag.config import Config
|
import pytest
|
||||||
|
|
||||||
|
from haiku.rag.config import AppConfig, Config
|
||||||
|
from haiku.rag.store.repositories.settings import ConfigMismatchError
|
||||||
|
|
||||||
|
|
||||||
def test_settings_table_populated_on_store_init(temp_db_path):
|
def test_settings_table_populated_on_store_init(temp_db_path):
|
||||||
|
|
@ -48,3 +51,108 @@ def test_monitor_filter_patterns_config():
|
||||||
assert isinstance(Config.monitor.ignore_patterns, list)
|
assert isinstance(Config.monitor.ignore_patterns, list)
|
||||||
assert isinstance(Config.monitor.include_patterns, list)
|
assert isinstance(Config.monitor.include_patterns, list)
|
||||||
assert isinstance(Config.monitor.directories, list)
|
assert isinstance(Config.monitor.directories, list)
|
||||||
|
|
||||||
|
|
||||||
|
class TestValidateConfigCompatibility:
|
||||||
|
"""Tests for validate_config_compatibility method."""
|
||||||
|
|
||||||
|
def test_empty_settings_saves_config(self, temp_db_path):
|
||||||
|
"""When settings row is missing, validation saves current config."""
|
||||||
|
from haiku.rag.store.engine import Store
|
||||||
|
from haiku.rag.store.repositories.settings import SettingsRepository
|
||||||
|
|
||||||
|
store = Store(temp_db_path, create=True, skip_validation=True)
|
||||||
|
settings_repo = SettingsRepository(store)
|
||||||
|
|
||||||
|
# Clear settings to simulate empty state
|
||||||
|
store.settings_table.delete("id = 'settings'")
|
||||||
|
assert settings_repo.get_current_settings() == {}
|
||||||
|
|
||||||
|
# Validation should save settings
|
||||||
|
settings_repo.validate_config_compatibility()
|
||||||
|
|
||||||
|
# Now settings should exist
|
||||||
|
saved = settings_repo.get_current_settings()
|
||||||
|
assert saved.get("embeddings", {}).get("model", {}).get("provider") is not None
|
||||||
|
store.close()
|
||||||
|
|
||||||
|
def test_compatible_config_no_error(self, temp_db_path):
|
||||||
|
"""Compatible config does not raise error."""
|
||||||
|
from haiku.rag.store.engine import Store
|
||||||
|
from haiku.rag.store.repositories.settings import SettingsRepository
|
||||||
|
|
||||||
|
store = Store(temp_db_path, create=True)
|
||||||
|
settings_repo = SettingsRepository(store)
|
||||||
|
|
||||||
|
# Should not raise - same config
|
||||||
|
settings_repo.validate_config_compatibility()
|
||||||
|
store.close()
|
||||||
|
|
||||||
|
def test_provider_mismatch_raises_error(self, temp_db_path):
|
||||||
|
"""Different embedding provider raises ConfigMismatchError."""
|
||||||
|
from haiku.rag.store.engine import Store
|
||||||
|
from haiku.rag.store.repositories.settings import SettingsRepository
|
||||||
|
|
||||||
|
# Create store with default config (ollama)
|
||||||
|
store = Store(temp_db_path, create=True)
|
||||||
|
store.close()
|
||||||
|
|
||||||
|
# Create new config with different provider
|
||||||
|
new_config = AppConfig()
|
||||||
|
new_config.embeddings.model.provider = "openai"
|
||||||
|
|
||||||
|
store2 = Store(temp_db_path, config=new_config, skip_validation=True)
|
||||||
|
settings_repo = SettingsRepository(store2)
|
||||||
|
|
||||||
|
with pytest.raises(ConfigMismatchError) as exc_info:
|
||||||
|
settings_repo.validate_config_compatibility()
|
||||||
|
|
||||||
|
assert "embedding provider" in str(exc_info.value)
|
||||||
|
assert "ollama" in str(exc_info.value)
|
||||||
|
assert "openai" in str(exc_info.value)
|
||||||
|
store2.close()
|
||||||
|
|
||||||
|
def test_model_mismatch_raises_error(self, temp_db_path):
|
||||||
|
"""Different embedding model raises ConfigMismatchError."""
|
||||||
|
from haiku.rag.store.engine import Store
|
||||||
|
from haiku.rag.store.repositories.settings import SettingsRepository
|
||||||
|
|
||||||
|
# Create store with default config
|
||||||
|
store = Store(temp_db_path, create=True)
|
||||||
|
store.close()
|
||||||
|
|
||||||
|
# Create new config with different model
|
||||||
|
new_config = AppConfig()
|
||||||
|
new_config.embeddings.model.name = "different-model"
|
||||||
|
|
||||||
|
store2 = Store(temp_db_path, config=new_config, skip_validation=True)
|
||||||
|
settings_repo = SettingsRepository(store2)
|
||||||
|
|
||||||
|
with pytest.raises(ConfigMismatchError) as exc_info:
|
||||||
|
settings_repo.validate_config_compatibility()
|
||||||
|
|
||||||
|
assert "embedding model" in str(exc_info.value)
|
||||||
|
store2.close()
|
||||||
|
|
||||||
|
def test_vector_dim_mismatch_raises_error(self, temp_db_path):
|
||||||
|
"""Different vector dimension raises ConfigMismatchError."""
|
||||||
|
from haiku.rag.store.engine import Store
|
||||||
|
from haiku.rag.store.repositories.settings import SettingsRepository
|
||||||
|
|
||||||
|
# Create store with default config
|
||||||
|
store = Store(temp_db_path, create=True)
|
||||||
|
store.close()
|
||||||
|
|
||||||
|
# Create new config with different vector dimension
|
||||||
|
new_config = AppConfig()
|
||||||
|
new_config.embeddings.model.vector_dim = 9999
|
||||||
|
|
||||||
|
store2 = Store(temp_db_path, config=new_config, skip_validation=True)
|
||||||
|
settings_repo = SettingsRepository(store2)
|
||||||
|
|
||||||
|
with pytest.raises(ConfigMismatchError) as exc_info:
|
||||||
|
settings_repo.validate_config_compatibility()
|
||||||
|
|
||||||
|
assert "vector dimension" in str(exc_info.value)
|
||||||
|
assert "9999" in str(exc_info.value)
|
||||||
|
store2.close()
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue