First Working Prototype

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.

pkg/database/database.go

@@ -0,0 +1,139 @@
+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)
+}