Skip to content

Commit

Permalink
feat: add configuration interface to support customer providers
Browse files Browse the repository at this point in the history
Signed-off-by: Aris Boutselis <aris.boutselis@senseon.io>
  • Loading branch information
Aris Boutselis committed Apr 20, 2023
1 parent bd4ab0e commit 84a3cc0
Show file tree
Hide file tree
Showing 5 changed files with 31 additions and 14 deletions.
15 changes: 14 additions & 1 deletion pkg/ai/iai.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,12 +5,17 @@ import (
)

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)
Parse(ctx context.Context, prompt []string, nocache bool) (string, error)
GetName() string
}

type IAIConfig interface {
GetPassword() string
GetModel() string
}

func NewClient(provider string) IAI {
switch provider {
case "openai":
Expand All @@ -31,3 +36,11 @@ type AIProvider struct {
Model string `mapstructure:"model"`
Password string `mapstructure:"password"`
}

func (p *AIProvider) GetPassword() string {
return p.Password
}

func (p *AIProvider) GetModel() string {
return p.Model
}
5 changes: 3 additions & 2 deletions pkg/ai/noopai.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,11 @@ type NoOpAIClient struct {
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.client = fmt.Sprintf("I am a noop client with the token %s ", token)
c.model = model
c.model = config.GetModel()
return nil
}

Expand Down
15 changes: 5 additions & 10 deletions pkg/ai/openai.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,27 +15,22 @@ import (
"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 {
client *openai.Client
language string
model string
}

func (c *OpenAIClient) Configure(token string, model string, language string) error {
client := openai.NewClient(token)
func (c *OpenAIClient) Configure(config IAIConfig, language string) error {
token := config.GetPassword()
defaultConfig := openai.DefaultConfig(token)
client := openai.NewClientWithConfig(defaultConfig)
if client == nil {
return errors.New("error creating OpenAI client")
}
c.language = language
c.client = client
c.model = model
c.model = config.GetModel()
return nil
}

Expand Down
8 changes: 8 additions & 0 deletions pkg/ai/prompts.go
Original file line number Diff line number Diff line change
@@ -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?"
)
2 changes: 1 addition & 1 deletion pkg/analysis/analysis.go
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,7 @@ func NewAnalysis(backend string, language string, filters []string, namespace st
}

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)
return nil, err
}
Expand Down

0 comments on commit 84a3cc0

Please sign in to comment.