Skip to content

Commit

Permalink
feat: support natural language search for k8s resources (#612)
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:

Integrating AI into our product to reduce the consumption of human
attention is a significant milestone for us in the upcoming versions.
This PR can be a very good start to build anything related to AI.

Purpose:
- Support an AI backend service that other modules can use to access AI
language models.
- Add a new search way: natural language search based on large language
models.

## Which issue(s) this PR fixes:

<!--
*Automatically closes linked issue when PR is merged.
Usage: `Fixes #<issue number>`, or `Fixes (paste link of issue)`.
_If PR is about `failing-tests or flakes`, please post the related
issues/tests in a comment and do not use `Fixes`_*
-->

Fixes #452

---------

Co-authored-by: 玦离 <[email protected]>
Co-authored-by: hai-tian <[email protected]>
  • Loading branch information
3 people authored Nov 27, 2024
1 parent ab3196b commit 2b45db9
Show file tree
Hide file tree
Showing 30 changed files with 1,456 additions and 436 deletions.
69 changes: 69 additions & 0 deletions cmd/karpor/app/options/ai.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,69 @@
// 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 options

import (
"github.com/KusionStack/karpor/pkg/kubernetes/registry"
"github.com/spf13/pflag"
)

type AIOptions struct {
AIBackend string
AIAuthToken string
AIBaseURL string
AIModel string
AITemperature float32
AITopP float32
}

const (
defaultBackend = "openai"
defaultModel = "gpt-3.5-turbo"
defaultTemperature = 1
defaultTopP = 1
)

func NewAIOptions() *AIOptions {
return &AIOptions{}
}

func (o *AIOptions) Validate() []error {
return nil
}

func (o *AIOptions) ApplyTo(config *registry.ExtraConfig) error {
// Apply the AIOptions to the provided config
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
}

// AddFlags adds flags for a specific Option to the specified FlagSet
func (o *AIOptions) AddFlags(fs *pflag.FlagSet) {
if o == nil {
return
}

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: 10 additions & 0 deletions cmd/karpor/app/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ type Options struct {
RecommendedOptions *options.RecommendedOptions
SearchStorageOptions *options.SearchStorageOptions
CoreOptions *options.CoreOptions
AIOptions *options.AIOptions

StdOut io.Writer
StdErr io.Writer
Expand All @@ -74,6 +75,7 @@ func NewOptions(out, errOut io.Writer) (*Options, error) {
),
SearchStorageOptions: options.NewSearchStorageOptions(),
CoreOptions: options.NewCoreOptions(),
AIOptions: options.NewAIOptions(),
StdOut: out,
StdErr: errOut,
}
Expand Down Expand Up @@ -105,6 +107,9 @@ func NewServerCommand(ctx context.Context) *cobra.Command {
expvar.Publish("StorageOptions", expvar.Func(func() interface{} {
return o.SearchStorageOptions
}))
expvar.Publish("AIOptions", expvar.Func(func() interface{} {
return o.AIOptions
}))
expvar.Publish("Version", expvar.Func(func() interface{} {
return version.GetVersion()
}))
Expand Down Expand Up @@ -140,13 +145,15 @@ func (o *Options) AddFlags(fs *pflag.FlagSet) {
o.RecommendedOptions.AddFlags(fs)
o.SearchStorageOptions.AddFlags(fs)
o.CoreOptions.AddFlags(fs)
o.AIOptions.AddFlags(fs)
}

// Validate validates Options
func (o *Options) Validate(args []string) error {
errors := []error{}
errors = append(errors, o.RecommendedOptions.Validate()...)
errors = append(errors, o.SearchStorageOptions.Validate()...)
errors = append(errors, o.AIOptions.Validate()...)
return utilerrors.NewAggregate(errors)
}

Expand Down Expand Up @@ -208,6 +215,9 @@ func (o *Options) Config() (*server.Config, error) {
if err := o.CoreOptions.ApplyTo(config.ExtraConfig); err != nil {
return nil, err
}
if err := o.AIOptions.ApplyTo(config.ExtraConfig); err != nil {
return nil, err
}

config.GenericConfig.BuildHandlerChainFunc = func(handler http.Handler, c *genericapiserver.Config) http.Handler {
handler = genericapiserver.DefaultBuildHandlerChain(handler, c)
Expand Down
4 changes: 3 additions & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,12 @@ require (
github.com/go-chi/render v1.0.3
github.com/go-logr/logr v1.2.3
github.com/google/gofuzz v1.2.0
github.com/google/uuid v1.3.0
github.com/google/uuid v1.4.0
github.com/hupe1980/go-huggingface v0.0.15
github.com/mitchellh/hashstructure/v2 v2.0.2
github.com/oliveagle/jsonpath v0.0.0-20180606110733-2e52cf6e6852
github.com/pkg/errors v0.9.1
github.com/sashabaranov/go-openai v1.27.0
github.com/spf13/cobra v1.6.1
github.com/spf13/pflag v1.0.5
github.com/stretchr/testify v1.8.4
Expand Down
8 changes: 6 additions & 2 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -254,8 +254,8 @@ github.com/google/pprof v0.0.0-20200708004538-1a94d8640e99/go.mod h1:ZgVRPoUq/hf
github.com/google/renameio v0.1.0/go.mod h1:KWCgfxg9yswjAJkECMjeO8J8rahYeXnNhOm40UhjYkI=
github.com/google/uuid v1.1.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/google/uuid v1.4.0 h1:MtMxsa51/r9yyhkyLsVeVt0B+BGQZzpQiTQ4eHZ8bc4=
github.com/google/uuid v1.4.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/googleapis/gax-go/v2 v2.0.4/go.mod h1:0Wqv26UfaUD9n4G6kQubkQ+KchISgw+vpHVxEJEs9eg=
github.com/googleapis/gax-go/v2 v2.0.5/go.mod h1:DWXyrwAJ9X0FpwwEdw+IPEYBICEFu5mhpdKc/us6bOk=
github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1 h1:EGx4pi6eqNxGaHF6qqu48+N2wcFQ5qg5FXgOdqsJ5d8=
Expand All @@ -275,6 +275,8 @@ github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ
github.com/huandu/xstrings v1.3.1/go.mod h1:y5/lhBue+AyNmUVz9RLU9xbLR0o4KIIExikq4ovT0aE=
github.com/huandu/xstrings v1.3.2 h1:L18LIDzqlW6xN2rEkpdV8+oL/IXWJ1APd+vsdYy4Wdw=
github.com/huandu/xstrings v1.3.2/go.mod h1:y5/lhBue+AyNmUVz9RLU9xbLR0o4KIIExikq4ovT0aE=
github.com/hupe1980/go-huggingface v0.0.15 h1:tTWmUGGunC/BYz4hrwS8SSVtMYVYjceG2uhL8HxeXvw=
github.com/hupe1980/go-huggingface v0.0.15/go.mod h1:IRvsik3+b9BJyw9hCfw1arI6gDObcVto1UA8f3kt8mM=
github.com/ianlancetaylor/demangle v0.0.0-20181102032728-5e5cf60278f6/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
github.com/imdario/mergo v0.3.11/go.mod h1:jmQim1M+e3UYxmgPu/WyfjB3N3VflVyUjjjwH0dnCYA=
github.com/imdario/mergo v0.3.12 h1:b6R2BslTbIEToALKP7LxUvijTsNI9TAe80pLWN2g/HU=
Expand Down Expand Up @@ -395,6 +397,8 @@ github.com/rogpeppe/go-internal v1.6.1 h1:/FiVV8dS/e+YqF2JvO3yXRFbBLTIuSDkuC7aBO
github.com/rogpeppe/go-internal v1.6.1/go.mod h1:xXDCJY+GAPziupqXw64V24skbSoqbTEfhy4qGm1nDQc=
github.com/russross/blackfriday/v2 v2.1.0 h1:JIOH55/0cWyOuilr9/qlrm0BSXldqnqwMsf35Ld67mk=
github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
github.com/sashabaranov/go-openai v1.27.0 h1:L3hO6650YUbKrbGUC6yCjsUluhKZ9h1/jcgbTItI8Mo=
github.com/sashabaranov/go-openai v1.27.0/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
github.com/shopspring/decimal v1.2.0 h1:abSATXmQEYyShuxI4/vyW3tV1MrKAJzCZ/0zLUXYbsQ=
github.com/shopspring/decimal v1.2.0/go.mod h1:DKyhrW/HYNuLGql+MJL6WCR6knT2jwCFRcu2hWCYk4o=
github.com/sirupsen/logrus v1.2.0/go.mod h1:LxeOpSwHxABJmUn/MG1IvRgCAasNZTLOkJPxbbu5VWo=
Expand Down
40 changes: 38 additions & 2 deletions pkg/core/handler/search/search.go
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ import (
"strconv"

"github.com/KusionStack/karpor/pkg/core/handler"
"github.com/KusionStack/karpor/pkg/core/manager/ai"
"github.com/KusionStack/karpor/pkg/core/manager/search"
"github.com/KusionStack/karpor/pkg/infra/search/storage"
"github.com/KusionStack/karpor/pkg/util/ctxutil"
Expand All @@ -34,7 +35,7 @@ import (
// @Tags search
// @Produce json
// @Param query query string true "The query to use for search. Required"
// @Param pattern query string true "The search pattern. Can be either sql or dsl. Required"
// @Param pattern query string true "The search pattern. Can be either sql, dsl or nl. Required"
// @Param pageSize query string false "The size of the page. Default to 10"
// @Param page query string false "The current page to fetch. Default to 1"
// @Success 200 {array} runtime.Object "Array of runtime.Object"
Expand All @@ -45,7 +46,7 @@ import (
// @Failure 429 {string} string "Too Many Requests"
// @Failure 500 {string} string "Internal Server Error"
// @Router /rest-api/v1/search [get]
func SearchForResource(searchMgr *search.SearchManager, searchStorage storage.SearchStorage) http.HandlerFunc {
func SearchForResource(searchMgr *search.SearchManager, aiMgr *ai.AIManager, searchStorage storage.SearchStorage) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
// Extract the context and logger from the request.
ctx := r.Context()
Expand All @@ -71,9 +72,43 @@ func SearchForResource(searchMgr *search.SearchManager, searchStorage storage.Se
searchPage = 1
}

query := searchQuery

if searchPattern == storage.NLPatternType {
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)
return
}
searchQuery = res
}

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 {
fixedQuery, fixErr := aiMgr.FixSQL(query, searchQuery, err.Error())
if fixErr != nil {
handler.FailureRender(ctx, w, r, err)
return
}
searchQuery = fixedQuery
res, err = searchStorage.Search(ctx, searchQuery, searchPattern, &storage.Pagination{Page: searchPage, PageSize: searchPageSize})
if err != nil {
handler.FailureRender(ctx, w, r, err)
return
}
} else {
handler.FailureRender(ctx, w, r, err)
return
}
}

if err != nil {
handler.FailureRender(ctx, w, r, err)
return
Expand All @@ -97,6 +132,7 @@ func SearchForResource(searchMgr *search.SearchManager, searchStorage storage.Se
Deleted: res.Deleted,
})
}
rt.SQLQuery = searchQuery
rt.Total = res.Total
rt.CurrentPage = searchPage
rt.PageSize = searchPageSize
Expand Down
47 changes: 47 additions & 0 deletions pkg/core/manager/ai/manager.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
// 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 (
"github.com/KusionStack/karpor/pkg/infra/ai"
"github.com/KusionStack/karpor/pkg/kubernetes/registry"
)

type AIManager struct {
client ai.AIProvider
}

// NewAIManager returns a new AIManager object
func NewAIManager(c registry.ExtraConfig) (*AIManager, error) {
if c.AIAuthToken == "" {
return nil, ErrMissingAuthToken
}
aiClient := ai.NewClient(c.AIBackend)
if err := aiClient.Configure(ai.ConvertToAIConfig(c)); err != nil {
return nil, err
}

return &AIManager{
client: aiClient,
}, nil
}

// CheckAIManager check if the AI manager is created
func CheckAIManager(aiMgr *AIManager) error {
if aiMgr == nil {
return ErrMissingAuthToken
}
return nil
}
47 changes: 47 additions & 0 deletions pkg/core/manager/ai/search.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
// 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 (
"context"
"fmt"

"github.com/KusionStack/karpor/pkg/infra/ai"
)

// ConvertTextToSQL converts natural language text to an SQL query
func (a *AIManager) ConvertTextToSQL(query string) (string, error) {
servicePrompt := ai.ServicePromptMap[ai.Text2sqlType]
prompt := fmt.Sprintf(servicePrompt, query)
res, err := a.client.Generate(context.Background(), prompt)
if err != nil {
return "", err
}
if IsInvalidQuery(res) {
return "", ErrInvalidQuery
}
return ExtractSelectSQL(res), nil
}

// FixSQL fix the error SQL
func (a *AIManager) FixSQL(sql string, query string, sqlErr string) (string, error) {
servicePrompt := ai.ServicePromptMap[ai.SQLFixType]
prompt := fmt.Sprintf(servicePrompt, query, sql, sqlErr)
res, err := a.client.Generate(context.Background(), prompt)
if err != nil {
return "", err
}
return ExtractSelectSQL(res), nil
}
22 changes: 22 additions & 0 deletions pkg/core/manager/ai/types.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
// 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")
ErrInvalidQuery = errors.New("query is invalid")
)
32 changes: 32 additions & 0 deletions pkg/core/manager/ai/util.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
// 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 (
"regexp"
"strings"
)

// IsInvalidQuery check if the query is invalid
func IsInvalidQuery(sql string) bool {
return strings.Contains(strings.ToLower(sql), "error")
}

// ExtractSelectSQL extracts SQL statements that start with "SELECT * FROM"
func ExtractSelectSQL(sql string) string {
res := regexp.MustCompile(`(?i)SELECT \* FROM [^;]+`)
match := res.FindString(sql)
return match
}
Loading

0 comments on commit 2b45db9

Please sign in to comment.