Skip to content

Commit

Permalink
Minimal implementation of the code to generate completions.
Browse files Browse the repository at this point in the history
* This is a pretty minimal implementation.
* The entire response from the AI is added; i.e. we potentially
  add multiple blocks to the document. I'm not sure that's the best UX.

* Implement some OpenAI utilities

  * Add a helper function to initiate OpenAI clients
  * Parse OpenAI errors for common cases like context length exceeded
  • Loading branch information
jlewi committed Apr 8, 2024
1 parent b9194cd commit f9cd8a8
Show file tree
Hide file tree
Showing 12 changed files with 385 additions and 9 deletions.
8 changes: 7 additions & 1 deletion app/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -12,16 +12,20 @@ require (
github.com/go-logr/logr v1.3.0
github.com/go-logr/zapr v1.3.0
github.com/google/go-cmp v0.6.0
github.com/gregjones/httpcache v0.0.0-20190611155906-901d90724c79
github.com/grpc-ecosystem/grpc-gateway/v2 v2.19.1
github.com/hashicorp/go-retryablehttp v0.7.5
github.com/jlewi/foyle/protos/go v0.0.0-00010101000000-000000000000
github.com/jlewi/hydros v0.0.6
github.com/jlewi/monogo v0.0.0-20240123191147-401afe194d74
github.com/pkg/errors v0.9.1
github.com/sashabaranov/go-openai v1.20.4
github.com/spf13/cobra v1.8.0
github.com/spf13/viper v1.18.2
github.com/timtadh/lexmachine v0.2.3
github.com/yuin/goldmark v1.4.13
go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.46.1
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.46.1
go.uber.org/zap v1.27.0
google.golang.org/grpc v1.62.1
gopkg.in/yaml.v3 v3.0.1
Expand Down Expand Up @@ -74,13 +78,15 @@ require (
github.com/gogo/protobuf v1.3.2 // indirect
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
github.com/golang/protobuf v1.5.3 // indirect
github.com/google/btree v1.0.1 // indirect
github.com/google/gnostic v0.6.9 // indirect
github.com/google/go-containerregistry v0.18.0 // indirect
github.com/google/gofuzz v1.2.0 // indirect
github.com/google/s2a-go v0.1.7 // indirect
github.com/google/uuid v1.6.0 // indirect
github.com/googleapis/enterprise-certificate-proxy v0.3.2 // indirect
github.com/googleapis/gax-go/v2 v2.12.0 // indirect
github.com/hashicorp/go-cleanhttp v0.5.2 // indirect
github.com/hashicorp/hcl v1.0.0 // indirect
github.com/imdario/mergo v0.3.13 // indirect
github.com/inconshreveable/mousetrap v1.1.0 // indirect
Expand All @@ -104,6 +110,7 @@ require (
github.com/opencontainers/go-digest v1.0.0 // indirect
github.com/opencontainers/image-spec v1.1.0-rc3 // indirect
github.com/pelletier/go-toml/v2 v2.2.0 // indirect
github.com/peterbourgon/diskv v2.0.1+incompatible // indirect
github.com/pjbgf/sha1cd v0.3.0 // indirect
github.com/rivo/uniseg v0.4.2 // indirect
github.com/sagikazarmark/locafero v0.4.0 // indirect
Expand All @@ -123,7 +130,6 @@ require (
github.com/xanzy/ssh-agent v0.3.3 // indirect
github.com/xlab/treeprint v1.1.0 // indirect
go.opencensus.io v0.24.0 // indirect
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.46.1 // indirect
go.opentelemetry.io/otel v1.21.0 // indirect
go.opentelemetry.io/otel/metric v1.21.0 // indirect
go.opentelemetry.io/otel/trace v1.21.0 // indirect
Expand Down
19 changes: 19 additions & 0 deletions app/go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,8 @@ github.com/envoyproxy/go-control-plane v0.9.9-0.20210512163311-63b5d3c536b0/go.m
github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c=
github.com/envoyproxy/protoc-gen-validate v1.0.4 h1:gVPz/FMfvh57HdSJQyvBtF00j8JU4zdyUgIUNhlgg0A=
github.com/envoyproxy/protoc-gen-validate v1.0.4/go.mod h1:qys6tmnRsYrQqIhm2bvKZH4Blx/1gTIZ2UKVY1M+Yew=
github.com/fatih/color v1.14.1 h1:qfhVLaG5s+nCROl1zJsZRxFeYrHLqWroPOQ8BWiNb4w=
github.com/fatih/color v1.14.1/go.mod h1:2oHN61fhTpgcxD3TSWCgKDiH1+x4OiDVVGH8WlgGZGg=
github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg=
github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
github.com/flowstack/go-jsonschema v0.1.1/go.mod h1:yL7fNggx1o8rm9RlgXv7hTBWxdBM0rVwpMwimd3F3N0=
Expand Down Expand Up @@ -174,6 +176,8 @@ github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaS
github.com/golang/protobuf v1.5.2/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY=
github.com/golang/protobuf v1.5.3 h1:KhyjKVUg7Usr/dYsdSqoFveMYd5ko72D+zANwlG1mmg=
github.com/golang/protobuf v1.5.3/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY=
github.com/google/btree v1.0.1 h1:gK4Kx5IaGY9CD5sPJ36FHiBJ6ZXl0kilRiiCj+jdYp4=
github.com/google/btree v1.0.1/go.mod h1:xXMiIv4Fb/0kKde4SpL7qlzvu5cMJDRkFDxJfI9uaxA=
github.com/google/gnostic v0.6.9 h1:ZK/5VhkoX835RikCHpSUJV9a+S3e1zLh59YnyWeBW+0=
github.com/google/gnostic v0.6.9/go.mod h1:Nm8234We1lq6iB9OmlgNv3nH91XLLVZHCDayfA3xq+E=
github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M=
Expand Down Expand Up @@ -202,9 +206,18 @@ github.com/googleapis/enterprise-certificate-proxy v0.3.2 h1:Vie5ybvEvT75RniqhfF
github.com/googleapis/enterprise-certificate-proxy v0.3.2/go.mod h1:VLSiSSBs/ksPL8kq3OBOQ6WRI2QnaFynd1DCjZ62+V0=
github.com/googleapis/gax-go/v2 v2.12.0 h1:A+gCJKdRfqXkr+BIRGtZLibNXf0m1f9E4HG56etFpas=
github.com/googleapis/gax-go/v2 v2.12.0/go.mod h1:y+aIqrI5eb1YGMVJfuV3185Ts/D7qKpsEkdD5+I6QGU=
github.com/gregjones/httpcache v0.0.0-20190611155906-901d90724c79 h1:+ngKgrYPPJrOjhax5N+uePQ0Fh1Z7PheYoUI/0nzkPA=
github.com/gregjones/httpcache v0.0.0-20190611155906-901d90724c79/go.mod h1:FecbI9+v66THATjSRHfNgh1IVFe/9kFxbXtjV0ctIMA=
github.com/grpc-ecosystem/grpc-gateway v1.16.0/go.mod h1:BDjrQk3hbvj6Nolgz8mAMFbcEtjT1g+wF4CSlocrBnw=
github.com/grpc-ecosystem/grpc-gateway/v2 v2.19.1 h1:/c3QmbOGMGTOumP2iT/rCwB7b0QDGLKzqOmktBjT+Is=
github.com/grpc-ecosystem/grpc-gateway/v2 v2.19.1/go.mod h1:5SN9VR2LTsRFsrEC6FHgRbTWrTHu6tqPeKxEQv15giM=
github.com/hashicorp/go-cleanhttp v0.5.2 h1:035FKYIWjmULyFRBKPs8TBQoi0x6d9G4xc9neXJWAZQ=
github.com/hashicorp/go-cleanhttp v0.5.2/go.mod h1:kO/YDlP8L1346E6Sodw+PrpBSV4/SoxCXGY6BqNFT48=
github.com/hashicorp/go-hclog v0.9.2/go.mod h1:5CU+agLiy3J7N7QjHK5d05KxGsuXiQLrjA0H7acj2lQ=
github.com/hashicorp/go-hclog v1.5.0 h1:bI2ocEMgcVlz55Oj1xZNBsVi900c7II+fWDyV9o+13c=
github.com/hashicorp/go-hclog v1.5.0/go.mod h1:W4Qnvbt70Wk/zYJryRzDRU/4r0kIg0PVHBcfoyhpF5M=
github.com/hashicorp/go-retryablehttp v0.7.5 h1:bJj+Pj19UZMIweq/iie+1u5YCdGrnxCT9yvm0e+Nd5M=
github.com/hashicorp/go-retryablehttp v0.7.5/go.mod h1:Jy/gPYAdjqffZ/yFGCFV2doI5wjtH1ewM9u8iYVjtX8=
github.com/hashicorp/hcl v1.0.0 h1:0Anlzjpi4vEasTeNFn2mLJgTSwt0+6sfsiTG8qcWGx4=
github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ=
github.com/imdario/mergo v0.3.13 h1:lFzP57bqS/wsqKssCGmtLAb8A0wKjLGrve2q3PPVcBk=
Expand Down Expand Up @@ -253,6 +266,8 @@ github.com/matryer/is v1.2.0 h1:92UTHpy8CDwaJ08GqLDzhhuixiBUUD1p3AU6PHddz4A=
github.com/matryer/is v1.2.0/go.mod h1:2fLPjFQM9rhQ15aVEtbuwhJinnOqrmgXPNdZsdwlWXA=
github.com/matryer/try v0.0.0-20161228173917-9ac251b645a2 h1:JAEbJn3j/FrhdWA9jW8B5ajsLIjeuEHLi8xE4fk997o=
github.com/matryer/try v0.0.0-20161228173917-9ac251b645a2/go.mod h1:0KeJpeMD6o+O4hW7qJOT7vyQPKrWmj26uf5wMc/IiIs=
github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
github.com/mattn/go-runewidth v0.0.13 h1:lTGmDsbAYt5DmK6OnoV7EuIF1wEIFAcxld6ypU4OSgU=
Expand All @@ -276,6 +291,8 @@ github.com/opencontainers/image-spec v1.1.0-rc3 h1:fzg1mXZFj8YdPeNkRXMg+zb88BFV0
github.com/opencontainers/image-spec v1.1.0-rc3/go.mod h1:X4pATf0uXsnn3g5aiGIsVnJBR4mxhKzfwmvK/B2NTm8=
github.com/pelletier/go-toml/v2 v2.2.0 h1:QLgLl2yMN7N+ruc31VynXs1vhMZa7CeHHejIeBAsoHo=
github.com/pelletier/go-toml/v2 v2.2.0/go.mod h1:1t835xjRzz80PqgE6HHgN2JOsmgYu/h4qDAS4n929Rs=
github.com/peterbourgon/diskv v2.0.1+incompatible h1:UBdAOUP5p4RWqPBg048CAvpKN+vxiaj6gdUUzhl4XmI=
github.com/peterbourgon/diskv v2.0.1+incompatible/go.mod h1:uqqh8zWWbv1HBMNONnaR/tNboyR3/BZd58JJSHlUSCU=
github.com/pjbgf/sha1cd v0.3.0 h1:4D5XXmUUBUl/xQ6IjCkEAbqXskkq/4O7LmGn0AqMDs4=
github.com/pjbgf/sha1cd v0.3.0/go.mod h1:nZ1rrWOcGJ5uZgEEVL1VUM9iRQiZvWdbZjkKyFzPPsI=
github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
Expand All @@ -295,6 +312,8 @@ github.com/sagikazarmark/locafero v0.4.0 h1:HApY1R9zGo4DBgr7dqsTH/JJxLTTsOt7u6ke
github.com/sagikazarmark/locafero v0.4.0/go.mod h1:Pe1W6UlPYUk/+wc/6KFhbORCfqzgYEpgQ3O5fPuL3H4=
github.com/sagikazarmark/slog-shim v0.1.0 h1:diDBnUNK9N/354PgrxMywXnAwEr1QZcOr6gto+ugjYE=
github.com/sagikazarmark/slog-shim v0.1.0/go.mod h1:SrcSrq8aKtyuqEI1uvTDTK1arOWRIczQRv+GVI1AkeQ=
github.com/sashabaranov/go-openai v1.20.4 h1:095xQ/fAtRa0+Rj21sezVJABgKfGPNbyx/sAN/hJUmg=
github.com/sashabaranov/go-openai v1.20.4/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
github.com/sergi/go-diff v1.1.0/go.mod h1:STckp+ISIX8hZLjrqAeVduY0gWCT9IjLuqbuNXdaHfM=
github.com/sergi/go-diff v1.2.0 h1:XU+rvMAioB0UC3q1MFrIQy4Vo5/4VsRDQQXHsEya6xQ=
github.com/sergi/go-diff v1.2.0/go.mod h1:STckp+ISIX8hZLjrqAeVduY0gWCT9IjLuqbuNXdaHfM=
Expand Down
126 changes: 119 additions & 7 deletions app/pkg/agent/agent.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,23 +2,135 @@ package agent

import (
"context"
"strings"

"github.com/jlewi/foyle/app/pkg/config"
"github.com/jlewi/foyle/app/pkg/docs"
"github.com/jlewi/foyle/app/pkg/logs"
"github.com/jlewi/foyle/app/pkg/oai"
"github.com/jlewi/foyle/protos/go/foyle/v1alpha1"
"github.com/pkg/errors"
"github.com/sashabaranov/go-openai"
)

const (
maxTries = 3
// MaxDocChars is an upper limit for the number of characters to include in prompts to avoid hitting
// OpenAI's context length limits. This can be an upper bound because if we get a context length exceeded
// error the code will automatically try to shrink the document even further.
// We use the heuristic 1 token ~ 2 characters
// We are currently using GPT3.5 which has a context window of 16385 tokens.
// (https://platform.openai.com/docs/models/gpt-3-5-turbo)
// If we use 50% of that's 16000 characters.
MaxDocChars = 16000
temperature = 0.9
)

// Agent is the agent.
type Agent struct {
v1alpha1.UnimplementedGenerateServiceServer
client *openai.Client
config config.Config
}

func NewAgent(cfg config.Config, client *openai.Client) (*Agent, error) {
return &Agent{
client: client,
config: cfg,
}, nil
}
func (a *Agent) Generate(ctx context.Context, req *v1alpha1.GenerateRequest) (*v1alpha1.GenerateResponse, error) {
blocks, err := a.completeWithRetries(ctx, req)

func (e *Agent) Generate(context.Context, *v1alpha1.GenerateRequest) (*v1alpha1.GenerateResponse, error) {
if err != nil {
// TODO(jeremy): Should we set a status code?
return nil, err
}
resp := &v1alpha1.GenerateResponse{
Blocks: []*v1alpha1.Block{
{
Kind: v1alpha1.BlockKind_MARKUP,
Contents: "Hello From The Foyle Server! Your generate request was recieved.",
},
},
Blocks: blocks,
}
return resp, nil
}

func (a *Agent) completeWithRetries(ctx context.Context, req *v1alpha1.GenerateRequest) ([]*v1alpha1.Block, error) {
log := logs.FromContext(ctx)

t := docs.NewTailer(req.Doc.GetBlocks(), MaxDocChars)
for try := 0; try < maxTries; try++ {

args := promptArgs{
Document: t.Text(),
}

var sb strings.Builder
if err := promptTemplate.Execute(&sb, args); err != nil {
return nil, errors.Wrapf(err, "Failed to execute prompt template")
}

messages := []openai.ChatCompletionMessage{
{Role: openai.ChatMessageRoleSystem,
Content: systemPrompt,
},
{Role: openai.ChatMessageRoleUser,
Content: sb.String(),
},
}
request := openai.ChatCompletionRequest{
Model: openai.GPT3Dot5Turbo0125,
Messages: messages,
MaxTokens: 2000,
Temperature: temperature,
}

log.Info("OpenAI:CreateChatCompletion", "request", request)
resp, err := a.client.CreateChatCompletion(ctx, request)

if err != nil {
if oai.ErrorIs(err, oai.ContextLengthExceededCode) {
log.Info("OpenAI:ContextLengthExceeded", "err", err)
if !t.Shorten() {
return nil, errors.Wrapf(err, "the document can't be shortened any further to fit within the context window")
}
continue
}
// TODO(jeremy): Should we surface the error to the user as blocks in the notebook
return nil, errors.Wrapf(err, "CreateChatCompletion failed")
}

log.Info("OpenAI:CreateChatCompletion response", "resp", resp)

blocks, err := a.parseResponse(ctx, &resp)
if err != nil {
return nil, errors.Wrapf(err, "Failed to parse response")
}

return blocks, nil
}
err := errors.Errorf("Failed to generate a chat completion after %d tries", maxTries)
log.Error(err, "Failed to generate a chat completion", "maxTries", maxTries)
return nil, err
}

func (a *Agent) parseResponse(ctx context.Context, resp *openai.ChatCompletionResponse) ([]*v1alpha1.Block, error) {
log := logs.FromContext(ctx)
allBlocks := make([]*v1alpha1.Block, 0, 10)
for _, choice := range resp.Choices {
if choice.Message.Content == "" {
continue
}

blocks, err := docs.MarkdownToBlocks(choice.Message.Content)
if err != nil {
log.Error(err, "Failed to parse markdown to blocks", "markdown", choice.Message.Content)
b := &v1alpha1.Block{
Kind: v1alpha1.BlockKind_MARKUP,
Contents: choice.Message.Content,
}
allBlocks = append(allBlocks, b)
continue
}

allBlocks = append(allBlocks, blocks...)
}
return allBlocks, nil
}
57 changes: 57 additions & 0 deletions app/pkg/agent/agent_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
package agent

import (
"context"
"os"
"testing"

"github.com/jlewi/foyle/app/pkg/config"
"github.com/jlewi/foyle/app/pkg/oai"
"github.com/jlewi/foyle/protos/go/foyle/v1alpha1"
"go.uber.org/zap"
)

func Test_Generate(t *testing.T) {
if os.Getenv("GITHUB_ACTIONS") != "" {
t.Skipf("Test is skipped in GitHub actions")
}

if err := config.InitViper(nil); err != nil {
t.Fatalf("Failed to initialize viper: %v", err)
}
cfg := config.GetConfig()

// Setup logs
c := zap.NewDevelopmentConfig()
log, err := c.Build()
if err != nil {
t.Fatalf("Error creating logger; %v", err)
}
zap.ReplaceGlobals(log)

client, err := oai.NewClient(*cfg)
if err != nil {
t.Fatalf("Error creating OpenAI client; %v", err)
}
agent, err := NewAgent(*cfg, client)

if err != nil {
t.Fatalf("Error creating agent; %v", err)
}

req := &v1alpha1.GenerateRequest{
Doc: &v1alpha1.Doc{
Blocks: []*v1alpha1.Block{
{
Contents: "Use gcloud to list all the cloud build jobs in project foyle",
},
},
},
}
resp, err := agent.Generate(context.Background(), req)
if err != nil {
t.Fatalf("Error generating; %v", err)
}

t.Logf("Response: %+v", resp)
}
30 changes: 30 additions & 0 deletions app/pkg/agent/prompt.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
package agent

import (
_ "embed"
"text/template"
)

const (
systemPrompt = `You are a helpful AI assistant for software developers. You are helping software engineers write markdown documents to deploy
and operate software. Your job is to help users reason about problems and tasks and come up with the appropriate
commands to accomplish them. You should never try to execute commands. You should always tell the user
to execute the commands themselves. To help the user place the commands inside a code block with the language set to
bash. Users can then execute the commands inside VSCode notebooks. The output will then be appended to the document.
You can then use that output to reason about the next steps.
You are only helping users with tasks related to building, deploying, and operating software. You should interpret
any questions or commands in that context.
`
)

//go:embed prompt.tmpl
var promptTemplateString string

var (
promptTemplate = template.Must(template.New("prompt").Parse(promptTemplateString))
)

type promptArgs struct {
Document string
}
12 changes: 12 additions & 0 deletions app/pkg/agent/prompt.tmpl
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
Please continue writing this markdown document to deal with any tasks or issues listed
in the document. The document is a markdown document. It will contain a description of the task
or problem, I need your help with. It will then contain one or more code blocks containing commands
to be executed to accomplish the task or obtain information needed to figure out the problem.
If a command has already been executed the output of the command will be provided in a code block
with the language `output`. Use the output to help you figure out the problem or complete the task.
If you need me to execute a command please provide the command in a code block and I will execute it
and then add the output to the document.

Here's the document:

{{.Document}}
6 changes: 6 additions & 0 deletions app/pkg/config/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ type Config struct {
Logging Logging `json:"logging" yaml:"logging"`
Server ServerConfig `json:"server" yaml:"server"`
Assets AssetConfig `json:"assets" yaml:"assets"`
OpenAI OpenAIConfig `json:"openai" yaml:"openai"`
}

// ServerConfig configures the server
Expand All @@ -65,6 +66,11 @@ type ServerConfig struct {
HttpMaxWriteTimeout time.Duration `json:"httpMaxWriteTimeout" yaml:"httpMaxWriteTimeout"`
}

type OpenAIConfig struct {
// APIKeyFile is the path to the file containing the API key
APIKeyFile string `json:"apiKeyFile" yaml:"apiKeyFile"`
}

type CorsConfig struct {
// AllowedOrigins is a list of origins allowed to make cross-origin requests.
AllowedOrigins []string `json:"allowedOrigins" yaml:"allowedOrigins"`
Expand Down
1 change: 1 addition & 0 deletions app/pkg/docs/converters.go
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,7 @@ func MarkdownToBlocks(mdText string) ([]*v1alpha1.Block, error) {
if err := renderer.Render(&sb, source, node); err != nil {
return ast.WalkStop, err
}

newBlock := &v1alpha1.Block{
Kind: v1alpha1.BlockKind_MARKUP,
Contents: sb.String(),
Expand Down
Loading

0 comments on commit f9cd8a8

Please sign in to comment.