Skip to content

Commit

Permalink
client: add llmClient that abstracts away the llm used to generate wo…
Browse files Browse the repository at this point in the history
…rkflows
  • Loading branch information
joshi4 committed Dec 21, 2024
1 parent 9a695f3 commit eee25ba
Show file tree
Hide file tree
Showing 3 changed files with 79 additions and 19 deletions.
32 changes: 14 additions & 18 deletions client/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -76,8 +76,9 @@ type FileInfo struct {
}

type client struct {
cl *http.Client
apiHost string
cl *http.Client
llmClient LLMClient
apiHost string
}

var _ Client = (*client)(nil)
Expand All @@ -97,7 +98,8 @@ func New() (Client, error) {
savvyVersion: config.Version(),
},
},
apiHost: config.APIHost(),
llmClient: NewLLMClient(cfg),
apiHost: config.APIHost(),
}

// validate token as early as possible
Expand Down Expand Up @@ -137,10 +139,14 @@ func (a *AuthorizedRoundTripper) RoundTrip(req *http.Request) (*http.Response, e
// path must start with a slash. e.g. /api/v1/whoami
// apiURL will add a slash if it's missing
func (c *client) apiURL(path string) string {
return genAPIURL(c.apiHost, path)
}

func genAPIURL(host, path string) string {
if !strings.HasPrefix(path, "/") {
path = "/" + path
}
return c.apiHost + path
return host + path
}

func (c *client) WhoAmI(ctx context.Context) (string, error) {
Expand Down Expand Up @@ -200,26 +206,16 @@ func (rb *Runbook) Commands() []string {
}

func (c *client) GenerateRunbookV2(ctx context.Context, commands []RecordedCommand) (*GeneratedRunbook, error) {
cl := c.cl
bs, err := json.Marshal(struct{ Commands []RecordedCommand }{Commands: commands})
if err != nil {
return nil, err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.apiURL("/api/v1/generate_runbookv2"), bytes.NewReader(bs))
generatedRunbook, err := c.llmClient.GenerateRunbook(ctx, commands)
if err != nil {
return nil, err
}
resp, err := cl.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()

var generatedRunbook GeneratedRunbook
if err := json.NewDecoder(resp.Body).Decode(&generatedRunbook); err != nil {
savedRunbook, err := c.SaveRunbook(ctx, generatedRunbook)
if err != nil {
return nil, err
}
return &generatedRunbook, nil
return savedRunbook, nil
}

func (c *client) SaveRunbook(ctx context.Context, runbook *Runbook) (*GeneratedRunbook, error) {
Expand Down
2 changes: 1 addition & 1 deletion client/guest.go
Original file line number Diff line number Diff line change
Expand Up @@ -153,7 +153,7 @@ func GetLoggedInClient() (Client, error) {
}

func getLoggedInClient() (Client, error) {
cl, err := New()
cl, err := New()
if err == nil {
return cl, nil
}
Expand Down
64 changes: 64 additions & 0 deletions client/llm.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
package client

import (
"bytes"
"context"
"encoding/json"
"net/http"

"github.com/getsavvyinc/savvy-cli/config"
)

type LLMClient interface {
GenerateRunbook(ctx context.Context, commands []RecordedCommand) (*Runbook, error)
}

func NewLLMClient(cfg *config.Config) LLMClient {
if cfg.OpenAIBaseURL == "" {
return &defaultLLMClient{
apiHost: config.APIHost(),
cl: &http.Client{
Transport: &AuthorizedRoundTripper{
token: cfg.Token,
savvyVersion: config.Version(),
},
},
}
}
return &customLLMClient{}
}

type defaultLLMClient struct {
cl *http.Client
apiHost string
}

func (dlc *defaultLLMClient) GenerateRunbook(ctx context.Context, commands []RecordedCommand) (*Runbook, error) {
cl := dlc.cl
bs, err := json.Marshal(struct{ Commands []RecordedCommand }{Commands: commands})
if err != nil {
return nil, err
}

req, err := http.NewRequestWithContext(ctx, http.MethodPost, genAPIURL(dlc.apiHost, "/api/v1/generate_workflow"), bytes.NewReader(bs))
if err != nil {
return nil, err
}
resp, err := cl.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()

var generatedRunbook Runbook
if err := json.NewDecoder(resp.Body).Decode(&generatedRunbook); err != nil {
return nil, err
}
return &generatedRunbook, nil
}

type customLLMClient struct{}

func (c *customLLMClient) GenerateRunbook(ctx context.Context, commands []RecordedCommand) (*Runbook, error) {
return nil, nil
}

0 comments on commit eee25ba

Please sign in to comment.