haiku.rag-slim 0.16.0__py3-none-any.whl → 0.24.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of haiku.rag-slim might be problematic. Click here for more details.

Files changed (94) hide show
  1. haiku/rag/app.py +430 -72
  2. haiku/rag/chunkers/__init__.py +31 -0
  3. haiku/rag/chunkers/base.py +31 -0
  4. haiku/rag/chunkers/docling_local.py +164 -0
  5. haiku/rag/chunkers/docling_serve.py +179 -0
  6. haiku/rag/cli.py +207 -24
  7. haiku/rag/cli_chat.py +489 -0
  8. haiku/rag/client.py +1251 -266
  9. haiku/rag/config/__init__.py +16 -10
  10. haiku/rag/config/loader.py +5 -44
  11. haiku/rag/config/models.py +126 -17
  12. haiku/rag/converters/__init__.py +31 -0
  13. haiku/rag/converters/base.py +63 -0
  14. haiku/rag/converters/docling_local.py +193 -0
  15. haiku/rag/converters/docling_serve.py +229 -0
  16. haiku/rag/converters/text_utils.py +237 -0
  17. haiku/rag/embeddings/__init__.py +123 -24
  18. haiku/rag/embeddings/voyageai.py +175 -20
  19. haiku/rag/graph/__init__.py +0 -11
  20. haiku/rag/graph/agui/__init__.py +8 -2
  21. haiku/rag/graph/agui/cli_renderer.py +1 -1
  22. haiku/rag/graph/agui/emitter.py +219 -31
  23. haiku/rag/graph/agui/server.py +20 -62
  24. haiku/rag/graph/agui/stream.py +1 -2
  25. haiku/rag/graph/research/__init__.py +5 -2
  26. haiku/rag/graph/research/dependencies.py +12 -126
  27. haiku/rag/graph/research/graph.py +390 -135
  28. haiku/rag/graph/research/models.py +91 -112
  29. haiku/rag/graph/research/prompts.py +99 -91
  30. haiku/rag/graph/research/state.py +35 -27
  31. haiku/rag/inspector/__init__.py +8 -0
  32. haiku/rag/inspector/app.py +259 -0
  33. haiku/rag/inspector/widgets/__init__.py +6 -0
  34. haiku/rag/inspector/widgets/chunk_list.py +100 -0
  35. haiku/rag/inspector/widgets/context_modal.py +89 -0
  36. haiku/rag/inspector/widgets/detail_view.py +130 -0
  37. haiku/rag/inspector/widgets/document_list.py +75 -0
  38. haiku/rag/inspector/widgets/info_modal.py +209 -0
  39. haiku/rag/inspector/widgets/search_modal.py +183 -0
  40. haiku/rag/inspector/widgets/visual_modal.py +126 -0
  41. haiku/rag/mcp.py +106 -102
  42. haiku/rag/monitor.py +33 -9
  43. haiku/rag/providers/__init__.py +5 -0
  44. haiku/rag/providers/docling_serve.py +108 -0
  45. haiku/rag/qa/__init__.py +12 -10
  46. haiku/rag/qa/agent.py +43 -61
  47. haiku/rag/qa/prompts.py +35 -57
  48. haiku/rag/reranking/__init__.py +9 -6
  49. haiku/rag/reranking/base.py +1 -1
  50. haiku/rag/reranking/cohere.py +5 -4
  51. haiku/rag/reranking/mxbai.py +5 -2
  52. haiku/rag/reranking/vllm.py +3 -4
  53. haiku/rag/reranking/zeroentropy.py +6 -5
  54. haiku/rag/store/__init__.py +2 -1
  55. haiku/rag/store/engine.py +242 -42
  56. haiku/rag/store/exceptions.py +4 -0
  57. haiku/rag/store/models/__init__.py +8 -2
  58. haiku/rag/store/models/chunk.py +190 -0
  59. haiku/rag/store/models/document.py +46 -0
  60. haiku/rag/store/repositories/chunk.py +141 -121
  61. haiku/rag/store/repositories/document.py +25 -84
  62. haiku/rag/store/repositories/settings.py +11 -14
  63. haiku/rag/store/upgrades/__init__.py +19 -3
  64. haiku/rag/store/upgrades/v0_10_1.py +1 -1
  65. haiku/rag/store/upgrades/v0_19_6.py +65 -0
  66. haiku/rag/store/upgrades/v0_20_0.py +68 -0
  67. haiku/rag/store/upgrades/v0_23_1.py +100 -0
  68. haiku/rag/store/upgrades/v0_9_3.py +3 -3
  69. haiku/rag/utils.py +371 -146
  70. {haiku_rag_slim-0.16.0.dist-info → haiku_rag_slim-0.24.0.dist-info}/METADATA +15 -12
  71. haiku_rag_slim-0.24.0.dist-info/RECORD +78 -0
  72. {haiku_rag_slim-0.16.0.dist-info → haiku_rag_slim-0.24.0.dist-info}/WHEEL +1 -1
  73. haiku/rag/chunker.py +0 -65
  74. haiku/rag/embeddings/base.py +0 -25
  75. haiku/rag/embeddings/ollama.py +0 -28
  76. haiku/rag/embeddings/openai.py +0 -26
  77. haiku/rag/embeddings/vllm.py +0 -29
  78. haiku/rag/graph/agui/events.py +0 -254
  79. haiku/rag/graph/common/__init__.py +0 -5
  80. haiku/rag/graph/common/models.py +0 -42
  81. haiku/rag/graph/common/nodes.py +0 -265
  82. haiku/rag/graph/common/prompts.py +0 -46
  83. haiku/rag/graph/common/utils.py +0 -44
  84. haiku/rag/graph/deep_qa/__init__.py +0 -1
  85. haiku/rag/graph/deep_qa/dependencies.py +0 -27
  86. haiku/rag/graph/deep_qa/graph.py +0 -243
  87. haiku/rag/graph/deep_qa/models.py +0 -20
  88. haiku/rag/graph/deep_qa/prompts.py +0 -59
  89. haiku/rag/graph/deep_qa/state.py +0 -56
  90. haiku/rag/graph/research/common.py +0 -87
  91. haiku/rag/reader.py +0 -135
  92. haiku_rag_slim-0.16.0.dist-info/RECORD +0 -71
  93. {haiku_rag_slim-0.16.0.dist-info → haiku_rag_slim-0.24.0.dist-info}/entry_points.txt +0 -0
  94. {haiku_rag_slim-0.16.0.dist-info → haiku_rag_slim-0.24.0.dist-info}/licenses/LICENSE +0 -0
