feat: enhacements token details for chat completions and added latency calculation in vertex

Author: Pratham-Mishra04

core/providers/vertex.go

@@ -13,6 +13,7 @@ import (
 	"net/http"
 	"strings"
 	"sync"
+	"time"
 
 	"golang.org/x/oauth2/google"
 
@@ -245,6 +246,8 @@ func (provider *VertexProvider) ChatCompletion(ctx context.Context, key schemas.
 		return nil, newBifrostOperationError("error creating auth client", err, schemas.Vertex)
 	}
 
+	startTime := time.Now()
+
 	// Make request
 	resp, err := client.Do(req)
 	if err != nil {
@@ -267,6 +270,8 @@ func (provider *VertexProvider) ChatCompletion(ctx context.Context, key schemas.
 	}
 	defer resp.Body.Close()
 
+	latency := time.Since(startTime)
+
 	// Handle error response
 	// Read response body
 	body, err := io.ReadAll(resp.Body)
@@ -314,6 +319,7 @@ func (provider *VertexProvider) ChatCompletion(ctx context.Context, key schemas.
 			RequestType:    schemas.ChatCompletionRequest,
 			Provider:       schemas.Vertex,
 			ModelRequested: request.Model,
+			Latency:        latency.Milliseconds(),
 		}
 
 		if provider.sendBackRawResponse {
@@ -322,10 +328,7 @@ func (provider *VertexProvider) ChatCompletion(ctx context.Context, key schemas.
 
 		return response, nil
 	} else {
-		// Pre-allocate response structs from pools
-		// response := acquireOpenAIResponse()
 		response := &schemas.BifrostChatResponse{}
-		// defer releaseOpenAIResponse(response)
 
 		// Use enhanced response handler with pre-allocated response
 		rawResponse, bifrostErr := handleProviderResponse(body, response, provider.sendBackRawResponse)
@@ -336,6 +339,7 @@ func (provider *VertexProvider) ChatCompletion(ctx context.Context, key schemas.
 		response.ExtraFields.RequestType = schemas.ChatCompletionRequest
 		response.ExtraFields.Provider = schemas.Vertex
 		response.ExtraFields.ModelRequested = request.Model
+		response.ExtraFields.Latency = latency.Milliseconds()
 
 		if provider.sendBackRawResponse {
 			response.ExtraFields.RawResponse = rawResponse
@@ -484,22 +488,15 @@ func (provider *VertexProvider) Embedding(ctx context.Context, key schemas.Key,
 		return nil, newConfigurationError("embedding input texts are empty", schemas.Vertex)
 	}
 
-	// All Vertex AI embedding models use the same native Vertex embedding API
-	return provider.handleVertexEmbedding(ctx, request.Model, key, reqBody, request.Params)
-}
-
-// handleVertexEmbedding handles embedding requests using Vertex's native embedding API
-// This is used for all Vertex AI embedding models as they all use the same response format
-func (provider *VertexProvider) handleVertexEmbedding(ctx context.Context, model string, key schemas.Key, vertexReq *vertex.VertexEmbeddingRequest, params *schemas.EmbeddingParameters) (*schemas.BifrostEmbeddingResponse, *schemas.BifrostError) {
 	// Use the typed request directly
-	jsonBody, err := sonic.Marshal(vertexReq)
+	jsonBody, err := sonic.Marshal(reqBody)
 	if err != nil {
 		return nil, newBifrostOperationError(schemas.ErrProviderJSONMarshaling, err, schemas.Vertex)
 	}
 
 	// Build the native Vertex embedding API endpoint
 	url := fmt.Sprintf("https://%s-aiplatform.googleapis.com/v1/projects/%s/locations/%s/publishers/google/models/%s:predict",
-		key.VertexKeyConfig.Region, key.VertexKeyConfig.ProjectID, key.VertexKeyConfig.Region, model)
+		key.VertexKeyConfig.Region, key.VertexKeyConfig.ProjectID, key.VertexKeyConfig.Region, request.Model)
 
 	// Create request
 	req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(jsonBody))
@@ -532,6 +529,8 @@ func (provider *VertexProvider) handleVertexEmbedding(ctx context.Context, model
 		return nil, newBifrostOperationError("error creating auth client", err, schemas.Vertex)
 	}
 
+	startTime := time.Now()
+
 	// Make request
 	resp, err := client.Do(req)
 	if err != nil {
@@ -554,6 +553,8 @@ func (provider *VertexProvider) handleVertexEmbedding(ctx context.Context, model
 	}
 	defer resp.Body.Close()
 
+	latency := time.Since(startTime)
+
 	// Handle error response
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
@@ -598,8 +599,9 @@ func (provider *VertexProvider) handleVertexEmbedding(ctx context.Context, model
 
 	// Set ExtraFields
 	bifrostResponse.ExtraFields.Provider = schemas.Vertex
-	bifrostResponse.ExtraFields.ModelRequested = model
+	bifrostResponse.ExtraFields.ModelRequested = request.Model
 	bifrostResponse.ExtraFields.RequestType = schemas.EmbeddingRequest
+	bifrostResponse.ExtraFields.Latency = latency.Milliseconds()
 
 	// Set raw response if enabled
 	if provider.sendBackRawResponse {

core/schemas/chatcompletions.go

@@ -545,7 +545,21 @@ type ContentLogProb struct {
 
 // BifrostLLMUsage represents token usage information
 type BifrostLLMUsage struct {
-	PromptTokens     int `json:"prompt_tokens,omitempty"`
-	CompletionTokens int `json:"completion_tokens,omitempty"`
-	TotalTokens      int `json:"total_tokens"`
+	PromptTokens            int                          `json:"prompt_tokens,omitempty"`
+	PromptTokensDetails     *ChatPromptTokensDetails     `json:"prompt_tokens_details,omitempty"`
+	CompletionTokens        int                          `json:"completion_tokens,omitempty"`
+	CompletionTokensDetails *ChatCompletionTokensDetails `json:"completion_tokens_details,omitempty"`
+	TotalTokens             int                          `json:"total_tokens"`
+}
+
+type ChatPromptTokensDetails struct {
+	AudioTokens  int `json:"audio_tokens,omitempty"`
+	CachedTokens int `json:"cached_tokens,omitempty"`
+}
+
+type ChatCompletionTokensDetails struct {
+	AcceptedPredictionTokens int `json:"accepted_prediction_tokens,omitempty"`
+	AudioTokens              int `json:"audio_tokens,omitempty"`
+	ReasoningTokens          int `json:"reasoning_tokens,omitempty"`
+	RejectedPredictionTokens int `json:"rejected_prediction_tokens,omitempty"`
 }

core/schemas/mux.go

@@ -615,6 +615,56 @@ func ToChatMessages(rms []ResponsesMessage) []ChatMessage {
 	return chatMessages
 }
 
+func (cu *BifrostLLMUsage) ToResponsesResponseUsage() *ResponsesResponseUsage {
+	if cu == nil {
+		return nil
+	}
+
+	usage := &ResponsesResponseUsage{
+		InputTokens:  cu.PromptTokens,
+		OutputTokens: cu.CompletionTokens,
+		TotalTokens:  cu.TotalTokens,
+	}
+
+	if cu.PromptTokensDetails != nil {
+		usage.InputTokensDetails = &ResponsesResponseInputTokens{
+			CachedTokens: cu.PromptTokensDetails.CachedTokens,
+		}
+	}
+	if cu.CompletionTokensDetails != nil {
+		usage.OutputTokensDetails = &ResponsesResponseOutputTokens{
+			ReasoningTokens: cu.CompletionTokensDetails.ReasoningTokens,
+		}
+	}
+
+	return usage
+}
+
+func (ru *ResponsesResponseUsage) ToBifrostLLMUsage() *BifrostLLMUsage {
+	if ru == nil {
+		return nil
+	}
+
+	usage := &BifrostLLMUsage{
+		PromptTokens:     ru.InputTokens,
+		CompletionTokens: ru.OutputTokens,
+		TotalTokens:      ru.TotalTokens,
+	}
+
+	if ru.InputTokensDetails != nil {
+		usage.PromptTokensDetails = &ChatPromptTokensDetails{
+			CachedTokens: ru.InputTokensDetails.CachedTokens,
+		}
+	}
+	if ru.OutputTokensDetails != nil {
+		usage.CompletionTokensDetails = &ChatCompletionTokensDetails{
+			ReasoningTokens: ru.OutputTokensDetails.ReasoningTokens,
+		}
+	}
+
+	return usage
+}
+
 // =============================================================================
 // REQUEST CONVERSION METHODS
 // =============================================================================
@@ -805,15 +855,7 @@ func (cr *BifrostChatResponse) ToBifrostResponsesResponse() *BifrostResponsesRes
 
 	// Convert Usage if needed
 	if cr.Usage != nil {
-		responsesResp.Usage = &ResponsesResponseUsage{
-			InputTokens:  cr.Usage.PromptTokens,
-			OutputTokens: cr.Usage.CompletionTokens,
-			TotalTokens:  cr.Usage.TotalTokens,
-		}
-
-		if responsesResp.Usage.TotalTokens == 0 {
-			responsesResp.Usage.TotalTokens = cr.Usage.PromptTokens + cr.Usage.CompletionTokens
-		}
+		responsesResp.Usage = cr.Usage.ToResponsesResponseUsage()
 	}
 
 	// Copy other relevant fields
@@ -859,15 +901,7 @@ func (responsesResp *BifrostResponsesResponse) ToBifrostChatResponse() *BifrostC
 	// Convert Usage if needed
 	if responsesResp.Usage != nil {
 		// Map Responses usage to Chat usage
-		chatResp.Usage = &BifrostLLMUsage{
-			PromptTokens:     responsesResp.Usage.InputTokens,
-			CompletionTokens: responsesResp.Usage.OutputTokens,
-			TotalTokens:      responsesResp.Usage.TotalTokens,
-		}
-
-		if chatResp.Usage.TotalTokens == 0 {
-			chatResp.Usage.TotalTokens = chatResp.Usage.PromptTokens + chatResp.Usage.CompletionTokens
-		}
+		chatResp.Usage = responsesResp.Usage.ToBifrostLLMUsage()
 	}
 
 	// Copy other relevant fields
@@ -976,11 +1010,7 @@ func (cr *BifrostChatResponse) ToBifrostResponsesStreamResponse() *BifrostRespon
 			// Add usage information if present in the response
 			if cr.Usage != nil {
 				streamResp.Response = &BifrostResponsesResponse{
-					Usage: &ResponsesResponseUsage{
-						InputTokens:  cr.Usage.PromptTokens,
-						OutputTokens: cr.Usage.CompletionTokens,
-						TotalTokens:  cr.Usage.TotalTokens,
-					},
+					Usage: cr.Usage.ToResponsesResponseUsage(),
 				}
 			}
 		} else {

core/schemas/providers/anthropic/chat.go

@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"time"
+
 	"github.com/maximhq/bifrost/core/schemas"
 )
 
@@ -350,7 +351,10 @@ func (response *AnthropicMessageResponse) ToBifrostChatResponse() *schemas.Bifro
 	// Convert usage information
 	if response.Usage != nil {
 		bifrostResponse.Usage = &schemas.BifrostLLMUsage{
-			PromptTokens:     response.Usage.InputTokens,
+			PromptTokens: response.Usage.InputTokens,
+			PromptTokensDetails: &schemas.ChatPromptTokensDetails{
+				CachedTokens: response.Usage.CacheCreationInputTokens + response.Usage.CacheReadInputTokens,
+			},
 			CompletionTokens: response.Usage.OutputTokens,
 			TotalTokens:      response.Usage.InputTokens + response.Usage.OutputTokens,
 		}
@@ -613,6 +617,11 @@ func ToAnthropicChatCompletionResponse(bifrostResp *schemas.BifrostChatResponse)
 			InputTokens:  bifrostResp.Usage.PromptTokens,
 			OutputTokens: bifrostResp.Usage.CompletionTokens,
 		}
+
+		//NOTE: We cannot segregate between cache creation and cache read tokens, so we will use the total cached tokens as the cache read tokens
+		if bifrostResp.Usage.PromptTokensDetails != nil && bifrostResp.Usage.PromptTokensDetails.CachedTokens > 0 {
+			anthropicResp.Usage.CacheReadInputTokens = bifrostResp.Usage.PromptTokensDetails.CachedTokens
+		}
 	}
 
 	// Convert choices to content

core/schemas/providers/cohere/chat.go

@@ -298,6 +298,11 @@ func (response *CohereChatResponse) ToBifrostChatResponse() *schemas.BifrostChat
 			if response.Usage.Tokens.OutputTokens != nil {
 				usage.CompletionTokens = int(*response.Usage.Tokens.OutputTokens)
 			}
+			if response.Usage.CachedTokens != nil {
+				usage.PromptTokensDetails = &schemas.ChatPromptTokensDetails{
+					CachedTokens: int(*response.Usage.CachedTokens),
+				}
+			}
 			usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
 		}
 

core/schemas/providers/cohere/types.go

@@ -500,14 +500,13 @@ func (c *CohereStreamCitationStruct) UnmarshalJSON(data []byte) error {
 	return fmt.Errorf("citations field is neither array nor object")
 }
 
-
 // CohereStreamMessage represents the message part of streaming deltas
 type CohereStreamMessage struct {
-	Role      *string                    `json:"role,omitempty"`       // For message-start
-	Content   *CohereStreamContentStruct `json:"content,omitempty"`    // For content events (object)
-	ToolPlan  *string                    `json:"tool_plan,omitempty"`  // For tool-plan-delta
-	ToolCalls *CohereStreamToolCallStruct      `json:"tool_calls,omitempty"` // For tool-call events (flexible)
-	Citations *CohereStreamCitationStruct            `json:"citations,omitempty"`  // For citation events
+	Role      *string                     `json:"role,omitempty"`       // For message-start
+	Content   *CohereStreamContentStruct  `json:"content,omitempty"`    // For content events (object)
+	ToolPlan  *string                     `json:"tool_plan,omitempty"`  // For tool-plan-delta
+	ToolCalls *CohereStreamToolCallStruct `json:"tool_calls,omitempty"` // For tool-call events (flexible)
+	Citations *CohereStreamCitationStruct `json:"citations,omitempty"`  // For citation events
 }
 
 // CohereStreamContent represents content in streaming events

core/schemas/providers/gemini/chat.go

@@ -335,7 +335,7 @@ func (response *GenerateContentResponse) ToBifrostChatResponse() *schemas.Bifros
 	}
 
 	// Extract usage metadata
-	inputTokens, outputTokens, totalTokens := response.extractUsageMetadata()
+	inputTokens, outputTokens, totalTokens, cachedTokens, reasoningTokens := response.extractUsageMetadata()
 
 	// Process candidates to extract text content
 	if len(response.Candidates) > 0 {
@@ -380,6 +380,12 @@ func (response *GenerateContentResponse) ToBifrostChatResponse() *schemas.Bifros
 		PromptTokens:     inputTokens,
 		CompletionTokens: outputTokens,
 		TotalTokens:      totalTokens,
+		PromptTokensDetails: &schemas.ChatPromptTokensDetails{
+			CachedTokens: cachedTokens,
+		},
+		CompletionTokensDetails: &schemas.ChatCompletionTokensDetails{
+			ReasoningTokens: reasoningTokens,
+		},
 	}
 
 	return bifrostResp
@@ -469,6 +475,12 @@ func ToGeminiChatResponse(bifrostResp *schemas.BifrostChatResponse) *GenerateCon
 			CandidatesTokenCount: int32(bifrostResp.Usage.CompletionTokens),
 			TotalTokenCount:      int32(bifrostResp.Usage.TotalTokens),
 		}
+		if bifrostResp.Usage.PromptTokensDetails != nil {
+			genaiResp.UsageMetadata.CachedContentTokenCount = int32(bifrostResp.Usage.PromptTokensDetails.CachedTokens)
+		}
+		if bifrostResp.Usage.CompletionTokensDetails != nil {
+			genaiResp.UsageMetadata.ThoughtsTokenCount = int32(bifrostResp.Usage.CompletionTokensDetails.ReasoningTokens)
+		}
 	}
 
 	return genaiResp

core/schemas/providers/gemini/transcription.go

@@ -77,7 +77,7 @@ func (response *GenerateContentResponse) ToBifrostTranscriptionResponse() *schem
 	bifrostResp := &schemas.BifrostTranscriptionResponse{}
 
 	// Extract usage metadata
-	inputTokens, outputTokens, totalTokens := response.extractUsageMetadata()
+	inputTokens, outputTokens, totalTokens, _, _ := response.extractUsageMetadata()
 
 	// Process candidates to extract text content
 	if len(response.Candidates) > 0 {

core/schemas/providers/gemini/utils.go

@@ -153,14 +153,16 @@ func ensureExtraParams(bifrostReq *schemas.BifrostChatRequest) {
 }
 
 // extractUsageMetadata extracts usage metadata from the Gemini response
-func (r *GenerateContentResponse) extractUsageMetadata() (int, int, int) {
-	var inputTokens, outputTokens, totalTokens int
+func (r *GenerateContentResponse) extractUsageMetadata() (int, int, int, int, int) {
+	var inputTokens, outputTokens, totalTokens, cachedTokens, reasoningTokens int
 	if r.UsageMetadata != nil {
 		inputTokens = int(r.UsageMetadata.PromptTokenCount)
 		outputTokens = int(r.UsageMetadata.CandidatesTokenCount)
 		totalTokens = int(r.UsageMetadata.TotalTokenCount)
+		cachedTokens = int(r.UsageMetadata.CachedContentTokenCount)
+		reasoningTokens = int(r.UsageMetadata.ThoughtsTokenCount)
 	}
-	return inputTokens, outputTokens, totalTokens
+	return inputTokens, outputTokens, totalTokens, cachedTokens, reasoningTokens
 }
 
 // convertParamsToGenerationConfig converts Bifrost parameters to Gemini GenerationConfig

framework/changelog.md

@@ -1,4 +1,5 @@
 <!-- The pattern we follow here is to keep the changelog for the latest version -->
 <!-- Old changelogs are automatically attached to the GitHub releases -->
 
-- chore: version update core to 1.2.12
+- chore: version update core to 1.2.12
+- feat: added support for vertex provider/model format in pricing lookup