fix price calculation when http responses are not okay

bricks-cloud · Mar 28, 2024 · a4b8991 · a4b8991
1 parent 9e4a3d6
commit a4b8991
Showing 1 changed file with 10 additions and 3 deletions.
diff --git a/internal/message/handler.go b/internal/message/handler.go
@@ -327,7 +327,9 @@ func (h *Handler) decorateEvent(m Message) error {
 		e.Event.PromptTokenCount = tks
 
 		e.Event.CompletionTokenCount = completiontks
-		e.Event.CostInUsd = completionCost + cost
+		if e.Event.Status == http.StatusOK {
+			e.Event.CostInUsd = completionCost + cost
+		}
 	}
 
 	if e.Event.Path == "/api/providers/azure/openai/deployments/:deployment_id/chat/completions" {
@@ -359,7 +361,10 @@ func (h *Handler) decorateEvent(m Message) error {
 
 			e.Event.PromptTokenCount = tks
 			e.Event.CompletionTokenCount = completiontks
-			e.Event.CostInUsd = cost + completionCost
+
+			if e.Event.Status == http.StatusOK {
+				e.Event.CostInUsd = cost + completionCost
+			}
 		}
 	}
 
@@ -386,7 +391,9 @@ func (h *Handler) decorateEvent(m Message) error {
 
 			e.Event.PromptTokenCount = tks
 			e.Event.CompletionTokenCount = completiontks
-			e.Event.CostInUsd = cost + completionCost
+			if e.Event.Status == http.StatusOK {
+				e.Event.CostInUsd = cost + completionCost
+			}
 		}
 	}