haiku/rag/app.py CHANGED
@@ -1,14 +1,24 @@
1
1
  import asyncio
2
2
  import json
3
3
  import logging
4
+ from datetime import datetime
4
5
  from importlib.metadata import version as pkg_version
5
6
  from pathlib import Path
7
+ from typing import TYPE_CHECKING
6
8
 
7
9
  from rich.console import Console
8
10
  from rich.markdown import Markdown
9
- from rich.progress import Progress
10
-
11
- from haiku.rag.client import HaikuRAG
11
+ from rich.progress import (
12
+ BarColumn,
13
+ DownloadColumn,
14
+ Progress,
15
+ SpinnerColumn,
16
+ TaskID,
17
+ TextColumn,
18
+ TransferSpeedColumn,
19
+ )
20
+
21
+ from haiku.rag.client import HaikuRAG, RebuildMode
12
22
  from haiku.rag.config import AppConfig, Config
13
23
  from haiku.rag.graph.agui import AGUIConsoleRenderer, stream_graph
14
24
  from haiku.rag.graph.research.dependencies import ResearchContext
@@ -16,18 +26,44 @@ from haiku.rag.graph.research.graph import build_research_graph
16
26
  from haiku.rag.graph.research.state import ResearchDeps, ResearchState
17
27
  from haiku.rag.mcp import create_mcp_server
18
28
  from haiku.rag.monitor import FileWatcher
19
- from haiku.rag.store.models.chunk import Chunk
20
29
  from haiku.rag.store.models.document import Document
21
30
 
31
+ if TYPE_CHECKING:
32
+ from haiku.rag.store.models import SearchResult
33
+ from haiku.rag.utils import format_bytes, format_citations_rich
34
+
22
35
  logger = logging.getLogger(__name__)
23
36
 
24
37
 
25
38
  class HaikuRAGApp:
26
- def __init__(self, db_path: Path, config: AppConfig = Config):
39
+ def __init__(
40
+ self,
41
+ db_path: Path,
42
+ config: AppConfig = Config,
43
+ read_only: bool = False,
44
+ before: datetime | None = None,
45
+ ):
27
46
  self.db_path = db_path
28
47
  self.config = config
48
+ self.read_only = read_only
49
+ self.before = before
29
50
  self.console = Console()
30
51
 
