From 7446f3463defa8f5bb4ef2ec771e72fcb73a7869 Mon Sep 17 00:00:00 2001
From: hobokenchicken <dustin@dustin.coffee>
Date: Mon, 27 Apr 2026 10:42:29 -0400
Subject: [PATCH] fix: add per-image cost tracking for DALL-E and Imagen

---
 .pi-lens/turn-state.json  |  2 +-
 internal/server/server.go | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/.pi-lens/turn-state.json b/.pi-lens/turn-state.json
index 41e962f9..72eb0e0e 100644
--- a/.pi-lens/turn-state.json
+++ b/.pi-lens/turn-state.json
@@ -2,5 +2,5 @@
   "files": {},
   "turnCycles": 0,
   "maxCycles": 3,
-  "lastUpdated": "2026-04-27T01:12:44.352Z"
+  "lastUpdated": "2026-04-27T14:41:46.671Z"
 }
\ No newline at end of file
diff --git a/internal/server/server.go b/internal/server/server.go
index d8a3bb50..ced8c25e 100644
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -564,15 +564,51 @@ func (s *Server) handleImageGenerations(c *gin.Context) {
 		promptTokens = 1
 	}
 
+	// Calculate per-image cost (not per-token like chat)
+	cost := imageGenCost(providerName, req.Model, req.Size, uint32(len(resp.Data)))
+
 	s.logRequest(startTime, clientID, providerName, req.Model, &models.Usage{
 		PromptTokens:     promptTokens,
 		CompletionTokens: uint32(len(resp.Data)),
 		TotalTokens:      promptTokens + uint32(len(resp.Data)),
 	}, nil, false)
 
+	// Update cost in DB — image gen is per-image, not per-token
+	if cost > 0 {
+		s.database.Exec("UPDATE llm_requests SET cost = ? WHERE id = (SELECT MAX(id) FROM llm_requests)", cost)
+	}
+
 	c.JSON(http.StatusOK, resp)
 }
 
+// imageGenCost returns per-image pricing for known image generation models.
+func imageGenCost(provider, model string, size *string, n uint32) float64 {
+	if n == 0 {
+		return 0
+	}
+	modelLower := strings.ToLower(model)
+	var perImage float64
+
+	switch {
+	case strings.Contains(modelLower, "dall-e-3"):
+		perImage = 0.040 // standard 1024x1024
+		if size != nil {
+			s := *size
+			if s == "1024x1792" || s == "1792x1024" {
+				perImage = 0.080
+			}
+		}
+	case strings.Contains(modelLower, "dall-e-2"):
+		perImage = 0.020
+	case strings.Contains(modelLower, "imagen"):
+		perImage = 0.040 // approximate
+	default:
+		return 0
+	}
+
+	return perImage * float64(n)
+}
+
 func (s *Server) logRequest(start time.Time, clientID, provider, model string, usage *models.Usage, err error, hasImages bool) {
 	entry := RequestLog{
 		Timestamp:  start,