fixes costing and token counts for streaming requests

core/bifrost.go

@@ -1135,7 +1135,18 @@ func (bifrost *Bifrost) tryStreamRequest(req *schemas.BifrostRequest, ctx contex
 		bifrost.releaseChannelMessage(msg)
 		return stream, nil
 	case bifrostErrVal := <-msg.Err:
+		bifrost.logger.Warn("error while executing stream request: %v", bifrostErrVal.Error.Message)
+		// Marking final chunk
+		ctx = context.WithValue(ctx, schemas.BifrostContextKeyStreamEndIndicator, true)
+		// On error we will complete post-hooks
+		recoveredResp, recoveredErr := pipeline.RunPostHooks(&ctx, nil, &bifrostErrVal, len(bifrost.plugins))
 		bifrost.releaseChannelMessage(msg)
+		if recoveredErr != nil {
+			return nil, recoveredErr
+		}
+		if recoveredResp != nil {
+			return newBifrostMessageChan(recoveredResp), nil
+		}
 		return nil, &bifrostErrVal
 	}
 }

core/changelog.md

@@ -1,2 +1,4 @@
 <!-- The pattern we follow here is to keep the changelog for the latest version -->
-<!-- Old changelogs are automatically attached to the GitHub releases -->
+<!-- Old changelogs are automatically attached to the GitHub releases -->
+
+- Fix: Updates token calculation for streaming responses. #520

core/providers/openai.go

@@ -499,9 +499,24 @@ func handleOpenAIStreaming(
 			}
 
 			// Handle usage-only chunks (when stream_options include_usage is true)
-			if len(response.Choices) == 0 && response.Usage != nil {
+			if response.Usage != nil {
 				// Collect usage information and send at the end of the stream
-				usage = response.Usage
+				// Here in some cases usage comes before final message
+				// So we need to check if the response.Usage is nil and then if usage != nil
+				// then add up all tokens
+				if response.Usage.PromptTokens > usage.PromptTokens {
+					usage.PromptTokens = response.Usage.PromptTokens
+				}
+				if response.Usage.CompletionTokens > usage.CompletionTokens {
+					usage.CompletionTokens = response.Usage.CompletionTokens
+				}
+				if response.Usage.TotalTokens > usage.TotalTokens {
+					usage.TotalTokens = response.Usage.TotalTokens
+				}
+				calculatedTotal := usage.PromptTokens + usage.CompletionTokens
+				if calculatedTotal > usage.TotalTokens {
+					usage.TotalTokens = calculatedTotal
+				}
 				response.Usage = nil
 			}
 

core/version

@@ -1 +1 @@
-1.1.37
+1.1.38

framework/changelog.md

@@ -1,4 +1,4 @@
 <!-- The pattern we follow here is to keep the changelog for the latest version -->
 <!-- Old changelogs are automatically attached to the GitHub releases -->
 
-- Pricing module now accommodates nested model names i.e. groq/openai/gpt-oss-20b was getting skipped while computing costs.
+- upgrade: core upgrades to 1.1.38

framework/logstore/sqlite.go

@@ -179,7 +179,7 @@ func (s *SQLiteLogStore) SearchLogs(filters SearchFilters, pagination Pagination
 		}
 		return nil, err
 	}
