diff --git a/app/haiku.rag.yaml.example b/app/haiku.rag.yaml.example index a60ea0d0..25d69e58 100644 --- a/app/haiku.rag.yaml.example +++ b/app/haiku.rag.yaml.example @@ -34,6 +34,7 @@ embeddings: search: limit: 5 context_radius: 0 + # context_expansion_mode: auto # auto, chunks, or disabled # Provider settings providers: diff --git a/docs/apps.md b/docs/apps.md index 2f0ac690..4d8eace4 100644 --- a/docs/apps.md +++ b/docs/apps.md @@ -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 diff --git a/docs/configuration/index.md b/docs/configuration/index.md index bdece61d..6417c97c 100644 --- a/docs/configuration/index.md +++ b/docs/configuration/index.md @@ -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 diff --git a/docs/configuration/qa-research.md b/docs/configuration/qa-research.md index 5e522082..78e2848c 100644 --- a/docs/configuration/qa-research.md +++ b/docs/configuration/qa-research.md @@ -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`. diff --git a/docs/python.md b/docs/python.md index f507c636..ccb60f53 100644 --- a/docs/python.md +++ b/docs/python.md @@ -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. diff --git a/docs/tuning.md b/docs/tuning.md index efcd0c7f..a5ef7c3a 100644 --- a/docs/tuning.md +++ b/docs/tuning.md @@ -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. diff --git a/haiku_rag_slim/haiku/rag/agents/qa/agent.py b/haiku_rag_slim/haiku/rag/agents/qa/agent.py index dc15da05..02737cd6 100644 --- a/haiku_rag_slim/haiku/rag/agents/qa/agent.py +++ b/haiku_rag_slim/haiku/rag/agents/qa/agent.py @@ -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 diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 1ea24c84..34bae47b 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -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, diff --git a/haiku_rag_slim/haiku/rag/store/models/document.py b/haiku_rag_slim/haiku/rag/store/models/document.py index 0e702cb1..861571f4 100644 --- a/haiku_rag_slim/haiku/rag/store/models/document.py +++ b/haiku_rag_slim/haiku/rag/store/models/document.py @@ -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, diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index 739de37c..a626a48a 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -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, diff --git a/haiku_rag_slim/haiku/rag/tools/search.py b/haiku_rag_slim/haiku/rag/tools/search.py index e76ddeea..50e25a3a 100644 --- a/haiku_rag_slim/haiku/rag/tools/search.py +++ b/haiku_rag_slim/haiku/rag/tools/search.py @@ -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, ) diff --git a/tests/test_context_enhancement.py b/tests/test_context_enhancement.py index f1a03792..10687aad 100644 --- a/tests/test_context_enhancement.py +++ b/tests/test_context_enhancement.py @@ -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