Configure the tool_choice to use a specific tool

Signed-off-by: MondayCha <[email protected]>

Author: MondayCha

pkg/common/tools_utils.go

@@ -21,6 +21,8 @@ import (
 	"fmt"
 
 	openaiserverapi "github.com/llm-d/llm-d-inference-sim/pkg/openai-server-api"
+	"github.com/openai/openai-go/v3"
+	"github.com/openai/openai-go/v3/packages/param"
 	"github.com/santhosh-tekuri/jsonschema/v5"
 )
 
@@ -52,24 +54,88 @@ var fakeStringArguments = []string{
 	`lifetime`,
 }
 
-// CreateToolCalls creates and returns response payload based on this request
-// (tool calls or nothing in case we randomly choose not to generate calls),
-// and the number of generated completion token sand the finish reason
-func CreateToolCalls(tools []openaiserverapi.Tool, toolChoice string, config *Configuration) ([]openaiserverapi.ToolCall, int, error) {
-	// This function is called if tool choice is either 'required' or 'auto'.
-	// In case of 'required' at least one tool call has to be created, and we randomly choose
-	// the number of calls starting from one. Otherwise, we start from 0, and in case we randomly
-	// choose the number of calls to be 0, response text will be generated instead of a tool call.
+// IsToolChoiceNone checks if the tool_choice is set to "none".
+func IsToolChoiceNone(toolChoice openai.ChatCompletionToolChoiceOptionUnionParam) bool {
+	if !param.IsOmitted(toolChoice.OfAuto) {
+		val := toolChoice.OfAuto.Or("")
+		return val == ToolChoiceNone
+	}
+	return false
+}
+
+// CreateToolCalls creates and returns tool calls based on the request's tool
+// definitions and tool_choice parameter.
+// The tool_choice parameter controls how the model responds to function calls:
+//   - "none": The model does not call any tools. This case should be handled
+//     before calling this function.
+//   - "auto": The model can choose to either generate a message or call one or
+//     more tools. This is the default behavior.
+//   - "required": The model must call one or more tools.
+//   - Specific function: A specific tool can be forced by providing an object
+//     like `{"type": "function", "function": {"name": "my_function"}}`. The
+//     model will be constrained to call that exact tool.
+//
+// This function returns the generated tool calls, the number of completion
+// tokens used, and an error if one occurs (e.g., if a specified tool is not found).
+func CreateToolCalls(
+	tools []openaiserverapi.Tool,
+	toolChoice openai.ChatCompletionToolChoiceOptionUnionParam,
+	config *Configuration,
+) ([]openaiserverapi.ToolCall, int, error) {
+	// If a specific function is required.
+	if functionChoice := toolChoice.GetFunction(); functionChoice != nil {
+		requiredFuncName := functionChoice.Name
+		var targetTool *openaiserverapi.Tool
+
+		// Find the specified tool in the list of available tools.
+		for i, tool := range tools {
+			if tool.Function.Name == requiredFuncName {
+				targetTool = &tools[i]
+				break
+			}
+		}
+
+		if targetTool == nil {
+			return nil, 0, fmt.Errorf("tool with name '%s' requested in tool_choice but not found in the tools list", requiredFuncName)
+		}
+
+		// Generate arguments for the specific tool.
+		args, err := generateToolArguments(*targetTool, config)
+		if err != nil {
+			return nil, 0, err
+		}
+		argsJson, err := json.Marshal(args)
+		if err != nil {
+			return nil, 0, err
+		}
+
+		call := openaiserverapi.ToolCall{
+			Function: openaiserverapi.FunctionCall{
+				Arguments:          string(argsJson),
+				TokenizedArguments: Tokenize(string(argsJson)),
+				Name:               &targetTool.Function.Name,
+			},
+			ID:    "chatcmpl-tool-" + RandomNumericString(10),
+			Type:  "function",
+			Index: 0,
+		}
+		calls := []openaiserverapi.ToolCall{call}
+		return calls, CountTokensForToolCalls(calls), nil
+	}
+
+	// Tool choice is 'auto' or 'required'.
+	// In 'required' mode, at least one tool call must be created.
+	// In 'auto' mode, the number of calls can be zero, leading to a text response instead.
 	min := 0
-	if toolChoice == ToolChoiceRequired {
+	if !param.IsOmitted(toolChoice.OfAuto) && toolChoice.OfAuto.Or("") == ToolChoiceRequired {
 		min = 1
 	}
 	numberOfCalls := RandomInt(min, len(tools))
 	if numberOfCalls == 0 {
 		return nil, 0, nil
 	}
 
-	calls := make([]openaiserverapi.ToolCall, 0)
+	calls := make([]openaiserverapi.ToolCall, 0, numberOfCalls)
 	for i := range numberOfCalls {
 		// Randomly choose which tools to call. We may call the same tool more than once.
 		index := RandomInt(0, len(tools)-1)

pkg/llm-d-inference-sim/worker.go

@@ -88,7 +88,7 @@ func (s *VllmSimulator) processRequest(reqCtx *openaiserverapi.CompletionReqCtx)
 	var toolCalls []openaiserverapi.ToolCall
 	var completionTokens int
 	if reqCtx.IsChatCompletion &&
-		req.GetToolChoice() != common.ToolChoiceNone &&
+		!common.IsToolChoiceNone(req.GetToolChoice()) &&
 		req.GetTools() != nil {
 		toolCalls, completionTokens, err =
 			common.CreateToolCalls(req.GetTools(), req.GetToolChoice(), s.config)

pkg/openai-server-api/request.go

@@ -20,6 +20,7 @@ package openaiserverapi
 import (
 	"sync"
 
+	"github.com/openai/openai-go/v3"
 	"github.com/valyala/fasthttp"
 )
 
@@ -46,10 +47,10 @@ type CompletionRequest interface {
 	SetNumberOfCachedPromptTokens(cachedPromptTokens int)
 	// GetPrompt returns the prompt
 	GetPrompt() string
-	// GetTools() returns tools to use (in chat completion)
+	// GetTools returns tools to use (in chat completion)
 	GetTools() []Tool
-	// GetToolChoice() returns tool choice (in chat completion)
-	GetToolChoice() string
+	// GetToolChoice returns tool choice (in chat completion)
+	GetToolChoice() openai.ChatCompletionToolChoiceOptionUnionParam
 	// GetMaxCompletionTokens returns the maximum completion tokens requested
 	GetMaxCompletionTokens() *int64
 	// GetIgnoreEOS returns true if the end-of-sequence tokens will be ignored
@@ -184,12 +185,13 @@ type ChatCompletionRequest struct {
 	// Tools is a list of tools the model may call.
 	Tools []Tool `json:"tools,omitempty"`
 
-	// ToolChoice controls which (if any) tool is called by the model,
-	// possible values: none, auto, required.
-	// Sending an object with a specific tool, is currently not supported.
-	ToolChoice string `json:"tool_choice,omitempty"`
+	// ToolChoice controls which (if any) tool is called by the model.
+	// It can be a string ("none", "auto", "required") or an object specifying the function.
+	ToolChoice openai.ChatCompletionToolChoiceOptionUnionParam `json:"tool_choice,omitzero"`
 }
 
+var _ CompletionRequest = (*ChatCompletionRequest)(nil)
+
 // function defines a tool
 type function struct {
 	// Name is the function's name
@@ -221,7 +223,7 @@ func (c *ChatCompletionRequest) GetTools() []Tool {
 	return c.Tools
 }
 
-func (c *ChatCompletionRequest) GetToolChoice() string {
+func (c *ChatCompletionRequest) GetToolChoice() openai.ChatCompletionToolChoiceOptionUnionParam {
 	return c.ToolChoice
 }
 
@@ -286,6 +288,8 @@ type TextCompletionRequest struct {
 	MaxTokens *int64 `json:"max_tokens"`
 }
 
+var _ CompletionRequest = (*TextCompletionRequest)(nil)
+
 func (t *TextCompletionRequest) GetPrompt() string {
 	return t.Prompt
 }
@@ -294,8 +298,8 @@ func (c *TextCompletionRequest) GetTools() []Tool {
 	return nil
 }
 
-func (c *TextCompletionRequest) GetToolChoice() string {
-	return ""
+func (c *TextCompletionRequest) GetToolChoice() openai.ChatCompletionToolChoiceOptionUnionParam {
+	return openai.ChatCompletionToolChoiceOptionUnionParam{}
 }
 
 func (c *TextCompletionRequest) GetMaxCompletionTokens() *int64 {