52
+ async def init(self):
53
+ """Initialize a new database."""
54
+ if self.db_path.exists():
55
+ self.console.print(
56
+ f"[yellow]Database already exists at {self.db_path}[/yellow]"
57
+ )
58
+ return
59
+
60
+ # Create the database
61
+ client = HaikuRAG(db_path=self.db_path, config=self.config, create=True)
62
+ client.close()
63
+ self.console.print(
64
+ f"[bold green]Database initialized at {self.db_path}[/bold green]"
65
+ )
66
+
31
67
  async def info(self):
32
68
  """Display read-only information about the database without modifying it."""
33
69
 
@@ -64,7 +100,13 @@ class HaikuRAGApp:
64
100
  except Exception:
65
101
  docling_version = "unknown"
66
102
 
67
- # Read settings (if present) to find stored haiku.rag version and embedding config
103
+ # Get comprehensive table statistics (this also runs migrations)
104
+ from haiku.rag.store.engine import Store
105
+
106
+ store = Store(self.db_path, config=self.config, skip_validation=True)
107
+ table_stats = store.get_stats()
108
+
109
+ # Read settings after Store init (migrations have run)
68
110
  stored_version = "unknown"
69
111
  embed_provider: str | None = None
70
112
  embed_model: str | None = None
@@ -79,14 +121,22 @@ class HaikuRAGApp:
79
121
  data = json.loads(raw) if isinstance(raw, str) else (raw or {})
80
122
  stored_version = str(data.get("version", stored_version))
81
123
  embeddings = data.get("embeddings", {})
82
- embed_provider = embeddings.get("provider")
83
- embed_model = embeddings.get("model")
84
- vector_dim = embeddings.get("vector_dim")
124
+ embed_model_obj = embeddings.get("model", {})
125
+ embed_provider = embed_model_obj.get("provider")
126
+ embed_model = embed_model_obj.get("name")
127
+ vector_dim = embed_model_obj.get("vector_dim")
128
+
129
+ store.close()
130
+
131
+ num_docs = table_stats["documents"].get("num_rows", 0)
132
+ doc_bytes = table_stats["documents"].get("total_bytes", 0)
133
+
134
+ num_chunks = table_stats["chunks"].get("num_rows", 0)
135
+ chunk_bytes = table_stats["chunks"].get("total_bytes", 0)
85
136
 
86
- num_docs = 0
87
- if "documents" in table_names:
88
- docs_tbl = db.open_table("documents")
89
- num_docs = int(docs_tbl.count_rows()) # type: ignore[attr-defined]
137
+ has_vector_index = table_stats["chunks"].get("has_vector_index", False)
138
+ num_indexed_rows = table_stats["chunks"].get("num_indexed_rows", 0)
139
+ num_unindexed_rows = table_stats["chunks"].get("num_unindexed_rows", 0)
90
140
 
91
141
  # Table versions per table (direct API)
92
142
  doc_versions = (
@@ -116,8 +166,43 @@ class HaikuRAGApp:
116
166
  " [repr.attrib_name]embeddings[/repr.attrib_name]: unknown"
117
167
  )
118
168
  self.console.print(
119
- f" [repr.attrib_name]documents[/repr.attrib_name]: {num_docs}"
169
+ f" [repr.attrib_name]documents[/repr.attrib_name]: {num_docs} "
170
+ f"({format_bytes(doc_bytes)})"
120
171
  )
172
+ self.console.print(
173
+ f" [repr.attrib_name]chunks[/repr.attrib_name]: {num_chunks} "
174
+ f"({format_bytes(chunk_bytes)})"
175
+ )
176
+
177
+ # Vector index information
178
+ if has_vector_index:
179
+ self.console.print(
180
+ " [repr.attrib_name]vector index[/repr.attrib_name]: ✓ exists"
181
+ )
182
+ self.console.print(
183
+ f" [repr.attrib_name]indexed chunks[/repr.attrib_name]: {num_indexed_rows}"
184
+ )
185
+ if num_unindexed_rows > 0:
186
+ self.console.print(
187
+ f" [repr.attrib_name]unindexed chunks[/repr.attrib_name]: [yellow]{num_unindexed_rows}[/yellow] "
188
+ "(consider running: haiku-rag create-index)"
189
+ )
190
+ else:
191
+ self.console.print(
192
+ f" [repr.attrib_name]unindexed chunks[/repr.attrib_name]: {num_unindexed_rows}"
193
+ )
194
+ else:
195
+ if num_chunks >= 256:
196
+ self.console.print(
197
+ " [repr.attrib_name]vector index[/repr.attrib_name]: [yellow]✗ not created[/yellow] "
198
+ "(run: haiku-rag create-index)"
199
+ )
200
+ else:
201
+ self.console.print(
202
+ f" [repr.attrib_name]vector index[/repr.attrib_name]: ✗ not created "
203
+ f"(need {256 - num_chunks} more chunks)"
204
+ )
205
+
121
206
  self.console.print(
122
207
  f" [repr.attrib_name]versions (documents)[/repr.attrib_name]: {doc_versions}"
123
208
  )
