Skip to content

Commit bdd0c5a

Browse files
csg-pr-botDev AgentRader
authored and
Jun Luo
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 ede0e44 commit bdd0c5a

File tree

6 files changed

+82
-79
lines changed

6 files changed

+82
-79
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_ce.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ 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
}
5151

‎aigateway/handler/openai_test.go‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -462,7 +462,7 @@ func TestOpenAIHandler_Chat(t *testing.T) {
462462
ClusterID: "test-cls",
463463
}, nil)
464464
tester.mocks.openAIComp.EXPECT().GetModelByID(mock.Anything, "testuser", "model1:svc1").Return(model, nil)
465-
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser").Return(nil)
465+
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser", "testuuid").Return(nil)
466466
expectReq := ChatCompletionRequest{}
467467
_ = json.Unmarshal(body, &expectReq)
468468
tester.mocks.moderationComp.EXPECT().CheckChatPrompts(mock.Anything, expectReq.Messages, "testuuid:"+model.ID).
@@ -507,7 +507,7 @@ func TestOpenAIHandler_Chat(t *testing.T) {
507507
ClusterID: "test-cls",
508508
}, nil)
509509
tester.mocks.openAIComp.EXPECT().GetModelByID(mock.Anything, "testuser", "model1:svc1").Return(model, nil)
510-
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser").Return(nil)
510+
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser", "testuuid").Return(nil)
511511
expectReq := ChatCompletionRequest{}
512512
_ = json.Unmarshal(body, &expectReq)
513513
tester.mocks.moderationComp.EXPECT().CheckChatPrompts(mock.Anything, expectReq.Messages, "testuuid:"+model.ID).
@@ -552,7 +552,7 @@ func TestOpenAIHandler_Chat(t *testing.T) {
552552
ClusterID: "test-cls",
553553
}, nil)
554554
tester.mocks.openAIComp.EXPECT().GetModelByID(mock.Anything, "testuser", "model1:svc1").Return(model, nil)
555-
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser").Return(nil)
555+
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser", "testuuid").Return(nil)
556556
expectReq := ChatCompletionRequest{}
557557
_ = json.Unmarshal(body, &expectReq)
558558
tester.mocks.moderationComp.EXPECT().CheckChatPrompts(mock.Anything, expectReq.Messages, "testuuid:"+model.ID).
@@ -614,7 +614,7 @@ func TestOpenAIHandler_Chat(t *testing.T) {
614614
ClusterID: "test-cls",
615615
}, nil)
616616
tester.mocks.openAIComp.EXPECT().GetModelByID(mock.Anything, "testuser", "model1:svc1").Return(model, nil)
617-
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser").Return(nil)
617+
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser", "testuuid").Return(nil)
618618
expectReq := ChatCompletionRequest{}
619619
_ = json.Unmarshal(body, &expectReq)
620620
tester.mocks.moderationComp.EXPECT().CheckChatPrompts(mock.Anything, expectReq.Messages, "testuuid:"+model.ID).
@@ -669,7 +669,7 @@ func TestOpenAIHandler_Chat(t *testing.T) {
669669
Endpoint: testServer.URL,
670670
}
671671
tester.mocks.openAIComp.EXPECT().GetModelByID(mock.Anything, "testuser", "external-model-id").Return(model, nil)
672-
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser").Return(nil)
672+
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser", "testuuid").Return(nil)
673673
expectReq := ChatCompletionRequest{}
674674
_ = json.Unmarshal(body, &expectReq)
675675
tester.mocks.moderationComp.EXPECT().CheckChatPrompts(mock.Anything, expectReq.Messages, "testuuid:"+model.ID).
@@ -871,7 +871,7 @@ func TestOpenAIHandler_Embedding(t *testing.T) {
871871
Return(tokenCounter).Once()
872872
tester.mocks.openAIComp.EXPECT().GetModelByID(mock.Anything, "testuser", "model1").
873873
Return(model, nil)
874-
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser").Return(nil)
874+
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser", "testuuid").Return(nil)
875875
tester.mocks.openAIComp.EXPECT().RecordUsage(mock.Anything, "testuuid", model, mock.Anything, mock.Anything).RunAndReturn(
876876
func(ctx context.Context, userID string, model *types.Model, counter token.Counter, sceneValue string) error {
877877
wg.Done()
@@ -1008,7 +1008,7 @@ func TestOpenAIHandler_GenerateImage(t *testing.T) {
10081008
Endpoint: "https://api.example.com/images/generations",
10091009
}
10101010
tester.mocks.openAIComp.EXPECT().GetModelByID(mock.Anything, "testuser", "test-model").Return(model, nil)
1011-
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser").Return(nil)
1011+
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser", "testuuid").Return(nil)
10121012
tester.mocks.moderationComp.EXPECT().CheckImagePrompts(mock.Anything, "sensitive prompt", "testuuid").Return(&rpc.CheckResult{IsSensitive: true}, nil)
10131013

10141014
tester.handler.GenerateImage(c)
@@ -1041,7 +1041,7 @@ func TestOpenAIHandler_GenerateImage(t *testing.T) {
10411041
Endpoint: "https://api.example.com/images/generations",
10421042
}
10431043
tester.mocks.openAIComp.EXPECT().GetModelByID(mock.Anything, "testuser", "test-model").Return(model, nil)
1044-
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser").Return(nil)
1044+
tester.mocks.openAIComp.EXPECT().CheckBalance(mock.Anything, "testuser", "testuuid").Return(nil)
10451045
tester.mocks.moderationComp.EXPECT().CheckImagePrompts(mock.Anything, "test prompt", "testuuid").Return(nil, errors.New("moderation service error"))
10461046

10471047
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)