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

Merge expanded context if possible

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-08-12 15:02:07 +02:00
parent e7c2a105ed
commit 99f7401817
No known key found for this signature in database
4 changed files with 220 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
docs/configuration.md
View file

@ -175,5 +175,7 @@ CHUNK_SIZE=256
# Number of adjacent chunks to include before/after retrieved chunks for context
# 0 = no expansion (default), 1 = include 1 chunk before and after, etc.
# When expanded chunks overlap or are adjacent, they are automatically merged
# into single chunks with continuous content to eliminate duplication
CONTEXT_CHUNK_RADIUS=0
```

2
docs/python.md
View file

@ -146,6 +146,8 @@ for chunk, score in expanded_results:
print(f"Expanded content: {chunk.content}") # Now includes before/after chunks
```
**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.
This is automatically used by the QA system when `CONTEXT_CHUNK_RADIUS > 0` to provide better answers with more complete context.
## Question Answering

133
src/haiku/rag/client.py
View file

@ -351,52 +351,129 @@ class HaikuRAG:
async def expand_context(
self, search_results: list[tuple[Chunk, float]]
) -> list[tuple[Chunk, float]]:
"""Expand search results with adjacent chunks based on CONTEXT_CHUNK_RADIUS.
"""Expand search results with adjacent chunks, merging overlapping chunks.
Args:
search_results: List of (chunk, score) tuples from search.
Returns:
List of (chunk, score) tuples with expanded context chunks.
List of (chunk, score) tuples with expanded and merged context chunks.
"""
if Config.CONTEXT_CHUNK_RADIUS == 0:
return search_results
# Group chunks by document_id to handle merging within documents
document_groups = {}
for chunk, score in search_results:
doc_id = chunk.document_id
if doc_id not in document_groups:
document_groups[doc_id] = []
document_groups[doc_id].append((chunk, score))
results = []
for chunk, score in search_results:
adjacent_chunks = await self.chunk_repository.get_adjacent_chunks(
chunk, Config.CONTEXT_CHUNK_RADIUS
)
for doc_id, doc_chunks in document_groups.items():
# Get all expanded ranges for this document
expanded_ranges = []
for chunk, score in doc_chunks:
adjacent_chunks = await self.chunk_repository.get_adjacent_chunks(
chunk, Config.CONTEXT_CHUNK_RADIUS
)
chunk_order = chunk.metadata.get("order", 0)
before_chunks = [
c for c in adjacent_chunks if c.metadata.get("order", 0) < chunk_order
]
after_chunks = [
c for c in adjacent_chunks if c.metadata.get("order", 0) > chunk_order
]
all_chunks = adjacent_chunks + [chunk]
combined_content_parts = (
[c.content for c in before_chunks]
+ [chunk.content]
+ [c.content for c in after_chunks]
)
# Get the range of orders for this expanded chunk
orders = [c.metadata.get("order", 0) for c in all_chunks]
min_order = min(orders)
max_order = max(orders)
# Create expanded chunk with combined content
expanded_chunk = Chunk(
id=chunk.id,
document_id=chunk.document_id,
content="".join(combined_content_parts),
metadata=chunk.metadata,
document_uri=chunk.document_uri,
document_meta=chunk.document_meta,
)
expanded_ranges.append(
{
"original_chunk": chunk,
"score": score,
"min_order": min_order,
"max_order": max_order,
"all_chunks": sorted(
all_chunks, key=lambda c: c.metadata.get("order", 0)
),
}
)
results.append((expanded_chunk, score))
# Merge overlapping/adjacent ranges
merged_ranges = self._merge_overlapping_ranges(expanded_ranges)
# Create merged chunks
for merged_range in merged_ranges:
combined_content_parts = [c.content for c in merged_range["all_chunks"]]
# Use the first original chunk for metadata
original_chunk = merged_range["original_chunks"][0]
merged_chunk = Chunk(
id=original_chunk.id,
document_id=original_chunk.document_id,
content="".join(combined_content_parts),
metadata=original_chunk.metadata,
document_uri=original_chunk.document_uri,
document_meta=original_chunk.document_meta,
)
# Use the highest score from merged chunks
best_score = max(merged_range["scores"])
results.append((merged_chunk, best_score))
return results
def _merge_overlapping_ranges(self, expanded_ranges):
"""Merge overlapping or adjacent expanded ranges."""
if not expanded_ranges:
return []
# Sort by min_order
sorted_ranges = sorted(expanded_ranges, key=lambda x: x["min_order"])
merged = []
current = {
"min_order": sorted_ranges[0]["min_order"],
"max_order": sorted_ranges[0]["max_order"],
"original_chunks": [sorted_ranges[0]["original_chunk"]],
"scores": [sorted_ranges[0]["score"]],
"all_chunks": sorted_ranges[0]["all_chunks"],
}
for range_info in sorted_ranges[1:]:
# Check if ranges overlap or are adjacent (max_order + 1 >= min_order)
if current["max_order"] >= range_info["min_order"] - 1:
# Merge ranges
current["max_order"] = max(
current["max_order"], range_info["max_order"]
)
current["original_chunks"].append(range_info["original_chunk"])
current["scores"].append(range_info["score"])
# Merge all_chunks and deduplicate by order
all_chunks_dict = {}
for chunk in current["all_chunks"] + range_info["all_chunks"]:
order = chunk.metadata.get("order", 0)
all_chunks_dict[order] = chunk
current["all_chunks"] = [
all_chunks_dict[order] for order in sorted(all_chunks_dict.keys())
]
else:
# No overlap, add current to merged and start new
merged.append(current)
current = {
"min_order": range_info["min_order"],
"max_order": range_info["max_order"],
"original_chunks": [range_info["original_chunk"]],
"scores": [range_info["score"]],
"all_chunks": range_info["all_chunks"],
}
# Add the last range
merged.append(current)
return merged
async def ask(self, question: str, cite: bool = False) -> str:
"""Ask a question using the configured QA agent.

