Merge remote-tracking branch 'origin/issue_6107_3' into issue_6107_3

.github/workflows/integrationci.yaml

@@ -47,6 +47,7 @@ jobs:
           - dashboard
           - ingestionkeys
           - inframonitoring
+          - llmpricingrules
           - logspipelines
           - passwordauthn
           - preference

docs/api/openapi.yml

@@ -12772,9 +12772,8 @@ paths:
     put:
       deprecated: false
       description: Single write endpoint used by both the user and the Zeus sync job.
-        Per-rule match is by id, then sourceId, then insert. Override rows (is_override=true)
-        are fully preserved when the request does not provide isOverride; only synced_at
-        is stamped.
+        Rules without isOverride are matched by sourceId and override rows (is_override=true)
+        are skipped. Rules with isOverride are matched by id and inserted when new.
       operationId: CreateOrUpdateLLMPricingRules
       requestBody:
         content:

frontend/src/api/generated/services/llmpricingrules/index.ts

@@ -150,7 +150,7 @@ export const invalidateListLLMPricingRules = async (
 };
 
 /**
- * Single write endpoint used by both the user and the Zeus sync job. Per-rule match is by id, then sourceId, then insert. Override rows (is_override=true) are fully preserved when the request does not provide isOverride; only synced_at is stamped.
+ * Single write endpoint used by both the user and the Zeus sync job. Rules without isOverride are matched by sourceId and override rows (is_override=true) are skipped. Rules with isOverride are matched by id and inserted when new.
  * @summary Create or update pricing rules
  */
 export const createOrUpdateLLMPricingRules = (

pkg/apiserver/signozapiserver/llmpricingrule.go

@@ -37,7 +37,7 @@ func (provider *provider) addLLMPricingRuleRoutes(router *mux.Router) error {
 			ID:                 "CreateOrUpdateLLMPricingRules",
 			Tags:               []string{"llmpricingrules"},
 			Summary:            "Create or update pricing rules",
-			Description:        "Single write endpoint used by both the user and the Zeus sync job. Per-rule match is by id, then sourceId, then insert. Override rows (is_override=true) are fully preserved when the request does not provide isOverride; only synced_at is stamped.",
+			Description:        "Single write endpoint used by both the user and the Zeus sync job. Rules without isOverride are matched by sourceId and override rows (is_override=true) are skipped. Rules with isOverride are matched by id and inserted when new.",
 			Request:            new(llmpricingruletypes.UpdatableLLMPricingRules),
 			RequestContentType: "application/json",
 			SuccessStatusCode:  http.StatusNoContent,

pkg/modules/llmpricingrule/impllmpricingrule/module.go

@@ -60,38 +60,29 @@ func (module *module) ListUnmappedModels(ctx context.Context, orgID valuer.UUID)
 	return unmapped, nil
 }
 
-// CreateOrUpdate applies a batch of pricing rule changes:
-//   - ID set       → match by id, overwrite fields.
-//   - SourceID set → match by source_id; if found overwrite, else insert.
-//   - neither set  → insert a new user-created row (is_override = true).
-//
-// When UpdatableLLMPricingRule.IsOverride is nil AND the matched row has
-// is_override = true, the row is fully preserved — only synced_at is stamped.
+// CreateOrUpdate saves a batch of pricing rules. isOverride decides how a rule
+// is matched, see UpdatableLLMPricingRule. New rules are inserted on either path.
 func (module *module) CreateOrUpdate(ctx context.Context, orgID valuer.UUID, userEmail string, rules []*llmpricingruletypes.UpdatableLLMPricingRule) error {
 	now := time.Now()
 
-	upsert := func(ctx context.Context, u *llmpricingruletypes.UpdatableLLMPricingRule) error {
+	var byID, bySourceID []*llmpricingruletypes.LLMPricingRule
+	for _, u := range rules {
 		if u == nil {
 			return errors.Newf(errors.TypeInvalidInput, llmpricingruletypes.ErrCodePricingRuleInvalidInput, "rule entry is null")
 		}
-		existing, err := module.findExisting(ctx, orgID, u)
-		if err != nil && errors.Ast(err, errors.TypeNotFound) {
-			return module.store.Create(ctx, llmpricingruletypes.NewLLMPricingRuleFromUpdatable(u, orgID, userEmail, now))
+		rule := llmpricingruletypes.NewLLMPricingRuleFromUpdatable(u, orgID, userEmail, now)
+		if u.IsOverride == nil {
+			bySourceID = append(bySourceID, rule)
+		} else {
+			byID = append(byID, rule)
 		}
-		if err != nil {
-			return err
-		}
-		existing.Update(u, userEmail, now)
-		return module.store.Update(ctx, existing)
 	}
 
 	err := module.store.RunInTx(ctx, func(ctx context.Context) error {
-		for _, u := range rules {
-			if err := upsert(ctx, u); err != nil {
-				return err
-			}
+		if err := module.store.UpsertByID(ctx, byID); err != nil {
+			return err
 		}
-		return nil
+		return module.store.UpsertBySourceID(ctx, bySourceID)
 	})
 	if err != nil {
 		return err
@@ -172,20 +163,6 @@ func (module *module) listAllRules(ctx context.Context, orgID valuer.UUID) ([]*l
 	return all, nil
 }
 
-// findExisting returns the row matching the updatable's ID or SourceID.
-// Returns a TypeNotFound error when neither matches; the caller treats that
-// as "insert new".
-func (module *module) findExisting(ctx context.Context, orgID valuer.UUID, u *llmpricingruletypes.UpdatableLLMPricingRule) (*llmpricingruletypes.LLMPricingRule, error) {
-	switch {
-	case u.ID != nil:
-		return module.store.Get(ctx, orgID, *u.ID)
-	case u.SourceID != nil:
-		return module.store.GetBySourceID(ctx, orgID, *u.SourceID)
-	default:
-		return nil, errors.Newf(errors.TypeNotFound, llmpricingruletypes.ErrCodePricingRuleNotFound, "rule has neither id nor sourceId")
-	}
-}
-
 // discoverModels runs a QBv5 traces aggregation grouped by gen_ai.request.model
 // over the lookback window and returns each distinct model with its span count.
 func (module *module) discoverModels(ctx context.Context, orgID valuer.UUID) ([]*llmpricingruletypes.UnmappedModel, error) {

pkg/modules/llmpricingrule/impllmpricingrule/store.go

@@ -7,8 +7,13 @@ import (
 	"github.com/SigNoz/signoz/pkg/sqlstore"
 	"github.com/SigNoz/signoz/pkg/types/llmpricingruletypes"
 	"github.com/SigNoz/signoz/pkg/valuer"
+	"github.com/uptrace/bun"
 )
 
+// Columns an existing row gets when a rule matches it. id, org_id, source_id,
+// created_at and created_by are never changed.
+var upsertColumns = []string{"model", "provider", "model_pattern", "unit", "pricing", "is_override", "enabled", "synced_at", "updated_at", "updated_by"}
+
 type store struct {
 	sqlstore sqlstore.SQLStore
 }
@@ -64,57 +69,60 @@ func (store *store) Get(ctx context.Context, orgID, id valuer.UUID) (*llmpricing
 	return rule, nil
 }
 
-func (store *store) GetBySourceID(ctx context.Context, orgID, sourceID valuer.UUID) (*llmpricingruletypes.LLMPricingRule, error) {
-	rule := new(llmpricingruletypes.LLMPricingRule)
-
-	err := store.sqlstore.
-		BunDBCtx(ctx).
-		NewSelect().
-		Model(rule).
-		Where("org_id = ?", orgID).
-		Where("source_id = ?", sourceID).
-		Scan(ctx)
-	if err != nil {
-		return nil, store.sqlstore.WrapNotFoundErrf(err, llmpricingruletypes.ErrCodePricingRuleNotFound, "pricing rule with source_id %s not found in the org", sourceID)
+// UpsertByID replaces the row with the same id, or inserts when there is
+// none. Rows of other orgs are left alone and reported as not found.
+func (store *store) UpsertByID(ctx context.Context, rules []*llmpricingruletypes.LLMPricingRule) error {
+	if len(rules) == 0 {
+		return nil
 	}
 
-	return rule, nil
-}
-
-func (store *store) Create(ctx context.Context, rule *llmpricingruletypes.LLMPricingRule) error {
-	_, err := store.sqlstore.
+	// bun can overwrite the rules slice with the rows it gets back, so count first.
+	expected := len(rules)
+	query := store.sqlstore.
 		BunDBCtx(ctx).
 		NewInsert().
-		Model(rule).
-		Exec(ctx)
-	if err != nil {
-		return store.sqlstore.WrapAlreadyExistsErrf(err, llmpricingruletypes.ErrCodePricingRuleAlreadyExists, "pricing rule with model %s already exists", rule.Model)
+		Model(&rules).
+		On("CONFLICT (id) DO UPDATE").
+		Where("llm_pricing_rule.org_id = EXCLUDED.org_id")
+	for _, col := range upsertColumns {
+		query = query.Set("? = EXCLUDED.?", bun.Ident(col), bun.Ident(col))
 	}
+
+	res, err := query.Exec(ctx)
+	if err != nil {
+		return err
+	}
+
+	affected, err := res.RowsAffected()
+	if err != nil {
+		return err
+	}
+	if int(affected) != expected {
+		return errors.Newf(errors.TypeNotFound, llmpricingruletypes.ErrCodePricingRuleNotFound, "one or more pricing rules not found in the org")
+	}
+
 	return nil
 }
 
-func (store *store) Update(ctx context.Context, rule *llmpricingruletypes.LLMPricingRule) error {
-	res, err := store.sqlstore.
+// UpsertBySourceID replaces the row with the same source_id, or inserts when
+// there is none. Rows the user has overridden are skipped.
+func (store *store) UpsertBySourceID(ctx context.Context, rules []*llmpricingruletypes.LLMPricingRule) error {
+	if len(rules) == 0 {
+		return nil
+	}
+
+	query := store.sqlstore.
 		BunDBCtx(ctx).
-		NewUpdate().
-		Model(rule).
-		Where("org_id = ?", rule.OrgID).
-		Where("id = ?", rule.ID).
-		ExcludeColumn("id", "org_id", "source_id", "created_at", "created_by").
-		Exec(ctx)
-	if err != nil {
-		return err
+		NewInsert().
+		Model(&rules).
+		On("CONFLICT (org_id, source_id) WHERE source_id IS NOT NULL DO UPDATE").
+		Where("NOT llm_pricing_rule.is_override")
+	for _, col := range upsertColumns {
+		query = query.Set("? = EXCLUDED.?", bun.Ident(col), bun.Ident(col))
 	}
 
-	rowsAffected, err := res.RowsAffected()
-	if err != nil {
-		return err
-	}
-	if rowsAffected == 0 {
-		return errors.Newf(errors.TypeNotFound, llmpricingruletypes.ErrCodePricingRuleNotFound, "pricing rule %s not found in the org", rule.ID)
-	}
-
-	return nil
+	_, err := query.Exec(ctx)
+	return err
 }
 
 func (store *store) Delete(ctx context.Context, orgID, id valuer.UUID) error {

pkg/types/llmpricingruletypes/pricing.go

@@ -77,11 +77,11 @@ type LLMPricingRule struct {
 	Provider     string             `bun:"provider,type:text,notnull" json:"provider" required:"true"`
 	ModelPattern StringSlice        `bun:"model_pattern,type:text,notnull" json:"modelPattern" required:"true"`
 	Unit         LLMPricingRuleUnit `bun:"unit,type:text,notnull" json:"unit" required:"true"`
-	Pricing      LLMRulePricing     `bun:"pricing,type:text,notnull,default:'{}'" json:"pricing" required:"true"`
+	Pricing      LLMRulePricing     `bun:"pricing,type:text,notnull" json:"pricing" required:"true"`
 	// IsOverride marks the row as user-pinned. When true, Zeus skips it entirely.
-	IsOverride bool       `bun:"is_override,notnull,default:false" json:"isOverride" required:"true"`
+	IsOverride bool       `bun:"is_override,notnull" json:"isOverride" required:"true"`
 	SyncedAt   *time.Time `bun:"synced_at" json:"syncedAt,omitempty"`
-	Enabled    bool       `bun:"enabled,notnull,default:true" json:"enabled" required:"true"`
+	Enabled    bool       `bun:"enabled,notnull" json:"enabled" required:"true"`
 }
 
 type GettableLLMPricingRule = LLMPricingRule
@@ -90,14 +90,9 @@ type StorableLLMPricingRule = LLMPricingRule
 
 // UpdatableLLMPricingRule is one entry in the bulk upsert batch.
 //
-// Identification:
-//   - ID set       → match by id (user editing a known row).
-//   - SourceID set → match by source_id (Zeus sync, or user editing a Zeus-synced row).
-//   - neither set  → insert a new row with source_id = NULL (user-created custom rule).
-//
-// IsOverride is a pointer so the caller can distinguish "not sent" from "set to false".
-// When IsOverride is nil AND the matched row has is_override = true, the row is fully
-// preserved — only synced_at is stamped.
+// IsOverride is a pointer so "not sent" differs from "false". Without it the
+// rule is matched on source_id and overridden rows are skipped. With it the
+// rule is matched on id and the value is stored.
 type UpdatableLLMPricingRule struct {
 	ID           *valuer.UUID       `json:"id,omitempty"`
 	SourceID     *valuer.UUID       `json:"sourceId,omitempty"`
@@ -214,6 +209,11 @@ func NewGettableUnmappedModels(items []*UnmappedModel) *GettableUnmappedModels {
 }
 
 func NewLLMPricingRuleFromUpdatable(u *UpdatableLLMPricingRule, orgID valuer.UUID, userEmail string, now time.Time) *LLMPricingRule {
+	id := valuer.GenerateUUID()
+	if u.ID != nil {
+		id = *u.ID
+	}
+
 	isOverride := true
 	if u.IsOverride != nil {
 		isOverride = *u.IsOverride
@@ -222,7 +222,7 @@ func NewLLMPricingRuleFromUpdatable(u *UpdatableLLMPricingRule, orgID valuer.UUI
 	}
 
 	return &LLMPricingRule{
-		Identifiable:  types.Identifiable{ID: valuer.GenerateUUID()},
+		Identifiable:  types.Identifiable{ID: id},
 		TimeAuditable: types.TimeAuditable{CreatedAt: now, UpdatedAt: now},
 		UserAuditable: types.UserAuditable{CreatedBy: userEmail, UpdatedBy: userEmail},
 		OrgID:         orgID,
@@ -238,26 +238,6 @@ func NewLLMPricingRuleFromUpdatable(u *UpdatableLLMPricingRule, orgID valuer.UUI
 	}
 }
 
-func (r *LLMPricingRule) Update(u *UpdatableLLMPricingRule, userEmail string, now time.Time) {
-	if u.IsOverride == nil && r.IsOverride {
-		r.SyncedAt = &now
-		return
-	}
-
-	r.Model = u.Model
-	r.Provider = u.Provider
-	r.ModelPattern = StringSlice(u.ModelPattern)
-	r.Unit = u.Unit
-	r.Pricing = u.Pricing
-	if u.IsOverride != nil {
-		r.IsOverride = *u.IsOverride
-	}
-	r.Enabled = u.Enabled
-	r.SyncedAt = &now
-	r.UpdatedAt = now
-	r.UpdatedBy = userEmail
-}
-
 func ModelMatchesAnyRule(model string, rules []*LLMPricingRule) bool {
 	for _, r := range rules {
 		for _, pattern := range r.ModelPattern {

pkg/types/llmpricingruletypes/store.go

@@ -9,9 +9,8 @@ import (
 type Store interface {
 	List(ctx context.Context, orgID valuer.UUID, offset, limit int, search string, isOverride *bool) ([]*LLMPricingRule, int, error)
 	Get(ctx context.Context, orgID, id valuer.UUID) (*LLMPricingRule, error)
-	GetBySourceID(ctx context.Context, orgID, sourceID valuer.UUID) (*LLMPricingRule, error)
-	Create(ctx context.Context, rule *LLMPricingRule) error
-	Update(ctx context.Context, rule *LLMPricingRule) error
+	UpsertByID(ctx context.Context, rules []*LLMPricingRule) error
+	UpsertBySourceID(ctx context.Context, rules []*LLMPricingRule) error
 	Delete(ctx context.Context, orgID, id valuer.UUID) error
 	RunInTx(ctx context.Context, cb func(ctx context.Context) error) error
 }

tests/fixtures/llmpricingrules.py

@@ -0,0 +1,42 @@
+from http import HTTPStatus
+
+import requests
+
+from fixtures import types
+
+LLM_PRICING_RULES_URL = "/api/v1/llm_pricing_rules"
+MAX_LIST_LIMIT = 100
+
+
+def upsert_llm_pricing_rules(signoz: types.SigNoz, token: str, rules: list[dict]) -> requests.Response:
+    return requests.put(
+        signoz.self.host_configs["8080"].get(LLM_PRICING_RULES_URL),
+        headers={"Authorization": f"Bearer {token}"},
+        json={"rules": rules},
+        timeout=10,
+    )
+
+
+def list_llm_pricing_rules(signoz: types.SigNoz, token: str) -> list[dict]:
+    items: list[dict] = []
+    while True:
+        response = requests.get(
+            signoz.self.host_configs["8080"].get(f"{LLM_PRICING_RULES_URL}?offset={len(items)}&limit={MAX_LIST_LIMIT}"),
+            headers={"Authorization": f"Bearer {token}"},
+            timeout=5,
+        )
+        assert response.status_code == HTTPStatus.OK, response.text
+        page = response.json()["data"]["items"]
+        items.extend(page)
+        if len(page) < MAX_LIST_LIMIT:
+            return items
+
+
+def delete_all_llm_pricing_rules(signoz: types.SigNoz, token: str) -> None:
+    for rule in list_llm_pricing_rules(signoz, token):
+        response = requests.delete(
+            signoz.self.host_configs["8080"].get(f"{LLM_PRICING_RULES_URL}/{rule['id']}"),
+            headers={"Authorization": f"Bearer {token}"},
+            timeout=5,
+        )
+        assert response.status_code == HTTPStatus.NO_CONTENT, response.text

tests/integration/tests/llmpricingrules/01_upsert.py

@@ -0,0 +1,132 @@
+from collections.abc import Callable
+from http import HTTPStatus
+
+from fixtures import types
+from fixtures.auth import USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD
+from fixtures.llmpricingrules import (
+    delete_all_llm_pricing_rules,
+    list_llm_pricing_rules,
+    upsert_llm_pricing_rules,
+)
+
+SOURCE_A = "11111111-1111-4111-8111-111111111101"
+SOURCE_B = "11111111-1111-4111-8111-111111111102"
+
+
+def zeus_rules(price: float) -> list[dict]:
+    return [
+        {
+            "sourceId": SOURCE_A,
+            "modelName": "zeus-a",
+            "provider": "OpenAI",
+            "modelPattern": ["zeus-a*"],
+            "unit": "per_million_tokens",
+            "pricing": {"input": price, "output": price * 2},
+            "enabled": True,
+        },
+        {
+            "sourceId": SOURCE_B,
+            "modelName": "zeus-b",
+            "provider": "OpenAI",
+            "modelPattern": ["zeus-b*"],
+            "unit": "per_million_tokens",
+            "pricing": {"input": price, "output": price * 2},
+            "enabled": False,
+        },
+    ]
+
+
+def test_sync_skips_overridden_rules(
+    signoz: types.SigNoz,
+    create_user_admin: types.Operation,  # pylint: disable=unused-argument
+    get_token: Callable[[str, str], str],
+):
+    token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
+    delete_all_llm_pricing_rules(signoz, token)
+
+    # first sync inserts, keeping the disabled flag
+    assert upsert_llm_pricing_rules(signoz, token, zeus_rules(10)).status_code == HTTPStatus.NO_CONTENT
+    rules = {r["sourceId"]: r for r in list_llm_pricing_rules(signoz, token)}
+    assert set(rules) == {SOURCE_A, SOURCE_B}
+    assert rules[SOURCE_B]["enabled"] is False
+    assert all(r["isOverride"] is False for r in rules.values())
+    rule_a_id = rules[SOURCE_A]["id"]
+
+    # replay updates in place
+    assert upsert_llm_pricing_rules(signoz, token, zeus_rules(20)).status_code == HTTPStatus.NO_CONTENT
+    rules = {r["sourceId"]: r for r in list_llm_pricing_rules(signoz, token)}
+    assert rules[SOURCE_A]["id"] == rule_a_id
+    assert rules[SOURCE_A]["pricing"]["input"] == 20
+    assert rules[SOURCE_B]["pricing"]["input"] == 20
+
+    # user overrides rule a
+    override = {**zeus_rules(99)[0], "id": rule_a_id, "isOverride": True}
+    assert upsert_llm_pricing_rules(signoz, token, [override]).status_code == HTTPStatus.NO_CONTENT
+    overridden = {r["sourceId"]: r for r in list_llm_pricing_rules(signoz, token)}[SOURCE_A]
+    assert overridden["isOverride"] is True
+    assert overridden["pricing"]["input"] == 99
+
+    # sync leaves the overridden row alone, updates the other
+    assert upsert_llm_pricing_rules(signoz, token, zeus_rules(30)).status_code == HTTPStatus.NO_CONTENT
+    rules = {r["sourceId"]: r for r in list_llm_pricing_rules(signoz, token)}
+    assert rules[SOURCE_A] == overridden
+    assert rules[SOURCE_B]["pricing"]["input"] == 30
+
+    # user hands rule a back, next sync reclaims it
+    assert upsert_llm_pricing_rules(signoz, token, [{**override, "isOverride": False}]).status_code == HTTPStatus.NO_CONTENT
+    assert upsert_llm_pricing_rules(signoz, token, zeus_rules(40)).status_code == HTTPStatus.NO_CONTENT
+    rules = {r["sourceId"]: r for r in list_llm_pricing_rules(signoz, token)}
+    assert rules[SOURCE_A]["isOverride"] is False
+    assert rules[SOURCE_A]["pricing"]["input"] == 40
+
+    # user-created rule has no source id and survives a sync
+    custom = {
+        "modelName": "custom",
+        "provider": "Anthropic",
+        "modelPattern": ["custom*"],
+        "unit": "per_million_tokens",
+        "pricing": {"input": 1, "output": 2},
+        "isOverride": True,
+        "enabled": True,
+    }
+    assert upsert_llm_pricing_rules(signoz, token, [custom]).status_code == HTTPStatus.NO_CONTENT
+    created = next(r for r in list_llm_pricing_rules(signoz, token) if r["modelName"] == "custom")
+    assert created["isOverride"] is True
+    assert created.get("sourceId") is None
+    assert upsert_llm_pricing_rules(signoz, token, zeus_rules(50)).status_code == HTTPStatus.NO_CONTENT
+    assert next(r for r in list_llm_pricing_rules(signoz, token) if r["modelName"] == "custom") == created
+
+    delete_all_llm_pricing_rules(signoz, token)
+
+
+def test_bulk_sync(
+    signoz: types.SigNoz,
+    create_user_admin: types.Operation,  # pylint: disable=unused-argument
+    get_token: Callable[[str, str], str],
+):
+    token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
+    delete_all_llm_pricing_rules(signoz, token)
+
+    rules = [
+        {
+            "sourceId": f"44444444-4444-4444-8444-{i:012d}",
+            "modelName": f"bulk-{i}",
+            "provider": "OpenAI",
+            "modelPattern": [f"bulk-{i}"],
+            "unit": "per_million_tokens",
+            "pricing": {"input": 1, "output": 2},
+            "enabled": True,
+        }
+        for i in range(300)
+    ]
+    assert upsert_llm_pricing_rules(signoz, token, rules).status_code == HTTPStatus.NO_CONTENT
+    assert len(list_llm_pricing_rules(signoz, token)) == 300
+
+    for rule in rules:
+        rule["pricing"] = {"input": 5, "output": 6}
+    assert upsert_llm_pricing_rules(signoz, token, rules).status_code == HTTPStatus.NO_CONTENT
+    stored = list_llm_pricing_rules(signoz, token)
+    assert len(stored) == 300
+    assert all(r["pricing"]["input"] == 5 for r in stored)
+
+    delete_all_llm_pricing_rules(signoz, token)