From 8bd795c50ec8caee9b91209c179258c230f56d20 Mon Sep 17 00:00:00 2001 From: Aakash Kattelu Date: Wed, 26 Aug 2026 16:41:04 -0400 Subject: [PATCH] fix: stop top_k=0 from reaching Turbopuffer on message search HONCHO-19Q: dreamer search_messages passed LLM limit=0 through to Turbopuffer (top_k must be 1..10000). #970 guarded documents; this closes the message path and floors tool limits at 1. --- src/crud/document.py | 3 ++ src/crud/message.py | 3 ++ src/utils/agent_tools.py | 27 +++++++------ src/utils/search.py | 3 ++ src/vector_store/lancedb.py | 3 ++ src/vector_store/turbopuffer.py | 3 ++ tests/crud/test_representation_manager.py | 45 +++++++++++++++++++++ tests/utils/test_agent_tools.py | 49 +++++++++++++++++++++++ tests/vector_store/test_lancedb.py | 21 ++++++++++ tests/vector_store/test_turbopuffer.py | 23 +++++++++++ 10 files changed, 168 insertions(+), 12 deletions(-) diff --git a/src/crud/document.py b/src/crud/document.py index 1b04bb0a..dd9a3f46 100644 --- a/src/crud/document.py +++ b/src/crud/document.py @@ -228,6 +228,9 @@ async def query_external_vector_document_ids( empty list when the external store has no results, or None when the pgvector (DB-only) path should be used instead. """ + if top_k <= 0: + return [] + if _uses_pgvector(): return None diff --git a/src/crud/message.py b/src/crud/message.py index 9759cb91..15b4902c 100644 --- a/src/crud/message.py +++ b/src/crud/message.py @@ -747,6 +747,9 @@ async def _search_messages_external( Multiple vector records can map to the same message (chunked embeddings), so we oversample from the vector store and deduplicate by message_id. """ + if limit <= 0: + return [] + external_vector_store = get_external_vector_store() if external_vector_store is None: return [] diff --git a/src/utils/agent_tools.py b/src/utils/agent_tools.py index e31b4d62..b753c462 100644 --- a/src/utils/agent_tools.py +++ b/src/utils/agent_tools.py @@ -292,6 +292,11 @@ def _safe_int(value: Any, default: int) -> int: return default +def _bounded_int(value: Any, default: int, *, lo: int = 1, hi: int) -> int: + """Coerce a tool int into ``[lo, hi]``, falling back to ``default`` on bad input.""" + return max(lo, min(_safe_int(value, default), hi)) + + # Module-level lock registry for thread-safe observation creation. # Keyed by (workspace_name, observer, observed) to ensure all tool executors # operating on the same data share the same lock. @@ -1883,7 +1888,7 @@ async def _handle_search_memory( """Handle search_memory tool.""" from src.utils.types import ToolResult - top_k = min(_safe_int(tool_input.get("top_k"), 20), 40) + top_k = _bounded_int(tool_input.get("top_k"), 20, hi=40) query = tool_input["query"] try: with embedding_call_purpose( @@ -1944,7 +1949,7 @@ async def _handle_search_memory( # information. zero_hit_meta = {**search_meta, "results_count": 0} if ctx.agent_type in ("dialectic", "workspace_dialectic"): - limit = min(_safe_int(tool_input.get("top_k"), 20), 20) + limit = _bounded_int(tool_input.get("top_k"), 20, hi=20) message_output = None snippets = await crud.search_messages( workspace_name=ctx.workspace_name, @@ -2021,7 +2026,7 @@ async def _handle_search_messages( from src.utils.types import ToolResult query = tool_input["query"] - limit = min(_safe_int(tool_input.get("limit"), 10), 20) # Cap at 20 + limit = _bounded_int(tool_input.get("limit"), 10, hi=20) # Pre-compute embedding outside DB session to avoid holding a connection # during the external API call (same pattern as _handle_search_memory). with embedding_call_purpose( @@ -2064,10 +2069,8 @@ async def _handle_grep_messages( text = tool_input.get("text", "") if not text: return "ERROR: 'text' parameter is required" - limit = min(_safe_int(tool_input.get("limit"), 10), 30) # Cap at 30 - context_window = min( - _safe_int(tool_input.get("context_window"), 2), 2 - ) # Cap context + limit = _bounded_int(tool_input.get("limit"), 10, hi=30) + context_window = _bounded_int(tool_input.get("context_window"), 2, lo=0, hi=2) snippets = await crud.grep_messages( workspace_name=ctx.workspace_name, @@ -2120,7 +2123,7 @@ async def _handle_get_messages_by_date_range( """Handle get_messages_by_date_range tool.""" after_date_str = tool_input.get("after_date") before_date_str = tool_input.get("before_date") - limit = min(_safe_int(tool_input.get("limit"), 20), 20) + limit = _bounded_int(tool_input.get("limit"), 20, hi=20) order = tool_input.get("order", "desc") after_date = _parse_date(after_date_str, "after_date") @@ -2186,8 +2189,8 @@ async def _handle_search_messages_temporal( after_date_str = tool_input.get("after_date") before_date_str = tool_input.get("before_date") - limit = min(_safe_int(tool_input.get("limit"), 10), 10) - context_window = min(_safe_int(tool_input.get("context_window"), 2), 2) + limit = _bounded_int(tool_input.get("limit"), 10, hi=10) + context_window = _bounded_int(tool_input.get("context_window"), 2, lo=0, hi=2) after_date = _parse_date(after_date_str, "after_date") if isinstance(after_date, str): @@ -2257,7 +2260,7 @@ async def _handle_get_recent_observations( workspace_name=ctx.workspace_name, observer=ctx.observer, observed=ctx.observed, - limit=min(_safe_int(tool_input.get("limit"), 10), 100), + limit=_bounded_int(tool_input.get("limit"), 10, hi=100), session_name=ctx.session_name if session_only else None, ) representation = Representation.from_documents(documents) @@ -2283,7 +2286,7 @@ async def _handle_get_most_derived_observations( workspace_name=ctx.workspace_name, observer=ctx.observer, observed=ctx.observed, - limit=min(_safe_int(tool_input.get("limit"), 10), 100), + limit=_bounded_int(tool_input.get("limit"), 10, hi=100), ) representation = Representation.from_documents(documents) total_count = representation.len() diff --git a/src/utils/search.py b/src/utils/search.py index 329e082b..76e86fd4 100644 --- a/src/utils/search.py +++ b/src/utils/search.py @@ -82,6 +82,9 @@ async def query_external_vector_message_ids( filters: dict[str, Any] | None = None, ) -> list[str]: """Query the external vector store and return ordered message IDs.""" + if limit <= 0: + return [] + external_vector_store = get_external_vector_store() if external_vector_store is None: return [] diff --git a/src/vector_store/lancedb.py b/src/vector_store/lancedb.py index 1b4c4880..ab44c1bc 100644 --- a/src/vector_store/lancedb.py +++ b/src/vector_store/lancedb.py @@ -214,6 +214,9 @@ class LanceDBVectorStore(VectorStore): Returns: List of VectorQueryResult objects, ordered by similarity (most similar first) """ + if top_k <= 0: + return [] + table = await self._get_table(namespace) if table is None: logger.debug(f"Table {namespace} does not exist, returning empty results") diff --git a/src/vector_store/turbopuffer.py b/src/vector_store/turbopuffer.py index e7825310..214c5f8a 100644 --- a/src/vector_store/turbopuffer.py +++ b/src/vector_store/turbopuffer.py @@ -143,6 +143,9 @@ class TurbopufferVectorStore(VectorStore): Returns: List of VectorQueryResult objects, ordered by similarity (most similar first) """ + if top_k <= 0: + return [] + ns = self._get_namespace(namespace) try: diff --git a/tests/crud/test_representation_manager.py b/tests/crud/test_representation_manager.py index 3c2f57e4..3f3d6f40 100644 --- a/tests/crud/test_representation_manager.py +++ b/tests/crud/test_representation_manager.py @@ -757,3 +757,48 @@ class TestVectorQueryTopKFloor: 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 + + @pytest.mark.asyncio + async def test_search_messages_external_returns_empty_without_querying_on_zero_limit( + self, + ): + """Message vector search is the remaining path that still hit Turbopuffer.""" + from src.crud import message as message_crud + + with patch( + "src.crud.message.get_external_vector_store", + return_value=AsyncMock(), + ) as mock_get_store: + for limit in (0, -1): + assert ( + await message_crud._search_messages_external( # pyright: ignore[reportPrivateUsage] + "workspace", + [0.1, 0.2, 0.3], + limit, + ) + == [] + ) + + mock_get_store.assert_not_called() + + @pytest.mark.asyncio + async def test_query_external_vector_message_ids_skips_store_on_zero_limit( + self, + ): + from src.utils import search as search_utils + + with patch( + "src.utils.search.get_external_vector_store", + return_value=AsyncMock(), + ) as mock_get_store: + for limit in (0, -1): + assert ( + await search_utils.query_external_vector_message_ids( + "workspace", + [0.1, 0.2, 0.3], + limit, + ) + == [] + ) + + mock_get_store.assert_not_called() diff --git a/tests/utils/test_agent_tools.py b/tests/utils/test_agent_tools.py index ac45cabf..346296ef 100644 --- a/tests/utils/test_agent_tools.py +++ b/tests/utils/test_agent_tools.py @@ -19,6 +19,7 @@ from src.utils.agent_tools import ( PEER_CARD_ALLOWED_PREFIXES, ObservationsCreatedResult, ToolContext, + _bounded_int, # pyright: ignore[reportPrivateUsage] _handle_create_observations, # pyright: ignore[reportPrivateUsage] _handle_delete_observations, # pyright: ignore[reportPrivateUsage] _handle_extract_preferences, # pyright: ignore[reportPrivateUsage] @@ -952,6 +953,21 @@ class TestSearchMemory: assert query_embeddings[0] == fallback_embeddings[0] +class TestBoundedInt: + """Unit tests for tool-input clamping.""" + + def test_floors_nonpositive_to_one(self) -> None: + assert _bounded_int(0, 10, hi=20) == 1 + assert _bounded_int(-5, 10, hi=20) == 1 + + def test_caps_at_hi(self) -> None: + assert _bounded_int(100, 10, hi=20) == 20 + + def test_falls_back_on_bad_input(self) -> None: + assert _bounded_int("Infinity", 10, hi=20) == 10 + assert _bounded_int(None, 10, hi=20) == 10 + + @pytest.mark.asyncio class TestSearchMessages: """Tests for _handle_search_messages.""" @@ -971,6 +987,39 @@ class TestSearchMessages: assert isinstance(result, str | ToolResult) + async def test_limit_zero_is_floored_to_one( + self, + make_tool_context: Callable[..., ToolContext], + monkeypatch: pytest.MonkeyPatch, + ): + """LLM-supplied limit=0 must not reach the vector store as top_k=0.""" + ctx = make_tool_context() + seen_limits: list[int] = [] + + async def fake_embed(query: str) -> list[float]: + _ = query + return [0.1, 0.2, 0.3] + + async def fake_search_messages( + workspace_name: str, + session_name: str | None, + query: str, + limit: int = 10, + **_kwargs: Any, + ) -> list[Any]: + _ = (workspace_name, session_name, query) + seen_limits.append(limit) + return [] + + monkeypatch.setattr("src.utils.agent_tools.embedding_client.embed", fake_embed) + monkeypatch.setattr( + "src.utils.agent_tools.crud.search_messages", fake_search_messages + ) + + await _handle_search_messages(ctx, {"query": "anything", "limit": 0}) + + assert seen_limits == [1] + @pytest.mark.asyncio class TestGrepMessages: diff --git a/tests/vector_store/test_lancedb.py b/tests/vector_store/test_lancedb.py index 9b52708a..0d36dc56 100644 --- a/tests/vector_store/test_lancedb.py +++ b/tests/vector_store/test_lancedb.py @@ -170,3 +170,24 @@ async def test_query_filters_by_max_distance(store: LanceDBVectorStore) -> None: ) assert [r.id for r in results] == ["vec_close"] + + +@pytest.mark.asyncio +async def test_query_returns_empty_without_opening_table_on_nonpositive_top_k( + store: LanceDBVectorStore, + monkeypatch: pytest.MonkeyPatch, +) -> None: + get_table = AsyncMock(return_value=MagicMock()) + monkeypatch.setattr(store, "_get_table", get_table) + + for top_k in (0, -1): + assert ( + await store.query( + "honcho.msg.test", + [0.1, 0.2, 0.3, 0.4], + top_k=top_k, + ) + == [] + ) + + get_table.assert_not_awaited() diff --git a/tests/vector_store/test_turbopuffer.py b/tests/vector_store/test_turbopuffer.py index cdae44d5..f7ada717 100644 --- a/tests/vector_store/test_turbopuffer.py +++ b/tests/vector_store/test_turbopuffer.py @@ -141,3 +141,26 @@ async def test_query_can_skip_attributes( namespace_mock.query.assert_awaited_once() assert namespace_mock.query.await_args.kwargs["include_attributes"] is False + + +@pytest.mark.asyncio +async def test_query_returns_empty_without_calling_api_on_nonpositive_top_k( + store: TurbopufferVectorStore, +) -> None: + """Turbopuffer rejects top_k < 1; never hit the network with a bad value.""" + namespace_mock = MagicMock() + namespace_mock.query = AsyncMock() + store._get_namespace = MagicMock(return_value=namespace_mock) # pyright: ignore[reportPrivateUsage] + + for top_k in (0, -1): + assert ( + await store.query( + "honcho.msg.test", + [0.1, 0.2, 0.3, 0.4], + top_k=top_k, + ) + == [] + ) + + store._get_namespace.assert_not_called() # pyright: ignore[reportPrivateUsage] + namespace_mock.query.assert_not_awaited()