111
tests/test_client.py
View file

@ -639,3 +639,114 @@ async def test_client_expand_context_multiple_chunks():
assert score2 == 0.7
assert "Doc2 Section X" in expanded2.content
assert "Doc2 Section Y" in expanded2.content
@pytest.mark.asyncio
async def test_client_expand_context_merges_overlapping_chunks():
"""Test that overlapping expanded chunks are merged into one."""
with patch("haiku.rag.client.Config.CONTEXT_CHUNK_RADIUS", 1):
async with HaikuRAG(":memory:") as client:
# Create document with 5 chunks
manual_chunks = [
Chunk(content="Chunk 0", metadata={"order": 0}),
Chunk(content="Chunk 1", metadata={"order": 1}),
Chunk(content="Chunk 2", metadata={"order": 2}),
Chunk(content="Chunk 3", metadata={"order": 3}),
Chunk(content="Chunk 4", metadata={"order": 4}),
]
doc = await client.create_document(
content="Full document content", chunks=manual_chunks
)
assert doc.id is not None
chunks = await client.chunk_repository.get_by_document_id(doc.id)
# Get adjacent chunks (orders 1 and 2) - these will overlap when expanded
chunk1 = next(c for c in chunks if c.metadata.get("order") == 1)
chunk2 = next(c for c in chunks if c.metadata.get("order") == 2)
# With radius=1:
# chunk1 expanded would be [0,1,2]
# chunk2 expanded would be [1,2,3]
# These should merge into one chunk containing [0,1,2,3]
search_results = [(chunk1, 0.8), (chunk2, 0.7)]
expanded_results = await client.expand_context(search_results)
# Should have only 1 merged result instead of 2 overlapping ones
assert len(expanded_results) == 1
merged_chunk, score = expanded_results[0]
# Should contain all chunks from 0 to 3
assert "Chunk 0" in merged_chunk.content
assert "Chunk 1" in merged_chunk.content
assert "Chunk 2" in merged_chunk.content
assert "Chunk 3" in merged_chunk.content
assert "Chunk 4" not in merged_chunk.content # Should not include chunk 4
# Should use the higher score (0.8)
assert score == 0.8
@pytest.mark.asyncio
async def test_client_expand_context_keeps_separate_non_overlapping():
"""Test that non-overlapping expanded chunks remain separate."""
with patch("haiku.rag.client.Config.CONTEXT_CHUNK_RADIUS", 1):
async with HaikuRAG(":memory:") as client:
# Create document with chunks far apart
manual_chunks = [
Chunk(content="Chunk 0", metadata={"order": 0}),
Chunk(content="Chunk 1", metadata={"order": 1}),
Chunk(content="Chunk 2", metadata={"order": 2}),
Chunk(content="Chunk 5", metadata={"order": 5}), # Gap here
Chunk(content="Chunk 6", metadata={"order": 6}),
Chunk(content="Chunk 7", metadata={"order": 7}),
]
doc = await client.create_document(
content="Full document content", chunks=manual_chunks
)
assert doc.id is not None
chunks = await client.chunk_repository.get_by_document_id(doc.id)
# Get chunks by index - they will have sequential orders 0,1,2,3,4,5
# So get chunk with order=0 and chunk with order=5 (far enough apart)
chunk0 = next(
c for c in chunks if c.metadata.get("order") == 0
) # Content: "Chunk 0"
chunk5 = next(
c for c in chunks if c.metadata.get("order") == 5
) # Content: "Chunk 7"
# chunk0 expanded: [0,1] with radius=1 (orders 0,1)
# chunk5 expanded: [4,5] with radius=1 (orders 4,5)
# These should remain separate (max_order 1 < min_order 4 - 1)
search_results = [(chunk0, 0.8), (chunk5, 0.7)]
expanded_results = await client.expand_context(search_results)
# Should have 2 separate results
assert len(expanded_results) == 2
# Sort by score to ensure predictable order
expanded_results.sort(key=lambda x: x[1], reverse=True)
chunk0_expanded, score1 = expanded_results[0]
chunk5_expanded, score2 = expanded_results[1]
# First chunk (order=0) expanded should contain orders [0,1]
# Content should be "Chunk 0" + "Chunk 1"
assert "Chunk 0" in chunk0_expanded.content
assert "Chunk 1" in chunk0_expanded.content
assert (
"Chunk 7" not in chunk0_expanded.content
) # Should not have chunk 7 content
assert score1 == 0.8
# Second chunk (order=5) expanded should contain orders [4,5]
# Content should be "Chunk 6" + "Chunk 7" (orders 4 and 5)
assert "Chunk 6" in chunk5_expanded.content # Order 4 content
assert "Chunk 7" in chunk5_expanded.content # Order 5 content
assert "Chunk 0" not in chunk5_expanded.content
assert score2 == 0.7