62 lines
2.0 KiB
Go
62 lines
2.0 KiB
Go
|
package workbench
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/models/embedding"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
)
|
||
|
|
|
||
|
|
type usageEmbedding struct{ calls int }
|
||
|
|
|
||
|
|
func (m *usageEmbedding) Embed(context.Context, string) ([]float32, error) {
|
||
|
|
m.calls++
|
||
|
|
return []float32{1, 2}, nil
|
||
|
|
}
|
||
|
|
func (m *usageEmbedding) BatchEmbed(_ context.Context, texts []string) ([][]float32, error) {
|
||
|
|
m.calls++
|
||
|
|
out := make([][]float32, len(texts))
|
||
|
|
for i := range out {
|
||
|
|
out[i] = []float32{1, 2}
|
||
|
|
}
|
||
|
|
return out, nil
|
||
|
|
}
|
||
|
|
func (m *usageEmbedding) BatchEmbedWithPool(ctx context.Context, model embedding.Embedder, texts []string) ([][]float32, error) {
|
||
|
|
return model.BatchEmbed(ctx, texts)
|
||
|
|
}
|
||
|
|
func (*usageEmbedding) GetModelID() string { return "embedding" }
|
||
|
|
func (*usageEmbedding) GetModelName() string { return "embedding model" }
|
||
|
|
func (*usageEmbedding) GetDimensions() int { return 2 }
|
||
|
|
|
||
|
|
func TestAuxiliaryUsageRecordsUnitsOnceWithoutInventedBilling(t *testing.T) {
|
||
|
|
s := sqliteWorkbench(t)
|
||
|
|
ctx := workbenchContext(7, "alice", types.TenantRoleContributor)
|
||
|
|
raw := &usageEmbedding{}
|
||
|
|
model := s.WrapEmbedding(raw)
|
||
|
|
if _, err := model.BatchEmbedWithPool(ctx, model, []string{"主题", "练习"}); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
var rows []Usage
|
||
|
|
if err := s.db.Find(&rows).Error; err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if raw.calls != 1 || len(rows) != 1 {
|
||
|
|
t.Fatalf("double counted: calls=%d rows=%d", raw.calls, len(rows))
|
||
|
|
}
|
||
|
|
r := rows[0]
|
||
|
|
if r.TenantID != 7 || r.OwnerID != "alice" || r.Points != 0 || r.InputTokens != 0 || r.OutputTokens != 0 {
|
||
|
|
t.Fatalf("wrong attribution or invented charge: %+v", r)
|
||
|
|
}
|
||
|
|
if r.Metadata["input_characters"] != float64(4) || r.Metadata["input_texts"] != float64(2) || r.Metadata["billing_status"] != "provider_usage_unavailable" {
|
||
|
|
t.Fatalf("missing measured units: %+v", r.Metadata)
|
||
|
|
}
|
||
|
|
var n int64
|
||
|
|
if err := s.db.Model(&Transaction{}).Count(&n).Error; err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if n != 0 {
|
||
|
|
t.Fatal("auxiliary metrics created an unsupported debit")
|
||
|
|
}
|
||
|
|
}
|