feat: allow chunks-only in recall (max_tokens=0) (#364)

* feat: allow chunks only in recall

* feat: fetch chunks independently of max_tokens filtering

Changes:
- Chunks now fetched BEFORE max_tokens filtering (Step 5.5)
- Implements batching: (max_chunk_tokens / retain_chunk_size) * 2
- Loop-based fetching until budget exhausted or no more chunks
- Handles varying chunk sizes across documents
- When max_tokens=0: returns 0 facts but still returns chunks
- When max_tokens>0: backward compatible (chunks match filtered facts)

Tests:
- Added test_recall_chunks_independence.py with 5 comprehensive tests
- Tests chunk independence, batching, ordering, and backward compat

Docs:
- Updated recall.mdx to explain new chunk behavior
- Updated memory_engine.py docstrings

Fixes chunk-related test failures by reordering chunks to match
filtered facts when max_tokens > 0 (backward compatibility).

* fix: fetch chunks after token filtering when max_tokens>0

Changes:
- When max_tokens=0: fetch chunks BEFORE token filtering (new behavior)
- When max_tokens>0: fetch chunks AFTER token filtering (backward compat)
- This ensures chunk ordering matches filtered facts for max_tokens>0
- Fixes test failures in test_chunks_and_entities_follow_fact_order,
  test_chunk_fact_mapping, test_chunk_ordering_preservation, etc.

The previous approach tried to reorder prefetched chunks, but that
caused issues when the chunk budget was exhausted before all facts
were processed. The new approach fetches chunks based on the correct
fact set for each scenario.

* fix: use ConfigResolver for bank-specific retain_chunk_size

Fixes error: Field 'retain_chunk_size' is bank-configurable and cannot
be accessed from global config.

Changed from:
- config.retain_chunk_size (global config, not allowed)

To:
- bank_config.retain_chunk_size (resolved from ConfigResolver)

This ensures the correct chunk size is used for each bank, respecting
any bank-specific overrides.

* fix: correct Budget import in test_recall_chunks_independence

Changed from:
- from hindsight_api.engine.interface import Budget (incorrect)

To:
- from hindsight_api.engine.memory_engine import Budget (correct)

This fixes the ImportError that was preventing the tests from running.

* fix: prevent infinite loop in chunk fetching and improve test content

- Add max(1, ...) to estimated_batch_size to prevent division resulting in 0
- Update test content to use more substantial examples that generate facts
- Add request_context parameter to all retain_async and recall_async test calls

* refactor: simplify chunk fetching to always use pre-filtering approach

Remove backward compatibility code that fetched chunks after token
filtering. Now chunks are always fetched from top-scored results
before max_tokens filtering, regardless of max_tokens value.

This simplifies the code by:
- Removing duplicate chunk fetching logic
- Eliminating conditional behavior based on max_tokens
- Making chunk fetching behavior consistent and predictable

Chunks are still fetched in batches and respect max_chunk_tokens limit.
This commit is contained in:
Nicolò Boschi 2026-02-13 16:56:35 +01:00 committed by GitHub
parent ff55283018
commit 7dad9da02d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 375 additions and 68 deletions

View file

@ -1682,15 +1682,21 @@ class MemoryEngine(MemoryEngineInterface):
max_entity_tokens: Maximum tokens for entity observations (default 500)
include_chunks: Whether to include raw chunks in the response
max_chunk_tokens: Maximum tokens for chunks (default 8192)
NOTE: Chunks are fetched independently of max_tokens filtering.
This means setting max_tokens=0 will return 0 facts but can still
return chunks from the top-scored (reranked) results.
Chunks are fetched in batches (estimated as (max_chunk_tokens // retain_chunk_size) * 2)
until the token budget is exhausted or all chunks are fetched.
This handles varying chunk sizes across documents.
tags: Optional list of tags for visibility filtering (OR matching - returns
memories that have at least one matching tag)
Returns:
RecallResultModel containing:
- results: List of MemoryFact objects
- results: List of MemoryFact objects (filtered by max_tokens)
- trace: Optional trace information for debugging
- entities: Optional dict of entity states (if include_entities=True)
- chunks: Optional dict of chunks (if include_chunks=True)
- chunks: Optional dict of chunks (if include_chunks=True, independent of max_tokens)
"""
# Authenticate tenant and set schema in context (for fq_table())
await self._authenticate_tenant(request_context)
@ -1918,7 +1924,8 @@ class MemoryEngine(MemoryEngineInterface):
2. Merge: RRF to combine ranked lists
3. Reranking: Pluggable strategy (heuristic or cross-encoder)
4. Diversity: MMR with λ=0.5
5. Token Filter: Limit results to max_tokens budget
5. Chunks: Fetch chunks from top-scored results (BEFORE token filtering)
6. Token Filter: Limit facts to max_tokens budget
Args:
bank_id: bank IDentifier
@ -1929,7 +1936,7 @@ class MemoryEngine(MemoryEngineInterface):
enable_trace: Whether to return search trace (deprecated)
include_entities: Whether to include entity observations
max_entity_tokens: Maximum tokens for entity observations
include_chunks: Whether to include raw chunks
include_chunks: Whether to include raw chunks (fetched before max_tokens filtering)
max_chunk_tokens: Maximum tokens for chunks
Returns:
@ -2352,6 +2359,85 @@ class MemoryEngine(MemoryEngineInterface):
top_scored = scored_results[:rerank_limit]
log_buffer.append(f" [5] Truncated to top {len(top_scored)} results")
# Step 5.5: Fetch chunks from top-scored results (before token filtering)
# Chunks are fetched independently of max_tokens filtering
chunks_dict = None
total_chunk_tokens = 0
if include_chunks and top_scored:
from .response_models import ChunkInfo
# Collect chunk_ids in order of fact relevance (preserving order from top_scored)
# Use a list to maintain order, but track seen chunks to avoid duplicates
chunk_ids_ordered = []
seen_chunk_ids = set()
for sr in top_scored:
chunk_id = sr.retrieval.chunk_id
if chunk_id and chunk_id not in seen_chunk_ids:
chunk_ids_ordered.append(chunk_id)
seen_chunk_ids.add(chunk_id)
if chunk_ids_ordered:
# Estimate batch size based on retain_chunk_size * 2 (rough estimate)
# Chunk sizes vary per document, so we fetch in batches until budget is exhausted
bank_config = await self._config_resolver.resolve_full_config(bank_id, request_context)
estimated_batch_size = max(1, (max_chunk_tokens // bank_config.retain_chunk_size) * 2)
chunks_dict = {}
encoding = _get_tiktoken_encoding()
chunk_offset = 0
# Fetch chunks in batches until we run out of budget or chunks
while chunk_offset < len(chunk_ids_ordered) and total_chunk_tokens < max_chunk_tokens:
# Get next batch of chunk IDs
batch_chunk_ids = chunk_ids_ordered[chunk_offset : chunk_offset + estimated_batch_size]
chunk_offset += estimated_batch_size
# Fetch chunk data from database
async with acquire_with_retry(pool) as conn:
chunks_rows = await conn.fetch(
f"""
SELECT chunk_id, chunk_text, chunk_index
FROM {fq_table("chunks")}
WHERE chunk_id = ANY($1::text[])
""",
batch_chunk_ids,
)
# Create a lookup dict for fast access (preserves order from batch_chunk_ids)
chunks_lookup = {row["chunk_id"]: row for row in chunks_rows}
# Process chunks in order, respecting token budget
for chunk_id in batch_chunk_ids:
if chunk_id not in chunks_lookup:
continue
row = chunks_lookup[chunk_id]
chunk_text = row["chunk_text"]
chunk_tokens = len(encoding.encode(chunk_text))
# Check if adding this chunk would exceed the limit
if total_chunk_tokens + chunk_tokens > max_chunk_tokens:
# Truncate the chunk to fit within the remaining budget
remaining_tokens = max_chunk_tokens - total_chunk_tokens
if remaining_tokens > 0:
# Truncate to remaining tokens
truncated_text = encoding.decode(encoding.encode(chunk_text)[:remaining_tokens])
chunks_dict[chunk_id] = ChunkInfo(
chunk_text=truncated_text, chunk_index=row["chunk_index"], truncated=True
)
total_chunk_tokens = max_chunk_tokens
# Budget exhausted - stop fetching more batches
break
else:
chunks_dict[chunk_id] = ChunkInfo(
chunk_text=chunk_text, chunk_index=row["chunk_index"], truncated=False
)
total_chunk_tokens += chunk_tokens
# If we hit the budget limit in this batch, stop fetching more batches
if total_chunk_tokens >= max_chunk_tokens:
break
# Step 6: Token budget filtering
step_start = time.time()
@ -2446,68 +2532,6 @@ class MemoryEngine(MemoryEngineInterface):
# Entity observations removed - always set to None
entities_dict = None
# Fetch chunks if requested
chunks_dict = None
total_chunk_tokens = 0
if include_chunks and top_scored:
from .response_models import ChunkInfo
# Collect chunk_ids in order of fact relevance (preserving order from top_scored)
# Use a list to maintain order, but track seen chunks to avoid duplicates
chunk_ids_ordered = []
seen_chunk_ids = set()
for sr in top_scored:
chunk_id = sr.retrieval.chunk_id
if chunk_id and chunk_id not in seen_chunk_ids:
chunk_ids_ordered.append(chunk_id)
seen_chunk_ids.add(chunk_id)
if chunk_ids_ordered:
# Fetch chunk data from database using chunk_ids (no ORDER BY to preserve input order)
async with acquire_with_retry(pool) as conn:
chunks_rows = await conn.fetch(
f"""
SELECT chunk_id, chunk_text, chunk_index
FROM {fq_table("chunks")}
WHERE chunk_id = ANY($1::text[])
""",
chunk_ids_ordered,
)
# Create a lookup dict for fast access
chunks_lookup = {row["chunk_id"]: row for row in chunks_rows}
# Apply token limit and build chunks_dict in the order of chunk_ids_ordered
chunks_dict = {}
encoding = _get_tiktoken_encoding()
for chunk_id in chunk_ids_ordered:
if chunk_id not in chunks_lookup:
continue
row = chunks_lookup[chunk_id]
chunk_text = row["chunk_text"]
chunk_tokens = len(encoding.encode(chunk_text))
# Check if adding this chunk would exceed the limit
if total_chunk_tokens + chunk_tokens > max_chunk_tokens:
# Truncate the chunk to fit within the remaining budget
remaining_tokens = max_chunk_tokens - total_chunk_tokens
if remaining_tokens > 0:
# Truncate to remaining tokens
truncated_text = encoding.decode(encoding.encode(chunk_text)[:remaining_tokens])
chunks_dict[chunk_id] = ChunkInfo(
chunk_text=truncated_text, chunk_index=row["chunk_index"], truncated=True
)
total_chunk_tokens = max_chunk_tokens
# Stop adding more chunks once we hit the limit
break
else:
chunks_dict[chunk_id] = ChunkInfo(
chunk_text=chunk_text, chunk_index=row["chunk_index"], truncated=False
)
total_chunk_tokens += chunk_tokens
# Finalize trace if enabled
trace_dict = None
if tracer:

View file

@ -0,0 +1,274 @@
"""
Test that recall chunks are fetched independently of max_tokens filtering.
This test verifies the new behavior where:
1. Chunks are fetched BEFORE max_tokens filtering
2. max_tokens=0 returns 0 facts but can still return chunks
3. Chunks are fetched in batches to handle varying chunk sizes
"""
import pytest
import pytest_asyncio
from hindsight_api.engine.memory_engine import Budget
@pytest.mark.asyncio
async def test_recall_chunks_independent_of_max_tokens(memory, request_context):
"""
Test that chunks are fetched independently of max_tokens.
When max_tokens=0, recall should:
- Return 0 memory facts
- Still return chunks (up to max_chunk_tokens)
- Chunks should come from top-scored results before token filtering
"""
bank_id = "test-chunks-independence"
try:
# Retain some test content with substantial size to generate chunks
test_content = """
The quantum computing research team at MIT has made significant breakthroughs.
Dr. Sarah Chen leads the team and focuses on quantum error correction.
The team published three papers in Nature Physics this year.
Their work on topological qubits shows promise for scalable quantum computers.
Collaborators include IBM Research and Google Quantum AI.
The research is funded by a $5M NSF grant running through 2026.
""" * 10 # Repeat to ensure we get multiple chunks
await memory.retain_async(
bank_id=bank_id,
content=test_content,
context="research notes",
request_context=request_context,
)
# Test 1: Normal recall with both facts and chunks
result_normal = await memory.recall_async(
bank_id=bank_id,
query="quantum computing",
max_tokens=4096, # Normal token budget
include_chunks=True,
max_chunk_tokens=2000,
budget=Budget.MID,
request_context=request_context,
)
assert len(result_normal.results) > 0, "Should return memory facts with normal max_tokens"
assert result_normal.chunks is not None, "Should include chunks when requested"
assert len(result_normal.chunks) > 0, "Should return at least one chunk"
# Test 2: Recall with max_tokens=0 but chunks enabled
result_chunks_only = await memory.recall_async(
bank_id=bank_id,
query="quantum computing",
max_tokens=0, # Zero token budget for facts
include_chunks=True,
max_chunk_tokens=2000, # But allow chunks
budget=Budget.MID,
request_context=request_context,
)
# Key assertions for new behavior
assert len(result_chunks_only.results) == 0, "max_tokens=0 should return 0 facts"
assert result_chunks_only.chunks is not None, "Should still include chunks dict"
assert len(result_chunks_only.chunks) > 0, "Should return chunks even with max_tokens=0"
# Verify chunks are from the same content (non-empty text)
for chunk_id, chunk_info in result_chunks_only.chunks.items():
assert len(chunk_info.chunk_text) > 0, "Chunks should contain text"
assert chunk_info.chunk_index >= 0, "Chunk should have valid index"
finally:
# Cleanup
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_recall_chunks_batching_with_varying_sizes(memory, request_context):
"""
Test that chunk batching works correctly with varying chunk sizes.
This verifies that:
1. Chunks are fetched in batches until token budget is exhausted
2. The system handles varying chunk sizes across documents
3. Token budget is respected across multiple batch fetches
"""
bank_id = "test-chunks-batching"
try:
# Retain multiple documents with different content sizes
# Document 1: Short content (small chunks)
await memory.retain_async(
bank_id=bank_id,
content="Alice is a software engineer who specializes in Python programming and machine learning.",
context="doc1",
request_context=request_context,
)
# Document 2: Medium content
content_bob = """
Bob works as a data scientist at a tech startup in San Francisco.
He has expertise in natural language processing and computer vision.
Bob completed his PhD at Stanford University in 2020.
He leads a team of five engineers working on AI-powered recommendation systems.
""" * 5
await memory.retain_async(
bank_id=bank_id,
content=content_bob,
context="doc2",
request_context=request_context,
)
# Document 3: Long content (large chunks)
content_charlie = """
Charlie is the CTO of a growing AI company focused on healthcare applications.
He has over 15 years of experience in software architecture and distributed systems.
Charlie's team builds machine learning models for medical image analysis and diagnosis.
The company recently raised $50 million in Series B funding.
They have partnerships with major hospitals in the United States and Europe.
Charlie holds several patents in medical imaging and deep learning.
""" * 20
await memory.retain_async(
bank_id=bank_id,
content=content_charlie,
context="doc3",
request_context=request_context,
)
# Recall with modest chunk token budget
result = await memory.recall_async(
bank_id=bank_id,
query="Alice Bob Charlie",
max_tokens=0, # No facts, only chunks
include_chunks=True,
max_chunk_tokens=1000, # Limited chunk budget
budget=Budget.MID,
request_context=request_context,
)
assert len(result.results) == 0, "Should return 0 facts with max_tokens=0"
assert result.chunks is not None, "Should include chunks"
# Verify we got chunks and respected the token budget
if len(result.chunks) > 0:
# Count total tokens (approximate)
total_chunk_chars = sum(len(chunk.chunk_text) for chunk in result.chunks.values())
# Very rough estimate: 1 token ≈ 4 characters
estimated_tokens = total_chunk_chars // 4
# Should be reasonably close to budget (within 2x due to estimation and batching)
assert estimated_tokens <= 1000 * 2, f"Should respect chunk token budget (got ~{estimated_tokens} tokens)"
finally:
# Cleanup
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_recall_chunks_ordering_by_relevance(memory, request_context):
"""
Test that chunks are returned in order of fact relevance.
Chunks should be ordered based on the top-scored (reranked) results,
not in document order or random order.
"""
bank_id = "test-chunks-ordering"
try:
# Retain content with different relevance to query
await memory.retain_async(
bank_id=bank_id,
content="The Python programming language is widely used for machine learning and data science applications.",
context="topic: Python",
request_context=request_context,
)
await memory.retain_async(
bank_id=bank_id,
content="JavaScript is commonly used for web development and frontend applications.",
context="topic: JavaScript",
request_context=request_context,
)
await memory.retain_async(
bank_id=bank_id,
content="Python's scikit-learn library is excellent for traditional machine learning tasks and model training.",
context="topic: Python ML",
request_context=request_context,
)
# Query specifically about Python - should rank Python facts higher
result = await memory.recall_async(
bank_id=bank_id,
query="Python machine learning",
max_tokens=0, # No facts
include_chunks=True,
max_chunk_tokens=5000, # Enough for all chunks
budget=Budget.HIGH, # Use high budget for better recall
request_context=request_context,
)
assert len(result.results) == 0, "Should return 0 facts with max_tokens=0"
assert result.chunks is not None, "Should include chunks"
# We should get chunks, and they should be ordered by relevance
# The exact ordering depends on the reranker, but we should have chunks
assert len(result.chunks) > 0, "Should return chunks from relevant facts"
# Verify chunks contain relevant content
all_chunk_text = " ".join(chunk.chunk_text for chunk in result.chunks.values())
# At least some chunks should mention Python (higher relevance)
# This is a soft check since exact ordering depends on scoring
assert "Python" in all_chunk_text or "python" in all_chunk_text.lower(), \
"Chunks should include content about Python (relevant to query)"
finally:
# Cleanup
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_recall_chunks_without_include_flag(memory, request_context):
"""
Test that chunks are NOT returned when include_chunks=False (default).
This ensures backward compatibility - chunks are only fetched when explicitly requested.
"""
bank_id = "test-chunks-no-include"
try:
# Retain content
test_content = """
Sarah is a product manager at a fintech company in New York.
She specializes in user experience design and agile methodologies.
Sarah graduated from MIT with a degree in computer science.
She has led the development of three successful mobile banking applications.
"""
await memory.retain_async(
bank_id=bank_id,
content=test_content,
request_context=request_context,
)
# Recall without include_chunks flag (default is False)
result = await memory.recall_async(
bank_id=bank_id,
query="Sarah product manager",
max_tokens=4096,
request_context=request_context,
# include_chunks=False is the default
)
# Should have facts but no chunks
assert len(result.results) > 0, "Should return facts"
assert result.chunks is None or len(result.chunks) == 0, \
"Should NOT return chunks when include_chunks=False"
finally:
# Cleanup
await memory.delete_bank(bank_id, request_context=request_context)

View file

@ -44,10 +44,10 @@ Make sure you've completed the [Quick Start](./quickstart) to install the client
| `query` | string | required | Natural language query |
| `types` | list | all | Filter: `world`, `experience`, `observation` |
| `budget` | string | "mid" | Budget level: `low`, `mid`, `high` |
| `max_tokens` | int | 4096 | Token budget for results |
| `max_tokens` | int | 4096 | Token budget for memory facts (text only) |
| `trace` | bool | false | Enable trace output for debugging |
| `include_chunks` | bool | false | Include raw text chunks that generated the memories |
| `max_chunk_tokens` | int | 500 | Token budget for chunks |
| `max_chunk_tokens` | int | 500 | Token budget for chunks (independent of `max_tokens`) |
| `tags` | list | None | Filter memories by tags (see [Tag Filtering](#filter-by-tags)) |
| `tags_match` | string | "any" | How to match tags: `any`, `all`, `any_strict`, `all_strict` |
@ -93,6 +93,15 @@ The `max_tokens` parameter lets you control how much of your agent's context bud
This design means you never have to guess whether 10 results or 50 results will fit your context. Just specify the token budget and Hindsight returns as many relevant memories as will fit.
:::note Chunks are Independent
When `include_chunks=True`, chunks are fetched **independently** of the `max_tokens` filtering. This means:
- Setting `max_tokens=0` will return **0 memory facts** but can still return **chunks** (up to `max_chunk_tokens`)
- Chunks are based on the top-scored (reranked) results **before** token filtering
- Chunks are fetched in batches (batch size estimated as `(max_chunk_tokens / retain_chunk_size) * 2`) until the token budget is exhausted
- This batching approach handles varying chunk sizes across documents efficiently
- This allows you to retrieve raw source text without memory facts when needed
:::
## Budget Levels
The `budget` parameter controls graph traversal depth: