feat: add configuration interface to support customer providers

Signed-off-by: Aris Boutselis <aris.boutselis@senseon.io>
This commit is contained in:
Aris Boutselis 2023-04-20 16:26:12 +01:00
parent bd4ab0e589
commit 84a3cc05fb
No known key found for this signature in database
GPG Key ID: 702CD37DB0D0FF51
5 changed files with 31 additions and 14 deletions

View File

@ -5,12 +5,17 @@ import (
) )
type IAI interface { type IAI interface {
Configure(token string, model string, language string) error Configure(config IAIConfig, language string) error
GetCompletion(ctx context.Context, prompt string) (string, error) GetCompletion(ctx context.Context, prompt string) (string, error)
Parse(ctx context.Context, prompt []string, nocache bool) (string, error) Parse(ctx context.Context, prompt []string, nocache bool) (string, error)
GetName() string GetName() string
} }
type IAIConfig interface {
GetPassword() string
GetModel() string
}
func NewClient(provider string) IAI { func NewClient(provider string) IAI {
switch provider { switch provider {
case "openai": case "openai":
@ -31,3 +36,11 @@ type AIProvider struct {
Model string `mapstructure:"model"` Model string `mapstructure:"model"`
Password string `mapstructure:"password"` Password string `mapstructure:"password"`
} }
func (p *AIProvider) GetPassword() string {
return p.Password
}
func (p *AIProvider) GetModel() string {
return p.Model
}

View File

@ -17,10 +17,11 @@ type NoOpAIClient struct {
model string model string
} }
func (c *NoOpAIClient) Configure(token string, model string, language string) error { func (c *NoOpAIClient) Configure(config IAIConfig, language string) error {
token := config.GetPassword()
c.language = language c.language = language
c.client = fmt.Sprintf("I am a noop client with the token %s ", token) c.client = fmt.Sprintf("I am a noop client with the token %s ", token)
c.model = model c.model = config.GetModel()
return nil return nil
} }

View File

@ -15,27 +15,22 @@ import (
"github.com/sashabaranov/go-openai" "github.com/sashabaranov/go-openai"
) )
const (
default_prompt = "Simplify the following Kubernetes error message and provide a solution in %s: %s"
prompt_a = "Read the following input %s and provide possible scenarios for remediation in %s"
prompt_b = "Considering the following input from the Kubernetes resource %s and the error message %s, provide possible scenarios for remediation in %s"
prompt_c = "Reading the following %s error message and it's accompanying log message %s, how would you simplify this message?"
)
type OpenAIClient struct { type OpenAIClient struct {
client *openai.Client client *openai.Client
language string language string
model string model string
} }
func (c *OpenAIClient) Configure(token string, model string, language string) error { func (c *OpenAIClient) Configure(config IAIConfig, language string) error {
client := openai.NewClient(token) token := config.GetPassword()
defaultConfig := openai.DefaultConfig(token)
client := openai.NewClientWithConfig(defaultConfig)
if client == nil { if client == nil {
return errors.New("error creating OpenAI client") return errors.New("error creating OpenAI client")
} }
c.language = language c.language = language
c.client = client c.client = client
c.model = model c.model = config.GetModel()
return nil return nil
} }

8
pkg/ai/prompts.go Normal file
View File

@ -0,0 +1,8 @@
package ai
const (
default_prompt = "Simplify the following Kubernetes error message and provide a solution in %s: %s"
prompt_a = "Read the following input %s and provide possible scenarios for remediation in %s"
prompt_b = "Considering the following input from the Kubernetes resource %s and the error message %s, provide possible scenarios for remediation in %s"
prompt_c = "Reading the following %s error message and it's accompanying log message %s, how would you simplify this message?"
)

View File

@ -69,7 +69,7 @@ func NewAnalysis(backend string, language string, filters []string, namespace st
} }
aiClient := ai.NewClient(aiProvider.Name) aiClient := ai.NewClient(aiProvider.Name)
if err := aiClient.Configure(aiProvider.Password, aiProvider.Model, language); err != nil { if err := aiClient.Configure(&aiProvider, language); err != nil {
color.Red("Error: %v", err) color.Red("Error: %v", err)
return nil, err return nil, err
} }