Michael Powers
4b8b8132fd
This application is a simple proof of concept demonstrating an agent capable of taking a prompt and generating a patch implementing code satisfying the prompt along with an accompanying unit test.
64 lines
1.4 KiB
Go
64 lines
1.4 KiB
Go
package chunks
|
|
|
|
import (
|
|
"ai-code-assistant/pkg/database"
|
|
"ai-code-assistant/pkg/llm"
|
|
"context"
|
|
"github.com/urfave/cli/v3"
|
|
"log/slog"
|
|
)
|
|
|
|
func Command() *cli.Command {
|
|
return &cli.Command{
|
|
Name: "chunks",
|
|
Usage: "this command tries to find relevant chunks of code in a git repository for a given query",
|
|
Action: (&chunks{}).run,
|
|
Flags: []cli.Flag{
|
|
&cli.StringFlag{
|
|
Name: "repo",
|
|
Usage: "path to git repository",
|
|
Required: true,
|
|
},
|
|
&cli.StringFlag{
|
|
Name: "query",
|
|
Usage: "query to search for",
|
|
Required: true,
|
|
},
|
|
&cli.IntFlag{
|
|
Name: "docs",
|
|
Usage: "number of docs to get",
|
|
Value: 10,
|
|
},
|
|
},
|
|
}
|
|
}
|
|
|
|
type chunks struct {
|
|
}
|
|
|
|
func (c *chunks) run(ctx context.Context, cmd *cli.Command) error {
|
|
db := database.FromContext(ctx)
|
|
llmRef := llm.FromContext(ctx)
|
|
|
|
repoID, err := db.RepoIDFromPath(ctx, cmd.String("repo"))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
relDocs := llm.NewGetRelevantDocs(db, llmRef, repoID, int(cmd.Int("docs")))
|
|
chunks, err := relDocs.GetRelevantFileChunks(ctx, cmd.String("query"))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if err := relDocs.RankChunks(ctx, cmd.String("query"), chunks); err != nil {
|
|
return err
|
|
}
|
|
|
|
for _, chunk := range chunks {
|
|
slog.Info("found relevant chunk", "name", chunk.Name, "start", chunk.Start, "end", chunk.End, "score", chunk.Score, "id", chunk.ChunkID)
|
|
}
|
|
|
|
return nil
|
|
}
|