wip

Author: sergei-bronnikov

internal/provider/openai/cost.go

@@ -9,6 +9,7 @@ import (
 	"strings"
 
 	"github.com/bricks-cloud/bricksllm/internal/util"
+	responsesOpenai "github.com/openai/openai-go/responses"
 	goopenai "github.com/sashabaranov/go-openai"
 )
 
@@ -484,6 +485,42 @@ func (ce *CostEstimator) EstimateEmbeddingsCost(r *goopenai.EmbeddingRequest) (f
 	return ce.EstimateEmbeddingsInputCost(string(r.Model), tks)
 }
 
+func (ce *CostEstimator) EstimateResponseApiTotalCost(model string, usage responsesOpenai.ResponseUsage) (float64, error) {
+	if len(model) == 0 {
+		return 0, errors.New("model is not provided")
+	}
+
+	inputTokens := usage.InputTokens
+	cachedInputTokens := usage.InputTokensDetails.CachedTokens
+	outputTokens := usage.OutputTokens
+
+	cachedInputCost, err := ce.estimateResponseApiTokensCost("cached-prompt", model, cachedInputTokens)
+	if err != nil {
+		cachedInputTokens = 0.0
+	}
+	inputCost, err := ce.estimateResponseApiTokensCost("prompt", model, inputTokens-cachedInputTokens)
+	if err != nil {
+		return 0.0, err
+	}
+
+	outputCost, err := ce.estimateResponseApiTokensCost("completion", model, outputTokens)
+
+	return math.Trunc((inputCost+outputCost+cachedInputCost)*100000) / 100000, err
+}
+
+func (ce *CostEstimator) estimateResponseApiTokensCost(costMapKey, model string, tks int64) (float64, error) {
+	costMap, ok := ce.tokenCostMap[costMapKey]
+	if !ok {
+		return 0, errors.New("cost map is not provided")
+	}
+	cost, ok := costMap[model]
+	if !ok {
+		return 0, fmt.Errorf("%s is not present in the cost map provided", model)
+	}
+	tksInFloat := float64(tks)
+	return tksInFloat / 1000 * cost, nil
+}
+
 func countFunctionTokens(model string, r *goopenai.ChatCompletionRequest, tc tokenCounter) (int, error) {
 	if len(r.Functions) == 0 {
 		return 0, nil

internal/server/web/proxy/middleware.go

@@ -4,13 +4,14 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
-	"github.com/bricks-cloud/bricksllm/internal/provider/xcustom"
 	"io"
 	"net/http"
 	"strconv"
 	"strings"
 	"time"
 
+	"github.com/bricks-cloud/bricksllm/internal/provider/xcustom"
+
 	"github.com/bricks-cloud/bricksllm/internal/event"
 	"github.com/bricks-cloud/bricksllm/internal/key"
 	"github.com/bricks-cloud/bricksllm/internal/message"
@@ -27,6 +28,7 @@ import (
 	"github.com/tidwall/sjson"
 	"go.uber.org/zap"
 
+	responsesOpenai "github.com/openai/openai-go/responses"
 	goopenai "github.com/sashabaranov/go-openai"
 )
 
@@ -52,6 +54,7 @@ type estimator interface {
 	EstimateTotalCost(model string, promptTks, completionTks int) (float64, error)
 	EstimateEmbeddingsInputCost(model string, tks int) (float64, error)
 	EstimateChatCompletionPromptTokenCounts(model string, r *goopenai.ChatCompletionRequest) (int, error)
+	EstimateResponseApiTotalCost(model string, usage responsesOpenai.ResponseUsage) (float64, error)
 }
 
 type azureEstimator interface {
@@ -792,6 +795,28 @@ func getMiddleware(cpm CustomProvidersManager, rm routeManager, pm PoliciesManag
 			policyInput = ccr
 		}
 
+		if strings.HasPrefix(c.FullPath(), "/api/providers/openai/v1/responses") {
+			responsesReq := &responsesOpenai.ResponseNewParams{}
+			err = json.Unmarshal(body, responsesReq)
+			if err != nil {
+				logError(logWithCid, "error when unmarshalling openai responses request", prod, err)
+				return
+			}
+
+			userId = responsesReq.User.String()
+			enrichedEvent.Request = responsesReq
+			c.Set("model", responsesReq.Model)
+
+			// TODO: log
+			//logRequest(logWithCid, prod, private, responsesReq)
+
+			if responsesReq.Metadata["stream"] == "true" {
+				c.Set("stream", true)
+			}
+
+			policyInput = responsesReq
+		}
+
 		if c.FullPath() == "/api/providers/openai/v1/embeddings" {
 			er := &goopenai.EmbeddingRequest{}
 			err = json.Unmarshal(body, er)

internal/server/web/proxy/responses.go

@@ -12,6 +12,7 @@ import (
 	"github.com/bricks-cloud/bricksllm/internal/util"
 	"github.com/gin-gonic/gin"
 
+	responsesOpenai "github.com/openai/openai-go/responses"
 	goopenai "github.com/sashabaranov/go-openai"
 )
 
@@ -42,11 +43,11 @@ func getResponsesHandler(prod, private bool, client http.Client, e estimator) gi
 
 		// TODO
 		isStreaming := c.GetBool("stream")
-		//if isStreaming {
-		//	req.Header.Set("Accept", "text/event-stream")
-		//	req.Header.Set("Cache-Control", "no-cache")
-		//	req.Header.Set("Connection", "keep-alive")
-		//}
+		if isStreaming {
+			req.Header.Set("Accept", "text/event-stream")
+			req.Header.Set("Cache-Control", "no-cache")
+			req.Header.Set("Connection", "keep-alive")
+		}
 
 		start := time.Now()
 		res, err := client.Do(req)
@@ -66,8 +67,47 @@ func getResponsesHandler(prod, private bool, client http.Client, e estimator) gi
 			}
 		}
 
+		model := c.GetString("model")
+
 		if res.StatusCode == http.StatusOK && !isStreaming {
+			dur := time.Since(start)
+			telemetry.Timing("bricksllm.proxy.get_responses_handler.latency", dur, nil, 1)
+
+			bytes, err := io.ReadAll(res.Body)
+			if err != nil {
+				logError(log, "error when reading openai http response api response body", prod, err)
+				JSON(c, http.StatusInternalServerError, "[BricksLLM] failed to read openai response body")
+				return
+			}
+
+			var cost float64 = 0
+			resp := &responsesOpenai.Response{}
+			telemetry.Incr("bricksllm.proxy.get_responses_handler.success", nil, 1)
+			telemetry.Timing("bricksllm.proxy.get_responses_handler.success_latency", dur, nil, 1)
+
+			err = json.Unmarshal(bytes, resp)
+			if err != nil {
+				logError(log, "error when unmarshalling openai http response api response body", prod, err)
+			}
 			//	TODO: implement non-streaming logic here
+
+			if err == nil {
+				// TODO log
+				//logChatCompletionResponse(log, prod, private, chatRes)
+				cost, err = e.EstimateResponseApiTotalCost(model, resp.Usage)
+				if err != nil {
+					telemetry.Incr("bricksllm.proxy.get_chat_completion_handler.estimate_total_cost_error", nil, 1)
+					logError(log, "error when estimating openai cost", prod, err)
+				}
+				//m, exists := c.Get("cost_map")
+			}
+
+			c.Set("costInUsd", cost)
+			c.Set("promptTokenCount", resp.Usage.InputTokens)
+			c.Set("completionTokenCount", resp.Usage.OutputTokens)
+
+			c.Data(res.StatusCode, "application/json", bytes)
+			return
 		}
 
 		if res.StatusCode != http.StatusOK {
@@ -94,7 +134,9 @@ func getResponsesHandler(prod, private bool, client http.Client, e estimator) gi
 			return
 		}
 
+		// handle streaming response
+		telemetry.Incr("bricksllm.proxy.get_responses_handler.streaming_requests", nil, 1)
 		// TODO: implement the actual streaming logic here
-
+		telemetry.Timing("bricksllm.proxy.get_chat_completion_handler.streaming_latency", time.Since(start), nil, 1)
 	}
 }