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.
- haiku/rag/app.py +430 -72
- haiku/rag/chunkers/__init__.py +31 -0
- haiku/rag/chunkers/base.py +31 -0
- haiku/rag/chunkers/docling_local.py +164 -0
- haiku/rag/chunkers/docling_serve.py +179 -0
- haiku/rag/cli.py +207 -24
- haiku/rag/cli_chat.py +489 -0
- haiku/rag/client.py +1251 -266
- haiku/rag/config/__init__.py +16 -10
- haiku/rag/config/loader.py +5 -44
- haiku/rag/config/models.py +126 -17
- haiku/rag/converters/__init__.py +31 -0
- haiku/rag/converters/base.py +63 -0
- haiku/rag/converters/docling_local.py +193 -0
- haiku/rag/converters/docling_serve.py +229 -0
- haiku/rag/converters/text_utils.py +237 -0
- haiku/rag/embeddings/__init__.py +123 -24
- haiku/rag/embeddings/voyageai.py +175 -20
- haiku/rag/graph/__init__.py +0 -11
- haiku/rag/graph/agui/__init__.py +8 -2
- haiku/rag/graph/agui/cli_renderer.py +1 -1
- haiku/rag/graph/agui/emitter.py +219 -31
- haiku/rag/graph/agui/server.py +20 -62
- haiku/rag/graph/agui/stream.py +1 -2
- haiku/rag/graph/research/__init__.py +5 -2
- haiku/rag/graph/research/dependencies.py +12 -126
- haiku/rag/graph/research/graph.py +390 -135
- haiku/rag/graph/research/models.py +91 -112
- haiku/rag/graph/research/prompts.py +99 -91
- haiku/rag/graph/research/state.py +35 -27
- haiku/rag/inspector/__init__.py +8 -0
- haiku/rag/inspector/app.py +259 -0
- haiku/rag/inspector/widgets/__init__.py +6 -0
- haiku/rag/inspector/widgets/chunk_list.py +100 -0
- haiku/rag/inspector/widgets/context_modal.py +89 -0
- haiku/rag/inspector/widgets/detail_view.py +130 -0
- haiku/rag/inspector/widgets/document_list.py +75 -0
- haiku/rag/inspector/widgets/info_modal.py +209 -0
- haiku/rag/inspector/widgets/search_modal.py +183 -0
- haiku/rag/inspector/widgets/visual_modal.py +126 -0
- haiku/rag/mcp.py +106 -102
- haiku/rag/monitor.py +33 -9
- haiku/rag/providers/__init__.py +5 -0
- haiku/rag/providers/docling_serve.py +108 -0
- haiku/rag/qa/__init__.py +12 -10
- haiku/rag/qa/agent.py +43 -61
- haiku/rag/qa/prompts.py +35 -57
- haiku/rag/reranking/__init__.py +9 -6
- haiku/rag/reranking/base.py +1 -1
- haiku/rag/reranking/cohere.py +5 -4
- haiku/rag/reranking/mxbai.py +5 -2
- haiku/rag/reranking/vllm.py +3 -4
- haiku/rag/reranking/zeroentropy.py +6 -5
- haiku/rag/store/__init__.py +2 -1
- haiku/rag/store/engine.py +242 -42
- haiku/rag/store/exceptions.py +4 -0
- haiku/rag/store/models/__init__.py +8 -2
- haiku/rag/store/models/chunk.py +190 -0
- haiku/rag/store/models/document.py +46 -0
- haiku/rag/store/repositories/chunk.py +141 -121
- haiku/rag/store/repositories/document.py +25 -84
- haiku/rag/store/repositories/settings.py +11 -14
- haiku/rag/store/upgrades/__init__.py +19 -3
- haiku/rag/store/upgrades/v0_10_1.py +1 -1
- haiku/rag/store/upgrades/v0_19_6.py +65 -0
- haiku/rag/store/upgrades/v0_20_0.py +68 -0
- haiku/rag/store/upgrades/v0_23_1.py +100 -0
- haiku/rag/store/upgrades/v0_9_3.py +3 -3
- haiku/rag/utils.py +371 -146
- {haiku_rag_slim-0.16.0.dist-info → haiku_rag_slim-0.24.0.dist-info}/METADATA +15 -12
- haiku_rag_slim-0.24.0.dist-info/RECORD +78 -0
- {haiku_rag_slim-0.16.0.dist-info → haiku_rag_slim-0.24.0.dist-info}/WHEEL +1 -1
- haiku/rag/chunker.py +0 -65
- haiku/rag/embeddings/base.py +0 -25
- haiku/rag/embeddings/ollama.py +0 -28
- haiku/rag/embeddings/openai.py +0 -26
- haiku/rag/embeddings/vllm.py +0 -29
- haiku/rag/graph/agui/events.py +0 -254
- haiku/rag/graph/common/__init__.py +0 -5
- haiku/rag/graph/common/models.py +0 -42
- haiku/rag/graph/common/nodes.py +0 -265
- haiku/rag/graph/common/prompts.py +0 -46
- haiku/rag/graph/common/utils.py +0 -44
- haiku/rag/graph/deep_qa/__init__.py +0 -1
- haiku/rag/graph/deep_qa/dependencies.py +0 -27
- haiku/rag/graph/deep_qa/graph.py +0 -243
- haiku/rag/graph/deep_qa/models.py +0 -20
- haiku/rag/graph/deep_qa/prompts.py +0 -59
- haiku/rag/graph/deep_qa/state.py +0 -56
- haiku/rag/graph/research/common.py +0 -87
- haiku/rag/reader.py +0 -135
- haiku_rag_slim-0.16.0.dist-info/RECORD +0 -71
- {haiku_rag_slim-0.16.0.dist-info → haiku_rag_slim-0.24.0.dist-info}/entry_points.txt +0 -0
- {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
|
|
10
|
-
|
|
11
|
-
|
|
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__(
|
|
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
|
-
#
|
|
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
|
-
|
|
83
|
-
|
|
84
|
-
|
|
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
|
-
|
|
87
|
-
|
|
88
|
-
|
|
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,
|
|
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(
|
|
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(
|
|
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,
|
|
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(
|
|
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(
|
|
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,
|
|
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
|
|
205
|
-
self._rich_print_search_result(
|
|
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,
|
|
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.
|
|
228
|
-
|
|
229
|
-
|
|
230
|
-
|
|
231
|
-
|
|
232
|
-
|
|
233
|
-
|
|
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 =
|
|
236
|
-
deps =
|
|
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
|
-
|
|
247
|
-
|
|
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
|
-
|
|
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,
|
|
476
|
+
answer, citations = await self.client.ask(question, filter=filter)
|
|
254
477
|
|
|
255
|
-
|
|
256
|
-
|
|
257
|
-
|
|
258
|
-
|
|
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(
|
|
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,
|
|
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,
|
|
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,
|
|
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,
|
|
451
|
-
"""Format a search result
|
|
452
|
-
content = Markdown(
|
|
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]: {
|
|
455
|
-
f"[repr.attrib_name]
|
|
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
|
|
458
|
-
self.console.print(
|
|
459
|
-
|
|
460
|
-
|
|
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(
|
|
463
|
-
if
|
|
464
|
-
self.console.print("[repr.attrib_name]
|
|
465
|
-
self.console.print(
|
|
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(
|
|
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
|
-
|
|
485
|
-
|
|
486
|
-
|
|
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(
|
|
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":
|