mirror of
https://github.com/SigNoz/signoz.git
synced 2026-09-22 11:20:43 +01:00
Compare commits
4 Commits
main
...
issue_6107
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
1f0918080a | ||
|
|
aa63d9ab1d | ||
|
|
e6572c10cf | ||
|
|
6e18a7937f |
1
.github/workflows/integrationci.yaml
vendored
1
.github/workflows/integrationci.yaml
vendored
@@ -47,6 +47,7 @@ jobs:
|
||||
- dashboard
|
||||
- ingestionkeys
|
||||
- inframonitoring
|
||||
- llmpricingrules
|
||||
- logspipelines
|
||||
- passwordauthn
|
||||
- preference
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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 = (
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -159,14 +159,12 @@ func (provider *provider) RotateToken(ctx context.Context, accessToken string, r
|
||||
var rotatedToken *authtypes.Token
|
||||
|
||||
if err := provider.tokenStore.GetOrUpdateByAccessTokenOrPrevAccessToken(ctx, accessToken, func(ctx context.Context, token *authtypes.StorableToken) error {
|
||||
currentAccessToken := token.AccessToken
|
||||
|
||||
if err := token.Rotate(accessToken, refreshToken, provider.config.Rotation.Duration, provider.config.Lifetime.Idle, provider.config.Lifetime.Max); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// If the token passed the Rotate method and is the same as the stored token, return the same token.
|
||||
if token.AccessToken == currentAccessToken {
|
||||
// If the token passed the Rotate method and is the same as the input token, return the same token.
|
||||
if token.AccessToken == accessToken && token.RefreshToken == refreshToken {
|
||||
rotatedToken = token
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -164,7 +164,7 @@ func (typ *Token) IsRotationRequired(rotationInterval time.Duration) error {
|
||||
func (typ *Token) Rotate(accessTokenOrPrevAccessToken string, refreshTokenOrPrevRefreshToken string, rotationDuration time.Duration, idleDuration time.Duration, maxDuration time.Duration) error {
|
||||
if typ.PrevAccessToken == accessTokenOrPrevAccessToken && typ.PrevRefreshToken == refreshTokenOrPrevRefreshToken {
|
||||
// If the token has been rotated within the rotation duration, do nothing and return the same token.
|
||||
if !typ.RotatedAt.IsZero() && typ.RotatedAt.After(time.Now().Add(-rotationDuration)) {
|
||||
if !typ.RotatedAt.IsZero() && typ.RotatedAt.Before(time.Now().Add(-rotationDuration)) {
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
@@ -1,62 +0,0 @@
|
||||
package authtypes
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/SigNoz/signoz/pkg/errors"
|
||||
"github.com/SigNoz/signoz/pkg/valuer"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
const (
|
||||
testRotationDuration = 60 * time.Second
|
||||
testIdleDuration = 7 * 24 * time.Hour
|
||||
testMaxDuration = 30 * 24 * time.Hour
|
||||
)
|
||||
|
||||
func newRotatedToken(t *testing.T, rotatedAt time.Time) (*Token, string, string) {
|
||||
t.Helper()
|
||||
|
||||
token, err := NewToken(map[string]string{}, valuer.GenerateUUID())
|
||||
require.NoError(t, err)
|
||||
|
||||
prevAccessToken, prevRefreshToken := token.AccessToken, token.RefreshToken
|
||||
require.NoError(t, token.Rotate(prevAccessToken, prevRefreshToken, testRotationDuration, testIdleDuration, testMaxDuration))
|
||||
token.RotatedAt = rotatedAt
|
||||
|
||||
return token, prevAccessToken, prevRefreshToken
|
||||
}
|
||||
|
||||
func TestTokenRotatePrevPair(t *testing.T) {
|
||||
testCases := []struct {
|
||||
name string
|
||||
rotatedAt time.Time
|
||||
wantErr bool
|
||||
}{
|
||||
{name: "InsideRotationDuration", rotatedAt: time.Now().Add(-testRotationDuration / 2), wantErr: false},
|
||||
{name: "OutsideRotationDuration", rotatedAt: time.Now().Add(-2 * testRotationDuration), wantErr: true},
|
||||
}
|
||||
|
||||
for _, testCase := range testCases {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
token, prevAccessToken, prevRefreshToken := newRotatedToken(t, testCase.rotatedAt)
|
||||
accessToken, refreshToken := token.AccessToken, token.RefreshToken
|
||||
|
||||
err := token.Rotate(prevAccessToken, prevRefreshToken, testRotationDuration, testIdleDuration, testMaxDuration)
|
||||
if testCase.wantErr {
|
||||
require.Error(t, err)
|
||||
assert.True(t, errors.Ast(err, errors.TypeUnauthenticated))
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
assert.Equal(t, accessToken, token.AccessToken)
|
||||
assert.Equal(t, refreshToken, token.RefreshToken)
|
||||
assert.Equal(t, prevAccessToken, token.PrevAccessToken)
|
||||
assert.Equal(t, prevRefreshToken, token.PrevRefreshToken)
|
||||
assert.Equal(t, testCase.rotatedAt, token.RotatedAt)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
42
tests/fixtures/llmpricingrules.py
vendored
Normal file
42
tests/fixtures/llmpricingrules.py
vendored
Normal file
@@ -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
|
||||
132
tests/integration/tests/llmpricingrules/01_upsert.py
Normal file
132
tests/integration/tests/llmpricingrules/01_upsert.py
Normal file
@@ -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)
|
||||
Reference in New Issue
Block a user