danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

docs: add docs

Browse source 
chore: switch logging to custom logger
This commit is contained in:
bryan davis 2026-04-06 15:32:37 -05:00
parent e26a452da5
commit 4a987ed6d1
12 changed files with 132 additions and 38 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
app/haiku.rag.yaml.example
View file

@ -34,6 +34,7 @@ embeddings:
search:
limit: 5
context_radius: 0
# context_expansion_mode: auto # auto, chunks, or disabled
# Provider settings
providers:

3
docs/apps.md
View file

@ -182,8 +182,9 @@ Search uses hybrid (vector + full-text) search across all chunks.
Press `c` while viewing a chunk to see the expanded context that would be provided to the QA agent:
- Type-aware expansion: tables, code blocks, and lists expand to their complete structures
- Type-aware expansion: tables, code blocks, and lists expand to their complete structures (when `search.context_expansion_mode` is `auto`)
- Text content expands based on `search.context_radius` setting
- Set `search.context_expansion_mode` to `chunks` for faster expansion on large corpora, or `disabled` to skip expansion
- Includes metadata like source document, content type, and relevance score
### Visual Grounding

1
docs/configuration/index.md
View file

@ -100,6 +100,7 @@ research:
search:
limit: 10 # Default number of results to return
context_radius: 0 # DocItems before/after to include for text content
context_expansion_mode: auto # auto, chunks, or disabled
max_context_items: 10 # Maximum items in expanded context
max_context_chars: 10000 # Maximum characters in expanded context
vector_index_metric: cosine # cosine, l2, or dot

7
docs/configuration/qa-research.md
View file

