fix: guard against top_k=0 reaching the vector store
Turbopuffer rejects top_k=0 with a 400 ('top_k must be between 1 and
10000'). The fix returns [] for a non-positive top_k, before the embedding call), and
floor the semantic budget at 1 so an explicitly requested search isn't
silently allocated zero.
This commit is contained in:
parent
4d3ab1c36b
commit
70f4dc5f48
|
|
@ -355,6 +355,13 @@ async def query_documents(
|
|||
Returns:
|
||||
Sequence of matching documents
|
||||
"""
|
||||
# A non-positive top_k means "no results wanted" -- return before embedding
|
||||
# or querying. Turbopuffer rejects top_k=0 with a 400 (pgvector would
|
||||
# silently do LIMIT 0), and callers derive top_k from budget arithmetic or
|
||||
# LLM tool input, so neither is guaranteed positive.
|
||||
if top_k <= 0:
|
||||
return []
|
||||
|
||||
# Use provided embedding or generate one
|
||||
if embedding is None:
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -318,11 +318,14 @@ class RepresentationManager:
|
|||
|
||||
total = max_observations
|
||||
|
||||
# Calculate how many observations to get from each source
|
||||
# Calculate how many observations to get from each source.
|
||||
# Floor of 1 when a semantic query was explicitly requested: `total // 3`
|
||||
# rounds to 0 for total < 3, which would allocate no budget to the source
|
||||
# the caller actually asked to curate around.
|
||||
semantic_observations = (
|
||||
min(
|
||||
max(
|
||||
0,
|
||||
1,
|
||||
semantic_search_top_k
|
||||
if semantic_search_top_k is not None
|
||||
else total // 3,
|
||||
|
|
|
|||
|
|
@ -629,3 +629,72 @@ class TestRepresentationManagerSave:
|
|||
assert len(saved.created_documents) == 0
|
||||
mock_embed.assert_not_awaited()
|
||||
mock_save.assert_not_awaited()
|
||||
|
||||
|
||||
class TestVectorQueryTopKFloor:
|
||||
"""Regression for HONCHO-19Q / HONCHO-4Q4.
|
||||
|
||||
A top_k of 0 reached Turbopuffer, which rejects it with a 400
|
||||
('top_k must be between 1 and 10000'). Two independent paths produced it:
|
||||
the working-representation budget split (``total // 3`` rounds to 0 for
|
||||
max_conclusions < 3) and the dialectic ``search_memory`` tool, whose
|
||||
LLM-supplied top_k has an upper clamp but no floor.
|
||||
"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_documents_returns_empty_without_querying_on_zero_top_k(self):
|
||||
"""The choke point every semantic document query routes through."""
|
||||
from src.crud.document import query_documents
|
||||
|
||||
with (
|
||||
patch(
|
||||
"src.crud.document.embedding_client.embed", new=AsyncMock()
|
||||
) as mock_embed,
|
||||
patch(
|
||||
"src.crud.document.query_external_vector_document_ids",
|
||||
new=AsyncMock(),
|
||||
) as mock_vector,
|
||||
):
|
||||
for top_k in (0, -1):
|
||||
assert (
|
||||
await query_documents(
|
||||
None,
|
||||
"workspace",
|
||||
"query",
|
||||
observer="observer",
|
||||
observed="observed",
|
||||
top_k=top_k,
|
||||
)
|
||||
== []
|
||||
)
|
||||
|
||||
mock_embed.assert_not_awaited()
|
||||
mock_vector.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_requested_semantic_search_always_gets_budget(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
"""max_conclusions < 3 must not allocate 0 to an explicitly requested search."""
|
||||
test_workspace, test_peer = sample_data
|
||||
manager = RepresentationManager(
|
||||
test_workspace.name, observer=test_peer.name, observed=test_peer.name
|
||||
)
|
||||
|
||||
for max_observations in (1, 2, 100):
|
||||
with patch(
|
||||
"src.crud.query_documents", new=AsyncMock(return_value=[])
|
||||
) as mock_query:
|
||||
await manager._get_working_representation_internal( # pyright: ignore[reportPrivateUsage]
|
||||
db_session,
|
||||
include_semantic_query="what do they like?",
|
||||
embedding=[0.1],
|
||||
max_observations=max_observations,
|
||||
)
|
||||
|
||||
assert mock_query.await_args is not None
|
||||
top_k = mock_query.await_args.kwargs["top_k"]
|
||||
assert top_k >= 1, f"max_observations={max_observations} gave top_k={top_k}"
|
||||
assert top_k <= max_observations
|
||||
|
|
|
|||
Loading…
Reference in New Issue