Просмотр исходного кода

feat: support price condition time range (#422)

zijiren 1 месяц назад
Родитель
Сommit
ea37a5242b
2 измененных файлов с 438 добавлено и 5 удалено
  1. 69 5
      core/model/usage.go
  2. 369 0
      core/model/usage_test.go

+ 69 - 5
core/model/usage.go

@@ -3,6 +3,7 @@ package model
 import (
 	"fmt"
 	"math"
+	"time"
 )
 
 type PriceCondition struct {
@@ -10,6 +11,8 @@ type PriceCondition struct {
 	InputTokenMax  int64 `json:"input_token_max,omitempty"`
 	OutputTokenMin int64 `json:"output_token_min,omitempty"`
 	OutputTokenMax int64 `json:"output_token_max,omitempty"`
+	StartTime      int64 `json:"start_time,omitempty"` // Unix timestamp, 0 means no start limit
+	EndTime        int64 `json:"end_time,omitempty"`   // Unix timestamp, 0 means no end limit
 }
 
 type ConditionalPrice struct {
@@ -80,6 +83,18 @@ func (p *Price) ValidateConditionalPrices() error {
 			}
 		}
 
+		// Validate time range
+		if condition.StartTime > 0 && condition.EndTime > 0 {
+			if condition.StartTime >= condition.EndTime {
+				return fmt.Errorf(
+					"conditional price %d: start time (%d) must be before end time (%d)",
+					i,
+					condition.StartTime,
+					condition.EndTime,
+				)
+			}
+		}
+
 		// Check for overlaps with other conditions
 		for j := i + 1; j < len(p.ConditionalPrices); j++ {
 			otherCondition := p.ConditionalPrices[j].Condition
@@ -94,11 +109,18 @@ func (p *Price) ValidateConditionalPrices() error {
 					condition.OutputTokenMin, condition.OutputTokenMax,
 					otherCondition.OutputTokenMin, otherCondition.OutputTokenMax,
 				) {
-					return fmt.Errorf(
-						"conditional prices %d and %d have overlapping conditions",
-						i,
-						j,
-					)
+					// If both token ranges overlap, check if time ranges also overlap
+					// If time ranges don't overlap, conditions are still valid
+					if hasTimeRangeOverlap(
+						condition.StartTime, condition.EndTime,
+						otherCondition.StartTime, otherCondition.EndTime,
+					) {
+						return fmt.Errorf(
+							"conditional prices %d and %d have overlapping conditions",
+							i,
+							j,
+						)
+					}
 				}
 			}
 		}
@@ -141,6 +163,37 @@ func hasRangeOverlap(min1, max1, min2, max2 int64) bool {
 	return actualMax1 >= actualMin2 && actualMin1 <= actualMax2
 }
 
+// hasTimeRangeOverlap checks if two time ranges overlap
+// Unlike hasRangeOverlap, this uses strict inequality to allow adjacent time ranges
+// Time range is defined by [start, end], where 0 means unbounded
+func hasTimeRangeOverlap(start1, end1, start2, end2 int64) bool {
+	// Convert 0 to appropriate bounds for comparison
+	actualStart1 := start1
+	actualEnd1 := end1
+	actualStart2 := start2
+	actualEnd2 := end2
+
+	if actualStart1 == 0 {
+		actualStart1 = 0
+	}
+
+	if actualEnd1 == 0 {
+		actualEnd1 = math.MaxInt64
+	}
+
+	if actualStart2 == 0 {
+		actualStart2 = 0
+	}
+
+	if actualEnd2 == 0 {
+		actualEnd2 = math.MaxInt64
+	}
+
+	// Check if ranges overlap with strict inequality: range1.end > range2.start && range1.start < range2.end
+	// This allows adjacent ranges like [t1, t2] and [t2, t3] to be considered non-overlapping
+	return actualEnd1 > actualStart2 && actualStart1 < actualEnd2
+}
+
 // validateConditionalPriceOrdering checks if conditional prices are properly ordered
 func (p *Price) validateConditionalPriceOrdering() error {
 	if len(p.ConditionalPrices) <= 1 {
@@ -195,10 +248,21 @@ func (p *Price) SelectConditionalPrice(usage Usage) Price {
 
 	inputTokens := int64(usage.InputTokens)
 	outputTokens := int64(usage.OutputTokens)
+	currentTime := time.Now().Unix()
 
 	for _, conditionalPrice := range p.ConditionalPrices {
 		condition := conditionalPrice.Condition
 
+		// Check time range
+		if condition.StartTime > 0 && currentTime < condition.StartTime {
+			continue
+		}
+
+		if condition.EndTime > 0 && currentTime > condition.EndTime {
+			continue
+		}
+
+		// Check token ranges
 		if condition.InputTokenMin > 0 && inputTokens < condition.InputTokenMin {
 			continue
 		}

+ 369 - 0
core/model/usage_test.go

@@ -2,6 +2,7 @@ package model_test
 
 import (
 	"testing"
+	"time"
 
 	"github.com/labring/aiproxy/core/model"
 )
@@ -349,3 +350,371 @@ func TestPrice_ValidateConditionalPrices(t *testing.T) {
 		})
 	}
 }
+
+func TestPrice_ValidateConditionalPrices_WithTime(t *testing.T) {
+	now := time.Now().Unix()
+	future := now + 3600 // 1 hour from now
+	past := now - 3600   // 1 hour ago
+
+	tests := []struct {
+		name    string
+		price   model.Price
+		wantErr bool
+	}{
+		{
+			name: "Valid time range - future time window",
+			price: model.Price{
+				ConditionalPrices: []model.ConditionalPrice{
+					{
+						Condition: model.PriceCondition{
+							InputTokenMin: 0,
+							InputTokenMax: 32000,
+							StartTime:     now,
+							EndTime:       future,
+						},
+						Price: model.Price{
+							InputPrice:  0.0008,
+							OutputPrice: 0.002,
+						},
+					},
+				},
+			},
+			wantErr: false,
+		},
+		{
+			name: "Valid time range - no end time",
+			price: model.Price{
+				ConditionalPrices: []model.ConditionalPrice{
+					{
+						Condition: model.PriceCondition{
+							InputTokenMin: 0,
+							InputTokenMax: 32000,
+							StartTime:     now,
+							EndTime:       0, // no end limit
+						},
+						Price: model.Price{
+							InputPrice:  0.0008,
+							OutputPrice: 0.002,
+						},
+					},
+				},
+			},
+			wantErr: false,
+		},
+		{
+			name: "Valid time range - no start time",
+			price: model.Price{
+				ConditionalPrices: []model.ConditionalPrice{
+					{
+						Condition: model.PriceCondition{
+							InputTokenMin: 0,
+							InputTokenMax: 32000,
+							StartTime:     0, // no start limit
+							EndTime:       future,
+						},
+						Price: model.Price{
+							InputPrice:  0.0008,
+							OutputPrice: 0.002,
+						},
+					},
+				},
+			},
+			wantErr: false,
+		},
+		{
+			name: "Invalid time range - start time >= end time",
+			price: model.Price{
+				ConditionalPrices: []model.ConditionalPrice{
+					{
+						Condition: model.PriceCondition{
+							InputTokenMin: 0,
+							InputTokenMax: 32000,
+							StartTime:     future,
+							EndTime:       now, // end before start
+						},
+						Price: model.Price{
+							InputPrice:  0.0008,
+							OutputPrice: 0.002,
+						},
+					},
+				},
+			},
+			wantErr: true,
+		},
+		{
+			name: "Invalid time range - start time equals end time",
+			price: model.Price{
+				ConditionalPrices: []model.ConditionalPrice{
+					{
+						Condition: model.PriceCondition{
+							InputTokenMin: 0,
+							InputTokenMax: 32000,
+							StartTime:     now,
+							EndTime:       now, // same time
+						},
+						Price: model.Price{
+							InputPrice:  0.0008,
+							OutputPrice: 0.002,
+						},
+					},
+				},
+			},
+			wantErr: true,
+		},
+		{
+			name: "Multiple conditions with different time ranges",
+			price: model.Price{
+				ConditionalPrices: []model.ConditionalPrice{
+					{
+						Condition: model.PriceCondition{
+							InputTokenMin: 0,
+							InputTokenMax: 32000,
+							StartTime:     past,
+							EndTime:       now,
+						},
+						Price: model.Price{
+							InputPrice:  0.0008,
+							OutputPrice: 0.002,
+						},
+					},
+					{
+						Condition: model.PriceCondition{
+							InputTokenMin: 0,
+							InputTokenMax: 32000,
+							StartTime:     now,
+							EndTime:       future,
+						},
+						Price: model.Price{
+							InputPrice:  0.001,
+							OutputPrice: 0.003,
+						},
+					},
+				},
+			},
+			wantErr: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			err := tt.price.ValidateConditionalPrices()
+
+			if tt.wantErr {
+				if err == nil {
+					t.Errorf("%s: ValidateConditionalPrices() expected error but got nil", tt.name)
+				}
+				return
+			}
+
+			if err != nil {
+				t.Errorf("%s: ValidateConditionalPrices() unexpected error = %v", tt.name, err)
+			}
+		})
+	}
+}
+
+func TestPrice_SelectConditionalPrice_WithTime(t *testing.T) {
+	now := time.Now().Unix()
+	past := now - 3600      // 1 hour ago
+	future := now + 3600    // 1 hour from now
+	farFuture := now + 7200 // 2 hours from now
+
+	tests := []struct {
+		name           string
+		price          model.Price
+		usage          model.Usage
+		expectedInput  float64
+		expectedOutput float64
+	}{
+		{
+			name: "Select price within active time range",
+			price: model.Price{
+				InputPrice:  0.001,
+				OutputPrice: 0.002,
+				ConditionalPrices: []model.ConditionalPrice{
+					{
+						Condition: model.PriceCondition{
+							InputTokenMin: 0,
+							InputTokenMax: 32000,
+							StartTime:     past,
+							EndTime:       future,
+						},
+						Price: model.Price{
+							InputPrice:  0.0005,
+							OutputPrice: 0.001,
+						},
+					},
+				},
+			},
+			usage: model.Usage{
+				InputTokens:  1000,
+				OutputTokens: 500,
+			},
+			expectedInput:  0.0005,
+			expectedOutput: 0.001,
+		},
+		{
+			name: "Fallback to default price when time range not active (before start)",
+			price: model.Price{
+				InputPrice:  0.001,
+				OutputPrice: 0.002,
+				ConditionalPrices: []model.ConditionalPrice{
+					{
+						Condition: model.PriceCondition{
+							InputTokenMin: 0,
+							InputTokenMax: 32000,
+							StartTime:     future,
+							EndTime:       farFuture,
+						},
+						Price: model.Price{
+							InputPrice:  0.0005,
+							OutputPrice: 0.001,
+						},
+					},
+				},
+			},
+			usage: model.Usage{
+				InputTokens:  1000,
+				OutputTokens: 500,
+			},
+			expectedInput:  0.001,
+			expectedOutput: 0.002,
+		},
+		{
+			name: "Fallback to default price when time range expired",
+			price: model.Price{
+				InputPrice:  0.001,
+				OutputPrice: 0.002,
+				ConditionalPrices: []model.ConditionalPrice{
+					{
+						Condition: model.PriceCondition{
+							InputTokenMin: 0,
+							InputTokenMax: 32000,
+							StartTime:     past - 7200, // 3 hours ago
+							EndTime:       past,        // 1 hour ago
+						},
+						Price: model.Price{
+							InputPrice:  0.0005,
+							OutputPrice: 0.001,
+						},
+					},
+				},
+			},
+			usage: model.Usage{
+				InputTokens:  1000,
+				OutputTokens: 500,
+			},
+			expectedInput:  0.001,
+			expectedOutput: 0.002,
+		},
+		{
+			name: "Select first matching price with multiple time-based conditions",
+			price: model.Price{
+				InputPrice:  0.001,
+				OutputPrice: 0.002,
+				ConditionalPrices: []model.ConditionalPrice{
+					{
+						Condition: model.PriceCondition{
+							InputTokenMin: 0,
+							InputTokenMax: 32000,
+							StartTime:     past,
+							EndTime:       future,
+						},
+						Price: model.Price{
+							InputPrice:  0.0005,
+							OutputPrice: 0.001,
+						},
+					},
+					{
+						Condition: model.PriceCondition{
+							InputTokenMin: 0,
+							InputTokenMax: 32000,
+							StartTime:     past,
+							EndTime:       farFuture, // broader time range
+						},
+						Price: model.Price{
+							InputPrice:  0.0008,
+							OutputPrice: 0.0015,
+						},
+					},
+				},
+			},
+			usage: model.Usage{
+				InputTokens:  1000,
+				OutputTokens: 500,
+			},
+			expectedInput:  0.0005,
+			expectedOutput: 0.001,
+		},
+		{
+			name: "Time range with no end time (ongoing promotion)",
+			price: model.Price{
+				InputPrice:  0.001,
+				OutputPrice: 0.002,
+				ConditionalPrices: []model.ConditionalPrice{
+					{
+						Condition: model.PriceCondition{
+							InputTokenMin: 0,
+							InputTokenMax: 32000,
+							StartTime:     past,
+							EndTime:       0, // no end time
+						},
+						Price: model.Price{
+							InputPrice:  0.0005,
+							OutputPrice: 0.001,
+						},
+					},
+				},
+			},
+			usage: model.Usage{
+				InputTokens:  1000,
+				OutputTokens: 500,
+			},
+			expectedInput:  0.0005,
+			expectedOutput: 0.001,
+		},
+		{
+			name: "Time range with no start time (promotion until end)",
+			price: model.Price{
+				InputPrice:  0.001,
+				OutputPrice: 0.002,
+				ConditionalPrices: []model.ConditionalPrice{
+					{
+						Condition: model.PriceCondition{
+							InputTokenMin: 0,
+							InputTokenMax: 32000,
+							StartTime:     0, // no start time
+							EndTime:       future,
+						},
+						Price: model.Price{
+							InputPrice:  0.0005,
+							OutputPrice: 0.001,
+						},
+					},
+				},
+			},
+			usage: model.Usage{
+				InputTokens:  1000,
+				OutputTokens: 500,
+			},
+			expectedInput:  0.0005,
+			expectedOutput: 0.001,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			selectedPrice := tt.price.SelectConditionalPrice(tt.usage)
+
+			if float64(selectedPrice.InputPrice) != tt.expectedInput {
+				t.Errorf("%s: expected input price %v, got %v",
+					tt.name, tt.expectedInput, float64(selectedPrice.InputPrice))
+			}
+
+			if float64(selectedPrice.OutputPrice) != tt.expectedOutput {
+				t.Errorf("%s: expected output price %v, got %v",
+					tt.name, tt.expectedOutput, float64(selectedPrice.OutputPrice))
+			}
+		})
+	}
+}