@ -8,16 +8,21 @@ Configure search behavior and context expansion:
search:
limit: 10 # Default number of results to return
context_radius: 0 # DocItems before/after to include for text content
context_expansion_mode: auto # auto, chunks, or disabled
max_context_items: 10 # Maximum items in expanded context
max_context_chars: 10000 # Maximum characters in expanded context
```
- **limit**: Default number of search results to return when no limit is specified. Used by CLI, MCP server, QA, and research workflows. Default: 10
- **context_radius**: For text content (paragraphs), includes N DocItems before and after. Set to 0 to disable expansion (default).
- **context_expansion_mode**: Controls which expansion strategy is used. Default: `auto`.
- `auto`: Prefer DoclingDocument-based expansion when doc_item_refs exist (structure-aware), fall back to chunk-based expansion.
- `chunks`: Always use chunk-based expansion. Skips DoclingDocument decompression, which can be significantly faster for large corpora.
- `disabled`: No expansion at all — return search results as-is.
- **max_context_items**: Limits how many document items (paragraphs, list items, etc.) can be included in expanded context. Default: 10.
- **max_context_chars**: Hard limit on total characters in expanded content. Default: 10000.
Structural content (tables, code blocks, lists) uses type-aware expansion that automatically includes the complete structure regardless of how it was chunked.
Structural content (tables, code blocks, lists) uses type-aware expansion that automatically includes the complete structure regardless of how it was chunked. This applies when `context_expansion_mode` is `auto`.
!!! note "Reranking behavior"
When a reranker is configured, search automatically retrieves 10x the requested limit, then reranks to return the final count. This improves result quality without requiring you to adjust `limit`.

3
docs/python.md
View file

@ -376,10 +376,11 @@ for result in expanded_results:
Context expansion uses your configuration settings:
- **search.context_radius**: For text content (paragraphs), includes N DocItems before and after
- **search.context_expansion_mode**: Controls the expansion strategy — `auto` (default, structure-aware with fallback), `chunks` (chunk-based only, faster for large corpora), or `disabled` (no expansion)
- **search.max_context_items**: Limits how many document items can be included
- **search.max_context_chars**: Hard limit on total characters
**Type-aware expansion**: Structural content (tables, code blocks, lists) automatically expands to include the complete structure, regardless of how it was split during chunking.
**Type-aware expansion**: Structural content (tables, code blocks, lists) automatically expands to include the complete structure, regardless of how it was split during chunking. This applies when `context_expansion_mode` is `auto`.
**Smart Merging**: When expanded chunks overlap or are adjacent within the same document, they are automatically merged into single chunks with continuous content. This eliminates duplication and provides coherent text blocks. The merged chunk uses the highest relevance score from the original chunks.

2
docs/tuning.md
View file

@ -28,6 +28,8 @@ When configured, a cross-encoder reranker re-scores 10x the requested candidates
`context_radius` expands text chunks with neighboring document items. Structural content (tables, code blocks, lists) expands automatically to include the complete structure. This setting matters most with small `chunk_size` values, where individual chunks may lack sufficient context. `max_context_items` and `max_context_chars` cap expansion to prevent context bloat.
`context_expansion_mode` controls which expansion strategy is used. The default `auto` uses DoclingDocument-based expansion (structure-aware) when available, falling back to chunk-based expansion. Set to `chunks` to skip DoclingDocument decompression for faster performance on large corpora, or `disabled` to skip expansion entirely.
## Tuning Generation
Model and temperature selection affect answer quality directly — see [Providers](configuration/providers.md#model-settings) for options.

7
haiku_rag_slim/haiku/rag/agents/qa/agent.py
View file

@ -18,6 +18,7 @@ from haiku.rag.tools.search import create_search_toolset
from haiku.rag.utils import get_model
logger = logging.getLogger(__name__)
perf_logger = logging.getLogger("haiku.rag.perf")
@dataclass
@ -81,15 +82,15 @@ class QuestionAnswerAgent:
t0 = time.perf_counter()
result = await agent.run(question, deps=deps)
agent_duration = time.perf_counter() - t0
logger.info("qa.agent_run took %.3fs", agent_duration)
perf_logger.debug("qa.agent_run took %.3fs", agent_duration)
t0 = time.perf_counter()
output = result.output
citations = resolve_citations(output.cited_chunks, accumulated_results)
logger.info(
perf_logger.debug(
"qa.resolve_citations count=%d took %.3fs",
len(citations),
time.perf_counter() - t0,
)
logger.info("qa.answer completed total=%.3fs", agent_duration)
perf_logger.debug("qa.answer completed total=%.3fs", agent_duration)
return output.answer, citations

27
haiku_rag_slim/haiku/rag/client.py
View file

@ -38,6 +38,7 @@ if TYPE_CHECKING:
from haiku.rag.agents.rlm.models import RLMResult
logger = logging.getLogger(__name__)
perf_logger = logging.getLogger("haiku.rag.perf")
class RebuildMode(Enum):
@ -1075,7 +1076,7 @@ class HaikuRAG:
# Step 1: Get reranker
t0 = time.perf_counter()
reranker = get_reranker(config=self._config)
logger.info(
perf_logger.debug(
"search.reranker_init took %.3fs",
time.perf_counter() - t0,
)
@ -1086,7 +1087,7 @@ class HaikuRAG:
chunk_results = await self.chunk_repository.search(
query, limit, search_type, filter
)
logger.info(
perf_logger.debug(
"search.chunk_search type=%s limit=%d results=%d took %.3fs",
search_type,
limit,
@ -1098,7 +1099,7 @@ class HaikuRAG:
raw_results = await self.chunk_repository.search(
query, search_limit, search_type, filter
)
logger.info(
perf_logger.debug(
"search.chunk_search type=%s limit=%d results=%d took %.3fs",
search_type,
search_limit,
@ -1110,7 +1111,7 @@ class HaikuRAG:
t0 = time.perf_counter()
chunks = [chunk for chunk, _ in raw_results]
chunk_results = await reranker.rerank(query, chunks, top_n=limit)
logger.info(
perf_logger.debug(
"search.rerank candidates=%d top_n=%d took %.3fs",
len(chunks),
limit,
@ -1122,14 +1123,14 @@ class HaikuRAG:
results = [
SearchResult.from_chunk(chunk, score) for chunk, score in chunk_results
]
logger.info(
perf_logger.debug(
"search.build_results count=%d took %.3fs",
len(results),
time.perf_counter() - t0,
)
duration_s = time.perf_counter() - search_start
logger.info(
perf_logger.debug(
"search completed query=%r type=%s results=%d duration=%.3fs",
query[:80],
search_type,
@ -1181,7 +1182,7 @@ class HaikuRAG:
max_chars = self._config.search.max_context_chars
if mode == "disabled":
logger.info("expand.disabled, skipping")
perf_logger.debug("expand.disabled, skipping")
return search_results
# Group by document_id for efficient processing
@ -1192,7 +1193,7 @@ class HaikuRAG:
document_groups[doc_id] = []
document_groups[doc_id].append(result)
logger.info(
perf_logger.debug(
"expand.groups docs=%d results=%d mode=%s",
len(document_groups),
len(search_results),
@ -1213,7 +1214,7 @@ class HaikuRAG:
# Only fetch docling data (skip content blob)
t0 = time.perf_counter()
doc = await self.document_repository.get_docling_data(doc_id)
logger.info(
perf_logger.debug(
"expand.fetch_docling_data doc=%s took %.3fs",
doc_id[:8],
time.perf_counter() - t0,
@ -1221,7 +1222,7 @@ class HaikuRAG:
if doc is not None:
t0 = time.perf_counter()
docling_doc = doc.get_docling_document()
logger.info(
perf_logger.debug(
"expand.get_docling_document doc=%s took %.3fs",
doc_id[:8],
time.perf_counter() - t0,
@ -1237,7 +1238,7 @@ class HaikuRAG:
max_items,
max_chars,
)
logger.info(
perf_logger.debug(
"expand._expand_with_docling doc=%s results=%d->%d took %.3fs",
doc_id[:8],
len(doc_results),
@ -1252,7 +1253,7 @@ class HaikuRAG:
expanded = await self._expand_with_chunks(
doc_id, doc_results, radius
)
logger.info(
perf_logger.debug(
"expand._expand_with_chunks doc=%s results=%d->%d took %.3fs",
doc_id[:8],
len(doc_results),
@ -1263,7 +1264,7 @@ class HaikuRAG:
else:
expanded_results.extend(doc_results)
logger.info(
perf_logger.debug(
"expand.total results=%d took %.3fs",
len(expanded_results),
time.perf_counter() - expand_start,

11
haiku_rag_slim/haiku/rag/store/models/document.py
View file

@ -12,6 +12,7 @@ if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument
logger = logging.getLogger(__name__)
perf_logger = logging.getLogger("haiku.rag.perf")
_docling_document_cache: LRUCache[str, "DoclingDocument"] = LRUCache(maxsize=100)
@ -29,7 +30,7 @@ def configure_docling_cache(maxsize: int) -> None:
# Copy existing entries (LRU order preserved by iteration)
for key in old:
_docling_document_cache[key] = old[key]
logger.info("docling.cache_resized maxsize=%d", maxsize)
perf_logger.debug("docling.cache_resized maxsize=%d", maxsize)
def _get_cached_docling_document(
@ -37,12 +38,12 @@ def _get_cached_docling_document(
) -> "DoclingDocument":
"""Get or parse DoclingDocument with LRU caching by document ID."""
if document_id in _docling_document_cache:
logger.info("docling.cache_hit doc=%s", document_id[:8])
perf_logger.debug("docling.cache_hit doc=%s", document_id[:8])
return _docling_document_cache[document_id]
from docling_core.types.doc.document import DoclingDocument
logger.info(
perf_logger.debug(
"docling.cache_miss doc=%s cache_size=%d/%d",
document_id[:8],
len(_docling_document_cache),
@ -52,7 +53,7 @@ def _get_cached_docling_document(
t0 = time.perf_counter()
json_str = decompress_json(compressed_data)
decompress_time = time.perf_counter() - t0
logger.info(
perf_logger.debug(
"docling.decompress doc=%s bytes=%d json_chars=%d took %.3fs",
document_id[:8],
len(compressed_data),
@ -63,7 +64,7 @@ def _get_cached_docling_document(
t0 = time.perf_counter()
doc = DoclingDocument.model_validate_json(json_str)
validate_time = time.perf_counter() - t0
logger.info(
perf_logger.debug(
"docling.model_validate doc=%s took %.3fs",
document_id[:8],
validate_time,

19
haiku_rag_slim/haiku/rag/store/repositories/chunk.py
View file

@ -17,6 +17,7 @@ from haiku.rag.store.engine import DocumentRecord, Store
from haiku.rag.store.models.chunk import Chunk
logger = logging.getLogger(__name__)
perf_logger = logging.getLogger("haiku.rag.perf")
class ChunkRepository:
@ -255,7 +256,7 @@ class ChunkRepository:
return []
# Keep as pandas Series for efficient vectorized operations
filtered_doc_ids = docs_df["id"]
logger.info(
perf_logger.debug(
"search.filter docs=%d took %.3fs",
len(filtered_doc_ids),
time.perf_counter() - t0,
@ -265,7 +266,7 @@ class ChunkRepository:
if search_type == "vector":
t0 = time.perf_counter()
query_embedding = await self.embedder.embed_query(query)
logger.info(
perf_logger.debug(
"search.embed took %.3fs", time.perf_counter() - t0
)
vector_query = cast(
@ -284,7 +285,7 @@ class ChunkRepository:
else: # hybrid (default)
t0 = time.perf_counter()
query_embedding = await self.embedder.embed_query(query)
logger.info(
perf_logger.debug(
"search.embed took %.3fs", time.perf_counter() - t0
)
# Create RRF reranker
@ -304,14 +305,14 @@ class ChunkRepository:
if filtered_doc_ids is not None:
t0 = time.perf_counter()
chunks_df = results.to_pandas()
logger.info(
perf_logger.debug(
"search.execute took %.3fs", time.perf_counter() - t0
)
t0 = time.perf_counter()
filtered_chunks_df = chunks_df.loc[
chunks_df["document_id"].isin(filtered_doc_ids)
].head(limit)
logger.info(
perf_logger.debug(
"search.doc_filter rows=%d->%d took %.3fs",
len(chunks_df),
len(filtered_chunks_df),
@ -506,7 +507,7 @@ class ChunkRepository:
# Convert LanceDB query result to DataFrame
# (this is where the actual DB query executes)
df = query_result.to_pandas()
logger.info(
perf_logger.debug(
"search.execute rows=%d took %.3fs",
len(df),
time.perf_counter() - t0,
@ -530,7 +531,7 @@ class ChunkRepository:
)
for row in rows
]
logger.info(
perf_logger.debug(
"search.to_records count=%d took %.3fs",
len(pydantic_results),
time.perf_counter() - t0,
@ -552,7 +553,7 @@ class ChunkRepository:
.to_pydantic(DocumentRecord)
)
documents_map = {doc.id: doc for doc in doc_results}
logger.info(
perf_logger.debug(
"search.doc_lookup docs=%d took %.3fs",
len(documents_map),
time.perf_counter() - t0,
@ -575,7 +576,7 @@ class ChunkRepository:
)
score = scores[i] if i < len(scores) else 1.0
chunks_with_scores.append((chunk, score))
logger.info(
perf_logger.debug(
"search.build_chunks count=%d took %.3fs",
len(chunks_with_scores),
time.perf_counter() - t0,

11
haiku_rag_slim/haiku/rag/tools/search.py
View file

@ -9,6 +9,7 @@ from haiku.rag.store.models import SearchResult
from haiku.rag.tools.context import RAGDeps
logger = logging.getLogger(__name__)
perf_logger = logging.getLogger("haiku.rag.perf")
def create_search_toolset(
@ -59,7 +60,7 @@ def create_search_toolset(
if _last_tool_return:
llm_think_time = tool_start - _last_tool_return[0]
logger.info(
perf_logger.debug(
"tool.llm_thinking took %.3fs", llm_think_time
)
@ -81,7 +82,7 @@ def create_search_toolset(
results = await client.search(
query, limit=effective_limit, filter=effective_filter
)
logger.info(
perf_logger.debug(
"tool.search query=%r took %.3fs",
query[:80],
time.perf_counter() - t0,
@ -90,7 +91,7 @@ def create_search_toolset(
if expand_context:
t0 = time.perf_counter()
results = await client.expand_context(results)
logger.info(
perf_logger.debug(
"tool.expand_context results=%d took %.3fs",
len(results),
time.perf_counter() - t0,
@ -112,14 +113,14 @@ def create_search_toolset(
for i, r in enumerate(results_list)
]
output = "\n\n".join(formatted)
logger.info(
perf_logger.debug(
"tool.format results=%d chars=%d took %.3fs",
total,
len(output),
time.perf_counter() - t0,
)
logger.info(
perf_logger.debug(
"tool.search_total took %.3fs",
time.perf_counter() - tool_start,
)

78
tests/test_context_enhancement.py
View file

@ -768,3 +768,81 @@ async def test_expand_context_no_base64_images_docling_serve(temp_db_path):
assert "data:image" not in result.content.lower(), (
f"Found 'data:image' in expanded content: {result.content[:500]}"
)
async def test_expand_context_disabled_mode(temp_db_path):
"""Test that context_expansion_mode='disabled' returns results unchanged."""
config = AppConfig()
config.search.context_radius = 3
config.search.context_expansion_mode = "disabled"
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
# Build SearchResults directly — no DB or embeddings needed
search_results = [
SearchResult(
content="Test content",
score=0.9,
chunk_id="chunk-1",
document_id="doc-1",
),
SearchResult(
content="Other content",
score=0.8,
chunk_id="chunk-2",
document_id="doc-1",
),
]
expanded = await client.expand_context(search_results)
# Should return the exact same list — no expansion performed
assert len(expanded) == 2
assert expanded[0].content == "Test content"
assert expanded[0].score == 0.9
assert expanded[1].content == "Other content"
assert expanded[1].score == 0.8
@pytest.mark.vcr()
async def test_expand_context_chunks_mode_skips_docling(temp_db_path):
"""Test that context_expansion_mode='chunks' uses chunk expansion."""
config = AppConfig()
config.search.context_radius = 1
config.search.context_expansion_mode = "chunks"
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
# Create document with manual chunks that have pre-set embeddings
# to avoid calling the embedding API
dim = 2560
docling_doc = DoclingDocument(name="test_chunks_mode")
docling_doc.add_text(label=DocItemLabel.TEXT, text="Full content")
manual_chunks = [
Chunk(content="Part A", order=0, embedding=[0.1] * dim),
Chunk(content="Part B", order=1, embedding=[0.2] * dim),
Chunk(content="Part C", order=2, embedding=[0.3] * dim),
]
doc = await client.import_document(
docling_document=docling_doc, chunks=manual_chunks
)
assert doc.id is not None
chunks = await client.chunk_repository.get_by_document_id(doc.id)
chunk_b = next(c for c in chunks if c.order == 1)
# Give it doc_item_refs so it would normally trigger docling path
search_results = [
SearchResult(
content=chunk_b.content,
score=0.9,
chunk_id=chunk_b.id,
document_id=doc.id,
doc_item_refs=["#/texts/0"],
),
]
expanded = await client.expand_context(search_results)
# Should use chunk-based expansion (radius=1 around order 1)
assert len(expanded) == 1
# Expanded content should include adjacent chunks
assert "Part A" in expanded[0].content
assert "Part B" in expanded[0].content
assert "Part C" in expanded[0].content