@@ -136,16 +221,76 @@ class HaikuRAGApp:
136
221
  f" [repr.attrib_name]docling[/repr.attrib_name]: {docling_version}"
137
222
  )
138
223
 
224
+ async def history(self, table: str | None = None, limit: int | None = None):
225
+ """Display version history for database tables.
226
+
227
+ Args:
228
+ table: Specific table to show history for (documents, chunks, settings).
229
+ If None, shows history for all tables.
230
+ limit: Maximum number of versions to show per table.
231
+ """
232
+ from haiku.rag.store.engine import Store
233
+
234
+ if not self.db_path.exists():
235
+ self.console.print("[red]Database path does not exist.[/red]")
236
+ return
237
+
238
+ store = Store(self.db_path, config=self.config, skip_validation=True)
239
+
240
+ tables = ["documents", "chunks", "settings"]
241
+ if table:
242
+ if table not in tables:
243
+ self.console.print(
244
+ f"[red]Unknown table: {table}. Must be one of: {', '.join(tables)}[/red]"
245
+ )
246
+ store.close()
247
+ return
248
+ tables = [table]
249
+
250
+ self.console.print("[bold]Version History[/bold]")
251
+
252
+ for table_name in tables:
253
+ versions = store.list_table_versions(table_name)
254
+
255
+ # Sort by version descending (newest first)
256
+ versions = sorted(versions, key=lambda v: v["version"], reverse=True)
257
+
258
+ if limit:
259
+ versions = versions[:limit]
260
+
261
+ self.console.print(f"\n[bold cyan]{table_name}[/bold cyan]")
262
+
263
+ if not versions:
264
+ self.console.print(" [dim]No versions found[/dim]")
265
+ continue
266
+
267
+ for v in versions:
268
+ version_num = v["version"]
269
+ timestamp = v["timestamp"]
270
+ self.console.print(
271
+ f" [repr.attrib_name]v{version_num}[/repr.attrib_name]: {timestamp}"
272
+ )
273
+
274
+ store.close()
275
+
139
276
  async def list_documents(self, filter: str | None = None):
140
277
  async with HaikuRAG(
141
- db_path=self.db_path, config=self.config, allow_create=False
278
+ db_path=self.db_path,
279
+ config=self.config,
280
+ read_only=self.read_only,
281
+ before=self.before,
142
282
  ) as self.client:
143
283
  documents = await self.client.list_documents(filter=filter)
144
284
  for doc in documents:
145
285
  self._rich_print_document(doc, truncate=True)
146
286
 
147
287
  async def add_document_from_text(self, text: str, metadata: dict | None = None):
148
- async with HaikuRAG(db_path=self.db_path, config=self.config) as self.client:
288
+ async with HaikuRAG(
289
+ db_path=self.db_path,
290
+ config=self.config,
291
+ read_only=self.read_only,
292
+ before=self.before,
293
+ ) as self.client:
149
294
  doc = await self.client.create_document(text, metadata=metadata)
150
295
  self._rich_print_document(doc, truncate=True)
