Skip to content

Commit

Permalink
feat: change the AI command line argument to optional (#617)
Browse files Browse the repository at this point in the history
## What type of PR is this?

/kind feature

## What this PR does / why we need it:

change the AI command line argument to optional

## Which issue(s) this PR fixes:

Fixes #551
  • Loading branch information
jueli12 authored Sep 24, 2024
1 parent 0ddb27a commit dbf281e
Show file tree
Hide file tree
Showing 9 changed files with 75 additions and 45 deletions.
36 changes: 18 additions & 18 deletions cmd/karpor/app/options/ai.go
Original file line number Diff line number Diff line change
Expand Up @@ -20,12 +20,12 @@ import (
)

type AIOptions struct {
Backend string
AuthToken string
BaseURL string
Model string
Temperature float32
TopP float32
AIBackend string
AIAuthToken string
AIBaseURL string
AIModel string
AITemperature float32
AITopP float32
}

const (
Expand All @@ -45,12 +45,12 @@ func (o *AIOptions) Validate() []error {

func (o *AIOptions) ApplyTo(config *registry.ExtraConfig) error {
// Apply the AIOptions to the provided config
config.Backend = o.Backend
config.AuthToken = o.AuthToken
config.BaseURL = o.BaseURL
config.Model = o.Model
config.Temperature = o.Temperature
config.TopP = o.TopP
config.AIBackend = o.AIBackend
config.AIAuthToken = o.AIAuthToken
config.AIBaseURL = o.AIBaseURL
config.AIModel = o.AIModel
config.AITemperature = o.AITemperature
config.AITopP = o.AITopP
return nil
}

Expand All @@ -60,10 +60,10 @@ func (o *AIOptions) AddFlags(fs *pflag.FlagSet) {
return
}

fs.StringVar(&o.Backend, "ai-backend", defaultBackend, "The ai backend")
fs.StringVar(&o.AuthToken, "ai-auth-token", "", "The ai auth token")
fs.StringVar(&o.BaseURL, "ai-base-url", "", "The ai base url")
fs.StringVar(&o.Model, "ai-model", defaultModel, "The ai model")
fs.Float32Var(&o.Temperature, "ai-temperature", defaultTemperature, "The ai temperature")
fs.Float32Var(&o.TopP, "ai-top-p", defaultTopP, "The ai top-p")
fs.StringVar(&o.AIBackend, "ai-backend", defaultBackend, "The ai backend")
fs.StringVar(&o.AIAuthToken, "ai-auth-token", "", "The ai auth token")
fs.StringVar(&o.AIBaseURL, "ai-base-url", "", "The ai base url")
fs.StringVar(&o.AIModel, "ai-model", defaultModel, "The ai model")
fs.Float32Var(&o.AITemperature, "ai-temperature", defaultTemperature, "The ai temperature")
fs.Float32Var(&o.AITopP, "ai-top-p", defaultTopP, "The ai top-p")
}
10 changes: 5 additions & 5 deletions pkg/core/handler/search/search.go
Original file line number Diff line number Diff line change
Expand Up @@ -52,8 +52,6 @@ func SearchForResource(searchMgr *search.SearchManager, aiMgr *ai.AIManager, sea
ctx := r.Context()
logger := ctxutil.GetLogger(ctx)

//res, nil := aiMgr.ConvertTextToSQL("搜索集群cluster中kind为namespace的")

// Extract URL query parameters with default value
searchQuery := r.URL.Query().Get("query")
searchPattern := r.URL.Query().Get("pattern")
Expand All @@ -69,7 +67,11 @@ func SearchForResource(searchMgr *search.SearchManager, aiMgr *ai.AIManager, sea
query := searchQuery

if searchPattern == storage.NLPatternType {
//logger.Info(searchQuery)
if err := ai.CheckAIManager(aiMgr); err != nil {
handler.FailureRender(ctx, w, r, err)
return
}

res, err := aiMgr.ConvertTextToSQL(searchQuery)
if err != nil {
handler.FailureRender(ctx, w, r, err)
Expand All @@ -78,13 +80,11 @@ func SearchForResource(searchMgr *search.SearchManager, aiMgr *ai.AIManager, sea
searchQuery = res
}

//logger.Info(searchQuery)
logger.Info("Searching for resources...", "page", searchPage, "pageSize", searchPageSize)

res, err := searchStorage.Search(ctx, searchQuery, searchPattern, &storage.Pagination{Page: searchPage, PageSize: searchPageSize})
if err != nil {
if searchPattern == storage.NLPatternType {
//logger.Info(err.Error())
fixedQuery, fixErr := aiMgr.FixSQL(query, searchQuery, err.Error())
if fixErr != nil {
handler.FailureRender(ctx, w, r, err)
Expand Down
13 changes: 12 additions & 1 deletion pkg/core/manager/ai/manager.go
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,10 @@ type AIManager struct {

// NewAIManager returns a new AIManager object
func NewAIManager(c registry.ExtraConfig) (*AIManager, error) {
aiClient := ai.NewClient(c.Backend)
if c.AIAuthToken == "" {
return nil, ErrMissingAuthToken
}
aiClient := ai.NewClient(c.AIBackend)
if err := aiClient.Configure(ai.ConvertToAIConfig(c)); err != nil {
return nil, err
}
Expand All @@ -34,3 +37,11 @@ func NewAIManager(c registry.ExtraConfig) (*AIManager, error) {
client: aiClient,
}, nil
}

// CheckAIManager check if the AI manager is created
func CheckAIManager(aiMgr *AIManager) error {
if aiMgr == nil {
return ErrMissingAuthToken
}
return nil
}
21 changes: 21 additions & 0 deletions pkg/core/manager/ai/types.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
// Copyright The Karpor Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package ai

import "errors"

var (
ErrMissingAuthToken = errors.New("auth token is required")
)
9 changes: 7 additions & 2 deletions pkg/core/route/route.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,8 @@
package route

import (
"errors"
"expvar"

docs "github.com/KusionStack/karpor/api/openapispec"
authnhandler "github.com/KusionStack/karpor/pkg/core/handler/authn"
clusterhandler "github.com/KusionStack/karpor/pkg/core/handler/cluster"
Expand Down Expand Up @@ -44,6 +44,7 @@ import (
"github.com/go-chi/chi/v5/middleware"
httpswagger "github.com/swaggo/http-swagger/v2"
genericapiserver "k8s.io/apiserver/pkg/server"
"k8s.io/klog/v2"
)

// NewCoreRoute creates and configures an instance of chi.Mux with the given
Expand Down Expand Up @@ -92,7 +93,11 @@ func NewCoreRoute(
}
aiMgr, err := aimanager.NewAIManager(*extraConfig)
if err != nil {
return nil, err
if errors.Is(err, aimanager.ErrMissingAuthToken) {
klog.Warning("Auth token is empty.")
} else {
return nil, err
}
}

clusterMgr := clustermanager.NewClusterManager()
Expand Down
4 changes: 0 additions & 4 deletions pkg/infra/ai/azureopenai.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,9 +27,6 @@ type AzureAIClient struct {
}

func (c *AzureAIClient) Configure(cfg AIConfig) error {
if cfg.AuthToken == "" {
return errors.New("auth token was not provided")
}
if cfg.BaseURL == "" {
return errors.New("base url was not provided")
}
Expand All @@ -48,7 +45,6 @@ func (c *AzureAIClient) Configure(cfg AIConfig) error {
}

func (c *AzureAIClient) Generate(ctx context.Context, prompt string) (string, error) {

resp, err := c.client.CreateChatCompletion(ctx, openai.ChatCompletionRequest{
Model: c.model,
Messages: []openai.ChatCompletionMessage{
Expand Down
3 changes: 0 additions & 3 deletions pkg/infra/ai/openai.go
Original file line number Diff line number Diff line change
Expand Up @@ -28,9 +28,6 @@ type OpenAIClient struct {
}

func (c *OpenAIClient) Configure(cfg AIConfig) error {
if cfg.AuthToken == "" {
return errors.New("auth token was not provided")
}
defaultConfig := openai.DefaultConfig(cfg.AuthToken)
if cfg.BaseURL != "" {
defaultConfig.BaseURL = cfg.BaseURL
Expand Down
12 changes: 6 additions & 6 deletions pkg/infra/ai/types.go
Original file line number Diff line number Diff line change
Expand Up @@ -57,12 +57,12 @@ type AIConfig struct {

func ConvertToAIConfig(c registry.ExtraConfig) AIConfig {
return AIConfig{
Name: c.Backend,
AuthToken: c.AuthToken,
BaseURL: c.BaseURL,
Model: c.Model,
Temperature: c.Temperature,
TopP: c.TopP,
Name: c.AIBackend,
AuthToken: c.AIAuthToken,
BaseURL: c.AIBaseURL,
Model: c.AIModel,
Temperature: c.AITemperature,
TopP: c.AITopP,
}
}

Expand Down
12 changes: 6 additions & 6 deletions pkg/kubernetes/registry/types.go
Original file line number Diff line number Diff line change
Expand Up @@ -47,10 +47,10 @@ type ExtraConfig struct {
ExtendExpiration bool

// AI configs
Backend string
AuthToken string
BaseURL string
Model string
Temperature float32
TopP float32
AIBackend string
AIAuthToken string
AIBaseURL string
AIModel string
AITemperature float32
AITopP float32
}

0 comments on commit dbf281e

Please sign in to comment.