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.
140 lines
3.3 KiB
Go
140 lines
3.3 KiB
Go
package database
|
|
|
|
import (
|
|
"ai-code-assistant/pkg/config"
|
|
"context"
|
|
"embed"
|
|
"errors"
|
|
"fmt"
|
|
"github.com/golang-migrate/migrate/v4"
|
|
_ "github.com/golang-migrate/migrate/v4/database/postgres"
|
|
"github.com/golang-migrate/migrate/v4/source/iofs"
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
_ "github.com/lib/pq"
|
|
"strconv"
|
|
)
|
|
|
|
//go:embed migrations
|
|
var migrations embed.FS
|
|
|
|
type contextKey string
|
|
|
|
const (
|
|
contextKeyDB contextKey = "db"
|
|
)
|
|
|
|
func preparedStatements() map[string]string {
|
|
return map[string]string{
|
|
"insert_repo": `INSERT INTO repos (repo_id, repo_hash, repo_path) VALUES ($1, $2, $3) ON CONFLICT DO NOTHING`,
|
|
"repo_from_path": `SELECT repo_id FROM repos WHERE repo_path = $1`,
|
|
"get_chunk": `SELECT document FROM langchain_pg_embedding WHERE JSON_EXTRACT_PATH_TEXT(cmetadata, 'chunk_id')=$1 AND JSON_EXTRACT_PATH_TEXT(cmetadata, 'file_path')=$2 AND JSON_EXTRACT_PATH_TEXT(cmetadata, 'repo_id')=$3`,
|
|
}
|
|
}
|
|
|
|
type Database struct {
|
|
pool *pgxpool.Pool
|
|
}
|
|
|
|
func (db *Database) DB(ctx context.Context) (*pgxpool.Conn, error) {
|
|
conn, err := db.pool.Acquire(ctx)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
for nm, query := range preparedStatements() {
|
|
if _, err := conn.Conn().Prepare(ctx, nm, query); err != nil {
|
|
return nil, fmt.Errorf("problem preparing statement %s: %w", nm, err)
|
|
}
|
|
}
|
|
|
|
return conn, nil
|
|
}
|
|
|
|
func (db *Database) RepoIDFromPath(ctx context.Context, path string) (string, error) {
|
|
conn, err := db.DB(ctx)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
defer conn.Release()
|
|
|
|
var repoID string
|
|
|
|
if err := conn.QueryRow(ctx, "repo_from_path", path).Scan(&repoID); err != nil {
|
|
return "", err
|
|
}
|
|
|
|
return repoID, nil
|
|
}
|
|
|
|
func (db *Database) GetChunk(ctx context.Context, chunkID int, path, repoID string) (string, error) {
|
|
conn, err := db.DB(ctx)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
defer conn.Release()
|
|
|
|
var chunk string
|
|
|
|
chunkIDStr := strconv.FormatInt(int64(chunkID), 10)
|
|
|
|
if err := conn.QueryRow(ctx, "get_chunk", chunkIDStr, path, repoID).Scan(&chunk); err != nil {
|
|
return "", err
|
|
}
|
|
|
|
return chunk, nil
|
|
}
|
|
|
|
func (db *Database) GetChunkContext(ctx context.Context, chunkID, distance int, path, repoID string) (string, error) {
|
|
minChunk := chunkID - distance
|
|
if minChunk < 0 {
|
|
minChunk = 0
|
|
}
|
|
|
|
chunkContext := ""
|
|
|
|
for chunkID := minChunk; chunkID < chunkID+(distance*2); chunkID++ {
|
|
chunk, err := db.GetChunk(ctx, chunkID, path, repoID)
|
|
if err == nil {
|
|
chunkContext += chunk
|
|
} else if !errors.Is(err, pgx.ErrNoRows) {
|
|
return "", err
|
|
} else {
|
|
break
|
|
}
|
|
}
|
|
|
|
return chunkContext, nil
|
|
}
|
|
|
|
func FromConfig(ctx context.Context, cfg *config.Configuration) (*Database, error) {
|
|
migFS, err := iofs.New(migrations, "migrations")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
mig, err := migrate.NewWithSourceInstance("iofs", migFS, cfg.Database.ConnString)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if err := mig.Up(); err != nil && !errors.Is(err, migrate.ErrNoChange) {
|
|
return nil, fmt.Errorf("problem performing database migrations: %w", err)
|
|
}
|
|
|
|
pool, err := pgxpool.New(ctx, cfg.Database.ConnString)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &Database{pool: pool}, nil
|
|
}
|
|
|
|
func FromContext(ctx context.Context) *Database {
|
|
return ctx.Value(contextKeyDB).(*Database)
|
|
}
|
|
|
|
func WrapContext(ctx context.Context, cfg *Database) context.Context {
|
|
return context.WithValue(ctx, contextKeyDB, cfg)
|
|
}
|