-
+	
 	return &SearchResult{
 		Logs:       logs,
 		Pagination: pagination,

framework/logstore/tables.go

@@ -233,6 +233,7 @@ func (l *Log) SerializeFields() error {
 	}
 
 	if l.ErrorDetailsParsed != nil {
+		l.ErrorDetailsParsed.Error.Error = nil
 		if data, err := json.Marshal(l.ErrorDetailsParsed); err != nil {
 			return err
 		} else {

framework/version

@@ -1 +1 @@
-1.0.23
+1.0.24

plugins/governance/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 -->
 
-- Upgrades framework to 1.0.23
+- upgrade: core to 1.1.38
+- upgrade: framework to 1.0.24

plugins/governance/version

@@ -1 +1 @@
-1.2.16
+1.2.17

plugins/jsonparser/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 -->
 
-- Upgrades framework to 1.0.23
+- upgrade: core to 1.1.38
+- upgrade: framework to 1.0.24

plugins/jsonparser/version

@@ -1 +1 @@
-1.2.15
+1.2.16

plugins/logging/changelog.md

@@ -1,5 +1,6 @@
 <!-- The pattern we follow here is to keep the changelog for the latest version -->
 <!-- Old changelogs are automatically attached to the GitHub releases -->
 
-- Upgrades framework to 1.0.23
-- Fixes pricing computation for nested model names.
+- fix: fixes error logging for streaming and non-streaming responses.
+- upgrade: core to 1.1.38
+- upgrade: framework to 1.0.24

plugins/logging/main.go

@@ -371,7 +371,7 @@ func (p *LoggerPlugin) PostHook(ctx *context.Context, result *schemas.BifrostRes
 	if !ok {
 		p.logger.Error("model not found in context")
 		return result, err, nil
-	}	
+	}
 	// Check if this is a streaming response
 	requestType, ok := (*ctx).Value(schemas.BifrostContextKeyRequestType).(schemas.RequestType)
 	if !ok {

plugins/logging/operations.go

@@ -37,7 +37,7 @@ func (p *LoggerPlugin) updateLogEntry(ctx context.Context, requestID string, tim
 	updates := make(map[string]interface{})
 	if !timestamp.IsZero() {
 		// Try to get original timestamp from context first for latency calculation
-		latency, err := p.calculateLatency(requestID, timestamp, ctx)
+		latency, err := p.calculateLatency(ctx, requestID, timestamp)
 		if err != nil {
 			return err
 		}
@@ -142,7 +142,7 @@ func (p *LoggerPlugin) processStreamUpdate(ctx context.Context, requestID string
 
 	// Handle error case first
 	if data.ErrorDetails != nil {
-		latency, err := p.calculateLatency(requestID, timestamp, ctx)
+		latency, err := p.calculateLatency(ctx, requestID, timestamp)
 		if err != nil {
 			// If we can't get created_at, just update status and error
 			tempEntry := &logstore.Log{}
@@ -182,7 +182,7 @@ func (p *LoggerPlugin) processStreamUpdate(ctx context.Context, requestID string
 	if isFinalChunk {
 		// Stream is finishing, calculate latency
 		var err error
-		latency, err = p.calculateLatency(requestID, timestamp, ctx)
+		latency, err = p.calculateLatency(ctx, requestID, timestamp)
 		if err != nil {
 			return fmt.Errorf("failed to get created_at for latency calculation: %w", err)
 		}
@@ -267,14 +267,17 @@ func (p *LoggerPlugin) processStreamUpdate(ctx context.Context, requestID string
 }
 
 // calculateLatency computes latency in milliseconds from creation time
-func (p *LoggerPlugin) calculateLatency(requestID string, currentTime time.Time, ctx context.Context) (float64, error) {
+func (p *LoggerPlugin) calculateLatency(ctx context.Context, requestID string, currentTime time.Time) (float64, error) {
 	// Try to get original timestamp from context first
 	if ctxTimestamp, ok := ctx.Value(CreatedTimestampKey).(time.Time); ok {
 		return float64(currentTime.Sub(ctxTimestamp).Nanoseconds()) / 1e6, nil
 	}
-
-	// Fallback to database query if not found in context
-	originalEntry, err := p.store.FindFirst(map[string]interface{}{"id": requestID}, "created_at")
+	var originalEntry *logstore.Log
+	err := retryOnNotFound(ctx, func() error {
+		var opErr error
+		originalEntry, opErr = p.store.FindFirst(map[string]interface{}{"id": requestID}, "created_at")
+		return opErr
+	})
 	if err != nil {
 		return 0, err
 	}

plugins/logging/streaming.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/bytedance/sonic"
 	bifrost "github.com/maximhq/bifrost/core"
 	"github.com/maximhq/bifrost/core/schemas"
 	"github.com/maximhq/bifrost/framework/logstore"
@@ -125,7 +126,7 @@ func (p *LoggerPlugin) addStreamChunk(requestID string, chunk *StreamChunk, obje
 }
 
 // processAccumulatedChunks processes all accumulated chunks in order
-func (p *LoggerPlugin) processAccumulatedChunks(requestID string) error {
+func (p *LoggerPlugin) processAccumulatedChunks(ctx context.Context, requestID string, respErr *schemas.BifrostError) error {
 	accumulator := p.getOrCreateStreamAccumulator(requestID)
 
 	accumulator.mu.Lock()
@@ -138,7 +139,7 @@ func (p *LoggerPlugin) processAccumulatedChunks(requestID string) error {
 	completeMessage := p.buildCompleteMessageFromChunks(accumulator.Chunks)
 
 	// Calculate final latency
-	latency, err := p.calculateLatency(requestID, accumulator.FinalTimestamp, context.Background())
+	latency, err := p.calculateLatency(ctx, requestID, accumulator.FinalTimestamp)
 	if err != nil {
 		p.logger.Error("failed to calculate latency for request %s: %v", requestID, err)
 		latency = 0
@@ -147,6 +148,9 @@ func (p *LoggerPlugin) processAccumulatedChunks(requestID string) error {
 	// Update database with complete message
 	updates := make(map[string]interface{})
 	updates["status"] = "success"
+	if respErr != nil {
+		updates["status"] = "error"
+	}
 	updates["stream"] = true
 	updates["latency"] = latency
 	updates["timestamp"] = accumulator.FinalTimestamp
@@ -158,13 +162,19 @@ func (p *LoggerPlugin) processAccumulatedChunks(requestID string) error {
 	if completeMessage.AssistantMessage != nil && completeMessage.AssistantMessage.ToolCalls != nil {
 		tempEntry.ToolCallsParsed = completeMessage.AssistantMessage.ToolCalls
 	}
-
 	if err := tempEntry.SerializeFields(); err != nil {
 		return fmt.Errorf("failed to serialize complete message: %w", err)
 	}
-
-	updates["output_message"] = tempEntry.OutputMessage
-	updates["content_summary"] = tempEntry.ContentSummary
+	if respErr != nil {
+		if b, mErr := sonic.Marshal(respErr); mErr == nil {
+			updates["error_details"] = string(b)
+		} else {
+			updates["error_details"] = fmt.Sprintf(`{"message":"failed to marshal error: %v"}`, mErr)
+		}
+	} else {
+		updates["output_message"] = tempEntry.OutputMessage
+		updates["content_summary"] = tempEntry.ContentSummary
+	}
 	if tempEntry.ToolCalls != "" {
 		updates["tool_calls"] = tempEntry.ToolCalls
 	}
@@ -402,7 +412,7 @@ func (p *LoggerPlugin) handleStreamingResponse(ctx *context.Context, result *sch
 
 			if shouldProcess {
 
-				if processErr := p.processAccumulatedChunks(requestID); processErr != nil {
+				if processErr := p.processAccumulatedChunks(*ctx, requestID, err); processErr != nil {
 					p.logger.Error("failed to process accumulated chunks for request %s: %v", requestID, processErr)
 				}
 

plugins/maxim/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 -->
 
-- Upgrades framework to 1.0.23
+- upgrade: core to 1.1.38
+- upgrade: framework to 1.0.24

plugins/maxim/version

@@ -1 +1 @@
-1.3.6
+1.3.7

plugins/mocker/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 -->
 
-- Upgrades framework to 1.0.23
+- upgrade: core to 1.1.38
+- upgrade: framework to 1.0.24

plugins/mocker/version

@@ -1 +1 @@
-1.2.14
+1.2.15

plugins/semanticcache/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 -->
 
-- Upgrades framework to 1.0.23
+- upgrade: core to 1.1.38
+- upgrade: framework to 1.0.24

plugins/semanticcache/version

@@ -1 +1 @@
-1.2.18
+1.2.19

plugins/telemetry/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 -->
 
-- Upgrades framework to 1.0.23
+- upgrade: core to 1.1.38
+- upgrade: framework to 1.0.24

plugins/telemetry/version

@@ -1 +1 @@
-1.2.15
+1.2.16

tests/core-providers/openrouter_test.go

@@ -1,7 +1,6 @@
 package tests
 
 import (
-	"os"
 	"testing"
 
 	"github.com/maximhq/bifrost/tests/core-providers/config"
@@ -10,9 +9,6 @@ import (
 )
 
 func TestOpenRouter(t *testing.T) {
-	if os.Getenv("OPENROUTER_API_KEY") == "" {
-		t.Skip("OPENROUTER_API_KEY not set; skipping OpenRouter tests")
-	}
 	client, ctx, cancel, err := config.SetupTest()
 	if err != nil {
 		t.Fatalf("Error initializing test setup: %v", err)
@@ -44,3 +40,71 @@ func TestOpenRouter(t *testing.T) {
 
 	runAllComprehensiveTests(t, client, ctx, testConfig)
 }
+
+// TestOpenRouterAnthropic tests Anthropic models via OpenRouter
+func TestOpenRouterAnthropic(t *testing.T) {
+	client, ctx, cancel, err := config.SetupTest()
+	if err != nil {
+		t.Fatalf("Error initializing test setup: %v", err)
+	}
+	defer cancel()
+	defer client.Shutdown()
+
+	testConfig := config.ComprehensiveTestConfig{
+		Provider:  schemas.OpenRouter,
+		ChatModel: "anthropic/claude-3.5-sonnet", // Using Claude 3.5 Sonnet via OpenRouter
+		TextModel: "", // Anthropic models don't support text completion
+		EmbeddingModel: "",
+		Scenarios: config.TestScenarios{
+			TextCompletion:        false, // Not supported by Anthropic
+			SimpleChat:            true,
+			ChatCompletionStream:  true,
+			MultiTurnConversation: true,
+			ToolCalls:             true,
+			MultipleToolCalls:     true,
+			End2EndToolCalling:    true,
+			AutomaticFunctionCall: true,
+			ImageURL:              true,
+			ImageBase64:           true,
+			MultipleImages:        true,
+			CompleteEnd2End:       true,
+			ProviderSpecific:      false, // Skip provider-specific tests
+		},
+	}
+
+	runAllComprehensiveTests(t, client, ctx, testConfig)
+}
+
+// TestOpenRouterMetaLlama tests Meta's Llama models via OpenRouter
+func TestOpenRouterMetaLlama(t *testing.T) {
+	client, ctx, cancel, err := config.SetupTest()
+	if err != nil {
+		t.Fatalf("Error initializing test setup: %v", err)
+	}
+	defer cancel()
+	defer client.Shutdown()
+
+	testConfig := config.ComprehensiveTestConfig{
+		Provider:  schemas.OpenRouter,
+		ChatModel: "meta-llama/llama-3.1-70b-instruct", // Using Llama 3.1 70B via OpenRouter
+		TextModel: "meta-llama/llama-3.1-8b-instruct",  // Using smaller model for text completion
+		EmbeddingModel: "",
+		Scenarios: config.TestScenarios{
+			TextCompletion:        true,
+			SimpleChat:            true,
+			ChatCompletionStream:  true,
+			MultiTurnConversation: true,
+			ToolCalls:             true,
+			MultipleToolCalls:     true,
+			End2EndToolCalling:    true,
+			AutomaticFunctionCall: true,
+			ImageURL:              false, // Llama models typically don't support image inputs
+			ImageBase64:           false,
+			MultipleImages:        false,
+			CompleteEnd2End:       true,
+			ProviderSpecific:      false,
+		},
+	}
+
+	runAllComprehensiveTests(t, client, ctx, testConfig)
+}