feat: allow to set a baseurl (#310)

* feat: allow to set a baseURL for OpenAI providers This allows to run local models that have a compatible OpenAI API, or for instance use a proxy. Signed-off-by: mudler <[email protected]> * feat: allow to set baseURL in the auth subcommand Signed-off-by: mudler <[email protected]> --------- Signed-off-by: mudler <[email protected]> Co-authored-by: Alex Jones <[email protected]> Co-authored-by: Matthis <[email protected]>
k8sgpt-ai · Apr 21, 2023 · cf797a6 · cf797a6
1 parent 754bf91
commit cf797a6
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 2 deletions.
diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go
@@ -29,6 +29,7 @@ import (
 var (
 	backend  string
 	password string
+	baseURL  string
 	model    string
 )
 
@@ -86,6 +87,7 @@ var AuthCmd = &cobra.Command{
 			Name:     backend,
 			Model:    model,
 			Password: password,
+			BaseURL:  baseURL,
 		}
 
 		if providerIndex == -1 {
@@ -113,4 +115,6 @@ func init() {
 	AuthCmd.Flags().StringVarP(&model, "model", "m", "gpt-3.5-turbo", "Backend AI model")
 	// add flag for password
 	AuthCmd.Flags().StringVarP(&password, "password", "p", "", "Backend AI password")
+	// add flag for url
+	AuthCmd.Flags().StringVarP(&baseURL, "baseurl", "u", "", "URL AI provider, (e.g `http://localhost:8080/v1`)")
 }
diff --git a/cmd/serve/serve.go b/cmd/serve/serve.go
@@ -47,13 +47,18 @@ var ServeCmd = &cobra.Command{
 			backend = os.Getenv("K8SGPT_BACKEND")
 			password := os.Getenv("K8SGPT_PASSWORD")
 			model := os.Getenv("K8SGPT_MODEL")
-			// If the envs are set, alocate in place to the aiProvider
+			baseURL := os.Getenv("K8SGPT_BASEURL")
+
+			// If the envs are set, allocate in place to the aiProvider
 			// else exit with error
-			if backend != "" || password != "" || model != "" {
+			envIsSet := backend != "" || password != "" || model != "" || baseURL != ""
+
+			if envIsSet {
 				aiProvider = &ai.AIProvider{
 					Name:     backend,
 					Password: password,
 					Model:    model,
+					BaseURL:  baseURL,
 				}
 
 				configAI.Providers = append(configAI.Providers, *aiProvider)

diff --git a/pkg/ai/iai.go b/pkg/ai/iai.go
@@ -27,6 +27,7 @@ type IAI interface {
 type IAIConfig interface {
 	GetPassword() string
 	GetModel() string
+	GetBaseURL() string
 }
 
 func NewClient(provider string) IAI {
@@ -48,6 +49,11 @@ type AIProvider struct {
 	Name     string `mapstructure:"name"`
 	Model    string `mapstructure:"model"`
 	Password string `mapstructure:"password"`
+	BaseURL  string `mapstructure:"base_url"`
+}
+
+func (p *AIProvider) GetBaseURL() string {
+	return p.BaseURL
 }
 
 func (p *AIProvider) GetPassword() string {

diff --git a/pkg/ai/openai.go b/pkg/ai/openai.go
@@ -37,6 +37,12 @@ type OpenAIClient struct {
 func (c *OpenAIClient) Configure(config IAIConfig, language string) error {
 	token := config.GetPassword()
 	defaultConfig := openai.DefaultConfig(token)
+
+	baseURL := config.GetBaseURL()
+	if baseURL != "" {
+		defaultConfig.BaseURL = baseURL
+	}
+
 	client := openai.NewClientWithConfig(defaultConfig)
 	if client == nil {
 		return errors.New("error creating OpenAI client")