danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
haiku.rag/tests/test_context_enhancement.py
Yiorgis Gozadinos 3603b00513
Fix base64 images leaking into expanded search results 
PictureItem.export_to_markdown() defaults to ImageRefMode.EMBEDDED,
which embeds base64 image data. When _extract_item_text() called this
method during context expansion, base64 data would leak into results.

Now explicitly handles PictureItem with ImageRefMode.PLACEHOLDER to
prevent base64 while still including VLM descriptions and captions.
2026-01-08 18:42:58 +02:00

820 lines
30 KiB
Python
Raw Blame History

import pytest
from docling_core.types.doc.document import DoclingDocument, TableData
from docling_core.types.doc.labels import DocItemLabel
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models import SearchResult
from haiku.rag.store.models.chunk import Chunk
async def create_document_with_docling(
client: HaikuRAG, docling_doc: DoclingDocument, title: str
):
"""Helper to create a document from a DoclingDocument using import_document."""
chunks = await client.chunk(docling_doc)
embedded_chunks = await client._ensure_chunks_embedded(chunks)
return await client.import_document(
docling_document=docling_doc,
chunks=embedded_chunks,
title=title,
)
def create_table_document() -> DoclingDocument:
"""Create a document with a table that will be split across chunks."""
doc = DoclingDocument(name="table_test")
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Introduction paragraph.")
doc.add_heading(text="Employee Data", level=1)
# Create a table with enough content to span multiple chunks
table_data = TableData(num_cols=3, num_rows=0)
table_data.add_row(["Name", "Age", "City"])
table_data.add_row(["Alice Smith", "30", "New York"])
table_data.add_row(["Bob Johnson", "25", "Los Angeles"])
table_data.add_row(["Charlie Brown", "35", "Chicago"])
table_data.add_row(["Diana Ross", "28", "Miami"])
doc.add_table(data=table_data)
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Conclusion paragraph.")
return doc
def create_list_document() -> DoclingDocument:
"""Create a document with list items that will be split across chunks."""
doc = DoclingDocument(name="list_test")
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Shopping list for the week:")
doc.add_list_item(
text="Fresh organic apples from the farmers market", enumerated=False
)
doc.add_list_item(text="Ripe yellow bananas for smoothies", enumerated=False)
doc.add_list_item(text="Valencia oranges for fresh juice", enumerated=False)
doc.add_list_item(text="Seedless red grapes as healthy snack", enumerated=False)
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Remember to bring reusable bags.")
return doc
def create_code_document() -> DoclingDocument:
"""Create a document with adjacent code blocks that will be split."""
doc = DoclingDocument(name="code_test")
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Here are several code snippets:")
# Multiple adjacent code blocks - type-aware expansion should group them
doc.add_text(label=DocItemLabel.CODE, text="# Part 1: Setup\nimport os\nimport sys")
doc.add_text(
label=DocItemLabel.CODE, text='# Part 2: Config\nCONFIG = {"debug": True}'
)
doc.add_text(
label=DocItemLabel.CODE, text="# Part 3: Main\ndef main():\n print(CONFIG)"
)
doc.add_text(label=DocItemLabel.PARAGRAPH, text="End of code examples.")
return doc
@pytest.fixture
def small_chunk_config() -> AppConfig:
"""Config with small chunk size to force splitting."""
config = AppConfig()
config.processing.chunk_size = 32
config.search.max_context_items = 25
config.search.max_context_chars = 10000
return config
@pytest.mark.vcr()
async def test_table_expansion_includes_split_rows(temp_db_path, small_chunk_config):
"""Verify that table expansion retrieves rows that were split into different chunks."""
docling_doc = create_table_document()
async with HaikuRAG(temp_db_path, config=small_chunk_config, create=True) as client:
doc = await create_document_with_docling(client, docling_doc, "Table Test")
assert doc.id is not None
# Search for table content
results = await client.search("Alice Smith New York employee", limit=5)
table_results = [r for r in results if "table" in r.labels]
assert len(table_results) > 0, (
f"No table results. Labels: {[r.labels for r in results]}"
)
original = table_results[0]
# Verify the original chunk does NOT contain all table data
# (proving we need expansion)
original_has_all = all(
name in original.content for name in ["Alice", "Bob", "Charlie", "Diana"]
)
# Expand context
expanded = await client.expand_context(table_results[:1])
assert len(expanded) == 1
expanded_content = expanded[0].content
# After expansion, we should have the complete table
assert "Alice" in expanded_content
assert "Bob" in expanded_content
assert "Charlie" in expanded_content
assert "Diana" in expanded_content
# Verify expansion actually added content (unless chunk already had everything)
if not original_has_all:
assert len(expanded_content) > len(original.content), (
"Expansion should have added content"
)
@pytest.mark.vcr()
async def test_list_expansion_includes_split_items(temp_db_path, small_chunk_config):
"""Verify that list expansion retrieves items that were split into different chunks."""
docling_doc = create_list_document()
async with HaikuRAG(temp_db_path, config=small_chunk_config, create=True) as client:
doc = await create_document_with_docling(client, docling_doc, "List Test")
assert doc.id is not None
# Search for a list item
results = await client.search("grapes healthy snack", limit=5)
list_results = [r for r in results if "list_item" in r.labels]
assert len(list_results) > 0, (
f"No list results. Labels: {[r.labels for r in results]}"
)
original = list_results[0]
# Check what the original chunk contains
original_items = sum(
1
for item in ["apples", "bananas", "oranges", "grapes"]
if item in original.content.lower()
)
# Expand context
expanded = await client.expand_context(list_results[:1])
assert len(expanded) == 1
expanded_content = expanded[0].content.lower()
# Count items after expansion
expanded_items = sum(
1
for item in ["apples", "bananas", "oranges", "grapes"]
if item in expanded_content
)
# Expansion should include at least as many items (more if split)
assert expanded_items >= original_items
# If original didn't have all items, expansion should have added some
if original_items < 4:
assert expanded_items > original_items, (
f"Expansion should have added items. Original: {original_items}, Expanded: {expanded_items}"
)
@pytest.mark.vcr()
async def test_code_expansion_includes_adjacent_blocks(
temp_db_path, small_chunk_config
):
"""Verify that code expansion retrieves adjacent code blocks split across chunks."""
docling_doc = create_code_document()
async with HaikuRAG(temp_db_path, config=small_chunk_config, create=True) as client:
doc = await create_document_with_docling(client, docling_doc, "Code Test")
assert doc.id is not None
# Search for middle code block (Part 2)
results = await client.search("CONFIG debug True", limit=5)
code_results = [r for r in results if "code" in r.labels]
assert len(code_results) > 0, (
f"No code results. Labels: {[r.labels for r in results]}"
)
original = code_results[0]
# Check what parts the original chunk has
original_parts = sum(
1 for part in ["Part 1", "Part 2", "Part 3"] if part in original.content
)
# Expand context
expanded = await client.expand_context(code_results[:1])
assert len(expanded) == 1
expanded_content = expanded[0].content
# Count parts after expansion
expanded_parts = sum(
1 for part in ["Part 1", "Part 2", "Part 3"] if part in expanded_content
)
# Expansion should include at least as many parts
assert expanded_parts >= original_parts
# If original didn't have all parts, expansion should have added some
if original_parts < 3:
assert expanded_parts > original_parts, (
f"Expansion should have added code blocks. Original: {original_parts}, Expanded: {expanded_parts}"
)
@pytest.mark.vcr()
async def test_text_expansion_uses_radius(temp_db_path):
"""Text content expansion should use radius, not structural boundaries."""
config = AppConfig()
config.processing.chunk_size = 32
config.search.context_radius = 1 # Small radius
# Create a document with longer paragraphs that will split
doc = DoclingDocument(name="text_test")
doc.add_text(
label=DocItemLabel.PARAGRAPH,
text="First paragraph with enough content to be its own chunk in the document.",
)
doc.add_text(
label=DocItemLabel.PARAGRAPH,
text="Second paragraph contains different information about software testing.",
)
doc.add_text(
label=DocItemLabel.PARAGRAPH,
text="Third paragraph discusses various topics and provides more details.",
)
doc.add_text(
label=DocItemLabel.PARAGRAPH,
text="Fourth paragraph concludes the document with final thoughts.",
)
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
document = await create_document_with_docling(client, doc, "Text Only")
assert document.id is not None
# Search for second paragraph
results = await client.search("software testing", limit=1)
assert len(results) > 0
original = results[0]
expanded = await client.expand_context(results)
# With radius=1, expansion should include adjacent paragraphs
# Content length should be >= original (may add adjacent content)
assert len(expanded[0].content) >= len(original.content)
@pytest.mark.vcr()
async def test_expansion_preserves_metadata(temp_db_path, small_chunk_config):
"""Expansion should preserve document metadata."""
docling_doc = create_table_document()
async with HaikuRAG(temp_db_path, config=small_chunk_config, create=True) as client:
doc = await create_document_with_docling(client, docling_doc, "Metadata Test")
assert doc.id is not None
results = await client.search("Introduction paragraph", limit=1)
assert len(results) > 0
expanded = await client.expand_context(results)
assert expanded[0].document_title == "Metadata Test"
assert expanded[0].chunk_id == results[0].chunk_id
assert expanded[0].document_id == results[0].document_id
@pytest.mark.vcr()
async def test_format_for_agent_output(temp_db_path, small_chunk_config):
"""format_for_agent should include source, type, and content sections."""
docling_doc = create_table_document()
async with HaikuRAG(temp_db_path, config=small_chunk_config, create=True) as client:
doc = await create_document_with_docling(client, docling_doc, "Format Test")
assert doc.id is not None
results = await client.search("Alice Smith employee data", limit=5)
table_results = [r for r in results if "table" in r.labels]
assert len(table_results) > 0
expanded = await client.expand_context(table_results[:1])
formatted = expanded[0].format_for_agent()
# Check format structure
assert "score:" in formatted
assert 'Source: "Format Test"' in formatted
assert "Type: table" in formatted
assert "Content:" in formatted
@pytest.mark.vcr()
async def test_max_items_limit_caps_expansion(temp_db_path):
"""Expansion should respect max_context_items limit."""
config = AppConfig()
config.processing.chunk_size = 32
config.search.max_context_items = 2 # Very restrictive
docling_doc = create_list_document()
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
doc = await create_document_with_docling(client, docling_doc, "Limit Test")
assert doc.id is not None
results = await client.search("grapes", limit=1)
assert len(results) > 0
expanded = await client.expand_context(results)
# With max_items=2, expansion should be limited
content = expanded[0].content.lower()
item_count = sum(
1 for item in ["apples", "bananas", "oranges", "grapes"] if item in content
)
# Should have at most 2 items (the limit)
assert item_count <= 2, f"Expected at most 2 items, got {item_count}"
@pytest.mark.vcr()
async def test_search_result_get_primary_label():
"""Test _get_primary_label prioritizes structural labels correctly."""
# Table should be prioritized
result = SearchResult(
content="test",
score=0.5,
chunk_id="c1",
document_id="d1",
labels=["paragraph", "table", "text"],
)
assert result._get_primary_label() == "table"
# Code should be prioritized over paragraph
result = SearchResult(
content="test",
score=0.5,
chunk_id="c2",
document_id="d2",
labels=["paragraph", "code"],
)
assert result._get_primary_label() == "code"
# list_item should be prioritized
result = SearchResult(
content="test",
score=0.5,
chunk_id="c3",
document_id="d3",
labels=["text", "list_item"],
)
assert result._get_primary_label() == "list_item"
# Returns first label when no priority match
result = SearchResult(
content="test",
score=0.5,
chunk_id="c4",
document_id="d4",
labels=["paragraph", "text"],
)
assert result._get_primary_label() == "paragraph"
# Returns None for empty labels
result = SearchResult(
content="test",
score=0.5,
chunk_id="c5",
document_id="d5",
labels=[],
)
assert result._get_primary_label() is None
@pytest.mark.vcr()
async def test_expand_context_radius_zero(temp_db_path):
"""Test expand_context with radius 0 returns original results."""
# Default config has context_radius=0
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(content="Simple test content")
assert doc.id is not None
chunks = await client.chunk_repository.get_by_document_id(doc.id)
search_results = [SearchResult.from_chunk(chunks[0], 0.9)]
expanded_results = await client.expand_context(search_results)
# Should return exactly the same results
assert len(expanded_results) == 1
assert expanded_results[0].content == search_results[0].content
assert expanded_results[0].score == search_results[0].score
@pytest.mark.vcr()
async def test_expand_context_multiple_documents(temp_db_path):
"""Test expand_context with results from multiple documents."""
config = AppConfig()
config.search.context_radius = 1
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
# Create first document with manual chunks
docling_doc1 = DoclingDocument(name="doc1")
docling_doc1.add_text(label=DocItemLabel.TEXT, text="Doc1 content")
doc1_chunks = [
Chunk(content="Doc1 Part A", order=0),
Chunk(content="Doc1 Part B", order=1),
Chunk(content="Doc1 Part C", order=2),
]
doc1 = await client.import_document(
docling_document=docling_doc1, chunks=doc1_chunks, uri="doc1.txt"
)
# Create second document with manual chunks
docling_doc2 = DoclingDocument(name="doc2")
docling_doc2.add_text(label=DocItemLabel.TEXT, text="Doc2 content")
doc2_chunks = [
Chunk(content="Doc2 Section X", order=0),
Chunk(content="Doc2 Section Y", order=1),
]
doc2 = await client.import_document(
docling_document=docling_doc2, chunks=doc2_chunks, uri="doc2.txt"
)
assert doc1.id is not None
assert doc2.id is not None
chunks1 = await client.chunk_repository.get_by_document_id(doc1.id)
chunks2 = await client.chunk_repository.get_by_document_id(doc2.id)
# Get middle chunk from doc1 (order=1) and first chunk from doc2 (order=0)
chunk1 = next(c for c in chunks1 if c.order == 1)
chunk2 = next(c for c in chunks2 if c.order == 0)
search_results = [
SearchResult.from_chunk(chunk1, 0.8),
SearchResult.from_chunk(chunk2, 0.7),
]
expanded_results = await client.expand_context(search_results)
assert len(expanded_results) == 2
# Check first expanded result (should include chunks 0,1,2 from doc1)
expanded1 = expanded_results[0]
assert expanded1.score == 0.8
assert "Doc1 Part A" in expanded1.content
assert "Doc1 Part B" in expanded1.content
assert "Doc1 Part C" in expanded1.content
# Check second expanded result (should include chunks 0,1 from doc2)
expanded2 = expanded_results[1]
assert expanded2.score == 0.7
assert "Doc2 Section X" in expanded2.content
assert "Doc2 Section Y" in expanded2.content
@pytest.mark.vcr()
async def test_expand_context_merges_overlapping_chunks(temp_db_path):
"""Test that overlapping expanded chunks are merged into one."""
config = AppConfig()
config.search.context_radius = 1
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
# Create document with 5 chunks
docling_doc = DoclingDocument(name="test")
docling_doc.add_text(label=DocItemLabel.TEXT, text="Full document content")
manual_chunks = [
Chunk(content="Chunk 0", order=0),
Chunk(content="Chunk 1", order=1),
Chunk(content="Chunk 2", order=2),
Chunk(content="Chunk 3", order=3),
Chunk(content="Chunk 4", order=4),
]
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)
# Get adjacent chunks (orders 1 and 2) - these will overlap when expanded
chunk1 = next(c for c in chunks if c.order == 1)
chunk2 = next(c for c in chunks if c.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 = [
SearchResult.from_chunk(chunk1, 0.8),
SearchResult.from_chunk(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 = expanded_results[0]
# Should contain all chunks from 0 to 3
assert "Chunk 0" in merged.content
assert "Chunk 1" in merged.content
assert "Chunk 2" in merged.content
assert "Chunk 3" in merged.content
assert "Chunk 4" not in merged.content # Should not include chunk 4
# Should use the higher score (0.8)
assert merged.score == 0.8
@pytest.mark.vcr()
async def test_expand_context_keeps_separate_non_overlapping(temp_db_path):
"""Test that non-overlapping expanded chunks remain separate."""
config = AppConfig()
config.search.context_radius = 1
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
# Create document with chunks far apart
docling_doc = DoclingDocument(name="test")
docling_doc.add_text(label=DocItemLabel.TEXT, text="Full document content")
manual_chunks = [
Chunk(content="Chunk 0", order=0),
Chunk(content="Chunk 1", order=1),
Chunk(content="Chunk 2", order=2),
Chunk(content="Chunk 5", order=5), # Gap here
Chunk(content="Chunk 6", order=6),
Chunk(content="Chunk 7", order=7),
]
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)
# 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.order == 0) # Content: "Chunk 0"
chunk5 = next(
c for c in chunks if c.order == 5
) # Content: "Chunk 7" but now at order 5
# chunk0 expanded: [0,1] with radius=1 (orders 0,1)
# chunk5 expanded: [4,5] with radius=1 (orders 4,5)
search_results = [
SearchResult.from_chunk(chunk0, 0.8),
SearchResult.from_chunk(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.score, reverse=True)
chunk0_expanded = expanded_results[0]
chunk5_expanded = 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 5" not in chunk0_expanded.content
assert chunk0_expanded.score == 0.8
# Second chunk (order=5) expanded should contain orders [4,5]
# Content should be "Chunk 6" (order 4) + "Chunk 7" (order 5)
assert "Chunk 6" in chunk5_expanded.content
assert "Chunk 7" in chunk5_expanded.content
assert "Chunk 0" not in chunk5_expanded.content
assert chunk5_expanded.score == 0.7
@pytest.mark.vcr()
async def test_expand_context_with_docling_merges_overlapping(temp_db_path):
"""Test that expand_context with DoclingDocument merges overlapping results."""
config = AppConfig()
config.search.context_radius = 3
markdown_content = """# Chapter 1
This is paragraph one about topic A.
This is paragraph two about topic A continued.
This is paragraph three about topic B.
# Chapter 2
This is paragraph four about topic C.
"""
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
doc = await client.create_document(
content=markdown_content,
uri="test://structured",
)
assert doc.id is not None
assert doc.docling_document_json is not None
# Get chunks which should have doc_item_refs
chunks = await client.chunk_repository.get_by_document_id(doc.id)
assert len(chunks) >= 1
# Find chunks that have doc_item_refs (from docling chunking)
chunks_with_refs = [c for c in chunks if c.get_chunk_metadata().doc_item_refs]
if len(chunks_with_refs) >= 2:
# Create search results from adjacent chunks
search_results = [
SearchResult.from_chunk(chunks_with_refs[0], 0.9),
SearchResult.from_chunk(chunks_with_refs[1], 0.8),
]
# Expand with configured radius that should cause overlap
expanded = await client.expand_context(search_results)
# If chunks were adjacent, they should be merged
# The expanded results should have merged metadata
assert len(expanded) >= 1
# Check that expanded result has page_numbers populated
for r in expanded:
# Should have doc_item_refs from expansion
assert r.doc_item_refs is not None
@pytest.mark.vcr()
async def test_expand_context_docling_merges_metadata(temp_db_path):
"""Test that expand_context properly merges metadata from multiple results."""
config = AppConfig()
config.search.context_radius = 10
markdown_content = """# Introduction
First paragraph of introduction.
Second paragraph of introduction.
# Methods
First paragraph of methods section.
Second paragraph of methods section.
# Results
First paragraph of results.
"""
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
doc = await client.create_document(
content=markdown_content,
uri="test://metadata-merge",
)
assert doc.id is not None
chunks = await client.chunk_repository.get_by_document_id(doc.id)
chunks_with_refs = [c for c in chunks if c.get_chunk_metadata().doc_item_refs]
if len(chunks_with_refs) >= 2:
# Get chunks with different headings if possible
chunk1 = chunks_with_refs[0]
chunk2 = chunks_with_refs[-1] # Last chunk likely has different heading
search_results = [
SearchResult.from_chunk(chunk1, 0.9),
SearchResult.from_chunk(chunk2, 0.8),
]
# Expand with large radius to potentially merge
expanded = await client.expand_context(search_results)
# Check that results have proper structure
for r in expanded:
# Content should be non-empty
assert len(r.content) > 0
# Score should be preserved (best score)
assert r.score in [0.9, 0.8]
def create_picture_document() -> DoclingDocument:
"""Create a document with picture items for testing image handling."""
from docling_core.types.doc.labels import DocItemLabel
doc = DoclingDocument(name="picture_test")
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Introduction with context.")
doc.add_heading(text="Figures Section", level=1)
# Add a picture item - this will have export_to_markdown method
# that defaults to EMBEDDED mode if called without image_mode parameter
doc.add_picture()
doc.add_text(label=DocItemLabel.CAPTION, text="Figure 1: Sample diagram")
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Conclusion paragraph.")
return doc
@pytest.mark.vcr()
async def test_expand_context_no_base64_images(temp_db_path):
"""Ensure expanded context does not contain base64 image data.
This test verifies that when expand_context includes PictureItem objects,
they are serialized with PLACEHOLDER mode (not EMBEDDED), preventing
base64 image data from leaking into the expanded content.
"""
config = AppConfig()
config.search.context_radius = 5 # Expand enough to include pictures
docling_doc = create_picture_document()
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
doc = await create_document_with_docling(client, docling_doc, "Picture Test")
assert doc.id is not None
# Search for content near the picture
results = await client.search("Figure diagram", limit=5)
assert len(results) > 0
# Expand context - this should include the picture area
expanded = await client.expand_context(results)
for result in expanded:
# Base64 image data should never appear in expanded content
assert "base64" not in result.content.lower(), (
f"Found 'base64' in expanded content: {result.content[:500]}"
)
assert "data:image" not in result.content.lower(), (
f"Found 'data:image' in expanded content: {result.content[:500]}"
)
@pytest.mark.integration
@pytest.mark.vcr()
async def test_expand_context_no_base64_images_docling_local(temp_db_path):
"""Ensure expanded context from real PDF does not contain base64 image data.
Tests end-to-end with doclaynet.pdf using docling-local converter.
"""
from pathlib import Path
config = AppConfig()
config.processing.converter = "docling-local"
config.processing.chunker = "docling-local"
config.search.context_radius = 5
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
pdf_path = Path(__file__).parent / "data" / "doclaynet.pdf"
result = await client.create_document_from_source(pdf_path)
doc = result if not isinstance(result, list) else result[0]
assert doc.id is not None
# Search for content that might include pictures
results = await client.search("figure table", limit=10)
if len(results) > 0:
expanded = await client.expand_context(results)
for result in expanded:
assert "base64" not in result.content.lower(), (
f"Found 'base64' in expanded content: {result.content[:500]}"
)
assert "data:image" not in result.content.lower(), (
f"Found 'data:image' in expanded content: {result.content[:500]}"
)
@pytest.mark.integration
@pytest.mark.vcr()
async def test_expand_context_no_base64_images_docling_serve(temp_db_path):
"""Ensure expanded context from real PDF does not contain base64 image data.
Tests end-to-end with doclaynet.pdf using docling-serve converter.
"""
from pathlib import Path
config = AppConfig()
config.processing.converter = "docling-serve"
config.processing.chunker = "docling-serve"
config.search.context_radius = 5
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
pdf_path = Path(__file__).parent / "data" / "doclaynet.pdf"
result = await client.create_document_from_source(pdf_path)
doc = result if not isinstance(result, list) else result[0]
assert doc.id is not None
# Search for content that might include pictures
results = await client.search("figure table", limit=10)
if len(results) > 0:
expanded = await client.expand_context(results)
for result in expanded:
assert "base64" not in result.content.lower(), (
f"Found 'base64' in expanded content: {result.content[:500]}"
)
assert "data:image" not in result.content.lower(), (
f"Found 'data:image' in expanded content: {result.content[:500]}"
)
Reference in a new issue View git blame Copy permalink