151
296
  self.console.print(
@@ -155,7 +300,12 @@ class HaikuRAGApp:
155
300
  async def add_document_from_source(
156
301
  self, source: str, title: str | None = None, metadata: dict | None = None
157
302
  ):
158
- async with HaikuRAG(db_path=self.db_path, config=self.config) as self.client:
303
+ async with HaikuRAG(
304
+ db_path=self.db_path,
305
+ config=self.config,
306
+ read_only=self.read_only,
307
+ before=self.before,
308
+ ) as self.client:
159
309
  result = await self.client.create_document_from_source(
160
310
  source, title=title, metadata=metadata
161
311
  )
@@ -173,7 +323,10 @@ class HaikuRAGApp:
173
323
 
174
324
  async def get_document(self, doc_id: str):
175
325
  async with HaikuRAG(
176
- db_path=self.db_path, config=self.config, allow_create=False
326
+ db_path=self.db_path,
327
+ config=self.config,
328
+ read_only=self.read_only,
329
+ before=self.before,
177
330
  ) as self.client:
178
331
  doc = await self.client.get_document_by_id(doc_id)
179
332
  if doc is None:
@@ -182,7 +335,12 @@ class HaikuRAGApp:
182
335
  self._rich_print_document(doc, truncate=False)
183
336
 
184
337
  async def delete_document(self, doc_id: str):
185
- async with HaikuRAG(db_path=self.db_path, config=self.config) as self.client:
338
+ async with HaikuRAG(
339
+ db_path=self.db_path,
340
+ config=self.config,
341
+ read_only=self.read_only,
342
+ before=self.before,
343
+ ) as self.client:
186
344
  deleted = await self.client.delete_document(doc_id)
187
345
  if deleted:
188
346
  self.console.print(
@@ -193,16 +351,58 @@ class HaikuRAGApp:
193
351
  f"[yellow]Document with id {doc_id} not found.[/yellow]"
194
352
  )
195
353
 
196
- async def search(self, query: str, limit: int = 5, filter: str | None = None):
354
+ async def search(
355
+ self, query: str, limit: int | None = None, filter: str | None = None
356
+ ):
197
357
  async with HaikuRAG(
198
- db_path=self.db_path, config=self.config, allow_create=False
358
+ db_path=self.db_path,
359
+ config=self.config,
360
+ read_only=self.read_only,
361
+ before=self.before,
199
362
  ) as self.client:
200
363
  results = await self.client.search(query, limit=limit, filter=filter)
201
364
  if not results:
202
365
  self.console.print("[yellow]No results found.[/yellow]")
203
366
  return
204
- for chunk, score in results:
205
- self._rich_print_search_result(chunk, score)
367
+ for result in results:
368
+ self._rich_print_search_result(result)
369
+
370
+ async def visualize_chunk(self, chunk_id: str):
371
+ """Display visual grounding images for a chunk."""
372
+ from textual_image.renderable import Image as RichImage
373
+
374
+ async with HaikuRAG(
375
+ db_path=self.db_path,
376
+ config=self.config,
377
+ read_only=self.read_only,
378
+ before=self.before,
379
+ ) as self.client:
380
+ chunk = await self.client.chunk_repository.get_by_id(chunk_id)
381
+ if not chunk:
382
+ self.console.print(f"[red]Chunk with id {chunk_id} not found.[/red]")
383
+ return
384
+
385
+ images = await self.client.visualize_chunk(chunk)
386
+ if not images:
387
+ self.console.print(
388
+ "[yellow]No visual grounding available for this chunk.[/yellow]"
389
+ )
390
+ self.console.print(
391
+ "This may be because the document was converted without page images."
392
+ )
393
+ return
394
+
395
+ self.console.print(f"[bold]Visual grounding for chunk {chunk_id}[/bold]")
396
+ if chunk.document_uri:
397
+ self.console.print(
398
+ f"[repr.attrib_name]document[/repr.attrib_name]: {chunk.document_uri}"
399
+ )
400
+
401
+ for i, img in enumerate(images):
402
+ self.console.print(
403
+ f"\n[bold cyan]Page {i + 1}/{len(images)}[/bold cyan]"
404
+ )
405
+ self.console.print(RichImage(img))
206
406
 
207
407
  async def ask(
208
408
  self,
@@ -210,6 +410,7 @@ class HaikuRAGApp:
210
410
  cite: bool = False,
211
411
  deep: bool = False,
212
412
  verbose: bool = False,
413
+ filter: str | None = None,
213
414
  ):
214
415
  """Ask a question using the RAG system.
215
416
 
@@ -218,56 +419,87 @@ class HaikuRAGApp:
218
419
  cite: Include citations in the answer
219
420
  deep: Use deep QA mode (multi-step reasoning)
220
421
  verbose: Show verbose output
422
+ filter: SQL WHERE clause to filter documents
221
423
  """
222
424
  async with HaikuRAG(
223
- db_path=self.db_path, config=self.config, allow_create=False
425
+ db_path=self.db_path,
426
+ config=self.config,
427
+ read_only=self.read_only,
428
+ before=self.before,
224
429
  ) as self.client:
225
430
  try:
431
+ citations = []
226
432
  if deep:
227
- from haiku.rag.graph.deep_qa.dependencies import DeepQAContext
228
- from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph
229
- from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState
230
-
231
- graph = build_deep_qa_graph(config=self.config)
232
- context = DeepQAContext(
233
- original_question=question, use_citations=cite
433
+ from haiku.rag.graph.research.models import ResearchReport
434
+
435
+ graph = build_research_graph(config=self.config)
436
+ context = ResearchContext(original_question=question)
437
+ state = ResearchState.from_config(
438
+ context=context,
439
+ config=self.config,
440
+ max_iterations=2,
441
+ confidence_threshold=0.0,
234
442
  )
235
- state = DeepQAState.from_config(context=context, config=self.config)
236
- deps = DeepQADeps(client=self.client)
443
+ state.search_filter = filter
444
+ deps = ResearchDeps(client=self.client)
237
445
 
238
446
  if verbose:
239
- # Use AG-UI renderer to process and display events
240
- from haiku.rag.graph.agui import AGUIConsoleRenderer
241
-
242
447
  renderer = AGUIConsoleRenderer(self.console)
243
448
  result_dict = await renderer.render(
244
449
  stream_graph(graph, state, deps)
245
450
  )
246
- # Result should be a dict with 'answer' key
247
- answer = result_dict.get("answer", "") if result_dict else ""
451
+ report = (
452
+ ResearchReport.model_validate(result_dict)
453
+ if result_dict
454
+ else None
455
+ )
456
+ else:
457
+ report = await graph.run(state=state, deps=deps)
458
+
459
+ self.console.print(f"[bold blue]Question:[/bold blue] {question}")
460
+ self.console.print()
461
+ if report:
462
+ self.console.print("[bold green]Answer:[/bold green]")
463
+ self.console.print(Markdown(report.executive_summary))
464
+ if report.main_findings:
465
+ self.console.print()
466
+ self.console.print("[bold cyan]Key Findings:[/bold cyan]")
467
+ for finding in report.main_findings:
468
+ self.console.print(f"• {finding}")
469
+ if report.sources_summary:
470
+ self.console.print()
471
+ self.console.print("[bold cyan]Sources:[/bold cyan]")
472
+ self.console.print(report.sources_summary)
248
473
  else:
249
- # Run without rendering events, just get the result
250
- result = await graph.run(state=state, deps=deps)
251
- answer = result.answer
474
+ self.console.print("[yellow]No answer generated.[/yellow]")
252
475
  else:
253
- answer = await self.client.ask(question, cite=cite)
476
+ answer, citations = await self.client.ask(question, filter=filter)
254
477
 
255
- self.console.print(f"[bold blue]Question:[/bold blue] {question}")
256
- self.console.print()
257
- self.console.print("[bold green]Answer:[/bold green]")
258
- self.console.print(Markdown(answer))
478
+ self.console.print(f"[bold blue]Question:[/bold blue] {question}")
479
+ self.console.print()
480
+ self.console.print("[bold green]Answer:[/bold green]")
481
+ self.console.print(Markdown(answer))
482
+ if cite and citations:
483
+ for renderable in format_citations_rich(citations):
484
+ self.console.print(renderable)
259
485
  except Exception as e:
260
486
  self.console.print(f"[red]Error: {e}[/red]")
261
487
 
262
- async def research(self, question: str, verbose: bool = False):
488
+ async def research(
489
+ self, question: str, verbose: bool = False, filter: str | None = None
490
+ ):
263
491
  """Run research via the pydantic-graph pipeline.
264
492
 
265
493
  Args:
266
494
  question: The research question
267
495
  verbose: Show AG-UI event stream during execution
496
+ filter: SQL WHERE clause to filter documents
268
497
  """
269
498
  async with HaikuRAG(
270
- db_path=self.db_path, config=self.config, allow_create=False
499
+ db_path=self.db_path,
500
+ config=self.config,
501
+ read_only=self.read_only,
502
+ before=self.before,
271
503
  ) as client:
272
504
  try:
273
505
  self.console.print("[bold cyan]Starting research[/bold cyan]")
@@ -277,6 +509,7 @@ class HaikuRAGApp:
277
509
  graph = build_research_graph(config=self.config)
278
510
  context = ResearchContext(original_question=question)
279
511
  state = ResearchState.from_config(context=context, config=self.config)
512
+ state.search_filter = filter
280
513
  deps = ResearchDeps(client=client)
281
514
 
282
515
  if verbose:
@@ -356,9 +589,13 @@ class HaikuRAGApp:
356
589
  except Exception as e:
357
590
  self.console.print(f"[red]Error during research: {e}[/red]")
358
591
 
359
- async def rebuild(self):
592
+ async def rebuild(self, mode: RebuildMode = RebuildMode.FULL):
360
593
  async with HaikuRAG(
361
- db_path=self.db_path, config=self.config, skip_validation=True
594
+ db_path=self.db_path,
595
+ config=self.config,
596
+ skip_validation=True,
597
+ read_only=self.read_only,
598
+ before=self.before,
362
599
  ) as client:
363
600
  try:
364
601
  documents = await client.list_documents()
@@ -370,12 +607,18 @@ class HaikuRAGApp:
370
607
  )
371
608
  return
372
609
 
610
+ mode_desc = {
611
+ RebuildMode.FULL: "full rebuild",
612
+ RebuildMode.RECHUNK: "rechunk",
613
+ RebuildMode.EMBED_ONLY: "embed only",
614
+ }[mode]
615
+
373
616
  self.console.print(
374
- f"[bold cyan]Rebuilding database with {total_docs} documents...[/bold cyan]"
617
+ f"[bold cyan]Rebuilding database ({mode_desc}) with {total_docs} documents...[/bold cyan]"
375
618
  )
376
619
  with Progress() as progress:
377
620
  task = progress.add_task("Rebuilding...", total=total_docs)
378
- async for _ in client.rebuild_database():
621
+ async for _ in client.rebuild_database(mode=mode):
379
622
  progress.update(task, advance=1)
380
623
 
381
624
  self.console.print(
@@ -388,7 +631,11 @@ class HaikuRAGApp:
388
631
  """Run database maintenance: optimize and cleanup table history."""
389
632
  try:
390
633
  async with HaikuRAG(
391
- db_path=self.db_path, config=self.config, skip_validation=True
634
+ db_path=self.db_path,
635
+ config=self.config,
636
+ skip_validation=True,
637
+ read_only=self.read_only,
638
+ before=self.before,
392
639
  ) as client:
393
640
  await client.vacuum()
394
641
  self.console.print(
@@ -397,6 +644,100 @@ class HaikuRAGApp:
397
644
  except Exception as e:
398
645
  self.console.print(f"[red]Error during vacuum: {e}[/red]")
399
646
 
647
+ async def create_index(self):
648
+ """Create vector index on the chunks table."""
649
+ try:
650
+ async with HaikuRAG(
651
+ db_path=self.db_path,
652
+ config=self.config,
653
+ skip_validation=True,
654
+ read_only=self.read_only,
655
+ before=self.before,
656
+ ) as client:
657
+ row_count = client.store.chunks_table.count_rows()
658
+ self.console.print(f"Chunks in database: {row_count}")
659
+
660
+ if row_count < 256:
661
+ self.console.print(
662
+ f"[yellow]Warning: Need at least 256 chunks to create an index (have {row_count})[/yellow]"
663
+ )
664
+ return
665
+
666
+ # Check if index already exists
667
+ indices = client.store.chunks_table.list_indices()
668
+ has_vector_index = any("vector" in str(idx).lower() for idx in indices)
669
+
670
+ if has_vector_index:
671
+ self.console.print(
672
+ "[yellow]Rebuilding existing vector index...[/yellow]"
673
+ )
674
+ else:
675
+ self.console.print("[bold]Creating vector index...[/bold]")
676
+
677
+ client.store._ensure_vector_index()
678
+ self.console.print(
679
+ "[bold green]Vector index created successfully.[/bold green]"
680
+ )
681
+ except Exception as e:
682
+ self.console.print(f"[red]Error creating index: {e}[/red]")
683
+
684
+ async def download_models(self):
685
+ """Download Docling, HuggingFace tokenizer, and Ollama models per config."""
686
+ from haiku.rag.client import HaikuRAG
687
+
688
+ client = HaikuRAG(db_path=None, config=self.config)
689
+
690
+ progress: Progress | None = None
691
+ task_id: TaskID | None = None
692
+ current_model = ""
693
+ current_digest = ""
694
+
695
+ async for event in client.download_models():
696
+ if event.status == "start":
697
+ self.console.print(
698
+ f"[bold blue]Downloading {event.model}...[/bold blue]"
699
+ )
700
+ elif event.status == "done":
701
+ if progress:
702
+ progress.stop()
703
+ progress = None
704
+ task_id = None
705
+ self.console.print(f"[green]✓[/green] {event.model}")
706
+ current_model = ""
707
+ current_digest = ""
708
+ elif event.status == "pulling":
709
+ self.console.print(f"[bold blue]Pulling {event.model}...[/bold blue]")
710
+ current_model = event.model
711
+ progress = Progress(
712
+ SpinnerColumn(),
713
+ TextColumn("[progress.description]{task.description}"),
714
+ BarColumn(),
715
+ DownloadColumn(),
716
+ TransferSpeedColumn(),
717
+ console=self.console,
718
+ transient=True,
719
+ auto_refresh=False,
720
+ )
721
+ progress.start()
722
+ task_id = progress.add_task(event.model, total=None)
723
+ elif event.status == "downloading" and progress and task_id is not None:
724
+ if event.digest != current_digest:
725
+ current_digest = event.digest
726
+ short_digest = event.digest[:19] if event.digest else ""
727
+ progress.update(
728
+ task_id,
729
+ description=f"{current_model} ({short_digest})",
730
+ total=event.total,
731
+ completed=0,
732
+ )
733
+ progress.update(task_id, completed=event.completed, refresh=True)
734
+ elif progress and task_id is not None:
735
+ progress.update(
736
+ task_id,
737
+ description=f"{current_model}: {event.status}",
738
+ refresh=True,
739
+ )
740
+
400
741
  def show_settings(self):
401
742
  """Display current configuration settings."""
402
743
  self.console.print("[bold]haiku.rag configuration[/bold]")
@@ -447,22 +788,27 @@ class HaikuRAGApp:
447
788
  self.console.print(content)
448
789
  self.console.rule()
449
790
 
450
- def _rich_print_search_result(self, chunk: Chunk, score: float):
451
- """Format a search result chunk for display."""
452
- content = Markdown(chunk.content)
791
+ def _rich_print_search_result(self, result: "SearchResult"):
792
+ """Format a search result for display."""
793
+ content = Markdown(result.content)
453
794
  self.console.print(
454
- f"[repr.attrib_name]document_id[/repr.attrib_name]: {chunk.document_id} "
455
- f"[repr.attrib_name]score[/repr.attrib_name]: {score:.4f}"
795
+ f"[repr.attrib_name]document_id[/repr.attrib_name]: {result.document_id} "
796
+ f"[repr.attrib_name]chunk_id[/repr.attrib_name]: {result.chunk_id} "
797
+ f"[repr.attrib_name]score[/repr.attrib_name]: {result.score:.4f}"
456
798
  )
457
- if chunk.document_uri:
458
- self.console.print("[repr.attrib_name]document uri[/repr.attrib_name]:")
459
- self.console.print(chunk.document_uri)
460
- if chunk.document_title:
799
+ if result.document_uri:
800
+ self.console.print(
801
+ f"[repr.attrib_name]document uri[/repr.attrib_name]: {result.document_uri}"
802
+ )
803
+ if result.document_title:
461
804
  self.console.print("[repr.attrib_name]document title[/repr.attrib_name]:")
462
- self.console.print(chunk.document_title)
463
- if chunk.document_meta:
464
- self.console.print("[repr.attrib_name]document meta[/repr.attrib_name]:")
465
- self.console.print(chunk.document_meta)
805
+ self.console.print(result.document_title)
806
+ if result.page_numbers:
807
+ self.console.print("[repr.attrib_name]pages[/repr.attrib_name]:")
808
+ self.console.print(", ".join(str(p) for p in result.page_numbers))
809
+ if result.headings:
810
+ self.console.print("[repr.attrib_name]headings[/repr.attrib_name]:")
811
+ self.console.print(" > ".join(result.headings))
466
812
  self.console.print("[repr.attrib_name]content[/repr.attrib_name]:")
467
813
  self.console.print(content)
468
814
  self.console.rule()
@@ -476,18 +822,30 @@ class HaikuRAGApp:
476
822
  enable_agui: bool = False,
477
823
  ):
478
824
  """Start the server with selected services."""
479
- async with HaikuRAG(self.db_path, config=self.config) as client:
825
+ async with HaikuRAG(
826
+ self.db_path,
827
+ config=self.config,
828
+ read_only=self.read_only,
829
+ before=self.before,
830
+ ) as client:
480
831
  tasks = []
481
832
 
482
- # Start file monitor if enabled
833
+ # Start file monitor if enabled (not available in read-only mode)
483
834
  if enable_monitor:
484
- monitor = FileWatcher(client=client, config=self.config)
485
- monitor_task = asyncio.create_task(monitor.observe())
486
- tasks.append(monitor_task)
835
+ if self.read_only:
836
+ logger.warning(
837
+ "File monitor disabled: cannot monitor files in read-only mode"
838
+ )
839
+ else:
840
+ monitor = FileWatcher(client=client, config=self.config)
841
+ monitor_task = asyncio.create_task(monitor.observe())
842
+ tasks.append(monitor_task)
487
843
 
488
844
  # Start MCP server if enabled
489
845
  if enable_mcp:
490
- server = create_mcp_server(self.db_path, config=self.config)
846
+ server = create_mcp_server(
847
+ self.db_path, config=self.config, read_only=self.read_only
848
+ )
491
849
 
492
850
  async def run_mcp():
493
851
  if mcp_transport == "stdio":