feat(llmcost): add (#3905)

Author: hekike

app/common/productcatalog.go

@@ -7,7 +7,11 @@ import (
 	"github.com/google/wire"
 
 	"github.com/openmeterio/openmeter/app/config"
+	"github.com/openmeterio/openmeter/openmeter/cost"
+	costadapter "github.com/openmeterio/openmeter/openmeter/cost/adapter"
+	costservice "github.com/openmeterio/openmeter/openmeter/cost/service"
 	entdb "github.com/openmeterio/openmeter/openmeter/ent/db"
+	"github.com/openmeterio/openmeter/openmeter/llmcost"
 	"github.com/openmeterio/openmeter/openmeter/meter"
 	productcatalogpgadapter "github.com/openmeterio/openmeter/openmeter/productcatalog/adapter"
 	"github.com/openmeterio/openmeter/openmeter/productcatalog/addon"
@@ -20,11 +24,13 @@ import (
 	"github.com/openmeterio/openmeter/openmeter/productcatalog/planaddon"
 	planaddonadapter "github.com/openmeterio/openmeter/openmeter/productcatalog/planaddon/adapter"
 	planaddonservice "github.com/openmeterio/openmeter/openmeter/productcatalog/planaddon/service"
+	"github.com/openmeterio/openmeter/openmeter/streaming"
 	"github.com/openmeterio/openmeter/openmeter/watermill/eventbus"
 )
 
 var ProductCatalog = wire.NewSet(
 	Feature,
+	Cost,
 	Plan,
 	Addon,
 	PlanAddon,
@@ -34,6 +40,10 @@ var Feature = wire.NewSet(
 	NewFeatureConnector,
 )
 
+var Cost = wire.NewSet(
+	NewCostService,
+)
+
 var Plan = wire.NewSet(
 	NewPlanService,
 )
@@ -56,6 +66,19 @@ func NewFeatureConnector(
 	return feature.NewFeatureConnector(featureRepo, meterService, publisher)
 }
 
+func NewCostService(
+	featureConnector feature.FeatureConnector,
+	meterService meter.Service,
+	streamingConnector streaming.Connector,
+	llmcostService llmcost.Service,
+) (cost.Service, error) {
+	adapter := costadapter.New(featureConnector, meterService, streamingConnector, llmcostService)
+
+	return costservice.New(costservice.Config{
+		Adapter: adapter,
+	})
+}
+
 func NewPlanService(
 	logger *slog.Logger,
 	db *entdb.Client,

cmd/server/main.go

@@ -148,12 +148,13 @@ func main() {
 			Billing:                     app.Billing,
 			BillingFeatureSwitches:      conf.Billing.FeatureSwitches,
 			CurrencyService:             app.CurrencyService,
+			CostService:                 app.CostService,
 			Customer:                    app.Customer,
 			DebugConnector:              debugConnector,
 			ErrorHandler:                errorsx.NewSlogHandler(logger),
 			EntitlementBalanceConnector: app.EntitlementRegistry.MeteredEntitlement,
 			EntitlementConnector:        app.EntitlementRegistry.Entitlement,
-			FeatureConnector:            app.EntitlementRegistry.Feature,
+			FeatureConnector:            app.FeatureConnector,
 			GrantConnector:              app.EntitlementRegistry.Grant,
 			GrantRepo:                   app.EntitlementRegistry.GrantRepo,
 			IngestService:               app.IngestService,

cmd/server/wire.go

@@ -15,6 +15,7 @@ import (
 	"github.com/openmeterio/openmeter/app/common"
 	"github.com/openmeterio/openmeter/app/config"
 	"github.com/openmeterio/openmeter/openmeter/billing"
+	"github.com/openmeterio/openmeter/openmeter/cost"
 	"github.com/openmeterio/openmeter/openmeter/currencies"
 	"github.com/openmeterio/openmeter/openmeter/customer"
 	"github.com/openmeterio/openmeter/openmeter/ent/db"
@@ -54,6 +55,7 @@ type Application struct {
 	CustomerEntitlementValidatorHook common.CustomerEntitlementValidatorHook
 	Billing                          billing.Service
 	CurrencyService                  currencies.CurrencyService
+	CostService                      cost.Service
 	EntClient                        *db.Client
 	EventPublisher                   eventbus.Publisher
 	EntitlementRegistry              *registry.Entitlement

cmd/server/wire_gen.go

@@ -7,11 +7,18 @@ input:
       let max_subjects = ${SEEDER_MAX_SUBJECTS:20}
       let event_type = "tokens"
       let source = "api-gateway"
-      let models = ["gpt-3", "gpt-4"]
+      let provider_models = [
+        {"provider": "openai", "models": ["gpt-4o", "gpt-4-turbo", "gpt-3.5-turbo"]},
+        {"provider": "anthropic", "models": ["claude-3-5-sonnet", "claude-3-opus", "claude-3-5-haiku"]}
+      ]
       let types = ["input", "output"]
       let subject = "customer-%d".format(random_int(seed: timestamp_unix_nano()) % $max_subjects)
       let time = now().ts_format()
-      let model = $models.index(random_int(seed: timestamp_unix_nano()) % $models.length())
+      # let time = (now().ts_sub_iso8601("P30D").ts_unix() + random_int(min: 60, max: 60 * 60 * 24 * 30)).ts_format()
+      let idx = random_int(seed: timestamp_unix_nano()) % $provider_models.length()
+      let selection = $provider_models.index($idx)
+      let provider = $selection.provider
+      let model = $selection.models.index(random_int(seed: timestamp_unix_nano()) % $selection.models.length())
       let type = $types.index(random_int(seed: timestamp_unix_nano()) % $types.length())
       let tokens = random_int(seed: timestamp_unix_nano(), max: 100)
       root = {
@@ -23,6 +30,7 @@ input:
         "time": $time,
         "data": {
           "model": $model,
+          "provider": $provider,
           "type": $type,
           "tokens": $tokens,
         },

etc/seed/seed.yaml

@@ -0,0 +1,8 @@
+package cost
+
+import "context"
+
+// Adapter provides the data access layer for cost computation.
+type Adapter interface {
+	QueryFeatureCost(ctx context.Context, input QueryFeatureCostInput) (*CostQueryResult, error)
+}