From 1621023958139584001762524f76ef0cb094c6e0 Mon Sep 17 00:00:00 2001 From: Michael Powers Date: Sun, 20 Apr 2025 20:22:57 -0400 Subject: [PATCH] Finished code cleanup, readme is mostly done. --- README.md | 161 ++++++++++++++++++++++++++++++++----- cmd/autopatch/autopatch.go | 14 +++- cmd/main.go | 4 + pkg/config/config.go | 1 + pkg/llm/llm.go | 12 ++- pkg/llm/relevent_docs.go | 4 + 6 files changed, 171 insertions(+), 25 deletions(-) diff --git a/README.md b/README.md index 0ac232b..531a13a 100644 --- a/README.md +++ b/README.md @@ -4,6 +4,34 @@ This repository contains a proof-of-concept of an AI coding assistant capable of examining the repository looking for the most appropriate place to update, and then generating code appropriate to the prompt, and then a unit test and commit message. +## Solutions Approach +The following approach is taken by this implementation: + * The repository is scanned, chunked, and embeddings are generated for all supported source code. + * The user provided prompt is also embedded and a vector search is performed to find the most appropriate chunk of code to modify. + * The identified chunk and surrounding context is provided to a code-appropriate LLM and the LLM is asked to write code appropriate to solve the prompt. + * The generated code is extracted from the LLM response and a diff patch is produced, the diff is applied to the source file. + * The context of the generated code and original prompt are provided to the code LLM and its asked to generate an appropriate unit test. + * The unit test is either appended to an existing test file or a new test file is generated. + * The changed files are staged and committed to Git after generating an appropriate Git commit message using a chat appropriate LLM. + +This implementation is very simple and has numerous drawbacks (see limitations / assumptions section below) but it does +work and could be iterated on to produce better results. + +## Code Structure + +| Package | Description | +|-------------------------|-------------------------------------------------------------------| +| cmd | Main entrypoint | +| cmd/autopatch | Command to generate a git patch automatically from a prompt. | +| cmd/indexer | Command to (re-)index a repository and regenerate embeddings. | +| cmd/chunks | Command to search chunks based on an embeddings prompt. | +| pkg/config | Configuration file definition. | +| pkg/database | Database abstraction, used to perform bootstrapping / migrations. | +| pkg/database/migrations | Database schema. | +| pkg/indexer | Embedding generator for Git repositories. | +| pkg/llm | Abstraction layer over LLM implementations. | +| pkg/llm/prompt | Prompt templates for the LLM. | + ## Requirements To use this application you must provide the following: * A PostgreSQL database server with the pgvector extension installed and configured. @@ -25,6 +53,27 @@ To see what got generated: git log --full-diff -p -n 1 ``` +## Limitations / Assumptions +The following shortcuts have been taken to reduce time to implementation: +* The application does not use an autonomous agentic approach as this would have taken implementing verification tools + for agent executed steps along with robust retry logic. Instead this implementation uses a much simpler rules based + approach. +* This application currently only supports modifying codebases written in Go. There's nothing fundamental about the approach + taken that prevents supporting other languages, but it would have added complexity to the implementation and testing. +* The application currently only supports modifying a single chunk of code. To support multi-chunk editing an autonomous + agentic approach would need to be used that would recursively identify interesting segements of code to patch. This would + have greatly increased the time to implementation. +* No attempt is made to verify the correctness of the code generated. This can be done by making sure the code + compiles, passes its generated test(s), and when run against a linter doesn't generate warnings. If an agentic approach + were used feedback from the tools could be passed back to the code generator to correct the code segments. +* The approach taken to generate the embeddings is very simple. Files are segmented into fixed size chunks and are + embedded. A more robust approach would be to generate an AST for all code in the repository along with an accompanying + reference graph. A coordinating agent would consume the symbols generated to produce a set of modification tasks and + then context from the coordinator could be passed to a patching agent along with the relevant section of code to + modify in order to produce a set of patches. +* No attempt was made to tune the models to the task, there are almost certainly better models to use than the ones + selected. This was the first set of models that produced a workable result. + ## Models Used in Testing The following models were used while developing the application. @@ -34,23 +83,95 @@ The following models were used while developing the application. | gemma2-9b-it | Used for generating code. | | llama3.2 | Used for conversational prompts and generating git commit messages. | -## Limitations / Assumptions -The following shortcuts have been taken to reduce time to implementation: - * The application does not use an autonomous agentic approach as this would have taken implementing verification tools - for agent executed steps along with robust retry logic. Instead this implementation uses a much simpler rules based - approach. - * This application currently only supports modifying codebases written in Go. There's nothing fundamental about the approach - taken that prevents supporting other languages, but it would have added complexity to the implementation and testing. - * The application currently only supports modifying a single chunk of code. To support multi-chunk editing an autonomous - agentic approach would need to be used that would recursively identify interesting segements of code to patch. This would - have greatly increased the time to implementation. - * No attempt is made to verify the correctness of the code generated. This can be done by making sure the code - compiles, passes its generated test(s), and when run against a linter doesn't generate warnings. If an agentic approach - were used feedback from the tools could be passed back to the code generator to correct the code segments. - * The approach taken to generate the embeddings is very simple. Files are segmented into fixed size chunks and are - embedded. A more robust approach would be to generate an AST for all code in the repository along with an accompanying - reference graph. A coordinating agent would consume the symbols generated to produce a set of modification tasks and - then context from the coordinator could be passed to a patching agent along with the relevant section of code to - modify in order to produce a set of patches. - * No attempt was made to tune the models to the task, there are almost certainly better models to use than the ones - selected. This was the first set of models that produced a workable result. \ No newline at end of file +## Run Example +An example of an execution: +``` +git clone https://github.com/andyjessop/simple-go-server + +go run ./cmd auto-patch \ + --repo /path/to/simple-go-server \ + --task "add a new api route to main that responds to GETs to /echo with 'hello world' with the current time appended at the end" \ + --execute +``` + +Log: +```json +{"time":"2025-04-20T10:16:53.742200259-04:00","level":"INFO","msg":"repo already indexed, skipping"} +{"time":"2025-04-20T10:16:53.793235826-04:00","level":"INFO","msg":"found most relevant file chunk","file":"~/Projects/simple-go-server/main.go","start":2048,"end":2461,"score":0.4782574772834778,"id":1} +{"time":"2025-04-20T10:16:55.908966322-04:00","level":"INFO","msg":"applying generated patch to file","file":"~/Projects/simple-go-server/main.go"} +{"time":"2025-04-20T10:16:56.412154285-04:00","level":"INFO","msg":"applying generated unit test to file","file":"~/Projects/simple-go-server/main_test.go"} +{"time":"2025-04-20T10:16:56.758200341-04:00","level":"INFO","msg":"committed changes to git repo","repo":"~/Projects/simple-go-server"} +``` + +Diff Generated: +``` +commit 29c5b5a2a2b71a05b4b782e94d482938cf51ba2b (HEAD -> main) +Author: Michael Powers +Date: Sun Apr 20 10:16:56 2025 -0400 + + "Added new API route to echo 'Hello World' with current time in response to GET requests to /echo" + +diff --git a/main.go b/main.go +index 9acae7f..b475345 100644 +--- a/main.go ++++ b/main.go +@@ -1,3 +1,4 @@ ++ + package main + + import ( +@@ -8,6 +9,7 @@ import ( + "net/http" + "strconv" + "sync" ++ "time" + ) + + type Post struct { +@@ -24,6 +26,7 @@ var ( + func main() { + http.HandleFunc("/posts", postsHandler) + http.HandleFunc("/posts/", postHandler) ++ http.HandleFunc("/echo", echoHandler) + + fmt.Println("Server is running at http://localhost:8080") + log.Fatal(http.ListenAndServe(":8080", nil)) +@@ -121,3 +124,7 @@ func handleDeletePost(w http.ResponseWriter, r *http.Request, id int) { + delete(posts, id) + w.WriteHeader(http.StatusOK) + } ++ ++func echoHandler(w http.ResponseWriter, r *http.Request) { ++ w.Write([]byte("hello world " + time.Now().Format(time.RFC3339))) ++} +diff --git a/main_test.go b/main_test.go +index 4ca1d78..a051498 100644 +--- a/main_test.go ++++ b/main_test.go +@@ -81,3 +81,25 @@ func TestHandlePostAndDeletePosts(t *testing.T) { + t.Errorf("handler returned wrong status code for delete: got %v want %v", status, http.StatusOK) + } + } ++ ++ ++func TestEchoHandler(t *testing.T) { ++ req, err := http.NewRequest("GET", "/echo", nil) ++ if err != nil { ++ t.Fatal(err) ++ } ++ ++ rr := httptest.NewRecorder() ++ echoHandler(rr, req) ++ ++ if rr.Code != http.StatusOK { ++ t.Errorf("handler returned wrong status code: got %v want %v", ++ rr.Code, http.StatusOK) ++ } ++ ++ expected := "hello world " + time.Now().Format(time.RFC3339) ++ if rr.Body.String() != expected { ++ t.Errorf("handler returned unexpected body: got %v want %v", ++ rr.Body.String(), expected) ++ } ++} +``` diff --git a/cmd/autopatch/autopatch.go b/cmd/autopatch/autopatch.go index c111ada..caf5d67 100644 --- a/cmd/autopatch/autopatch.go +++ b/cmd/autopatch/autopatch.go @@ -21,6 +21,7 @@ import ( "strings" ) +// Command defines the autopatch command along with command line flags. func Command() *cli.Command { return &cli.Command{ Name: "auto-patch", @@ -34,7 +35,7 @@ func Command() *cli.Command { }, &cli.StringFlag{ Name: "task", - Usage: "task to perform, e.g. \"add a test for a function\"", + Usage: "task to perform, e.g. \"add a http route for a new health check endpoint\"", Required: true, }, &cli.BoolFlag{ @@ -45,11 +46,13 @@ func Command() *cli.Command { } } +// autoPatch is a struct implementing the auto patcher. type autoPatch struct { llm *llm.LLM execute bool } +// run gets executed when the command is run. func (a *autoPatch) run(ctx context.Context, cmd *cli.Command) error { llmRef := llm.FromContext(ctx) a.llm = llmRef @@ -70,6 +73,8 @@ func (a *autoPatch) run(ctx context.Context, cmd *cli.Command) error { return nil } +// generateGitCommit will generate the code patch, the unit test, and will stage and commit the changes with an +// appropriate commit message. func (a *autoPatch) generateGitCommit(ctx context.Context, repoPath, prompt string) error { var affectedFiles []string @@ -97,6 +102,7 @@ func (a *autoPatch) generateGitCommit(ctx context.Context, repoPath, prompt stri return nil } +// commit stages the changed files and commits them with a commit message. func (a *autoPatch) commit(ctx context.Context, prompt, repoPath string, files ...string) error { gitPath := osfs.New(filepath.Join(repoPath, ".git")) @@ -146,6 +152,7 @@ func (a *autoPatch) commit(ctx context.Context, prompt, repoPath string, files . return nil } +// generateCodePatch generates an appropriate code patch given a repository and prompt. func (a *autoPatch) generateCodePatch(ctx context.Context, repoPath, prompt string) (string, string, error) { db := database.FromContext(ctx) cfg := config.FromContext(ctx) @@ -209,6 +216,8 @@ func (a *autoPatch) generateCodePatch(ctx context.Context, repoPath, prompt stri return fileName, codeBlock, err } +// generateUnitTest passes the new code and a prompt to the LLM to generate an appropriate unit test. It will add the +// new test to the bottom of an existing test file, or generate a new one if no unit test file already exists. func (a *autoPatch) generateUnitTest(ctx context.Context, prompt, fileName, newCode string) (string, error) { // Check to see if a test file for this already exists. testFileExists := false @@ -263,6 +272,8 @@ func (a *autoPatch) generateUnitTest(ctx context.Context, prompt, fileName, newC return testFile, nil } +// generateCode takes a prompt and tries to extract code from it. The prompt should try to get the LLM to structure the +// code as a single yaml chunk. func (a *autoPatch) generateCode(ctx context.Context, prompt string) (string, error) { rsp, err := a.llm.CodePrompt(ctx, prompt) if err != nil { @@ -284,6 +295,7 @@ func (a *autoPatch) generateCode(ctx context.Context, prompt string) (string, er return codeBlock, nil } +// patchFile attempts to write the merged diff to the original file so it can be staged and committed. func patchFile(fileName string, diffs []diffmatchpatch.Diff) error { var buff bytes.Buffer for _, diff := range diffs { diff --git a/cmd/main.go b/cmd/main.go index d0a3474..e341a4f 100644 --- a/cmd/main.go +++ b/cmd/main.go @@ -16,6 +16,9 @@ import ( "os" ) +// main is the entry point to the application. It's mostly responsible for bootstrapping the database, configuration, +// logging, and llms. It passes everything to subcommands by injecting state into the context to avoid coupling during +// initialization. func main() { app := &cli.Command{ Name: "ai-coding-assistant", @@ -60,6 +63,7 @@ func main() { } } +// readConfig attempts to read the configuration file. func readConfig(ctx context.Context, cmd *cli.Command) (context.Context, error) { cfgFile := cmd.String("config") cfgHandle, err := os.Open(cfgFile) diff --git a/pkg/config/config.go b/pkg/config/config.go index 40f7625..6b9eacb 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -15,6 +15,7 @@ const ( contextKeyConfig contextKey = "config" ) +// Configuration is a simple configuration that can be loaded from a YAML file. type Configuration struct { Database struct { ConnString string `yaml:"conn_string"` diff --git a/pkg/llm/llm.go b/pkg/llm/llm.go index cf0486c..117be46 100644 --- a/pkg/llm/llm.go +++ b/pkg/llm/llm.go @@ -18,12 +18,14 @@ const contextKeyLLM contextKey = "llm" //go:embed prompts var prompts embed.FS +// LLM is responsible for abstracting the configuration and implementations of the LLMs used. type LLM struct { code llms.Model chat llms.Model embedder embeddings.Embedder } +// FromConfig bootstraps the LLM from a passed in configuration. func FromConfig(cfg *config.Configuration) (*LLM, error) { embedLLM, err := cfg.Embedding.GetEmbedding() if err != nil { @@ -52,30 +54,32 @@ func FromConfig(cfg *config.Configuration) (*LLM, error) { }, nil } +// FromContext retrieves an LLM from a passed in context wrapped with WrapContext. func FromContext(ctx context.Context) *LLM { return ctx.Value(contextKeyLLM).(*LLM) } +// WrapContext embeds an LLM inside a context so it can be retrieved with FromContext. func WrapContext(ctx context.Context, llmRef *LLM) context.Context { return context.WithValue(ctx, contextKeyLLM, llmRef) } -func (llm *LLM) GetEmbedding(ctx context.Context, texts ...string) ([][]float32, error) { - return llm.embedder.EmbedDocuments(ctx, texts) -} - +// Embedder gets an embedder that can be used to store and retrieve embeddings. func (llm *LLM) Embedder() embeddings.Embedder { return llm.embedder } +// CodePrompt passes a prompt to the code LLM and returns the response. func (llm *LLM) CodePrompt(ctx context.Context, prompt string) (string, error) { return llm.code.Call(ctx, prompt) } +// ChatPrompt passes a prompt to the chat LLM and returns the response. func (llm *LLM) ChatPrompt(ctx context.Context, prompt string) (string, error) { return llm.chat.Call(ctx, prompt) } +// GetPrompt loads a LLM prompt template and injects variables into it. Uses the go template format. func GetPrompt(name string, data any) (string, error) { tmplText, err := prompts.ReadFile("prompts/" + name + ".tmpl") if err != nil { diff --git a/pkg/llm/relevent_docs.go b/pkg/llm/relevent_docs.go index 57dcd09..a8907e9 100644 --- a/pkg/llm/relevent_docs.go +++ b/pkg/llm/relevent_docs.go @@ -9,6 +9,7 @@ import ( "strconv" ) +// RelevantDocs attempts to find the most relevant file chunks based on context from the prompt. type RelevantDocs struct { CallbacksHandler callbacks.Handler db *database.Database @@ -17,6 +18,7 @@ type RelevantDocs struct { size int } +// FileChunkID is a pointer to a repository file chunk that has been indexed. type FileChunkID struct { Name string ChunkID int @@ -26,6 +28,7 @@ type FileChunkID struct { Doc *schema.Document } +// NewGetRelevantDocs creates a new RelevantDocs scanner. func NewGetRelevantDocs(db *database.Database, llm *LLM, repoID string, size int) *RelevantDocs { return &RelevantDocs{ db: db, @@ -35,6 +38,7 @@ func NewGetRelevantDocs(db *database.Database, llm *LLM, repoID string, size int } } +// GetRelevantFileChunks will scan for relevant documents based on a prompt. func (rd *RelevantDocs) GetRelevantFileChunks(ctx context.Context, query string) ([]*FileChunkID, error) { vectorStore, closeFunc, err := rd.db.GetVectorStore(ctx, rd.llm.Embedder()) if err != nil {