From 3ffa5d96d28d50cb14d95f4222065995e94ecbf5 Mon Sep 17 00:00:00 2001 From: "lilong.129" Date: Tue, 29 Apr 2025 22:33:18 +0800 Subject: [PATCH] refactor: config llm env --- internal/version/VERSION | 2 +- uixt/ai/ai.go | 54 +++++++++++++++++++++++++++++++++-- uixt/ai/ai_ark.go | 55 ------------------------------------ uixt/ai/ai_openai.go | 58 -------------------------------------- uixt/ai/asserter.go | 44 ++++++++--------------------- uixt/ai/planner_ui_tars.go | 6 ++-- 6 files changed, 68 insertions(+), 151 deletions(-) delete mode 100644 uixt/ai/ai_ark.go delete mode 100644 uixt/ai/ai_openai.go diff --git a/internal/version/VERSION b/internal/version/VERSION index 7aea0823..fe2287ae 100644 --- a/internal/version/VERSION +++ b/internal/version/VERSION @@ -1 +1 @@ -v5.0.0-beta-2504292203 +v5.0.0-beta-2504292233 diff --git a/uixt/ai/ai.go b/uixt/ai/ai.go index e29a8931..39cd4f9e 100644 --- a/uixt/ai/ai.go +++ b/uixt/ai/ai.go @@ -4,9 +4,11 @@ import ( "context" "os" - "github.com/rs/zerolog/log" - + "github.com/cloudwego/eino-ext/components/model/openai" "github.com/httprunner/httprunner/v5/code" + "github.com/httprunner/httprunner/v5/internal/config" + "github.com/pkg/errors" + "github.com/rs/zerolog/log" ) func NewAIService(opts ...AIServiceOption) *AIServices { @@ -105,3 +107,51 @@ func (c *combinedLLMService) Call(opts *PlanningOptions) (*PlanningResult, error func (c *combinedLLMService) Assert(opts *AssertOptions) (*AssertionResponse, error) { return c.asserter.Assert(opts) } + +// LLM model config env variables +const ( + EnvOpenAIBaseURL = "OPENAI_BASE_URL" + EnvOpenAIAPIKey = "OPENAI_API_KEY" + EnvModelName = "LLM_MODEL_NAME" +) + +// GetOpenAIModelConfig get OpenAI config +func GetOpenAIModelConfig() (*openai.ChatModelConfig, error) { + if err := config.LoadEnv(); err != nil { + return nil, errors.Wrap(code.LoadEnvError, err.Error()) + } + + openaiBaseURL := os.Getenv(EnvOpenAIBaseURL) + if openaiBaseURL == "" { + return nil, errors.Wrapf(code.LLMEnvMissedError, + "env %s missed", EnvOpenAIBaseURL) + } + openaiAPIKey := os.Getenv(EnvOpenAIAPIKey) + if openaiAPIKey == "" { + return nil, errors.Wrapf(code.LLMEnvMissedError, + "env %s missed", EnvOpenAIAPIKey) + } + modelName := os.Getenv(EnvModelName) + if modelName == "" { + return nil, errors.Wrapf(code.LLMEnvMissedError, + "env %s missed", EnvModelName) + } + + temperature := float32(0.01) + modelConfig := &openai.ChatModelConfig{ + BaseURL: openaiBaseURL, + APIKey: openaiAPIKey, + Model: modelName, + Timeout: defaultTimeout, + Temperature: &temperature, + } + + // log config info + log.Info().Str("model", modelConfig.Model). + Str("baseURL", modelConfig.BaseURL). + Str("apiKey", maskAPIKey(modelConfig.APIKey)). + Str("timeout", defaultTimeout.String()). + Msg("get model config") + + return modelConfig, nil +} diff --git a/uixt/ai/ai_ark.go b/uixt/ai/ai_ark.go deleted file mode 100644 index 72b7a578..00000000 --- a/uixt/ai/ai_ark.go +++ /dev/null @@ -1,55 +0,0 @@ -package ai - -import ( - "os" - - "github.com/cloudwego/eino-ext/components/model/ark" - "github.com/httprunner/httprunner/v5/code" - "github.com/httprunner/httprunner/v5/internal/config" - "github.com/pkg/errors" - "github.com/rs/zerolog/log" -) - -const ( - EnvArkBaseURL = "ARK_BASE_URL" - EnvArkAPIKey = "ARK_API_KEY" - EnvArkModelID = "ARK_MODEL_ID" -) - -func GetArkModelConfig() (*ark.ChatModelConfig, error) { - if err := config.LoadEnv(); err != nil { - return nil, errors.Wrap(code.LoadEnvError, err.Error()) - } - - arkBaseURL := os.Getenv(EnvArkBaseURL) - arkAPIKey := os.Getenv(EnvArkAPIKey) - if arkAPIKey == "" { - return nil, errors.Wrapf(code.LLMEnvMissedError, - "env %s missed", EnvArkAPIKey) - } - modelName := os.Getenv(EnvArkModelID) - if modelName == "" { - return nil, errors.Wrapf(code.LLMEnvMissedError, - "env %s missed", EnvArkModelID) - } - timeout := defaultTimeout - - // https://www.volcengine.com/docs/82379/1494384?redirect=1 - temperature := float32(0.01) - modelConfig := &ark.ChatModelConfig{ - BaseURL: arkBaseURL, - APIKey: arkAPIKey, - Model: modelName, - Timeout: &timeout, - Temperature: &temperature, - } - - // log config info - log.Info().Str("model", modelConfig.Model). - Str("baseURL", modelConfig.BaseURL). - Str("apiKey", maskAPIKey(modelConfig.APIKey)). - Str("timeout", defaultTimeout.String()). - Msg("get model config") - - return modelConfig, nil -} diff --git a/uixt/ai/ai_openai.go b/uixt/ai/ai_openai.go deleted file mode 100644 index e275cac1..00000000 --- a/uixt/ai/ai_openai.go +++ /dev/null @@ -1,58 +0,0 @@ -package ai - -import ( - "os" - - "github.com/cloudwego/eino-ext/components/model/openai" - "github.com/httprunner/httprunner/v5/code" - "github.com/httprunner/httprunner/v5/internal/config" - "github.com/pkg/errors" - "github.com/rs/zerolog/log" -) - -const ( - EnvOpenAIBaseURL = "OPENAI_BASE_URL" - EnvOpenAIAPIKey = "OPENAI_API_KEY" - EnvModelName = "LLM_MODEL_NAME" -) - -// GetOpenAIModelConfig get OpenAI config -func GetOpenAIModelConfig() (*openai.ChatModelConfig, error) { - if err := config.LoadEnv(); err != nil { - return nil, errors.Wrap(code.LoadEnvError, err.Error()) - } - - openaiBaseURL := os.Getenv(EnvOpenAIBaseURL) - if openaiBaseURL == "" { - return nil, errors.Wrapf(code.LLMEnvMissedError, - "env %s missed", EnvOpenAIBaseURL) - } - openaiAPIKey := os.Getenv(EnvOpenAIAPIKey) - if openaiAPIKey == "" { - return nil, errors.Wrapf(code.LLMEnvMissedError, - "env %s missed", EnvOpenAIAPIKey) - } - modelName := os.Getenv(EnvModelName) - if modelName == "" { - return nil, errors.Wrapf(code.LLMEnvMissedError, - "env %s missed", EnvModelName) - } - - temperature := float32(0.01) - modelConfig := &openai.ChatModelConfig{ - BaseURL: openaiBaseURL, - APIKey: openaiAPIKey, - Model: modelName, - Timeout: defaultTimeout, - Temperature: &temperature, - } - - // log config info - log.Info().Str("model", modelConfig.Model). - Str("baseURL", modelConfig.BaseURL). - Str("apiKey", maskAPIKey(modelConfig.APIKey)). - Str("timeout", defaultTimeout.String()). - Msg("get model config") - - return modelConfig, nil -} diff --git a/uixt/ai/asserter.go b/uixt/ai/asserter.go index 750f2a32..cd01971c 100644 --- a/uixt/ai/asserter.go +++ b/uixt/ai/asserter.go @@ -7,7 +7,6 @@ import ( "strings" "time" - "github.com/cloudwego/eino-ext/components/model/ark" "github.com/cloudwego/eino-ext/components/model/openai" openai2 "github.com/cloudwego/eino-ext/libs/acl/openai" "github.com/cloudwego/eino/components/model" @@ -55,24 +54,17 @@ func NewAsserter(ctx context.Context, modelType LLMServiceType) (*Asserter, erro systemPrompt: defaultAssertionPrompt, } + config, err := GetOpenAIModelConfig() + if err != nil { + return nil, err + } + switch modelType { case LLMServiceTypeUITARS: - config, err := GetArkModelConfig() - if err != nil { - return nil, err - } asserter.systemPrompt += "\n\n" + uiTarsAssertionResponseFormat - asserter.model, err = ark.NewChatModel(ctx, config) - if err != nil { - return nil, err - } - + case LLMServiceTypeQwenVL: + asserter.systemPrompt += "\n\n" + defaultAssertionResponseJsonFormat case LLMServiceTypeGPT4Vision, LLMServiceTypeGPT4o: - config, err := GetOpenAIModelConfig() - if err != nil { - return nil, err - } - // define output format type OutputFormat struct { Thought string `json:"thought"` @@ -94,27 +86,15 @@ func NewAsserter(ctx context.Context, modelType LLMServiceType) (*Asserter, erro Strict: false, }, } - - asserter.model, err = openai.NewChatModel(ctx, config) - if err != nil { - return nil, err - } - - case LLMServiceTypeQwenVL: - config, err := GetOpenAIModelConfig() - if err != nil { - return nil, err - } - asserter.systemPrompt += "\n\n" + defaultAssertionResponseJsonFormat - asserter.model, err = openai.NewChatModel(ctx, config) - if err != nil { - return nil, err - } - default: return nil, errors.New("not supported model type for asserter") } + asserter.model, err = openai.NewChatModel(ctx, config) + if err != nil { + return nil, err + } + return asserter, nil } diff --git a/uixt/ai/planner_ui_tars.go b/uixt/ai/planner_ui_tars.go index 83184233..bbd22dfa 100644 --- a/uixt/ai/planner_ui_tars.go +++ b/uixt/ai/planner_ui_tars.go @@ -9,7 +9,7 @@ import ( "strings" "time" - "github.com/cloudwego/eino-ext/components/model/ark" + "github.com/cloudwego/eino-ext/components/model/openai" "github.com/cloudwego/eino/components/model" "github.com/cloudwego/eino/schema" "github.com/httprunner/httprunner/v5/code" @@ -20,11 +20,11 @@ import ( ) func NewUITarsPlanner(ctx context.Context) (*UITarsPlanner, error) { - config, err := GetArkModelConfig() + config, err := GetOpenAIModelConfig() if err != nil { return nil, err } - chatModel, err := ark.NewChatModel(ctx, config) + chatModel, err := openai.NewChatModel(ctx, config) if err != nil { return nil, err }