Support local model and embedding in calculatePrice()

JiawenZou · Mar 8, 2024 · 2a8522e · 2a8522e
1 parent 0b8e230
commit 2a8522e
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/embedding/local.go b/embedding/local.go
@@ -73,6 +73,8 @@ func (p *LocalEmbeddingProvider) calculatePrice(res *EmbeddingResult) error {
 		pricePerThousandTokens = 0.00002
 	case strings.Contains(embeddingModel, "text-embedding-3-large"):
 		pricePerThousandTokens = 0.00013
+	case embeddingModel == "custom-embedding":
+		pricePerThousandTokens = 0.0001
 	default:
 		return fmt.Errorf("calculatePrice() error: unknown model type: %s", embeddingModel)
 	}

diff --git a/model/local.go b/model/local.go
@@ -147,6 +147,11 @@ func (p *LocalModelProvider) calculatePrice(modelResult *ModelResult) error {
 			outputPricePerThousandTokens = 0.06
 		}
 
+	// local custom model:
+	case model == "custom-model":
+		inputPricePerThousandTokens = 0.001
+		outputPricePerThousandTokens = 0.002
+
 	// dall-e model
 	case strings.Contains(model, "dall-e-3"):
 		modelResult.TotalPrice = float64(modelResult.ImageCount) * 0.08