Start of Readme, More Cleanup

pkg/database/database.go

@@ -17,6 +17,9 @@ import (
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgxpool"
 	_ "github.com/lib/pq"
+	"github.com/tmc/langchaingo/embeddings"
+	"github.com/tmc/langchaingo/vectorstores"
+	"github.com/tmc/langchaingo/vectorstores/pgvector"
 	"path/filepath"
 	"strconv"
 )
@@ -39,11 +42,13 @@ func preparedStatements() map[string]string {
 	}
 }
 
+// Database is a helper that abstracts access to the database as well as handles migrations.
 type Database struct {
 	pool *pgxpool.Pool
 }
 
-func (db *Database) DB(ctx context.Context) (*pgxpool.Conn, error) {
+// db gets a references to the connection pool.
+func (db *Database) db(ctx context.Context) (*pgxpool.Conn, error) {
 	conn, err := db.pool.Acquire(ctx)
 	if err != nil {
 		return nil, err
@@ -58,6 +63,7 @@ func (db *Database) DB(ctx context.Context) (*pgxpool.Conn, error) {
 	return conn, nil
 }
 
+// UpsertRepo will ensure we have a repository reference for this particular repository path and HEAD reference.
 func (db *Database) UpsertRepo(ctx context.Context, repoPath string) (string, bool, error) {
 	gitPath := osfs.New(filepath.Join(repoPath, ".git"))
 
@@ -71,7 +77,7 @@ func (db *Database) UpsertRepo(ctx context.Context, repoPath string) (string, bo
 		return "", false, err
 	}
 
-	conn, err := db.DB(ctx)
+	conn, err := db.db(ctx)
 	if err != nil {
 		return "", false, err
 	}
@@ -94,8 +100,9 @@ func (db *Database) UpsertRepo(ctx context.Context, repoPath string) (string, bo
 	return id, false, nil
 }
 
+// GetChunk will get a specified chunk from the database.
 func (db *Database) GetChunk(ctx context.Context, chunkID int, path, repoID string) (string, error) {
-	conn, err := db.DB(ctx)
+	conn, err := db.db(ctx)
 	if err != nil {
 		return "", err
 	}
@@ -112,6 +119,43 @@ func (db *Database) GetChunk(ctx context.Context, chunkID int, path, repoID stri
 	return chunk, nil
 }
 
+// ClearChunkIndex will clear all embeddings for a particular repo ID.
+func (db *Database) ClearChunkIndex(ctx context.Context, repoID string) error {
+	conn, err := db.db(ctx)
+	if err != nil {
+		return err
+	}
+	defer conn.Release()
+
+	if _, err := conn.Exec(ctx, "clear_chunks_for_repo", repoID); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// GetVectorStore gets a vector store capable of storing and searching embeddings.
+func (db *Database) GetVectorStore(ctx context.Context, embedder embeddings.Embedder) (vectorstores.VectorStore, func(), error) {
+	conn, err := db.db(ctx)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	vectorStore, err := pgvector.New(ctx,
+		pgvector.WithConn(conn),
+		pgvector.WithEmbedder(embedder),
+		pgvector.WithCollectionName("file_chunks"),
+	)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return vectorStore, func() {
+		conn.Release()
+	}, nil
+}
+
+// GetChunkContext will get a specified chunk and surrounding context based on some amount of distance.
 func (db *Database) GetChunkContext(ctx context.Context, chunkID, distance int, path, repoID string) (string, error) {
 	minChunk := chunkID - distance
 	if minChunk < 0 {
@@ -134,6 +178,7 @@ func (db *Database) GetChunkContext(ctx context.Context, chunkID, distance int,
 	return chunkContext, nil
 }
 
+// FromConfig generates a new Database from a passed in configuration. Used for bootstrapping.
 func FromConfig(ctx context.Context, cfg *config.Configuration) (*Database, error) {
 	migFS, err := iofs.New(migrations, "migrations")
 	if err != nil {
@@ -157,10 +202,12 @@ func FromConfig(ctx context.Context, cfg *config.Configuration) (*Database, erro
 	return &Database{pool: pool}, nil
 }
 
+// FromContext will retrieve the Database from a Context where WrapContext was used to embed the database in the context.
 func FromContext(ctx context.Context) *Database {
 	return ctx.Value(contextKeyDB).(*Database)
 }
 
+// WrapContext embeds a Database inside a Context so it can be passed to functions.
 func WrapContext(ctx context.Context, cfg *Database) context.Context {
 	return context.WithValue(ctx, contextKeyDB, cfg)
 }