tools cost

Author: sergei-bronnikov

internal/provider/openai/cost.go

@@ -189,6 +189,17 @@ var OpenAiPerThousandTokenCost = map[string]map[string]float64{
 	},
 }
 
+var OpenAiPerThousandCallsToolCost = map[string]float64{
+	"web_search":                   10.0,
+	"web_search_preview":           25.0,
+	"web_search_preview_reasoning": 10.0,
+}
+
+var AllowedTools = []string{
+	"web_search",
+	"web_search_preview",
+}
+
 type tokenCounter interface {
 	Count(model string, input string) (int, error)
 }
@@ -527,6 +538,34 @@ func (ce *CostEstimator) EstimateResponseApiTotalCost(model string, usage respon
 	return inputCost + cachedInputCost + outputCost, err
 }
 
+func (ce *CostEstimator) EstimateResponseApiToolCallsCost(tools []responsesOpenai.ToolUnion, model string) (float64, error) {
+	if len(tools) == 0 {
+		return 0, nil
+	}
+	totalCost := 0.0
+	for _, tool := range tools {
+		toolType := tool.Type
+		cost, ok := OpenAiPerThousandCallsToolCost[extendedToolType(toolType, model)]
+		if !ok {
+			return 0, fmt.Errorf("tool type %s is not present in the tool cost map provided", toolType)
+		}
+		totalCost += cost
+	}
+	return totalCost / 1000, nil
+}
+
+var reasoningModelPrefix = []string{"gpt-5", "o1", "o2", "o3"}
+
+func extendedToolType(toolType, model string) string {
+	if toolType != "web_search_preview" {
+		return toolType
+	}
+	if slices.ContainsFunc(reasoningModelPrefix, func(s string) bool { return strings.HasPrefix(model, s) }) {
+		return "web_search_preview_reasoning"
+	}
+	return toolType
+}
+
 func (ce *CostEstimator) estimateResponseApiTokensCost(costMapKey, model string, tks int64) (float64, error) {
 	costMap, ok := ce.tokenCostMap[costMapKey]
 	if !ok {

internal/provider/openai/types.go

@@ -1,31 +1,35 @@
 package openai
 
 type ResponseRequest struct {
-	Background         *bool              `json:"background,omitzero"`
-	Conversation       *any               `json:"conversation,omitzero"`
-	Include            []string           `json:"include,omitzero"`
-	Input              *any               `json:"input,omitzero"`
-	Instructions       *string            `json:"instructions,omitzero"`
-	MaxOutputTokens    *int               `json:"max_output_tokens,omitzero"`
-	MaxToolCalls       *int               `json:"max_tool_calls,omitzero"`
-	Metadata           *map[string]string `json:"metadata,omitzero"`
-	Model              *string            `json:"model,omitzero"`
-	ParallelToolCalls  *bool              `json:"parallel_tool_calls,omitzero"`
-	PreviousResponseId *string            `json:"previous_response_id,omitzero"`
-	Prompt             *any               `json:"prompt,omitzero"`
-	PromptCacheKey     *string            `json:"prompt_cache_key,omitzero"`
-	Reasoning          *any               `json:"reasoning,omitzero"`
-	SafetyIdentifier   *string            `json:"safety_identifier,omitzero"`
-	ServiceTier        *string            `json:"service_tier,omitzero"`
-	Store              *bool              `json:"store,omitzero"`
-	Stream             *bool              `json:"stream,omitzero"`
-	StreamOptions      *any               `json:"stream_options,omitzero"`
-	Temperature        *float32           `json:"temperature,omitzero"`
-	Text               *any               `json:"text,omitzero"`
-	ToolChoice         *any               `json:"tool_choice,omitzero"`
-	Tools              []any              `json:"tools,omitzero"`
-	TopLogprobs        *int               `json:"top_logprobs,omitzero"`
-	TopP               *float32           `json:"top_p,omitzero"`
-	Truncation         *string            `json:"truncation,omitzero"`
+	Background         *bool                      `json:"background,omitzero"`
+	Conversation       *any                       `json:"conversation,omitzero"`
+	Include            []string                   `json:"include,omitzero"`
+	Input              *any                       `json:"input,omitzero"`
+	Instructions       *string                    `json:"instructions,omitzero"`
+	MaxOutputTokens    *int                       `json:"max_output_tokens,omitzero"`
+	MaxToolCalls       *int                       `json:"max_tool_calls,omitzero"`
+	Metadata           *map[string]string         `json:"metadata,omitzero"`
+	Model              *string                    `json:"model,omitzero"`
+	ParallelToolCalls  *bool                      `json:"parallel_tool_calls,omitzero"`
+	PreviousResponseId *string                    `json:"previous_response_id,omitzero"`
+	Prompt             *any                       `json:"prompt,omitzero"`
+	PromptCacheKey     *string                    `json:"prompt_cache_key,omitzero"`
+	Reasoning          *any                       `json:"reasoning,omitzero"`
+	SafetyIdentifier   *string                    `json:"safety_identifier,omitzero"`
+	ServiceTier        *string                    `json:"service_tier,omitzero"`
+	Store              *bool                      `json:"store,omitzero"`
+	Stream             *bool                      `json:"stream,omitzero"`
+	StreamOptions      *any                       `json:"stream_options,omitzero"`
+	Temperature        *float32                   `json:"temperature,omitzero"`
+	Text               *any                       `json:"text,omitzero"`
+	ToolChoice         *any                       `json:"tool_choice,omitzero"`
+	Tools              []ResponseRequestToolUnion `json:"tools,omitzero"`
+	TopLogprobs        *int                       `json:"top_logprobs,omitzero"`
+	TopP               *float32                   `json:"top_p,omitzero"`
+	Truncation         *string                    `json:"truncation,omitzero"`
 	//User            *string  `json:"user,omitzero"` //Deprecated
 }
+
+type ResponseRequestToolUnion struct {
+	Type string `json:"type"`
+}

internal/server/web/proxy/middleware.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"slices"
 	"strconv"
 	"strings"
 	"time"
@@ -57,6 +58,7 @@ type estimator interface {
 	EstimateEmbeddingsInputCost(model string, tks int) (float64, error)
 	EstimateChatCompletionPromptTokenCounts(model string, r *goopenai.ChatCompletionRequest) (int, error)
 	EstimateResponseApiTotalCost(model string, usage responsesOpenai.ResponseUsage) (float64, error)
+	EstimateResponseApiToolCallsCost(tools []responsesOpenai.ToolUnion, model string) (float64, error)
 }
 
 type azureEstimator interface {
@@ -812,6 +814,21 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
 				return
 			}
 
+			hasNotAllowedTools := false
+			for _, tool := range responsesReq.Tools {
+				if !slices.Contains(openai.AllowedTools, tool.Type) {
+					hasNotAllowedTools = true
+					break
+				}
+			}
+
+			if hasNotAllowedTools {
+				telemetry.Incr("bricksllm.proxy.get_middleware.tool_not_allowed", nil, 1)
+				JSON(c, http.StatusForbidden, "[BricksLLM] one of the tools is not allowed")
+				c.Abort()
+				return
+			}
+
 			userId = gopointer.ToValueOrDefault(responsesReq.SafetyIdentifier, "")
 			enrichedEvent.Request = responsesReq
 			c.Set("model", gopointer.ToValueOrDefault(responsesReq.Model, ""))

internal/server/web/proxy/responses.go

@@ -98,6 +98,12 @@ func getResponsesHandler(prod, private bool, client http.Client, e estimator) gi
 					telemetry.Incr("bricksllm.proxy.get_chat_completion_handler.estimate_total_cost_error", nil, 1)
 					logError(log, "error when estimating openai cost", prod, err)
 				}
+				toolsCost, err := e.EstimateResponseApiToolCallsCost(resp.Tools, model)
+				if err != nil {
+					telemetry.Incr("bricksllm.proxy.get_chat_completion_handler.estimate_tool_calls_cost_error", nil, 1)
+					logError(log, "error when estimating openai tool calls cost", prod, err)
+				}
+				cost += toolsCost
 			}
 
 			c.Set("costInUsd", cost)
@@ -231,6 +237,12 @@ func getResponsesHandler(prod, private bool, client http.Client, e estimator) gi
 						telemetry.Incr("bricksllm.proxy.get_chat_completion_handler.estimate_total_cost_error", nil, 1)
 						logError(log, "error when estimating openai cost", prod, err)
 					}
+					toolsCost, err := e.EstimateResponseApiToolCallsCost(responsesStreamResp.Response.Tools, model)
+					if err != nil {
+						telemetry.Incr("bricksllm.proxy.get_chat_completion_handler.estimate_tool_calls_cost_error", nil, 1)
+						logError(log, "error when estimating openai tool calls cost", prod, err)
+					}
+					streamCost += toolsCost
 					streamPromptTokenCount, err = int64ToInt(responsesStreamResp.Response.Usage.InputTokens)
 					if err != nil {
 						telemetry.Incr("bricksllm.proxy.get_responses_handler.int64_to_int_error", nil, 1)