From 93122e8017be74edabbd5250485dbab3164234ac Mon Sep 17 00:00:00 2001 From: aguest-kc Date: Mon, 19 Oct 2026 14:14:12 -0500 Subject: [PATCH 1/4] [DEV-16408] Add toolChoice --- usaspending_api/llm/assistants/filter_search.py | 8 ++++++-- .../integration/test_filter_search_assistant.py | 13 +++++++++++++ 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/usaspending_api/llm/assistants/filter_search.py b/usaspending_api/llm/assistants/filter_search.py index 6f710af748..c9bcddbb18 100644 --- a/usaspending_api/llm/assistants/filter_search.py +++ b/usaspending_api/llm/assistants/filter_search.py @@ -112,7 +112,7 @@ async def _create_message_from_response(self, response: dict) -> Message: self.messages.append(output_message) return message - async def aget_tool_config(self) -> dict[str, list[dict]]: + async def aget_tool_config(self) -> dict[str, Any]: if self._tool_config is None: specs = [tool.description.model_dump() for tool in self.tools] for spec in specs: @@ -122,7 +122,11 @@ async def aget_tool_config(self) -> dict[str, list[dict]]: self._tool_config = { "tools": [{"toolSpec": {"inputSchema": {"json": spec.pop("input_schema")}, **spec}} for spec in specs] } - return self._tool_config + + if self.tool_iterations >= self.MAX_TOOL_ITERATIONS - 1: + return {**self._tool_config, "toolChoice": {"tool": {"name": self.COMPLETION_TOOL_NAME}}} + + return {**self._tool_config, "toolChoice": {"auto": {}}} @staticmethod def _fiscal_year_date_context() -> str: diff --git a/usaspending_api/llm/tests/integration/test_filter_search_assistant.py b/usaspending_api/llm/tests/integration/test_filter_search_assistant.py index 3aa054f3b0..dfa07747e4 100644 --- a/usaspending_api/llm/tests/integration/test_filter_search_assistant.py +++ b/usaspending_api/llm/tests/integration/test_filter_search_assistant.py @@ -382,6 +382,19 @@ async def test_tool_config_property(self, assistant, mock_tool): assert "toolSpec" in config["tools"][0] assert "inputSchema" in config["tools"][0]["toolSpec"] + async def test_tool_config_defaults_to_auto_tool_choice(self, assistant): + """On a normal turn the model is free to pick which tool (if any) to call.""" + config = await assistant.aget_tool_config() + + assert config["toolChoice"] == {"auto": {}} + + async def test_tool_config_forces_completion_tool_on_final_iteration(self, assistant): + """On the last tool-use iteration, execute_filter is forced so the search always completes.""" + assistant.tool_iterations = assistant.MAX_TOOL_ITERATIONS - 1 + config = await assistant.aget_tool_config() + + assert config["toolChoice"] == {"tool": {"name": assistant.COMPLETION_TOOL_NAME}} + @pytest.mark.django_db async def test_tool_config_patches_defcodes_description_from_db(self, mock_assistant, mock_session): """tool_config injects a DB-sourced defCodes description into execute_filter's schema. From 95d84b3624f9180e7631c074fcf93ca3c6b9a0eb Mon Sep 17 00:00:00 2001 From: aguest-kc Date: Mon, 19 Oct 2026 14:45:49 -0500 Subject: [PATCH 2/4] [DEV-16408] Set toolChoice to any instead of auto --- usaspending_api/llm/assistants/filter_search.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/usaspending_api/llm/assistants/filter_search.py b/usaspending_api/llm/assistants/filter_search.py index c9bcddbb18..eb55bf42e0 100644 --- a/usaspending_api/llm/assistants/filter_search.py +++ b/usaspending_api/llm/assistants/filter_search.py @@ -126,7 +126,7 @@ async def aget_tool_config(self) -> dict[str, Any]: if self.tool_iterations >= self.MAX_TOOL_ITERATIONS - 1: return {**self._tool_config, "toolChoice": {"tool": {"name": self.COMPLETION_TOOL_NAME}}} - return {**self._tool_config, "toolChoice": {"auto": {}}} + return {**self._tool_config, "toolChoice": {"any": {}}} @staticmethod def _fiscal_year_date_context() -> str: From 4b6ec9556c1f25baabe1861ca3e7efeecfc33232 Mon Sep 17 00:00:00 2001 From: aguest-kc Date: Tue, 20 Oct 2026 05:50:26 -0500 Subject: [PATCH 3/4] [DEV-16408] Fix test --- .../llm/tests/integration/test_filter_search_assistant.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/usaspending_api/llm/tests/integration/test_filter_search_assistant.py b/usaspending_api/llm/tests/integration/test_filter_search_assistant.py index dfa07747e4..07e57df098 100644 --- a/usaspending_api/llm/tests/integration/test_filter_search_assistant.py +++ b/usaspending_api/llm/tests/integration/test_filter_search_assistant.py @@ -382,11 +382,11 @@ async def test_tool_config_property(self, assistant, mock_tool): assert "toolSpec" in config["tools"][0] assert "inputSchema" in config["tools"][0]["toolSpec"] - async def test_tool_config_defaults_to_auto_tool_choice(self, assistant): - """On a normal turn the model is free to pick which tool (if any) to call.""" + async def test_tool_config_defaults_to_any_tool_choice(self, assistant): + """Force the model to call one of the configured tools.""" config = await assistant.aget_tool_config() - assert config["toolChoice"] == {"auto": {}} + assert config["toolChoice"] == {"any": {}} async def test_tool_config_forces_completion_tool_on_final_iteration(self, assistant): """On the last tool-use iteration, execute_filter is forced so the search always completes.""" From 0be1d4fe8e727d046ab12e60ec7e9f4c17b79a7e Mon Sep 17 00:00:00 2001 From: aguest-kc Date: Tue, 27 Oct 2026 10:55:27 -0500 Subject: [PATCH 4/4] [DEV-16408] Use execute_filter on last loop --- .../llm/assistants/filter_search.py | 65 ++++++++++++-- .../test_filter_search_assistant.py | 88 ++++++++++++++++--- 2 files changed, 133 insertions(+), 20 deletions(-) diff --git a/usaspending_api/llm/assistants/filter_search.py b/usaspending_api/llm/assistants/filter_search.py index 3e0e1352eb..069e476a21 100644 --- a/usaspending_api/llm/assistants/filter_search.py +++ b/usaspending_api/llm/assistants/filter_search.py @@ -124,10 +124,7 @@ async def aget_tool_config(self) -> dict[str, Any]: "tools": [{"toolSpec": {"inputSchema": {"json": spec.pop("input_schema")}, **spec}} for spec in specs] } - if self.tool_iterations >= self.MAX_TOOL_ITERATIONS - 1: - return {**self._tool_config, "toolChoice": {"tool": {"name": self.COMPLETION_TOOL_NAME}}} - - return {**self._tool_config, "toolChoice": {"any": {}}} + return {**self._tool_config, "toolChoice": {"auto": {}}} @staticmethod def _fiscal_year_date_context() -> str: @@ -256,6 +253,12 @@ async def search(self, query: str) -> AsyncGenerator[dict[str, str], None]: }, ) + if not search_complete: + async for event in self._aforce_completion(client, query, response, m): + yield event + if event.get("type") == "search_complete": + search_complete = True + # Calculate total token usage for this search. totals = await self.session.messages.aaggregate( input_tokens=Sum("input_tokens"), output_tokens=Sum("output_tokens") @@ -264,9 +267,6 @@ async def search(self, query: str) -> AsyncGenerator[dict[str, str], None]: total_output_tokens = totals["output_tokens"] or 0 total_tokens = total_input_tokens + total_output_tokens - if not search_complete: - raise RuntimeError(f"Filter search session {self.session.id} ended without completing") - # Log each search. logger.info( f"Search completed for session {self.session.id}", @@ -280,6 +280,57 @@ async def search(self, query: str) -> AsyncGenerator[dict[str, str], None]: }, ) + if not search_complete: + raise RuntimeError(f"Filter search session {self.session.id} ended without completing") + + async def _aforce_completion( + self, client: Any, query: str, response: dict, message: Message + ) -> AsyncGenerator[dict, None]: + """Force the completion tool after the main loop ends without it.""" + if response["stopReason"] == "tool_use": + pending = [request for request in response["output"]["message"]["content"] if "toolUse" in request] + + async for event in self.handle_tool_use(pending, message): + yield event + if event.get("type") == "search_complete": + return + + logger.info( + f"Search not complete after {self.tool_iterations} tool iterations. " + f"Force calling {self.COMPLETION_TOOL_NAME}." + ) + model_message = ( + f"Search is not complete and {self.COMPLETION_TOOL_NAME} was never called. As a last resort, " + f'extract a single relevant keyword from the user\'s original query ("{query}") and call ' + f"{self.COMPLETION_TOOL_NAME} using that keyword." + ) + + await Message.objects.acreate( + session=self.session, role="user", message=model_message, order=self.message_order + ) + self.message_order += 1 + + # Converse requires alternating roles, so merge into a trailing toolResult user turn. + if self.messages[-1]["role"] == "user": + self.messages[-1]["content"].append({"text": model_message}) + else: + self.messages.append({"role": "user", "content": [{"text": model_message}]}) + + response = await client.converse( + modelId=self.assistant.ai_model.model_id, + messages=self.messages, + toolConfig={**await self.aget_tool_config(), "toolChoice": {"tool": {"name": self.COMPLETION_TOOL_NAME}}}, + system=[{"text": self.system_message}], + inferenceConfig=self.inference_config, + ) + self._sanitize_tool_use_input(response["output"]["message"]["content"]) + + forced_message = await self._create_message_from_response(response) + tool_requests = [request for request in response["output"]["message"]["content"] if "toolUse" in request] + + async for event in self.handle_tool_use(tool_requests, forced_message): + yield event + async def handle_tool_use(self, tool_requests: list[dict], message: Message) -> AsyncGenerator[dict, None]: tool_result_message = {"role": "user", "content": []} for tool_request in tool_requests: diff --git a/usaspending_api/llm/tests/integration/test_filter_search_assistant.py b/usaspending_api/llm/tests/integration/test_filter_search_assistant.py index 250ad8d8af..b45f5abd05 100644 --- a/usaspending_api/llm/tests/integration/test_filter_search_assistant.py +++ b/usaspending_api/llm/tests/integration/test_filter_search_assistant.py @@ -167,8 +167,9 @@ def create_message(**kwargs): assert len(results) == 1 assert results[0]["type"] == "search_start" - assert assistant.client.converse.call_count == 1 - assert mock_message_create.call_count == 2 # User message + assistant message + assert assistant.client.converse.call_count == 2 # Initial call + forced last-ditch call + # User query, assistant reply, last-ditch instruction, last-ditch assistant reply + assert mock_message_create.call_count == 4 @patch("usaspending_api.llm.models.db_models.Message.objects.acreate", new_callable=AsyncMock) async def test_assistant_message_text_is_persisted(self, mock_message_create, assistant): @@ -378,7 +379,74 @@ def create_message(**kwargs): _ = [event async for event in assistant.search("test query")] assert assistant.tool_iterations == assistant.MAX_TOOL_ITERATIONS - assert assistant.client.converse.call_count == assistant.MAX_TOOL_ITERATIONS + 1 + assert assistant.client.converse.call_count == assistant.MAX_TOOL_ITERATIONS + 2 + + # The pending toolUse from the final loop response is answered, and the last-ditch + # instruction shares that user turn instead of creating two consecutive user messages. + instruction_index = next( + i + for i, msg in enumerate(assistant.messages) + if msg["role"] == "user" and any("last resort" in block.get("text", "") for block in msg["content"]) + ) + instruction_turn = assistant.messages[instruction_index] + assert any("toolResult" in block for block in instruction_turn["content"]) + assert assistant.messages[instruction_index - 1]["role"] == "assistant" + + roles = [msg["role"] for msg in assistant.messages] + assert all(a != b for a, b in zip(roles, roles[1:], strict=False)) + + @patch("usaspending_api.llm.models.db_models.ToolUse.objects.acreate", new_callable=AsyncMock) + @patch("usaspending_api.llm.models.db_models.Message.objects.acreate", new_callable=AsyncMock) + async def test_forces_completion_tool(self, mock_message_create, mock_tool_use_create, assistant): + """If the model never calls execute_filter, one final call forces it with a keyword instruction.""" + + def create_message(**kwargs): + mock_message = Mock() + mock_message.input_tokens = kwargs.get("input_tokens", 10) + mock_message.output_tokens = kwargs.get("output_tokens", 20) + mock_message.id = len(assistant.session._test_messages) + 1 + mock_message.tool_uses = Mock() + mock_message.tool_uses.count.return_value = 0 + assistant.session._test_messages.append(mock_message) + return mock_message + + mock_message_create.side_effect = create_message + mock_tool_use_create.return_value.asave = AsyncMock() + + text_only_response = { + "output": {"message": {"role": "assistant", "content": [{"text": "I need more details."}]}}, + "usage": {"inputTokens": 10, "outputTokens": 20}, + "metrics": {"latencyMs": 100}, + "stopReason": "end_turn", + } + forced_response = { + "output": { + "message": { + "role": "assistant", + "content": [ + {"toolUse": {"toolUseId": "tool-999", "name": "execute_filter", "input": {"query": "test"}}} + ], + } + }, + "usage": {"inputTokens": 10, "outputTokens": 20}, + "metrics": {"latencyMs": 100}, + "stopReason": "tool_use", + } + assistant.client.converse.side_effect = [text_only_response, forced_response] + + results = [event async for event in assistant.search("test query")] + + assert assistant.client.converse.call_count == 2 + first_call, last_ditch_call = assistant.client.converse.call_args_list + assert first_call.kwargs["toolConfig"]["toolChoice"] == {"auto": {}} + assert last_ditch_call.kwargs["toolConfig"]["toolChoice"] == {"tool": {"name": assistant.COMPLETION_TOOL_NAME}} + assert any( + "last resort" in block.get("text", "") + for msg in assistant.messages + if msg["role"] == "user" + for block in msg["content"] + ) + assert any(r["type"] == "search_complete" and r["result"] == "abc123" for r in results) async def test_tool_config_property(self, assistant, mock_tool): """Test that tool_config is properly formatted.""" @@ -389,18 +457,12 @@ async def test_tool_config_property(self, assistant, mock_tool): assert "toolSpec" in config["tools"][0] assert "inputSchema" in config["tools"][0]["toolSpec"] - async def test_tool_config_defaults_to_any_tool_choice(self, assistant): - """Force the model to call one of the configured tools.""" - config = await assistant.aget_tool_config() - - assert config["toolChoice"] == {"any": {}} + async def test_tool_config_always_uses_auto_tool_choice(self, assistant): + """The main loop leaves tool selection to the model; forcing only happens in the last-ditch call.""" + assert (await assistant.aget_tool_config())["toolChoice"] == {"auto": {}} - async def test_tool_config_forces_completion_tool_on_final_iteration(self, assistant): - """On the last tool-use iteration, execute_filter is forced so the search always completes.""" assistant.tool_iterations = assistant.MAX_TOOL_ITERATIONS - 1 - config = await assistant.aget_tool_config() - - assert config["toolChoice"] == {"tool": {"name": assistant.COMPLETION_TOOL_NAME}} + assert (await assistant.aget_tool_config())["toolChoice"] == {"auto": {}} @pytest.mark.django_db async def test_tool_config_patches_defcodes_description_from_db(self, mock_assistant, mock_session):