Skip to content

Commit 65b7bd1

Browse files
csg-pr-botDev AgentRader
committed
update(aigateway): refresh user balance cache after recharge or consuming (#968)
* update(aigateway): refresh user balance cache after recharge or consuming * fix build error --------- Co-authored-by: Dev Agent <dev-agent@example.com> Co-authored-by: rader <ld_leida@hotmail.com>
1 parent 2ea3d03 commit 65b7bd1

File tree

8 files changed

+86
-83
lines changed

8 files changed

+86
-83
lines changed

‎_mocks/opencsg.com/csghub-server/aigateway/component/mock_OpenAIComponent.go‎

Lines changed: 11 additions & 10 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎_mocks/opencsg.com/csghub-server/builder/store/cache/mock_RedisClient.go‎

Lines changed: 59 additions & 59 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎aigateway/component/openai.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ type OpenAIComponent interface {
3030
ListModels(c context.Context, user string, req types.ListModelsReq) (types.ModelList, error)
3131
GetModelByID(c context.Context, username, modelID string) (*types.Model, error)
3232
RecordUsage(c context.Context, userUUID string, model *types.Model, tokenCounter token.Counter, sceneValue string) error
33-
CheckBalance(ctx context.Context, username string) error
33+
CheckBalance(ctx context.Context, username, userUUID string) error
3434
}
3535

3636
type openaiComponentImpl struct {

‎aigateway/component/openai_ce.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,6 @@ func parseScene(sceneValue string) common_types.SceneType {
4545
return common_types.SceneModelServerless
4646
}
4747

48-
func (e *extendOpenai) CheckBalance(ctx context.Context, username string) error {
48+
func (e *extendOpenai) CheckBalance(ctx context.Context, username, userUUID string) error {
4949
return nil
5050
}

‎aigateway/handler/openai.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -337,7 +337,7 @@ func (h *OpenAIHandlerImpl) Chat(c *gin.Context) {
337337

338338
sceneValue := c.Request.Header.Get(commonType.SceneHeaderKey)
339339
// Check balance before processing request
340-
if err := h.openaiComponent.CheckBalance(c.Request.Context(), username); err != nil {
340+
if err := h.openaiComponent.CheckBalance(c.Request.Context(), username, userUUID); err != nil {
341341
h.handleInsufficientBalance(c, chatReq.Stream, username, modelID, err)
342342
return
343343
}
@@ -510,7 +510,7 @@ func (h *OpenAIHandlerImpl) GenerateImage(c *gin.Context) {
510510
}
511511

512512
sceneValue := c.Request.Header.Get(commonType.SceneHeaderKey)
513-
if err := h.openaiComponent.CheckBalance(ctx, username); err != nil {
513+
if err := h.openaiComponent.CheckBalance(ctx, username, userUUID); err != nil {
514514
h.handleInsufficientBalance(c, false, username, modelID, err)
515515
return
516516
}
@@ -690,7 +690,7 @@ func (h *OpenAIHandlerImpl) Embedding(c *gin.Context) {
690690

691691
sceneValue := c.Request.Header.Get(commonType.SceneHeaderKey)
692692
// Check balance before processing request
693-
if err := h.openaiComponent.CheckBalance(c.Request.Context(), username); err != nil {
693+
if err := h.openaiComponent.CheckBalance(c.Request.Context(), username, userUUID); err != nil {
694694
h.handleInsufficientBalance(c, false, username, modelID, err)
695695
return
696696
}

‎aigateway/handler/openai_test.go‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -393,7 +393,7 @@ func TestOpenAIHandler_Chat(t *testing.T) {
393393
ClusterID: "test-cls",
394394
}, nil)
395395
tester.mocks.openAIComp.EXPECT().GetModelByID(mock.Anything, "testuser", "model1:svc1").Return(model, nil)
396-
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser").Return(nil)
396+
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser", "testuuid").Return(nil)
397397
expectReq := ChatCompletionRequest{}
398398
_ = json.Unmarshal(body, &expectReq)
399399
tester.mocks.moderationComp.EXPECT().CheckChatPrompts(mock.Anything, expectReq.Messages, "testuuid:"+model.ID).
@@ -438,7 +438,7 @@ func TestOpenAIHandler_Chat(t *testing.T) {
438438
ClusterID: "test-cls",
439439
}, nil)
440440
tester.mocks.openAIComp.EXPECT().GetModelByID(mock.Anything, "testuser", "model1:svc1").Return(model, nil)
441-
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser").Return(nil)
441+
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser", "testuuid").Return(nil)
442442
expectReq := ChatCompletionRequest{}
443443
_ = json.Unmarshal(body, &expectReq)
444444
tester.mocks.moderationComp.EXPECT().CheckChatPrompts(mock.Anything, expectReq.Messages, "testuuid:"+model.ID).
@@ -483,7 +483,7 @@ func TestOpenAIHandler_Chat(t *testing.T) {
483483
ClusterID: "test-cls",
484484
}, nil)
485485
tester.mocks.openAIComp.EXPECT().GetModelByID(mock.Anything, "testuser", "model1:svc1").Return(model, nil)
486-
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser").Return(nil)
486+
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser", "testuuid").Return(nil)
487487
expectReq := ChatCompletionRequest{}
488488
_ = json.Unmarshal(body, &expectReq)
489489
tester.mocks.moderationComp.EXPECT().CheckChatPrompts(mock.Anything, expectReq.Messages, "testuuid:"+model.ID).
@@ -545,7 +545,7 @@ func TestOpenAIHandler_Chat(t *testing.T) {
545545
ClusterID: "test-cls",
546546
}, nil)
547547
tester.mocks.openAIComp.EXPECT().GetModelByID(mock.Anything, "testuser", "model1:svc1").Return(model, nil)
548-
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser").Return(nil)
548+
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser", "testuuid").Return(nil)
549549
expectReq := ChatCompletionRequest{}
550550
_ = json.Unmarshal(body, &expectReq)
551551
tester.mocks.moderationComp.EXPECT().CheckChatPrompts(mock.Anything, expectReq.Messages, "testuuid:"+model.ID).
@@ -600,7 +600,7 @@ func TestOpenAIHandler_Chat(t *testing.T) {
600600
Endpoint: testServer.URL,
601601
}
602602
tester.mocks.openAIComp.EXPECT().GetModelByID(mock.Anything, "testuser", "external-model-id").Return(model, nil)
603-
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser").Return(nil)
603+
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser", "testuuid").Return(nil)
604604
expectReq := ChatCompletionRequest{}
605605
_ = json.Unmarshal(body, &expectReq)
606606
tester.mocks.moderationComp.EXPECT().CheckChatPrompts(mock.Anything, expectReq.Messages, "testuuid:"+model.ID).
@@ -802,7 +802,7 @@ func TestOpenAIHandler_Embedding(t *testing.T) {
802802
Return(tokenCounter).Once()
803803
tester.mocks.openAIComp.EXPECT().GetModelByID(mock.Anything, "testuser", "model1").
804804
Return(model, nil)
805-
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser").Return(nil)
805+
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser", "testuuid").Return(nil)
806806
tester.mocks.openAIComp.EXPECT().RecordUsage(mock.Anything, "testuuid", model, mock.Anything, mock.Anything).RunAndReturn(
807807
func(ctx context.Context, userID string, model *types.Model, counter token.Counter, sceneValue string) error {
808808
wg.Done()
@@ -939,7 +939,7 @@ func TestOpenAIHandler_GenerateImage(t *testing.T) {
939939
Endpoint: "https://api.example.com/images/generations",
940940
}
941941
tester.mocks.openAIComp.EXPECT().GetModelByID(mock.Anything, "testuser", "test-model").Return(model, nil)
942-
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser").Return(nil)
942+
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser", "testuuid").Return(nil)
943943
tester.mocks.moderationComp.EXPECT().CheckImagePrompts(mock.Anything, "sensitive prompt", "testuuid").Return(&rpc.CheckResult{IsSensitive: true}, nil)
944944

945945
tester.handler.GenerateImage(c)
@@ -972,7 +972,7 @@ func TestOpenAIHandler_GenerateImage(t *testing.T) {
972972
Endpoint: "https://api.example.com/images/generations",
973973
}
974974
tester.mocks.openAIComp.EXPECT().GetModelByID(mock.Anything, "testuser", "test-model").Return(model, nil)
975-
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser").Return(nil)
975+
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser", "testuuid").Return(nil)
976976
tester.mocks.moderationComp.EXPECT().CheckImagePrompts(mock.Anything, "test prompt", "testuuid").Return(nil, errors.New("moderation service error"))
977977

978978
tester.handler.GenerateImage(c)

‎builder/store/database/agent_knowledge_base_test.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -388,6 +388,7 @@ func TestAgentKnowledgeBaseStore_List_OrderByUpdatedAt(t *testing.T) {
388388

389389
// Create knowledge bases
390390
kb1 := &database.AgentKnowledgeBase{
391+
ID: 1,
391392
UserUUID: userUUID,
392393
Name: "First Knowledge Base",
393394
Description: "First description",
@@ -398,6 +399,7 @@ func TestAgentKnowledgeBaseStore_List_OrderByUpdatedAt(t *testing.T) {
398399
require.NoError(t, err)
399400

400401
kb2 := &database.AgentKnowledgeBase{
402+
ID: 2,
401403
UserUUID: userUUID,
402404
Name: "Second Knowledge Base",
403405
Description: "Second description",

‎common/config/config.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -222,7 +222,7 @@ type Config struct {
222222
ExpiredPresentCronExpression string `env:"OPENCSG_ACCOUNTING_EXPIRED_PRESENT_CRON_EXPRESSION" default:"0 0 * * *"`
223223
ThresholdOfStopDeploy int `env:"OPENCSG_ACCOUNTING_THRESHOLD_OF_STOP_DEPLOY" default:"5000"`
224224
ThresholdOfStopLLMInference int `env:"OPENCSG_ACCOUNTING_THRESHOLD_OF_STOP_LLM_INFERENCE" default:"5000"`
225-
LLMBalanceCheckCacheTTL int `env:"OPENCSG_ACCOUNTING_LLM_BALANCE_CHECK_CACHE_TTL" default:"86400"`
225+
LLMBalanceCheckCacheTTL int `env:"OPENCSG_ACCOUNTING_LLM_BALANCE_CHECK_CACHE_TTL" default:"120"`
226226
}
227227

228228
User struct {

0 commit comments

Comments
 (0)