feat: force tool calling support (#751)

Co-authored-by: Konrad Czarnota <[email protected]>
Co-authored-by: GlockPL <[email protected]>
Co-authored-by: GlockPL <[email protected]>

Author: GlockPL

docs/how-to/agents/define_and_use_agents.md

@@ -49,6 +49,14 @@ The result is an [AgentResult][ragbits.agents.AgentResult], which includes the m
 
 You can find the complete code example in the Ragbits repository [here](https://github.com/deepsense-ai/ragbits/blob/main/examples/agents/tool_use.py).
 
+## Tool choice
+To control what tool is used at first call you could use `tool_choice` parameter. There are the following options:
+- "auto": let model decide if tool call is needed
+- "none": do not call tool
+- "required: enforce tool usage (model decides which one)
+- Callable: one of provided tools
+
+
 ## Conversation history
 [`Agent`][ragbits.agents.Agent]s can retain conversation context across multiple interactions by enabling the `keep_history` flag when initializing the agent. This is useful when you want the agent to understand follow-up questions without needing the user to repeat earlier details.
 

examples/agents/tool_use.py

@@ -82,7 +82,7 @@ async def main() -> None:
         tools=[get_weather],
         default_options=AgentOptions(max_total_tokens=500, max_turns=5),
     )
-    response = await agent.run(WeatherPromptInput(location="Paris"))
+    response = await agent.run(WeatherPromptInput(location="Paris"), tool_choice=get_weather)
     print(response)
 
 

packages/ragbits-agents/CHANGELOG.md

@@ -1,6 +1,7 @@
 # CHANGELOG
 
 ## Unreleased
+- Add tool_choice parameter to agent interface (#738)
 
 ## 1.2.1 (2025-08-04)
 
@@ -13,7 +14,6 @@
 ### Changed
 
 - ragbits-core updated to version v1.2.0
-
 - Add native openai tools support (#621)
 - add Context to Agents (#715)
 

packages/ragbits-agents/src/ragbits/agents/_main.py

@@ -22,7 +22,7 @@
 )
 from ragbits.agents.mcp.server import MCPServer
 from ragbits.agents.mcp.utils import get_tools
-from ragbits.agents.tool import Tool, ToolCallResult
+from ragbits.agents.tool import Tool, ToolCallResult, ToolChoice
 from ragbits.core.audit.traces import trace
 from ragbits.core.llms.base import LLM, LLMClientOptionsT, LLMResponseWithMetadata, ToolCall, Usage
 from ragbits.core.options import Options
@@ -192,6 +192,7 @@ async def run(
         input: str | None = None,
         options: AgentOptions[LLMClientOptionsT] | None = None,
         context: AgentRunContext | None = None,
+        tool_choice: ToolChoice | None = None,
     ) -> AgentResult[PromptOutputT]: ...
 
     @overload
@@ -200,13 +201,15 @@ async def run(
         input: PromptInputT,
         options: AgentOptions[LLMClientOptionsT] | None = None,
         context: AgentRunContext | None = None,
+        tool_choice: ToolChoice | None = None,
     ) -> AgentResult[PromptOutputT]: ...
 
     async def run(
         self,
         input: str | PromptInputT | None = None,
         options: AgentOptions[LLMClientOptionsT] | None = None,
         context: AgentRunContext | None = None,
+        tool_choice: ToolChoice | None = None,
     ) -> AgentResult[PromptOutputT]:
         """
         Run the agent. The method is experimental, inputs and outputs may change in the future.
@@ -218,6 +221,11 @@ async def run(
                 - None: No input. Only valid when a string prompt was provided during initialization.
             options: The options for the agent run.
             context: The context for the agent run.
+            tool_choice: Parameter that allows to control what tool is used at first call. Can be one of:
+                - "auto": let model decide if tool call is needed
+                - "none": do not call tool
+                - "required: enforce tool usage (model decides which one)
+                - Callable: one of provided tools
 
         Returns:
             The result of the agent run.
@@ -251,6 +259,7 @@ async def run(
                     await self.llm.generate_with_metadata(
                         prompt=prompt_with_history,
                         tools=[tool.to_function_schema() for tool in tools_mapping.values()],
+                        tool_choice=tool_choice if tool_choice and turn_count == 0 else None,
                         options=self._get_llm_options(llm_options, merged_options, context.usage),
                     ),
                 )
@@ -294,6 +303,7 @@ def run_streaming(
         input: str | None = None,
         options: AgentOptions[LLMClientOptionsT] | None = None,
         context: AgentRunContext | None = None,
+        tool_choice: ToolChoice | None = None,
     ) -> AgentResultStreaming: ...
 
     @overload
@@ -302,13 +312,15 @@ def run_streaming(
         input: PromptInputT,
         options: AgentOptions[LLMClientOptionsT] | None = None,
         context: AgentRunContext | None = None,
+        tool_choice: ToolChoice | None = None,
     ) -> AgentResultStreaming: ...
 
     def run_streaming(
         self,
         input: str | PromptInputT | None = None,
         options: AgentOptions[LLMClientOptionsT] | None = None,
         context: AgentRunContext | None = None,
+        tool_choice: ToolChoice | None = None,
     ) -> AgentResultStreaming:
         """
         This method returns an `AgentResultStreaming` object that can be asynchronously
@@ -318,6 +330,11 @@ def run_streaming(
             input: The input for the agent run.
             options: The options for the agent run.
             context: The context for the agent run.
+            tool_choice: Parameter that allows to control what tool is used at first call. Can be one of:
+                - "auto": let model decide if tool call is needed
+                - "none": do not call tool
+                - "required: enforce tool usage (model decides which one)
+                - Callable: one of provided tools
 
         Returns:
             A `StreamingResult` object for iteration and collection.
@@ -329,14 +346,15 @@ def run_streaming(
             AgentInvalidPromptInputError: If the prompt/input combination is invalid.
             AgentMaxTurnsExceededError: If the maximum number of turns is exceeded.
         """
-        generator = self._stream_internal(input, options, context)
+        generator = self._stream_internal(input, options, context, tool_choice)
         return AgentResultStreaming(generator)
 
     async def _stream_internal(
         self,
         input: str | PromptInputT | None = None,
         options: AgentOptions[LLMClientOptionsT] | None = None,
         context: AgentRunContext | None = None,
+        tool_choice: ToolChoice | None = None,
     ) -> AsyncGenerator[str | ToolCall | ToolCallResult | SimpleNamespace | BasePrompt | Usage]:
         if context is None:
             context = AgentRunContext()
@@ -357,6 +375,7 @@ async def _stream_internal(
                 streaming_result = self.llm.generate_streaming(
                     prompt=prompt_with_history,
                     tools=[tool.to_function_schema() for tool in tools_mapping.values()],
+                    tool_choice=tool_choice if tool_choice and turn_count == 0 else None,
                     options=self._get_llm_options(llm_options, merged_options, context.usage),
                 )
                 async for chunk in streaming_result:

packages/ragbits-agents/src/ragbits/agents/tool.py

@@ -1,6 +1,6 @@
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Literal
 
 from typing_extensions import Self
 
@@ -76,3 +76,6 @@ def to_function_schema(self) -> dict[str, Any]:
                 "parameters": self.parameters,
             },
         }
+
+
+ToolChoice = Literal["auto", "none", "required"] | Callable

packages/ragbits-agents/tests/unit/test_agent.py

@@ -6,7 +6,7 @@
 from pydantic import BaseModel
 
 from ragbits.agents import Agent, AgentRunContext
-from ragbits.agents._main import AgentOptions, AgentResult, AgentResultStreaming, ToolCallResult
+from ragbits.agents._main import AgentOptions, AgentResult, AgentResultStreaming, ToolCallResult, ToolChoice
 from ragbits.agents.exceptions import (
     AgentInvalidPromptInputError,
     AgentMaxTurnsExceededError,
@@ -143,22 +143,92 @@ def llm_with_tool_call_context() -> MockLLM:
     return MockLLM(default_options=options)
 
 
+def get_time() -> str:
+    """
+    Returns the current time.
+
+    Returns:
+        The current time as a string.
+    """
+    return "12:00 PM"
+
+
+@pytest.fixture
+def llm_no_tool_call_when_none() -> MockLLM:
+    """LLM that doesn't call tools when tool_choice is 'none'."""
+    options = MockLLMOptions(response="I cannot call tools right now.")
+    return MockLLM(default_options=options)
+
+
+@pytest.fixture
+def llm_auto_tool_call() -> MockLLM:
+    """LLM that automatically decides to call a tool."""
+    options = MockLLMOptions(
+        response="Let me check the weather for you.",
+        tool_calls=[
+            {
+                "name": "get_weather",
+                "arguments": '{"location": "New York"}',
+                "id": "auto_test",
+                "type": "function",
+            }
+        ],
+    )
+    return MockLLM(default_options=options)
+
+
+@pytest.fixture
+def llm_required_tool_call() -> MockLLM:
+    """LLM that is forced to call a tool when tool_choice is 'required'."""
+    options = MockLLMOptions(
+        response="",
+        tool_calls=[
+            {
+                "name": "get_weather",
+                "arguments": '{"location": "Boston"}',
+                "id": "required_test",
+                "type": "function",
+            }
+        ],
+    )
+    return MockLLM(default_options=options)
+
+
+@pytest.fixture
+def llm_specific_tool_call() -> MockLLM:
+    """LLM that calls a specific tool when tool_choice is a specific function."""
+    options = MockLLMOptions(
+        response="",
+        tool_calls=[
+            {
+                "name": "get_time",
+                "arguments": "{}",
+                "id": "specific_test",
+                "type": "function",
+            }
+        ],
+    )
+    return MockLLM(default_options=options)
+
+
 async def _run(
     agent: Agent,
     input: str | BaseModel | None = None,
     options: AgentOptions | None = None,
     context: AgentRunContext | None = None,
+    tool_choice: ToolChoice | None = None,
 ) -> AgentResult:
-    return await agent.run(input, options=options, context=context)
+    return await agent.run(input, options=options, context=context, tool_choice=tool_choice)
 
 
 async def _run_streaming(
     agent: Agent,
     input: str | BaseModel | None = None,
     options: AgentOptions | None = None,
     context: AgentRunContext | None = None,
+    tool_choice: ToolChoice | None = None,
 ) -> AgentResultStreaming:
-    result = agent.run_streaming(input, options=options, context=context)
+    result = agent.run_streaming(input, options=options, context=context, tool_choice=tool_choice)
     async for _chunk in result:
         pass
     return result
@@ -588,3 +658,114 @@ async def test_max_turns_not_exeeded_with_many_tool_calls(llm_multiple_tool_call
 
     assert result.content == "Final response after multiple tool calls"
     assert len(result.tool_calls) == 3
+
+
+@pytest.mark.parametrize("method", [_run, _run_streaming])
+async def test_agent_run_with_tool_choice_none(llm_no_tool_call_when_none: MockLLM, method: Callable):
+    """Test agent run with tool_choice set to 'none'."""
+    agent = Agent(
+        llm=llm_no_tool_call_when_none,
+        prompt=CustomPrompt,
+        tools=[get_weather],
+    )
+    result = await method(agent, tool_choice="none")
+
+    assert result.content == "I cannot call tools right now."
+    assert result.tool_calls is None
+
+
+@pytest.mark.parametrize("method", [_run, _run_streaming])
+async def test_agent_run_with_auto_tool_call(llm_auto_tool_call: MockLLM, method: Callable):
+    """Test agent run with automatic tool call."""
+    agent = Agent(
+        llm=llm_auto_tool_call,
+        prompt=CustomPrompt,
+        tools=[get_weather],
+    )
+    result = await method(agent)
+
+    assert result.content == "Let me check the weather for you."
+    assert len(result.tool_calls) == 1
+    assert result.tool_calls[0].id == "auto_test"
+
+
+@pytest.mark.parametrize("method", [_run, _run_streaming])
+async def test_agent_run_with_required_tool_call(llm_required_tool_call: MockLLM, method: Callable):
+    """Test agent run with required tool call."""
+    agent = Agent(
+        llm=llm_required_tool_call,
+        prompt=CustomPrompt,
+        tools=[get_weather],
+    )
+    result = await method(agent, tool_choice="required")
+
+    assert result.content == ""
+    assert len(result.tool_calls) == 1
+    assert result.tool_calls[0].id == "required_test"
+
+
+@pytest.mark.parametrize("method", [_run, _run_streaming])
+async def test_agent_run_with_specific_tool_call(llm_specific_tool_call: MockLLM, method: Callable):
+    """Test agent run with specific tool call."""
+    agent = Agent(
+        llm=llm_specific_tool_call,
+        prompt=CustomPrompt,
+        tools=[get_weather, get_time],
+    )
+    result = await method(agent, tool_choice=get_time)
+
+    assert result.content == ""
+    assert len(result.tool_calls) == 1
+    assert result.tool_calls[0].id == "specific_test"
+    assert result.tool_calls[0].name == "get_time"
+    assert result.tool_calls[0].result == "12:00 PM"
+
+
+@pytest.mark.parametrize("method", [_run, _run_streaming])
+async def test_agent_run_with_tool_choice_auto_explicit(llm_auto_tool_call: MockLLM, method: Callable):
+    """Test agent run with tool_choice explicitly set to 'auto'."""
+    agent = Agent(
+        llm=llm_auto_tool_call,
+        prompt=CustomPrompt,
+        tools=[get_weather],
+    )
+    result = await method(agent, tool_choice="auto")
+
+    assert result.content == "Let me check the weather for you."
+    assert len(result.tool_calls) == 1
+    assert result.tool_calls[0].name == "get_weather"
+    assert result.tool_calls[0].arguments == {"location": "New York"}
+
+
+@pytest.mark.parametrize("method", [_run, _run_streaming])
+async def test_tool_choice_with_multiple_tools_available(llm_auto_tool_call: MockLLM, method: Callable):
+    """Test tool_choice behavior when multiple tools are available."""
+    agent = Agent(
+        llm=llm_auto_tool_call,
+        prompt=CustomPrompt,
+        tools=[get_weather, get_time],  # Multiple tools available
+    )
+
+    result = await method(agent, tool_choice="auto")
+
+    assert result.content == "Let me check the weather for you."
+    assert len(result.tool_calls) == 1
+    # The LLM chose to call get_weather based on its configuration
+    assert result.tool_calls[0].name == "get_weather"
+
+
+@pytest.mark.parametrize("method", [_run, _run_streaming])
+async def test_tool_choice_history_preservation(llm_with_tool_call: MockLLM, method: Callable):
+    """Test that tool_choice works correctly with history preservation."""
+    agent: Agent = Agent(
+        llm=llm_with_tool_call,
+        prompt="You are a helpful assistant",
+        tools=[get_weather],
+        keep_history=True,
+    )
+
+    await method(agent, input="Check weather", tool_choice="auto")
+    assert len(agent.history) >= 3  # At least system, user, assistant messages
+    # Should include tool call in history
+    tool_call_messages = [msg for msg in agent.history if msg.get("role") == "tool"]
+    assert len(tool_call_messages) >= 1

packages/ragbits-core/CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Add tool_choice parameter to LLM interface (#738)
+
 ## 1.2.1 (2025-08-04)
 
 ## 1.2.0 (2025-08-01)