[DEV-16408] Set toolChoice to "any"

usaspending_api/llm/assistants/filter_search.py

@@ -113,7 +113,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:
@@ -123,7 +123,8 @@ 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
+
+        return {**self._tool_config, "toolChoice": {"auto": {}}}
 
     @staticmethod
     def _fiscal_year_date_context() -> str:
@@ -252,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")
@@ -260,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}",
@@ -276,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:

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,6 +457,13 @@ 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_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": {}}
+
+        assistant.tool_iterations = assistant.MAX_TOOL_ITERATIONS - 1
+        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):
         """tool_config injects a DB-sourced defCodes description into execute_filter's schema.